From c673be410f355a9c3db9008023fb4a849db06710 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Mon, 31 Aug 2026 23:25:36 -0400 Subject: [PATCH 01/24] renovate: bump toolchain/nix, macOS support, dependency refresh (#509) Co-authored-by: Xiangpeng --- .github/workflows/ci.yml | 54 +- .github/workflows/fuzz.yml | 6 +- .github/workflows/prepare-release.yml | 4 +- .github/workflows/publish.yml | 2 +- Cargo.lock | 2201 ++++++++--------- Cargo.toml | 26 +- README.md | 2 +- benchmark/Cargo.toml | 14 +- benchmark/src/inprocess_runner.rs | 23 + benchmark/src/observability.rs | 4 +- dev/README.md | 6 + dev/dev-tools/Cargo.toml | 6 + .../src/components/cache_state_view.rs | 2 +- flake.lock | 12 +- flake.nix | 9 +- src/core/Cargo.toml | 6 +- .../byte_view_array/comparisons.rs | 4 +- .../byte_view_array/serialization.rs | 4 +- src/core/src/liquid_array/raw/fsst_buffer.rs | 3 +- src/datafusion-local/src/tests/mod.rs | 7 + src/datafusion-server/Cargo.toml | 4 +- src/datafusion/src/utils.rs | 1 + 22 files changed, 1160 insertions(+), 1240 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 11e509b35..541591729 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: name: Basic check runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable with: @@ -56,11 +56,29 @@ jobs: - name: Check for unused dependencies run: cargo shear + macos: + name: macOS check + runs-on: macos-latest + steps: + - uses: actions/checkout@v7 + - uses: dtolnay/rust-toolchain@stable + with: + components: clippy + - uses: Swatinem/rust-cache@v2 + with: + shared-key: ci-${{ runner.os }} + save-if: ${{ github.ref == 'refs/heads/main' }} + # dev-tools needs a generated tailwind.css, covered by the Linux jobs. + - name: Run clippy + run: cargo clippy --workspace --exclude dev-tools --all-targets --all-features -- -D warnings + - name: Run unit tests + run: cargo test --workspace --exclude dev-tools + dev_tools: name: Dev Tools runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: cachix/install-nix-action@v31 with: @@ -76,7 +94,7 @@ jobs: name: Unit Test runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 @@ -96,7 +114,7 @@ jobs: - name: Generate code coverage run: cargo llvm-cov --workspace --codecov --output-path codecov.json - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} files: codecov.json @@ -106,7 +124,7 @@ jobs: name: Shuttle Test runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 @@ -122,7 +140,7 @@ jobs: name: Address Sanitizer runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space # Sanitizers can only run on nightly - uses: dtolnay/rust-toolchain@nightly @@ -139,7 +157,7 @@ jobs: name: ClickBench runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - run: sudo apt-get update && sudo apt-get install -y wget @@ -174,7 +192,7 @@ jobs: env RUST_LOG=info cargo run --bin in_process -- --manifest benchmark/clickbench/benchmark_manifest.json --bench-mode liquid --max-memory-mb 256 cargo llvm-cov report --codecov --output-path codecov_clickbench.json - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} files: codecov_clickbench.json @@ -184,7 +202,7 @@ jobs: name: TPC-H runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - run: sudo apt-get update && sudo apt-get install -y wget @@ -214,7 +232,7 @@ jobs: env RUST_LOG=info cargo run --bin in_process -- --manifest benchmark/tpch/manifest.json --bench-mode liquid --max-memory-mb 256 cargo llvm-cov report --codecov --output-path codecov_tpch.json - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} files: codecov_tpch.json @@ -224,7 +242,7 @@ jobs: name: TPC-DS runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - run: sudo apt-get update && sudo apt-get install -y wget @@ -254,7 +272,7 @@ jobs: env RUST_LOG=info cargo run --bin in_process -- --manifest benchmark/tpcds/manifest.json --bench-mode liquid --max-memory-mb 256 cargo llvm-cov report --codecov --output-path codecov_tpcds.json - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} files: codecov_tpcds.json @@ -264,7 +282,7 @@ jobs: name: StackOverflow runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space - uses: dtolnay/rust-toolchain@stable - name: Install system dependencies @@ -276,7 +294,7 @@ jobs: mkdir -p benchmark/stackoverflow/data/dba mkdir -p benchmark/stackoverflow/downloads - name: Cache StackOverflow dataset - uses: actions/cache@v4 + uses: actions/cache@v6 with: path: | benchmark/stackoverflow/data/dba @@ -309,7 +327,7 @@ jobs: env RUST_LOG=info cargo run --bin in_process -- --manifest benchmark/stackoverflow/manifest.dba.json --bench-mode liquid --max-memory-mb 10 cargo llvm-cov report --codecov --output-path codecov_stackoverflow.json - name: Upload coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} files: codecov_stackoverflow.json @@ -322,7 +340,7 @@ jobs: contents: write pull-requests: write steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 with: fetch-depth: 0 token: ${{ secrets.GITHUB_TOKEN }} @@ -395,7 +413,7 @@ jobs: - name: Comment PR with benchmark results if: steps.compare.outputs.COMPARISON_AVAILABLE == 'true' && github.event_name == 'pull_request' - uses: actions/github-script@v7 + uses: actions/github-script@v9 with: script: | const fs = require('fs'); @@ -457,7 +475,7 @@ jobs: name: Run client/server/inprocess examples runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 with: diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index bf4bc824c..9d6513793 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -18,7 +18,7 @@ jobs: target: [fsst_view_fuzz] steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v7 - uses: ./.github/actions/free-disk-space @@ -95,7 +95,7 @@ jobs: } >> "$GITHUB_STEP_SUMMARY" - name: Upload coverage text summary - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 if: always() with: name: fuzz-coverage-text-${{ matrix.target }} @@ -104,7 +104,7 @@ jobs: if-no-files-found: warn - name: Upload fuzz artifacts (crashes, reduces, etc.) - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 if: always() with: name: fuzz-artifacts-${{ matrix.target }} diff --git a/.github/workflows/prepare-release.yml b/.github/workflows/prepare-release.yml index 822ec9dfe..af042bcfe 100644 --- a/.github/workflows/prepare-release.yml +++ b/.github/workflows/prepare-release.yml @@ -30,7 +30,7 @@ jobs: pull-requests: write steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v7 with: fetch-depth: 0 token: ${{ secrets.GITHUB_TOKEN }} @@ -63,7 +63,7 @@ jobs: fi - name: Create Pull Request - uses: peter-evans/create-pull-request@v7 + uses: peter-evans/create-pull-request@v8 with: token: ${{ secrets.PAT_TOKEN }} branch: release/v${{ steps.bump.outputs.new_version }} diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 022a5a3ff..0022c14b1 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -22,7 +22,7 @@ jobs: contents: write steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v7 with: ref: main fetch-depth: 0 diff --git a/Cargo.lock b/Cargo.lock index 20b7f2aeb..d496de5ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -33,9 +33,9 @@ dependencies = [ [[package]] name = "aho-corasick" -version = "1.1.4" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" dependencies = [ "memchr", ] @@ -57,9 +57,9 @@ checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" [[package]] name = "alloc-stdlib" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" dependencies = [ "alloc-no-stdlib", ] @@ -72,9 +72,9 @@ checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" [[package]] name = "android_system_properties" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" dependencies = [ "libc", ] @@ -131,17 +131,17 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "ar_archive_writer" -version = "0.5.1" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eb93bbb63b9c227414f6eb3a0adfddca591a8ce1e9b60661bb08969b87e340b" +checksum = "73cd58deff2140a0a8eae87e417bd01db68a33e148aa93d1e8cd837e55e312b6" dependencies = [ - "object", + "object 0.39.1", ] [[package]] @@ -153,23 +153,17 @@ dependencies = [ "derive_arbitrary", ] -[[package]] -name = "arrayref" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" - [[package]] name = "arrayvec" -version = "0.7.6" +version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" [[package]] name = "arrow" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" +checksum = "6cfdd0833e32a9874d2b55089333ad310c0be208aafa277385ce2461dec90be3" dependencies = [ "arrow-arith", "arrow-array", @@ -188,9 +182,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" +checksum = "0a41203398f0eaa6f7ec8e62c0da742a21abf282c148fc157f6c35c90e29981a" dependencies = [ "arrow-array", "arrow-buffer", @@ -202,9 +196,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" +checksum = "ae33dad492b7df00a217563a7b0ef2874df68a0deea1b1a3acf628152f7f7a69" dependencies = [ "ahash", "arrow-buffer", @@ -221,9 +215,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" +checksum = "b9552f96391c005e6ab449fa941420935e7e062489b12b8b1b08879b2163f5b5" dependencies = [ "bytes", "half", @@ -233,9 +227,9 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" +checksum = "3a8a327c9649f30d8406995f27642b68df354713cca3baaaf100f076f18d5f34" dependencies = [ "arrow-array", "arrow-buffer", @@ -244,7 +238,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.22.1", + "base64", "chrono", "comfy-table", "half", @@ -255,9 +249,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" +checksum = "af0dd6d90d1955e9f9a014c1e563ee8aeffc21909085d25623e1da44d96eca26" dependencies = [ "arrow-array", "arrow-cast", @@ -270,9 +264,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" +checksum = "2b24852db04738907e06c04ea61e42fe7fda962a34513022dc0d0e754fb7976b" dependencies = [ "arrow-buffer", "arrow-schema", @@ -283,9 +277,9 @@ dependencies = [ [[package]] name = "arrow-flight" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28abfe8bf9f124e5fc83b334af4fa58f8d0323ad25312ccb2d1da50178415704" +checksum = "b2dbe34824c639e43136af8f106992792ab456540d54b880bc320a3192502d2e" dependencies = [ "arrow-arith", "arrow-array", @@ -298,7 +292,7 @@ dependencies = [ "arrow-schema", "arrow-select", "arrow-string", - "base64 0.22.1", + "base64", "bytes", "futures", "once_cell", @@ -311,9 +305,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" +checksum = "29a908a11fcfb3fb2f6730f4ac15e367bc644e419155e96238f68cf3adde572b" dependencies = [ "arrow-array", "arrow-buffer", @@ -327,9 +321,9 @@ dependencies = [ [[package]] name = "arrow-json" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" +checksum = "b8a96aed3931c076adee39ec2a40d8219fc7f09e79bcdaca1df16272993e1e14" dependencies = [ "arrow-array", "arrow-buffer", @@ -339,7 +333,7 @@ dependencies = [ "arrow-select", "chrono", "half", - "indexmap 2.14.0", + "indexmap", "itoa", "lexical-core", "memchr", @@ -352,9 +346,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" +checksum = "63a083ec750f5c043f02946b4baf05fcdbb55f4560a3277055caca5cc99f3eb0" dependencies = [ "arrow-array", "arrow-buffer", @@ -365,9 +359,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" +checksum = "514ba0ef0d4c5896202dae736251ce415abb43a950bed570fb7981b8716c0e4c" dependencies = [ "arrow-array", "arrow-buffer", @@ -378,9 +372,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" +checksum = "21ca356ad6425cecb6eb7b28e4f659f1ee7880fbb1a16127de7dd62901efee9e" dependencies = [ "serde", "serde_core", @@ -389,9 +383,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" +checksum = "c58da39eb3d8350ad4a549e5c2bc49284dac554016c69829310350f1731b0aad" dependencies = [ "ahash", "arrow-array", @@ -403,9 +397,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" +checksum = "b6789b388467525e3271326b6b4915666ecfdf5142aef09779445c954b67543c" dependencies = [ "arrow-array", "arrow-buffer", @@ -432,9 +426,9 @@ checksum = "bfdc70193dadb9d7287fa4b633f15f90c876915b31f6af17da307fc59c9859a8" [[package]] name = "async-compression" -version = "0.4.42" +version = "0.4.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e79b3f8a79cccc2898f31920fc69f304859b3bd567490f75ebf51ae1c792a9ac" +checksum = "3976abdc8fe7d1133d43d304afd42abdf5bc3e1319d263d223bde07b5efc4be8" dependencies = [ "compression-codecs", "compression-core", @@ -461,18 +455,18 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "async-trait" -version = "0.1.89" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -508,9 +502,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "axum" @@ -520,7 +514,7 @@ checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ "axum-core", "axum-macros", - "base64 0.22.1", + "base64", "bytes", "form_urlencoded", "futures-util", @@ -600,7 +594,7 @@ checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -612,8 +606,8 @@ dependencies = [ "addr2line", "cfg-if", "libc", - "miniz_oxide", - "object", + "miniz_oxide 0.8.9", + "object 0.37.3", "rustc-demangle", "windows-link", ] @@ -624,12 +618,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d27c3610c36aee21ce8ac510e6224498de4228ad772a171ed65643a24693a5a8" -[[package]] -name = "base64" -version = "0.21.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" - [[package]] name = "base64" version = "0.22.1" @@ -657,18 +645,18 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.11.1" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" dependencies = [ "serde_core", ] [[package]] name = "bitvec" -version = "1.0.1" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c" +checksum = "ddcec3d12c579d40898fe0a9a358a803c23e9c52ca3c425707f81c9436211837" dependencies = [ "funty", "radium", @@ -687,16 +675,15 @@ dependencies = [ [[package]] name = "blake3" -version = "1.8.5" +version = "1.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +checksum = "6d9e454fc11f76977dc803893aff6304ed33d6a26efae8696573bea74baa27ae" dependencies = [ - "arrayref", "arrayvec", "cc", "cfg-if", "constant_time_eq", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", ] [[package]] @@ -710,18 +697,18 @@ dependencies = [ [[package]] name = "block-buffer" -version = "0.12.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ "hybrid-array", ] [[package]] name = "brotli" -version = "8.0.2" +version = "8.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -730,9 +717,9 @@ dependencies = [ [[package]] name = "brotli-decompressor" -version = "5.0.0" +version = "5.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -740,15 +727,15 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.20.2" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytemuck" -version = "1.25.0" +version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" [[package]] name = "byteorder" @@ -758,9 +745,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" dependencies = [ "serde", ] @@ -782,9 +769,9 @@ checksum = "cd17eb909a8c6a894926bfcc3400a4bb0e732f5a57d37b1f14e8b29e329bace8" [[package]] name = "cc" -version = "1.2.61" +version = "1.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" +checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273" dependencies = [ "find-msvc-tools", "jobserver", @@ -817,18 +804,18 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "rand_core 0.10.1", ] @@ -838,7 +825,7 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1f927b07c74ba84c7e5fe4db2baeb3e996ab2688992e39ac68ce3220a677c7e" dependencies = [ - "base64 0.22.1", + "base64", "encoding_rs", ] @@ -895,9 +882,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.1" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" dependencies = [ "clap_builder", "clap_derive", @@ -905,9 +892,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.6.0" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" dependencies = [ "anstream", "anstyle", @@ -918,14 +905,14 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.6.1" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -940,11 +927,20 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "colored" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "combine" -version = "4.6.7" +version = "4.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", "memchr", @@ -998,9 +994,9 @@ dependencies = [ [[package]] name = "console" -version = "0.16.3" +version = "0.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" dependencies = [ "encode_unicode", "libc", @@ -1049,7 +1045,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e42cd5aabba86f128b3763da1fec1491c0f728ce99245062cd49b6f9e6d235b" dependencies = [ "const-serialize 0.7.2", - "const-serialize-macro 0.8.0-alpha.0", + "const-serialize-macro 0.8.0-alpha.1", "serde", ] @@ -1061,18 +1057,18 @@ checksum = "4f160aad86b4343e8d4e261fee9965c3005b2fd6bc117d172ab65948779e4acf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "const-serialize-macro" -version = "0.8.0-alpha.0" +version = "0.8.0-alpha.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42571ed01eb46d2e1adcf99c8ca576f081e46f2623d13500eba70d1d99a4c439" +checksum = "8e4c3b1c2ce89797adff100c510b92c7cf32983fdbc632e703253f2ef516ef56" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1089,9 +1085,9 @@ checksum = "b0664d2867b4a32697dfe655557f5c3b187e9b605b38612a748e5ec99811d160" [[package]] name = "const_for" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9c50fcfdf972929aff202c16b80086aa3cfc6a3a820af714096c58c7c1d0582" +checksum = "988d3bd6bf67b6d7ae2b519c55296fa57411c27c312575e47b2975f8d1d8590b" [[package]] name = "const_format" @@ -1129,6 +1125,12 @@ dependencies = [ "charset", ] +[[package]] +name = "convert_case" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e" + [[package]] name = "convert_case" version = "0.8.0" @@ -1149,9 +1151,9 @@ dependencies = [ [[package]] name = "cookie" -version = "0.18.1" +version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" +checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87" dependencies = [ "percent-encoding", "time", @@ -1202,17 +1204,11 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core_detect" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f8f80099a98041a3d1622845c271458a2d73e688351bf3cb999266764b81d48" - [[package]] name = "corosensei" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c54787b605c7df106ceccf798df23da4f2e09918defad66705d1cedf3bb914f" +checksum = "6886a0c0f263965933c438626e7179139a62b978a33aa18281cbf0cd5a975f34" dependencies = [ "autocfg", "cfg-if", @@ -1241,45 +1237,45 @@ dependencies = [ [[package]] name = "cpufeatures" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" dependencies = [ "libc", ] [[package]] name = "crc32fast" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" dependencies = [ "cfg-if", ] [[package]] name = "crossbeam-channel" -version = "0.5.15" +version = "0.5.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" dependencies = [ "crossbeam-utils", ] [[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-utils" -version = "0.8.21" +version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" [[package]] name = "crunchy" @@ -1347,7 +1343,7 @@ dependencies = [ "ident_case", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1358,14 +1354,14 @@ checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" dependencies = [ "darling_core", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "dashmap" -version = "6.1.0" +version = "6.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" dependencies = [ "cfg-if", "crossbeam-utils", @@ -1377,15 +1373,15 @@ dependencies = [ [[package]] name = "data-encoding" -version = "2.11.0" +version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" [[package]] name = "datafusion" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040" +checksum = "754ef4e8f073922a26f5b23133b9db4829342362b09be0bc94309cf261c2f098" dependencies = [ "arrow", "arrow-schema", @@ -1419,7 +1415,7 @@ dependencies = [ "datafusion-sql", "flate2", "futures", - "indexmap 2.14.0", + "indexmap", "itertools", "liblzma", "log", @@ -1436,9 +1432,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139" +checksum = "06afd1e38dd27bbb1258685a1fc6524df6aff4e07b25b393a47de59635178d99" dependencies = [ "arrow", "async-trait", @@ -1461,9 +1457,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891" +checksum = "f0668fb32c12065ec242be0e5b4bc62bd7a06a0be3ecd83791ef877e4be67e02" dependencies = [ "arrow", "async-trait", @@ -1484,9 +1480,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65" +checksum = "ca43b263cdff57042cfa8fb817fb3469f4878933380dccff25f5e793580abbf9" dependencies = [ "arrow", "arrow-ipc", @@ -1495,7 +1491,7 @@ dependencies = [ "foldhash 0.2.0", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap", "itertools", "libc", "log", @@ -1510,9 +1506,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "981b9dae74f78ee3d9f714fb49b01919eab975461b56149510c3ba9ea11287d1" +checksum = "05f0ba2b864792bdca4d76c59a1de0ab6e1b61946596b9936888dbd6360035f2" dependencies = [ "futures", "log", @@ -1521,9 +1517,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe" +checksum = "b840a8bce0bcbf5afad02946d438591e7c373f7afccaf3d874c04485772514dd" dependencies = [ "arrow", "async-compression", @@ -1548,7 +1544,7 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.4", + "rand 0.9.5", "tokio", "tokio-util", "url", @@ -1557,9 +1553,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f" +checksum = "a24cc0b9cf6e367f27f27406eff13abf48a11b72446aaa40b3105c0ded5c17d9" dependencies = [ "arrow", "arrow-ipc", @@ -1581,9 +1577,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada" +checksum = "e1abe56b2a7a2d1d6de5117dd1a203181e267f28529faa5da546947621b697d7" dependencies = [ "arrow", "async-trait", @@ -1604,9 +1600,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f" +checksum = "9c3e467f0611ad7bdd5aad17c63c9bb6182d04e5282e5496d897ea2b49c024ba" dependencies = [ "arrow", "async-trait", @@ -1627,9 +1623,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffd2499c1bee0eeccf6a57156105700eeeb17bc701899ac719183c4e74231450" +checksum = "4cc35b92cd560082155e80d9c826929c852d3c51543f4affd3a51c464a0aab3a" dependencies = [ "arrow", "async-trait", @@ -1658,15 +1654,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb9e7e5d11130c48c8bd4e80c79a9772dd28ce6dc330baca9246205d245b9e2e" +checksum = "d69bb69d8769e34f76839c960dbde24c1ac0c885a79b6c3c2287bdc56ec67891" [[package]] name = "datafusion-execution" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0" +checksum = "d8eac0a09bc8d263f52025cad9e001da4d8138d633fa288edda4d06b1772eae6" dependencies = [ "arrow", "arrow-buffer", @@ -1679,16 +1675,16 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.4", + "rand 0.9.5", "tempfile", "url", ] [[package]] name = "datafusion-expr" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0" +checksum = "eeb14d374767ee0fc62dc79a5ba8bcf8a63c14e993c7d992d0e63adfa23d77d3" dependencies = [ "arrow", "arrow-schema", @@ -1700,7 +1696,7 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr-common", - "indexmap 2.14.0", + "indexmap", "itertools", "recursive", "serde_json", @@ -1709,25 +1705,25 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2" +checksum = "b7b19a8c95522bee8cbb313d74263b85e355d2b52f42e67ef5694bf5de9e9356" dependencies = [ "arrow", "datafusion-common", - "indexmap 2.14.0", + "indexmap", "itertools", ] [[package]] name = "datafusion-functions" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" +checksum = "5f64c983bbbdcb729d921a2b2ac3375598719b5cc0c30345ad664936f3176fc7" dependencies = [ "arrow", "arrow-buffer", - "base64 0.22.1", + "base64", "blake2", "blake3", "chrono", @@ -1745,7 +1741,7 @@ dependencies = [ "md-5 0.11.0", "memchr", "num-traits", - "rand 0.9.4", + "rand 0.9.5", "regex", "sha2 0.11.0", "uuid", @@ -1753,9 +1749,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09" +checksum = "89bc17041e424a47ed062f43df24d84aab8b57c4c3221e5c1a5eef46d6c5718b" dependencies = [ "arrow", "datafusion-common", @@ -1774,9 +1770,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c" +checksum = "97dd2a9e865c6108059f5b37b77934f84b50bfb108f837bd0e5c9536e03f0545" dependencies = [ "arrow", "datafusion-common", @@ -1786,9 +1782,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5" +checksum = "75f0bdfeef16d96417b9632ef855645376b242e9006a126dfd0bedfc54a93f5f" dependencies = [ "arrow", "arrow-ord", @@ -1811,9 +1807,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26" +checksum = "f4e4941673c917819616877e9993da4503e4f4739812be0bc32c5356184c6383" dependencies = [ "arrow", "async-trait", @@ -1827,9 +1823,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78" +checksum = "12dd2e16c12b84b6f6b41b19f55b366dd1c46876bb35b86896c6349067379e8d" dependencies = [ "arrow", "datafusion-common", @@ -1844,9 +1840,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa84836dc2392df6f43d6a29d37fb56a8ebdc8b3f4e10ae8dc15861fd20278fb" +checksum = "4cdc5e4b6f8b6ef823cc1c761f85088ad4c884fe8df64df3cbcc6b2b84698441" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1854,20 +1850,20 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5" +checksum = "1a3614234dd93578c92428cb4f408e020874f0d2b7e6c90c928d9d28b5df2ceb" dependencies = [ "datafusion-doc", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "datafusion-optimizer" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281" +checksum = "a0635620b050b81bb92764e99250868f654e2cd5ad1bece413283b3f73c83179" dependencies = [ "arrow", "chrono", @@ -1875,7 +1871,7 @@ dependencies = [ "datafusion-expr", "datafusion-expr-common", "datafusion-physical-expr", - "indexmap 2.14.0", + "indexmap", "itertools", "log", "recursive", @@ -1885,9 +1881,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c" +checksum = "8cabf7a86eb70b816729e33c81bf7767c936ee1226f607a114f5dac2decac8d0" dependencies = [ "arrow", "datafusion-common", @@ -1897,7 +1893,7 @@ dependencies = [ "datafusion-physical-expr-common", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap", "itertools", "parking_lot", "petgraph", @@ -1907,9 +1903,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235" +checksum = "de222e04f7e6744501555a54ab0abe26bfdfebee380af79a9bdc175704246859" dependencies = [ "arrow", "datafusion-common", @@ -1922,16 +1918,16 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8" +checksum = "72d0d0057fc5a502d45c870cb6d47c66eb7bdd5edb1bd71ad6f3f724975ac2a8" dependencies = [ "arrow", "chrono", "datafusion-common", "datafusion-expr-common", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap", "itertools", "parking_lot", "pin-project", @@ -1939,9 +1935,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36" +checksum = "86046eed10950c5f9aaed9acfd148e9bd2e1dfdfe4f9aef607d1447b271e4183" dependencies = [ "arrow", "datafusion-common", @@ -1958,9 +1954,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc" +checksum = "9bc84da934c903407ba297971ebcc020c4c1a38aafd765d6c144c76eff3fa6a1" dependencies = [ "arrow", "arrow-data", @@ -1980,7 +1976,7 @@ dependencies = [ "futures", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap", "itertools", "log", "num-traits", @@ -1991,9 +1987,9 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dd15a1ba5d3af93808241065c6c44dbca8296a189845e8a587c45c07bf0ffae" +checksum = "67791dcfacd142a9d95f8f73c2a161094cc936d231d80c235f5017dacf24e84e" dependencies = [ "arrow", "chrono", @@ -2018,9 +2014,9 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90042982cf9462eb06a0b81f92efa4188dae871e7ea3ab8dc61aa9c9349b2530" +checksum = "b8cd9e80d637891645d074db0f6c650b591117367247deb313fdfb78dff559cb" dependencies = [ "arrow", "datafusion-common", @@ -2029,9 +2025,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3" +checksum = "eb63eeac6de19be40f487b65dd84e546195f783c5a9928618e0c4f2a3569b0d7" dependencies = [ "arrow", "datafusion-common", @@ -2045,9 +2041,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "15ce715fa2a61f4623cc234bcc14a3ef6a91f189128d5b14b468a6a17cdfc417" +checksum = "5f961d209177f91bd014db5cbb2c33b7d28a2597b9003e77f17aeb712964315a" dependencies = [ "async-trait", "datafusion-common", @@ -2059,9 +2055,9 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "54.0.0" +version = "54.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099" +checksum = "b1d71cb454da682b2af7488e1fc1ddd72ee1b28f19297b8ccad73f0a21ee9a69" dependencies = [ "arrow", "bigdecimal", @@ -2069,7 +2065,7 @@ dependencies = [ "datafusion-common", "datafusion-expr", "datafusion-functions-nested", - "indexmap 2.14.0", + "indexmap", "log", "recursive", "regex", @@ -2085,14 +2081,42 @@ dependencies = [ "uuid", ] +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.20", +] + [[package]] name = "deranged" version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" -dependencies = [ - "powerfmt", -] [[package]] name = "derive_arbitrary" @@ -2102,7 +2126,7 @@ checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2124,7 +2148,7 @@ dependencies = [ "proc-macro2", "quote", "rustc_version", - "syn", + "syn 2.0.119", "unicode-xid", ] @@ -2133,6 +2157,7 @@ name = "dev-tools" version = "0.1.3" dependencies = [ "dioxus", + "wasm-bindgen", ] [[package]] @@ -2152,7 +2177,7 @@ version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ - "block-buffer 0.12.0", + "block-buffer 0.12.1", "const-oid", "crypto-common 0.2.2", ] @@ -2192,9 +2217,9 @@ dependencies = [ [[package]] name = "dioxus-asset-resolver" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8b546050ecfc7fcd310be344b2f3a2a79f21554c5a9e8df28e7a07e9b36009b" +checksum = "84235ff0e272f15d4537603cdd50df6d45d50e51bb96326dbc0bdfd01d825226" dependencies = [ "dioxus-cli-config", "http", @@ -2205,7 +2230,7 @@ dependencies = [ "ndk-context", "ndk-sys", "percent-encoding", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "wasm-bindgen-futures", "web-sys", @@ -2213,18 +2238,18 @@ dependencies = [ [[package]] name = "dioxus-cli-config" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4ad73f0ff638cd27466d389cd57f0975f909b66130dc1c25d5212d4041e5352" +checksum = "322d10f47effebd85ee0662a2cb083ff920d40dd6295a4afd8cfa4f9bc2e05c6" dependencies = [ "wasm-bindgen", ] [[package]] name = "dioxus-config-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e004bc8b958031117d373db2b5e0ab9d7e763751de129dd36f00ef7e318333cd" +checksum = "9f3a906c8219e94fa48189a793c760a4e3193cbb7791f799e9d37633fe41161d" dependencies = [ "proc-macro2", "quote", @@ -2232,15 +2257,15 @@ dependencies = [ [[package]] name = "dioxus-config-macros" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "808a9994a9a2623e6b6890b6cc68def24bd669177ec4713684447fb46418c256" +checksum = "20de015bc89c20f4ffc7b4f849943420308365fa66f991250b5c075261990a21" [[package]] name = "dioxus-core" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "247ed8d679a13232641f1c84ba22246623fae01320b4c22db225c0b4f2fa7398" +checksum = "1b232228ada232c0adc151e41ebd9ef70cc04683db097fa0dc6015979fa01722" dependencies = [ "anyhow", "const_format", @@ -2249,7 +2274,7 @@ dependencies = [ "futures-util", "generational-box", "longest-increasing-subsequence", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "rustversion", "serde", "slab", @@ -2260,28 +2285,28 @@ dependencies = [ [[package]] name = "dioxus-core-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75fbe64029b90144041f8521300c7b3508e6c48caea3244de2ff5d1ade15390c" +checksum = "6256a68462222f0600f5fd7d9542cbd9f299bc19355d900f7827dd3fc9e09fdb" dependencies = [ "convert_case 0.8.0", "dioxus-rsx", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "dioxus-core-types" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbfea5c8946e0745b254b5c33c515d81b3ba638f33c2a532ed06730392394d4d" +checksum = "96007ed6cbad951dfed82d83c581aabd95bc591f351f5666d454ddbe7845b324" [[package]] name = "dioxus-devtools" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d30370fa78266aed3f3d9119dea2de9e92a0348941dbfa777f7a669f2ea375" +checksum = "434cebb282b3f820a341582ebfd37dab907d5a25c3eeafd87802468afcd2a959" dependencies = [ "dioxus-cli-config", "dioxus-core", @@ -2292,16 +2317,16 @@ dependencies = [ "serde", "serde_json", "subsecond", - "thiserror 2.0.18", + "thiserror 2.0.20", "tracing", "tungstenite 0.28.0", ] [[package]] name = "dioxus-devtools-types" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3907a2b61cf56039f047da6a37317a03d9e15411753bc40e5e34a27e405ac320" +checksum = "dfd6f3475e38c93be245a664a9ce11679a8f2b99832620705cf7837f42aa39ce" dependencies = [ "dioxus-core", "serde", @@ -2310,9 +2335,9 @@ dependencies = [ [[package]] name = "dioxus-document" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3b75a1809af7c13546ae4487c8b02ab80cc4d059e7db5a5d090374ceaa71d5b" +checksum = "2f6fadd5c6d886b48b762820abae1a08142a2159f8a1136c0466b4b4ed9f547e" dependencies = [ "dioxus-core", "dioxus-core-macro", @@ -2329,9 +2354,9 @@ dependencies = [ [[package]] name = "dioxus-fullstack" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8ee56dd65fbf1222fa6a2749c3f821df28facf1832854b43d480b547a096f6d" +checksum = "a91d5acabd1470145bcf164cab9fd2edcf2e4fba0fa6fb0fd6dca23079fcc169" dependencies = [ "anyhow", "async-stream", @@ -2339,7 +2364,7 @@ dependencies = [ "axum", "axum-core", "axum-extra", - "base64 0.22.1", + "base64", "bytes", "ciborium", "const-str", @@ -2374,13 +2399,13 @@ dependencies = [ "serde_json", "serde_qs", "serde_urlencoded", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tokio-stream", "tokio-tungstenite 0.28.0", "tokio-util", "tower", - "tower-http", + "tower-http 0.6.11", "tower-layer", "tracing", "tungstenite 0.27.0", @@ -2394,13 +2419,13 @@ dependencies = [ [[package]] name = "dioxus-fullstack-core" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d40d33a447cb158acdb61787b2ff52dd8a0f9a9f20e95e5c5fe9873f01c2b55b" +checksum = "0011e9bca8da2ae6bf02ba9ef93ba6c5fa539e402de3aeb316ad86ca77c59079" dependencies = [ "anyhow", "axum-core", - "base64 0.22.1", + "base64", "ciborium", "dioxus-core", "dioxus-document", @@ -2415,30 +2440,30 @@ dependencies = [ "parking_lot", "serde", "serde_json", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tracing", ] [[package]] name = "dioxus-fullstack-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c06eb66bce50d5f47b793e6af5fc2e0a511bf2b4fa2423cf86e35023d8f17e6" +checksum = "681f39665de36af7bb724b1cc0ca9c2b7c36a2b93febf1b2911716479eac3db0" dependencies = [ "const_format", "convert_case 0.8.0", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "xxhash-rust", ] [[package]] name = "dioxus-history" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1d8024afd482956eadae2c43d0b1e73e584adb724ac09be87f268d52002387b" +checksum = "2404b77d4441e694a5e93afb5c9a729efb26a3f6920f769cfe6cfedf5c7ac210" dependencies = [ "dioxus-core", "tracing", @@ -2446,9 +2471,9 @@ dependencies = [ [[package]] name = "dioxus-hooks" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "233b5e168a7c38c4bf96d0f390221a72a5a28bb31ce2dcf6d2af879dc561ef42" +checksum = "9d3807a1bc039299cd37ad85f1c56ae7088210e8a38e3e877f01034d68c15fbb" dependencies = [ "dioxus-core", "dioxus-signals", @@ -2462,9 +2487,9 @@ dependencies = [ [[package]] name = "dioxus-html" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4abf4ad27eee650d1ab8ebe13591e8b0ee595fa5a5dd236be13a5b7b3fab678d" +checksum = "40b9a85679c4dfd8699407a9ad5041a6a4fe6be0ee386353951cd03b57a12030" dependencies = [ "async-trait", "bytes", @@ -2489,28 +2514,28 @@ dependencies = [ [[package]] name = "dioxus-html-internal-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "025e107e677f790f4ed648a189e36f51ae014e5341902b97313e4eae626cffa2" +checksum = "500282dad5b62e91f93e127cbf41c587bf1c457d52514a75f3ff2a55fd0ea9f6" dependencies = [ "convert_case 0.8.0", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "dioxus-interpreter-js" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57caa76427d8ec4105ccca44ff8c511055688af732a094b9fe9ef3547d2a11b2" +checksum = "faa22cc3431e7efdbae28069542a1763c977d5edbe4d4fbfb0b7667b9c9966e9" dependencies = [ "dioxus-core", "dioxus-core-types", "dioxus-html", "js-sys", "lazy-js-bundle", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "sledgehammer_bindgen", "sledgehammer_utils", "wasm-bindgen", @@ -2520,9 +2545,9 @@ dependencies = [ [[package]] name = "dioxus-liveview" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64e64d86ad604897c796fdcc1f1f42cab838258647dd47664ab637a8eb97f08e" +checksum = "8b9cfecfc0037cb85b649739acc49a09dd303d1fa1daa7af609f5796b491b71c" dependencies = [ "axum", "dioxus-cli-config", @@ -2535,11 +2560,11 @@ dependencies = [ "futures-channel", "futures-util", "generational-box", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "serde", "serde_json", "slab", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tokio-stream", "tokio-util", @@ -2548,9 +2573,9 @@ dependencies = [ [[package]] name = "dioxus-logger" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cbbee192b1b12fccb444a5b04d809710cfce4d27b792129fea6c845fae7f329" +checksum = "0f36b2dac4df8c747431f77f2a6644dd033ea98e7f8ac80271ff47fd4ba8a548" dependencies = [ "dioxus-cli-config", "tracing", @@ -2560,9 +2585,9 @@ dependencies = [ [[package]] name = "dioxus-router" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ecdb19d7a1489ba252be9b3a07f9db92814020eb4ba8c1306f04242a44d17e66" +checksum = "e9652a5bcead34f687d7dbb12d5c8f6a7ecb342786273395ca0078e51b24320f" dependencies = [ "dioxus-cli-config", "dioxus-core", @@ -2581,9 +2606,9 @@ dependencies = [ [[package]] name = "dioxus-router-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b525ab775585f1dc4850178de9d0cb6bb37f7b9c1a1a0c121eab7449d7021480" +checksum = "56d20cf2173183946fdf8ab069b606a840a9d06e949050619355fc8cb97aafb6" dependencies = [ "base16", "digest 0.10.7", @@ -2591,32 +2616,32 @@ dependencies = [ "quote", "sha2 0.10.9", "slab", - "syn", + "syn 2.0.119", ] [[package]] name = "dioxus-rsx" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37fb07e40e9734946511659668ea3675ed214b60889e7aa7f0a5a85271518475" +checksum = "aea662c9e13a91279d8acbc0d02a80430c7081704c20e8c791e4c4a151268596" dependencies = [ "proc-macro2", "proc-macro2-diagnostics", "quote", "rustversion", - "syn", + "syn 2.0.119", ] [[package]] name = "dioxus-server" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aa22ad381073c68a70cdb28152485abb5f2dcf20dabc68c631e26ce7ac046dd" +checksum = "ad8d8cfcc78453d8f1eba4350126161844a0b10b4dc898bac7649b8615f92c48" dependencies = [ "anyhow", "async-trait", "axum", - "base64 0.22.1", + "base64", "bytes", "chrono", "ciborium", @@ -2648,17 +2673,17 @@ dependencies = [ "lru", "parking_lot", "pin-project", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "serde", "serde_json", "serde_qs", "subsecond", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tokio-tungstenite 0.28.0", "tokio-util", "tower", - "tower-http", + "tower-http 0.6.11", "tracing", "tracing-futures", "url", @@ -2667,37 +2692,37 @@ dependencies = [ [[package]] name = "dioxus-signals" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5393fd579f42c6547bf47ec0a2dedf2a366cd541d31deedc1059a096e6c35798" +checksum = "95d85e36fcaf7abdc986836801bf248776da7d5961f5fac65fdbf5ed491afb4f" dependencies = [ "dioxus-core", "futures-channel", "futures-util", "generational-box", "parking_lot", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "tracing", "warnings", ] [[package]] name = "dioxus-ssr" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d17c75a43e218012b63a97f55cf585747bd0ca37840ac2a0cf40add06ffcf9fe" +checksum = "4544ecf908962909de403bfdf94f771b036e9e1857e371cd3279841a15a892e6" dependencies = [ "askama_escape", "dioxus-core", "dioxus-core-types", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", ] [[package]] name = "dioxus-stores" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1c59f52e8194439604dd35f8c540456c22b4a4b076930e424d0289f98ea3cb4" +checksum = "15eaa9f2ccc5ce962d056515e61345c846a54b13a2b880cf7e8ac93faa53a178" dependencies = [ "dioxus-core", "dioxus-signals", @@ -2707,21 +2732,21 @@ dependencies = [ [[package]] name = "dioxus-stores-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "737600865572cecf60ff934f88252bd4144f4ca189e0e570b7a780e8f0e01a1f" +checksum = "a13bd8a693d7042544fa3ac79698a08783bbc698d1ed257e88c5d4059e9ac476" dependencies = [ "convert_case 0.8.0", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "dioxus-web" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "176eb0a5ee8251203a816b413a64fdc014b43e53d4245f769b1b0c2035b88ac3" +checksum = "481093233d789c5ffbdde66e7970e803b01a199a7ab0c0f8ef84c65d30b7384f" dependencies = [ "dioxus-cli-config", "dioxus-core", @@ -2739,7 +2764,7 @@ dependencies = [ "gloo-timers", "js-sys", "lazy-js-bundle", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "send_wrapper", "serde", "serde-wasm-bindgen", @@ -2751,15 +2776,25 @@ dependencies = [ "web-sys", ] +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags 2.13.1", + "objc2", +] + [[package]] name = "displaydoc" -version = "0.2.5" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -2784,7 +2819,7 @@ checksum = "9556bc800956545d6420a640173e5ba7dfa82f38d3ea5a167eb555bc69ac3323" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2806,7 +2841,7 @@ dependencies = [ "pretty-hex", "serde", "serde_json", - "thiserror 2.0.18", + "thiserror 2.0.20", "zerocopy", ] @@ -2818,7 +2853,7 @@ checksum = "dc09b90bda5770641457f1c0a42c8203c48f5a3d9799dcf1bafbd84e30ccf080" dependencies = [ "pest", "pest_derive", - "thiserror 2.0.18", + "thiserror 2.0.20", ] [[package]] @@ -2829,9 +2864,9 @@ checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" [[package]] name = "either" -version = "1.15.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" [[package]] name = "encode_unicode" @@ -2850,23 +2885,23 @@ dependencies = [ [[package]] name = "enumset" -version = "1.1.10" +version = "1.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25b07a8dfbbbfc0064c0a6bdf9edcf966de6b1c33ce344bdeca3b41615452634" +checksum = "ccc5801fd11762e24d1e420d01d2ac518f2a2ca4329d4fbb6639f2412b6204e0" dependencies = [ "enumset_derive", ] [[package]] name = "enumset_derive" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f43e744e4ea338060faee68ed933e46e722fb7f3617e722a5772d7e856d8b3ce" +checksum = "4bd536557b58c682b217b8fb199afdff47cd3eff260623f19e77074eb073d63a" dependencies = [ "darling", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2886,7 +2921,7 @@ checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2895,17 +2930,6 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" -[[package]] -name = "erased-serde" -version = "0.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2add8a07dd6a8d93ff627029c51de145e12686fbc36ecb298ac22e74cf02dec" -dependencies = [ - "serde", - "serde_core", - "typeid", -] - [[package]] name = "errno" version = "0.3.14" @@ -2957,50 +2981,47 @@ dependencies = [ [[package]] name = "fastlanes" -version = "0.5.1" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20c597e23b8ec8506f589d18bc701ca83a3def6086748f628ad23092e1dfe577" +checksum = "34f6c951d711d8a10f08524071f6171bc95c5f857e8b39d91e36ad6353fc4f4f" dependencies = [ - "arrayref", "const_for", - "core_detect", "num-traits", - "paste", + "pastey", "seq-macro", ] [[package]] name = "fastrace" -version = "0.7.17" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2130caec636d7a1d23b173576674ced1af967228642ceaeb6a1b4705c282b00e" +checksum = "dcef89a7c5b37f7c13551af01f509cbd5d32a8644c27423623639eb923cc2973" dependencies = [ "fastant", "fastrace-macro", "parking_lot", "pin-project", - "rand 0.9.4", + "rand 0.10.2", "rtrb", "serde", ] [[package]] name = "fastrace-macro" -version = "0.7.17" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b35f67e02527fca6515ff61f922360df781f477daf6a806fff16bd59525dee5" +checksum = "2f1fdb97ada2fde7912c9cab29b560d550ccc467ca1f39c952ce4b05f8d152b3" dependencies = [ - "proc-macro-error2", "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] name = "fastrace-opentelemetry" -version = "0.16.0" +version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4644a8a7ce6d20d83e73d9562f323388e4d817470d40bcb51fa2fe328db58cd2" +checksum = "67d87b05a82ea59d6d2a11066f4e0cc1e169ce50bc8c276279249ba374936d91" dependencies = [ "fastrace", "log", @@ -3023,15 +3044,15 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "find-msvc-tools" -version = "0.1.9" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +checksum = "d45db016d36b838f563236e9193d0ee6ce38f3f68b6c94e914b4929c96bbb890" [[package]] name = "findshlibs" @@ -3057,18 +3078,18 @@ version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "rustc_version", ] [[package]] name = "flate2" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.9.1", "zlib-rs", ] @@ -3101,11 +3122,11 @@ dependencies = [ [[package]] name = "fsst-rs" -version = "0.5.11" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b13ac798afc0d9194eb4efefef8b9332efbd80b43f302a968cb8cb23b9d5360" +checksum = "03b614a2a7f2efc50d76c1e47d64b57c55acf6e0e718b0cae78236d147c487d0" dependencies = [ - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", ] [[package]] @@ -3116,9 +3137,9 @@ checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" [[package]] name = "futures" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" dependencies = [ "futures-channel", "futures-core", @@ -3131,9 +3152,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", "futures-sink", @@ -3141,15 +3162,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" [[package]] name = "futures-executor" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" dependencies = [ "futures-core", "futures-task", @@ -3158,38 +3179,38 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" [[package]] name = "futures-macro" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] name = "futures-sink" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-channel", "futures-core", @@ -3204,9 +3225,9 @@ dependencies = [ [[package]] name = "generational-box" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c68d74be1fbe3bba37604bdfd61403f26af9f6324cf325053abd89d60c22e799" +checksum = "b8b2dc9b873cd1a8adb80aa5dcef9a8205f781a8a313d5ffb867248cb3bcb764" dependencies = [ "parking_lot", "tracing", @@ -3242,25 +3263,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", - "js-sys", "libc", "r-efi 5.3.0", "wasip2", - "wasm-bindgen", ] [[package]] name = "getrandom" -version = "0.4.2" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", "rand_core 0.10.1", - "wasip2", - "wasip3", + "wasm-bindgen", ] [[package]] @@ -3271,9 +3290,9 @@ checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" [[package]] name = "glob" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] name = "gloo-net" @@ -3323,9 +3342,9 @@ dependencies = [ [[package]] name = "goblin" -version = "0.10.5" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "983a6aafb3b12d4c41ea78d39e189af4298ce747353945ff5105b54a056e5cd9" +checksum = "17582616a7718cca54cec18e534a76c7c4aec11a8b9a85695712f262fd15a4c8" dependencies = [ "log", "plain", @@ -3334,9 +3353,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16" dependencies = [ "atomic-waker", "bytes", @@ -3344,7 +3363,7 @@ dependencies = [ "futures-core", "futures-sink", "http", - "indexmap 2.14.0", + "indexmap", "slab", "tokio", "tokio-util", @@ -3363,12 +3382,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "hashbrown" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" - [[package]] name = "hashbrown" version = "0.14.5" @@ -3408,11 +3421,11 @@ dependencies = [ [[package]] name = "hdrhistogram" -version = "7.5.4" +version = "7.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" +checksum = "f49d1053f4708f0af3cf9fc5bffc7e68a914a3c45becb231c80068c9c3f78bea" dependencies = [ - "base64 0.21.7", + "base64", "byteorder", "crossbeam-channel", "flate2", @@ -3426,7 +3439,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "headers-core", "http", @@ -3452,9 +3465,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" [[package]] name = "hex" @@ -3464,9 +3477,9 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "http" -version = "1.4.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" dependencies = [ "bytes", "itoa", @@ -3474,9 +3487,9 @@ dependencies = [ [[package]] name = "http-body" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", "http", @@ -3484,9 +3497,9 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.3" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c" dependencies = [ "bytes", "futures-core", @@ -3515,24 +3528,24 @@ 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" -version = "0.4.12" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" dependencies = [ "typenum", ] [[package]] name = "hyper" -version = "1.9.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" +checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43" dependencies = [ "atomic-waker", "bytes", @@ -3586,7 +3599,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "futures-channel", "futures-util", @@ -3632,9 +3645,9 @@ dependencies = [ [[package]] name = "icu_collections" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513" dependencies = [ "displaydoc", "potential_utf", @@ -3646,9 +3659,9 @@ dependencies = [ [[package]] name = "icu_locale_core" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" dependencies = [ "displaydoc", "litemap", @@ -3659,9 +3672,9 @@ dependencies = [ [[package]] name = "icu_normalizer" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f" dependencies = [ "icu_collections", "icu_normalizer_data", @@ -3673,16 +3686,17 @@ dependencies = [ [[package]] name = "icu_normalizer_data" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" +checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0" [[package]] name = "icu_properties" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148" dependencies = [ + "displaydoc", "icu_collections", "icu_locale_core", "icu_properties_data", @@ -3693,15 +3707,15 @@ dependencies = [ [[package]] name = "icu_properties_data" -version = "2.2.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" +checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" [[package]] name = "icu_provider" -version = "2.2.0" +version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73" dependencies = [ "displaydoc", "icu_locale_core", @@ -3712,12 +3726,6 @@ dependencies = [ "zerovec", ] -[[package]] -name = "id-arena" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" - [[package]] name = "ident_case" version = "1.0.1" @@ -3747,24 +3755,12 @@ dependencies = [ [[package]] name = "indexmap" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" -dependencies = [ - "autocfg", - "hashbrown 0.12.3", -] - -[[package]] -name = "indexmap" -version = "2.14.0" +version = "2.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "07aa2048142242915a31d35844fb311e0e53fcca590c3a0a40dcf1b841fa09eb" dependencies = [ "equivalent", "hashbrown 0.17.1", - "serde", - "serde_core", ] [[package]] @@ -3783,7 +3779,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "232929e1d75fe899576a3d5c7416ad0d88dbfbb3c3d6aa00873a7408a50ddb88" dependencies = [ "ahash", - "indexmap 2.14.0", + "indexmap", "is-terminal", "itoa", "log", @@ -3796,9 +3792,9 @@ dependencies = [ [[package]] name = "insta" -version = "1.47.2" +version = "1.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b4a6248eb93a4401ed2f37dfe8ea592d3cf05b7cf4f8efa867b6895af7e094e" +checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" dependencies = [ "console", "once_cell", @@ -3823,20 +3819,20 @@ dependencies = [ [[package]] name = "io-uring" -version = "0.7.12" +version = "0.7.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d09b98f7eace8982db770e4408e7470b028ce513ac28fecdc6bf4c30fe92b62" +checksum = "d64d8ca234d152948ceaede1f419b6a83983a5ecccaac05fb337a809c96d3aa6" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "cfg-if", "libc", ] [[package]] name = "ipnet" -version = "2.12.0" +version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +checksum = "6a756c3fac73139e83f14c2d742155dd2b78d3ee56597b419a0579b7bdd6dd78" [[package]] name = "is-terminal" @@ -3872,35 +3868,47 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jiff" -version = "0.2.24" +version = "0.2.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f00b5dbd620d61dfdcb6007c9c1f6054ebd75319f163d886a9055cec1155073d" +checksum = "668b7183bd07af9a4885f5c35b0cc5c83c4607a913c16b7e17291832910d2dcc" dependencies = [ + "defmt", + "jiff-core", "jiff-static", "jiff-tzdb-platform", "log", "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.61.2", + "windows-link", +] + +[[package]] +name = "jiff-core" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7feca88439efe53da3754500c1851dedf3cb36c524dd5cf8225cc0794de95d09" +dependencies = [ + "defmt", ] [[package]] name = "jiff-static" -version = "0.2.24" +version = "0.2.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e000de030ff8022ea1da3f466fbb0f3a809f5e51ed31f6dd931c35181ad8e6d7" +checksum = "3a69dcb3a21cfb32ce1cd056169337ca284af0766dd766e7878819b251a49204" dependencies = [ + "jiff-core", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "jiff-tzdb" -version = "0.1.6" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076" +checksum = "142bd39932ad231f10513df9ab62661fead8719872150b7ad02a2df79f4e141e" [[package]] name = "jiff-tzdb-platform" @@ -3952,28 +3960,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" dependencies = [ "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "jobserver" -version = "0.1.34" +version = "0.1.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" dependencies = [ - "getrandom 0.3.4", + "getrandom 0.4.3", "libc", ] [[package]] name = "js-sys" -version = "0.3.95" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", - "once_cell", "wasm-bindgen", ] @@ -3983,7 +3990,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b750dcadc39a09dbadd74e118f6dd6598df77fa01df0cfcdc52c28dece74528a" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "serde", ] @@ -4004,9 +4011,9 @@ checksum = "a4933f3f57a8e9d9da04db23fb153356ecaf00cbd14aee46279c33dc80925c37" [[package]] name = "lazy-js-bundle" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebbde2c5796719fbd82d6b8ec0be3dacf1f70c2876dee0f2c001632794d6641f" +checksum = "51cdaa5abc885e2d606e6985281ace220778bde190be20aab1ec2346ed8bd1fa" [[package]] name = "lazy_static" @@ -4014,12 +4021,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "leb128fmt" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" - [[package]] name = "lexical-core" version = "1.0.6" @@ -4079,21 +4080,21 @@ dependencies = [ [[package]] name = "libbz2-rs-sys" -version = "0.2.3" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3a6a8c165077efc8f3a971534c50ea6a1a18b329ef4a66e897a7e3a1494565f" +checksum = "34b357333733e8260735ba5894eb928c02ecc69c78715f01a8019e7fa7f2db4c" [[package]] name = "libc" -version = "0.2.186" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libfuzzer-sys" -version = "0.4.12" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d" +checksum = "a9fd2f41a1cba099f79a0b6b6c35656cf7c03351a7bae8ff0f28f25270f929d2" dependencies = [ "arbitrary", "cc", @@ -4111,18 +4112,18 @@ dependencies = [ [[package]] name = "liblzma" -version = "0.4.6" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" +checksum = "2fe0a34ca854fd4f20c07f696fc8675aec78f87d88d29f5e10257a7490a1b2e1" dependencies = [ "liblzma-sys", ] [[package]] name = "liblzma-sys" -version = "0.4.6" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a60851d15cd8c5346eca4ab8babff585be2ae4bc8097c067291d3ffe2add3b6" +checksum = "a0dad045e4b1b7b170be4b60b54b780cafb4490165461bac7d1cf7b703f61d5f" dependencies = [ "cc", "libc", @@ -4180,7 +4181,7 @@ dependencies = [ "object_store", "parquet", "parquet-variant-compute", - "rand 0.10.1", + "rand 0.10.2", "serde", "serde_json", "shuttle", @@ -4245,7 +4246,7 @@ dependencies = [ "prost", "serde", "tempfile", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "url", ] @@ -4268,7 +4269,7 @@ dependencies = [ "object_store", "parquet", "parquet-variant-json", - "rand 0.10.1", + "rand 0.10.2", "serde_json", "shuttle", "t4", @@ -4345,7 +4346,7 @@ dependencies = [ "tempfile", "tokio", "tonic", - "tower-http", + "tower-http 0.7.1", "url", "uuid", ] @@ -4362,9 +4363,9 @@ dependencies = [ [[package]] name = "litemap" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" [[package]] name = "litrs" @@ -4383,44 +4384,40 @@ dependencies = [ [[package]] name = "log" -version = "0.4.32" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" -dependencies = [ - "sval", - "sval_ref", - "value-bag", -] +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "logforth" -version = "0.29.1" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40c105c59828d07aeb95b06f9a345b12869ddc249d44a7302697a66da439076f" +checksum = "522000d3921e4b089de59204d2d3ca8792cd53f8ce5a54b5ac8b9a6e867259f0" dependencies = [ - "logforth-append-opentelemetry", + "log", + "logforth-append-file", "logforth-bridge-log", "logforth-core", + "logforth-filter-rustlog", + "logforth-layout-json", + "logforth-layout-text", ] [[package]] -name = "logforth-append-opentelemetry" -version = "0.3.1" +name = "logforth-append-file" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "544f950997c23b8b0a8c324af05aad29db446738fb04ae57cbcdf233d67b1865" +checksum = "80f4ed78a03a30c12285135d98330f0f5d53cb0019bd1ac6d66863dae72d8d52" dependencies = [ + "jiff", "logforth-core", - "logforth-layout-json", - "opentelemetry", - "opentelemetry-otlp", - "opentelemetry_sdk", ] [[package]] name = "logforth-bridge-log" -version = "0.3.0" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4aa6ca548389fd166a995b5940e15b0dacbdd5a30f2f24eac9aa4bf664bda5c" +checksum = "9c7224c78547e542572ae4d1f787f96c4395a4c3f24513bf221bd8e49888f610" dependencies = [ "log", "logforth-core", @@ -4428,19 +4425,28 @@ dependencies = [ [[package]] name = "logforth-core" -version = "0.3.1" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a77869b8dba38c67ed19e1753e59d9faefdcc60557bc4e84db0348606a304ac5" +checksum = "400ba5305e0cb6819efa6c2c503020562eb5b47fd53c91c935be55af1ffd374e" dependencies = [ "anyhow", - "value-bag", + "serde", +] + +[[package]] +name = "logforth-filter-rustlog" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e8431cfa1d3eeca479c363651320a66c7003350528f678b30e8c1474fb0d802" +dependencies = [ + "logforth-core", ] [[package]] name = "logforth-layout-json" -version = "0.3.0" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01b80d310e0670560404a825f64dbd78a8761c5bb7da952513e90ba9dd525bd2" +checksum = "9d86a105f8f32151ca1e0a6d4097d478badb02d5715239ba0fa431a188a5d966" dependencies = [ "jiff", "logforth-core", @@ -4448,6 +4454,17 @@ dependencies = [ "serde_json", ] +[[package]] +name = "logforth-layout-text" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4eafe007ac55293d807e8c7f5a23002e2518d32e476e195443cde23e7845416a" +dependencies = [ + "colored", + "jiff", + "logforth-core", +] + [[package]] name = "longest-increasing-subsequence" version = "0.1.0" @@ -4471,9 +4488,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.0" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" dependencies = [ "twox-hash", ] @@ -4486,14 +4503,14 @@ checksum = "1b27834086c65ec3f9387b096d66e99f221cf081c2b738042aa252bcd41204e3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "manganis" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e06225f29a781d86afdfafa562de09621f2ace377136ae2ae9ca9e72a29b920" +checksum = "1ed2af894a724a0b6e521f9efe84f98cfdc74a4de6828efef77ffd3b1679fd52" dependencies = [ "const-serialize 0.7.2", "const-serialize 0.8.0-alpha.0", @@ -4502,14 +4519,14 @@ dependencies = [ "manganis-macro", "ndk-context", "objc2", - "thiserror 2.0.18", + "thiserror 2.0.20", ] [[package]] name = "manganis-core" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "774ddc382b4fb30f3fdcf2418131cbac5e6111f00e6c7ddf9e598ef3b2b4cd91" +checksum = "e258c71b137bf48deaffb48b8a8a2f994af9e5c3ca201daa20d41e26b379da11" dependencies = [ "const-serialize 0.7.2", "const-serialize 0.8.0-alpha.0", @@ -4521,16 +4538,16 @@ dependencies = [ [[package]] name = "manganis-macro" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "731c83c89d831f341fb46eba0aefdfb3433f77a3f78a8b08d9d88746613a8f8b" +checksum = "0abae8346bbf637e6d21ac3831863bd5a98657517f5cd2d69d651046b95e7ce9" dependencies = [ "dunce", "macro-string", "manganis-core", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -4570,9 +4587,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.0" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "memfd" @@ -4585,9 +4602,9 @@ dependencies = [ [[package]] name = "memmap2" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", ] @@ -4626,12 +4643,6 @@ dependencies = [ "unicase", ] -[[package]] -name = "minimal-lexical" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" - [[package]] name = "miniz_oxide" version = "0.8.9" @@ -4639,14 +4650,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" dependencies = [ "adler2", - "simd-adler32", +] + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", ] [[package]] name = "mio" -version = "1.2.0" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", @@ -4666,7 +4686,7 @@ dependencies = [ "httparse", "memchr", "mime", - "spin 0.9.8", + "spin 0.9.9", "version_check", ] @@ -4676,7 +4696,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "jni-sys 0.3.1", "log", "ndk-sys", @@ -4713,12 +4733,11 @@ dependencies = [ [[package]] name = "nom" -version = "7.1.3" +version = "8.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" dependencies = [ "memchr", - "minimal-lexical", ] [[package]] @@ -4741,9 +4760,9 @@ dependencies = [ [[package]] name = "num-bigint" -version = "0.4.6" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" dependencies = [ "num-integer", "num-traits", @@ -4760,9 +4779,9 @@ dependencies = [ [[package]] name = "num-conv" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" [[package]] name = "num-format" @@ -4776,9 +4795,9 @@ dependencies = [ [[package]] name = "num-integer" -version = "0.1.46" +version = "0.1.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" dependencies = [ "num-traits", ] @@ -4812,7 +4831,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -4830,7 +4849,9 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", + "dispatch2", + "objc2", ] [[package]] @@ -4839,6 +4860,16 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" +[[package]] +name = "objc2-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" +dependencies = [ + "bitflags 2.13.1", + "objc2", +] + [[package]] name = "objc2-io-kit" version = "0.3.2" @@ -4849,6 +4880,17 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "objc2-open-directory" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb82bed227edf5201dfedf072bba4015a33d3d4a98519837295a90f0a23f676d" +dependencies = [ + "objc2", + "objc2-core-foundation", + "objc2-foundation", +] + [[package]] name = "object" version = "0.37.3" @@ -4858,6 +4900,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "object" +version = "0.39.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e5a6c098c7a3b6547378093f5cc30bc54fd361ce711e05293a5cc589562739b" +dependencies = [ + "memchr", +] + [[package]] name = "object_store" version = "0.13.2" @@ -4865,7 +4916,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" dependencies = [ "async-trait", - "base64 0.22.1", + "base64", "bytes", "chrono", "form_urlencoded", @@ -4880,14 +4931,14 @@ dependencies = [ "md-5 0.10.6", "parking_lot", "percent-encoding", - "quick-xml 0.39.2", - "rand 0.10.1", + "quick-xml 0.39.4", + "rand 0.10.2", "reqwest 0.12.28", "ring", "serde", "serde_json", "serde_urlencoded", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tracing", "url", @@ -4916,36 +4967,36 @@ checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" [[package]] name = "opentelemetry" -version = "0.31.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b84bcd6ae87133e903af7ef497404dda70c60d0ea14895fc8a5e6722754fc2a0" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" dependencies = [ "futures-core", "futures-sink", "js-sys", "pin-project-lite", - "thiserror 2.0.18", + "thiserror 2.0.20", "tracing", ] [[package]] name = "opentelemetry-http" -version = "0.31.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7a6d09a73194e6b66df7c8f1b680f156d916a1a942abf2de06823dd02b7855d" +checksum = "5683015d09e2df236ef005b17f6f196f0d5f6313c4fa43a7b6a53b52776e4331" dependencies = [ "async-trait", "bytes", "http", "opentelemetry", - "reqwest 0.12.28", + "reqwest 0.13.4", ] [[package]] name = "opentelemetry-otlp" -version = "0.31.1" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f69cd6acbb9af919df949cd1ec9e5e7fdc2ef15d234b6b795aaa525cc02f71f" +checksum = "9966929966d17620d7c316c643ba62631826e10021409357772d5eea84f62c35" dependencies = [ "http", "opentelemetry", @@ -4953,18 +5004,18 @@ dependencies = [ "opentelemetry-proto", "opentelemetry_sdk", "prost", - "reqwest 0.12.28", - "thiserror 2.0.18", + "reqwest 0.13.4", + "thiserror 2.0.20", "tokio", "tonic", - "tracing", + "tonic-types", ] [[package]] name = "opentelemetry-proto" -version = "0.31.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" +checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" dependencies = [ "opentelemetry", "opentelemetry_sdk", @@ -4975,17 +5026,18 @@ dependencies = [ [[package]] name = "opentelemetry_sdk" -version = "0.31.0" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14ae4f5991976fd48df6d843de219ca6d31b01daaab2dad5af2badeded372bd" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" dependencies = [ "futures-channel", "futures-executor", "futures-util", "opentelemetry", "percent-encoding", - "rand 0.9.4", - "thiserror 2.0.18", + "portable-atomic", + "rand 0.9.5", + "thiserror 2.0.20", ] [[package]] @@ -4999,9 +5051,9 @@ dependencies = [ [[package]] name = "owo-colors" -version = "3.5.0" +version = "4.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f" +checksum = "13c45bb4a6ae1280ec0803b1ef9d3455eb50f01efbbe1447ab020f1d54fba9d8" [[package]] name = "parking_lot" @@ -5028,9 +5080,9 @@ dependencies = [ [[package]] name = "parquet" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" +checksum = "d298093b2dec60289dce0684c986d0f7679e9dd15771c2c65406e1aaf604a704" dependencies = [ "ahash", "arrow-array", @@ -5039,7 +5091,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64", "brotli", "bytes", "chrono", @@ -5067,15 +5119,15 @@ dependencies = [ [[package]] name = "parquet-variant" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74c8db065291f088a2aad8ab831853eae1871c0d311c8d0b83bbc3b7e735d0fc" +checksum = "3fc70e87931167a4a3fde2ee923023a0624367691e3fd503476a1084dda5a054" dependencies = [ "arrow", "arrow-schema", "chrono", "half", - "indexmap 2.14.0", + "indexmap", "num-traits", "simdutf8", "uuid", @@ -5083,15 +5135,15 @@ dependencies = [ [[package]] name = "parquet-variant-compute" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a530e8d5b5e14efcb39c9a6ec55432ad11f6afb7dc4455a79be0dc615fe3cc31" +checksum = "823a9ecee8fd83a68f7165ef13acc840c4d7ed995838ba99d4714b20a2b5e780" dependencies = [ "arrow", "arrow-schema", "chrono", "half", - "indexmap 2.14.0", + "indexmap", "parquet-variant", "parquet-variant-json", "serde_json", @@ -5100,12 +5152,12 @@ dependencies = [ [[package]] name = "parquet-variant-json" -version = "58.3.0" +version = "58.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00ed89908289f67caa2ca078f9ff9aacd6229a313ec92b12bf4f48f613dc2b97" +checksum = "6f37a91177e2dddb10333546952fcf2b7674b97ebd9449432f24b883d6ea8108" dependencies = [ "arrow-schema", - "base64 0.22.1", + "base64", "chrono", "parquet-variant", "serde_json", @@ -5118,6 +5170,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -5130,7 +5188,7 @@ version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "575828d9d7d205188048eb1508560607a03d21eafdbba47b8cade1736c1c28e1" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "c-enum", "perf-event-open-sys2", ] @@ -5151,7 +5209,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0939b8fad77dfaeb29ebbd35faaeaadbf833167f30975f1b8993bbba09ea0a0f" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "c-enum", "libc", "memmap2", @@ -5161,9 +5219,9 @@ dependencies = [ [[package]] name = "pest" -version = "2.8.6" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0848c601009d37dfa3430c4666e147e49cdcf1b92ecd3e63657d8a5f19da662" +checksum = "5a07a60cc7a4d00c91f95c685609d1d2f79050e6804b70ebedd7650f0b839bcf" dependencies = [ "memchr", "ucd-trie", @@ -5171,9 +5229,9 @@ dependencies = [ [[package]] name = "pest_derive" -version = "2.8.6" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11f486f1ea21e6c10ed15d5a7c77165d0ee443402f0780849d1768e7d9d6fe77" +checksum = "b3a83744a5c8455b8b3e0dc5031362780a347c878bdd11584d1a8984228cc88d" dependencies = [ "pest", "pest_generator", @@ -5181,25 +5239,24 @@ dependencies = [ [[package]] name = "pest_generator" -version = "2.8.6" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8040c4647b13b210a963c1ed407c1ff4fdfa01c31d6d2a098218702e6664f94f" +checksum = "e0cd3451aa3de60d4b9a1e736885e4dea6b31617598026f12256ad566d63304a" dependencies = [ "pest", "pest_meta", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "pest_meta" -version = "2.8.6" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89815c69d36021a140146f26659a81d6c2afa33d216d736dd4be5381a7362220" +checksum = "e04d3a0849e241d7dfce834c83b1c5edc8622009e8dd51a12ba1927c32f05496" dependencies = [ "pest", - "sha2 0.10.9", ] [[package]] @@ -5210,7 +5267,7 @@ checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", "hashbrown 0.15.5", - "indexmap 2.14.0", + "indexmap", "serde", ] @@ -5234,22 +5291,22 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -5260,9 +5317,9 @@ checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "pkg-config" -version = "0.3.33" +version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" [[package]] name = "plain" @@ -5278,9 +5335,9 @@ checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" [[package]] name = "portable-atomic" -version = "1.13.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" [[package]] name = "portable-atomic-util" @@ -5293,9 +5350,9 @@ dependencies = [ [[package]] name = "potential_utf" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ "zerovec", ] @@ -5322,10 +5379,10 @@ dependencies = [ "nix", "once_cell", "smallvec", - "spin 0.10.0", + "spin 0.10.1", "symbolic-demangle", "tempfile", - "thiserror 2.0.18", + "thiserror 2.0.20", ] [[package]] @@ -5343,16 +5400,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a65843dfefbafd3c879c683306959a6de478443ffe9c9adf02f5976432402d7" -[[package]] -name = "prettyplease" -version = "0.2.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" -dependencies = [ - "proc-macro2", - "syn", -] - [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -5362,33 +5409,11 @@ dependencies = [ "toml_edit", ] -[[package]] -name = "proc-macro-error-attr2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" -dependencies = [ - "proc-macro2", - "quote", -] - -[[package]] -name = "proc-macro-error2" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" -dependencies = [ - "proc-macro-error-attr2", - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -5401,7 +5426,7 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "version_check", ] @@ -5425,14 +5450,14 @@ dependencies = [ "itertools", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "prost-types" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" dependencies = [ "prost", ] @@ -5445,9 +5470,9 @@ checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" [[package]] name = "psm" -version = "0.1.31" +version = "0.1.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "645dbe486e346d9b5de3ef16ede18c26e6c70ad97418f4874b8b1889d6e761ea" +checksum = "4dcd034599e63b970727f70d79e02d62390a4a84f7c6b827c27c46d5ac3fa622" dependencies = [ "ar_archive_writer", "cc", @@ -5474,9 +5499,9 @@ dependencies = [ [[package]] name = "quick-xml" -version = "0.39.2" +version = "0.39.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "958f21e8e7ceb5a1aa7fa87fab28e7c75976e0bfe7e23ff069e0a260f894067d" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" dependencies = [ "memchr", "serde", @@ -5484,19 +5509,19 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.9" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" dependencies = [ "bytes", "cfg_aliases", "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "rustls", "socket2", - "thiserror 2.0.18", + "thiserror 2.0.20", "tokio", "tracing", "web-time", @@ -5504,20 +5529,21 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "04759210543be93709136e28212294a659ef5001836ff4eab4d663e4529bba83" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg 0.10.2", "ring", - "rustc-hash 2.1.2", + "rustc-hash 2.1.3", "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.18", + "thiserror 2.0.20", "tinyvec", "tracing", "web-time", @@ -5525,23 +5551,23 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", "socket2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] name = "quote" -version = "1.0.45" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -5566,9 +5592,9 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" dependencies = [ "libc", "rand_chacha 0.3.1", @@ -5577,9 +5603,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.5", @@ -5587,12 +5613,12 @@ dependencies = [ [[package]] name = "rand" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ "chacha20", - "getrandom 0.4.2", + "getrandom 0.4.3", "rand_core 0.10.1", ] @@ -5649,6 +5675,15 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "raw-window-handle" version = "0.6.2" @@ -5672,7 +5707,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -5681,14 +5716,14 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", ] [[package]] name = "regex" -version = "1.12.4" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -5698,9 +5733,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.14" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" dependencies = [ "aho-corasick", "memchr", @@ -5725,11 +5760,10 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "cookie", "cookie_store", - "futures-channel", "futures-core", "futures-util", "h2", @@ -5756,7 +5790,7 @@ dependencies = [ "tokio-rustls", "tokio-util", "tower", - "tower-http", + "tower-http 0.6.11", "tower-service", "url", "wasm-bindgen", @@ -5772,9 +5806,11 @@ version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ - "base64 0.22.1", + "base64", "bytes", + "futures-channel", "futures-core", + "futures-util", "http", "http-body", "http-body-util", @@ -5789,7 +5825,7 @@ dependencies = [ "sync_wrapper", "tokio", "tower", - "tower-http", + "tower-http 0.6.11", "tower-service", "url", "wasm-bindgen", @@ -5822,15 +5858,15 @@ dependencies = [ [[package]] name = "rtrb" -version = "0.3.4" +version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ade083ccbb4bf536df69d1f6432cc23deb7acccff86b183f3923a6fd56a1153" +checksum = "fae8ee26b0371a29a77d2b2d6b3ae13aa81def6f9bf1b1b92a32d279a5e709b7" [[package]] name = "rustc-demangle" -version = "0.1.27" +version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" +checksum = "b74b56ffa8bb2830709a538c2cbcae9aa062db0d2a42563bfb09bdaae44020eb" [[package]] name = "rustc-hash" @@ -5840,9 +5876,9 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" [[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" @@ -5859,7 +5895,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys", @@ -5868,9 +5904,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.40" +version = "0.23.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" dependencies = [ "once_cell", "ring", @@ -5882,9 +5918,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -5894,9 +5930,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" dependencies = [ "web-time", "zeroize", @@ -5904,9 +5940,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "ring", "rustls-pki-types", @@ -5915,9 +5951,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" @@ -5966,13 +6002,13 @@ dependencies = [ [[package]] name = "scroll_derive" -version = "0.13.1" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed76efe62313ab6610570951494bdaa81568026e0318eaa55f167de70eeea67d" +checksum = "e1a36a382ed65dbcc0ab47fd5e9a94112417ccd34560a392ef3b7b0f0ec39148" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -5981,7 +6017,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "core-foundation 0.10.1", "core-foundation-sys", "libc", @@ -6021,9 +6057,9 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -6040,49 +6076,31 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "serde_buf" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc948de1bbead18a61be0b33182636603ea0239ca2577b9704fc39eba900e4e5" -dependencies = [ - "serde_core", -] - [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", -] - -[[package]] -name = "serde_fmt" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e497af288b3b95d067a23a4f749f2861121ffcb2f6d8379310dcda040c345ed" -dependencies = [ - "serde_core", + "syn 3.0.4", ] [[package]] name = "serde_json" -version = "1.0.149" +version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ "itoa", "memchr", @@ -6110,18 +6128,18 @@ checksum = "f3faaf9e727533a19351a43cc5a8de957372163c7d35cc48c90b75cdda13c352" dependencies = [ "percent-encoding", "serde", - "thiserror 2.0.18", + "thiserror 2.0.20", ] [[package]] name = "serde_repr" -version = "0.1.20" +version = "0.1.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" +checksum = "8d3b1629de253c70a0508c3899572da79ca359fdab27c7920ff00406df418906" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -6133,7 +6151,7 @@ dependencies = [ "proc-macro2", "quote", "serde", - "syn", + "syn 2.0.119", ] [[package]] @@ -6150,9 +6168,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.6" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", "cpufeatures 0.2.17", @@ -6177,7 +6195,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "digest 0.11.3", ] @@ -6192,15 +6210,37 @@ dependencies = [ [[package]] name = "shlex" -version = "1.3.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "shuttle" -version = "0.9.1" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba93071c1b720be2505f4c8ce2863502cb9a26a3819e268df1932458a755152c" +checksum = "786792d8bc94c770b53a938f0ae68726387a3a0a3762f9cc2c38d5e1bb40f0c0" +dependencies = [ + "bitvec", + "cfg-if", + "const-siphasher", + "corosensei", + "hex", + "owo-colors", + "rand 0.8.8", + "rand_core 0.6.4", + "rand_pcg 0.3.1", + "scoped-tls", + "shuttle-engine", + "shuttle-schedulers", + "shuttle-std", + "tracing", +] + +[[package]] +name = "shuttle-engine" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb03d6daa3cf319f53bef382ee4431203c770912209703452033d19fb71e813b" dependencies = [ "assoc", "bitvec", @@ -6209,19 +6249,45 @@ dependencies = [ "corosensei", "hex", "owo-colors", - "rand 0.8.6", + "rand 0.8.8", "rand_core 0.6.4", - "rand_pcg", + "rand_pcg 0.3.1", "scoped-tls", "smallvec", "tracing", ] +[[package]] +name = "shuttle-schedulers" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7673a579e1660404067f5aa1344d6530f15b59f53b0fe98767d7ab117d14067" +dependencies = [ + "rand 0.8.8", + "rand_pcg 0.3.1", + "shuttle-engine", + "smallvec", + "tracing", +] + +[[package]] +name = "shuttle-std" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf4a3088fd2d19e5ebc24bc3779bc85a68228fafde2d87d138e475b8591dc721" +dependencies = [ + "assoc", + "owo-colors", + "shuttle-engine", + "smallvec", + "tracing", +] + [[package]] name = "simd-adler32" -version = "0.3.9" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" [[package]] name = "simdutf8" @@ -6237,9 +6303,9 @@ checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" [[package]] name = "siphasher" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" [[package]] name = "slab" @@ -6264,7 +6330,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb251b407f50028476a600541542b605bb864d35d9ee1de4f6cab45d88475e6d" dependencies = [ "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -6294,21 +6360,21 @@ checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" [[package]] name = "smallvec" -version = "1.15.1" +version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" [[package]] name = "snap" -version = "1.1.1" +version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" [[package]] name = "socket2" -version = "0.6.3" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", @@ -6316,15 +6382,15 @@ dependencies = [ [[package]] name = "spin" -version = "0.9.8" +version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" [[package]] name = "spin" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5fe4ccb98d9c292d56fec89a5e07da7fc4cf0dc11e156b41793132775d3e591" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" dependencies = [ "lock_api", ] @@ -6348,7 +6414,7 @@ checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -6359,9 +6425,9 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stacker" -version = "0.1.24" +version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "640c8cdd92b6b12f5bcb1803ca3bbf5ab96e5e6b6b96b9ab77dabe9e880b3190" +checksum = "707f49d46706bacf8a2b00d51dace3f9de527c13eec3778f570c411f89e69967" dependencies = [ "cc", "cfg-if", @@ -6372,9 +6438,9 @@ dependencies = [ [[package]] name = "str_stack" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9091b6114800a5f2141aee1d1b9d6ca3592ac062dc5decb3764ec5895a47b4eb" +checksum = "7f446288b699d66d0fd2e30d1cfe7869194312524b3b9252594868ed26ef056a" [[package]] name = "strsim" @@ -6384,9 +6450,9 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "subsecond" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "feae81a4a7ca6d0bcf70c385a43b7dbacbff527f0805cb0a4043ce2c2c559a2c" +checksum = "d350d5788fa94d560d92269266a50efc77b036bfca6675e420b64df7b3211f37" dependencies = [ "js-sys", "libc", @@ -6395,7 +6461,7 @@ dependencies = [ "memmap2", "serde", "subsecond-types", - "thiserror 2.0.18", + "thiserror 2.0.20", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", @@ -6403,9 +6469,9 @@ dependencies = [ [[package]] name = "subsecond-types" -version = "0.7.6" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85256ee192cbdf00473e48e6133863b125dd4f772fddfbc97287ec7a61458c25" +checksum = "dcf32d66269b5fbb8558334e8d22b657b2f7af0dd2ef56c231541099a862e464" dependencies = [ "serde", ] @@ -6416,84 +6482,6 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" -[[package]] -name = "sval" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2eb9318255ebd817902d7e279d8f8e39b35b1b9954decd5eb9ea0e30e5fd2b6a" - -[[package]] -name = "sval_buffer" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12571299185e653fdb0fbfe36cd7f6529d39d4e747a60b15a3f34574b7b97c61" -dependencies = [ - "sval", - "sval_ref", -] - -[[package]] -name = "sval_dynamic" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39526f24e997706c0de7f03fb7371f7f5638b66a504ded508e20ad173d0a3677" -dependencies = [ - "sval", -] - -[[package]] -name = "sval_fmt" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "933dd3bb26965d682280fcc49400ac2a05036f4ee1e6dbd61bf8402d5a5c3a54" -dependencies = [ - "itoa", - "ryu", - "sval", -] - -[[package]] -name = "sval_json" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0cda08f6d5c9948024a6551077557b1fdcc3880ff2f20ae839667d2ec2d87ed" -dependencies = [ - "itoa", - "ryu", - "sval", -] - -[[package]] -name = "sval_nested" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88d49d5e6c1f9fd0e53515819b03a97ca4eb1bff5c8ee097c43391c09ecfb19f" -dependencies = [ - "sval", - "sval_buffer", - "sval_ref", -] - -[[package]] -name = "sval_ref" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14f876c5a78405375b4e19cbb9554407513b59c93dea12dc6a4af4e1d30899ca" -dependencies = [ - "sval", -] - -[[package]] -name = "sval_serde" -version = "2.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f9ccd3b7f7200239a655e517dd3fd48d960b9111ad24bd6a5e055bef17607c7" -dependencies = [ - "serde_core", - "sval", - "sval_nested", -] - [[package]] name = "symbolic-common" version = "12.18.3" @@ -6519,9 +6507,20 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.117" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" dependencies = [ "proc-macro2", "quote", @@ -6545,20 +6544,21 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "sysinfo" -version = "0.38.4" +version = "0.39.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ab6a2f8bfe508deb3c6406578252e491d299cbbf3bc0529ecc3313aee4a52f" +checksum = "d2071df9448915b71c4fe6d25deaf1c22f12bd234f01540b77312bb8e41361e6" dependencies = [ "libc", "memchr", "ntapi", "objc2-core-foundation", "objc2-io-kit", + "objc2-open-directory", "windows", ] @@ -6568,7 +6568,7 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "core-foundation 0.9.4", "system-configuration-sys", ] @@ -6585,9 +6585,9 @@ dependencies = [ [[package]] name = "t4" -version = "0.1.7" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3105026c5fb7e1dcf8864c2193387439aa650b690e1a5513ab27ae7de7999a9" +checksum = "5d116b5de6db7e42465e0b9807038609248e536e84bc76167b65622dac1dc635" dependencies = [ "cfg-if", "crossbeam-channel", @@ -6601,9 +6601,9 @@ dependencies = [ [[package]] name = "t4-verified" -version = "0.1.7" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "494c2f2c4cfb21784ef9b380c378fbfdbbe2f19f1ad606e6729816ea2a0f4297" +checksum = "9610c43f95b308dc0f5d717b84698783869980c9ca23ce19f78ec7b72a5dbf2b" dependencies = [ "vstd", ] @@ -6621,7 +6621,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -6648,11 +6648,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.18", + "thiserror-impl 2.0.20", ] [[package]] @@ -6663,18 +6663,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -6689,9 +6689,9 @@ dependencies = [ [[package]] name = "thread_local" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" dependencies = [ "cfg-if", ] @@ -6709,12 +6709,11 @@ dependencies = [ [[package]] name = "time" -version = "0.3.47" +version = "0.3.55" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134" dependencies = [ "deranged", - "itoa", "num-conv", "powerfmt", "serde_core", @@ -6724,15 +6723,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.8" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" [[package]] name = "time-macros" -version = "0.2.27" +version = "0.2.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +checksum = "7e689342a48d2ea927c87ea50cabf8594854bf940e9310208848d680d668ed85" dependencies = [ "num-conv", "time-core", @@ -6749,9 +6748,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" dependencies = [ "displaydoc", "zerovec", @@ -6759,9 +6758,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -6774,9 +6773,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" dependencies = [ "bytes", "libc", @@ -6789,13 +6788,13 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.7.0" +version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] @@ -6810,9 +6809,9 @@ dependencies = [ [[package]] name = "tokio-stream" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" dependencies = [ "futures-core", "pin-project-lite", @@ -6857,15 +6856,16 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.18" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" dependencies = [ "bytes", "futures-core", "futures-io", "futures-sink", "futures-util", + "libc", "pin-project-lite", "tokio", ] @@ -6881,23 +6881,23 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.25.11+spec-1.1.0" +version = "0.25.13+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" dependencies = [ - "indexmap 2.14.0", + "indexmap", "toml_datetime", "toml_parser", - "winnow 1.0.2", + "winnow 1.0.4", ] [[package]] name = "toml_parser" -version = "1.1.2+spec-1.1.0" +version = "1.1.3+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" dependencies = [ - "winnow 1.0.2", + "winnow 1.0.4", ] [[package]] @@ -6908,7 +6908,7 @@ checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", - "base64 0.22.1", + "base64", "bytes", "h2", "http", @@ -6931,15 +6931,26 @@ dependencies = [ [[package]] name = "tonic-prost" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", "prost", "tonic", ] +[[package]] +name = "tonic-types" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73ab1b02061f83d519bba3caa167f88f261ef05720ab8ebc954ade70de3348e8" +dependencies = [ + "prost", + "prost-types", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -6948,7 +6959,7 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "indexmap 2.14.0", + "indexmap", "pin-project-lite", "slab", "sync_wrapper", @@ -6965,7 +6976,7 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ - "bitflags 2.11.1", + "bitflags 2.13.1", "bytes", "futures-core", "futures-util", @@ -6986,6 +6997,21 @@ dependencies = [ "url", ] +[[package]] +name = "tower-http" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" +dependencies = [ + "bitflags 2.13.1", + "bytes", + "http", + "percent-encoding", + "pin-project-lite", + "tower-layer", + "tower-service", +] + [[package]] name = "tower-layer" version = "0.3.3" @@ -7018,7 +7044,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -7098,9 +7124,9 @@ dependencies = [ "http", "httparse", "log", - "rand 0.9.4", + "rand 0.9.5", "sha1", - "thiserror 2.0.18", + "thiserror 2.0.20", "utf-8", ] @@ -7115,9 +7141,9 @@ dependencies = [ "http", "httparse", "log", - "rand 0.9.4", + "rand 0.9.5", "sha1", - "thiserror 2.0.18", + "thiserror 2.0.20", "utf-8", ] @@ -7132,28 +7158,22 @@ dependencies = [ "http", "httparse", "log", - "rand 0.9.4", + "rand 0.9.5", "sha1", - "thiserror 2.0.18", + "thiserror 2.0.20", ] [[package]] name = "twox-hash" -version = "2.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" - -[[package]] -name = "typeid" -version = "1.0.3" +version = "2.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" +checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a" [[package]] name = "typenum" -version = "1.20.0" +version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" [[package]] name = "ucd-trie" @@ -7175,9 +7195,9 @@ checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" -version = "1.13.2" +version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" [[package]] name = "unicode-width" @@ -7234,7 +7254,7 @@ dependencies = [ "proc-macro2", "quote", "serde_tokenstream", - "syn", + "syn 2.0.119", "usdt-impl", ] @@ -7252,8 +7272,8 @@ dependencies = [ "quote", "serde", "serde_json", - "syn", - "thiserror 2.0.18", + "syn 2.0.119", + "thiserror 2.0.20", "thread-id", ] @@ -7267,7 +7287,7 @@ dependencies = [ "proc-macro2", "quote", "serde_tokenstream", - "syn", + "syn 2.0.119", "usdt-impl", ] @@ -7291,11 +7311,11 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.3" +version = "1.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7" +checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812" dependencies = [ - "getrandom 0.4.2", + "getrandom 0.4.3", "js-sys", "wasm-bindgen", ] @@ -7306,43 +7326,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" -[[package]] -name = "value-bag" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ba6f5989077681266825251a52748b8c1d8a4ad098cc37e440103d0ea717fc0" -dependencies = [ - "value-bag-serde1", - "value-bag-sval2", -] - -[[package]] -name = "value-bag-serde1" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16530907bfe2999a1773ca5900a65101e092c70f642f25cc23ca0c43573262c5" -dependencies = [ - "erased-serde", - "serde_buf", - "serde_core", - "serde_fmt", -] - -[[package]] -name = "value-bag-sval2" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d00ae130edd690eaa877e4f40605d534790d1cf1d651e7685bd6a144521b251f" -dependencies = [ - "sval", - "sval_buffer", - "sval_dynamic", - "sval_fmt", - "sval_json", - "sval_ref", - "sval_serde", -] - [[package]] name = "version_check" version = "0.9.5" @@ -7351,19 +7334,20 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "verus_builtin" -version = "0.0.0-2026-05-06-1803" +version = "0.0.0-2026-08-30-0159" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "650bcc71ef90cbc79bc2a2494c54dd159edad63a72677be791954edc993883fe" +checksum = "300d269a2e06dbe54cb53084464065912ad7accbb52ae8d6c8aa7a6c521974c5" [[package]] name = "verus_builtin_macros" -version = "0.0.0-2026-05-10-0145" +version = "0.0.0-2026-08-30-0159" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef946d78c84f284991d59035ca252f29cd8071526003e794be80f0beee18f2d" +checksum = "cc040bfded82e79a708c1819c52b386ee23c0c823239334f39ced152bbdf3624" dependencies = [ + "convert_case 0.4.0", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", "verus_prettyplease", "verus_syn", @@ -7371,9 +7355,9 @@ dependencies = [ [[package]] name = "verus_prettyplease" -version = "0.0.0-2026-05-10-0145" +version = "0.0.0-2026-08-09-0044" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a227e7eaa03f51ac4d659641192bc1b9fe4eda0a509a6734752cf72f5d7daaf" +checksum = "51fc115de5fb3806bc362060bb683024660ed2bc13c1183d1f01d464e4537139" dependencies = [ "proc-macro2", "verus_syn", @@ -7381,11 +7365,11 @@ dependencies = [ [[package]] name = "verus_state_machines_macros" -version = "0.0.0-2026-05-10-0145" +version = "0.0.0-2026-08-02-0125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ff46046bffd2d55757503feef0b72a74a60ead2b58aa211861267485b955136" +checksum = "93f77ff8d121edb1bf2651335769a80a2f611da8573c3b498434a66b3162805f" dependencies = [ - "indexmap 1.9.3", + "indexmap", "proc-macro2", "quote", "verus_syn", @@ -7393,9 +7377,9 @@ dependencies = [ [[package]] name = "verus_syn" -version = "0.0.0-2026-05-10-0145" +version = "0.0.0-2026-08-02-0125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a82174e474e1f06418dd304f714785ae1c3e4d2fbff415a0d0c05a02cd995819" +checksum = "f17237ea6d267e457d53ce36f55b315fc9edd26eb88c765dd95e035c0b415869" dependencies = [ "proc-macro2", "quote", @@ -7404,9 +7388,9 @@ dependencies = [ [[package]] name = "vstd" -version = "0.0.0-2026-05-10-0145" +version = "0.0.0-2026-08-30-0159" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6958222eaf7bbe565c93890676bb1f96c4f5e25d90a81bb9993fdc19fc0d3f72" +checksum = "7169790c92857fa9441c8ba7e0c469f4fc62104caf1dfc632c318577a1a08d7a" dependencies = [ "verus_builtin", "verus_builtin_macros", @@ -7451,7 +7435,7 @@ checksum = "59195a1db0e95b920366d949ba5e0d3fc0e70b67c09be15ce5abb790106b0571" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -7462,27 +7446,18 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasip2" -version = "1.0.3+wasi-0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" -dependencies = [ - "wit-bindgen 0.57.1", -] - -[[package]] -name = "wasip3" -version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ - "wit-bindgen 0.51.0", + "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -7493,9 +7468,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.68" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f371d383f2fb139252e0bfac3b81b265689bf45b6874af544ffa4c975ac1ebf8" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -7503,9 +7478,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -7513,48 +7488,26 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.118" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] -[[package]] -name = "wasm-encoder" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" -dependencies = [ - "leb128fmt", - "wasmparser", -] - -[[package]] -name = "wasm-metadata" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" -dependencies = [ - "anyhow", - "indexmap 2.14.0", - "wasm-encoder", - "wasmparser", -] - [[package]] name = "wasm-streams" version = "0.4.2" @@ -7568,23 +7521,11 @@ dependencies = [ "web-sys", ] -[[package]] -name = "wasmparser" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" -dependencies = [ - "bitflags 2.11.1", - "hashbrown 0.15.5", - "indexmap 2.14.0", - "semver", -] - [[package]] name = "web-sys" -version = "0.3.95" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" dependencies = [ "js-sys", "wasm-bindgen", @@ -7602,9 +7543,9 @@ dependencies = [ [[package]] name = "webpki-roots" -version = "1.0.7" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52f5ee44c96cf55f1b349600768e3ece3a8f26010c05265ab73f945bb1a2eb9d" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -7693,7 +7634,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -7704,7 +7645,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -7929,112 +7870,24 @@ dependencies = [ [[package]] name = "winnow" -version = "1.0.2" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" dependencies = [ "memchr", ] -[[package]] -name = "wit-bindgen" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" -dependencies = [ - "wit-bindgen-rust-macro", -] - [[package]] name = "wit-bindgen" version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" -[[package]] -name = "wit-bindgen-core" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" -dependencies = [ - "anyhow", - "heck", - "wit-parser", -] - -[[package]] -name = "wit-bindgen-rust" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" -dependencies = [ - "anyhow", - "heck", - "indexmap 2.14.0", - "prettyplease", - "syn", - "wasm-metadata", - "wit-bindgen-core", - "wit-component", -] - -[[package]] -name = "wit-bindgen-rust-macro" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" -dependencies = [ - "anyhow", - "prettyplease", - "proc-macro2", - "quote", - "syn", - "wit-bindgen-core", - "wit-bindgen-rust", -] - -[[package]] -name = "wit-component" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" -dependencies = [ - "anyhow", - "bitflags 2.11.1", - "indexmap 2.14.0", - "log", - "serde", - "serde_derive", - "serde_json", - "wasm-encoder", - "wasm-metadata", - "wasmparser", - "wit-parser", -] - -[[package]] -name = "wit-parser" -version = "0.244.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" -dependencies = [ - "anyhow", - "id-arena", - "indexmap 2.14.0", - "log", - "semver", - "serde", - "serde_derive", - "serde_json", - "unicode-xid", - "wasmparser", -] - [[package]] name = "writeable" -version = "0.6.3" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" [[package]] name = "wyz" @@ -8047,15 +7900,15 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" [[package]] name = "yoke" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -8070,35 +7923,35 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.48" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.48" +version = "0.8.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "zerofrom" -version = "0.1.7" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" dependencies = [ "zerofrom-derive", ] @@ -8111,21 +7964,21 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zeroize" -version = "1.8.2" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" [[package]] name = "zerotrie" -version = "0.2.4" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f" dependencies = [ "displaydoc", "yoke", @@ -8134,9 +7987,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.6" +version = "0.11.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8" dependencies = [ "yoke", "zerofrom", @@ -8145,26 +7998,26 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.3" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.4", ] [[package]] name = "zlib-rs" -version = "0.6.3" +version = "0.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" [[package]] name = "zstd" diff --git a/Cargo.toml b/Cargo.toml index 9fcc379c5..c0c4a235e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,25 +30,25 @@ liquid-cache-datafusion = { path = "src/datafusion", version = "0.1.13" } liquid-cache-common = { path = "src/common", version = "0.1.13" } liquid-cache = { path = "src/core", version = "0.1.13" } liquid-cache-datafusion-local = { path = "src/datafusion-local", version = "0.1.13" } -arrow = { version = "58.3.0", default-features = false, features = [ +arrow = { version = "58.4.0", default-features = false, features = [ "prettyprint", "ipc", ] } -arrow-flight = { version = "58.3.0", features = ["flight-sql-experimental"] } -arrow-schema = { version = "58.3.0", features = ["serde"] } -parquet = { version = "58.3.0", features = [ +arrow-flight = { version = "58.4.0", features = ["flight-sql-experimental"] } +arrow-schema = { version = "58.4.0", features = ["serde"] } +parquet = { version = "58.4.0", features = [ "async", "experimental", "variant_experimental", ] } -parquet-variant-json = { version = "58.3.0" } -parquet-variant-compute = { version = "58.3.0" } -datafusion = { version = "54.0.0" } -datafusion-common = { version = "54.0.0" } -datafusion-expr-common = { version = "54.0.0" } -datafusion-physical-expr = { version = "54.0.0" } -datafusion-physical-expr-common = { version = "54.0.0" } -datafusion-proto = { version = "54.0.0" } +parquet-variant-json = { version = "58.4.0" } +parquet-variant-compute = { version = "58.4.0" } +datafusion = { version = "54.1.0" } +datafusion-common = { version = "54.1.0" } +datafusion-expr-common = { version = "54.1.0" } +datafusion-physical-expr = { version = "54.1.0" } +datafusion-physical-expr-common = { version = "54.1.0" } +datafusion-proto = { version = "54.1.0" } async-trait = "0.1.89" futures = { version = "0.3.32", default-features = false, features = ["std"] } tokio = { version = "1.52.3", features = ["rt-multi-thread"] } @@ -67,7 +67,7 @@ fastrace = "0.7" fastrace-tonic = "0.2" congee = "0.4.1" insta = "1.47.2" -t4 = "0.1.7" +t4 = "0.1.9" [profile.dev.package] insta.opt-level = 3 diff --git a/README.md b/README.md index 98dcfa526..6515ed31f 100644 --- a/README.md +++ b/README.md @@ -93,7 +93,7 @@ tokio_test::block_on(async { ### LiquidCache uses DIRECT I/O -By default, LiquidCache uses [DIRECT I/O](https://man7.org/linux/man-pages/man2/open.2.html#:~:text=O_DIRECT). This means that it bypasses the OS page cache, this avoids double-caching and bound memory usage. +By default, LiquidCache bypasses the OS page cache using [O_DIRECT](https://man7.org/linux/man-pages/man2/open.2.html#:~:text=O_DIRECT) on Linux and `F_NOCACHE` on macOS. This avoids double-caching and bounds memory usage. This also means LiquidCache can *appear slower* than other caches when most data fits in OS page cache, which is common in dev environments but unrealistic in production. diff --git a/benchmark/Cargo.toml b/benchmark/Cargo.toml index 25eb33de4..ce7e7ad00 100644 --- a/benchmark/Cargo.toml +++ b/benchmark/Cargo.toml @@ -21,7 +21,7 @@ url = { workspace = true } mimalloc = "0.1.52" serde_json.workspace = true serde.workspace = true -sysinfo = { version = "0.38.4", default-features = false, features = [ +sysinfo = { version = "0.39.6", default-features = false, features = [ "network", "disk", ] } @@ -30,17 +30,19 @@ parquet = { workspace = true } arrow = { workspace = true } fastrace = { version = "0.7.17" } fastrace-tonic = { workspace = true } -fastrace-opentelemetry = "0.16" -opentelemetry = "0.31.0" -opentelemetry_sdk = "0.31.0" -opentelemetry-otlp = { version = "0.31.1", features = ["trace", "grpc-tonic"] } -logforth = { version = "0.29.1", features = ["append-opentelemetry", "bridge-log"] } +fastrace-opentelemetry = "0.18" +opentelemetry = "0.32.0" +opentelemetry_sdk = "0.32.0" +opentelemetry-otlp = { version = "0.32.0", features = ["trace", "grpc-tonic"] } +logforth = { version = "0.30.1", features = ["starter-log", "filter-rustlog"] } reqwest = { version = "0.13.4", default-features = false, features = ["json"] } uuid = { version = "1.23.3", features = ["v4"] } pprof = { version = "0.15.0", features = ["flamegraph"] } anyhow = "1.0" usdt = "0.6" regex = "1.12.4" + +[target.'cfg(target_os = "linux")'.dependencies] perf-event2 = "0.7.4" [features] diff --git a/benchmark/src/inprocess_runner.rs b/benchmark/src/inprocess_runner.rs index 4ef6f57db..0aac7d495 100644 --- a/benchmark/src/inprocess_runner.rs +++ b/benchmark/src/inprocess_runner.rs @@ -14,6 +14,7 @@ use liquid_cache::cache_policies::LiquidPolicy; use liquid_cache_datafusion::{LiquidCacheParquetRef, extract_execution_metrics}; use liquid_cache_datafusion_local::LiquidCacheLocalBuilder; use log::{info, warn}; +#[cfg(target_os = "linux")] use perf_event::{ Builder as PerfBuilder, Counter, Group, events::{Hardware, Software}, @@ -72,6 +73,27 @@ impl DiskIoGuard { } } +// perf events are backed by Linux's perf_event_open; on other platforms the +// collector fails to initialize and the runner skips the counters. +#[cfg(not(target_os = "linux"))] +struct PerfEventCollector; + +#[cfg(not(target_os = "linux"))] +impl PerfEventCollector { + fn new() -> io::Result { + Err(io::Error::other("perf events are only supported on Linux")) + } + + fn start(&mut self) -> io::Result<()> { + Ok(()) + } + + fn stop(self) -> io::Result { + Err(io::Error::other("perf events are only supported on Linux")) + } +} + +#[cfg(target_os = "linux")] struct PerfEventCollector { group: Group, cycles: Counter, @@ -82,6 +104,7 @@ struct PerfEventCollector { page_faults: Counter, } +#[cfg(target_os = "linux")] impl PerfEventCollector { fn new() -> io::Result { let mut group = Group::new()?; diff --git a/benchmark/src/observability.rs b/benchmark/src/observability.rs index d33d3d455..aff918c84 100644 --- a/benchmark/src/observability.rs +++ b/benchmark/src/observability.rs @@ -7,7 +7,7 @@ use datafusion::datasource::source::DataSource; use datafusion::physical_plan::ExecutionPlan; use fastrace_opentelemetry::OpenTelemetryReporter; use liquid_cache_datafusion::LiquidParquetSource; -use logforth::filter::env_filter::EnvFilterBuilder; +use logforth::filter::rustlog::RustLogFilterBuilder; use opentelemetry::InstrumentationScope; use opentelemetry::KeyValue; use opentelemetry_otlp::SpanExporter; @@ -54,7 +54,7 @@ pub fn instrument_liquid_source_with_span( pub fn setup_observability(service_name: &str, jaeger_endpoint: Option<&str>) { logforth::starter_log::builder() .dispatch(|d| { - d.filter(EnvFilterBuilder::from_default_env().build()) + d.filter(RustLogFilterBuilder::from_default_env().build()) .append(logforth::append::Stdout::default()) }) .apply(); diff --git a/dev/README.md b/dev/README.md index 72aa3c031..8257cb6b8 100644 --- a/dev/README.md +++ b/dev/README.md @@ -9,6 +9,12 @@ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh ``` +Alternatively, use the Nix dev shell (works on both Linux and macOS; also provides the tooling for `dev-tools`, e.g. dioxus-cli, tailwindcss, wasm-bindgen): + +```bash +nix develop +``` + Run tests: ```bash diff --git a/dev/dev-tools/Cargo.toml b/dev/dev-tools/Cargo.toml index 879412d3e..b4667e1b8 100644 --- a/dev/dev-tools/Cargo.toml +++ b/dev/dev-tools/Cargo.toml @@ -7,6 +7,12 @@ edition = "2024" [dependencies] dioxus = { version = "=0.7.5", features = ["router", "fullstack"] } +# Must match the wasm-bindgen-cli version pinned in flake.nix. +# Only used through dioxus macros, hence the cargo-shear ignore. +wasm-bindgen = "=0.2.126" + +[package.metadata.cargo-shear] +ignored = ["wasm-bindgen"] [features] default = ["web"] diff --git a/dev/dev-tools/src/components/cache_state_view.rs b/dev/dev-tools/src/components/cache_state_view.rs index e909f1756..fe8d4a6d9 100644 --- a/dev/dev-tools/src/components/cache_state_view.rs +++ b/dev/dev-tools/src/components/cache_state_view.rs @@ -50,7 +50,7 @@ pub fn CacheStateView(simulator: Signal) -> Element { // Create a unified sorted list of all entry IDs (both actual and failed) let mut all_entry_ids: Vec = entries.iter().map(|e| e.entry_id).collect(); - for (entry_id, _) in state.failed_inserts.iter() { + for entry_id in state.failed_inserts.keys() { if !all_entry_ids.contains(entry_id) { all_entry_ids.push(*entry_id); } diff --git a/flake.lock b/flake.lock index 91d09a14a..3a7310330 100644 --- a/flake.lock +++ b/flake.lock @@ -20,11 +20,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1777268161, - "narHash": "sha256-bxrdOn8SCOv8tN4JbTF/TXq7kjo9ag4M+C8yzzIRYbE=", + "lastModified": 1788039129, + "narHash": "sha256-pa4Q0qErvCvzCaaUph7Sm37RhR4xvPrYI8Lgz6k85+A=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "1c3fe55ad329cbcb28471bb30f05c9827f724c76", + "rev": "d2f67949798825fe853f7c5d0492b8bf016d3f88", "type": "github" }, "original": { @@ -62,11 +62,11 @@ "nixpkgs": "nixpkgs_2" }, "locked": { - "lastModified": 1777432579, - "narHash": "sha256-Ce11TStDsqCge2vAAfLKe2+4lDI5cSX5ZYZOuKJBKKQ=", + "lastModified": 1788165049, + "narHash": "sha256-en4IoUeCqvq9F66YhwOrUFw1nc70OBhrJwrzfalezvY=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "3ecb5e6ab380ced3272ef7fcfe398bffbcc0f152", + "rev": "d03cd474bd97389dcc2e8cd3b3bb6b8c6e346b1a", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index ef8360ae8..7281d6b15 100644 --- a/flake.nix +++ b/flake.nix @@ -41,20 +41,21 @@ llvmPackages.bintools lldb cargo-fuzz - bpftrace - perf nixd inferno cargo-flamegraph nodejs tailwindcss_4 dioxus-cli - wasm-bindgen-cli_0_2_118 + wasm-bindgen-cli_0_2_126 binaryen (rust-bin.selectLatestNightlyWith (toolchain: toolchain.default.override { extensions = [ "rust-src" "llvm-tools-preview" ]; - targets = [ "x86_64-unknown-linux-gnu" "wasm32-unknown-unknown" ]; + targets = [ "wasm32-unknown-unknown" ]; })) + ] ++ lib.optionals stdenv.hostPlatform.isLinux [ + bpftrace + perf ]; shellHook = '' diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml index d8e2b9f49..1a83b0e2f 100644 --- a/src/core/Cargo.toml +++ b/src/core/Cargo.toml @@ -21,9 +21,9 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } arrow = { workspace = true } arrow-schema = { workspace = true } -fastlanes = "0.5.1" +fastlanes = "0.7.0" num-traits = "0.2.19" -fsst-rs = "0.5.11" +fsst-rs = "0.6.0" ahash = { workspace = true } tempfile = { workspace = true } congee = { workspace = true } @@ -34,7 +34,7 @@ parquet-variant-compute = { workspace = true } fastrace = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -sysinfo = { version = "0.38.4", default-features = false, features = ["system"] } +sysinfo = { version = "0.39.6", default-features = false, features = ["system"] } [dev-dependencies] tempfile = { workspace = true } diff --git a/src/core/src/liquid_array/byte_view_array/comparisons.rs b/src/core/src/liquid_array/byte_view_array/comparisons.rs index 2fd21e6d2..38b5efc7d 100644 --- a/src/core/src/liquid_array/byte_view_array/comparisons.rs +++ b/src/core/src/liquid_array/byte_view_array/comparisons.rs @@ -540,8 +540,10 @@ fn compare_with_arrow_inner( fn compress_needle(compressor: &Compressor, needle: &[u8]) -> Vec { let mut compressed = Vec::with_capacity(needle.len().saturating_mul(2)); + // SAFETY: the largest compressed size is all escapes == 2 * plaintext_len. unsafe { - compressor.compress_into(needle, &mut compressed); + let len = compressor.compress_into(needle, compressed.spare_capacity_mut()); + compressed.set_len(len); } compressed } diff --git a/src/core/src/liquid_array/byte_view_array/serialization.rs b/src/core/src/liquid_array/byte_view_array/serialization.rs index 94777f8c1..f37870a57 100644 --- a/src/core/src/liquid_array/byte_view_array/serialization.rs +++ b/src/core/src/liquid_array/byte_view_array/serialization.rs @@ -308,8 +308,8 @@ impl LiquidByteViewArray { panic!("Fingerprint data size does not match dictionary size"); } let mut fingerprints = Vec::with_capacity(view_header.fingerprint_size as usize / 4); - for chunk in bytes[cursor..fingerprint_end].chunks_exact(4) { - fingerprints.push(u32::from_le_bytes(chunk.try_into().unwrap())); + for chunk in bytes[cursor..fingerprint_end].as_chunks::<4>().0 { + fingerprints.push(u32::from_le_bytes(*chunk)); } Some(Arc::from(fingerprints.into_boxed_slice())) }; diff --git a/src/core/src/liquid_array/raw/fsst_buffer.rs b/src/core/src/liquid_array/raw/fsst_buffer.rs index 4e64b6f4d..0789623ff 100644 --- a/src/core/src/liquid_array/raw/fsst_buffer.rs +++ b/src/core/src/liquid_array/raw/fsst_buffer.rs @@ -70,7 +70,8 @@ impl RawFsstBuffer { // (all bytes escaped) which is `2 * plaintext_len`. compress_buffer.reserve(bytes.len().saturating_mul(2)); unsafe { - compressor.compress_into(bytes, compress_buffer); + let len = compressor.compress_into(bytes, compress_buffer.spare_capacity_mut()); + compress_buffer.set_len(len); } values_buffer.extend_from_slice(compress_buffer); diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index f92b9852d..2bc020214 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -413,6 +413,13 @@ async fn test_provide_schema2() { } } + // FSST breaks equal-gain symbol ties using target-specific HashMap iteration order. + // Canonicalize the known AArch64 totals to x86_64 while leaving unexpected totals visible. + #[cfg(target_arch = "aarch64")] + let snapshot = snapshot + .replace("usage.memory_bytes: 999980", "usage.memory_bytes: 1000915") + .replace("usage.memory_bytes: 1035369", "usage.memory_bytes: 1036304"); + insta::assert_snapshot!(snapshot); } diff --git a/src/datafusion-server/Cargo.toml b/src/datafusion-server/Cargo.toml index 192d1563e..67549800c 100644 --- a/src/datafusion-server/Cargo.toml +++ b/src/datafusion-server/Cargo.toml @@ -25,8 +25,8 @@ liquid-cache-common = { workspace = true } tempfile = { workspace = true } axum = "0.8.9" serde = { workspace = true } -tower-http = { version = "0.6.11", features = ["cors"] } -sysinfo = { version = "0.38.4", default-features = false, features = [ +tower-http = { version = "0.7.1", features = ["cors"] } +sysinfo = { version = "0.39.6", default-features = false, features = [ "component", "disk", "network", diff --git a/src/datafusion/src/utils.rs b/src/datafusion/src/utils.rs index 7cc144271..e6d903ccf 100644 --- a/src/datafusion/src/utils.rs +++ b/src/datafusion/src/utils.rs @@ -310,6 +310,7 @@ pub fn extract_execution_metrics( #[cfg(test)] mod tests { + #[cfg(target_arch = "x86_64")] use super::*; #[test] From 1909d083e4691b3f30e3bae6dae5cf9089297923 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Tue, 1 Sep 2026 11:07:19 -0400 Subject: [PATCH 02/24] datafusion 55 (#510) --- Cargo.lock | 395 ++++++++++-------- Cargo.toml | 24 +- dev/dev-tools/Cargo.toml | 2 +- src/core/src/cache/policies/squeeze.rs | 4 +- src/datafusion-client/src/client_exec.rs | 60 ++- src/datafusion-client/src/lib.rs | 3 +- src/datafusion-client/src/optimizer.rs | 15 +- ...che_datafusion_client__tests__tpch_q1.snap | 48 +-- ...he_datafusion_client__tests__tpch_q10.snap | 37 +- ...he_datafusion_client__tests__tpch_q11.snap | 18 +- ...he_datafusion_client__tests__tpch_q12.snap | 10 +- ...he_datafusion_client__tests__tpch_q13.snap | 12 +- ...he_datafusion_client__tests__tpch_q14.snap | 12 +- ...he_datafusion_client__tests__tpch_q16.snap | 31 +- ...he_datafusion_client__tests__tpch_q17.snap | 6 +- ...he_datafusion_client__tests__tpch_q18.snap | 2 +- ...he_datafusion_client__tests__tpch_q19.snap | 70 ++-- ...che_datafusion_client__tests__tpch_q2.snap | 6 +- ...he_datafusion_client__tests__tpch_q20.snap | 6 +- ...he_datafusion_client__tests__tpch_q21.snap | 168 ++++---- ...he_datafusion_client__tests__tpch_q22.snap | 24 +- ...che_datafusion_client__tests__tpch_q3.snap | 41 +- ...che_datafusion_client__tests__tpch_q4.snap | 12 +- ...che_datafusion_client__tests__tpch_q5.snap | 33 +- ...che_datafusion_client__tests__tpch_q6.snap | 7 +- ...che_datafusion_client__tests__tpch_q7.snap | 20 +- ...che_datafusion_client__tests__tpch_q8.snap | 28 +- ...che_datafusion_client__tests__tpch_q9.snap | 16 +- src/datafusion-local/src/lib.rs | 4 +- src/datafusion-local/src/tests/mod.rs | 68 ++- ...afusion_local__tests__provide_schema2.snap | 21 +- ...al__tests__provide_schema_with_filter.snap | 7 + src/datafusion-local/src/tests/squeeze.rs | 14 +- .../src/admin_server/handlers.rs | 25 +- src/datafusion-server/src/lib.rs | 3 +- src/datafusion/src/reader/plantime/opener.rs | 8 +- src/datafusion/src/reader/plantime/source.rs | 75 +++- 37 files changed, 751 insertions(+), 584 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d496de5ac..ef27b66fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -161,9 +161,9 @@ checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" [[package]] name = "arrow" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6cfdd0833e32a9874d2b55089333ad310c0be208aafa277385ce2461dec90be3" +checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" dependencies = [ "arrow-arith", "arrow-array", @@ -182,9 +182,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a41203398f0eaa6f7ec8e62c0da742a21abf282c148fc157f6c35c90e29981a" +checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -196,9 +196,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae33dad492b7df00a217563a7b0ef2874df68a0deea1b1a3acf628152f7f7a69" +checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" dependencies = [ "ahash", "arrow-buffer", @@ -215,21 +215,21 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9552f96391c005e6ab449fa941420935e7e062489b12b8b1b08879b2163f5b5" +checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.5.1", "num-traits", ] [[package]] name = "arrow-cast" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a8a327c9649f30d8406995f27642b68df354713cca3baaaf100f076f18d5f34" +checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" dependencies = [ "arrow-array", "arrow-buffer", @@ -238,7 +238,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.23.1", "chrono", "comfy-table", "half", @@ -249,9 +249,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af0dd6d90d1955e9f9a014c1e563ee8aeffc21909085d25623e1da44d96eca26" +checksum = "25011b52b346407d497ef0030e12b45e4f2d0cc279efc09c4f3d09106db30e36" dependencies = [ "arrow-array", "arrow-cast", @@ -264,9 +264,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b24852db04738907e06c04ea61e42fe7fda962a34513022dc0d0e754fb7976b" +checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" dependencies = [ "arrow-buffer", "arrow-schema", @@ -277,9 +277,9 @@ dependencies = [ [[package]] name = "arrow-flight" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2dbe34824c639e43136af8f106992792ab456540d54b880bc320a3192502d2e" +checksum = "2bebfacc9d71f0728f6774164e4d4254b5e504d2b46812d0512d8290ec119a64" dependencies = [ "arrow-arith", "arrow-array", @@ -292,11 +292,10 @@ dependencies = [ "arrow-schema", "arrow-select", "arrow-string", - "base64", + "base64 0.23.1", "bytes", "futures", "once_cell", - "paste", "prost", "prost-types", "tonic", @@ -305,9 +304,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29a908a11fcfb3fb2f6730f4ac15e367bc644e419155e96238f68cf3adde572b" +checksum = "149437b14371f5b9ec60f5ddc751483ae99d7a7072653c0075e5e469156eea7b" dependencies = [ "arrow-array", "arrow-buffer", @@ -321,9 +320,9 @@ dependencies = [ [[package]] name = "arrow-json" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8a96aed3931c076adee39ec2a40d8219fc7f09e79bcdaca1df16272993e1e14" +checksum = "f18b9123ccfec418a663f821c9a034af339711678c11ffe00d3ec07da5ff9f7e" dependencies = [ "arrow-array", "arrow-buffer", @@ -346,9 +345,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63a083ec750f5c043f02946b4baf05fcdbb55f4560a3277055caca5cc99f3eb0" +checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" dependencies = [ "arrow-array", "arrow-buffer", @@ -359,9 +358,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "514ba0ef0d4c5896202dae736251ce415abb43a950bed570fb7981b8716c0e4c" +checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" dependencies = [ "arrow-array", "arrow-buffer", @@ -372,9 +371,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21ca356ad6425cecb6eb7b28e4f659f1ee7880fbb1a16127de7dd62901efee9e" +checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" dependencies = [ "serde", "serde_core", @@ -383,9 +382,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c58da39eb3d8350ad4a549e5c2bc49284dac554016c69829310350f1731b0aad" +checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" dependencies = [ "ahash", "arrow-array", @@ -397,9 +396,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6789b388467525e3271326b6b4915666ecfdf5142aef09779445c954b67543c" +checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" dependencies = [ "arrow-array", "arrow-buffer", @@ -514,7 +513,7 @@ checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ "axum-core", "axum-macros", - "base64", + "base64 0.22.1", "bytes", "form_urlencoded", "futures-util", @@ -624,6 +623,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "bigdecimal" version = "0.4.10" @@ -632,7 +637,7 @@ checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.8", "num-integer", "num-traits", ] @@ -825,7 +830,7 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1f927b07c74ba84c7e5fe4db2baeb3e996ab2688992e39ac68ce3220a677c7e" dependencies = [ - "base64", + "base64 0.22.1", "encoding_rs", ] @@ -1379,9 +1384,9 @@ checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" [[package]] name = "datafusion" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "754ef4e8f073922a26f5b23133b9db4829342362b09be0bc94309cf261c2f098" +checksum = "96f76f0167ed0842b29a3d1e41be3c034c0a46409a3a703cc4cc84ee8c24abf4" dependencies = [ "arrow", "arrow-schema", @@ -1416,7 +1421,7 @@ dependencies = [ "flate2", "futures", "indexmap", - "itertools", + "itertools 0.15.0", "liblzma", "log", "object_store", @@ -1432,9 +1437,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06afd1e38dd27bbb1258685a1fc6524df6aff4e07b25b393a47de59635178d99" +checksum = "d79ec3460f6ed5c58f9b3f2d873fbc77748b82653bff1b4cdaf06de33bb4e05f" dependencies = [ "arrow", "async-trait", @@ -1448,7 +1453,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -1457,9 +1462,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0668fb32c12065ec242be0e5b4bc62bd7a06a0be3ecd83791ef877e4be67e02" +checksum = "b48cef241e2efcfd496fe05ae4d0d5de20793451862faefe406c397a467e12d4" dependencies = [ "arrow", "async-trait", @@ -1473,16 +1478,17 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", + "percent-encoding", ] [[package]] name = "datafusion-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca43b263cdff57042cfa8fb817fb3469f4878933380dccff25f5e793580abbf9" +checksum = "3f72810485975c258f1b4d00baab31728470676c60c5546f366ebd0d99f05ab6" dependencies = [ "arrow", "arrow-ipc", @@ -1492,9 +1498,10 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "libc", "log", + "num-traits", "object_store", "parquet", "recursive", @@ -1506,9 +1513,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f0ba2b864792bdca4d76c59a1de0ab6e1b61946596b9936888dbd6360035f2" +checksum = "533c28e75dba52f41bde187d23a1cb24ab91c7c097966824fa471e67b60320ea" dependencies = [ "futures", "log", @@ -1517,9 +1524,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b840a8bce0bcbf5afad02946d438591e7c373f7afccaf3d874c04485772514dd" +checksum = "5b00a1fa0da26f6087136a82fea7f13c76a672cbab452d4086952a7cf770a19b" dependencies = [ "arrow", "async-compression", @@ -1535,11 +1542,12 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", "glob", - "itertools", + "itertools 0.15.0", "liblzma", "log", "object_store", @@ -1553,9 +1561,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a24cc0b9cf6e367f27f27406eff13abf48a11b72446aaa40b3105c0ded5c17d9" +checksum = "5ad17ec881bff2ed7768b4bfe971d3efbf3473f2fd1f9d365447bccbdf908678" dependencies = [ "arrow", "arrow-ipc", @@ -1568,18 +1576,19 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "object_store", "tokio", ] [[package]] name = "datafusion-datasource-csv" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1abe56b2a7a2d1d6de5117dd1a203181e267f28529faa5da546947621b697d7" +checksum = "b5345285b0c3eaab412e7539b706973c083bd7e5bce575de5e0a3da488d08d1d" dependencies = [ "arrow", "async-trait", @@ -1591,6 +1600,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1600,9 +1610,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3e467f0611ad7bdd5aad17c63c9bb6182d04e5282e5496d897ea2b49c024ba" +checksum = "da02fb9324f56bd8c53f1ee2e949547425cb66f76adc6832b10d44f80a1221d2" dependencies = [ "arrow", "async-trait", @@ -1614,6 +1624,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1623,11 +1634,12 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cc35b92cd560082155e80d9c826929c852d3c51543f4affd3a51c464a0aab3a" +checksum = "3c0b0dc1453952952fd5c69ad1c7f6042176e69ed233011d47e07cf74ed0949e" dependencies = [ "arrow", + "arrow-schema", "async-trait", "bytes", "datafusion-common", @@ -1641,10 +1653,11 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", - "itertools", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -1654,19 +1667,20 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d69bb69d8769e34f76839c960dbde24c1ac0c885a79b6c3c2287bdc56ec67891" +checksum = "a88fd985bc0550c36f557db69543cc9d6393b1509783520b30e902f23c555da6" [[package]] name = "datafusion-execution" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8eac0a09bc8d263f52025cad9e001da4d8138d633fa288edda4d06b1772eae6" +checksum = "a98f1052f91b4991f0bf2ce1e4e36dfbdcda454a956b8c8d562c7c845e8fce1d" dependencies = [ "arrow", "arrow-buffer", "async-trait", + "bytes", "dashmap", "datafusion-common", "datafusion-expr", @@ -1675,16 +1689,19 @@ dependencies = [ "log", "object_store", "parking_lot", + "pin-project-lite", "rand 0.9.5", "tempfile", + "tokio", + "tokio-util", "url", ] [[package]] name = "datafusion-expr" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eeb14d374767ee0fc62dc79a5ba8bcf8a63c14e993c7d992d0e63adfa23d77d3" +checksum = "464625a1f0e4b9df552d894fafcc8aac953ebbc8b0fa0acdaf20975fd615040e" dependencies = [ "arrow", "arrow-schema", @@ -1696,8 +1713,10 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "indexmap", - "itertools", + "itertools 0.15.0", "recursive", "serde_json", "sqlparser", @@ -1705,25 +1724,25 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7b19a8c95522bee8cbb313d74263b85e355d2b52f42e67ef5694bf5de9e9356" +checksum = "2604994999d5aeca1d1df645ffc98bc787447aaff05dde27aad0342b48fc1fe0" dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools", + "itertools 0.15.0", ] [[package]] name = "datafusion-functions" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f64c983bbbdcb729d921a2b2ac3375598719b5cc0c30345ad664936f3176fc7" +checksum = "051e97533e6af53e4aa0a0667cadc886abcaf36c4a5925019c55c0aa4c218fde" dependencies = [ "arrow", "arrow-buffer", - "base64", + "base64 0.23.1", "blake2", "blake3", "chrono", @@ -1736,7 +1755,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hex", - "itertools", + "itertools 0.15.0", "log", "md-5 0.11.0", "memchr", @@ -1749,9 +1768,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89bc17041e424a47ed062f43df24d84aab8b57c4c3221e5c1a5eef46d6c5718b" +checksum = "2d0f1bb166d3572b6ed40e1afb2faaacade962abc08c2fcf04babee74681c56b" dependencies = [ "arrow", "datafusion-common", @@ -1762,17 +1781,17 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr", "datafusion-physical-expr-common", - "foldhash 0.2.0", "half", + "hashbrown 0.17.1", "log", "num-traits", ] [[package]] name = "datafusion-functions-aggregate-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97dd2a9e865c6108059f5b37b77934f84b50bfb108f837bd0e5c9536e03f0545" +checksum = "7ed756770f5f98369e181d692fd5ee6b1127ffd7322caba92f3730f9f5c92333" dependencies = [ "arrow", "datafusion-common", @@ -1782,9 +1801,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75f0bdfeef16d96417b9632ef855645376b242e9006a126dfd0bedfc54a93f5f" +checksum = "91173fdb5c0ff2a41169a8ffa1b385b8844f18728747bb0a37e35ad7d5772a4f" dependencies = [ "arrow", "arrow-ord", @@ -1799,7 +1818,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools", + "itertools 0.15.0", "itoa", "log", "memchr", @@ -1807,9 +1826,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4e4941673c917819616877e9993da4503e4f4739812be0bc32c5356184c6383" +checksum = "b1bcdfb286a745461b126719c32700777e83df4f17cc44db5d71ebce5731e840" dependencies = [ "arrow", "async-trait", @@ -1823,9 +1842,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12dd2e16c12b84b6f6b41b19f55b366dd1c46876bb35b86896c6349067379e8d" +checksum = "9ec4b508f1f93f00038ba3e737e894ec6c775528b4369413386655ae6125f0fc" dependencies = [ "arrow", "datafusion-common", @@ -1840,9 +1859,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cdc5e4b6f8b6ef823cc1c761f85088ad4c884fe8df64df3cbcc6b2b84698441" +checksum = "0b352020834140073fbf5b46ee0ceb926e5074a9d0bcae1dbd91d0586d999cde" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1850,20 +1869,20 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a3614234dd93578c92428cb4f408e020874f0d2b7e6c90c928d9d28b5df2ceb" +checksum = "15192effab05d38cce10e92a6fb48c967b5f166b27b7195a165a72b232569c58" dependencies = [ "datafusion-doc", "quote", - "syn 2.0.119", + "syn 3.0.4", ] [[package]] name = "datafusion-optimizer" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0635620b050b81bb92764e99250868f654e2cd5ad1bece413283b3f73c83179" +checksum = "854445d9f7847e1e46089cf61b8d341a64382f14484e912c83a0f23b31216896" dependencies = [ "arrow", "chrono", @@ -1872,7 +1891,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools", + "itertools 0.15.0", "log", "recursive", "regex", @@ -1881,9 +1900,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cabf7a86eb70b816729e33c81bf7767c936ee1226f607a114f5dac2decac8d0" +checksum = "671558dad1d2aa253c39c0a4c52515958b99eb91abf649f4b88d5e69cc55282f" dependencies = [ "arrow", "datafusion-common", @@ -1891,10 +1910,11 @@ dependencies = [ "datafusion-expr-common", "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", + "datafusion-proto-models", "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "parking_lot", "petgraph", "recursive", @@ -1903,9 +1923,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de222e04f7e6744501555a54ab0abe26bfdfebee380af79a9bdc175704246859" +checksum = "ffae3d78c2da80ecc829cb58536cc5aca2e99cf1365eda694fc75bfe288861e0" dependencies = [ "arrow", "datafusion-common", @@ -1913,31 +1933,32 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools", + "itertools 0.15.0", ] [[package]] name = "datafusion-physical-expr-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72d0d0057fc5a502d45c870cb6d47c66eb7bdd5edb1bd71ad6f3f724975ac2a8" +checksum = "3d9092ed15e7203fbd0903215172f7c9d18f10d94cba35137f3b3836f7c46f16" dependencies = [ "arrow", "chrono", "datafusion-common", "datafusion-expr-common", + "datafusion-proto-models", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "parking_lot", "pin-project", ] [[package]] name = "datafusion-physical-optimizer" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86046eed10950c5f9aaed9acfd148e9bd2e1dfdfe4f9aef607d1447b271e4183" +checksum = "9005b6cf50b57b72d476c6ed4662b04be7ca6be5320ba9127c6d0b7e4218095b" dependencies = [ "arrow", "datafusion-common", @@ -1948,15 +1969,16 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools", + "datafusion-session", + "itertools 0.15.0", "recursive", ] [[package]] name = "datafusion-physical-plan" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bc84da934c903407ba297971ebcc020c4c1a38aafd765d6c144c76eff3fa6a1" +checksum = "5787e4fcff4adc4fce8948441103a99705018b49c8dff0720b650bd7a15da112" dependencies = [ "arrow", "arrow-data", @@ -1964,6 +1986,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "async-trait", + "bytes", "datafusion-common", "datafusion-common-runtime", "datafusion-execution", @@ -1973,26 +1996,28 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "futures", "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.15.0", "log", "num-traits", "parking_lot", "pin-project-lite", + "serde_json", "tokio", ] [[package]] name = "datafusion-proto" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67791dcfacd142a9d95f8f73c2a161094cc936d231d80c235f5017dacf24e84e" +checksum = "0df0504eb9028d5e01f481af3519cde3f97ab650fd50ce7b787e94b69a7a193c" dependencies = [ "arrow", - "chrono", "datafusion-catalog", "datafusion-catalog-listing", "datafusion-common", @@ -2008,26 +2033,38 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-common", + "datafusion-proto-models", "object_store", "prost", ] [[package]] name = "datafusion-proto-common" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8cd9e80d637891645d074db0f6c650b591117367247deb313fdfb78dff559cb" +checksum = "6b92415a2442964f180d39cdcd8ff1edd99f9260c1499ba80a396499c2154d11" dependencies = [ "arrow", "datafusion-common", "prost", ] +[[package]] +name = "datafusion-proto-models" +version = "55.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62e4c0bd6af4fcabdbe201ee86fbeef2ac423a44e1d8fda56d994c9e2d1d3ad2" +dependencies = [ + "datafusion-common", + "datafusion-proto-common", + "prost", +] + [[package]] name = "datafusion-pruning" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb63eeac6de19be40f487b65dd84e546195f783c5a9928618e0c4f2a3569b0d7" +checksum = "9e651c8df0b90daed6a7be5921ec0ee379e6909705f063eeff70fd4e35010e4c" dependencies = [ "arrow", "datafusion-common", @@ -2041,10 +2078,11 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f961d209177f91bd014db5cbb2c33b7d28a2597b9003e77f17aeb712964315a" +checksum = "fb56667ee38217efab19b895d9a936052cfb47ed438a19663351bdc42a6214a1" dependencies = [ + "arrow-schema", "async-trait", "datafusion-common", "datafusion-execution", @@ -2055,9 +2093,9 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "54.1.0" +version = "55.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1d71cb454da682b2af7488e1fc1ddd72ee1b28f19297b8ccad73f0a21ee9a69" +checksum = "9c29067cb9d32f8e603c45e15d61ea18f1069f96ceafeceb4e18466b8e5b31d9" dependencies = [ "arrow", "bigdecimal", @@ -2070,6 +2108,7 @@ dependencies = [ "recursive", "regex", "sqlparser", + "stacker", ] [[package]] @@ -2184,9 +2223,9 @@ dependencies = [ [[package]] name = "dioxus" -version = "0.7.5" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44c550c06b6785e16258ad620d5b559f5bbcbcc50e3c18c08aa6af2604a4c32" +checksum = "9320593eda4f01858f046698b42441603d2af5fcb62a6c1f9046a0c2eebd6d9b" dependencies = [ "dioxus-asset-resolver", "dioxus-cli-config", @@ -2364,7 +2403,7 @@ dependencies = [ "axum", "axum-core", "axum-extra", - "base64", + "base64 0.22.1", "bytes", "ciborium", "const-str", @@ -2425,7 +2464,7 @@ checksum = "0011e9bca8da2ae6bf02ba9ef93ba6c5fa539e402de3aeb316ad86ca77c59079" dependencies = [ "anyhow", "axum-core", - "base64", + "base64 0.22.1", "ciborium", "dioxus-core", "dioxus-document", @@ -2641,7 +2680,7 @@ dependencies = [ "anyhow", "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "chrono", "ciborium", @@ -3425,7 +3464,7 @@ version = "7.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f49d1053f4708f0af3cf9fc5bffc7e68a914a3c45becb231c80068c9c3f78bea" dependencies = [ - "base64", + "base64 0.22.1", "byteorder", "crossbeam-channel", "flate2", @@ -3439,7 +3478,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "headers-core", "http", @@ -3599,7 +3638,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -3802,12 +3841,6 @@ dependencies = [ "tempfile", ] -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "inventory" version = "0.3.24" @@ -3860,6 +3893,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -4488,9 +4530,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.1" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" dependencies = [ "twox-hash", ] @@ -4768,6 +4810,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -4916,7 +4968,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "form_urlencoded", @@ -4927,7 +4979,7 @@ dependencies = [ "http-body-util", "humantime", "hyper", - "itertools", + "itertools 0.14.0", "md-5 0.10.6", "parking_lot", "percent-encoding", @@ -5040,15 +5092,6 @@ dependencies = [ "thiserror 2.0.20", ] -[[package]] -name = "ordered-float" -version = "2.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] - [[package]] name = "owo-colors" version = "4.4.0" @@ -5080,9 +5123,9 @@ dependencies = [ [[package]] name = "parquet" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d298093b2dec60289dce0684c986d0f7679e9dd15771c2c65406e1aaf604a704" +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" dependencies = [ "ahash", "arrow-array", @@ -5091,7 +5134,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64", + "base64 0.23.1", "brotli", "bytes", "chrono", @@ -5100,18 +5143,16 @@ dependencies = [ "half", "hashbrown 0.17.1", "lz4_flex", - "num-bigint", + "num-bigint 0.5.1", "num-integer", "num-traits", "object_store", "parquet-variant", "parquet-variant-compute", "parquet-variant-json", - "paste", "seq-macro", "simdutf8", "snap", - "thrift", "tokio", "twox-hash", "zstd", @@ -5119,9 +5160,9 @@ dependencies = [ [[package]] name = "parquet-variant" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fc70e87931167a4a3fde2ee923023a0624367691e3fd503476a1084dda5a054" +checksum = "3f7e5fff3ed0c07514a7fb8bee3f2ea5a53f36939410ecac4a466620213539a8" dependencies = [ "arrow", "arrow-schema", @@ -5135,9 +5176,9 @@ dependencies = [ [[package]] name = "parquet-variant-compute" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "823a9ecee8fd83a68f7165ef13acc840c4d7ed995838ba99d4714b20a2b5e780" +checksum = "ba4d3de89dab8d1aaaf601ae8d71bd07ea88cfca9efc1df5815b982c30f631e1" dependencies = [ "arrow", "arrow-schema", @@ -5152,24 +5193,18 @@ dependencies = [ [[package]] name = "parquet-variant-json" -version = "58.4.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f37a91177e2dddb10333546952fcf2b7674b97ebd9449432f24b883d6ea8108" +checksum = "fb19dfe1bd24c17addd761ba4f7000f615e2fa12525871c7baa835dbb3d7f147" dependencies = [ "arrow-schema", - "base64", + "base64 0.23.1", "chrono", "parquet-variant", "serde_json", "uuid", ] -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pastey" version = "0.2.3" @@ -5447,7 +5482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.119", @@ -5760,7 +5795,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "cookie", "cookie_store", @@ -5806,7 +5841,7 @@ version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -6102,6 +6137,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ + "indexmap", "itoa", "memchr", "serde", @@ -6696,17 +6732,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float", -] - [[package]] name = "time" version = "0.3.55" @@ -6908,7 +6933,7 @@ checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "h2", "http", diff --git a/Cargo.toml b/Cargo.toml index c0c4a235e..ab7fe172c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,25 +30,25 @@ liquid-cache-datafusion = { path = "src/datafusion", version = "0.1.13" } liquid-cache-common = { path = "src/common", version = "0.1.13" } liquid-cache = { path = "src/core", version = "0.1.13" } liquid-cache-datafusion-local = { path = "src/datafusion-local", version = "0.1.13" } -arrow = { version = "58.4.0", default-features = false, features = [ +arrow = { version = "59.2.0", default-features = false, features = [ "prettyprint", "ipc", ] } -arrow-flight = { version = "58.4.0", features = ["flight-sql-experimental"] } -arrow-schema = { version = "58.4.0", features = ["serde"] } -parquet = { version = "58.4.0", features = [ +arrow-flight = { version = "59.2.0", features = ["flight-sql-experimental"] } +arrow-schema = { version = "59.2.0", features = ["serde"] } +parquet = { version = "59.2.0", features = [ "async", "experimental", "variant_experimental", ] } -parquet-variant-json = { version = "58.4.0" } -parquet-variant-compute = { version = "58.4.0" } -datafusion = { version = "54.1.0" } -datafusion-common = { version = "54.1.0" } -datafusion-expr-common = { version = "54.1.0" } -datafusion-physical-expr = { version = "54.1.0" } -datafusion-physical-expr-common = { version = "54.1.0" } -datafusion-proto = { version = "54.1.0" } +parquet-variant-json = { version = "59.2.0" } +parquet-variant-compute = { version = "59.2.0" } +datafusion = { version = "55.0.0" } +datafusion-common = { version = "55.0.0" } +datafusion-expr-common = { version = "55.0.0" } +datafusion-physical-expr = { version = "55.0.0" } +datafusion-physical-expr-common = { version = "55.0.0" } +datafusion-proto = { version = "55.0.0" } async-trait = "0.1.89" futures = { version = "0.3.32", default-features = false, features = ["std"] } tokio = { version = "1.52.3", features = ["rt-multi-thread"] } diff --git a/dev/dev-tools/Cargo.toml b/dev/dev-tools/Cargo.toml index b4667e1b8..854f00155 100644 --- a/dev/dev-tools/Cargo.toml +++ b/dev/dev-tools/Cargo.toml @@ -6,7 +6,7 @@ edition = "2024" [dependencies] -dioxus = { version = "=0.7.5", features = ["router", "fullstack"] } +dioxus = { version = "=0.7.10", features = ["router", "fullstack"] } # Must match the wasm-bindgen-cli version pinned in flake.nix. # Only used through dioxus macros, hence the cargo-shear ignore. wasm-bindgen = "=0.2.126" diff --git a/src/core/src/cache/policies/squeeze.rs b/src/core/src/cache/policies/squeeze.rs index 53f489135..7941a1feb 100644 --- a/src/core/src/cache/policies/squeeze.rs +++ b/src/core/src/cache/policies/squeeze.rs @@ -244,7 +244,7 @@ pub(crate) fn try_variant_squeeze( shredded_array = Some(shredded_struct); } - let typed_root = variant_array.typed_value_field()?; + let typed_root = variant_array.typed_value_column()?; let typed_root = typed_root.as_any().downcast_ref::()?; let mut collected = Vec::new(); @@ -676,7 +676,7 @@ mod tests { inner .column_by_name("metadata") .cloned() - .unwrap_or_else(|| Arc::new(base_variant.metadata_field().clone()) as ArrayRef), + .unwrap_or_else(|| base_variant.metadata_column().clone()), inner.column_by_name("value").cloned().unwrap_or_else(|| { Arc::new(BinaryViewArray::from(vec![None::<&[u8]>; inner.len()])) as ArrayRef }), diff --git a/src/datafusion-client/src/client_exec.rs b/src/datafusion-client/src/client_exec.rs index 658c86dc1..13ee28f55 100644 --- a/src/datafusion-client/src/client_exec.rs +++ b/src/datafusion-client/src/client_exec.rs @@ -10,7 +10,7 @@ use arrow_flight::flight_service_client::FlightServiceClient; use arrow_schema::{Schema, SchemaRef}; use datafusion::catalog::memory::DataSourceExec; use datafusion::common::internal_err; -use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::config::ConfigOptions; use datafusion::datasource::physical_plan::{FileSource, ParquetSource}; use datafusion::execution::object_store::ObjectStoreUrl; @@ -23,7 +23,10 @@ use datafusion::physical_plan::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion::physical_plan::{ExecutionPlanProperties, PhysicalExpr, PlanProperties}; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, ExecutionPlanProperties, PhysicalExpr, PlanProperties, + ReplaceChildrenOptions, +}; use datafusion::{ error::Result, execution::{RecordBatchStream, SendableRecordBatchStream}, @@ -75,18 +78,22 @@ impl std::fmt::Debug for LiquidCacheClientExec { } impl LiquidCacheClientExec { + fn plan_properties(remote_plan: &Arc) -> Arc { + Arc::new(PlanProperties::new( + remote_plan.equivalence_properties().clone(), + remote_plan.output_partitioning().clone(), + remote_plan.pipeline_behavior(), + remote_plan.boundedness(), + )) + } + pub(crate) fn new( remote_plan: Arc, cache_server: String, object_stores: Vec<(ObjectStoreUrl, HashMap)>, squeeze_hints: ColumnSqueezeHints, ) -> Self { - let properties = Arc::new(PlanProperties::new( - remote_plan.equivalence_properties().clone(), // Equivalence Properties - remote_plan.output_partitioning().clone(), // Output Partitioning - remote_plan.pipeline_behavior(), - remote_plan.boundedness(), - )); + let properties = Self::plan_properties(&remote_plan); let uuid = Uuid::new_v4(); Self { remote_plan, @@ -140,22 +147,51 @@ impl ExecutionPlan for LiquidCacheClientExec { vec![&self.remote_plan] } - fn with_new_children( + fn replace_children( self: Arc, - children: Vec>, + mut children: Vec>, + options: ReplaceChildrenOptions, ) -> datafusion::error::Result> { + if children.len() != 1 { + return internal_err!( + "LiquidCacheClientExec expects one child, received {}", + children.len() + ); + } + let remote_plan = children.swap_remove(0); + let properties = match options.children_properties { + ChildrenPropertiesMode::Keep => Arc::clone(&self.properties), + ChildrenPropertiesMode::Recompute => Self::plan_properties(&remote_plan), + }; Ok(Arc::new(Self { - remote_plan: children.first().unwrap().clone(), + remote_plan, cache_server: self.cache_server.clone(), plan_registered: self.plan_registered.clone(), object_stores: self.object_stores.clone(), metrics: self.metrics.clone(), uuid: self.uuid, - properties: self.properties.clone(), + properties, squeeze_hints: self.squeeze_hints.clone(), })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion::error::Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, partition: usize, diff --git a/src/datafusion-client/src/lib.rs b/src/datafusion-client/src/lib.rs index 079c7611a..02d8aef9b 100644 --- a/src/datafusion-client/src/lib.rs +++ b/src/datafusion-client/src/lib.rs @@ -9,6 +9,7 @@ mod metrics; mod optimizer; pub use client_exec::LiquidCacheClientExec; use datafusion::{ + common::config::ConfigNonZeroUsize, error::{DataFusionError, Result}, execution::{SessionStateBuilder, object_store::ObjectStoreUrl, runtime_env::RuntimeEnv}, prelude::*, @@ -91,7 +92,7 @@ impl LiquidCacheClientBuilder { .execution .parquet .binary_as_string = true; - session_config.options_mut().execution.batch_size = 8192 * 2; + session_config.options_mut().execution.batch_size = ConfigNonZeroUsize::try_new(8192 * 2)?; // Dynamic filters (e.g. a hash join's runtime build-side filter) are pushed // into scan predicates by DataFusion. In distributed mode those scans are // serialized and executed on a remote server that can never receive the diff --git a/src/datafusion-client/src/optimizer.rs b/src/datafusion-client/src/optimizer.rs index 228e75270..ad75324ba 100644 --- a/src/datafusion-client/src/optimizer.rs +++ b/src/datafusion-client/src/optimizer.rs @@ -1,10 +1,15 @@ use std::{collections::HashMap, sync::Arc}; use datafusion::{ - config::ConfigOptions, datasource::source::DataSourceExec, error::Result, - execution::object_store::ObjectStoreUrl, physical_optimizer::PhysicalOptimizerRule, - physical_plan::ExecutionPlan, physical_plan::aggregates::AggregateExec, - physical_plan::aggregates::AggregateMode, physical_plan::repartition::RepartitionExec, + config::ConfigOptions, + datasource::source::DataSourceExec, + error::Result, + execution::object_store::ObjectStoreUrl, + physical_optimizer::PhysicalOptimizerRule, + physical_plan::aggregates::AggregateExec, + physical_plan::aggregates::AggregateMode, + physical_plan::repartition::RepartitionExec, + physical_plan::{ExecutionPlan, execution_plan::replace_children_if_necessary}, }; use liquid_cache_datafusion::optimizers::SqueezeHintMap; @@ -81,7 +86,7 @@ impl PushdownOptimizer { // If any children were changed, create a new plan with the updated children if children_changed { - plan.with_new_children(new_children) + replace_children_if_necessary(plan, new_children) } else { Ok(plan) } diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q1.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q1.snap index 7590f16a0..f112cc27d 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q1.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q1.snap @@ -10,13 +10,6 @@ expression: displayable.tree_render().to_string() │ ASC NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ l_returnflag@0 ASC NULLS │ -│ LAST, l_linestatus@1 │ -│ ASC NULLS LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ avg_disc: │ @@ -51,6 +44,13 @@ expression: displayable.tree_render().to_string() │ ... │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ l_returnflag@0 ASC NULLS │ +│ LAST, l_linestatus@1 │ +│ ASC NULLS LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ @@ -62,14 +62,13 @@ expression: displayable.tree_render().to_string() │ .l_extendedprice * │ │ Int64(1) - lineitem │ │ .l_discount), sum │ -│ (__common_expr_1 * │ -│ Some(1),20,0 + │ -│ lineitem.l_tax) as │ -│ sum(lineitem │ -│ .l_extendedpri │ -│ ce * Int64(1) - lineitem │ -│ .l_discount * Int64(1) │ -│ + lineitem.l_tax), avg │ +│ (__common_expr_1 * 1 + │ +│ lineitem.l_tax) as sum │ +│ (lineitem │ +│ .l_extendedp │ +│ rice * Int64(1) - lineitem│ +│ .l_discount * Int64(1) + │ +│ lineitem.l_tax), avg │ │ (lineitem.l_quantity), │ │ avg(lineitem │ │ .l_extendedpri │ @@ -105,14 +104,13 @@ expression: displayable.tree_render().to_string() │ .l_extendedprice * │ │ Int64(1) - lineitem │ │ .l_discount), sum │ -│ (__common_expr_1 * │ -│ Some(1),20,0 + │ -│ lineitem.l_tax) as │ -│ sum(lineitem │ -│ .l_extendedpri │ -│ ce * Int64(1) - lineitem │ -│ .l_discount * Int64(1) │ -│ + lineitem.l_tax), avg │ +│ (__common_expr_1 * 1 + │ +│ lineitem.l_tax) as sum │ +│ (lineitem │ +│ .l_extendedp │ +│ rice * Int64(1) - lineitem│ +│ .l_discount * Int64(1) + │ +│ lineitem.l_tax), avg │ │ (lineitem.l_quantity), │ │ avg(lineitem │ │ .l_extendedpri │ @@ -129,8 +127,8 @@ expression: displayable.tree_render().to_string() │ ProjectionExec │ │ -------------------- │ │ __common_expr_1: │ -│ l_extendedprice * (Some(1)│ -│ ,20,0 - l_discount) │ +│ l_extendedprice * (1 - │ +│ l_discount) │ │ │ │ l_discount: │ │ l_discount │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q10.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q10.snap index 8643f5637..2f6a455e4 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q10.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q10.snap @@ -10,13 +10,6 @@ expression: displayable.tree_render().to_string() │ revenue DESC │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec(TopK) │ -│ -------------------- │ -│ limit: 20 │ -│ │ -│ revenue@2 DESC │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ c_acctbal: c_acctbal │ @@ -34,16 +27,26 @@ expression: displayable.tree_render().to_string() │ .l_discount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec(TopK) │ +│ -------------------- │ +│ limit: 20 │ +│ │ +│ sum(lineitem │ +│ .l_extendedp │ +│ rice * Int64(1) - lineitem│ +│ .l_discount)@7 DESC │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: │ @@ -73,11 +76,11 @@ expression: displayable.tree_render().to_string() │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q11.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q11.snap index cbc05809d..b12a8e0db 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q11.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q11.snap @@ -19,23 +19,21 @@ expression: displayable.tree_render().to_string() │ │ │ Decimal128(38, 15)) │ └─────────────┬─────────────┘ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ -│ SortExec │ │ AggregateExec │ +│ ProjectionExec │ │ AggregateExec │ │ -------------------- │ │ -------------------- │ -│ value@1 DESC │ │ aggr: │ -│ │ │ sum(partsupp.ps_supplycost│ +│ ps_partkey: │ │ aggr: │ +│ ps_partkey │ │ sum(partsupp.ps_supplycost│ │ │ │ * partsupp.ps_availqty) │ -│ │ │ │ -│ │ │ mode: Final │ +│ value: │ │ │ +│ sum(partsupp.ps_supplycost│ │ mode: Final │ +│ * partsupp.ps_availqty) │ │ │ └─────────────┬─────────────┘ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ -│ ProjectionExec │ │ CoalescePartitionsExec │ +│ SortExec │ │ CoalescePartitionsExec │ │ -------------------- │ │ │ -│ ps_partkey: │ │ │ -│ ps_partkey │ │ │ -│ │ │ │ -│ value: │ │ │ │ sum(partsupp.ps_supplycost│ │ │ │ * partsupp.ps_availqty) │ │ │ +│ @1 DESC │ │ │ └─────────────┬─────────────┘ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ │ FilterExec │ │ AggregateExec │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q12.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q12.snap index 97026caa9..708c98429 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q12.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q12.snap @@ -8,11 +8,6 @@ expression: displayable.tree_render().to_string() │ l_shipmode ASC NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│l_shipmode@0 ASC NULLS LAST│ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ high_line_count: │ @@ -37,6 +32,11 @@ expression: displayable.tree_render().to_string() │ END) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│l_shipmode@0 ASC NULLS LAST│ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q13.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q13.snap index 3002872f9..14919ac99 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q13.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q13.snap @@ -8,12 +8,6 @@ expression: displayable.tree_render().to_string() │custdist DESC, c_count DESC│ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ custdist@1 DESC, c_count@0│ -│ DESC │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ c_count: c_count │ @@ -22,6 +16,12 @@ expression: displayable.tree_render().to_string() │ count(Int64(1)) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ count(Int64(1))@1 DESC, │ +│ c_count@0 DESC │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: count(1) │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q14.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q14.snap index 9f8cd9219..ec8ab519e 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q14.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q14.snap @@ -26,8 +26,8 @@ expression: displayable.tree_render().to_string() │ sum(CASE WHEN part.p_type │ │ LIKE PROMO% THEN │ │ __common_expr_1 │ -│ ELSE Some(0),38,4 END) │ -│ as sum(CASE WHEN part │ +│ ELSE 0.0000 END) as │ +│ sum(CASE WHEN part │ │ .p_type LIKE Utf8( │ │ "PROMO%") THEN │ │ lineitem │ @@ -53,8 +53,8 @@ expression: displayable.tree_render().to_string() │ sum(CASE WHEN part.p_type │ │ LIKE PROMO% THEN │ │ __common_expr_1 │ -│ ELSE Some(0),38,4 END) │ -│ as sum(CASE WHEN part │ +│ ELSE 0.0000 END) as │ +│ sum(CASE WHEN part │ │ .p_type LIKE Utf8( │ │ "PROMO%") THEN │ │ lineitem │ @@ -74,8 +74,8 @@ expression: displayable.tree_render().to_string() │ ProjectionExec │ │ -------------------- │ │ __common_expr_1: │ -│ l_extendedprice * (Some(1)│ -│ ,20,0 - l_discount) │ +│ l_extendedprice * (1 - │ +│ l_discount) │ │ │ │ p_type: p_type │ └─────────────┬─────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q16.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q16.snap index 6c2e06f76..b7c8c7f59 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q16.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q16.snap @@ -11,16 +11,6 @@ expression: displayable.tree_render().to_string() │ ASC NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ supplier_cnt@3 DESC, │ -│ p_brand@0 ASC │ -│ NULLS LAST, p_type │ -│ @1 ASC NULLS LAST, │ -│ p_size@2 ASC NULLS │ -│ LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ p_brand: p_brand │ @@ -31,6 +21,15 @@ expression: displayable.tree_render().to_string() │ count(alias1) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ count(alias1)@3 DESC, │ +│ p_brand@0 ASC NULLS │ +│ LAST, p_type@1 ASC NULLS │ +│ LAST, p_size@2 ASC NULLS │ +│ LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: count(alias1) │ @@ -94,7 +93,9 @@ expression: displayable.tree_render().to_string() │ HashJoinExec │ │ -------------------- │ │ join_type: LeftAnti │ -│ ├───────────────────────────────────────────┐ +│ │ +│ null_aware ├───────────────────────────────────────────┐ +│ │ │ │ on: │ │ │ (ps_suppkey = s_suppkey) │ │ └─────────────┬─────────────┘ │ @@ -122,8 +123,8 @@ expression: displayable.tree_render().to_string() │ ││ ││ │ │ predicate: ││ ││ predicate: │ │ p_brand != Brand#45 AND ││ ││ s_comment LIKE %Customer │ -│ p_type NOT LIKE ││ ││ %Complaints% │ -│ MEDIUM POLISHED% ││ ││ │ -│ AND p_size IN (49, 14, ││ ││ │ -│ 23, 45, 19, 3, 36, 9) ││ ││ │ +│ p_size IN (49, 14, 23 ││ ││ %Complaints% │ +│ , 45, 19, 3, 36, 9) AND ││ ││ │ +│ p_type NOT LIKE MEDIUM ││ ││ │ +│ POLISHED% ││ ││ │ └───────────────────────────┘└───────────────────────────┘└───────────────────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q17.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q17.snap index 78d6cade3..2c05107f6 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q17.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q17.snap @@ -39,8 +39,10 @@ expression: displayable.tree_render().to_string() │ filter: │ │ CAST(l_quantity AS │ │ Decimal128(30, │ -│ 15)) < Float64(0.2) * ├──────────────┐ -│ avg(lineitem.l_quantity) │ │ +│ 15)) < Float64(0.2) * │ +│ avg(lineitem.l_quantity) ├──────────────┐ +│ │ │ +│ join_type: RightSemi │ │ │ │ │ │ on: │ │ │ (l_partkey = p_partkey) │ │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q18.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q18.snap index f6b222226..a03daf342 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q18.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q18.snap @@ -79,7 +79,7 @@ expression: displayable.tree_render().to_string() │ -------------------- ││ -------------------- │ │ predicate: ││ on: ├───────────────────────────────────────────┐ │ sum(lineitem.l_quantity) >││ (o_orderkey = l_orderkey) │ │ -│ Some(30000),25,2 ││ │ │ +│ 300.00 ││ │ │ └─────────────┬─────────────┘└─────────────┬─────────────┘ │ ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ │ AggregateExec ││ HashJoinExec │ │ LiquidCacheClientExec │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q19.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q19.snap index 080a4932c..b65eaaf77 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q19.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q19.snap @@ -17,11 +17,11 @@ expression: displayable.tree_render().to_string() │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ mode: Final │ @@ -35,11 +35,11 @@ expression: displayable.tree_render().to_string() │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ mode: Partial │ @@ -52,26 +52,21 @@ expression: displayable.tree_render().to_string() │ p_container IN (SM │ │ CASE, SM BOX, SM PACK, │ │ SM PKG) AND l_quantity > │ -│ = Some(100),15,2 AND │ -│ l_quantity <= Some │ -│ (1100),15,2 AND p_size < │ -│ = 5 OR p_brand = Brand#23 │ +│ = 1.00 AND l_quantity <= │ +│ 11.00 AND p_size <= 5 │ +│ OR p_brand = Brand#23 │ │ AND p_container IN (MED │ │ BAG, MED BOX, MED PKG, │ -│ MED PACK) AND │ -│ l_quantity >= │ -│ Some(1000),15,2 AND ├──────────────┐ -│ l_quantity <= Some │ │ -│ (2000),15,2 AND p_size │ │ -│ <= 10 OR p_brand = Brand │ │ -│ #34 AND p_container IN │ │ -│ (LG CASE, LG BOX, LG │ │ -│ PACK, LG PKG) AND │ │ -│ l_quantity >= Some │ │ -│ (2000),15,2 AND │ │ -│ l_quantity <= │ │ -│ Some(3000),15,2 AND │ │ -│ p_size <= 15 │ │ +│ MED PACK) AND ├──────────────┐ +│ l_quantity >= │ │ +│ 10.00 AND l_quantity < │ │ +│ = 20.00 AND p_size <= 10 │ │ +│ OR p_brand = Brand#34 │ │ +│ AND p_container IN (LG │ │ +│ CASE, LG BOX, LG PACK, │ │ +│ LG PKG) AND l_quantity │ │ +│ >= 20.00 AND l_quantity │ │ +│ <= 30.00 AND p_size <= 15 │ │ │ │ │ │ on: │ │ │ (p_partkey = l_partkey) │ │ @@ -118,16 +113,11 @@ expression: displayable.tree_render().to_string() │ REG) AND │ │ l_shipinstr │ │ uct = DELIVER IN PERSON │ - │ AND (l_quantity >= │ - │ Some(100),15,2 AND │ - │ l_quantity <= Some │ - │ (1100),15,2 OR │ - │ l_quantity >= │ - │ Some(1000),15,2 AND │ - │ l_quantity <= Some │ - │ (2000),15,2 OR │ - │ l_quantity >= │ - │ Some(2000),15,2 AND │ - │ l_quantity <= Some(3000 │ - │ ),15,2) │ + │ AND (l_quantity >= 1 │ + │ .00 AND l_quantity <= 11 │ + │ .00 OR l_quantity >= 10 │ + │ .00 AND l_quantity <= │ + │ 20.00 OR l_quantity >= │ + │ 20.00 AND l_quantity <= │ + │ 30.00) │ └───────────────────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q2.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q2.snap index 633cea2c9..1cc295cdd 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q2.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q2.snap @@ -22,9 +22,11 @@ expression: displayable.tree_render().to_string() ┌─────────────┴─────────────┐ │ HashJoinExec │ │ -------------------- │ +│ join_type: LeftSemi │ +│ │ │ on: │ -│ (p_partkey = ps_partkey), │ -│ (ps_supplycost = min ├──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ (p_partkey = ps_partkey), ├──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ (ps_supplycost = min │ │ │ (partsupp │ │ │ .ps_supplycost │ │ │ )) │ │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q20.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q20.snap index 7a145d7b9..9f69961d4 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q20.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q20.snap @@ -27,8 +27,10 @@ expression: displayable.tree_render().to_string() │(n_nationkey = s_nationkey)│ │ CAST(ps_availqty AS │ │ │ │ Float64) > │ │ │ │ Float64(0.5) * │ -│ ├──────────────┐ │ sum(lineitem ├───────────────────────────────────────────┐ -│ │ │ │ .l_quantity) │ │ +│ │ │ sum(lineitem │ +│ ├──────────────┐ │ .l_quantity) ├───────────────────────────────────────────┐ +│ │ │ │ │ │ +│ │ │ │ join_type: LeftSemi │ │ │ │ │ │ │ │ │ │ │ │ on: │ │ │ │ │ │ (ps_partkey = l_partkey), │ │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q21.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q21.snap index d2a6f6f5b..859314c85 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q21.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q21.snap @@ -9,12 +9,6 @@ expression: displayable.tree_render().to_string() │ NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ numwait@1 DESC, s_name@0 │ -│ ASC NULLS LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ numwait: │ @@ -23,6 +17,13 @@ expression: displayable.tree_render().to_string() │ s_name: s_name │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ count(Int64(1))@1 DESC, │ +│ s_name@0 ASC NULLS │ +│ LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: count(1) │ @@ -52,83 +53,86 @@ expression: displayable.tree_render().to_string() │ -------------------- │ │ filter: │ │ l_suppkey != l_suppkey │ -│ ├──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ -│ join_type: LeftAnti │ │ -│ │ │ -│ on: │ │ -│ (l_orderkey = l_orderkey) │ │ -└─────────────┬─────────────┘ │ -┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ -│ HashJoinExec │ │ LiquidCacheClientExec │ -│ -------------------- │ │ -------------------- │ -│ filter: │ │ server: │ -│ l_suppkey != l_suppkey │ │ http://localhost:50051, │ -│ ├─────────────────────────────────────────────────────────────────────────────────────────────────────┐ │ object_stores=[] │ -│ join_type: LeftSemi │ │ │ │ -│ │ │ │ │ -│ on: │ │ │ │ -│ (l_orderkey = l_orderkey) │ │ │ │ -└─────────────┬─────────────┘ │ └─────────────┬─────────────┘ -┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ -│ CoalescePartitionsExec │ │ LiquidCacheClientExec ││ RepartitionExec │ -│ │ │ -------------------- ││ -------------------- │ -│ │ │ server: ││ partition_count(in->out): │ -│ │ │ http://localhost:50051, ││ 1 -> 8 │ -│ │ │ object_stores=[] ││ │ -│ │ │ ││ partitioning_scheme: │ -│ │ │ ││ RoundRobinBatch(8) │ -└─────────────┬─────────────┘ └─────────────┬─────────────┘└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ -│ HashJoinExec │ │ DataSourceExec ││ DataSourceExec │ -│ -------------------- │ │ -------------------- ││ -------------------- │ -│ on: │ │ files: 1 ││ files: 1 │ -│(n_nationkey = s_nationkey)│ │ format: parquet ││ format: parquet │ -│ ├──────────────┐ │ ││ │ -│ │ │ │ ││ predicate: │ -│ │ │ │ ││ l_receiptdate > │ -│ │ │ │ ││ l_commitdate │ -└─────────────┬─────────────┘ │ └───────────────────────────┘└───────────────────────────┘ +│ ├──────────────┐ +│ join_type: RightAnti │ │ +│ │ │ +│ on: │ │ +│ (l_orderkey = l_orderkey) │ │ +└─────────────┬─────────────┘ │ +┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +│ LiquidCacheClientExec ││ RepartitionExec │ +│ -------------------- ││ -------------------- │ +│ server: ││ partition_count(in->out): │ +│ http://localhost:50051, ││ 1 -> 8 │ +│ object_stores=[] ││ │ +│ ││ partitioning_scheme: │ +│ ││ RoundRobinBatch(8) │ +└─────────────┬─────────────┘└─────────────┬─────────────┘ ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ -│ LiquidCacheClientExec ││ HashJoinExec │ +│ DataSourceExec ││ HashJoinExec │ │ -------------------- ││ -------------------- │ -│ server: ││ on: ├──────────────┐ -│ http://localhost:50051, ││ (o_orderkey = l_orderkey) │ │ -│ object_stores=[] ││ │ │ -└─────────────┬─────────────┘└─────────────┬─────────────┘ │ -┌─────────────┴─────────────┐┌─────────────┴─────────────┐┌─────────────┴─────────────┐ -│ DataSourceExec ││ LiquidCacheClientExec ││ HashJoinExec │ -│ -------------------- ││ -------------------- ││ -------------------- │ -│ files: 1 ││ server: ││ on: │ -│ format: parquet ││ http://localhost:50051, ││ (s_suppkey = l_suppkey) ├──────────────┐ -│ ││ object_stores=[] ││ │ │ -│ predicate: ││ ││ │ │ -│ n_name = SAUDI ARABIA ││ ││ │ │ -└───────────────────────────┘└─────────────┬─────────────┘└─────────────┬─────────────┘ │ +│ files: 1 ││ filter: │ +│ format: parquet ││ l_suppkey != l_suppkey │ +│ ││ ├─────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ predicate: ││ join_type: LeftSemi │ │ +│ l_receiptdate > ││ │ │ +│ l_commitdate ││ on: │ │ +│ ││ (l_orderkey = l_orderkey) │ │ +└───────────────────────────┘└─────────────┬─────────────┘ │ + ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ + │ CoalescePartitionsExec │ │ LiquidCacheClientExec │ + │ │ │ -------------------- │ + │ │ │ server: │ + │ │ │ http://localhost:50051, │ + │ │ │ object_stores=[] │ + └─────────────┬─────────────┘ └─────────────┬─────────────┘ + ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ + │ HashJoinExec │ │ DataSourceExec │ + │ -------------------- │ │ -------------------- │ + │ on: ├──────────────┐ │ files: 1 │ + │(n_nationkey = s_nationkey)│ │ │ format: parquet │ + └─────────────┬─────────────┘ │ └───────────────────────────┘ + ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ + │ LiquidCacheClientExec ││ HashJoinExec │ + │ -------------------- ││ -------------------- │ + │ server: ││ on: ├──────────────┐ + │ http://localhost:50051, ││ (o_orderkey = l_orderkey) │ │ + │ object_stores=[] ││ │ │ + └─────────────┬─────────────┘└─────────────┬─────────────┘ │ ┌─────────────┴─────────────┐┌─────────────┴─────────────┐┌─────────────┴─────────────┐ - │ DataSourceExec ││ LiquidCacheClientExec ││ LiquidCacheClientExec │ + │ DataSourceExec ││ LiquidCacheClientExec ││ HashJoinExec │ │ -------------------- ││ -------------------- ││ -------------------- │ - │ files: 1 ││ server: ││ server: │ - │ format: parquet ││ http://localhost:50051, ││ http://localhost:50051, │ - │ ││ object_stores=[] ││ object_stores=[] │ - │ predicate: ││ ││ │ - │ o_orderstatus = F ││ ││ │ - └───────────────────────────┘└─────────────┬─────────────┘└─────────────┬─────────────┘ - ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ - │ DataSourceExec ││ RepartitionExec │ - │ -------------------- ││ -------------------- │ - │ files: 1 ││ partition_count(in->out): │ - │ format: parquet ││ 1 -> 8 │ - │ ││ │ - │ ││ partitioning_scheme: │ - │ ││ RoundRobinBatch(8) │ - └───────────────────────────┘└─────────────┬─────────────┘ - ┌─────────────┴─────────────┐ - │ DataSourceExec │ - │ -------------------- │ - │ files: 1 │ - │ format: parquet │ - │ │ - │ predicate: │ - │ l_receiptdate > │ - │ l_commitdate │ - └───────────────────────────┘ + │ files: 1 ││ server: ││ on: │ + │ format: parquet ││ http://localhost:50051, ││ (s_suppkey = l_suppkey) ├──────────────┐ + │ ││ object_stores=[] ││ │ │ + │ predicate: ││ ││ │ │ + │ n_name = SAUDI ARABIA ││ ││ │ │ + └───────────────────────────┘└─────────────┬─────────────┘└─────────────┬─────────────┘ │ + ┌─────────────┴─────────────┐┌─────────────┴─────────────┐┌─────────────┴─────────────┐ + │ DataSourceExec ││ LiquidCacheClientExec ││ LiquidCacheClientExec │ + │ -------------------- ││ -------------------- ││ -------------------- │ + │ files: 1 ││ server: ││ server: │ + │ format: parquet ││ http://localhost:50051, ││ http://localhost:50051, │ + │ ││ object_stores=[] ││ object_stores=[] │ + │ predicate: ││ ││ │ + │ o_orderstatus = F ││ ││ │ + └───────────────────────────┘└─────────────┬─────────────┘└─────────────┬─────────────┘ + ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ + │ DataSourceExec ││ RepartitionExec │ + │ -------------------- ││ -------------------- │ + │ files: 1 ││ partition_count(in->out): │ + │ format: parquet ││ 1 -> 8 │ + │ ││ │ + │ ││ partitioning_scheme: │ + │ ││ RoundRobinBatch(8) │ + └───────────────────────────┘└─────────────┬─────────────┘ + ┌─────────────┴─────────────┐ + │ DataSourceExec │ + │ -------------------- │ + │ files: 1 │ + │ format: parquet │ + │ │ + │ predicate: │ + │ l_receiptdate > │ + │ l_commitdate │ + └───────────────────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q22.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q22.snap index 740987243..89d275d25 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q22.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q22.snap @@ -14,22 +14,24 @@ expression: displayable.tree_render().to_string() │ │ │ mode: Final │ └─────────────┬─────────────┘ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ -│ SortExec │ │ CoalescePartitionsExec │ +│ ProjectionExec │ │ CoalescePartitionsExec │ │ -------------------- │ │ │ -│ cntrycode@0 ASC NULLS LAST│ │ │ -└─────────────┬─────────────┘ └─────────────┬─────────────┘ -┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ -│ ProjectionExec │ │ LiquidCacheClientExec │ -│ -------------------- │ │ -------------------- │ -│ cntrycode: cntrycode │ │ server: │ -│ │ │ http://localhost:50051, │ -│ numcust: │ │ object_stores=[] │ +│ cntrycode: cntrycode │ │ │ +│ │ │ │ +│ numcust: │ │ │ │ count(Int64(1)) │ │ │ │ │ │ │ │ totacctbal: │ │ │ │ sum(custsale.c_acctbal) │ │ │ └─────────────┬─────────────┘ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ +│ SortExec │ │ LiquidCacheClientExec │ +│ -------------------- │ │ -------------------- │ +│ cntrycode@0 ASC NULLS LAST│ │ server: │ +│ │ │ http://localhost:50051, │ +│ │ │ object_stores=[] │ +└─────────────┬─────────────┘ └─────────────┬─────────────┘ +┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ │ AggregateExec │ │ AggregateExec │ │ -------------------- │ │ -------------------- │ │ aggr: │ │ aggr: │ @@ -57,8 +59,8 @@ expression: displayable.tree_render().to_string() │ count(1), sum(custsale │ │ format: parquet │ │ .c_acctbal) │ │ │ │ │ │ predicate: │ -│ group_by: cntrycode │ │ c_acctbal > Some(0),15,2 │ -│ mode: Partial │ │ AND substr(c_phone, 1, │ +│ group_by: cntrycode │ │ c_acctbal > 0.00 AND │ +│ mode: Partial │ │ substr(c_phone, 1, │ │ │ │ 2) IN (13, 31, 23, 29, │ │ │ │ 30, 18, 17) │ └─────────────┬─────────────┘ └───────────────────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q3.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q3.snap index 9387f57b6..76f7ce9ed 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q3.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q3.snap @@ -11,15 +11,6 @@ expression: displayable.tree_render().to_string() │ ASC NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec(TopK) │ -│ -------------------- │ -│ limit: 10 │ -│ │ -│ revenue@1 DESC, │ -│ o_orderdate@2 │ -│ ASC NULLS LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ l_orderkey: │ @@ -38,16 +29,28 @@ expression: displayable.tree_render().to_string() │ .l_discount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec(TopK) │ +│ -------------------- │ +│ limit: 10 │ +│ │ +│ sum(lineitem │ +│ .l_extendedp │ +│ rice * Int64(1) - lineitem│ +│ .l_discount)@3 DESC, │ +│ o_orderdate@1 ASC │ +│ NULLS LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: │ @@ -74,11 +77,11 @@ expression: displayable.tree_render().to_string() │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q4.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q4.snap index 7974f5a26..a727ea5fc 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q4.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q4.snap @@ -9,12 +9,6 @@ expression: displayable.tree_render().to_string() │ LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ o_orderpriority@0 ASC │ -│ NULLS LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ o_orderpriority: │ @@ -24,6 +18,12 @@ expression: displayable.tree_render().to_string() │ count(Int64(1)) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ o_orderpriority@0 ASC │ +│ NULLS LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: count(1) │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q5.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q5.snap index 0c5df4436..bd87dd29b 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q5.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q5.snap @@ -8,11 +8,6 @@ expression: displayable.tree_render().to_string() │ revenue DESC │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ revenue@1 DESC │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ n_name: n_name │ @@ -24,16 +19,24 @@ expression: displayable.tree_render().to_string() │ .l_discount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ sum(lineitem │ +│ .l_extendedp │ +│ rice * Int64(1) - lineitem│ +│ .l_discount)@1 DESC │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: n_name │ @@ -56,11 +59,11 @@ expression: displayable.tree_render().to_string() │ aggr: │ │ sum(lineitem │ │ .l_extendedp │ -│ rice * Some(1),20,0 - │ -│ lineitem.l_discount │ -│ ) as sum(lineitem │ -│ .l_extendedprice │ -│ * Int64(1) - lineitem │ +│ rice * 1 - lineitem │ +│ .l_discount) as │ +│ sum(lineitem │ +│ .l_extendedprice │ +│ * Int64(1) - lineitem │ │ .l_discount) │ │ │ │ group_by: n_name │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q6.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q6.snap index 4e5df56ea..2b52c8267 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q6.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q6.snap @@ -59,9 +59,6 @@ expression: displayable.tree_render().to_string() │ l_shipdate >= 1994-01-01 │ │ AND l_shipdate < 1995 │ │ -01-01 AND l_discount >= │ -│ Some(5),15,2 AND │ -│ l_discount <= │ -│ Some(7),15,2 AND │ -│ l_quantity < Some │ -│ (2400),15,2 │ +│ 0.05 AND l_discount <= 0 │ +│ .07 AND l_quantity < 24.00│ └───────────────────────────┘ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q7.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q7.snap index 66132d2ea..dc6c9b082 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q7.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q7.snap @@ -11,14 +11,6 @@ expression: displayable.tree_render().to_string() │ LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ supp_nation@0 ASC NULLS │ -│ LAST, cust_nation@1 │ -│ ASC NULLS LAST, l_year │ -│ @2 ASC NULLS LAST │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ cust_nation: │ @@ -33,6 +25,14 @@ expression: displayable.tree_render().to_string() │ supp_nation │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ supp_nation@0 ASC NULLS │ +│ LAST, cust_nation@1 │ +│ ASC NULLS LAST, l_year │ +│ @2 ASC NULLS LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ @@ -79,8 +79,8 @@ expression: displayable.tree_render().to_string() │ supp_nation: n_name │ │ │ │ volume: │ -│ l_extendedprice * (Some(1)│ -│ ,20,0 - l_discount) │ +│ l_extendedprice * (1 - │ +│ l_discount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ │ HashJoinExec │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q8.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q8.snap index 7f2733a23..c133378a9 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q8.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q8.snap @@ -32,16 +32,14 @@ expression: displayable.tree_render().to_string() │ sum(CASE WHEN all_nations │ │ .nation = BRAZIL THEN │ │ all_nations.volume │ -│ ELSE Some(0),38,4 END │ -│ ) as sum(CASE WHEN │ +│ ELSE 0.0000 END) as │ +│ sum(CASE WHEN │ │ all_nations │ │ .nation = Utf8( │ │ "BRAZIL") THEN │ -│ all_nations │ -│ .volume ELSE Int64 │ -│ (0) END), sum │ -│ (all_nations │ -│ .volume) │ +│ all_nations.volume │ +│ ELSE Int64(0) END), sum │ +│ (all_nations.volume) │ │ │ │ group_by: o_year │ │ │ @@ -64,16 +62,14 @@ expression: displayable.tree_render().to_string() │ sum(CASE WHEN all_nations │ │ .nation = BRAZIL THEN │ │ all_nations.volume │ -│ ELSE Some(0),38,4 END │ -│ ) as sum(CASE WHEN │ +│ ELSE 0.0000 END) as │ +│ sum(CASE WHEN │ │ all_nations │ │ .nation = Utf8( │ │ "BRAZIL") THEN │ -│ all_nations │ -│ .volume ELSE Int64 │ -│ (0) END), sum │ -│ (all_nations │ -│ .volume) │ +│ all_nations.volume │ +│ ELSE Int64(0) END), sum │ +│ (all_nations.volume) │ │ │ │ group_by: o_year │ │ mode: Partial │ @@ -88,8 +84,8 @@ expression: displayable.tree_render().to_string() │ o_orderdate) │ │ │ │ volume: │ -│ l_extendedprice * (Some(1)│ -│ ,20,0 - l_discount) │ +│ l_extendedprice * (1 - │ +│ l_discount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ │ HashJoinExec │ diff --git a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q9.snap b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q9.snap index 54d9618e5..7957aaf08 100644 --- a/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q9.snap +++ b/src/datafusion-client/src/tests/snapshots/liquid_cache_datafusion_client__tests__tpch_q9.snap @@ -9,12 +9,6 @@ expression: displayable.tree_render().to_string() │ o_year DESC │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│ nation@0 ASC NULLS LAST, │ -│ o_year@1 DESC │ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ nation: nation │ @@ -24,6 +18,12 @@ expression: displayable.tree_render().to_string() │ sum(profit.amount) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ nation@0 ASC NULLS LAST, │ +│ o_year@1 DESC │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: │ @@ -60,8 +60,8 @@ expression: displayable.tree_render().to_string() │ ProjectionExec │ │ -------------------- │ │ amount: │ -│ l_extendedprice * (Some(1)│ -│ ,20,0 - l_discount) - │ +│ l_extendedprice * (1 - │ +│ l_discount) - │ │ ps_supplycost * │ │ l_quantity │ │ │ diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 073864d65..dfdcb91fa 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -7,9 +7,9 @@ mod tests; use std::path::PathBuf; use std::sync::Arc; -use datafusion::error::Result; use datafusion::logical_expr::ScalarUDF; use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::{common::config::ConfigNonZeroUsize, error::Result}; use liquid_cache::cache::squeeze_policies::{SqueezePolicy, TranscodeSqueezeEvict}; use liquid_cache::cache::{AlwaysHydrate, HydrationPolicy, default_max_memory_bytes}; use liquid_cache::cache_policies::{CachePolicy, LiquidPolicy}; @@ -159,7 +159,7 @@ impl LiquidCacheLocalBuilder { .schema_force_view_types = false; config.options_mut().execution.parquet.skip_arrow_metadata = false; config.options_mut().execution.parquet.skip_metadata = false; - config.options_mut().execution.batch_size = self.batch_size; + config.options_mut().execution.batch_size = ConfigNonZeroUsize::try_new(self.batch_size)?; let store = t4::mount(self.cache_dir.join("liquid_cache.t4")) .await diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 2bc020214..776dd3a7e 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -100,7 +100,11 @@ async fn create_session_context_with_liquid_cache( cache_size_bytes: usize, cache_dir: &Path, ) -> Result<(SessionContext, LiquidCacheParquetRef)> { - let mut config = SessionConfig::new(); + // These tests snapshot exact cache contents and counters. A repartitioned + // file scan populates the cache concurrently, so insertion order (and, for + // LIMIT queries, which partitions finish before cancellation) is not a + // stable property to snapshot. + let mut config = SessionConfig::new().with_repartition_file_scans(false); config.options_mut().execution.target_partitions = 4; let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(cache_size_bytes) @@ -124,6 +128,12 @@ async fn get_physical_plan(sql: &str, ctx: &SessionContext) -> Arc String { + let plan = get_physical_plan(sql, ctx).await; + let batches = collect(plan, ctx.task_ctx()).await.unwrap(); + pretty_format_batches(&batches).unwrap().to_string() +} + async fn run_sql_with_cache( sql: &str, squeeze_policy: Box, @@ -139,12 +149,6 @@ async fn run_sql_with_cache( let displayable = DisplayableExecutionPlan::new(plan.as_ref()); let plan_string = format!("{}", displayable.tree_render()); - async fn get_result(ctx: &SessionContext, sql: &str) -> String { - let plan = get_physical_plan(sql, ctx).await; - let batches = collect(plan, ctx.task_ctx()).await.unwrap(); - pretty_format_batches(&batches).unwrap().to_string() - } - // Clear any historical runtime counters before warming the cache. cache.storage().stats(); @@ -486,3 +490,53 @@ async fn test_provide_schema_with_filter() { } assert_eq!(formatted_results, reference); } + +#[tokio::test] +async fn test_repartitioned_file_scan_cache_correctness() { + let reference_cache_dir = TempDir::new().unwrap(); + let parallel_cache_dir = TempDir::new().unwrap(); + let sql = r#"select "WatchID", "OS", "EventTime" from hits where "OS" <> 2 order by "WatchID" desc limit 10"#; + + let reference = run_sql_with_cache( + sql, + Box::new(TranscodeSqueezeEvict), + 1024 * 1024, + reference_cache_dir.path(), + ) + .await + .values; + + // DataFusion 55 lowered repartition_file_min_size from 10 MiB to 1 MiB, + // which splits the 2.3 MiB fixture into four concurrent scan partitions. + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = 4; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_max_memory_bytes(1024 * 1024) + .with_cache_dir(parallel_cache_dir.path().to_path_buf()) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build(config) + .await + .unwrap(); + ctx.register_parquet("hits", TEST_FILE, ParquetReadOptions::default()) + .await + .unwrap(); + + let plan = get_physical_plan(sql, &ctx).await; + let plan = format!( + "{}", + DisplayableExecutionPlan::new(plan.as_ref()).tree_render() + ); + assert!( + plan.contains("files: 4"), + "expected a repartitioned scan:\n{plan}" + ); + + assert_eq!(get_result(&ctx, sql).await, reference); + let entries_after_first_run = cache.storage().stats().total_entries; + assert_eq!(get_result(&ctx, sql).await, reference); + + let stats = cache.storage().stats(); + assert!(stats.runtime.get_with_selection > 0); + assert!(stats.total_entries >= entries_after_first_run); +} diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap index 3b21a3f31..0a535fb1a 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap @@ -54,11 +54,6 @@ plan: │ zo_sql_key ASC NULLS LAST │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ -│ SortExec │ -│ -------------------- │ -│zo_sql_key@0 ASC NULLS LAST│ -└─────────────┬─────────────┘ -┌─────────────┴─────────────┐ │ ProjectionExec │ │ -------------------- │ │ zo_sql_key: │ @@ -78,6 +73,22 @@ plan: │ count(Int64(1)) │ └─────────────┬─────────────┘ ┌─────────────┴─────────────┐ +│ SortExec │ +│ -------------------- │ +│ date_bin │ +│ (Interva │ +│ lMonthDayNano( │ +│ "IntervalMonth │ +│ DayNano { months: 0, days:│ +│ 0, nanoseconds: │ +│ 10000000000 }") │ +│ ,to_timestamp_micros │ +│ (default._timestamp) │ +│ ,to_timestamp(Utf8("2001 │ +│ -01-01T00:00:00")))@0 │ +│ ASC NULLS LAST │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ AggregateExec │ │ -------------------- │ │ aggr: count(1) │ diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap index e518b4d6c..0b214cc49 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap @@ -4,6 +4,13 @@ expression: "format!(\"plan: \\n{}\\nvalues: \\n{}\\nstats:\\n{}\", plan, values --- plan: ┌───────────────────────────┐ +│ ProjectionExec │ +│ -------------------- │ +│ EventTime: EventTime │ +│ OS: OS │ +│ WatchID: WatchID │ +└─────────────┬─────────────┘ +┌─────────────┴─────────────┐ │ SortExec(TopK) │ │ -------------------- │ │ WatchID@0 DESC │ diff --git a/src/datafusion-local/src/tests/squeeze.rs b/src/datafusion-local/src/tests/squeeze.rs index 89b2eba80..7a50e3497 100644 --- a/src/datafusion-local/src/tests/squeeze.rs +++ b/src/datafusion-local/src/tests/squeeze.rs @@ -6,13 +6,17 @@ use crate::LiquidCacheLocalBuilder; const TEST_FILE: &str = "../../examples/nano_hits.parquet"; +fn squeeze_test_config() -> SessionConfig { + SessionConfig::new().with_repartition_file_scans(false) +} + #[tokio::test] async fn basic_squeeze() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 128) .with_cache_dir(cache_dir.path().to_path_buf()) - .build(SessionConfig::new()) + .build(squeeze_test_config()) .await .unwrap(); ctx.register_parquet("hits", TEST_FILE, Default::default()) @@ -39,7 +43,7 @@ async fn squeeze_strings() { let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 1024) .with_cache_dir(cache_dir.path().to_path_buf()) - .build(SessionConfig::new()) + .build(squeeze_test_config()) .await .unwrap(); ctx.register_parquet("hits", TEST_FILE, Default::default()) @@ -66,7 +70,7 @@ async fn squeeze_substrings_search() { let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 256) .with_cache_dir(cache_dir.path().to_path_buf()) - .build(SessionConfig::new()) + .build(squeeze_test_config()) .await .unwrap(); ctx.register_parquet("hits", TEST_FILE, Default::default()) @@ -90,7 +94,7 @@ async fn squeeze_substrings_search_title() { let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 1024 * 4) .with_cache_dir(cache_dir.path().to_path_buf()) - .build(SessionConfig::new()) + .build(squeeze_test_config()) .await .unwrap(); ctx.register_parquet("hits", TEST_FILE, Default::default()) @@ -115,7 +119,7 @@ async fn squeeze_distinct_search_phase() { let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 256) .with_cache_dir(cache_dir.path().to_path_buf()) - .build(SessionConfig::new()) + .build(squeeze_test_config()) .await .unwrap(); ctx.register_parquet("hits", TEST_FILE, Default::default()) diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 8fcdf9c5b..029f928aa 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -15,7 +15,7 @@ use datafusion::{ tree_node::{TreeNode, TreeNodeRecursion}, }, datasource::physical_plan::FileScanConfig, - physical_plan::ExecutionPlan, + physical_plan::{ExecutionPlan, StatisticsArgs, StatisticsContext}, }; use liquid_cache_common::rpc::ExecutionMetricsResponse; use liquid_cache_datafusion::LiquidParquetSource; @@ -315,6 +315,9 @@ pub(crate) async fn start_flamegraph_handler( impl From<&Arc> for ExecutionPlanWithStats { fn from(plan: &Arc) -> Self { let metrics = plan.metrics().unwrap().aggregate_by_name(); + let statistics = StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap(); let mut metric_values = Vec::new(); for metric in metrics.iter() { metric_values.push(MetricValues { @@ -324,13 +327,7 @@ impl From<&Arc> for ExecutionPlanWithStats { } let mut column_statistics = Vec::new(); - for (i, cs) in plan - .partition_statistics(None) - .unwrap() - .column_statistics - .iter() - .enumerate() - { + for (i, cs) in statistics.column_statistics.iter().enumerate() { let min = if cs.min_value != Precision::Absent { Some(cs.min_value.to_string()) } else { @@ -378,16 +375,8 @@ impl From<&Arc> for ExecutionPlanWithStats { }) .collect(), statistics: Statistics { - num_rows: plan - .partition_statistics(None) - .unwrap() - .num_rows - .to_string(), - total_byte_size: plan - .partition_statistics(None) - .unwrap() - .total_byte_size - .to_string(), + num_rows: statistics.num_rows.to_string(), + total_byte_size: statistics.total_byte_size.to_string(), column_statistics, }, metrics: metric_values, diff --git a/src/datafusion-server/src/lib.rs b/src/datafusion-server/src/lib.rs index 4da84881f..a12b6dd4d 100644 --- a/src/datafusion-server/src/lib.rs +++ b/src/datafusion-server/src/lib.rs @@ -28,6 +28,7 @@ use arrow_flight::{ }, }; use datafusion::{ + common::config::ConfigNonZeroUsize, error::DataFusionError, execution::{SessionStateBuilder, object_store::ObjectStoreUrl}, prelude::{SessionConfig, SessionContext}, @@ -167,7 +168,7 @@ impl LiquidCacheService { let mut session_config = SessionConfig::from_env()?; let options_mut = session_config.options_mut(); options_mut.execution.parquet.pushdown_filters = true; - options_mut.execution.batch_size = 8192 * 2; + options_mut.execution.batch_size = ConfigNonZeroUsize::try_new(8192 * 2)?; { // view types cause excessive memory usage because they are not gced. diff --git a/src/datafusion/src/reader/plantime/opener.rs b/src/datafusion/src/reader/plantime/opener.rs index b1b7e8362..f54890d01 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -20,11 +20,10 @@ use datafusion::{ table_schema::TableSchema, }, error::DataFusionError, - physical_expr::PhysicalExprSimplifier, physical_expr::projection::ProjectionExprs, physical_expr::utils::reassign_expr_columns, + physical_expr::{DynamicFilterTracking, PhysicalExprSimplifier}, physical_expr_adapter::{PhysicalExprAdapterFactory, replace_columns_with_literals}, - physical_expr_common::physical_expr::is_dynamic_physical_expr, physical_optimizer::pruning::{FilePruner, PruningPredicate, build_pruning_predicate}, physical_plan::{ PhysicalExpr, @@ -157,7 +156,10 @@ impl FileOpener for LiquidParquetOpener { // we can end the stream early. let mut file_pruner = predicate .as_ref() - .filter(|p| is_dynamic_physical_expr(p) || partitioned_file.has_statistics()) + .filter(|p| { + DynamicFilterTracking::classify(p).contains_dynamic_filter() + || partitioned_file.has_statistics() + }) .and_then(|p| { FilePruner::try_new( Arc::clone(p), diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 79d715b7d..f95963b03 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -4,6 +4,7 @@ use ahash::{HashMap, HashMapExt}; use arrow_schema::Schema; use bytes::Bytes; use datafusion::{ + common::tree_node::TreeNodeRecursion, config::TableParquetOptions, datasource::{ listing::PartitionedFile, @@ -16,19 +17,17 @@ use datafusion::{ error::Result, physical_expr::projection::ProjectionExprs, physical_expr_adapter::DefaultPhysicalExprAdapterFactory, - physical_optimizer::pruning::PruningPredicate, + physical_optimizer::pruning::{PruningPredicate, PruningPredicateBuilder}, physical_plan::{ - PhysicalExpr, + PhysicalExpr, apply_expression_roots, metrics::{ExecutionPlanMetricsSet, MetricBuilder}, }, }; use futures::{FutureExt, future::BoxFuture}; -use object_store::{ObjectStore, path::Path}; +use object_store::{ObjectStore, ObjectStoreExt, path::Path}; use parquet::{ - arrow::{ - arrow_reader::ArrowReaderOptions, - async_reader::{AsyncFileReader, ParquetObjectReader}, - }, + arrow::{arrow_reader::ArrowReaderOptions, async_reader::AsyncFileReader}, + errors::ParquetError, file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}, }; use std::{ @@ -57,17 +56,12 @@ impl CachedMetaReaderFactory { metrics: &ExecutionPlanMetricsSet, ) -> ParquetMetadataCacheReader { let path = partitioned_file.object_meta.location.clone(); - let store = Arc::clone(&self.store); - let mut inner = ParquetObjectReader::new(store, path.clone()) - .with_file_size(partitioned_file.object_meta.size); - - if let Some(hint) = metadata_size_hint { - inner = inner.with_footer_size_hint(hint); - } ParquetMetadataCacheReader { file_metrics: ParquetFileMetrics::new(partition_index, path.as_ref(), metrics), - inner, + store: Arc::clone(&self.store), + file_size: partitioned_file.object_meta.size, + metadata_size_hint, path, } } @@ -106,10 +100,16 @@ impl MetadataCache { #[derive(Clone)] pub struct ParquetMetadataCacheReader { file_metrics: ParquetFileMetrics, - inner: ParquetObjectReader, + store: Arc, + file_size: u64, + metadata_size_hint: Option, path: Path, } +fn to_parquet_err(error: object_store::Error) -> ParquetError { + ParquetError::External(Box::new(error)) +} + impl AsyncFileReader for ParquetMetadataCacheReader { fn get_byte_ranges( &mut self, @@ -117,14 +117,26 @@ impl AsyncFileReader for ParquetMetadataCacheReader { ) -> BoxFuture<'_, parquet::errors::Result>> { let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); self.file_metrics.bytes_scanned.add(total as usize); - self.inner.get_byte_ranges(ranges) + async move { + self.store + .get_ranges(&self.path, &ranges) + .await + .map_err(to_parquet_err) + } + .boxed() } fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { self.file_metrics .bytes_scanned .add((range.end - range.start) as usize); - self.inner.get_bytes(range) + async move { + self.store + .get_range(&self.path, range) + .await + .map_err(to_parquet_err) + } + .boxed() } fn get_metadata( @@ -147,11 +159,15 @@ impl AsyncFileReader for ParquetMetadataCacheReader { match cache.entry(path.clone()) { std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()), std::collections::hash_map::Entry::Vacant(entry) => { - let meta = self.inner.get_metadata(options.as_ref()).await?; - let meta = Arc::try_unwrap(meta).unwrap_or_else(|e| e.as_ref().clone()); + let file_size = self.file_size; + let meta = ParquetMetaDataReader::new() + .with_arrow_reader_options(options.as_ref()) + .with_prefetch_hint(self.metadata_size_hint) + .load_and_finish(&mut *self, file_size) + .await?; let mut reader = ParquetMetaDataReader::new_with_metadata(meta.clone()) .with_page_index_policy(PageIndexPolicy::Optional); - reader.load_page_index(&mut self.inner).await?; + reader.load_page_index(&mut *self).await?; let meta = Arc::new(reader.finish()?); entry.insert(meta.clone()); Ok(meta) @@ -226,7 +242,10 @@ impl LiquidParquetSource { self.metrics = metrics; self.predicate = Some(Arc::clone(&predicate)); - match PruningPredicate::try_new(Arc::clone(&predicate), Arc::clone(&file_schema)) { + match PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&file_schema)) + .try_build(Arc::clone(&predicate)) + { Ok(pruning_predicate) => { if !pruning_predicate.always_true() { self.pruning_predicate = Some(Arc::new(pruning_predicate)); @@ -355,4 +374,16 @@ impl FileSource for LiquidParquetSource { fn file_type(&self) -> &str { "liquid_parquet" } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots( + self.predicate + .iter() + .chain(self.projection.iter().map(|projection| &projection.expr)), + f, + ) + } } From 4761d52187451a7dcc896c3819ab0e4003b57b7d Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Tue, 1 Sep 2026 21:33:11 -0400 Subject: [PATCH 03/24] Make CI benchmark comparison robust (median warm metrics + filesystem warmup + tests) (#511) ### Motivation - The CI benchmark comparison showed large, order-dependent slowdowns because LiquidCache runs first and populates the OS page cache, making the baseline look artificially faster. - Single outlier warm iterations from CI pauses can make the mean-driven warm metric noisy and report false regressions. - Make minimal, targeted changes to the comparison and CI workflow so results reflect real performance differences. ### Description - Replace mean with median for warm-iteration metrics in `.github/compare_benchmarks.py` by importing `statistics` and updating `get_warm_metrics` to return median values. - Make the configured `--threshold` consistently used when highlighting regressions and computing the warm-time summary in `format_change_percentage` and the report generation. - Add a lightweight unit test file `.github/test_compare_benchmarks.py` that verifies median warm metrics and threshold-based highlighting. - Update CI in `.github/workflows/ci.yml` to build the release `in_process` binary once, run a small unmeasured DataFusion warmup to populate the filesystem page cache, and invoke `target/release/in_process` for measured runs to remove order-dependent bias. ### Testing - Ran unit tests `python3 .github/test_compare_benchmarks.py`, and they passed (2 tests, OK). - Linted and checked Python files with `ruff` and `python -m py_compile`, both passed. - Verified `cargo check -p liquid-cache-benchmarks` for the benchmark crate succeeded. - A full workspace `cargo check` encountered a pre-existing environment-specific failure due to a missing generated asset (`dev/dev-tools/assets/tailwind.css`) that is unrelated to these benchmark changes. ------ [Codex Task](https://chatgpt.com/codex/cloud/tasks/task_e_6a96f83694ec8332ae8d6fade89739ef) --- .github/compare_benchmarks.py | 41 +++++++++++++++++++----------- .github/test_compare_benchmarks.py | 38 +++++++++++++++++++++++++++ .github/workflows/ci.yml | 16 +++++++++--- 3 files changed, 77 insertions(+), 18 deletions(-) create mode 100644 .github/test_compare_benchmarks.py diff --git a/.github/compare_benchmarks.py b/.github/compare_benchmarks.py index b8ea342be..5840e493e 100644 --- a/.github/compare_benchmarks.py +++ b/.github/compare_benchmarks.py @@ -9,6 +9,7 @@ import json import sys import argparse +import statistics from typing import Dict, List, Any @@ -47,18 +48,22 @@ def get_cold_metrics(iteration_results: List[Dict[str, Any]]) -> Dict[str, float def get_warm_metrics(iteration_results: List[Dict[str, Any]]) -> Dict[str, float]: - """Calculate average metrics from warm iterations (excluding first).""" + """Calculate median metrics from warm iterations (excluding first). + + CI runners occasionally pause a process while it is being timed. A median + keeps one such pause from turning into a reported regression. + """ warm_results = iteration_results[1:] if len(iteration_results) > 1 else iteration_results if not warm_results: return {"time_millis": 0, "cache_cpu_time": 0} - avg_time = sum(r["time_millis"] for r in warm_results) / len(warm_results) - avg_cpu_time = sum(r.get("cache_cpu_time", 0) for r in warm_results) / len(warm_results) + median_time = statistics.median(r["time_millis"] for r in warm_results) + median_cpu_time = statistics.median(r.get("cache_cpu_time", 0) for r in warm_results) # No memory column in report return { - "time_millis": avg_time, - "cache_cpu_time": avg_cpu_time, + "time_millis": median_time, + "cache_cpu_time": median_cpu_time, } @@ -79,15 +84,17 @@ def format_metric_with_baseline(current: float, baseline: float, formatter_func) return f"{formatter_func(current)} *({formatter_func(baseline)})*" -def format_change_percentage(current: float, baseline: float, highlight_mode: str = "none") -> str: +def format_change_percentage( + current: float, baseline: float, threshold: float, highlight_mode: str = "none" +) -> str: """Format percentage change and optionally highlight when slower. highlight_mode: - "none": never bold - - "slower_only": bold only if current > baseline (i.e., slower) and ≥15% + - "slower_only": bold only if current exceeds baseline by the threshold """ change_pct = calculate_change(baseline, current) - if highlight_mode == "slower_only" and change_pct > 0 and abs(change_pct) >= 15.0: + if highlight_mode == "slower_only" and change_pct >= threshold: return f"**{change_pct:+.1f}%**" return f"{change_pct:+.1f}%" @@ -195,21 +202,21 @@ def extract_mode(d: Dict[str, Any]) -> str: comp['curr_cold_time'], comp['baseline_cold_time'], format_time ) cold_change_str = format_change_percentage( - comp['curr_cold_time'], comp['baseline_cold_time'], highlight_mode="none" + comp['curr_cold_time'], comp['baseline_cold_time'], threshold, highlight_mode="none" ) warm_time_str = format_metric_with_baseline( comp['curr_warm_time'], comp['baseline_warm_time'], format_time ) warm_change_str = format_change_percentage( - comp['curr_warm_time'], comp['baseline_warm_time'], highlight_mode="slower_only" + comp['curr_warm_time'], comp['baseline_warm_time'], threshold, highlight_mode="slower_only" ) cpu_time_str = format_metric_with_baseline( comp['curr_cpu_time'], comp['baseline_cpu_time'], format_time ) cpu_change_str = format_change_percentage( - comp['curr_cpu_time'], comp['baseline_cpu_time'], highlight_mode="none" + comp['curr_cpu_time'], comp['baseline_cpu_time'], threshold, highlight_mode="none" ) lines.append( @@ -220,7 +227,7 @@ def extract_mode(d: Dict[str, Any]) -> str: ) # Summary focused on LiquidCache being slower than DataFusion (warm time) - slower_warm = [c for c in comparison if c["warm_time_change"] > 0] + slower_warm = [c for c in comparison if c["warm_time_change"] >= threshold] lines.append("") if slower_warm: lines.append(f"**⚠️ LiquidCache is slower on {len(slower_warm)} queries (warm)**") @@ -230,18 +237,22 @@ def extract_mode(d: Dict[str, Any]) -> str: slower_warm, key=lambda x: x["warm_time_change"], reverse=True ) for c in slower_warm_sorted: - curr = c["curr_warm_time"]; base = c["baseline_warm_time"] + curr = c["curr_warm_time"] + base = c["baseline_warm_time"] pct = calculate_change(base, curr) lines.append( f"- Q{c['query']}: warm {pct:+.1f}% " f"({format_time(curr)} vs {format_time(base)})" ) else: - lines.append("✅ LiquidCache is faster or equal on warm time for all queries") + lines.append(f"✅ No warm-time regression met the {threshold:.0f}% threshold") lines.append("") lines.append(f"*Compared {current_mode} vs {baseline_mode} on the same runner*") - lines.append("*Cold Time: first iteration; Warm Time: average of remaining iterations.*") + lines.append( + f"*Regressions: warm-time increases of at least {threshold:.0f}%. " + "Cold Time: first iteration; Warm Time: median of remaining iterations.*" + ) return "\n".join(lines) diff --git a/.github/test_compare_benchmarks.py b/.github/test_compare_benchmarks.py new file mode 100644 index 000000000..9fc0845ff --- /dev/null +++ b/.github/test_compare_benchmarks.py @@ -0,0 +1,38 @@ +import importlib.util +import pathlib +import unittest + + +SCRIPT = pathlib.Path(__file__).with_name("compare_benchmarks.py") +SPEC = importlib.util.spec_from_file_location("compare_benchmarks", SCRIPT) +compare = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(compare) + + +class CompareBenchmarksTest(unittest.TestCase): + def test_warm_metrics_use_median(self): + iterations = [ + {"time_millis": 100, "cache_cpu_time": 100}, + {"time_millis": 10, "cache_cpu_time": 1}, + {"time_millis": 11, "cache_cpu_time": 2}, + {"time_millis": 500, "cache_cpu_time": 100}, + ] + + self.assertEqual( + compare.get_warm_metrics(iterations), + {"time_millis": 11, "cache_cpu_time": 2}, + ) + + def test_highlight_respects_configured_threshold(self): + self.assertEqual( + compare.format_change_percentage(114, 100, 15, "slower_only"), + "+14.0%", + ) + self.assertEqual( + compare.format_change_percentage(114, 100, 10, "slower_only"), + "**+14.0%**", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 541591729..7949e230b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -363,10 +363,20 @@ jobs: - name: Build benchmark binary run: cargo build --release --bin in_process + # The LiquidCache run reads every query's columns before the DataFusion + # run. Without an explicit warmup, Linux's page cache makes DataFusion's + # "cold" measurements warm and creates a large, order-dependent bias. + - name: Warm filesystem cache + run: | + env RUST_LOG=warn target/release/in_process \ + --manifest benchmark/clickbench/benchmark_manifest.json \ + --iteration 1 \ + --bench-mode datafusion-default + - name: Run LiquidCache benchmark (in-process) run: | mkdir -p benchmark_results - env RUST_LOG=info cargo run --release --bin in_process -- \ + env RUST_LOG=info target/release/in_process \ --manifest benchmark/clickbench/benchmark_manifest.json \ --output benchmark_results/liquid.json \ --iteration 5 \ @@ -376,7 +386,7 @@ jobs: - name: Run DataFusion benchmark (plain parquet) run: | - env RUST_LOG=info cargo run --release --bin in_process -- \ + env RUST_LOG=info target/release/in_process \ --manifest benchmark/clickbench/benchmark_manifest.json \ --output benchmark_results/parquet.json \ --iteration 5 \ @@ -384,7 +394,7 @@ jobs: - name: Run DataFusion benchmark (default config) run: | - env RUST_LOG=info cargo run --release --bin in_process -- \ + env RUST_LOG=info target/release/in_process \ --manifest benchmark/clickbench/benchmark_manifest.json \ --output benchmark_results/df_default.json \ --iteration 5 \ From cc9ed21493a543536200c608f28bcf538bbc1ea3 Mon Sep 17 00:00:00 2001 From: Eddie A Tejeda Date: Tue, 1 Sep 2026 18:58:49 -0700 Subject: [PATCH 04/24] fix(reader): do not apply limit before the row filter (#512) ### Problem `LiquidStreamBuilder::build` gives `limit` and `offset` to the reader. `plan_row_group` then cuts the row selection to the first `limit + offset` physical rows **before** the pushed-down row filter runs. So the limit counts scanned rows, not matched rows. For example, a query like `SELECT id FROM t WHERE tag = 'MATCH' LIMIT 10` drops every match that comes after the first 10 physical rows of a file. If the matches are later in the file, the query returns 0 rows, but there are matches. Parquet counts the limit against post-filter matches. DataFusion depends on this when it pushes `fetch` into a scan whose filters were absorbed (`pushdown_filters = true`, which local mode sets). To reproduce on `main`: write a 20-row file where rows 15-19 have `tag = 'MATCH'`, then run `SELECT id FROM t WHERE tag = 'MATCH' LIMIT 10`. You get 0 rows. Expected: 5 rows. ### Fix Use `limit` and `offset` only when the scan has no row filter. In that case scanned rows equal emitted rows, so the cut is correct. Filtered scans still get a limit. DataFusion's `FileStream` cuts the emitted batches with `FileScanConfig::limit` after the filter. ### Tests New file `src/datafusion-local/src/tests/filter_limit.rs`. Every test puts the matches at the physical end of the data, so a cut-then-filter scan returns too few rows: * one row group, * a limit that spans row groups (this also touches row-group statistics pruning), * a scan over two files. Each test runs its queries twice, so both the cold (parquet) path and the warm (liquid cache) path are covered. All three tests fail on `main` and pass with this change. Co-authored-by: Eddie A Tejeda <669988+eddietejeda@users.noreply.github.com> Co-authored-by: Xiangpeng Hao --- .../src/tests/filter_limit.rs | 225 ++++++++++++++++++ src/datafusion-local/src/tests/mod.rs | 1 + .../src/reader/runtime/liquid_stream.rs | 22 +- 3 files changed, 246 insertions(+), 2 deletions(-) create mode 100644 src/datafusion-local/src/tests/filter_limit.rs diff --git a/src/datafusion-local/src/tests/filter_limit.rs b/src/datafusion-local/src/tests/filter_limit.rs new file mode 100644 index 000000000..7f75d414a --- /dev/null +++ b/src/datafusion-local/src/tests/filter_limit.rs @@ -0,0 +1,225 @@ +//! Regression tests for `WHERE LIMIT n` (no ORDER BY). +//! +//! The scan-level limit used to truncate the row selection *before* the +//! pushed-down row filter ran, so any match past the first `limit + offset` +//! physical rows was silently dropped — `SELECT .. WHERE tag='MATCH' LIMIT 10` +//! returned 0 rows even with 5 real matches. These tables place every match at +//! the physical tail of the data, past any LIMIT-sized prefix, so a +//! prefix-then-filter scan returns strictly fewer rows than expected. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array, StringArray}; +use arrow::record_batch::RecordBatch; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::error::Result; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// 20 rows: `tag='other'` for id 0-14, `tag='MATCH'` for id 15-19 — all 5 +/// matches sit at the physical end of the file. +fn write_tail_match_file(path: &Path, max_row_group_size: Option) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("tag", DataType::Utf8, false), + ])); + let ids = Int64Array::from((0..20i64).collect::>()); + let tags = StringArray::from( + (0..20) + .map(|i| if i >= 15 { "MATCH" } else { "other" }) + .collect::>(), + ); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(tags)]).unwrap(); + let props = max_row_group_size.map(|n| { + WriterProperties::builder() + .set_max_row_group_row_count(Some(n)) + .build() + }); + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, props).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +async fn liquid_ctx(cache_dir: &Path) -> Result { + std::fs::create_dir_all(cache_dir)?; + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = 4; + let (ctx, _cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir.to_path_buf()) + .build(config) + .await?; + Ok(ctx) +} + +/// Runs `sql` twice and asserts the row count both times, so the second +/// execution exercises the cached path where one exists. (Only the first +/// `assert_rows` against a table starts from a cold cache; later calls in the +/// same test run warm/warm.) +async fn assert_rows(ctx: &SessionContext, sql: &str, expected: usize) { + for run in ["cold", "warm"] { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, expected, "{run} run returned wrong row count: {sql}"); + } +} + +/// All returned ids must be actual matches (>= 15), not prefix rows. +async fn assert_ids_are_matches(ctx: &SessionContext, sql: &str, expected_len: usize) { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|b| { + let col = b + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + (0..col.len()).map(move |i| col.value(i)) + }) + .collect(); + ids.sort_unstable(); + assert_eq!(ids.len(), expected_len, "wrong number of rows: {sql}"); + assert!( + ids.iter().all(|id| *id >= 15), + "returned non-matching rows {ids:?}: {sql}" + ); +} + +#[tokio::test] +async fn filter_with_limit_single_row_group() { + let dir = TempDir::new().unwrap(); + let file = dir.path().join("tail20.parquet"); + write_tail_match_file(&file, None); + let ctx = liquid_ctx(&dir.path().join("cache")).await.unwrap(); + ctx.register_parquet( + "tail20", + file.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + // Baselines. + assert_rows(&ctx, "SELECT id FROM tail20 WHERE tag='MATCH'", 5).await; + // The bug: LIMIT larger than the match count must still return every + // match, even though all matches sit past the first LIMIT physical rows. + assert_rows(&ctx, "SELECT id FROM tail20 WHERE tag='MATCH' LIMIT 10", 5).await; + // LIMIT smaller than the match count caps *matches*, not scanned rows. + assert_rows(&ctx, "SELECT id FROM tail20 WHERE tag='MATCH' LIMIT 3", 3).await; + // ORDER BY variant was never broken (fetch stays in the TopK sort). + assert_rows( + &ctx, + "SELECT id FROM tail20 WHERE tag='MATCH' ORDER BY id LIMIT 10", + 5, + ) + .await; + // OFFSET is applied to filtered rows: skip 2 of the 5 matches. + assert_rows( + &ctx, + "SELECT id FROM tail20 WHERE tag='MATCH' LIMIT 18 OFFSET 2", + 3, + ) + .await; + // The rows themselves must be matches, not prefix rows. + assert_ids_are_matches(&ctx, "SELECT id FROM tail20 WHERE tag='MATCH' LIMIT 10", 5).await; + assert_ids_are_matches(&ctx, "SELECT id FROM tail20 WHERE tag='MATCH' LIMIT 3", 3).await; + + // Numeric predicate, same shape. + assert_rows(&ctx, "SELECT id FROM tail20 WHERE id >= 15 LIMIT 3", 3).await; +} + +#[tokio::test] +async fn filter_with_limit_across_row_groups() { + let dir = TempDir::new().unwrap(); + let file = dir.path().join("tail20_rg8.parquet"); + // Row groups of 8 rows: matches (rows 15-19) span the 2nd and 3rd groups. + write_tail_match_file(&file, Some(8)); + let ctx = liquid_ctx(&dir.path().join("cache")).await.unwrap(); + ctx.register_parquet( + "tail20_rg8", + file.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + assert_rows(&ctx, "SELECT id FROM tail20_rg8 WHERE tag='MATCH'", 5).await; + assert_rows( + &ctx, + "SELECT id FROM tail20_rg8 WHERE tag='MATCH' LIMIT 10", + 5, + ) + .await; + assert_rows( + &ctx, + "SELECT id FROM tail20_rg8 WHERE tag='MATCH' LIMIT 3", + 3, + ) + .await; + assert_ids_are_matches( + &ctx, + "SELECT id FROM tail20_rg8 WHERE tag='MATCH' LIMIT 10", + 5, + ) + .await; +} + +#[tokio::test] +async fn filter_with_limit_across_files() { + let dir = TempDir::new().unwrap(); + let table_dir = dir.path().join("two_files"); + std::fs::create_dir_all(&table_dir).unwrap(); + // Two identical files, 5 tail matches each: 10 matches total. The broken + // behavior applied the limit per file pre-filter, so LIMIT 17 returned 4 + // rows (2 per file) instead of 10. + write_tail_match_file(&table_dir.join("a.parquet"), None); + write_tail_match_file(&table_dir.join("b.parquet"), None); + let ctx = liquid_ctx(&dir.path().join("cache")).await.unwrap(); + ctx.register_parquet( + "two_files", + &format!("{}/", table_dir.to_str().unwrap()), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + assert_rows(&ctx, "SELECT id FROM two_files WHERE tag='MATCH'", 10).await; + assert_rows( + &ctx, + "SELECT id FROM two_files WHERE tag='MATCH' LIMIT 10", + 10, + ) + .await; + assert_rows( + &ctx, + "SELECT id FROM two_files WHERE tag='MATCH' LIMIT 17", + 10, + ) + .await; + assert_rows( + &ctx, + "SELECT id FROM two_files WHERE tag='MATCH' LIMIT 40", + 10, + ) + .await; + // Global cap still enforced on filtered rows. + assert_rows( + &ctx, + "SELECT id FROM two_files WHERE tag='MATCH' LIMIT 7", + 7, + ) + .await; + assert_ids_are_matches( + &ctx, + "SELECT id FROM two_files WHERE tag='MATCH' LIMIT 7", + 7, + ) + .await; +} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 776dd3a7e..998052ce1 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -23,6 +23,7 @@ use datafusion::{ use crate::LiquidCacheLocalBuilder; mod date_optimizer; +mod filter_limit; mod squeeze; mod variants; diff --git a/src/datafusion/src/reader/runtime/liquid_stream.rs b/src/datafusion/src/reader/runtime/liquid_stream.rs index 2933d5d2a..aa0b7c678 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -293,12 +293,30 @@ impl LiquidStreamBuilder { let file_schema = liquid_cache.schema(); let schema = build_projection_schema(&file_schema, &projection_column_ids); + // `plan_row_group` applies limit/offset by truncating the row + // selection BEFORE the row filter runs, so combining them with a + // filter caps the rows *scanned* rather than the rows *matched* — + // silently dropping matches that sit past the first `limit + offset` + // physical rows (upstream parquet counts the limit against + // post-filter matches instead). Until limit accounting moves after + // predicate evaluation, only honor limit/offset for unfiltered + // scans, where scanned rows == emitted rows and truncation is + // exact. Filtered scans still get capped post-filter by + // DataFusion's FileStream, which slices emitted batches against + // `FileScanConfig::limit`; all that is lost is scan-internal early + // termination. + let (limit, offset) = if self.filter.is_some() { + (None, None) + } else { + (self.limit, self.offset) + }; + let reader = ReaderFactory { metadata: Arc::clone(&self.metadata), input: self.input, filter: self.filter, - limit: self.limit, - offset: self.offset, + limit, + offset, cached_file: liquid_cache, }; From a0dadd78fa63fb256f6e8d3a8e0369e61adb5b3f Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Wed, 2 Sep 2026 20:47:26 -0400 Subject: [PATCH 05/24] Morsel based parquet loader, part 1 (#513) switch to the new api --- Cargo.lock | 1 + Cargo.toml | 1 + src/datafusion/Cargo.toml | 1 + src/datafusion/src/optimizers/mod.rs | 166 ++- src/datafusion/src/reader/plantime/mod.rs | 4 +- .../src/reader/plantime/morselizer.rs | 1138 +++++++++++++++++ src/datafusion/src/reader/plantime/opener.rs | 415 ------ .../src/reader/plantime/row_filter.rs | 26 +- .../src/reader/plantime/row_group_filter.rs | 427 ------- src/datafusion/src/reader/plantime/source.rs | 184 +-- .../src/reader/runtime/liquid_cache_reader.rs | 48 +- .../src/reader/runtime/liquid_stream.rs | 902 ------------- src/datafusion/src/reader/runtime/mod.rs | 5 +- src/datafusion/src/reader/runtime/morsel.rs | 171 +++ src/datafusion/src/reader/runtime/utils.rs | 63 +- 15 files changed, 1601 insertions(+), 1951 deletions(-) create mode 100644 src/datafusion/src/reader/plantime/morselizer.rs delete mode 100644 src/datafusion/src/reader/plantime/opener.rs delete mode 100644 src/datafusion/src/reader/plantime/row_group_filter.rs delete mode 100644 src/datafusion/src/reader/runtime/liquid_stream.rs create mode 100644 src/datafusion/src/reader/runtime/morsel.rs diff --git a/Cargo.lock b/Cargo.lock index ef27b66fb..76b296915 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4302,6 +4302,7 @@ dependencies = [ "arrow-schema", "bytes", "datafusion", + "datafusion-datasource", "divan", "fastrace", "futures", diff --git a/Cargo.toml b/Cargo.toml index ab7fe172c..28f79e68e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,7 @@ parquet = { version = "59.2.0", features = [ parquet-variant-json = { version = "59.2.0" } parquet-variant-compute = { version = "59.2.0" } datafusion = { version = "55.0.0" } +datafusion-datasource = { version = "55.0.0" } datafusion-common = { version = "55.0.0" } datafusion-expr-common = { version = "55.0.0" } datafusion-physical-expr = { version = "55.0.0" } diff --git a/src/datafusion/Cargo.toml b/src/datafusion/Cargo.toml index 0dfd8d590..4fc534077 100644 --- a/src/datafusion/Cargo.toml +++ b/src/datafusion/Cargo.toml @@ -12,6 +12,7 @@ arrow = { workspace = true } arrow-schema = { workspace = true } parquet = { workspace = true } datafusion = { workspace = true } +datafusion-datasource = { workspace = true } futures = { workspace = true } tokio = { workspace = true } ahash = { workspace = true } diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 47b842bbe..597df2113 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -123,23 +123,33 @@ fn convert_parquet_scan( #[cfg(test)] mod tests { - use datafusion::{datasource::physical_plan::FileScanConfig, prelude::SessionContext}; + use std::{fs::File, path::Path}; + + use arrow::{array::Int32Array, record_batch::RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::{ + common::{ScalarValue, stats::Precision}, + datasource::physical_plan::{FileScanConfig, FileSource}, + logical_expr::Operator, + physical_expr::expressions::{BinaryExpr, Column, Literal}, + physical_plan::{ + PhysicalExpr, collect, display::DisplayableExecutionPlan, filter_pushdown::PushedDown, + }, + prelude::SessionContext, + }; use liquid_cache::{ cache::{AlwaysHydrate, squeeze_policies::TranscodeSqueezeEvict}, cache_policies::LiquidPolicy, }; + use parquet::{arrow::ArrowWriter, file::properties::WriterProperties}; use crate::LiquidCacheParquet; use super::*; - async fn rewrite_plan_inner(plan: Arc) { - let expected_schema = plan.schema(); - let tmp_dir = tempfile::tempdir().unwrap(); - let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(); - let liquid_cache = Arc::new( + async fn make_cache(path: &Path) -> LiquidCacheParquetRef { + let store = t4::mount(path.join("liquid_cache.t4")).await.unwrap(); + Arc::new( LiquidCacheParquet::new( 8192, 1000000, @@ -150,7 +160,32 @@ mod tests { Box::new(AlwaysHydrate::new()), ) .await, - ); + ) + } + + fn liquid_source(plan: &Arc) -> LiquidParquetSource { + let mut source = None; + plan.apply(|node| { + if let Some(plan) = node.downcast_ref::() { + let config = plan.data_source().downcast_ref::().unwrap(); + source = Some( + config + .file_source() + .downcast_ref::() + .unwrap() + .clone(), + ); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + source.unwrap() + } + + async fn rewrite_plan_inner(plan: Arc) -> Arc { + let expected_schema = plan.schema(); + let tmp_dir = tempfile::tempdir().unwrap(); + let liquid_cache = make_cache(tmp_dir.path()).await; let rewritten = rewrite_data_source_plan(plan, &liquid_cache); rewritten @@ -167,6 +202,8 @@ mod tests { Ok(TreeNodeRecursion::Continue) }) .unwrap(); + + rewritten } #[tokio::test] @@ -184,6 +221,115 @@ mod tests { .await .unwrap(); let plan = df.create_physical_plan().await.unwrap(); - rewrite_plan_inner(plan.clone()).await; + let rewritten = rewrite_plan_inner(plan).await; + + let displayed = DisplayableExecutionPlan::new(rewritten.as_ref()) + .indent(true) + .to_string(); + assert!(displayed.contains("predicate="), "{displayed}"); + + rewritten + .apply(|node| { + if let Some(plan) = node.downcast_ref::() { + let statistics = plan.data_source().partition_statistics(None)?; + assert!(!matches!(statistics.num_rows, Precision::Exact(_))); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + + // Supported filters are conjoined onto the predicate; unsupported ones + // are handed back to the parent. + let source = liquid_source(&rewritten); + let url_index = source.table_schema().file_schema().index_of("URL").unwrap(); + let supported: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("URL", url_index)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "https://example.com".into(), + )))), + )); + let unsupported: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("missing", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("value".into())))), + )); + let result = source + .try_pushdown_filters( + vec![supported, unsupported], + &datafusion::config::ConfigOptions::new(), + ) + .unwrap(); + assert!(matches!( + result.filters.as_slice(), + [PushedDown::Yes, PushedDown::No] + )); + let predicate = result.updated_node.unwrap().filter().unwrap().to_string(); + assert!(predicate.contains(" AND "), "{predicate}"); + assert!(predicate.contains("https://example.com"), "{predicate}"); + assert!(!predicate.contains("missing"), "{predicate}"); + } + + fn write_bloom_file(path: &Path) { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let properties = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + let mut writer = ArrowWriter::try_new( + File::create(path).unwrap(), + schema.clone(), + Some(properties), + ) + .unwrap(); + for values in [[1, 2, 4], [1, 3, 4]] { + writer + .write( + &RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(values.to_vec()))], + ) + .unwrap(), + ) + .unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + } + + #[tokio::test] + async fn prunes_row_group_with_bloom_filter() { + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("bloom.parquet"); + write_bloom_file(&parquet_path); + + let ctx = SessionContext::new(); + ctx.register_parquet("t", parquet_path.to_str().unwrap(), Default::default()) + .await + .unwrap(); + let plan = ctx + .sql("SELECT * FROM t WHERE a = 2") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let cache = make_cache(tmp_dir.path()).await; + let rewritten = rewrite_data_source_plan(plan, &cache); + let metrics = liquid_source(&rewritten).metrics().clone(); + + let batches = collect(rewritten, ctx.task_ctx()).await.unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + + let metric = metrics + .clone_inner() + .sum_by_name("row_groups_pruned_bloom_filter") + .unwrap(); + let datafusion::physical_plan::metrics::MetricValue::PruningMetrics { + pruning_metrics, .. + } = metric + else { + panic!("unexpected metric: {metric:?}"); + }; + assert_eq!(pruning_metrics.pruned(), 1); } } diff --git a/src/datafusion/src/reader/plantime/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index b52a9bcf3..e87b02006 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -3,9 +3,9 @@ pub(crate) use source::CachedMetaReaderFactory; pub use source::LiquidParquetSource; pub(crate) use source::ParquetMetadataCacheReader; -mod opener; +mod morselizer; mod row_filter; -mod row_group_filter; mod source; +pub(crate) use morselizer::LiquidMorselizer; pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs new file mode 100644 index 000000000..210002e85 --- /dev/null +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -0,0 +1,1138 @@ +use std::{fmt, future::Future, sync::Arc}; + +use arrow_schema::SchemaRef; +use datafusion::{ + common::exec_err, + datasource::{ + listing::{FileRange, PartitionedFile}, + physical_plan::{ + ParquetFileMetrics, + parquet::{ + BloomFilterStatistics, PagePruningAccessPlanFilter, ParquetAccessPlan, + RowGroupAccessPlanFilter, + }, + }, + table_schema::TableSchema, + }, + error::Result, + physical_expr::{ + DynamicFilterTracking, PhysicalExpr, PhysicalExprSimplifier, projection::ProjectionExprs, + utils::reassign_expr_columns, + }, + physical_expr_adapter::{PhysicalExprAdapterFactory, replace_columns_with_literals}, + physical_optimizer::pruning::{FilePruner, PruningPredicate, build_pruning_predicate}, + physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}, +}; +use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselizer}; +use futures::{FutureExt, future::BoxFuture}; +use log::debug; +use parquet::{ + arrow::{ + ParquetRecordBatchStreamBuilder, ProjectionMask, + arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}, + parquet_column, + }, + file::metadata::PageIndexPolicy, +}; + +use crate::{ + cache::{ColumnSqueezeHints, LiquidCacheParquetRef}, + reader::{ + plantime::row_filter::build_row_filter, + runtime::{LiquidRowGroupPlanner, build_projection_schema, get_root_column_ids}, + }, +}; + +use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; + +pub(crate) struct LiquidMorselizer { + pub(crate) partition_index: usize, + pub(crate) projection: ProjectionExprs, + pub(crate) batch_size: usize, + pub(crate) predicate: Option>, + pub(crate) table_schema: TableSchema, + pub(crate) metrics: ExecutionPlanMetricsSet, + pub(crate) parquet_file_reader_factory: Arc, + pub(crate) reorder_filters: bool, + pub(crate) liquid_cache: LiquidCacheParquetRef, + pub(crate) expr_adapter_factory: Arc, + pub(crate) span: Option>, + pub(crate) squeeze_hints: Arc, +} + +impl fmt::Debug for LiquidMorselizer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LiquidMorselizer") + .field("partition_index", &self.partition_index) + .field("batch_size", &self.batch_size) + .finish_non_exhaustive() + } +} + +impl Morselizer for LiquidMorselizer { + fn plan_file(&self, partitioned_file: PartitionedFile) -> Result> { + let file_range = partitioned_file.range.clone(); + let access_plan = partitioned_file.extensions.get_arc::(); + let file_name = partitioned_file.object_meta.location.to_string(); + let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); + let metadata_size_hint = partitioned_file.metadata_size_hint; + let file_location = partitioned_file.object_meta.location.to_string(); + let reader = self.parquet_file_reader_factory.create_liquid_reader( + self.partition_index, + partitioned_file.clone(), + metadata_size_hint, + &self.metrics, + ); + + let logical_file_schema = Arc::clone(self.table_schema.file_schema()); + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); + let mut projection = self.projection.clone(); + let mut predicate = self.predicate.clone(); + let mut literal_columns = std::collections::HashMap::new(); + for (field, value) in self + .table_schema + .table_partition_cols() + .iter() + .zip(&partitioned_file.partition_values) + { + literal_columns.insert(field.name().clone(), value.clone()); + } + if !literal_columns.is_empty() { + projection = projection.try_map_exprs(|expr| { + replace_columns_with_literals(Arc::clone(&expr), &literal_columns) + })?; + predicate = predicate + .map(|predicate| replace_columns_with_literals(predicate, &literal_columns)) + .transpose()?; + } + + let predicate_creation_errors = + MetricBuilder::new(&self.metrics).global_counter("num_predicate_creation_errors"); + let file_pruner = predicate + .as_ref() + .filter(|predicate| { + DynamicFilterTracking::classify(predicate).contains_dynamic_filter() + || partitioned_file.has_statistics() + }) + .and_then(|predicate| { + FilePruner::try_new( + Arc::clone(predicate), + &logical_file_schema, + &partitioned_file, + predicate_creation_errors.clone(), + ) + }); + let span = self.span.as_ref().map(|span| { + Arc::new(fastrace::Span::enter_with_parent( + format!("file_{file_name}"), + span, + )) + }); + + Ok(Box::new(LiquidFilePlanner { + state: LiquidOpenState::PruneFile(Box::new(PreparedLiquidOpen { + file_range, + access_plan, + file_name, + file_metrics, + file_pruner, + reader, + batch_size: self.batch_size, + logical_file_schema, + output_schema, + projection, + predicate, + predicate_creation_errors, + reorder_filters: self.reorder_filters, + liquid_cache: self.liquid_cache.clone(), + expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), + file_location, + span, + squeeze_hints: Arc::clone(&self.squeeze_hints), + })), + })) + } +} + +struct PreparedLiquidOpen { + file_range: Option, + access_plan: Option>, + file_name: String, + file_metrics: ParquetFileMetrics, + file_pruner: Option, + reader: ParquetMetadataCacheReader, + batch_size: usize, + logical_file_schema: SchemaRef, + output_schema: SchemaRef, + projection: ProjectionExprs, + predicate: Option>, + predicate_creation_errors: Count, + reorder_filters: bool, + liquid_cache: LiquidCacheParquetRef, + expr_adapter_factory: Arc, + file_location: String, + span: Option>, + squeeze_hints: Arc, +} + +struct MetadataLoadedLiquidOpen { + prepared: Box, + reader_metadata: ArrowReaderMetadata, + options: ArrowReaderOptions, +} + +struct PreparedRowGroups { + context: RowGroupPlanningContext, + row_groups: RowGroupAccessPlanFilter, +} + +struct RowGroupPlanningContext { + prepared: Box, + reader_metadata: ArrowReaderMetadata, + physical_file_schema: SchemaRef, + cache_full_schema: SchemaRef, + builder: ParquetRecordBatchStreamBuilder, + projection_mask: ProjectionMask, + row_filter: Option, + pruning_predicate: Option>, + page_pruning_predicate: Option>, +} + +struct BloomFiltersLoadedLiquidOpen { + prepared: PreparedRowGroups, + bloom_filters: Vec, +} + +struct PlannedRowGroups { + context: RowGroupPlanningContext, + access_plan: ParquetAccessPlan, +} + +enum LiquidOpenState { + PruneFile(Box), + LoadMetadata(BoxFuture<'static, Result>), + PrepareAndPruneByStats(Box), + LoadBloomFilters(BoxFuture<'static, Result>), + PruneBloomAndPages(Box), + PlanRowGroups(Box), + Done, +} + +impl fmt::Debug for LiquidOpenState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::PruneFile(_) => "PruneFile", + Self::LoadMetadata(_) => "LoadMetadata", + Self::PrepareAndPruneByStats(_) => "PrepareAndPruneByStats", + Self::LoadBloomFilters(_) => "LoadBloomFilters", + Self::PruneBloomAndPages(_) => "PruneBloomAndPages", + Self::PlanRowGroups(_) => "PlanRowGroups", + Self::Done => "Done", + }) + } +} + +impl LiquidOpenState { + fn transition(self) -> Result { + match self { + Self::PruneFile(mut prepared) => { + if let Some(file_pruner) = &mut prepared.file_pruner + && file_pruner.should_prune()? + { + prepared + .file_metrics + .files_ranges_pruned_statistics + .add_pruned(1); + return Ok(Self::Done); + } + + prepared + .file_metrics + .files_ranges_pruned_statistics + .add_matched(1); + Ok(Self::LoadMetadata( + async move { + let options = ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let metadata_load_time = prepared.file_metrics.metadata_load_time.clone(); + let mut timer = metadata_load_time.timer(); + let reader_metadata = + ArrowReaderMetadata::load_async(&mut prepared.reader, options.clone()) + .await?; + timer.stop(); + Ok(MetadataLoadedLiquidOpen { + prepared, + reader_metadata, + options, + }) + } + .boxed(), + )) + } + Self::LoadMetadata(future) => Ok(Self::LoadMetadata(future)), + Self::PrepareAndPruneByStats(loaded) => prepare_and_prune_by_stats(*loaded), + Self::LoadBloomFilters(future) => Ok(Self::LoadBloomFilters(future)), + Self::PruneBloomAndPages(loaded) => { + let mut prepared = loaded.prepared; + let predicate = prepared + .context + .pruning_predicate + .as_deref() + .expect("bloom filters are loaded only with a pruning predicate"); + prepared.row_groups.prune_by_bloom_filters( + predicate, + &prepared.context.prepared.file_metrics, + &loaded.bloom_filters, + ); + Ok(Self::PlanRowGroups(Box::new(prune_pages(prepared)))) + } + Self::PlanRowGroups(planned) => Ok(Self::PlanRowGroups(planned)), + Self::Done => Ok(Self::Done), + } + } +} + +fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result { + let metadata_load_time = loaded.prepared.file_metrics.metadata_load_time.clone(); + let mut metadata_timer = metadata_load_time.timer(); + let physical_file_schema = Arc::clone(loaded.reader_metadata.schema()); + let cache_full_schema = Arc::clone(&physical_file_schema); + loaded.options = loaded + .options + .with_schema(Arc::clone(&physical_file_schema)); + loaded.reader_metadata = ArrowReaderMetadata::try_new( + Arc::clone(loaded.reader_metadata.metadata()), + loaded.options, + )?; + debug_assert!( + Arc::strong_count(loaded.reader_metadata.metadata()) > 1, + "meta data must be cached already" + ); + + let rewriter = loaded.prepared.expr_adapter_factory.create( + Arc::clone(&loaded.prepared.logical_file_schema), + Arc::clone(&physical_file_schema), + )?; + let simplifier = PhysicalExprSimplifier::new(&physical_file_schema); + loaded.prepared.predicate = loaded + .prepared + .predicate + .take() + .map(|predicate| simplifier.simplify(rewriter.rewrite(predicate)?)) + .transpose()?; + loaded.prepared.projection = loaded + .prepared + .projection + .try_map_exprs(|expr| simplifier.simplify(rewriter.rewrite(expr)?))?; + + let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates( + loaded.prepared.predicate.as_ref(), + &physical_file_schema, + &loaded.prepared.predicate_creation_errors, + ); + metadata_timer.stop(); + let builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + loaded.prepared.reader.clone(), + loaded.reader_metadata.clone(), + ); + let projection_mask = ProjectionMask::roots( + builder.parquet_schema(), + loaded.prepared.projection.column_indices(), + ); + let row_filter = + loaded.prepared.predicate.as_ref().and_then(|predicate| { + match build_row_filter( + predicate, + &physical_file_schema, + loaded.reader_metadata.metadata(), + loaded.prepared.reorder_filters, + &loaded.prepared.file_metrics, + ) { + Ok(filter) => filter, + Err(error) => { + debug!( + "Ignoring error building row filter for '{:?}': {error:?}", + loaded.prepared.predicate + ); + None + } + } + }); + + let metadata = builder.metadata(); + let row_group_metadata = metadata.row_groups(); + let access_plan = create_initial_plan( + &loaded.prepared.file_name, + loaded.prepared.access_plan.take(), + row_group_metadata.len(), + )?; + let mut row_groups = RowGroupAccessPlanFilter::new(access_plan); + if let Some(range) = &loaded.prepared.file_range { + row_groups.prune_by_range(row_group_metadata, range); + } + if let Some(predicate) = pruning_predicate.as_deref() { + row_groups.prune_by_statistics( + &physical_file_schema, + builder.parquet_schema(), + row_group_metadata, + predicate, + &loaded.prepared.file_metrics, + ); + } + + let prepared = PreparedRowGroups { + context: RowGroupPlanningContext { + prepared: loaded.prepared, + reader_metadata: loaded.reader_metadata, + physical_file_schema, + cache_full_schema, + builder, + projection_mask, + row_filter, + pruning_predicate, + page_pruning_predicate, + }, + row_groups, + }; + if prepared.context.pruning_predicate.is_some() && !prepared.row_groups.is_empty() { + Ok(LiquidOpenState::LoadBloomFilters( + async move { + let mut prepared = prepared; + let predicate = Arc::clone( + prepared + .context + .pruning_predicate + .as_ref() + .expect("pruning predicate was checked before scheduling bloom I/O"), + ); + let bloom_filters = load_bloom_filters( + &mut prepared.context.builder, + predicate.as_ref(), + &prepared.context.prepared.file_metrics, + &prepared.row_groups, + ) + .await; + Ok(BloomFiltersLoadedLiquidOpen { + prepared, + bloom_filters, + }) + } + .boxed(), + )) + } else { + Ok(LiquidOpenState::PlanRowGroups(Box::new(prune_pages( + prepared, + )))) + } +} + +fn prune_pages(prepared: PreparedRowGroups) -> PlannedRowGroups { + let PreparedRowGroups { + context, + row_groups, + } = prepared; + let mut access_plan = row_groups.build(); + if !access_plan.is_empty() + && let Some(predicate) = &context.page_pruning_predicate + { + access_plan = predicate.prune_plan_with_page_index( + access_plan, + &context.physical_file_schema, + context.builder.parquet_schema(), + context.builder.metadata().as_ref(), + &context.prepared.file_metrics, + ); + } + PlannedRowGroups { + context, + access_plan, + } +} + +struct LiquidFilePlanner { + state: LiquidOpenState, +} + +impl fmt::Debug for LiquidFilePlanner { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("LiquidFilePlanner") + .field(&self.state) + .finish() + } +} + +impl LiquidFilePlanner { + fn schedule_io(future: F) -> MorselPlan + where + F: Future> + Send + 'static, + { + let future = async move { + let state = future.await?; + Ok(Box::new(Self { state }) as Box) + }; + MorselPlan::new().with_pending_planner(future) + } +} + +impl MorselPlanner for LiquidFilePlanner { + fn plan(self: Box) -> Result> { + let state = self.state.transition()?; + match state { + LiquidOpenState::LoadMetadata(future) => Ok(Some(Self::schedule_io(async move { + Ok(LiquidOpenState::PrepareAndPruneByStats(Box::new( + future.await?, + ))) + }))), + LiquidOpenState::LoadBloomFilters(future) => Ok(Some(Self::schedule_io(async move { + Ok(LiquidOpenState::PruneBloomAndPages(Box::new(future.await?))) + }))), + LiquidOpenState::PlanRowGroups(planned) => plan_row_group_morsels(*planned), + LiquidOpenState::Done => Ok(None), + cpu_state => Ok(Some( + MorselPlan::new().with_planners(vec![Box::new(Self { state: cpu_state })]), + )), + } + } +} + +fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result> { + let PlannedRowGroups { + context, + access_plan, + } = planned; + let cached_file = context + .prepared + .liquid_cache + .register_or_get_file_with_hints( + context.prepared.file_location.clone(), + Arc::clone(&context.cache_full_schema), + Arc::clone(&context.prepared.squeeze_hints), + ); + let metadata = Arc::clone(context.reader_metadata.metadata()); + let schema_descriptor = metadata.file_metadata().schema_descr(); + let projection_column_ids = get_root_column_ids(schema_descriptor, &context.projection_mask); + let stream_schema = build_projection_schema(&cached_file.schema(), &projection_column_ids); + let replace_schema = !stream_schema.eq(&context.prepared.output_schema); + let projection = context + .prepared + .projection + .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; + let projector = Arc::new(projection.make_projector(&stream_schema)?); + let row_group_planner = LiquidRowGroupPlanner { + metadata: Arc::clone(&metadata), + input: context.prepared.reader.clone(), + row_filter: context.row_filter, + cached_file, + projection: context.projection_mask, + batch_size: context.prepared.batch_size, + stream_schema, + output_schema: Arc::clone(&context.prepared.output_schema), + projector, + replace_schema, + span: context.prepared.span, + }; + + let row_group_indexes = access_plan.row_group_indexes(); + let row_group_metadata = metadata.row_groups(); + let mut selection = access_plan.into_overall_row_selection(row_group_metadata)?; + let mut morsels: Vec> = Vec::with_capacity(row_group_indexes.len()); + for row_group_idx in row_group_indexes { + let row_count = row_group_metadata[row_group_idx].num_rows() as usize; + let row_group_selection = selection + .as_mut() + .map(|selection| selection.split_off(row_count)); + if let Some(morsel) = row_group_planner.plan(row_group_idx, row_group_selection) { + morsels.push(Box::new(morsel)); + } + } + + Ok((!morsels.is_empty()).then(|| MorselPlan::new().with_morsels(morsels))) +} + +async fn load_bloom_filters( + builder: &mut ParquetRecordBatchStreamBuilder, + predicate: &PruningPredicate, + file_metrics: &ParquetFileMetrics, + row_groups: &RowGroupAccessPlanFilter, +) -> Vec { + let mut row_group_bloom_filters = + vec![BloomFilterStatistics::new(); builder.metadata().num_row_groups()]; + let parquet_columns = predicate + .literal_columns() + .into_iter() + .filter_map(|column_name| { + let parquet_schema = builder.parquet_schema(); + let (column_idx, _) = parquet_column(parquet_schema, predicate.schema(), &column_name)?; + let column = parquet_schema.column(column_idx); + Some(( + column_name, + column_idx, + column.physical_type(), + column.type_length(), + )) + }) + .collect::>(); + + for row_group_idx in row_groups.row_group_indexes() { + let mut bloom_filters = BloomFilterStatistics::with_capacity(parquet_columns.len()); + for (column_name, column_idx, physical_type, type_length) in &parquet_columns { + let bloom_filter = match builder + .get_row_group_column_bloom_filter(row_group_idx, *column_idx) + .await + { + Ok(Some(bloom_filter)) => bloom_filter, + Ok(None) => continue, + Err(error) => { + debug!("Ignoring error reading bloom filter: {error}"); + file_metrics.predicate_evaluation_errors.add(1); + continue; + } + }; + bloom_filters.insert(column_name, bloom_filter, *physical_type, *type_length); + } + row_group_bloom_filters[row_group_idx] = bloom_filters; + } + + row_group_bloom_filters +} + +fn create_initial_plan( + file_name: &str, + access_plan: Option>, + row_group_count: usize, +) -> Result { + if let Some(access_plan) = access_plan { + let plan_len = access_plan.len(); + if plan_len != row_group_count { + return exec_err!( + "Invalid ParquetAccessPlan for {file_name}. Specified {plan_len} row groups, but file has {row_group_count}" + ); + } + return Ok(access_plan.as_ref().clone()); + } + + Ok(ParquetAccessPlan::new_all(row_group_count)) +} + +pub(crate) fn build_pruning_predicates( + predicate: Option<&Arc>, + file_schema: &SchemaRef, + predicate_creation_errors: &Count, +) -> ( + Option>, + Option>, +) { + let Some(predicate) = predicate else { + return (None, None); + }; + let pruning_predicate = build_pruning_predicate( + Arc::clone(predicate), + file_schema, + predicate_creation_errors, + ); + let page_pruning_predicate = build_page_pruning_predicate(predicate, file_schema); + (pruning_predicate, Some(page_pruning_predicate)) +} + +pub(crate) fn build_page_pruning_predicate( + predicate: &Arc, + file_schema: &SchemaRef, +) -> Arc { + Arc::new(PagePruningAccessPlanFilter::new( + predicate, + Arc::clone(file_schema), + )) +} + +#[cfg(test)] +mod tests { + use std::{ + collections::VecDeque, + fs::File, + sync::atomic::{AtomicUsize, Ordering}, + }; + + use arrow::{ + array::{Array, ArrayRef, Int32Array, RecordBatch}, + datatypes::{DataType, Field, Schema}, + }; + use datafusion::{ + common::ScalarValue, + datasource::{ + listing::PartitionedFile, + physical_plan::{FileScanConfigBuilder, FileSource, ParquetSource}, + }, + execution::object_store::ObjectStoreUrl, + logical_expr::Operator, + physical_expr::{ + PhysicalExpr, + expressions::{BinaryExpr, Column, Literal}, + projection::ProjectionExprs, + }, + physical_expr_adapter::DefaultPhysicalExprAdapterFactory, + physical_plan::metrics::ExecutionPlanMetricsSet, + }; + use futures::StreamExt; + use liquid_cache::{ + cache::{AlwaysHydrate, squeeze_policies::Evict}, + cache_policies::LiquidPolicy, + }; + use object_store::local::LocalFileSystem; + use parquet::arrow::{ArrowWriter, async_reader::AsyncFileReader}; + + use crate::{ + cache::{BatchID, CachedFileRef, CachedRowGroupRef, LiquidCacheParquet}, + reader::LiquidParquetSource, + }; + + use super::*; + + static NEXT_FILE_ID: AtomicUsize = AtomicUsize::new(0); + + struct PlannedTestFile { + morsels: Vec>, + _cache: Arc, + cached_file: CachedFileRef, + _tmp_dir: tempfile::TempDir, + } + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])) + } + + fn write_two_row_group_file(path: &std::path::Path, schema: SchemaRef) { + let file = File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer + .write( + &RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 11, 12, 13])), + ], + ) + .unwrap(), + ) + .unwrap(); + writer.flush().unwrap(); + writer + .write( + &RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6, 7])), + Arc::new(Int32Array::from(vec![14, 15, 16, 17])), + ], + ) + .unwrap(), + ) + .unwrap(); + writer.close().unwrap(); + } + + fn write_single_row_group_file(path: &std::path::Path, schema: SchemaRef, a: Vec) { + let file = File::create(path).unwrap(); + let b = a.iter().map(|value| value + 1000).collect::>(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(a)), Arc::new(Int32Array::from(b))], + ) + .unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + async fn drive_planner(planner: Box) -> Vec> { + let mut planners = VecDeque::from([planner]); + let mut morsels = Vec::new(); + while let Some(planner) = planners.pop_front() { + let Some(mut plan) = planner.plan().unwrap() else { + continue; + }; + morsels.extend(plan.take_morsels()); + planners.extend(plan.take_ready_planners()); + if let Some(pending) = plan.take_pending_planner() { + planners.push_back(pending.await.unwrap()); + } + } + morsels + } + + struct PlanOptions { + max_memory_bytes: usize, + max_disk_bytes: usize, + predicate: Option>, + projection_columns: Vec, + single_row_group_values: Option>, + } + + impl Default for PlanOptions { + fn default() -> Self { + Self { + max_memory_bytes: usize::MAX, + max_disk_bytes: usize::MAX, + predicate: None, + projection_columns: vec![0, 1], + single_row_group_values: None, + } + } + } + + async fn create_test_cache( + path: &std::path::Path, + max_memory_bytes: usize, + max_disk_bytes: usize, + ) -> Arc { + let store = t4::mount(path.join("liquid_cache.t4")).await.unwrap(); + Arc::new( + LiquidCacheParquet::new( + 4, + max_memory_bytes, + max_disk_bytes, + store, + Box::new(LiquidPolicy::new()), + Box::new(Evict), + Box::new(AlwaysHydrate::new()), + ) + .await, + ) + } + + async fn plan_test_file(options: PlanOptions) -> PlannedTestFile { + let schema = schema(); + let tmp_dir = tempfile::tempdir().unwrap(); + let file_id = NEXT_FILE_ID.fetch_add(1, Ordering::Relaxed); + let file_name = "data.parquet".to_string(); + let parquet_path = tmp_dir.path().join(&file_name); + if let Some(values) = options.single_row_group_values { + write_single_row_group_file(&parquet_path, Arc::clone(&schema), values); + } else { + write_two_row_group_file(&parquet_path, Arc::clone(&schema)); + } + let partitioned_file = PartitionedFile::new( + file_name.clone(), + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let cache = create_test_cache( + tmp_dir.path(), + options.max_memory_bytes, + options.max_disk_bytes, + ) + .await; + let metrics = ExecutionPlanMetricsSet::new(); + let morselizer = LiquidMorselizer { + partition_index: 0, + projection: ProjectionExprs::from_indices(&options.projection_columns, schema.as_ref()), + batch_size: 4, + predicate: options.predicate, + table_schema: TableSchema::from(Arc::clone(&schema)), + metrics: metrics.clone(), + parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( + object_store, + ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(), + )), + reorder_filters: false, + liquid_cache: cache.clone(), + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + span: None, + squeeze_hints: Arc::default(), + }; + let morsels = drive_planner(morselizer.plan_file(partitioned_file).unwrap()).await; + let cached_file = cache.register_or_get_file(file_name, schema); + PlannedTestFile { + morsels, + _cache: cache, + cached_file, + _tmp_dir: tmp_dir, + } + } + + fn gt_expr(column_name: &str, column_index: usize, literal: i32) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new(column_name, column_index)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(literal)))), + )) + } + + #[tokio::test] + async fn metadata_cache_is_scoped_to_object_store() { + let schema = schema(); + let dir_a = tempfile::tempdir().unwrap(); + let dir_b = tempfile::tempdir().unwrap(); + let path_a = dir_a.path().join("data.parquet"); + let path_b = dir_b.path().join("data.parquet"); + write_single_row_group_file(&path_a, schema.clone(), vec![1]); + write_single_row_group_file(&path_b, schema, vec![1, 2]); + let metrics = ExecutionPlanMetricsSet::new(); + let mut reader_a = CachedMetaReaderFactory::new( + Arc::new(LocalFileSystem::new_with_prefix(dir_a.path()).unwrap()), + ObjectStoreUrl::parse("store-a:///").unwrap(), + ) + .create_liquid_reader( + 0, + PartitionedFile::new("data.parquet", std::fs::metadata(path_a).unwrap().len()), + None, + &metrics, + ); + let mut reader_b = CachedMetaReaderFactory::new( + Arc::new(LocalFileSystem::new_with_prefix(dir_b.path()).unwrap()), + ObjectStoreUrl::parse("store-b:///").unwrap(), + ) + .create_liquid_reader( + 0, + PartitionedFile::new("data.parquet", std::fs::metadata(path_b).unwrap().len()), + None, + &metrics, + ); + + let metadata_a = reader_a.get_metadata(None).await.unwrap(); + let metadata_b = reader_b.get_metadata(None).await.unwrap(); + + assert_eq!(metadata_a.file_metadata().num_rows(), 1); + assert_eq!(metadata_b.file_metadata().num_rows(), 2); + } + + async fn collect_columns(morsels: Vec>) -> (Vec, Vec) { + let mut a = Vec::new(); + let mut b = Vec::new(); + for morsel in morsels { + let batches = morsel.into_stream().collect::>().await; + for batch in batches { + let batch = batch.unwrap(); + a.extend( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + if batch.num_columns() > 1 { + b.extend( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + } + } + } + (a, b) + } + + async fn insert_batches( + row_group: &CachedRowGroupRef, + column_id: usize, + batches: &[(u16, &[i32])], + ) { + let column = row_group.get_column(column_id as u64).unwrap(); + for (batch_idx, values) in batches { + let array: ArrayRef = Arc::new(Int32Array::from(values.to_vec())); + column + .insert(BatchID::from_raw(*batch_idx), array) + .await + .unwrap(); + } + } + + async fn is_cached(row_group: &CachedRowGroupRef, column_id: usize, batch_idx: u16) -> bool { + row_group + .get_column(column_id as u64) + .unwrap() + .get_arrow_array_test_only(BatchID::from_raw(batch_idx)) + .await + .is_some() + } + + #[tokio::test] + async fn plans_one_morsel_per_selected_row_group() { + let all = plan_test_file(PlanOptions { + ..Default::default() + }) + .await; + assert_eq!(all.morsels.len(), 2); + assert_eq!( + collect_columns(all.morsels).await.0, + vec![0, 1, 2, 3, 4, 5, 6, 7] + ); + + let pruned = plan_test_file(PlanOptions { + predicate: Some(gt_expr("a", 0, 3)), + ..Default::default() + }) + .await; + assert_eq!(pruned.morsels.len(), 1); + assert_eq!(collect_columns(pruned.morsels).await.0, vec![4, 5, 6, 7]); + } + + #[tokio::test] + async fn cache_full_keeps_inserted_batches_and_skips_failed_inserts() { + let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); + let planned = plan_test_file(PlanOptions { + max_memory_bytes: one_array_memory * 3, + max_disk_bytes: 0, + ..Default::default() + }) + .await; + let row_group0 = planned.cached_file.create_row_group(0, vec![]); + let row_group1 = planned.cached_file.create_row_group(1, vec![]); + + let (a, b) = collect_columns(planned.morsels).await; + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + assert!(is_cached(&row_group0, 0, 0).await); + assert!(is_cached(&row_group0, 1, 0).await); + assert!(is_cached(&row_group1, 0, 0).await); + assert!(!is_cached(&row_group1, 1, 0).await); + } + + #[tokio::test] + async fn cache_full_with_filter_keeps_results_correct() { + let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); + let planned = plan_test_file(PlanOptions { + max_memory_bytes: one_array_memory * 3, + max_disk_bytes: 0, + predicate: Some(gt_expr("a", 0, 2)), + ..Default::default() + }) + .await; + let row_group0 = planned.cached_file.create_row_group(0, vec![]); + let row_group1 = planned.cached_file.create_row_group(1, vec![]); + let (a, b) = collect_columns(planned.morsels).await; + assert_eq!(a, vec![3, 4, 5, 6, 7]); + assert_eq!(b, vec![13, 14, 15, 16, 17]); + assert!(is_cached(&row_group0, 0, 0).await); + assert!(is_cached(&row_group0, 1, 0).await); + assert!(is_cached(&row_group1, 0, 0).await); + assert!(!is_cached(&row_group1, 1, 0).await); + } + + #[tokio::test] + async fn mid_scan_eviction_recovers() { + let planned = plan_test_file(PlanOptions { + max_memory_bytes: 0, + max_disk_bytes: 0, + ..Default::default() + }) + .await; + let row_group0 = planned.cached_file.create_row_group(0, vec![]); + let row_group1 = planned.cached_file.create_row_group(1, vec![]); + let (a, b) = collect_columns(planned.morsels).await; + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + for row_group in [&row_group0, &row_group1] { + assert!(!is_cached(row_group, 0, 0).await); + assert!(!is_cached(row_group, 1, 0).await); + } + } + + #[tokio::test] + async fn predicate_fallback_uses_predicate_projection() { + let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); + let planned = plan_test_file(PlanOptions { + max_memory_bytes: one_array_memory * 3, + max_disk_bytes: 0, + predicate: Some(gt_expr("b", 1, 12)), + projection_columns: vec![0], + ..Default::default() + }) + .await; + let row_group0 = planned.cached_file.create_row_group(0, vec![]); + let row_group1 = planned.cached_file.create_row_group(1, vec![]); + assert_eq!( + collect_columns(planned.morsels).await.0, + vec![3, 4, 5, 6, 7] + ); + assert!(is_cached(&row_group0, 0, 0).await); + assert!(is_cached(&row_group0, 1, 0).await); + assert!(is_cached(&row_group1, 0, 0).await); + assert!(!is_cached(&row_group1, 1, 0).await); + } + + #[tokio::test] + async fn missing_column_falls_back_to_parquet() { + let planned = plan_test_file(PlanOptions { + ..Default::default() + }) + .await; + let row_group0 = planned.cached_file.create_row_group(0, vec![]); + let row_group1 = planned.cached_file.create_row_group(1, vec![]); + insert_batches(&row_group0, 0, &[(0, &[0, 1, 2, 3])]).await; + insert_batches(&row_group1, 0, &[(0, &[4, 5, 6, 7])]).await; + + let (a, b) = collect_columns(planned.morsels).await; + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + assert!(is_cached(&row_group0, 1, 0).await); + assert!(is_cached(&row_group1, 1, 0).await); + } + + #[tokio::test] + async fn fallback_stream_advances_across_misses() { + let parquet_a = vec![ + 100, 101, 102, 103, 4, 5, 6, 7, 200, 201, 202, 203, 12, 13, 14, 15, + ]; + let planned = plan_test_file(PlanOptions { + projection_columns: vec![0], + single_row_group_values: Some(parquet_a), + ..Default::default() + }) + .await; + let row_group = planned.cached_file.create_row_group(0, vec![]); + insert_batches(&row_group, 0, &[(0, &[0, 1, 2, 3]), (2, &[8, 9, 10, 11])]).await; + + assert_eq!( + collect_columns(planned.morsels).await.0, + (0..16).collect::>() + ); + for batch_idx in 0..4 { + assert!(is_cached(&row_group, 0, batch_idx).await); + } + } + + #[tokio::test] + async fn source_uses_native_morsel_api() { + let schema = schema(); + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("data.parquet"); + write_two_row_group_file(&parquet_path, Arc::clone(&schema)); + let file = PartitionedFile::new( + "data.parquet", + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let cache = create_test_cache(tmp_dir.path(), usize::MAX, usize::MAX).await; + let source = LiquidParquetSource::from_parquet_source( + ParquetSource::new(Arc::clone(&schema)), + cache, + ); + let base_config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()), + ) + .with_file(file.clone()) + .build(); + let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + + assert!( + source + .create_file_opener(object_store.clone(), &base_config, 0) + .is_err() + ); + let morselizer = source + .create_morselizer(object_store, &base_config, 0) + .unwrap(); + assert!(morselizer.plan_file(file).is_ok()); + } +} diff --git a/src/datafusion/src/reader/plantime/opener.rs b/src/datafusion/src/reader/plantime/opener.rs deleted file mode 100644 index f54890d01..000000000 --- a/src/datafusion/src/reader/plantime/opener.rs +++ /dev/null @@ -1,415 +0,0 @@ -use std::sync::Arc; - -use crate::{ - cache::{ColumnSqueezeHints, LiquidCacheParquetRef}, - reader::{ - plantime::{row_filter::build_row_filter, row_group_filter::RowGroupAccessPlanFilter}, - runtime::LiquidStreamBuilder, - }, -}; -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow_schema::SchemaRef; -use datafusion::{ - common::exec_err, - datasource::{ - listing::PartitionedFile, - physical_plan::{ - FileOpenFuture, FileOpener, ParquetFileMetrics, - parquet::{PagePruningAccessPlanFilter, ParquetAccessPlan}, - }, - table_schema::TableSchema, - }, - error::DataFusionError, - physical_expr::projection::ProjectionExprs, - physical_expr::utils::reassign_expr_columns, - physical_expr::{DynamicFilterTracking, PhysicalExprSimplifier}, - physical_expr_adapter::{PhysicalExprAdapterFactory, replace_columns_with_literals}, - physical_optimizer::pruning::{FilePruner, PruningPredicate, build_pruning_predicate}, - physical_plan::{ - PhysicalExpr, - metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}, - }, -}; -use futures::StreamExt; -use futures::TryStreamExt; -use log::debug; -use parquet::arrow::{ - ParquetRecordBatchStreamBuilder, ProjectionMask, - arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}, -}; -use parquet::file::metadata::ParquetMetaData; - -use super::source::CachedMetaReaderFactory; - -pub struct LiquidParquetOpener { - partition_index: usize, - projection: ProjectionExprs, - batch_size: usize, - limit: Option, - predicate: Option>, - table_schema: TableSchema, - metrics: ExecutionPlanMetricsSet, - parquet_file_reader_factory: Arc, - reorder_filters: bool, - liquid_cache: LiquidCacheParquetRef, - expr_adapter_factory: Arc, - span: Option>, - squeeze_hints: Arc, -} - -impl LiquidParquetOpener { - #[allow(clippy::too_many_arguments)] - pub fn new( - partition_index: usize, - projection: ProjectionExprs, - batch_size: usize, - limit: Option, - predicate: Option>, - table_schema: TableSchema, - metrics: ExecutionPlanMetricsSet, - liquid_cache: LiquidCacheParquetRef, - parquet_file_reader_factory: Arc, - reorder_filters: bool, - expr_adapter_factory: Arc, - span: Option>, - squeeze_hints: Arc, - ) -> Self { - Self { - partition_index, - projection, - batch_size, - limit, - predicate, - table_schema, - metrics, - liquid_cache, - parquet_file_reader_factory, - reorder_filters, - expr_adapter_factory, - span, - squeeze_hints, - } - } -} - -impl FileOpener for LiquidParquetOpener { - fn open(&self, partitioned_file: PartitionedFile) -> Result { - let file_range = partitioned_file.range.clone(); - let access_plan_ext = partitioned_file.extensions.get_arc::(); - let file_name = partitioned_file.object_meta.location.to_string(); - let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); - - let metadata_size_hint = partitioned_file.metadata_size_hint; - - let lc = self.liquid_cache.clone(); - let file_loc = partitioned_file.object_meta.location.to_string(); - - let mut async_file_reader = self.parquet_file_reader_factory.create_liquid_reader( - self.partition_index, - partitioned_file.clone(), - metadata_size_hint, - &self.metrics, - ); - - let batch_size = self.batch_size; - let logical_file_schema = Arc::clone(self.table_schema.file_schema()); - let output_schema = Arc::new( - self.projection - .project_schema(self.table_schema.table_schema())?, - ); - let mut projection = self.projection.clone(); - let mut predicate = self.predicate.clone(); - let mut literal_columns = std::collections::HashMap::new(); - for (field, value) in self - .table_schema - .table_partition_cols() - .iter() - .zip(partitioned_file.partition_values.iter()) - { - literal_columns.insert(field.name().clone(), value.clone()); - } - if !literal_columns.is_empty() { - projection = projection.try_map_exprs(|expr| { - replace_columns_with_literals(Arc::clone(&expr), &literal_columns) - })?; - predicate = predicate - .map(|p| replace_columns_with_literals(p, &literal_columns)) - .transpose()?; - } - let reorder_predicates = self.reorder_filters; - let limit = self.limit; - - let predicate_creation_errors = - MetricBuilder::new(&self.metrics).global_counter("num_predicate_creation_errors"); - - let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); - let span = self.span.clone(); - let squeeze_hints = Arc::clone(&self.squeeze_hints); - Ok(Box::pin(async move { - // Prune this file using the file level statistics and partition values. - // Since dynamic filters may have been updated since planning it is possible that we are able - // to prune files now that we couldn't prune at planning time. - // It is assumed that there is no point in doing pruning here if the predicate is not dynamic, - // as it would have been done at planning time. - // We'll also check this after every record batch we read, - // and if at some point we are able to prove we can prune the file using just the file level statistics - // we can end the stream early. - let mut file_pruner = predicate - .as_ref() - .filter(|p| { - DynamicFilterTracking::classify(p).contains_dynamic_filter() - || partitioned_file.has_statistics() - }) - .and_then(|p| { - FilePruner::try_new( - Arc::clone(p), - &logical_file_schema, - &partitioned_file, - predicate_creation_errors.clone(), - ) - }); - - if let Some(file_pruner) = &mut file_pruner - && file_pruner.should_prune()? - { - file_metrics.files_ranges_pruned_statistics.add_pruned(1); - return Ok(futures::stream::empty().boxed()); - } - - file_metrics.files_ranges_pruned_statistics.add_matched(1); - - let mut options = ArrowReaderOptions::new() - .with_page_index_policy(parquet::file::metadata::PageIndexPolicy::Required); - let mut metadata_timer = file_metrics.metadata_load_time.timer(); - - // Begin by loading the metadata from the underlying reader (note - // the returned metadata may actually include page indexes as some - // readers may return page indexes even when not requested -- for - // example when they are cached) - let mut reader_metadata = - ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await?; - - // Note about schemas: we are actually dealing with **3 different schemas** here: - // - The table schema as defined by the TableProvider. - // This is what the user sees, what they get when they `SELECT * FROM table`, etc. - // - The logical file schema: this is the table schema minus any hive partition columns and projections. - // This is what the physical file schema is coerced to. - // - The physical file schema: this is the schema as defined by the parquet file. This is what the parquet file actually contains. - let physical_file_schema = Arc::clone(reader_metadata.schema()); - let cache_full_schema = Arc::clone(&physical_file_schema); - options = options.with_schema(Arc::clone(&physical_file_schema)); - reader_metadata = - ArrowReaderMetadata::try_new(Arc::clone(reader_metadata.metadata()), options)?; - debug_assert!( - Arc::strong_count(reader_metadata.metadata()) > 1, - "meta data must be cached already" - ); - - let rewriter = expr_adapter_factory.create( - Arc::clone(&logical_file_schema), - Arc::clone(&physical_file_schema), - )?; - let simplifier = PhysicalExprSimplifier::new(&physical_file_schema); - predicate = predicate - .map(|p| simplifier.simplify(rewriter.rewrite(p)?)) - .transpose()?; - projection = projection.try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?; - - let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates( - predicate.as_ref(), - &physical_file_schema, - &predicate_creation_errors, - ); - - metadata_timer.stop(); - - let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata( - async_file_reader.clone(), - reader_metadata.clone(), - ); - let indices = projection.column_indices(); - let mask = ProjectionMask::roots(builder.parquet_schema(), indices); - - // Filter pushdown: evaluate predicates during scan - let row_filter = predicate.as_ref().and_then(|p| { - let row_filter = build_row_filter( - p, - &physical_file_schema, - reader_metadata.metadata(), - reorder_predicates, - &file_metrics, - ); - - match row_filter { - Ok(Some(filter)) => Some(filter), - Ok(None) => None, - Err(e) => { - debug!("Ignoring error building row filter for '{predicate:?}': {e:?}"); - None - } - } - }); - - // Determine which row groups to actually read. The idea is to skip - // as many row groups as possible based on the metadata and query - let file_metadata: Arc = Arc::clone(builder.metadata()); - let predicate = pruning_predicate.as_ref().map(|p| p.as_ref()); - let rg_metadata = file_metadata.row_groups(); - // track which row groups to actually read - let access_plan = create_initial_plan(&file_name, access_plan_ext, rg_metadata.len())?; - let mut row_groups = RowGroupAccessPlanFilter::new(access_plan); - // if there is a range restricting what parts of the file to read - if let Some(range) = file_range.as_ref() { - row_groups.prune_by_range(rg_metadata, range); - } - // If there is a predicate that can be evaluated against the metadata - if let Some(predicate) = predicate.as_ref() { - row_groups.prune_by_statistics( - &physical_file_schema, - builder.parquet_schema(), - rg_metadata, - predicate, - &file_metrics, - ); - - if !row_groups.is_empty() { - row_groups - .prune_by_bloom_filters( - &physical_file_schema, - &mut builder, - predicate, - &file_metrics, - ) - .await; - } - } - - let mut access_plan = row_groups.build(); - - // page index pruning: if all data on individual pages can - // be ruled using page metadata, rows from other columns - // with that range can be skipped as well - if !access_plan.is_empty() - && let Some(p) = page_pruning_predicate - { - access_plan = p.prune_plan_with_page_index( - access_plan, - &physical_file_schema, - builder.parquet_schema(), - file_metadata.as_ref(), - &file_metrics, - ); - } - - let row_group_indexes = access_plan.row_group_indexes(); - let row_selection = access_plan.into_overall_row_selection(rg_metadata)?; - - let mut liquid_builder = - LiquidStreamBuilder::new(async_file_reader, Arc::clone(reader_metadata.metadata())) - .with_batch_size(batch_size) - .with_row_groups(row_group_indexes) - .with_projection(mask) - .with_selection(row_selection) - .with_limit(limit); - - if let Some(row_filter) = row_filter { - liquid_builder = liquid_builder.with_row_filter(row_filter); - } - - if let Some(s) = &span { - let span = fastrace::Span::enter_with_parent("liquid_stream", s); - liquid_builder = liquid_builder.with_span(span); - } - - let liquid_cache = lc.register_or_get_file_with_hints( - file_loc, - Arc::clone(&cache_full_schema), - squeeze_hints, - ); - - let stream = liquid_builder.build(liquid_cache)?; - - let stream_schema = Arc::clone(stream.schema()); - let replace_schema = !stream_schema.eq(&output_schema); - let projection = - projection.try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; - let projector = projection.make_projector(&stream_schema)?; - - let adapted = stream - .map_err(|e| DataFusionError::External(Box::new(e))) - .map(move |batch| { - batch.and_then(|batch| { - let batch = projector.project_batch(&batch)?; - if replace_schema { - let (_schema, arrays, num_rows) = batch.into_parts(); - let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); - RecordBatch::try_new_with_options( - Arc::clone(&output_schema), - arrays, - &options, - ) - .map_err(Into::into) - } else { - Ok(batch) - } - }) - }); - - Ok(adapted.boxed()) - })) - } -} - -fn create_initial_plan( - file_name: &str, - access_plan: Option>, - row_group_count: usize, -) -> Result { - if let Some(access_plan) = access_plan { - let plan_len = access_plan.len(); - if plan_len != row_group_count { - return exec_err!( - "Invalid ParquetAccessPlan for {file_name}. Specified {plan_len} row groups, but file has {row_group_count}" - ); - } - - // check row group count matches the plan - return Ok(access_plan.as_ref().clone()); - } - - // default to scanning all row groups - Ok(ParquetAccessPlan::new_all(row_group_count)) -} - -pub(crate) fn build_pruning_predicates( - predicate: Option<&Arc>, - file_schema: &SchemaRef, - predicate_creation_errors: &Count, -) -> ( - Option>, - Option>, -) { - let Some(predicate) = predicate.as_ref() else { - return (None, None); - }; - let pruning_predicate = build_pruning_predicate( - Arc::clone(predicate), - file_schema, - predicate_creation_errors, - ); - let page_pruning_predicate = build_page_pruning_predicate(predicate, file_schema); - (pruning_predicate, Some(page_pruning_predicate)) -} - -/// Build a page pruning predicate from an optional predicate expression. -/// If the predicate is None or the predicate cannot be converted to a page pruning -/// predicate, return None. -pub(crate) fn build_page_pruning_predicate( - predicate: &Arc, - file_schema: &SchemaRef, -) -> Arc { - Arc::new(PagePruningAccessPlanFilter::new( - predicate, - Arc::clone(file_schema), - )) -} diff --git a/src/datafusion/src/reader/plantime/row_filter.rs b/src/datafusion/src/reader/plantime/row_filter.rs index fe9c5ea16..d0c97c81a 100644 --- a/src/datafusion/src/reader/plantime/row_filter.rs +++ b/src/datafusion/src/reader/plantime/row_filter.rs @@ -84,6 +84,7 @@ use datafusion::physical_expr::expressions::Column; use datafusion::physical_expr::{PhysicalExpr, split_conjunction}; /// A row filter that can be used to filter rows from a parquet file. +#[derive(Clone)] pub struct LiquidRowFilter { predicates: Vec, } @@ -105,24 +106,6 @@ impl LiquidRowFilter { } } -pub(crate) fn get_predicate_column_id(projection: &parquet::arrow::ProjectionMask) -> Vec { - #[derive(Debug, Clone)] - struct ProjectionMaskLiquid { - mask: Option>, - } - let project_inner: &ProjectionMaskLiquid = unsafe { std::mem::transmute(projection) }; - project_inner - .mask - .as_ref() - .map(|m| { - m.iter() - .enumerate() - .filter_map(|(pos, &x)| if x { Some(pos) } else { None }) - .collect::>() - }) - .unwrap_or_default() -} - /// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform /// row-level filtering during parquet decoding. /// @@ -142,6 +125,9 @@ pub struct LiquidPredicate { /// Path to the columns in the parquet schema required to evaluate the /// expression projection_mask: ProjectionMask, + /// Indices into the file schema of the columns required to evaluate the + /// expression, in `filter_schema` order + column_ids: Vec, /// how many rows were filtered out by this predicate rows_pruned: metrics::Count, /// how many rows passed this predicate @@ -172,6 +158,7 @@ impl LiquidPredicate { physical_expr, physical_expr_physical_column_index: candidate.expr, projection_mask: projection, + column_ids: candidate.projection, rows_pruned, rows_matched, time, @@ -201,8 +188,7 @@ impl LiquidPredicate { /// Get the column ids of the predicate. pub fn predicate_column_ids(&self) -> Vec { - let projection = self.projection(); - get_predicate_column_id(projection) + self.column_ids.clone() } } diff --git a/src/datafusion/src/reader/plantime/row_group_filter.rs b/src/datafusion/src/reader/plantime/row_group_filter.rs deleted file mode 100644 index eb26b34b0..000000000 --- a/src/datafusion/src/reader/plantime/row_group_filter.rs +++ /dev/null @@ -1,427 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use arrow::{array::ArrayRef, array::BooleanArray, array::UInt64Array, datatypes::Schema}; -use datafusion::common::{Column, Result, ScalarValue}; -use datafusion::datasource::listing::FileRange; -use datafusion::datasource::physical_plan::ParquetFileMetrics; -use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; -use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; -use parquet::arrow::arrow_reader::statistics::StatisticsConverter; -use parquet::arrow::parquet_column; -use parquet::basic::Type; -use parquet::data_type::Decimal; -use parquet::schema::types::SchemaDescriptor; -use parquet::{ - arrow::{ParquetRecordBatchStreamBuilder, async_reader::AsyncFileReader}, - bloom_filter::Sbbf, - file::metadata::RowGroupMetaData, -}; -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -/// Reduces the [`ParquetAccessPlan`] based on row group level metadata. -/// -/// This struct implements the various types of pruning that are applied to a -/// set of row groups within a parquet file, progressively narrowing down the -/// set of row groups (and ranges/selections within those row groups) that -/// should be scanned, based on the available metadata. -#[derive(Debug, Clone, PartialEq)] -pub struct RowGroupAccessPlanFilter { - /// which row groups should be accessed - access_plan: ParquetAccessPlan, -} - -impl RowGroupAccessPlanFilter { - /// Create a new `RowGroupPlanBuilder` for pruning out the groups to scan - /// based on metadata and statistics - pub fn new(access_plan: ParquetAccessPlan) -> Self { - Self { access_plan } - } - - /// Return true if there are no row groups - pub fn is_empty(&self) -> bool { - self.access_plan.is_empty() - } - - /// Returns the inner access plan - pub fn build(self) -> ParquetAccessPlan { - self.access_plan - } - - /// Prune remaining row groups to only those within the specified range. - /// - /// Updates this set to mark row groups that should not be scanned - /// - /// # Panics - /// if `groups.len() != self.len()` - pub fn prune_by_range(&mut self, groups: &[RowGroupMetaData], range: &FileRange) { - assert_eq!(groups.len(), self.access_plan.len()); - for (idx, metadata) in groups.iter().enumerate() { - if !self.access_plan.should_scan(idx) { - continue; - } - - // Skip the row group if the first dictionary/data page are not - // within the range. - // - // note don't use the location of metadata - // - let col = metadata.column(0); - let offset = col - .dictionary_page_offset() - .unwrap_or_else(|| col.data_page_offset()); - if !range.contains(offset) { - self.access_plan.skip(idx); - } - } - } - /// Prune remaining row groups using min/max/null_count statistics and - /// the [`PruningPredicate`] to determine if the predicate can not be true. - /// - /// Updates this set to mark row groups that should not be scanned - /// - /// Note: This method currently ignores ColumnOrder - /// - /// - /// # Panics - /// if `groups.len() != self.len()` - pub fn prune_by_statistics( - &mut self, - arrow_schema: &Schema, - parquet_schema: &SchemaDescriptor, - groups: &[RowGroupMetaData], - predicate: &PruningPredicate, - metrics: &ParquetFileMetrics, - ) { - // scoped timer updates on drop - let _timer_guard = metrics.statistics_eval_time.timer(); - - assert_eq!(groups.len(), self.access_plan.len()); - // Indexes of row groups still to scan - let row_group_indexes = self.access_plan.row_group_indexes(); - let row_group_metadatas = row_group_indexes - .iter() - .map(|&i| &groups[i]) - .collect::>(); - - let pruning_stats = RowGroupPruningStatistics { - parquet_schema, - row_group_metadatas, - arrow_schema, - }; - - // try to prune the row groups in a single call - match predicate.prune(&pruning_stats) { - Ok(values) => { - // values[i] is false means the predicate could not be true for row group i - for (idx, &value) in row_group_indexes.iter().zip(values.iter()) { - if !value { - self.access_plan.skip(*idx); - metrics.row_groups_pruned_statistics.add_pruned(1); - } else { - metrics.row_groups_pruned_statistics.add_matched(1); - } - } - } - // stats filter array could not be built, so we can't prune - Err(e) => { - log::debug!("Error evaluating row group predicate values {e}"); - metrics.predicate_evaluation_errors.add(1); - } - } - } - - /// Prune remaining row groups using available bloom filters and the - /// [`PruningPredicate`]. - /// - /// Updates this set with row groups that should not be scanned - /// - /// # Panics - /// if the builder does not have the same number of row groups as this set - pub async fn prune_by_bloom_filters( - &mut self, - arrow_schema: &Schema, - builder: &mut ParquetRecordBatchStreamBuilder, - predicate: &PruningPredicate, - metrics: &ParquetFileMetrics, - ) { - // scoped timer updates on drop - let _timer_guard = metrics.bloom_filter_eval_time.timer(); - - assert_eq!(builder.metadata().num_row_groups(), self.access_plan.len()); - for idx in 0..self.access_plan.len() { - if !self.access_plan.should_scan(idx) { - continue; - } - - // Attempt to find bloom filters for filtering this row group - let literal_columns = predicate.literal_columns(); - let mut column_sbbf = HashMap::with_capacity(literal_columns.len()); - - for column_name in literal_columns { - let Some((column_idx, _field)) = - parquet_column(builder.parquet_schema(), arrow_schema, &column_name) - else { - continue; - }; - - let bf = match builder - .get_row_group_column_bloom_filter(idx, column_idx) - .await - { - Ok(Some(bf)) => bf, - Ok(None) => continue, // no bloom filter for this column - Err(e) => { - log::debug!("Ignoring error reading bloom filter: {e}"); - metrics.predicate_evaluation_errors.add(1); - continue; - } - }; - let physical_type = builder.parquet_schema().column(column_idx).physical_type(); - - column_sbbf.insert(column_name.to_string(), (bf, physical_type)); - } - - let stats = BloomFilterStatistics { column_sbbf }; - - // Can this group be pruned? - let prune_group = match predicate.prune(&stats) { - Ok(values) => !values[0], - Err(e) => { - log::debug!("Error evaluating row group predicate on bloom filter: {e}"); - metrics.predicate_evaluation_errors.add(1); - false - } - }; - - if prune_group { - metrics.row_groups_pruned_bloom_filter.add_pruned(1); - self.access_plan.skip(idx) - } else if !stats.column_sbbf.is_empty() { - metrics.row_groups_pruned_bloom_filter.add_matched(1); - } - } - } -} -/// Implements [`PruningStatistics`] for Parquet Split Block Bloom Filters (SBBF) -struct BloomFilterStatistics { - /// Maps column name to the parquet bloom filter and parquet physical type - column_sbbf: HashMap, -} - -impl BloomFilterStatistics { - /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. - /// - /// In case the type of scalar is not supported, returns `true`, assuming that the - /// value may be present. - fn check_scalar(sbbf: &Sbbf, value: &ScalarValue, parquet_type: &Type) -> bool { - match value { - ScalarValue::Utf8(Some(v)) - | ScalarValue::Utf8View(Some(v)) - | ScalarValue::LargeUtf8(Some(v)) => sbbf.check(&v.as_str()), - ScalarValue::Binary(Some(v)) - | ScalarValue::BinaryView(Some(v)) - | ScalarValue::LargeBinary(Some(v)) => sbbf.check(v), - ScalarValue::FixedSizeBinary(_size, Some(v)) => sbbf.check(v), - ScalarValue::Boolean(Some(v)) => sbbf.check(v), - ScalarValue::Float64(Some(v)) => sbbf.check(v), - ScalarValue::Float32(Some(v)) => sbbf.check(v), - ScalarValue::Int64(Some(v)) => sbbf.check(v), - ScalarValue::Int32(Some(v)) => sbbf.check(v), - ScalarValue::UInt64(Some(v)) => sbbf.check(v), - ScalarValue::UInt32(Some(v)) => sbbf.check(v), - ScalarValue::Decimal128(Some(v), p, s) => match parquet_type { - Type::INT32 => { - //https://github.com/apache/parquet-format/blob/eb4b31c1d64a01088d02a2f9aefc6c17c54cc6fc/Encodings.md?plain=1#L35-L42 - // All physical type are little-endian - if *p > 9 { - //DECIMAL can be used to annotate the following types: - // - // int32: for 1 <= precision <= 9 - // int64: for 1 <= precision <= 18 - return true; - } - let b = (*v as i32).to_le_bytes(); - // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 - let decimal = Decimal::Int32 { - value: b, - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - Type::INT64 => { - if *p > 18 { - return true; - } - let b = (*v as i64).to_le_bytes(); - let decimal = Decimal::Int64 { - value: b, - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - Type::FIXED_LEN_BYTE_ARRAY => { - // keep with from_bytes_to_i128 - let b = v.to_be_bytes().to_vec(); - // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 - let decimal = Decimal::Bytes { - value: b.into(), - precision: *p as i32, - scale: *s as i32, - }; - sbbf.check(&decimal) - } - _ => true, - }, - // One more pattern matching since not all data types are supported - // inside of a Dictionary - ScalarValue::Dictionary(_, inner) => match inner.as_ref() { - ScalarValue::Int32(_) - | ScalarValue::Int64(_) - | ScalarValue::UInt32(_) - | ScalarValue::UInt64(_) - | ScalarValue::Float32(_) - | ScalarValue::Float64(_) - | ScalarValue::Utf8(_) - | ScalarValue::LargeUtf8(_) - | ScalarValue::Binary(_) - | ScalarValue::LargeBinary(_) => { - BloomFilterStatistics::check_scalar(sbbf, inner, parquet_type) - } - _ => true, - }, - _ => true, - } - } -} - -impl PruningStatistics for BloomFilterStatistics { - fn min_values(&self, _column: &Column) -> Option { - None - } - - fn max_values(&self, _column: &Column) -> Option { - None - } - - fn num_containers(&self) -> usize { - 1 - } - - fn null_counts(&self, _column: &Column) -> Option { - None - } - - fn row_counts(&self) -> Option { - None - } - - /// Use bloom filters to determine if we are sure this column can not - /// possibly contain `values` - /// - /// The `contained` API returns false if the bloom filters knows that *ALL* - /// of the values in a column are not present. - fn contained(&self, column: &Column, values: &HashSet) -> Option { - let (sbbf, parquet_type) = self.column_sbbf.get(column.name.as_str())?; - - // Bloom filters are probabilistic data structures that can return false - // positives (i.e. it might return true even if the value is not - // present) however, the bloom filter will return `false` if the value is - // definitely not present. - - let known_not_present = values - .iter() - .map(|value| BloomFilterStatistics::check_scalar(sbbf, value, parquet_type)) - // The row group doesn't contain any of the values if - // all the checks are false - .all(|v| !v); - - let contains = if known_not_present { - Some(false) - } else { - // Given the bloom filter is probabilistic, we can't be sure that - // the row group actually contains the values. Return `None` to - // indicate this uncertainty - None - }; - - Some(BooleanArray::from(vec![contains])) - } -} - -/// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`] -struct RowGroupPruningStatistics<'a> { - parquet_schema: &'a SchemaDescriptor, - row_group_metadatas: Vec<&'a RowGroupMetaData>, - arrow_schema: &'a Schema, -} - -impl<'a> RowGroupPruningStatistics<'a> { - /// Return an iterator over the row group metadata - fn metadata_iter(&'a self) -> impl Iterator + 'a { - self.row_group_metadatas.iter().copied() - } - - fn statistics_converter<'b>(&'a self, column: &'b Column) -> Result> { - Ok(StatisticsConverter::try_new( - &column.name, - self.arrow_schema, - self.parquet_schema, - )?) - } -} - -impl PruningStatistics for RowGroupPruningStatistics<'_> { - fn min_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_mins(self.metadata_iter())?)) - .ok() - } - - fn max_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_maxes(self.metadata_iter())?)) - .ok() - } - - fn num_containers(&self) -> usize { - self.row_group_metadatas.len() - } - - fn null_counts(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_null_counts(self.metadata_iter())?)) - .ok() - .map(|counts| Arc::new(counts) as ArrayRef) - } - - fn row_counts(&self) -> Option { - // Row counts are container-level — read directly from row group metadata. - let counts: UInt64Array = self - .metadata_iter() - .map(|rg| Some(rg.num_rows() as u64)) - .collect(); - Some(Arc::new(counts) as ArrayRef) - } - - fn contained(&self, _column: &Column, _values: &HashSet) -> Option { - None - } -} diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index f95963b03..2ec80e48a 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -1,28 +1,30 @@ -use super::opener::LiquidParquetOpener; +use super::LiquidMorselizer; use crate::cache::{ColumnSqueezeHints, LiquidCacheParquetRef}; use ahash::{HashMap, HashMapExt}; -use arrow_schema::Schema; use bytes::Bytes; use datafusion::{ - common::tree_node::TreeNodeRecursion, - config::TableParquetOptions, + common::{internal_err, tree_node::TreeNodeRecursion}, + config::{ConfigOptions, TableParquetOptions}, datasource::{ listing::PartitionedFile, physical_plan::{ FileScanConfig, FileSource, ParquetFileMetrics, ParquetFileReaderFactory, - ParquetSource, parquet::PagePruningAccessPlanFilter, + ParquetSource, parquet::can_expr_be_pushed_down_with_schemas, }, table_schema::TableSchema, }, error::Result, + execution::object_store::ObjectStoreUrl, physical_expr::projection::ProjectionExprs, + physical_expr::utils::conjunction, physical_expr_adapter::DefaultPhysicalExprAdapterFactory, - physical_optimizer::pruning::{PruningPredicate, PruningPredicateBuilder}, physical_plan::{ - PhysicalExpr, apply_expression_roots, - metrics::{ExecutionPlanMetricsSet, MetricBuilder}, + DisplayFormatType, PhysicalExpr, apply_expression_roots, + filter_pushdown::{FilterPushdownPropagation, PushedDown, PushedDownPredicate}, + metrics::ExecutionPlanMetricsSet, }, }; +use datafusion_datasource::morsel::Morselizer; use futures::{FutureExt, future::BoxFuture}; use object_store::{ObjectStore, ObjectStoreExt, path::Path}; use parquet::{ @@ -31,6 +33,7 @@ use parquet::{ file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}, }; use std::{ + fmt::{self, Formatter}, ops::Range, sync::{Arc, LazyLock}, }; @@ -41,11 +44,12 @@ static META_CACHE: LazyLock = LazyLock::new(MetadataCache::new); #[derive(Debug)] pub(crate) struct CachedMetaReaderFactory { store: Arc, + store_url: ObjectStoreUrl, } impl CachedMetaReaderFactory { - pub(crate) fn new(store: Arc) -> Self { - Self { store } + pub(crate) fn new(store: Arc, store_url: ObjectStoreUrl) -> Self { + Self { store, store_url } } pub(crate) fn create_liquid_reader( @@ -60,6 +64,7 @@ impl CachedMetaReaderFactory { ParquetMetadataCacheReader { file_metrics: ParquetFileMetrics::new(partition_index, path.as_ref(), metrics), store: Arc::clone(&self.store), + store_url: self.store_url.clone(), file_size: partitioned_file.object_meta.size, metadata_size_hint, path, @@ -86,7 +91,7 @@ impl ParquetFileReaderFactory for CachedMetaReaderFactory { } struct MetadataCache { - val: RwLock>>, + val: RwLock>>, } impl MetadataCache { @@ -101,6 +106,7 @@ impl MetadataCache { pub struct ParquetMetadataCacheReader { file_metrics: ParquetFileMetrics, store: Arc, + store_url: ObjectStoreUrl, file_size: u64, metadata_size_hint: Option, path: Path, @@ -143,20 +149,20 @@ impl AsyncFileReader for ParquetMetadataCacheReader { &mut self, options: Option<&ArrowReaderOptions>, ) -> BoxFuture<'_, parquet::errors::Result>> { - let path = self.path.clone(); + let cache_key = (self.store_url.clone(), self.path.clone()); let options = options.cloned(); async move { // First check with read lock { let cache = META_CACHE.val.read().await; - if let Some(meta) = cache.get(&path) { + if let Some(meta) = cache.get(&cache_key) { return Ok(meta.clone()); } } // Upgrade to write lock and double-check let mut cache = META_CACHE.val.write().await; - match cache.entry(path.clone()) { + match cache.entry(cache_key) { std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()), std::collections::hash_map::Entry::Vacant(entry) => { let file_size = self.file_size; @@ -183,8 +189,6 @@ impl AsyncFileReader for ParquetMetadataCacheReader { pub struct LiquidParquetSource { metrics: ExecutionPlanMetricsSet, predicate: Option>, - pruning_predicate: Option>, - page_pruning_predicate: Option>, table_parquet_options: TableParquetOptions, liquid_cache: LiquidCacheParquetRef, batch_size: Option, @@ -229,40 +233,9 @@ impl LiquidParquetSource { &self.squeeze_hints } - /// Set predicate information, also sets pruning_predicate and page_pruning_predicate attributes - pub fn with_predicate( - mut self, - file_schema: Arc, - predicate: Arc, - ) -> Self { - let metrics = ExecutionPlanMetricsSet::new(); - let predicate_creation_errors = - MetricBuilder::new(&metrics).global_counter("num_predicate_creation_errors"); - - self.metrics = metrics; - self.predicate = Some(Arc::clone(&predicate)); - - match PruningPredicateBuilder::new() - .with_file_schema(Arc::clone(&file_schema)) - .try_build(Arc::clone(&predicate)) - { - Ok(pruning_predicate) => { - if !pruning_predicate.always_true() { - self.pruning_predicate = Some(Arc::new(pruning_predicate)); - } - } - Err(e) => { - log::debug!("Could not create pruning predicate for: {e}"); - predicate_creation_errors.add(1); - } - }; - - let page_pruning_predicate = Arc::new(PagePruningAccessPlanFilter::new( - &predicate, - Arc::clone(&file_schema), - )); - self.page_pruning_predicate = Some(page_pruning_predicate); - + /// Set predicate information. + pub fn with_predicate(mut self, predicate: Arc) -> Self { + self.predicate = Some(predicate); self } @@ -271,7 +244,6 @@ impl LiquidParquetSource { let predicate = source.filter(); let table_schema = source.table_schema().clone(); - let file_schema = table_schema.file_schema().clone(); let projection = source.projection().cloned().unwrap_or_else(|| { let table_schema = table_schema.table_schema(); ProjectionExprs::from_indices( @@ -287,14 +259,12 @@ impl LiquidParquetSource { projection, metrics: source.metrics().clone(), predicate: None, - pruning_predicate: None, - page_pruning_predicate: None, span: None, squeeze_hints: Arc::default(), }; if let Some(predicate) = predicate { - v = v.with_predicate(file_schema, predicate); + v = v.with_predicate(predicate); } v @@ -308,40 +278,52 @@ impl LiquidParquetSource { impl FileSource for LiquidParquetSource { fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!( + "LiquidParquetSource::create_file_opener called but it supports the Morsel API, please use that instead" + ) + } + + fn create_morselizer( &self, object_store: Arc, base_config: &FileScanConfig, partition: usize, - ) -> Result> { + ) -> Result> { let expr_adapter_factory = base_config .expr_adapter_factory .clone() .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); - let reader_factory = Arc::new(CachedMetaReaderFactory::new(object_store)); + let reader_factory = Arc::new(CachedMetaReaderFactory::new( + object_store, + base_config.object_store_url.clone(), + )); let execution_span = self .span .clone() .map(|span| fastrace::Span::enter_with_parent(format!("opener_{partition}"), &span)); - let opener = LiquidParquetOpener::new( - partition, - self.projection.clone(), - self.batch_size - .expect("Batch size must be set before creating LiquidParquetOpener"), - base_config.limit, - self.predicate.clone(), - self.table_schema.clone(), - self.metrics.clone(), - self.liquid_cache.clone(), - reader_factory, - self.reorder_filters(), + Ok(Box::new(LiquidMorselizer { + partition_index: partition, + projection: self.projection.clone(), + batch_size: self + .batch_size + .expect("Batch size must be set before creating LiquidMorselizer"), + predicate: self.predicate.clone(), + table_schema: self.table_schema.clone(), + metrics: self.metrics.clone(), + liquid_cache: self.liquid_cache.clone(), + parquet_file_reader_factory: reader_factory, + reorder_filters: self.reorder_filters(), expr_adapter_factory, - execution_span.map(Arc::new), - Arc::clone(&self.squeeze_hints), - ); - - Ok(Arc::new(opener)) + span: execution_span.map(Arc::new), + squeeze_hints: Arc::clone(&self.squeeze_hints), + })) } fn with_batch_size(&self, batch_size: usize) -> Arc { @@ -350,6 +332,10 @@ impl FileSource for LiquidParquetSource { Arc::new(conf) } + fn filter(&self) -> Option> { + self.predicate.clone() + } + fn table_schema(&self) -> &TableSchema { &self.table_schema } @@ -375,6 +361,58 @@ impl FileSource for LiquidParquetSource { "liquid_parquet" } + fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + if let Some(predicate) = self.filter() { + write!(f, ", predicate={predicate}")?; + } + Ok(()) + } + DisplayFormatType::TreeRender => Ok(()), + } + } + + fn try_pushdown_filters( + &self, + filters: Vec>, + _config: &ConfigOptions, + ) -> Result>> { + let filters: Vec<_> = filters + .into_iter() + .map(|filter| { + if can_expr_be_pushed_down_with_schemas(&filter, self.table_schema.file_schema()) { + PushedDownPredicate::supported(filter) + } else { + PushedDownPredicate::unsupported(filter) + } + }) + .collect(); + + if filters + .iter() + .all(|filter| matches!(filter.discriminant, PushedDown::No)) + { + return Ok(FilterPushdownPropagation::with_parent_pushdown_result( + vec![PushedDown::No; filters.len()], + )); + } + + let supported = filters + .iter() + .filter_map(|filter| match filter.discriminant { + PushedDown::Yes => Some(Arc::clone(&filter.predicate)), + PushedDown::No => None, + }); + let predicate = conjunction(self.predicate.iter().cloned().chain(supported)); + let source = Arc::new(self.clone().with_predicate(predicate)); + + Ok(FilterPushdownPropagation::with_parent_pushdown_result( + filters.iter().map(|filter| filter.discriminant).collect(), + ) + .with_updated_node(source)) + } + fn apply_expressions( &self, f: &mut dyn FnMut(&Arc) -> Result, diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 4a5a6b27b..40c42a2b4 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -92,6 +92,11 @@ struct ParquetFallback { impl LiquidCacheReader { pub(crate) fn new(config: LiquidCacheReaderConfig) -> Self { + debug_assert_eq!( + config.batch_size, + config.cached_row_group.batch_size(), + "DataFusion and LiquidCache batch sizes must agree" + ); let inner = LiquidCacheReaderInner::new( config.batch_size, config.selection, @@ -105,14 +110,6 @@ impl LiquidCacheReader { row_filter: config.row_filter, } } - - pub(crate) fn into_filter(self) -> Option { - debug_assert!( - matches!(self.state, ReaderState::Finished), - "cannot extract filter before reader completes" - ); - self.row_filter - } } impl Stream for LiquidCacheReader { @@ -591,12 +588,11 @@ mod tests { std::fs::metadata(&parquet_path).unwrap().len(), ); let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( - 0, - partitioned_file, - None, - &metrics, - ); + let input = CachedMetaReaderFactory::new( + object_store, + datafusion::execution::object_store::ObjectStoreUrl::parse("test-runtime:///").unwrap(), + ) + .create_liquid_reader(0, partitioned_file, None, &metrics); let projection = ProjectionMask::roots( reader_metadata.metadata().file_metadata().schema_descr(), [0], @@ -755,30 +751,6 @@ mod tests { assert_eq!(batch.num_rows(), 2); } - #[tokio::test] - async fn into_filter_returns_stored_filter_after_completion() { - let batch_size = 2; - let test = make_row_group(batch_size, &[vec![1, 2]]).await; - let selection = RowSelection::from(Vec::::new()); - let filter = LiquidRowFilter::new(Vec::new()); - - let mut reader = test.reader(ReaderRequest { - selection, - row_filter: Some(filter), - projection_columns: vec![0], - schema: Arc::clone(&test.schema), - }); - - let waker = futures::task::noop_waker(); - let mut cx = Context::from_waker(&waker); - assert!(matches!( - Pin::new(&mut reader).poll_next(&mut cx), - Poll::Ready(None) - )); - - assert!(reader.into_filter().is_some()); - } - #[tokio::test] async fn predicate_filters_rows_across_batches() { let batches = vec![vec![1, 2], vec![3, 4]]; diff --git a/src/datafusion/src/reader/runtime/liquid_stream.rs b/src/datafusion/src/reader/runtime/liquid_stream.rs deleted file mode 100644 index aa0b7c678..000000000 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ /dev/null @@ -1,902 +0,0 @@ -use crate::cache::{CachedFileRef, CachedRowGroupRef}; -use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; -use arrow::array::RecordBatch; -use arrow_schema::{Schema, SchemaRef}; -use fastrace::Event; -use fastrace::local::LocalSpan; -use futures::Stream; -use parquet::{ - arrow::{ - ProjectionMask, - arrow_reader::{ArrowPredicate, RowSelection, RowSelector}, - }, - errors::ParquetError, - file::metadata::ParquetMetaData, -}; -use std::{ - collections::VecDeque, - fmt::Formatter, - pin::Pin, - sync::Arc, - task::{Context, Poll}, -}; - -use super::liquid_cache_reader::{ - LiquidCacheReader, LiquidCacheReaderConfig, ParquetFallbackConfig, -}; -use super::utils::{get_root_column_ids, limit_row_selection, offset_row_selection}; - -type PlanResult = Option; - -struct ReaderFactory { - metadata: Arc, - - input: ParquetMetadataCacheReader, - - filter: Option, - - limit: Option, - - offset: Option, - - cached_file: CachedFileRef, -} - -impl ReaderFactory { - /// Plans what to read from cache vs parquet for the next row group - fn plan_row_group( - &mut self, - row_group_idx: usize, - selection: Option, - projection: ProjectionMask, - batch_size: usize, - ) -> PlanResult { - let meta = self.metadata.row_group(row_group_idx); - - let mut predicate_projection: Option = None; - if let Some(filter) = self.filter.as_mut() { - for predicate in filter.predicates_mut() { - let p_projection = predicate.projection(); - if let Some(ref mut p) = predicate_projection { - p.union(p_projection); - } else { - predicate_projection = Some(p_projection.clone()); - } - } - } - - let mut selection = - selection.unwrap_or_else(|| vec![RowSelector::select(meta.num_rows() as usize)].into()); - - let rows_before = selection.row_count(); - - if rows_before == 0 { - return None; - } - - if let Some(offset) = self.offset { - selection = offset_row_selection(selection, offset); - } - - if let Some(limit) = self.limit { - selection = limit_row_selection(selection, limit); - } - - let rows_after = selection.row_count(); - - // Update offset if necessary - if let Some(offset) = &mut self.offset { - // Reduction is either because of offset or limit, as limit is applied - // after offset has been "exhausted" can just use saturating sub here - *offset = offset.saturating_sub(rows_before - rows_after) - } - - if rows_after == 0 { - return None; - } - - if let Some(limit) = &mut self.limit { - *limit -= rows_after; - } - - let mut cache_projection = projection.clone(); - if let Some(ref predicate_projection) = predicate_projection { - cache_projection.union(predicate_projection); - } - - let schema_descr = self.metadata.file_metadata().schema_descr(); - let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); - let predicate_column_ids = if let Some(ref predicate_projection) = predicate_projection { - get_root_column_ids(schema_descr, predicate_projection) - } else { - Vec::new() - }; - let cached_row_group = self - .cached_file - .create_row_group(row_group_idx as u64, predicate_column_ids); - - let projection_column_ids = get_root_column_ids(schema_descr, &projection); - - let context = PlanningContext { - row_group_idx, - selection, - batch_size, - cached_row_group, - cache_projection, - projection_column_ids, - cache_column_ids, - }; - - Some(context) - } -} - -fn build_projection_schema(file_schema: &SchemaRef, projection_column_ids: &[usize]) -> SchemaRef { - let fields: Vec<_> = projection_column_ids - .iter() - .filter_map(|column_id| file_schema.fields().get(*column_id)) - .map(|field_ref| field_ref.as_ref().clone()) - .collect(); - Arc::new(Schema::new(fields)) -} - -/// Context for planning what to read from cache vs parquet -struct PlanningContext { - row_group_idx: usize, - selection: RowSelection, - batch_size: usize, - cached_row_group: CachedRowGroupRef, - cache_projection: ProjectionMask, - projection_column_ids: Vec, - cache_column_ids: Vec, -} - -fn build_liquid_cache_reader( - reader_factory: &mut ReaderFactory, - context: PlanningContext, - schema: SchemaRef, -) -> LiquidCacheReader { - let row_count = reader_factory - .metadata - .row_group(context.row_group_idx) - .num_rows() as usize; - let cache_batch_size = context.cached_row_group.batch_size(); - LiquidCacheReader::new(LiquidCacheReaderConfig { - batch_size: context.batch_size, - selection: context.selection, - row_filter: reader_factory.filter.take(), - cached_row_group: context.cached_row_group, - projection_columns: context.projection_column_ids, - schema, - parquet_fallback: ParquetFallbackConfig { - row_group_idx: context.row_group_idx, - metadata: Arc::clone(&reader_factory.metadata), - input: reader_factory.input.clone(), - cache_projection: context.cache_projection, - cache_column_ids: context.cache_column_ids, - cache_batch_size, - row_count, - }, - }) -} - -enum StreamState { - /// At the start of a new row group, or the end of the parquet stream - Init, - /// Decoding a batch from cache - ReadFromCache(Box), -} - -impl std::fmt::Debug for StreamState { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self { - StreamState::Init => write!(f, "StreamState::Init"), - StreamState::ReadFromCache(_) => write!(f, "StreamState::Decoding"), - } - } -} - -pub struct LiquidStreamBuilder { - pub(crate) input: ParquetMetadataCacheReader, - - pub(crate) metadata: Arc, - - pub(crate) batch_size: usize, - - pub(crate) row_groups: Option>, - - pub(crate) projection: ProjectionMask, - - pub(crate) filter: Option, - - pub(crate) selection: Option, - - pub(crate) limit: Option, - - pub(crate) offset: Option, - - pub(crate) span: Option, -} - -impl LiquidStreamBuilder { - pub fn new(input: ParquetMetadataCacheReader, metadata: Arc) -> Self { - Self { - input, - metadata, - batch_size: 1024, - row_groups: None, - projection: ProjectionMask::all(), - filter: None, - selection: None, - limit: None, - offset: None, - span: None, - } - } - - pub fn with_batch_size(mut self, batch_size: usize) -> Self { - self.batch_size = batch_size; - self - } - - pub fn with_row_groups(mut self, row_groups: Vec) -> Self { - self.row_groups = Some(row_groups); - self - } - - pub fn with_projection(mut self, projection: ProjectionMask) -> Self { - self.projection = projection; - self - } - - pub fn with_selection(mut self, selection: Option) -> Self { - self.selection = selection; - self - } - - pub fn with_limit(mut self, limit: Option) -> Self { - self.limit = limit; - self - } - - pub fn with_row_filter(mut self, filter: LiquidRowFilter) -> Self { - self.filter = Some(filter); - self - } - - pub fn with_span(mut self, span: fastrace::Span) -> Self { - self.span = Some(span); - self - } - - pub fn build(self, liquid_cache: CachedFileRef) -> Result { - let num_row_groups = self.metadata.row_groups().len(); - - let row_groups: VecDeque = match self.row_groups { - Some(row_groups) => { - if let Some(col) = row_groups.iter().find(|x| **x >= num_row_groups) { - return Err(ParquetError::ArrowError(format!( - "row group {col} out of bounds 0..{num_row_groups}" - ))); - } - row_groups.into() - } - None => (0..self.metadata.row_groups().len()).collect(), - }; - - let batch_size = self - .batch_size - .min(self.metadata.file_metadata().num_rows() as usize); - - let schema_descr = self.metadata.file_metadata().schema_descr(); - let projection_column_ids = get_root_column_ids(schema_descr, &self.projection); - let file_schema = liquid_cache.schema(); - let schema = build_projection_schema(&file_schema, &projection_column_ids); - - // `plan_row_group` applies limit/offset by truncating the row - // selection BEFORE the row filter runs, so combining them with a - // filter caps the rows *scanned* rather than the rows *matched* — - // silently dropping matches that sit past the first `limit + offset` - // physical rows (upstream parquet counts the limit against - // post-filter matches instead). Until limit accounting moves after - // predicate evaluation, only honor limit/offset for unfiltered - // scans, where scanned rows == emitted rows and truncation is - // exact. Filtered scans still get capped post-filter by - // DataFusion's FileStream, which slices emitted batches against - // `FileScanConfig::limit`; all that is lost is scan-internal early - // termination. - let (limit, offset) = if self.filter.is_some() { - (None, None) - } else { - (self.limit, self.offset) - }; - - let reader = ReaderFactory { - metadata: Arc::clone(&self.metadata), - input: self.input, - filter: self.filter, - limit, - offset, - cached_file: liquid_cache, - }; - - Ok(LiquidStream { - metadata: self.metadata, - schema, - row_groups, - projection: self.projection, - batch_size, - selection: self.selection, - reader: Some(reader), - state: StreamState::Init, - span: self.span, - }) - } -} - -pub struct LiquidStream { - metadata: Arc, - - schema: SchemaRef, - - row_groups: VecDeque, - - projection: ProjectionMask, - - batch_size: usize, - - selection: Option, - - /// This is an option so it can be moved into a future - reader: Option, - - state: StreamState, - - span: Option, -} - -impl std::fmt::Debug for LiquidStream { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ParquetRecordBatchStream") - .field("metadata", &self.metadata) - .field("schema", &self.schema) - .field("batch_size", &self.batch_size) - .field("projection", &self.projection) - .field("state", &self.state) - .finish() - } -} - -impl LiquidStream { - pub fn schema(&self) -> &SchemaRef { - &self.schema - } -} - -impl Stream for LiquidStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let _guard = self.span.as_ref().map(|s| s.set_local_parent()); - loop { - let state = std::mem::replace(&mut self.state, StreamState::Init); - - match state { - StreamState::ReadFromCache(mut batch_reader) => { - match Pin::new(&mut *batch_reader).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => { - self.state = StreamState::ReadFromCache(batch_reader); - return Poll::Ready(Some(Ok(batch))); - } - Poll::Ready(Some(Err(e))) => { - panic!("Decoding next batch error: {e:?}"); - } - Poll::Ready(None) => { - let batch_reader = *batch_reader; - let filter = batch_reader.into_filter(); - self.reader.as_mut().unwrap().filter = filter; - // state left as Init, continue loop to plan next row group - } - Poll::Pending => { - self.state = StreamState::ReadFromCache(batch_reader); - return Poll::Pending; - } - } - } - StreamState::Init => { - let row_group_idx = match self.row_groups.pop_front() { - Some(idx) => idx, - None => return Poll::Ready(None), - }; - - let row_count = self.metadata.row_group(row_group_idx).num_rows() as usize; - - let selection = self.selection.as_mut().map(|s| s.split_off(row_count)); - - LocalSpan::add_event(Event::new("LiquidStream::plan_row_group")); - let projection = self.projection.clone(); - let batch_size = self.batch_size; - let maybe_context = self.reader.as_mut().expect("lost reader").plan_row_group( - row_group_idx, - selection, - projection, - batch_size, - ); - match maybe_context { - Some(context) => { - LocalSpan::add_event(Event::new("LiquidStream::read_from_cache")); - let schema = Arc::clone(&self.schema); - let reader_factory = self.reader.as_mut().unwrap(); - let batch_reader = - build_liquid_cache_reader(reader_factory, context, schema); - self.state = StreamState::ReadFromCache(Box::new(batch_reader)); - } - None => { - self.state = StreamState::Init; - } - } - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::cache::{BatchID, CachedFileRef, LiquidCacheParquet}; - use crate::reader::plantime::{ - CachedMetaReaderFactory, FilterCandidateBuilder, LiquidPredicate, - }; - use arrow::array::{Array, ArrayRef, Int32Array}; - use arrow_schema::{DataType, Field, Schema}; - use datafusion::common::ScalarValue; - use datafusion::datasource::listing::PartitionedFile; - use datafusion::logical_expr::Operator; - use datafusion::physical_expr::PhysicalExpr; - use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; - use futures::StreamExt; - use liquid_cache::cache::AlwaysHydrate; - use liquid_cache::cache::squeeze_policies::Evict; - use liquid_cache::cache_policies::LiquidPolicy; - use object_store::local::LocalFileSystem; - use parquet::arrow::ArrowWriter; - use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; - use std::fs::File; - use std::sync::Arc; - - fn write_two_row_group_file(path: &std::path::Path, schema: SchemaRef) { - let file = File::create(path).unwrap(); - let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap(); - let batch0 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![0, 1, 2, 3])), - Arc::new(Int32Array::from(vec![10, 11, 12, 13])), - ], - ) - .unwrap(); - let batch1 = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![4, 5, 6, 7])), - Arc::new(Int32Array::from(vec![14, 15, 16, 17])), - ], - ) - .unwrap(); - writer.write(&batch0).unwrap(); - writer.flush().unwrap(); - writer.write(&batch1).unwrap(); - writer.close().unwrap(); - } - - fn write_single_row_group_file(path: &std::path::Path, schema: SchemaRef, a: Vec) { - let file = File::create(path).unwrap(); - let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap(); - let b: Vec<_> = a.iter().map(|value| value + 1000).collect(); - let batch = RecordBatch::try_new( - schema, - vec![Arc::new(Int32Array::from(a)), Arc::new(Int32Array::from(b))], - ) - .unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); - } - - async fn make_liquid_stream( - max_memory_bytes: usize, - max_disk_bytes: usize, - row_filter: Option, - ) -> ( - LiquidStream, - Arc, - CachedFileRef, - tempfile::TempDir, - ) { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let tmp_dir = tempfile::tempdir().unwrap(); - let parquet_path = tmp_dir.path().join("data.parquet"); - write_two_row_group_file(&parquet_path, schema.clone()); - let metadata_file = File::open(&parquet_path).unwrap(); - let reader_metadata = - ArrowReaderMetadata::load(&metadata_file, ArrowReaderOptions::new()).unwrap(); - let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); - let partitioned_file = PartitionedFile::new( - "data.parquet", - std::fs::metadata(&parquet_path).unwrap().len(), - ); - let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( - 0, - partitioned_file, - None, - &metrics, - ); - - let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(); - let cache = Arc::new( - LiquidCacheParquet::new( - 4, - max_memory_bytes, - max_disk_bytes, - store, - Box::new(LiquidPolicy::new()), - Box::new(Evict), - Box::new(AlwaysHydrate::new()), - ) - .await, - ); - let cached_file = cache.register_or_get_file("data.parquet".to_string(), schema); - let projection = ProjectionMask::roots( - reader_metadata.metadata().file_metadata().schema_descr(), - [0, 1], - ); - let mut builder = LiquidStreamBuilder::new(input, Arc::clone(reader_metadata.metadata())) - .with_batch_size(4) - .with_row_groups(vec![0, 1]) - .with_projection(projection); - if let Some(row_filter) = row_filter { - builder = builder.with_row_filter(row_filter); - } - let stream = builder.build(cached_file.clone()).unwrap(); - (stream, cache, cached_file, tmp_dir) - } - - async fn collect_liquid_values(stream: LiquidStream) -> (Vec, Vec) { - let batches = stream - .map(|batch| batch.expect("valid liquid stream batch")) - .collect::>() - .await; - let mut a = Vec::new(); - let mut b = Vec::new(); - for batch in batches { - let a_array = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let b_array = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - a.extend(a_array.iter().map(|value| value.unwrap())); - b.extend(b_array.iter().map(|value| value.unwrap())); - } - (a, b) - } - - async fn collect_projected_a(stream: LiquidStream) -> Vec { - let batches = stream - .map(|batch| batch.expect("valid liquid stream batch")) - .collect::>() - .await; - let mut a = Vec::new(); - for batch in batches { - let a_array = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - a.extend(a_array.iter().map(|value| value.unwrap())); - } - a - } - - fn gt_filter(schema: SchemaRef, literal: i32) -> LiquidRowFilter { - gt_filter_on(schema, "a", 0, literal) - } - - fn gt_filter_on( - schema: SchemaRef, - col_name: &str, - col_idx: usize, - literal: i32, - ) -> LiquidRowFilter { - let expr: Arc = Arc::new(BinaryExpr::new( - Arc::new(Column::new(col_name, col_idx)), - Operator::Gt, - Arc::new(Literal::new(ScalarValue::Int32(Some(literal)))), - )); - let tmp_meta = tempfile::NamedTempFile::new().unwrap(); - write_two_row_group_file(tmp_meta.path(), schema.clone()); - let file = File::open(tmp_meta.path()).unwrap(); - let metadata = ArrowReaderMetadata::load(&file, ArrowReaderOptions::new()).unwrap(); - let builder = FilterCandidateBuilder::new(expr, schema); - let candidate = builder.build(metadata.metadata()).unwrap().unwrap(); - let projection = candidate.projection(metadata.metadata()); - let predicate = LiquidPredicate::try_new(candidate, projection).unwrap(); - LiquidRowFilter::new(vec![predicate]) - } - - async fn make_liquid_stream_with_projection( - max_memory_bytes: usize, - max_disk_bytes: usize, - row_filter: Option, - projection_columns: Vec, - ) -> ( - LiquidStream, - Arc, - CachedFileRef, - tempfile::TempDir, - ) { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let tmp_dir = tempfile::tempdir().unwrap(); - let parquet_path = tmp_dir.path().join("data.parquet"); - write_two_row_group_file(&parquet_path, schema.clone()); - let metadata_file = File::open(&parquet_path).unwrap(); - let reader_metadata = - ArrowReaderMetadata::load(&metadata_file, ArrowReaderOptions::new()).unwrap(); - let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); - let partitioned_file = PartitionedFile::new( - "data.parquet", - std::fs::metadata(&parquet_path).unwrap().len(), - ); - let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( - 0, - partitioned_file, - None, - &metrics, - ); - - let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(); - let cache = Arc::new( - LiquidCacheParquet::new( - 4, - max_memory_bytes, - max_disk_bytes, - store, - Box::new(LiquidPolicy::new()), - Box::new(Evict), - Box::new(AlwaysHydrate::new()), - ) - .await, - ); - let cached_file = cache.register_or_get_file("data.parquet".to_string(), schema); - let projection = ProjectionMask::roots( - reader_metadata.metadata().file_metadata().schema_descr(), - projection_columns, - ); - let mut builder = LiquidStreamBuilder::new(input, Arc::clone(reader_metadata.metadata())) - .with_batch_size(4) - .with_row_groups(vec![0, 1]) - .with_projection(projection); - if let Some(row_filter) = row_filter { - builder = builder.with_row_filter(row_filter); - } - let stream = builder.build(cached_file.clone()).unwrap(); - (stream, cache, cached_file, tmp_dir) - } - - async fn make_single_row_group_stream( - parquet_a: Vec, - projection_columns: Vec, - ) -> (LiquidStream, CachedFileRef, tempfile::TempDir) { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let tmp_dir = tempfile::tempdir().unwrap(); - let parquet_path = tmp_dir.path().join("data.parquet"); - write_single_row_group_file(&parquet_path, schema.clone(), parquet_a); - let metadata_file = File::open(&parquet_path).unwrap(); - let reader_metadata = - ArrowReaderMetadata::load(&metadata_file, ArrowReaderOptions::new()).unwrap(); - let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); - let partitioned_file = PartitionedFile::new( - "data.parquet", - std::fs::metadata(&parquet_path).unwrap().len(), - ); - let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( - 0, - partitioned_file, - None, - &metrics, - ); - - let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(); - let cache = Arc::new( - LiquidCacheParquet::new( - 4, - usize::MAX, - usize::MAX, - store, - Box::new(LiquidPolicy::new()), - Box::new(Evict), - Box::new(AlwaysHydrate::new()), - ) - .await, - ); - let cached_file = cache.register_or_get_file("data.parquet".to_string(), schema); - let projection = ProjectionMask::roots( - reader_metadata.metadata().file_metadata().schema_descr(), - projection_columns, - ); - let stream = LiquidStreamBuilder::new(input, Arc::clone(reader_metadata.metadata())) - .with_batch_size(4) - .with_row_groups(vec![0]) - .with_projection(projection) - .build(cached_file.clone()) - .unwrap(); - (stream, cached_file, tmp_dir) - } - - async fn insert_batches( - row_group: &CachedRowGroupRef, - column_id: usize, - batch_payloads: &[(u16, &[i32])], - ) { - let column = row_group.get_column(column_id as u64).unwrap(); - for (batch_idx, values) in batch_payloads.iter() { - let array: ArrayRef = Arc::new(Int32Array::from(values.to_vec())); - column - .insert(BatchID::from_raw(*batch_idx), array) - .await - .unwrap(); - } - } - - async fn is_cached(row_group: &CachedRowGroupRef, column_id: usize, batch_idx: u16) -> bool { - row_group - .get_column(column_id as u64) - .unwrap() - .get_arrow_array_test_only(BatchID::from_raw(batch_idx)) - .await - .is_some() - } - - #[tokio::test] - async fn cache_full_keeps_inserted_batches_and_skips_failed_inserts() { - let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); - let (stream, _cache, cached_file, _tmp_dir) = - make_liquid_stream(one_array_memory * 3, 0, None).await; - - let (a, b) = collect_liquid_values(stream).await; - - assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); - assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); - - let row_group0 = cached_file.create_row_group(0, vec![]); - let row_group1 = cached_file.create_row_group(1, vec![]); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); - } - - #[tokio::test] - async fn cache_full_with_row_filter_keeps_lookaside_results_correct() { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); - let filter = gt_filter(schema, 2); - let (stream, _cache, cached_file, _tmp_dir) = - make_liquid_stream(one_array_memory * 3, 0, Some(filter)).await; - - let (a, b) = collect_liquid_values(stream).await; - - assert_eq!(a, vec![3, 4, 5, 6, 7]); - assert_eq!(b, vec![13, 14, 15, 16, 17]); - - let row_group0 = cached_file.create_row_group(0, vec![]); - let row_group1 = cached_file.create_row_group(1, vec![]); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); - } - - #[tokio::test] - async fn mid_scan_eviction_recovers() { - let (stream, _cache, cached_file, _tmp_dir) = make_liquid_stream(0, 0, None).await; - - let (a, b) = collect_liquid_values(stream).await; - - assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); - assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); - - let row_group0 = cached_file.create_row_group(0, vec![]); - let row_group1 = cached_file.create_row_group(1, vec![]); - assert!(!is_cached(&row_group0, 0, 0).await); - assert!(!is_cached(&row_group0, 1, 0).await); - assert!(!is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); - } - - #[tokio::test] - async fn predicate_fallback_uses_predicate_projection() { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - ])); - let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); - let filter = gt_filter_on(schema, "b", 1, 13); - let (stream, _cache, cached_file, _tmp_dir) = - make_liquid_stream_with_projection(one_array_memory * 3, 0, Some(filter), vec![0]) - .await; - - let a_values = collect_projected_a(stream).await; - - assert_eq!(a_values, vec![4, 5, 6, 7]); - - let row_group0 = cached_file.create_row_group(0, vec![]); - let row_group1 = cached_file.create_row_group(1, vec![]); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); - } - - #[tokio::test] - async fn missing_column_falls_back_to_parquet() { - let (stream, _cache, cached_file, _tmp_dir) = - make_liquid_stream(usize::MAX, usize::MAX, None).await; - let row_group0 = cached_file.create_row_group(0, vec![]); - let row_group1 = cached_file.create_row_group(1, vec![]); - insert_batches(&row_group0, 0, &[(0, &[0, 1, 2, 3])]).await; - insert_batches(&row_group1, 0, &[(0, &[4, 5, 6, 7])]).await; - - let (a, b) = collect_liquid_values(stream).await; - - assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); - assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 1, 0).await); - } - - #[tokio::test] - async fn fallback_stream_advances_across_misses() { - let parquet_a = vec![ - 100, 101, 102, 103, 4, 5, 6, 7, 200, 201, 202, 203, 12, 13, 14, 15, - ]; - let (stream, cached_file, _tmp_dir) = - make_single_row_group_stream(parquet_a, vec![0]).await; - let row_group = cached_file.create_row_group(0, vec![]); - insert_batches(&row_group, 0, &[(0, &[0, 1, 2, 3]), (2, &[8, 9, 10, 11])]).await; - - let a_values = collect_projected_a(stream).await; - - assert_eq!(a_values, (0..16).collect::>()); - assert!(is_cached(&row_group, 0, 0).await); - assert!(is_cached(&row_group, 0, 1).await); - assert!(is_cached(&row_group, 0, 2).await); - assert!(is_cached(&row_group, 0, 3).await); - } -} diff --git a/src/datafusion/src/reader/runtime/mod.rs b/src/datafusion/src/reader/runtime/mod.rs index bfde28dcc..68be7a68c 100644 --- a/src/datafusion/src/reader/runtime/mod.rs +++ b/src/datafusion/src/reader/runtime/mod.rs @@ -1,7 +1,8 @@ pub(crate) use liquid_predicate::extract_multi_column_or; -pub(crate) use liquid_stream::LiquidStreamBuilder; +pub(crate) use morsel::{LiquidRowGroupPlanner, build_projection_schema}; +pub(crate) use utils::get_root_column_ids; mod liquid_cache_reader; mod liquid_predicate; -mod liquid_stream; +mod morsel; mod utils; diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs new file mode 100644 index 000000000..ee3f611ec --- /dev/null +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -0,0 +1,171 @@ +use std::{fmt, pin::Pin, sync::Arc}; + +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow_schema::{Schema, SchemaRef}; +use datafusion::{error::DataFusionError, physical_expr::projection::Projector}; +use datafusion_datasource::morsel::Morsel; +use futures::{Stream, StreamExt, stream::BoxStream}; +use parquet::{ + arrow::{ + ProjectionMask, + arrow_reader::{ArrowPredicate, RowSelection, RowSelector}, + }, + file::metadata::ParquetMetaData, +}; + +use crate::{ + cache::CachedFileRef, + reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}, +}; + +use super::{ + liquid_cache_reader::{LiquidCacheReader, LiquidCacheReaderConfig, ParquetFallbackConfig}, + utils::get_root_column_ids, +}; + +pub(crate) struct LiquidRowGroupPlanner { + pub(crate) metadata: Arc, + pub(crate) input: ParquetMetadataCacheReader, + pub(crate) row_filter: Option, + pub(crate) cached_file: CachedFileRef, + pub(crate) projection: ProjectionMask, + pub(crate) batch_size: usize, + pub(crate) stream_schema: SchemaRef, + pub(crate) output_schema: SchemaRef, + pub(crate) projector: Arc, + pub(crate) replace_schema: bool, + pub(crate) span: Option>, +} + +impl LiquidRowGroupPlanner { + pub(crate) fn plan( + &self, + row_group_idx: usize, + selection: Option, + ) -> Option { + let metadata = self.metadata.row_group(row_group_idx); + let mut predicate_projection: Option = None; + if let Some(filter) = &self.row_filter { + for predicate in filter.predicates() { + let projection = predicate.projection(); + if let Some(predicate_projection) = &mut predicate_projection { + predicate_projection.union(projection); + } else { + predicate_projection = Some(projection.clone()); + } + } + } + + let selection = selection + .unwrap_or_else(|| vec![RowSelector::select(metadata.num_rows() as usize)].into()); + if selection.row_count() == 0 { + return None; + } + + let mut cache_projection = self.projection.clone(); + if let Some(predicate_projection) = &predicate_projection { + cache_projection.union(predicate_projection); + } + + let schema_descr = self.metadata.file_metadata().schema_descr(); + let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); + let predicate_column_ids = predicate_projection + .as_ref() + .map(|projection| get_root_column_ids(schema_descr, projection)) + .unwrap_or_default(); + let projection_columns = get_root_column_ids(schema_descr, &self.projection); + let cached_row_group = self + .cached_file + .create_row_group(row_group_idx as u64, predicate_column_ids); + let cache_batch_size = cached_row_group.batch_size(); + + Some(LiquidRowGroupMorsel { + config: LiquidCacheReaderConfig { + batch_size: self.batch_size, + selection, + row_filter: self.row_filter.clone(), + cached_row_group, + projection_columns, + schema: Arc::clone(&self.stream_schema), + parquet_fallback: ParquetFallbackConfig { + row_group_idx, + metadata: Arc::clone(&self.metadata), + input: self.input.clone(), + cache_projection, + cache_column_ids, + cache_batch_size, + row_count: metadata.num_rows() as usize, + }, + }, + output_schema: Arc::clone(&self.output_schema), + projector: Arc::clone(&self.projector), + replace_schema: self.replace_schema, + span: self.span.clone(), + }) + } +} + +pub(crate) fn build_projection_schema( + file_schema: &SchemaRef, + projection_column_ids: &[usize], +) -> SchemaRef { + let fields = projection_column_ids + .iter() + .filter_map(|column_id| file_schema.fields().get(*column_id)) + .map(|field| field.as_ref().clone()) + .collect::>(); + Arc::new(Schema::new(fields)) +} + +pub(crate) struct LiquidRowGroupMorsel { + config: LiquidCacheReaderConfig, + output_schema: SchemaRef, + projector: Arc, + replace_schema: bool, + span: Option>, +} + +impl fmt::Debug for LiquidRowGroupMorsel { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LiquidRowGroupMorsel") + .finish_non_exhaustive() + } +} + +impl Morsel for LiquidRowGroupMorsel { + fn into_stream(self: Box) -> BoxStream<'static, datafusion::error::Result> { + let Self { + config, + output_schema, + projector, + replace_schema, + span, + } = *self; + let mut reader = LiquidCacheReader::new(config); + let stream = futures::stream::poll_fn(move |cx| { + let _guard = span.as_ref().map(|span| span.set_local_parent()); + Pin::new(&mut reader).poll_next(cx) + }); + + stream + .map(|batch| batch.map_err(|error| DataFusionError::External(Box::new(error)))) + .map(move |batch| { + batch.and_then(|batch| { + let batch = projector.project_batch(&batch)?; + if replace_schema { + let (_schema, arrays, num_rows) = batch.into_parts(); + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + arrays, + &options, + ) + .map_err(Into::into) + } else { + Ok(batch) + } + }) + }) + .boxed() + } +} diff --git a/src/datafusion/src/reader/runtime/utils.rs b/src/datafusion/src/reader/runtime/utils.rs index a05b79942..629ffcfd7 100644 --- a/src/datafusion/src/reader/runtime/utils.rs +++ b/src/datafusion/src/reader/runtime/utils.rs @@ -1,10 +1,7 @@ use std::collections::VecDeque; use parquet::{ - arrow::{ - ProjectionMask, - arrow_reader::{RowSelection, RowSelector}, - }, + arrow::{ProjectionMask, arrow_reader::RowSelector}, schema::types::SchemaDescriptor, }; @@ -28,64 +25,6 @@ pub(crate) fn get_root_column_ids( .collect() } -pub(crate) fn offset_row_selection(selection: RowSelection, offset: usize) -> RowSelection { - if offset == 0 { - return selection; - } - - let mut selected_count = 0; - let mut skipped_count = 0; - - let mut selectors: Vec = selection.into(); - - let find = selectors.iter().position(|selector| match selector.skip { - true => { - skipped_count += selector.row_count; - false - } - false => { - selected_count += selector.row_count; - selected_count > offset - } - }); - - let split_idx = match find { - Some(idx) => idx, - None => { - selectors.clear(); - return RowSelection::from(selectors); - } - }; - - let mut new_selectors = Vec::with_capacity(selectors.len() - split_idx + 1); - new_selectors.push(RowSelector::skip(skipped_count + offset)); - new_selectors.push(RowSelector::select(selected_count - offset)); - new_selectors.extend_from_slice(&selectors[split_idx + 1..]); - - RowSelection::from(new_selectors) -} - -pub(crate) fn limit_row_selection(selection: RowSelection, mut limit: usize) -> RowSelection { - let mut selectors: Vec = selection.into(); - - if limit == 0 { - selectors.clear(); - } - - for (idx, selection) in selectors.iter_mut().enumerate() { - if !selection.skip { - if selection.row_count >= limit { - selection.row_count = limit; - selectors.truncate(idx + 1); - break; - } else { - limit -= selection.row_count; - } - } - } - RowSelection::from(selectors) -} - /// Take the next batch from the selection queue. /// The returning selection will have exactly the batch size, or less if the selection is exhausted. pub(super) fn take_next_batch( From 7554acd17842b0c9f4c5d47d9b4cbaf6e2c9fd4b Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Thu, 3 Sep 2026 01:15:45 -0400 Subject: [PATCH 06/24] implement morsel reader (#514) Doesn't have performance impact yet, just make it nicer to look at. --- src/core/README.md | 2 +- src/core/src/cache/builders.rs | 1 - src/core/src/cache/core.rs | 191 ++++-- src/core/src/cache/index.rs | 16 +- src/core/src/cache/mod.rs | 2 +- src/core/study/fsst_selectivity.rs | 2 +- src/datafusion-local/src/lib.rs | 10 +- src/datafusion-local/src/tests/mod.rs | 48 +- ...afusion_local__tests__provide_schema2.snap | 3 +- src/datafusion-local/src/tests/squeeze.rs | 5 + src/datafusion/src/cache/column.rs | 82 ++- src/datafusion/src/cache/mod.rs | 228 ++++++- src/datafusion/src/optimizers/mod.rs | 23 +- src/datafusion/src/reader/plantime/mod.rs | 2 +- .../src/reader/plantime/morselizer.rs | 613 ++++++++++++++++-- src/datafusion/src/reader/plantime/source.rs | 9 + .../src/reader/runtime/liquid_cache_reader.rs | 161 +++-- .../src/reader/runtime/liquid_predicate.rs | 20 +- src/datafusion/src/reader/runtime/mod.rs | 7 +- src/datafusion/src/reader/runtime/morsel.rs | 140 +++- src/datafusion/src/reader/runtime/utils.rs | 2 +- 21 files changed, 1292 insertions(+), 275 deletions(-) diff --git a/src/core/README.md b/src/core/README.md index d263aab51..1338bc2e6 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -27,7 +27,7 @@ let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); // Insert once; replacement/placement is handled by the cache policy storage.insert(entry_id, arrow_array.clone()).await; -assert!(storage.is_cached(&entry_id)); +assert!(storage.contains(&entry_id)); }); ``` diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 18db151f3..6881f0c46 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -268,7 +268,6 @@ impl<'a> Get<'a> { /// Materialize the cached array as [`ArrayRef`]. pub async fn read(self) -> Option { - self.storage.observer().on_get(self.selection.is_some()); self.storage .read_arrow_array( self.entry_id, diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index f26f2877f..4fc02f9f7 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -62,11 +62,21 @@ pub struct LiquidCache { squeeze_victims_concurrently: bool, } +/// Outcome of [`LiquidCache::prefetch`]. +pub enum PrefetchResult { + /// A memory-form snapshot of the entry (Arrow or Liquid), ready to hand to a reader. + Snapshot(Arc), + /// The entry is squeezed; prefetch leaves it alone. + Squeezed, + /// The entry is not in the index, or its disk blob is gone. + Absent, +} + /// Builder returned by [`LiquidCache::insert`] for configuring cache writes. impl LiquidCache { /// Return current cache statistics: counts and resource usage. pub fn stats(&self) -> CacheStats { - // Count entries by residency/format + // Count entries by storage tier and format let total_entries = self.index.entry_count(); let mut memory_arrow_entries = 0usize; @@ -141,6 +151,35 @@ impl LiquidCache { EvaluatePredicate::new(self, entry_id, predicate) } + /// Prefetch an entry into a memory-form snapshot without recording an access. + pub async fn prefetch(&self, entry_id: &EntryID) -> PrefetchResult { + let Some(entry) = self.index.get(entry_id) else { + return PrefetchResult::Absent; + }; + match entry.as_ref() { + CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => { + PrefetchResult::Snapshot(entry) + } + disk @ CacheEntry::DiskArrow { .. } => { + let Some(array) = self.read_disk_arrow_array(entry_id).await else { + return PrefetchResult::Absent; + }; + self.maybe_hydrate(entry_id, disk, MaterializedEntry::Arrow(&array), None) + .await; + PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_arrow(array))) + } + disk @ CacheEntry::DiskLiquid { .. } => { + let Some(array) = self.read_disk_liquid_array(entry_id).await else { + return PrefetchResult::Absent; + }; + self.maybe_hydrate(entry_id, disk, MaterializedEntry::Liquid(&array), None) + .await; + PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_liquid(array))) + } + CacheEntry::MemorySqueezedLiquid(_) => PrefetchResult::Squeezed, + } + } + /// Try to read a liquid array from the cache. /// Returns None if the cached data is not in liquid format. pub async fn try_read_liquid( @@ -156,14 +195,14 @@ impl LiquidCache { match batch.as_ref() { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id).await?; self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) .await; Some(liquid) } CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { SqueezedBacking::Liquid(_) => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id).await?; Some(liquid) } SqueezedBacking::Arrow(_) => None, @@ -185,9 +224,9 @@ impl LiquidCache { self.budget.reset_usage(); } - /// Check if a batch is cached. - pub fn is_cached(&self, entry_id: &EntryID) -> bool { - self.index.is_cached(entry_id) + /// Check whether the cache contains a batch. + pub fn contains(&self, entry_id: &EntryID) -> bool { + self.index.contains(entry_id) } /// Get the config of the cache. @@ -598,19 +637,44 @@ impl LiquidCache { selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { - use arrow::array::BooleanArray; - + self.observer.on_get(selection.is_some()); let batch = self.index.get(entry_id)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + self.read_entry_inner(entry_id, batch.as_ref(), selection, expression) + .await + } + + /// Read an already-looked-up cache entry. + pub async fn read_entry( + &self, + entry_id: &EntryID, + entry: &CacheEntry, + selection: Option<&BooleanBuffer>, + expression: Option<&CacheExpression>, + ) -> Option { + self.observer.on_get(selection.is_some()); + self.read_entry_inner(entry_id, entry, selection, expression) + .await + } + + async fn read_entry_inner( + &self, + entry_id: &EntryID, + entry: &CacheEntry, + selection: Option<&BooleanBuffer>, + expression: Option<&CacheExpression>, + ) -> Option { + use arrow::array::BooleanArray; + self.trace(InternalEvent::Read { entry: *entry_id, selection: selection.is_some(), expr: expression.cloned(), - cached: CachedBatchType::from(batch.as_ref()), + cached: CachedBatchType::from(entry), }); - match batch.as_ref() { + match entry { CacheEntry::MemoryArrow(array) => match selection { Some(selection) => { let selection_array = BooleanArray::new(selection.clone(), None); @@ -623,7 +687,7 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(batch.as_ref(), entry_id, expression, selection) + self.read_disk_array(entry, entry_id, expression, selection) .await } CacheEntry::MemorySqueezedLiquid(array) => { @@ -647,7 +711,7 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await; + let full_array = self.read_disk_arrow_array(entry_id).await?; self.maybe_hydrate( entry_id, entry, @@ -669,7 +733,7 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id).await?; self.maybe_hydrate( entry_id, entry, @@ -762,7 +826,7 @@ impl LiquidCache { let full_array = if !all_paths_present { let batch = CacheEntry::MemorySqueezedLiquid(array.clone()); self.observer.on_get_squeezed_needs_io(); - let full_array = self.read_disk_arrow_array(entry_id).await; + let full_array = self.read_disk_arrow_array(entry_id).await?; self.maybe_hydrate( entry_id, &batch, @@ -818,12 +882,12 @@ impl LiquidCache { Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> ArrayRef { - let bytes = self - .store - .get(&entry_id_to_key(entry_id)) - .await - .expect("read failed"); + async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> Option { + let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { + Ok(bytes) => bytes, + Err(t4::Error::NotFound) => return None, + Err(error) => panic!("read failed: {error}"), + }; let bytes_len = bytes.len(); let cursor = std::io::Cursor::new(bytes); let mut reader = @@ -834,18 +898,18 @@ impl LiquidCache { entry: *entry_id, bytes: bytes_len, }); - array + Some(array) } async fn read_disk_liquid_array( &self, entry_id: &EntryID, - ) -> crate::liquid_array::LiquidArrayRef { - let bytes = self - .store - .get(&entry_id_to_key(entry_id)) - .await - .expect("read failed"); + ) -> Option { + let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { + Ok(bytes) => bytes, + Err(t4::Error::NotFound) => return None, + Err(error) => panic!("read failed: {error}"), + }; self.trace(InternalEvent::IoReadLiquid { entry: *entry_id, bytes: bytes.len(), @@ -853,10 +917,12 @@ impl LiquidCache { let compressor_states = self.metadata.get_compressor(entry_id); let compressor = compressor_states.fsst_compressor(); - (crate::liquid_array::ipc::read_from_bytes( - Bytes::from(bytes), - &crate::liquid_array::ipc::LiquidIPCContext::new(compressor), - )) as _ + Some( + (crate::liquid_array::ipc::read_from_bytes( + Bytes::from(bytes), + &crate::liquid_array::ipc::LiquidIPCContext::new(compressor), + )) as _, + ) } pub(crate) async fn eval_predicate_internal( @@ -865,19 +931,41 @@ impl LiquidCache { selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { - use arrow::array::BooleanArray; - self.observer.on_eval_predicate(); let batch = self.index.get(entry_id)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + self.eval_predicate_on_entry_inner(entry_id, batch.as_ref(), selection_opt, predicate) + .await + } + + /// Evaluate a predicate on an already-looked-up cache entry. + pub async fn eval_predicate_on_entry( + &self, + entry_id: &EntryID, + entry: &CacheEntry, + selection_opt: Option<&BooleanBuffer>, + predicate: &LiquidExpr, + ) -> Option { + self.observer.on_eval_predicate(); + self.eval_predicate_on_entry_inner(entry_id, entry, selection_opt, predicate) + .await + } + + async fn eval_predicate_on_entry_inner( + &self, + entry_id: &EntryID, + entry: &CacheEntry, + selection_opt: Option<&BooleanBuffer>, + predicate: &LiquidExpr, + ) -> Option { self.trace(InternalEvent::EvalPredicate { entry: *entry_id, selection: selection_opt.is_some(), - cached: CachedBatchType::from(batch.as_ref()), + cached: CachedBatchType::from(entry), }); - match batch.as_ref() { + match entry { CacheEntry::MemoryArrow(array) => { let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { @@ -890,7 +978,7 @@ impl LiquidCache { Some(self.eval_predicate_on_array(filtered, predicate)) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await; + let array = self.read_disk_arrow_array(entry_id).await?; self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) .await; let mut owned = None; @@ -912,7 +1000,7 @@ impl LiquidCache { Some(array.try_eval_predicate(predicate, selection)) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id).await?; self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) .await; let mut owned = None; @@ -1229,6 +1317,23 @@ mod tests { } } + #[tokio::test] + async fn missing_disk_blob_is_a_cache_miss() { + let directory = tempfile::tempdir().unwrap(); + let store = t4::mount(directory.path().join("cache.t4")).await.unwrap(); + let cache = LiquidCacheBuilder::new() + .with_store(store.clone()) + .build() + .await; + let id = EntryID::from(320usize); + + cache.insert(id, create_test_arrow_array(8)).await.unwrap(); + cache.flush_all_to_disk().await.unwrap(); + store.remove(&entry_id_to_key(&id)).await.unwrap(); + + assert!(cache.get(&id).await.is_none()); + } + #[tokio::test] async fn hydrate_disk_liquid_on_get_promotes_to_memory_liquid() { let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; @@ -1268,7 +1373,7 @@ mod tests { let err = cache.insert(EntryID::from(900usize), array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.is_cached(&EntryID::from(900usize))); + assert!(!cache.contains(&EntryID::from(900usize))); } #[tokio::test] @@ -1289,12 +1394,12 @@ mod tests { let second = EntryID::from(911usize); cache.insert(first, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.is_cached(&first)); + assert!(cache.contains(&first)); cache.insert(second, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first)); + assert!(!cache.contains(&first)); assert!(matches!( cache.index().get(&second).unwrap().as_ref(), CacheEntry::DiskArrow { .. } @@ -1321,7 +1426,7 @@ mod tests { cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&first) || !cache.is_cached(&second)); + assert!(!cache.contains(&first) || !cache.contains(&second)); } #[tokio::test] @@ -1343,7 +1448,7 @@ mod tests { cache.remove_disk_entry(entry).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.is_cached(&entry)); + assert!(!cache.contains(&entry)); } #[tokio::test] @@ -1361,6 +1466,6 @@ mod tests { let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.is_cached(&entry_id)); + assert!(!cache.contains(&entry_id)); } } diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index a25fec751..62acca365 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -35,7 +35,7 @@ impl ArtIndex { Some(batch) } - pub(crate) fn is_cached(&self, entry_id: &EntryID) -> bool { + pub(crate) fn contains(&self, entry_id: &EntryID) -> bool { let guard = self.art.pin(); self.art.get(*entry_id, &guard).is_some() } @@ -99,15 +99,15 @@ mod tests { use super::*; #[test] - fn test_get_and_is_cached() { + fn test_get_and_contains() { let store = ArtIndex::new(); let entry_id1: EntryID = EntryID::from(1); let entry_id2: EntryID = EntryID::from(2); let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(!store.contains(&entry_id1)); + assert!(!store.contains(&entry_id2)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached @@ -115,8 +115,8 @@ mod tests { store.insert(&entry_id1, array1.clone()); } - assert!(store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(store.contains(&entry_id1)); + assert!(!store.contains(&entry_id2)); // Get should return the cached value match store.get(&entry_id1) { @@ -137,10 +137,10 @@ mod tests { store.insert(&entry_id, array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.is_cached(&entry_id)); + assert!(store.contains(&entry_id)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.is_cached(&entry_id)); + assert!(!store.contains(&entry_id)); } } diff --git a/src/core/src/cache/mod.rs b/src/core/src/cache/mod.rs index 4f52abe56..19ae5dc13 100644 --- a/src/core/src/cache/mod.rs +++ b/src/core/src/cache/mod.rs @@ -15,7 +15,7 @@ mod utils; pub use builders::{EvaluatePredicate, Get, Insert, LiquidCacheBuilder, default_max_memory_bytes}; pub use cached_batch::{CacheEntry, CachedBatchType}; -pub use core::LiquidCache; +pub use core::{LiquidCache, PrefetchResult}; pub use expressions::{CacheExpression, VariantRequest}; #[cfg(test)] pub(crate) use io_context::TestSqueezeIo; diff --git a/src/core/study/fsst_selectivity.rs b/src/core/study/fsst_selectivity.rs index ed349858e..9e0ae16e4 100644 --- a/src/core/study/fsst_selectivity.rs +++ b/src/core/study/fsst_selectivity.rs @@ -99,7 +99,7 @@ async fn main() { continue; } - // Warm up once to reduce cold-start noise. + // Run once to reduce cold-start noise. std::hint::black_box(fsst.to_uncompressed_selected(&selection.indices)); let mut total = 0.0; diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index dfdcb91fa..32e2f0d58 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,6 +69,7 @@ pub struct LiquidCacheLocalBuilder { squeeze_policy: Box, /// Hydration policy hydration_policy: Box, + prefetch: bool, span: fastrace::Span, } @@ -84,6 +85,7 @@ impl Default for LiquidCacheLocalBuilder { cache_policy: Box::new(LiquidPolicy::new()), squeeze_policy: Box::new(TranscodeSqueezeEvict), hydration_policy: Box::new(AlwaysHydrate::new()), + prefetch: true, span: fastrace::Span::enter_with_local_parent("liquid_cache_datafusion_local_builder"), } } @@ -139,6 +141,12 @@ impl LiquidCacheLocalBuilder { self } + /// Enable or disable row-group prefetching. + pub fn with_prefetch(mut self, prefetch: bool) -> Self { + self.prefetch = prefetch; + self + } + /// Set fastrace span pub fn with_span(mut self, span: fastrace::Span) -> Self { self.span = span; @@ -190,7 +198,7 @@ impl LiquidCacheLocalBuilder { .await; let cache_ref = Arc::new(cache); - let optimizer = LocalModeOptimizer::new(cache_ref.clone()); + let optimizer = LocalModeOptimizer::new(cache_ref.clone()).with_prefetch(self.prefetch); let state = datafusion::execution::SessionStateBuilder::new() .with_config(config) diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 998052ce1..366669cf1 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -108,6 +108,7 @@ async fn create_session_context_with_liquid_cache( let mut config = SessionConfig::new().with_repartition_file_scans(false); config.options_mut().execution.target_partitions = 4; let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(cache_size_bytes) .with_cache_dir(cache_dir.to_path_buf()) .with_squeeze_policy(squeeze_policy) @@ -135,6 +136,48 @@ async fn get_result(ctx: &SessionContext, sql: &str) -> String { pretty_format_batches(&batches).unwrap().to_string() } +async fn run_io_profile(prefetch: bool, cache_dir: &Path) -> (String, u64, u64, u64) { + let config = SessionConfig::new().with_repartition_file_scans(false); + let builder = LiquidCacheLocalBuilder::new() + .with_max_memory_bytes(64 * 1024 * 1024) + .with_cache_dir(cache_dir.to_path_buf()); + let builder = if prefetch { + builder + } else { + builder.with_prefetch(false) + }; + let (ctx, cache) = builder.build(config).await.unwrap(); + ctx.register_parquet("hits", TEST_FILE, ParquetReadOptions::default()) + .await + .unwrap(); + let sql = r#"SELECT "WatchID" FROM hits WHERE "SearchPhrase" LIKE '%abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789%'"#; + + let first = get_result(&ctx, sql).await; + cache.flush_data().await.unwrap(); + cache.storage().stats(); + let second = get_result(&ctx, sql).await; + let runtime = cache.storage().stats().runtime; + assert_eq!(first, second); + + ( + second, + runtime.read_io_count, + runtime.get, + runtime.eval_predicate, + ) +} + +#[tokio::test] +async fn prefetch_matches_lazy_io() { + let lazy_dir = TempDir::new().unwrap(); + let prefetch_dir = TempDir::new().unwrap(); + + let lazy = run_io_profile(false, lazy_dir.path()).await; + let prefetch = run_io_profile(true, prefetch_dir.path()).await; + + assert_eq!(lazy, prefetch); +} + async fn run_sql_with_cache( sql: &str, squeeze_policy: Box, @@ -150,7 +193,7 @@ async fn run_sql_with_cache( let displayable = DisplayableExecutionPlan::new(plan.as_ref()); let plan_string = format!("{}", displayable.tree_render()); - // Clear any historical runtime counters before warming the cache. + // Clear any historical runtime counters before prefetching the cache. cache.storage().stats(); let first_run = get_result(&ctx, sql).await; @@ -346,6 +389,7 @@ async fn test_provide_schema2() { let mut config = SessionConfig::new(); config.options_mut().execution.target_partitions = 4; let (liquid_ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_cache_dir(cache_dir.path().to_path_buf()) .with_max_memory_bytes(1024 * 1024) .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) @@ -394,7 +438,7 @@ async fn test_provide_schema2() { let displayable = DisplayableExecutionPlan::new(plan.as_ref()); let plan_string = format!("{}", displayable.tree_render()); - // Reset runtime counters so we measure hits from the warm run onwards. + // Reset runtime counters so we measure hits from the prefetch run onwards. cache.storage().stats(); let first_liquid_run = liquid_ctx.sql(sql).await.unwrap().collect().await.unwrap(); diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap index 0a535fb1a..c127d3fe1 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap @@ -1,5 +1,6 @@ --- source: src/datafusion-local/src/tests/mod.rs +assertion_line: 428 expression: snapshot --- query[0]: SELECT * from default where log like '%hhj%' order by _timestamp @@ -190,7 +191,7 @@ RuntimeStatsSnapshot: eval_predicate: 0 get_squeezed_success: 0 get_squeezed_needs_io: 0 - try_read_liquid_calls: 0 + try_read_liquid_calls: 3 hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 diff --git a/src/datafusion-local/src/tests/squeeze.rs b/src/datafusion-local/src/tests/squeeze.rs index 7a50e3497..f2c40dfc6 100644 --- a/src/datafusion-local/src/tests/squeeze.rs +++ b/src/datafusion-local/src/tests/squeeze.rs @@ -14,6 +14,7 @@ fn squeeze_test_config() -> SessionConfig { async fn basic_squeeze() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(1024 * 128) .with_cache_dir(cache_dir.path().to_path_buf()) .build(squeeze_test_config()) @@ -41,6 +42,7 @@ async fn basic_squeeze() { async fn squeeze_strings() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(1024 * 1024) .with_cache_dir(cache_dir.path().to_path_buf()) .build(squeeze_test_config()) @@ -68,6 +70,7 @@ async fn squeeze_strings() { async fn squeeze_substrings_search() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(1024 * 256) .with_cache_dir(cache_dir.path().to_path_buf()) .build(squeeze_test_config()) @@ -92,6 +95,7 @@ async fn squeeze_substrings_search() { async fn squeeze_substrings_search_title() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(1024 * 1024 * 4) .with_cache_dir(cache_dir.path().to_path_buf()) .build(squeeze_test_config()) @@ -117,6 +121,7 @@ async fn squeeze_substrings_search_title() { async fn squeeze_distinct_search_phase() { let cache_dir = TempDir::new().unwrap(); let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_prefetch(false) .with_max_memory_bytes(1024 * 256) .with_cache_dir(cache_dir.path().to_path_buf()) .build(squeeze_test_config()) diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 1f582fd86..fc02ccebc 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -5,12 +5,14 @@ use arrow::{ record_batch::RecordBatch, }; use arrow_schema::{ArrowError, Field, Schema}; -use liquid_cache::cache::{CacheExpression, CacheFull, LiquidCache, LiquidExpr}; +use liquid_cache::cache::{ + CacheEntry, CacheExpression, CacheFull, LiquidCache, LiquidExpr, PrefetchResult, +}; use parquet::arrow::arrow_reader::ArrowPredicate; use crate::{ LiquidPredicate, - cache::{BatchID, ColumnAccessPath, ParquetArrayID}, + cache::{BatchID, ColumnAccessPath, ParquetArrayID, RowGroupSnapshots}, }; use std::sync::Arc; @@ -21,6 +23,7 @@ pub struct CachedColumn { field: Arc, column_path: ColumnAccessPath, expression: Option>, + snapshots: Arc, } /// A reference to a cached column. @@ -35,6 +38,13 @@ pub enum InsertArrowArrayError { CacheFull, } +pub(crate) enum PrefetchOutcome { + Snapshotted, + AlreadySnapshotted, + Squeezed, + Missing, +} + impl From for InsertArrowArrayError { fn from(_: CacheFull) -> Self { Self::CacheFull @@ -48,6 +58,7 @@ impl CachedColumn { column_access_path: ColumnAccessPath, expression: Option>, is_predicate_column: bool, + snapshots: Arc, ) -> Self { // Register the column's squeeze hint. Squeeze hints are column-scoped; // `ParquetCacheMetadata` keys them by column (the batch id is masked @@ -70,6 +81,7 @@ impl CachedColumn { cache_store, column_path: column_access_path, expression, + snapshots, } } @@ -78,8 +90,12 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } - pub(crate) fn is_cached(&self, batch_id: BatchID) -> bool { - self.cache_store.is_cached(&self.entry_id(batch_id).into()) + pub(crate) fn contains(&self, batch_id: BatchID) -> bool { + self.cache_store.contains(&self.entry_id(batch_id).into()) + } + + pub(crate) fn snapshot_entry(&self, batch_id: BatchID) -> Option> { + self.snapshots.get(&self.entry_id(batch_id).into()) } /// Returns the Arrow field metadata for this cached column. @@ -112,11 +128,24 @@ impl CachedColumn { ); if let Some(liquid_expr) = liquid_expr - && let Some(boolean_array) = self - .cache_store - .eval_predicate(&entry_id, &liquid_expr) - .with_selection(filter) - .await + && let Some(boolean_array) = match self.snapshots.get(&entry_id) { + Some(entry) => { + self.cache_store + .eval_predicate_on_entry( + &entry_id, + entry.as_ref(), + Some(filter), + &liquid_expr, + ) + .await + } + None => { + self.cache_store + .eval_predicate(&entry_id, &liquid_expr) + .with_selection(filter) + .await + } + } { let predicate_filter = match boolean_array.null_count() { 0 => boolean_array, @@ -159,6 +188,17 @@ impl CachedColumn { filter: &BooleanBuffer, ) -> Option { let entry_id = self.entry_id(batch_id).into(); + if let Some(entry) = self.snapshots.get(&entry_id) { + return self + .cache_store + .read_entry( + &entry_id, + entry.as_ref(), + Some(filter), + self.expression.as_deref(), + ) + .await; + } self.cache_store .get(&entry_id) .with_selection(filter) @@ -179,7 +219,7 @@ impl CachedColumn { batch_id: BatchID, array: ArrayRef, ) -> Result<(), InsertArrowArrayError> { - if self.is_cached(batch_id) { + if self.contains(batch_id) { return Err(InsertArrowArrayError::AlreadyCached); } @@ -188,4 +228,26 @@ impl CachedColumn { .await?; Ok(()) } + + pub(crate) fn insert_snapshot(&self, batch_id: BatchID, array: ArrayRef) { + self.snapshots.insert( + self.entry_id(batch_id).into(), + Arc::new(CacheEntry::memory_arrow(array)), + ); + } + + pub(crate) async fn prefetch_snapshot(&self, batch_id: BatchID) -> PrefetchOutcome { + let entry_id = self.entry_id(batch_id).into(); + if self.snapshots.get(&entry_id).is_some() { + return PrefetchOutcome::AlreadySnapshotted; + } + match self.cache_store.prefetch(&entry_id).await { + PrefetchResult::Snapshot(entry) => { + self.snapshots.insert(entry_id, entry); + PrefetchOutcome::Snapshotted + } + PrefetchResult::Squeezed => PrefetchOutcome::Squeezed, + PrefetchResult::Absent => PrefetchOutcome::Missing, + } + } } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 7c2fda5b0..1984ac782 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -3,14 +3,18 @@ use crate::io::ParquetCacheMetadata; use crate::reader::{LiquidPredicate, extract_multi_column_or}; -use crate::sync::Mutex; +use crate::sync::{Mutex, RwLock}; use ahash::AHashMap; use arrow::array::{BooleanArray, RecordBatch}; use arrow::buffer::BooleanBuffer; use arrow_schema::{ArrowError, Field, Schema, SchemaRef}; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::Column; use liquid_cache::cache::squeeze_policies::SqueezePolicy; use liquid_cache::cache::{ - CacheExpression, CachePolicy, EventTrace, HydrationPolicy, LiquidCache, LiquidCacheBuilder, + CacheEntry, CacheExpression, CachePolicy, EntryID, EventTrace, HydrationPolicy, LiquidCache, + LiquidCacheBuilder, }; use parquet::arrow::arrow_reader::ArrowPredicate; use std::collections::HashMap; @@ -22,8 +26,8 @@ mod column; mod id; mod stats; -pub(crate) use column::InsertArrowArrayError; pub use column::{CachedColumn, CachedColumnRef}; +pub(crate) use column::{InsertArrowArrayError, PrefetchOutcome}; pub(crate) use id::ColumnAccessPath; pub use id::{BatchID, ParquetArrayID}; @@ -37,6 +41,30 @@ pub type ColumnSqueezeHints = HashMap>; /// One column of a row group: (file column index, field, squeeze hint, is-predicate). type CachedColumnSpec = (u64, Arc, Option>, bool); +#[derive(Default, Debug)] +pub(crate) struct RowGroupSnapshots { + entries: RwLock>>, + selections: RwLock>, +} + +impl RowGroupSnapshots { + pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { + self.entries.read().unwrap().get(entry_id).cloned() + } + + pub(crate) fn insert(&self, entry_id: EntryID, entry: Arc) { + self.entries.write().unwrap().insert(entry_id, entry); + } + + pub(crate) fn selection(&self, batch_id: BatchID) -> Option { + self.selections.read().unwrap().get(&batch_id).cloned() + } + + pub(crate) fn insert_selection(&self, batch_id: BatchID, selection: BooleanBuffer) { + self.selections.write().unwrap().insert(batch_id, selection); + } +} + #[derive(Default, Debug)] struct ColumnMaps { // invariant: Arc::ptr_eq(map[field.name()], map[field.id()]) @@ -49,6 +77,7 @@ struct ColumnMaps { pub struct CachedRowGroup { columns: ColumnMaps, cache_store: Arc, + snapshots: Arc, } impl CachedRowGroup { @@ -60,6 +89,7 @@ impl CachedRowGroup { row_group_idx: u64, file_idx: u64, columns: &[CachedColumnSpec], + snapshots: Arc, ) -> Self { let mut column_maps = ColumnMaps::default(); for (column_id, field, expression, is_predicate_column) in columns { @@ -70,6 +100,7 @@ impl CachedRowGroup { column_access_path, expression.clone(), *is_predicate_column, + Arc::clone(&snapshots), )); column_maps.by_id.insert(*column_id, column.clone()); column_maps.by_name.insert(field.name().to_string(), column); @@ -78,6 +109,7 @@ impl CachedRowGroup { Self { columns: column_maps, cache_store, + snapshots, } } @@ -103,6 +135,10 @@ impl CachedRowGroup { self.columns.by_name.get(unqualified).cloned() } + pub(crate) fn snapshot_selection(&self, batch_id: BatchID) -> Option { + self.snapshots.selection(batch_id) + } + /// Evaluate a predicate on a row group. #[fastrace::trace] pub async fn evaluate_selection_with_predicate( @@ -129,7 +165,28 @@ impl CachedRowGroup { for (col_name, expr) in column_exprs { let column = self.get_column_by_name(col_name)?; - let liquid_expr = column.liquid_expr_for_predicate(Arc::clone(&expr)); + let snapshot_liquid = match column.snapshot_entry(batch_id) { + Some(entry) => match entry.as_ref() { + CacheEntry::MemoryLiquid(array) => Some(Arc::clone(array)), + _ => { + combined_buffer = None; + break; + } + }, + None => None, + }; + let expr = expr + .transform_up(|expr| { + if let Some(column) = expr.downcast_ref::() { + Ok(Transformed::yes(Arc::new(Column::new(column.name(), 0)) + as Arc)) + } else { + Ok(Transformed::no(expr)) + } + }) + .ok()? + .data; + let liquid_expr = column.liquid_expr_for_predicate(expr); let liquid_expr = match liquid_expr { Some(expr) => expr, None => { @@ -138,7 +195,10 @@ impl CachedRowGroup { } }; let entry_id = column.entry_id(batch_id).into(); - let liquid_array = self.cache_store.try_read_liquid(&entry_id).await; + let liquid_array = match snapshot_liquid { + Some(array) => Some(array), + None => self.cache_store.try_read_liquid(&entry_id).await, + }; let liquid_array = match liquid_array { None => { combined_buffer = None; @@ -210,6 +270,15 @@ impl CachedFile { &self, row_group_id: u64, predicate_column_ids: Vec, + ) -> CachedRowGroupRef { + self.create_row_group_with_snapshots(row_group_id, predicate_column_ids, Arc::default()) + } + + pub(crate) fn create_row_group_with_snapshots( + &self, + row_group_id: u64, + predicate_column_ids: Vec, + snapshots: Arc, ) -> CachedRowGroupRef { let columns: Vec = self .file_schema @@ -233,6 +302,7 @@ impl CachedFile { row_group_id, self.file_id, &columns, + snapshots, )) } @@ -432,7 +502,7 @@ mod tests { use super::*; use crate::cache::{CachedRowGroupRef, LiquidCacheParquet}; use crate::reader::FilterCandidateBuilder; - use arrow::array::Int32Array; + use arrow::array::{Array, ArrayRef, Int32Array, StringViewArray}; use arrow::buffer::BooleanBuffer; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -467,6 +537,152 @@ mod tests { file.create_row_group(0, vec![]) } + async fn setup_liquid_cache( + batch_size: usize, + schema: SchemaRef, + max_memory_bytes: usize, + ) -> CachedRowGroupRef { + let tmp_dir = tempfile::tempdir().unwrap(); + let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) + .await + .unwrap(); + let cache = LiquidCacheParquet::new( + batch_size, + max_memory_bytes, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + cache + .register_or_get_file("test".to_string(), schema) + .create_row_group(0, vec![0, 1]) + } + + fn build_predicate( + schema: &SchemaRef, + arrays: Vec, + expr: Arc, + ) -> LiquidPredicate { + let tmp_meta = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(tmp_meta.reopen().unwrap(), Arc::clone(schema), None).unwrap(); + writer + .write(&RecordBatch::try_new(Arc::clone(schema), arrays).unwrap()) + .unwrap(); + writer.close().unwrap(); + let reader = std::fs::File::open(tmp_meta.path()).unwrap(); + let metadata = ArrowReaderMetadata::load(&reader, ArrowReaderOptions::new()).unwrap(); + let candidate = FilterCandidateBuilder::new(expr, Arc::clone(schema)) + .build(metadata.metadata()) + .unwrap() + .unwrap(); + let projection = candidate.projection(metadata.metadata()); + LiquidPredicate::try_new(candidate, projection).unwrap() + } + + fn equals(name: &str, index: usize, value: ScalarValue) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new(name, index)), + Operator::Eq, + Arc::new(Literal::new(value)), + )) + } + + fn dummy(len: usize) -> ArrayRef { + Arc::new(Int32Array::from(vec![0; len])) + } + + async fn insert_liquid_columns( + row_group: &CachedRowGroupRef, + batch_id: BatchID, + arrays: [ArrayRef; 3], + ) { + for (column_id, array) in arrays.into_iter().enumerate() { + row_group + .get_column(column_id as u64) + .unwrap() + .insert(batch_id, array) + .await + .unwrap(); + } + assert_eq!(row_group.cache_store.stats().memory_liquid_entries, 2); + } + + #[tokio::test] + async fn or_fast_path_evaluates_on_liquid() { + let batch_size = 1024; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("dummy", DataType::Int32, false), + ])); + let a: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..batch_size).map(|i| (i % 4) as i32 + 1), + )); + let b: ArrayRef = Arc::new(Int32Array::from_iter_values( + (0..batch_size).map(|i| (i % 5) as i32 * 10), + )); + let budget = a.get_array_memory_size() + b.get_array_memory_size(); + let row_group = setup_liquid_cache(batch_size, schema.clone(), budget).await; + let batch_id = BatchID::from_row_id(0, batch_size); + insert_liquid_columns(&row_group, batch_id, [a.clone(), b.clone(), dummy(1)]).await; + let expr = Arc::new(BinaryExpr::new( + equals("a", 0, ScalarValue::Int32(Some(3))), + Operator::Or, + equals("b", 1, ScalarValue::Int32(Some(20))), + )); + let mut predicate = build_predicate(&schema, vec![a, b, dummy(batch_size)], expr); + row_group.cache_store.stats(); + let selection = BooleanBuffer::new_set(batch_size); + let result = row_group + .evaluate_selection_with_predicate(batch_id, &selection, &mut predicate) + .await + .unwrap() + .unwrap(); + let expected = BooleanBuffer::collect_bool(batch_size, |i| i % 4 == 2 || i % 5 == 2); + assert_eq!(result, BooleanArray::new(expected, None)); + assert!(row_group.cache_store.stats().runtime.try_read_liquid_calls >= 2); + } + + #[tokio::test] + async fn or_fast_path_on_strings() { + let batch_size = 128; + let schema = Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8View, false), + Field::new("name", DataType::Utf8View, false), + Field::new("dummy", DataType::Int32, false), + ])); + let city: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..batch_size).map(|i| if i % 5 == 2 { "Tokyo" } else { "Paris" }), + )); + let name: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..batch_size).map(|i| if i % 4 == 1 { "Bob" } else { "Alice" }), + )); + let budget = city.get_array_memory_size() + name.get_array_memory_size(); + let row_group = setup_liquid_cache(batch_size, schema.clone(), budget).await; + let batch_id = BatchID::from_row_id(0, batch_size); + insert_liquid_columns(&row_group, batch_id, [city.clone(), name.clone(), dummy(1)]).await; + let expr = Arc::new(BinaryExpr::new( + equals("name", 1, ScalarValue::Utf8View(Some("Bob".into()))), + Operator::Or, + equals("city", 0, ScalarValue::Utf8View(Some("Tokyo".into()))), + )); + let mut predicate = build_predicate(&schema, vec![city, name, dummy(batch_size)], expr); + row_group.cache_store.stats(); + let selection = BooleanBuffer::new_set(batch_size); + let result = row_group + .evaluate_selection_with_predicate(batch_id, &selection, &mut predicate) + .await + .unwrap() + .unwrap(); + let expected = BooleanBuffer::collect_bool(batch_size, |i| i % 4 == 1 || i % 5 == 2); + assert_eq!(result, BooleanArray::new(expected, None)); + assert!(row_group.cache_store.stats().runtime.try_read_liquid_calls >= 2); + } + #[tokio::test] async fn evaluate_or_on_cached_columns() { let batch_size = 4; diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 597df2113..3d3ae3dcb 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -26,17 +26,27 @@ use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHint #[derive(Debug)] pub struct LocalModeOptimizer { cache: LiquidCacheParquetRef, + prefetch: bool, } impl LocalModeOptimizer { /// Create an optimizer with an existing cache instance pub fn new(cache: LiquidCacheParquetRef) -> Self { - Self { cache } + Self { + cache, + prefetch: true, + } } /// Create an optimizer with an existing cache instance pub fn with_cache(cache: LiquidCacheParquetRef) -> Self { - Self { cache } + Self::new(cache) + } + + /// Enable or disable row-group prefetching. + pub fn with_prefetch(mut self, prefetch: bool) -> Self { + self.prefetch = prefetch; + self } } @@ -48,8 +58,9 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { ) -> Result, datafusion::error::DataFusionError> { let analysis = HintAnalyzer::analyze(&plan); let cache = self.cache.clone(); + let prefetch = self.prefetch; let mut convert = |node: &Arc, hints: ColumnSqueezeHints| { - convert_parquet_scan(node, &cache, hints) + convert_parquet_scan(node, &cache, hints, prefetch) }; Ok(squeeze_hint::rewrite_with_hints( plan, @@ -79,7 +90,7 @@ pub fn rewrite_data_source_plan_with_hints( hints: &ColumnSqueezeHints, ) -> Arc { plan.transform_up( - |node| match convert_parquet_scan(&node, cache, hints.clone()) { + |node| match convert_parquet_scan(&node, cache, hints.clone(), true) { Some(new_node) => Ok(Transformed::new( new_node, true, @@ -106,6 +117,7 @@ fn convert_parquet_scan( node: &Arc, cache: &LiquidCacheParquetRef, hints: ColumnSqueezeHints, + prefetch: bool, ) -> Option> { let data_source_exec = node.downcast_ref::()?; let (file_scan_config, parquet_source) = @@ -113,7 +125,8 @@ fn convert_parquet_scan( let new_source = LiquidParquetSource::from_parquet_source(parquet_source.clone(), cache.clone()) - .with_squeeze_hints(Arc::new(hints)); + .with_squeeze_hints(Arc::new(hints)) + .with_prefetch(prefetch); let mut new_config = file_scan_config.clone(); new_config.file_source = Arc::new(new_source); diff --git a/src/datafusion/src/reader/plantime/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index e87b02006..ddc34eadc 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -7,5 +7,5 @@ mod morselizer; mod row_filter; mod source; -pub(crate) use morselizer::LiquidMorselizer; +pub(crate) use morselizer::{LiquidFileMetrics, LiquidMorselizer}; pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 210002e85..6098ed4de 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -1,8 +1,8 @@ -use std::{fmt, future::Future, sync::Arc}; +use std::{collections::VecDeque, fmt, future::Future, sync::Arc}; use arrow_schema::SchemaRef; use datafusion::{ - common::exec_err, + common::{exec_err, internal_err}, datasource::{ listing::{FileRange, PartitionedFile}, physical_plan::{ @@ -23,27 +23,37 @@ use datafusion::{ physical_optimizer::pruning::{FilePruner, PruningPredicate, build_pruning_predicate}, physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}, }; -use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselizer}; +#[cfg(test)] +use datafusion_datasource::morsel::Morsel; +use datafusion_datasource::morsel::{MorselPlan, MorselPlanner, Morselizer}; use futures::{FutureExt, future::BoxFuture}; use log::debug; use parquet::{ arrow::{ ParquetRecordBatchStreamBuilder, ProjectionMask, - arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}, + arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection}, parquet_column, }, file::metadata::PageIndexPolicy, }; +use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; use crate::{ - cache::{ColumnSqueezeHints, LiquidCacheParquetRef}, + cache::{ + BatchID, ColumnSqueezeHints, InsertArrowArrayError, LiquidCacheParquetRef, PrefetchOutcome, + RowGroupSnapshots, + }, reader::{ plantime::row_filter::build_row_filter, - runtime::{LiquidRowGroupPlanner, build_projection_schema, get_root_column_ids}, + runtime::{ + LiquidRowGroupPlanner, apply_predicates, build_projection_schema, get_root_column_ids, + take_next_batch, + }, }, + utils::row_selector_to_boolean_buffer, }; - -use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; +#[cfg(test)] +use liquid_cache::cache::{CachedBatchType, EntryID, LiquidCache}; pub(crate) struct LiquidMorselizer { pub(crate) partition_index: usize, @@ -58,6 +68,7 @@ pub(crate) struct LiquidMorselizer { pub(crate) expr_adapter_factory: Arc, pub(crate) span: Option>, pub(crate) squeeze_hints: Arc, + pub(crate) prefetch: bool, } impl fmt::Debug for LiquidMorselizer { @@ -74,7 +85,7 @@ impl Morselizer for LiquidMorselizer { let file_range = partitioned_file.range.clone(); let access_plan = partitioned_file.extensions.get_arc::(); let file_name = partitioned_file.object_meta.location.to_string(); - let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); + let metrics = LiquidFileMetrics::new(self.partition_index, &file_name, &self.metrics); let metadata_size_hint = partitioned_file.metadata_size_hint; let file_location = partitioned_file.object_meta.location.to_string(); let reader = self.parquet_file_reader_factory.create_liquid_reader( @@ -109,8 +120,6 @@ impl Morselizer for LiquidMorselizer { .transpose()?; } - let predicate_creation_errors = - MetricBuilder::new(&self.metrics).global_counter("num_predicate_creation_errors"); let file_pruner = predicate .as_ref() .filter(|predicate| { @@ -122,7 +131,7 @@ impl Morselizer for LiquidMorselizer { Arc::clone(predicate), &logical_file_schema, &partitioned_file, - predicate_creation_errors.clone(), + metrics.predicate_creation_errors.clone(), ) }); let span = self.span.as_ref().map(|span| { @@ -137,7 +146,7 @@ impl Morselizer for LiquidMorselizer { file_range, access_plan, file_name, - file_metrics, + metrics, file_pruner, reader, batch_size: self.batch_size, @@ -145,23 +154,49 @@ impl Morselizer for LiquidMorselizer { output_schema, projection, predicate, - predicate_creation_errors, reorder_filters: self.reorder_filters, liquid_cache: self.liquid_cache.clone(), expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), file_location, span, squeeze_hints: Arc::clone(&self.squeeze_hints), + prefetch: self.prefetch, })), })) } } +#[derive(Clone)] +pub(crate) struct LiquidFileMetrics { + pub(crate) file_metrics: ParquetFileMetrics, + pub(crate) predicate_creation_errors: Count, + pub(crate) batches_prefetched: Count, + pub(crate) prefetch_skipped: Count, +} + +impl LiquidFileMetrics { + pub(crate) fn new( + partition_index: usize, + file_name: &str, + metrics: &ExecutionPlanMetricsSet, + ) -> Self { + Self { + file_metrics: ParquetFileMetrics::new(partition_index, file_name, metrics), + predicate_creation_errors: MetricBuilder::new(metrics) + .global_counter("num_predicate_creation_errors"), + batches_prefetched: MetricBuilder::new(metrics) + .counter("batches_prefetched", partition_index), + prefetch_skipped: MetricBuilder::new(metrics) + .counter("prefetch_skipped", partition_index), + } + } +} + struct PreparedLiquidOpen { file_range: Option, access_plan: Option>, file_name: String, - file_metrics: ParquetFileMetrics, + metrics: LiquidFileMetrics, file_pruner: Option, reader: ParquetMetadataCacheReader, batch_size: usize, @@ -169,13 +204,13 @@ struct PreparedLiquidOpen { output_schema: SchemaRef, projection: ProjectionExprs, predicate: Option>, - predicate_creation_errors: Count, reorder_filters: bool, liquid_cache: LiquidCacheParquetRef, expr_adapter_factory: Arc, file_location: String, span: Option>, squeeze_hints: Arc, + prefetch: bool, } struct MetadataLoadedLiquidOpen { @@ -243,6 +278,7 @@ impl LiquidOpenState { && file_pruner.should_prune()? { prepared + .metrics .file_metrics .files_ranges_pruned_statistics .add_pruned(1); @@ -250,6 +286,7 @@ impl LiquidOpenState { } prepared + .metrics .file_metrics .files_ranges_pruned_statistics .add_matched(1); @@ -257,7 +294,8 @@ impl LiquidOpenState { async move { let options = ArrowReaderOptions::new() .with_page_index_policy(PageIndexPolicy::Required); - let metadata_load_time = prepared.file_metrics.metadata_load_time.clone(); + let metadata_load_time = + prepared.metrics.file_metrics.metadata_load_time.clone(); let mut timer = metadata_load_time.timer(); let reader_metadata = ArrowReaderMetadata::load_async(&mut prepared.reader, options.clone()) @@ -284,7 +322,7 @@ impl LiquidOpenState { .expect("bloom filters are loaded only with a pruning predicate"); prepared.row_groups.prune_by_bloom_filters( predicate, - &prepared.context.prepared.file_metrics, + &prepared.context.prepared.metrics.file_metrics, &loaded.bloom_filters, ); Ok(Self::PlanRowGroups(Box::new(prune_pages(prepared)))) @@ -296,7 +334,12 @@ impl LiquidOpenState { } fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result { - let metadata_load_time = loaded.prepared.file_metrics.metadata_load_time.clone(); + let metadata_load_time = loaded + .prepared + .metrics + .file_metrics + .metadata_load_time + .clone(); let mut metadata_timer = metadata_load_time.timer(); let physical_file_schema = Arc::clone(loaded.reader_metadata.schema()); let cache_full_schema = Arc::clone(&physical_file_schema); @@ -331,7 +374,7 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result
  • filter, Err(error) => { @@ -379,7 +422,7 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result
  • PlannedRowGroups { &context.physical_file_schema, context.builder.parquet_schema(), context.builder.metadata().as_ref(), - &context.prepared.file_metrics, + &context.prepared.metrics.file_metrics, ); } PlannedRowGroups { @@ -503,6 +546,7 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result Result Result> = Vec::with_capacity(row_group_indexes.len()); + let mut queue = VecDeque::with_capacity(row_group_indexes.len()); for row_group_idx in row_group_indexes { let row_count = row_group_metadata[row_group_idx].num_rows() as usize; let row_group_selection = selection .as_mut() .map(|selection| selection.split_off(row_count)); - if let Some(morsel) = row_group_planner.plan(row_group_idx, row_group_selection) { - morsels.push(Box::new(morsel)); + let row_group_selection = row_group_selection.unwrap_or_else(|| { + vec![parquet::arrow::arrow_reader::RowSelector::select(row_count)].into() + }); + if row_group_selection.row_count() > 0 { + queue.push_back((row_group_idx, row_group_selection)); } } - Ok((!morsels.is_empty()).then(|| MorselPlan::new().with_morsels(morsels))) + if queue.is_empty() { + return Ok(None); + } + + let chain = LiquidRowGroupChain { + planner: row_group_planner, + queue, + snapshots: Arc::default(), + prefetch, + }; + let plan = if prefetch { + MorselPlan::new().with_pending_planner(prefetch_future(chain)) + } else { + MorselPlan::new().with_planners(vec![Box::new(chain)]) + }; + Ok(Some(plan)) +} + +struct LiquidRowGroupChain { + planner: Arc, + queue: VecDeque<(usize, RowSelection)>, + snapshots: Arc, + prefetch: bool, +} + +impl fmt::Debug for LiquidRowGroupChain { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LiquidRowGroupChain") + .field("remaining_row_groups", &self.queue.len()) + .finish_non_exhaustive() + } +} + +impl MorselPlanner for LiquidRowGroupChain { + fn plan(mut self: Box) -> Result> { + let (row_group_idx, selection) = self + .queue + .pop_front() + .expect("a row group chain is never empty"); + let snapshots = std::mem::take(&mut self.snapshots); + let Some(morsel) = self.planner.plan(row_group_idx, Some(selection), snapshots) else { + return internal_err!("selected row group {row_group_idx} produced no morsel"); + }; + let mut plan = MorselPlan::new().with_morsels(vec![Box::new(morsel)]); + if self.queue.is_empty() { + return Ok(Some(plan)); + } + + if !self.prefetch { + return Ok(Some(plan.with_planners(vec![self]))); + } + + let next_row_group = self.queue.front().unwrap().0; + let estimate = self.planner.estimated_bytes(next_row_group); + let headroom = self + .planner + .liquid_cache + .max_memory_bytes() + .saturating_sub(self.planner.liquid_cache.memory_usage_bytes()); + if headroom >= estimate { + plan = plan.with_pending_planner(prefetch_future(*self)); + } else { + self.planner.metrics.prefetch_skipped.add(1); + plan = plan.with_planners(vec![self]); + } + Ok(Some(plan)) + } +} + +async fn prefetch_future(chain: LiquidRowGroupChain) -> Result> { + Ok(Box::new(prefetch_front(chain).await) as Box) +} + +async fn prefetch_front(chain: LiquidRowGroupChain) -> LiquidRowGroupChain { + let (row_group_idx, selection) = chain.queue.front().expect("prefetch chain has work"); + let mut selectors: VecDeque<_> = selection.clone().into(); + let batch_size = chain.planner.cached_file.batch_size(); + let selected_batches = std::iter::from_fn(|| take_next_batch(&mut selectors, batch_size)) + .enumerate() + .filter_map(|(idx, selection)| { + let selection = row_selector_to_boolean_buffer(&selection); + (selection.count_set_bits() > 0).then_some((BatchID::from_raw(idx as u16), selection)) + }) + .collect::>(); + let estimate = chain.planner.estimated_bytes(*row_group_idx); + let per_batch_estimate = estimate / selected_batches.len().max(1); + let snapshots = Arc::clone(&chain.snapshots); + let mut context = chain + .planner + .prefetch_context(*row_group_idx, Arc::clone(&snapshots)); + + for (batch_id, input_selection) in selected_batches { + let mut produced_snapshots = false; + let predicate_summary = prefetch_columns( + &context.cached_row_group, + batch_id, + &context.predicate_column_ids, + ) + .await; + produced_snapshots |= predicate_summary.any_snapshotted; + + if predicate_summary.any_missing { + match materialize_prefetch_batch(&mut context, batch_id).await { + Ok(()) => produced_snapshots = true, + Err(error) => { + debug!("Stopping row group {row_group_idx} prefetch: {error}"); + break; + } + } + } + + let filtered_selection = if let Some(filter) = context.row_filter.as_mut() { + match apply_predicates(&context.cached_row_group, batch_id, input_selection, filter) + .await + { + Ok(selection) => selection, + Err(error) => { + debug!("Stopping row group {row_group_idx} prefetch: {error}"); + break; + } + } + } else { + Some(input_selection) + }; + + if let Some(filtered_selection) = filtered_selection { + snapshots.insert_selection(batch_id, filtered_selection.clone()); + + if filtered_selection.count_set_bits() > 0 { + let projection_summary = prefetch_columns( + &context.cached_row_group, + batch_id, + &context.projection_column_ids, + ) + .await; + produced_snapshots |= projection_summary.any_snapshotted; + if projection_summary.any_missing { + match materialize_prefetch_batch(&mut context, batch_id).await { + Ok(()) => produced_snapshots = true, + Err(error) => { + debug!("Stopping row group {row_group_idx} prefetch: {error}"); + break; + } + } + } + } + } + + if produced_snapshots { + chain.planner.metrics.batches_prefetched.add(1); + } + let headroom = chain + .planner + .liquid_cache + .max_memory_bytes() + .saturating_sub(chain.planner.liquid_cache.memory_usage_bytes()); + if headroom < per_batch_estimate { + break; + } + } + + chain +} + +struct PrefetchColumnsSummary { + any_missing: bool, + any_snapshotted: bool, +} + +async fn prefetch_columns( + row_group: &crate::cache::CachedRowGroupRef, + batch_id: BatchID, + column_ids: &[usize], +) -> PrefetchColumnsSummary { + let mut summary = PrefetchColumnsSummary { + any_missing: false, + any_snapshotted: false, + }; + for column_id in column_ids { + let column = row_group.get_column(*column_id as u64).unwrap(); + match column.prefetch_snapshot(batch_id).await { + PrefetchOutcome::Snapshotted => summary.any_snapshotted = true, + PrefetchOutcome::Missing => summary.any_missing = true, + PrefetchOutcome::AlreadySnapshotted | PrefetchOutcome::Squeezed => {} + } + } + summary +} + +async fn materialize_prefetch_batch( + context: &mut crate::reader::runtime::LiquidRowGroupPrefetchContext, + batch_id: BatchID, +) -> std::result::Result<(), parquet::errors::ParquetError> { + let record_batch = context.fallback.fetch_batch(batch_id).await?; + for (position, column_id) in context.cache_column_ids.iter().enumerate() { + let column = context + .cached_row_group + .get_column(*column_id as u64) + .unwrap(); + let array = Arc::clone(record_batch.column(position)); + match column.insert(batch_id, Arc::clone(&array)).await { + Ok(()) | Err(InsertArrowArrayError::AlreadyCached) => {} + Err(InsertArrowArrayError::CacheFull) => {} + } + column.insert_snapshot(batch_id, array); + } + Ok(()) } async fn load_bloom_filters( @@ -685,7 +940,7 @@ mod tests { use crate::{ cache::{BatchID, CachedFileRef, CachedRowGroupRef, LiquidCacheParquet}, - reader::LiquidParquetSource, + reader::{LiquidParquetSource, extract_multi_column_or}, }; use super::*; @@ -699,6 +954,13 @@ mod tests { _tmp_dir: tempfile::TempDir, } + struct TestFilePlanner { + planner: Box, + cache: Arc, + cached_file: CachedFileRef, + tmp_dir: tempfile::TempDir, + } + fn schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), @@ -806,7 +1068,7 @@ mod tests { ) } - async fn plan_test_file(options: PlanOptions) -> PlannedTestFile { + async fn prepare_test_file(options: PlanOptions) -> TestFilePlanner { let schema = schema(); let tmp_dir = tempfile::tempdir().unwrap(); let file_id = NEXT_FILE_ID.fetch_add(1, Ordering::Relaxed); @@ -845,14 +1107,43 @@ mod tests { expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), span: None, squeeze_hints: Arc::default(), + prefetch: true, }; - let morsels = drive_planner(morselizer.plan_file(partitioned_file).unwrap()).await; let cached_file = cache.register_or_get_file(file_name, schema); + TestFilePlanner { + planner: morselizer.plan_file(partitioned_file).unwrap(), + cache, + cached_file, + tmp_dir, + } + } + + async fn plan_test_file(options: PlanOptions) -> PlannedTestFile { + let prepared = prepare_test_file(options).await; + let morsels = drive_planner(prepared.planner).await; PlannedTestFile { morsels, - _cache: cache, - cached_file, - _tmp_dir: tmp_dir, + _cache: prepared.cache, + cached_file: prepared.cached_file, + _tmp_dir: prepared.tmp_dir, + } + } + + async fn advance_to_row_group_chain( + mut planner: Box, + ) -> Box { + loop { + let mut plan = planner.plan().unwrap().expect("file has row groups"); + assert!(plan.take_morsels().is_empty()); + if let Some(ready) = plan.take_ready_planners().pop() { + planner = ready; + continue; + } + let pending = plan.take_pending_planner().expect("planner has more work"); + planner = pending.await.unwrap(); + if format!("{planner:?}").contains("LiquidRowGroupChain") { + return planner; + } } } @@ -864,6 +1155,14 @@ mod tests { )) } + fn eq_expr(column_name: &str, column_index: usize, literal: i32) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new(column_name, column_index)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(literal)))), + )) + } + #[tokio::test] async fn metadata_cache_is_scoped_to_object_store() { let schema = schema(); @@ -947,7 +1246,7 @@ mod tests { } } - async fn is_cached(row_group: &CachedRowGroupRef, column_id: usize, batch_idx: u16) -> bool { + async fn contains(row_group: &CachedRowGroupRef, column_id: usize, batch_idx: u16) -> bool { row_group .get_column(column_id as u64) .unwrap() @@ -956,6 +1255,16 @@ mod tests { .is_some() } + fn kind_of(cache: &LiquidCache, id: &EntryID) -> Option { + let mut kind = None; + cache.for_each_entry(|entry_id, entry| { + if entry_id == id { + kind = Some(CachedBatchType::from(entry)); + } + }); + kind + } + #[tokio::test] async fn plans_one_morsel_per_selected_row_group() { let all = plan_test_file(PlanOptions { @@ -977,6 +1286,185 @@ mod tests { assert_eq!(collect_columns(pruned.morsels).await.0, vec![4, 5, 6, 7]); } + #[tokio::test] + async fn prefetch_hands_snapshots_to_next_morsel() { + let file = prepare_test_file(PlanOptions::default()).await; + let chain = advance_to_row_group_chain(file.planner).await; + let mut first_plan = chain.plan().unwrap().unwrap(); + let mut morsels = first_plan.take_morsels(); + assert_eq!(morsels.len(), 1); + let next = first_plan.take_pending_planner().unwrap().await.unwrap(); + + let row_group = file.cached_file.create_row_group(1, vec![]); + for column_id in 0..2 { + let id = row_group + .get_column(column_id) + .unwrap() + .entry_id(BatchID::from_raw(0)) + .into(); + assert_eq!( + kind_of(file.cache.storage(), &id), + Some(CachedBatchType::MemoryArrow) + ); + } + + morsels.extend(next.plan().unwrap().unwrap().take_morsels()); + assert_eq!( + collect_columns(morsels).await.0, + vec![0, 1, 2, 3, 4, 5, 6, 7] + ); + } + + #[tokio::test] + async fn prefetched_multi_column_or_uses_snapshots() { + let predicate: Arc = Arc::new(BinaryExpr::new( + eq_expr("a", 0, 3), + Operator::Or, + eq_expr("b", 1, 20), + )); + assert!(extract_multi_column_or(&predicate).is_some()); + let file = prepare_test_file(PlanOptions { + predicate: Some(predicate), + ..Default::default() + }) + .await; + file.cache.storage().stats(); + + let chain = advance_to_row_group_chain(file.planner).await; + let morsels = chain.plan().unwrap().unwrap().take_morsels(); + assert_eq!(collect_columns(morsels).await, (vec![3], vec![13])); + assert_eq!( + file.cache.storage().stats().runtime.try_read_liquid_calls, + 0 + ); + } + + #[tokio::test] + async fn prefetch_fetches_absent_predicate_columns() { + let file = prepare_test_file(PlanOptions { + predicate: Some(gt_expr("a", 0, 7)), + projection_columns: vec![1], + single_row_group_values: Some((0..12).collect()), + ..Default::default() + }) + .await; + + let chain = advance_to_row_group_chain(file.planner).await; + let predicate = file + .cached_file + .create_row_group(0, vec![0]) + .get_column(0) + .unwrap(); + for batch_idx in 0..3 { + let entry_id = predicate.entry_id(BatchID::from_raw(batch_idx)).into(); + assert_eq!( + kind_of(file.cache.storage(), &entry_id), + Some(CachedBatchType::MemoryArrow) + ); + } + + let morsels = chain.plan().unwrap().unwrap().take_morsels(); + assert_eq!( + collect_columns(morsels).await.0, + vec![1008, 1009, 1010, 1011] + ); + } + + #[tokio::test] + async fn prefetch_skips_projection_for_filtered_batches() { + let file = prepare_test_file(PlanOptions { + predicate: Some(gt_expr("a", 0, 7)), + projection_columns: vec![1], + single_row_group_values: Some((0..12).collect()), + ..Default::default() + }) + .await; + let row_group = file.cached_file.create_row_group(0, vec![0]); + insert_batches( + &row_group, + 0, + &[(0, &[0, 1, 2, 3]), (1, &[4, 5, 6, 7]), (2, &[8, 9, 10, 11])], + ) + .await; + insert_batches( + &row_group, + 1, + &[ + (0, &[1000, 1001, 1002, 1003]), + (1, &[1004, 1005, 1006, 1007]), + (2, &[1008, 1009, 1010, 1011]), + ], + ) + .await; + file.cache.flush_data().await.unwrap(); + + let chain = advance_to_row_group_chain(file.planner).await; + let projection = row_group.get_column(1).unwrap(); + for batch_idx in 0..2 { + let entry_id = projection.entry_id(BatchID::from_raw(batch_idx)).into(); + assert_eq!( + kind_of(file.cache.storage(), &entry_id), + Some(CachedBatchType::DiskArrow) + ); + } + let surviving_entry = projection.entry_id(BatchID::from_raw(2)).into(); + assert_eq!( + kind_of(file.cache.storage(), &surviving_entry), + Some(CachedBatchType::MemoryArrow) + ); + + let morsels = chain.plan().unwrap().unwrap().take_morsels(); + assert_eq!( + collect_columns(morsels).await.0, + vec![1008, 1009, 1010, 1011] + ); + } + + #[tokio::test] + async fn snapshots_survive_eviction() { + let file = prepare_test_file(PlanOptions::default()).await; + let chain = advance_to_row_group_chain(file.planner).await; + let mut first = chain.plan().unwrap().unwrap(); + let next = first.take_pending_planner().unwrap().await.unwrap(); + let second = next.plan().unwrap().unwrap().take_morsels(); + file.cache.flush_data().await.unwrap(); + + let row_group = file.cached_file.create_row_group(1, vec![]); + assert_eq!(collect_columns(second).await.0, vec![4, 5, 6, 7]); + for column_id in 0..2 { + let id = row_group + .get_column(column_id) + .unwrap() + .entry_id(BatchID::from_raw(0)) + .into(); + assert_eq!( + kind_of(file.cache.storage(), &id), + Some(CachedBatchType::DiskArrow) + ); + } + } + + #[tokio::test] + async fn headroom_gate_skips_prefetch() { + let file = prepare_test_file(PlanOptions { + max_memory_bytes: 1, + max_disk_bytes: 0, + ..Default::default() + }) + .await; + let chain = advance_to_row_group_chain(file.planner).await; + let mut first = chain.plan().unwrap().unwrap(); + assert!(first.take_pending_planner().is_none()); + let next = first.take_ready_planners().pop().unwrap(); + + let mut morsels = first.take_morsels(); + morsels.extend(next.plan().unwrap().unwrap().take_morsels()); + assert_eq!( + collect_columns(morsels).await.0, + vec![0, 1, 2, 3, 4, 5, 6, 7] + ); + } + #[tokio::test] async fn cache_full_keeps_inserted_batches_and_skips_failed_inserts() { let one_array_memory = Arc::new(Int32Array::from(vec![0, 1, 2, 3])).get_array_memory_size(); @@ -992,10 +1480,10 @@ mod tests { let (a, b) = collect_columns(planned.morsels).await; assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); + assert!(contains(&row_group0, 0, 0).await); + assert!(contains(&row_group0, 1, 0).await); + assert!(contains(&row_group1, 0, 0).await); + assert!(!contains(&row_group1, 1, 0).await); } #[tokio::test] @@ -1013,10 +1501,10 @@ mod tests { let (a, b) = collect_columns(planned.morsels).await; assert_eq!(a, vec![3, 4, 5, 6, 7]); assert_eq!(b, vec![13, 14, 15, 16, 17]); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); + assert!(contains(&row_group0, 0, 0).await); + assert!(contains(&row_group0, 1, 0).await); + assert!(contains(&row_group1, 0, 0).await); + assert!(!contains(&row_group1, 1, 0).await); } #[tokio::test] @@ -1033,8 +1521,8 @@ mod tests { assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); for row_group in [&row_group0, &row_group1] { - assert!(!is_cached(row_group, 0, 0).await); - assert!(!is_cached(row_group, 1, 0).await); + assert!(!contains(row_group, 0, 0).await); + assert!(!contains(row_group, 1, 0).await); } } @@ -1055,28 +1543,25 @@ mod tests { collect_columns(planned.morsels).await.0, vec![3, 4, 5, 6, 7] ); - assert!(is_cached(&row_group0, 0, 0).await); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 0, 0).await); - assert!(!is_cached(&row_group1, 1, 0).await); + assert!(contains(&row_group0, 0, 0).await); + assert!(contains(&row_group0, 1, 0).await); + assert!(contains(&row_group1, 0, 0).await); + assert!(!contains(&row_group1, 1, 0).await); } #[tokio::test] async fn missing_column_falls_back_to_parquet() { - let planned = plan_test_file(PlanOptions { - ..Default::default() - }) - .await; - let row_group0 = planned.cached_file.create_row_group(0, vec![]); - let row_group1 = planned.cached_file.create_row_group(1, vec![]); + let file = prepare_test_file(PlanOptions::default()).await; + let row_group0 = file.cached_file.create_row_group(0, vec![]); + let row_group1 = file.cached_file.create_row_group(1, vec![]); insert_batches(&row_group0, 0, &[(0, &[0, 1, 2, 3])]).await; insert_batches(&row_group1, 0, &[(0, &[4, 5, 6, 7])]).await; - let (a, b) = collect_columns(planned.morsels).await; + let (a, b) = collect_columns(drive_planner(file.planner).await).await; assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); - assert!(is_cached(&row_group0, 1, 0).await); - assert!(is_cached(&row_group1, 1, 0).await); + assert!(contains(&row_group0, 1, 0).await); + assert!(contains(&row_group1, 1, 0).await); } #[tokio::test] @@ -1084,21 +1569,21 @@ mod tests { let parquet_a = vec![ 100, 101, 102, 103, 4, 5, 6, 7, 200, 201, 202, 203, 12, 13, 14, 15, ]; - let planned = plan_test_file(PlanOptions { + let file = prepare_test_file(PlanOptions { projection_columns: vec![0], single_row_group_values: Some(parquet_a), ..Default::default() }) .await; - let row_group = planned.cached_file.create_row_group(0, vec![]); + let row_group = file.cached_file.create_row_group(0, vec![]); insert_batches(&row_group, 0, &[(0, &[0, 1, 2, 3]), (2, &[8, 9, 10, 11])]).await; assert_eq!( - collect_columns(planned.morsels).await.0, + collect_columns(drive_planner(file.planner).await).await.0, (0..16).collect::>() ); for batch_idx in 0..4 { - assert!(is_cached(&row_group, 0, batch_idx).await); + assert!(contains(&row_group, 0, batch_idx).await); } } diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 2ec80e48a..50623c01b 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -196,6 +196,7 @@ pub struct LiquidParquetSource { table_schema: TableSchema, span: Option>, squeeze_hints: Arc, + prefetch: bool, } impl LiquidParquetSource { @@ -228,6 +229,12 @@ impl LiquidParquetSource { } } + /// Enable or disable row-group prefetching. + pub fn with_prefetch(mut self, prefetch: bool) -> Self { + self.prefetch = prefetch; + self + } + /// The typed squeeze hints currently attached to this source. pub fn squeeze_hints(&self) -> &Arc { &self.squeeze_hints @@ -261,6 +268,7 @@ impl LiquidParquetSource { predicate: None, span: None, squeeze_hints: Arc::default(), + prefetch: true, }; if let Some(predicate) = predicate { @@ -323,6 +331,7 @@ impl FileSource for LiquidParquetSource { expr_adapter_factory, span: execution_span.map(Arc::new), squeeze_hints: Arc::clone(&self.squeeze_hints), + prefetch: self.prefetch, })) } diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 40c42a2b4..54e887bc9 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -7,10 +7,10 @@ use arrow::array::{Array, ArrayRef, BooleanArray, RecordBatch}; use arrow::buffer::BooleanBuffer; use arrow::compute::prep_null_mask_filter; use arrow::record_batch::RecordBatchOptions; -use arrow_schema::{ArrowError, Schema, SchemaRef}; +use arrow_schema::{ArrowError, SchemaRef}; use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; use parquet::arrow::arrow_reader::{ - ArrowPredicate, ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector, + ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector, }; use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; use parquet::errors::ParquetError; @@ -78,7 +78,7 @@ pub(crate) struct ParquetFallbackConfig { pub(crate) row_count: usize, } -struct ParquetFallback { +pub(crate) struct ParquetFallback { row_group_idx: usize, metadata: Arc, input: ParquetMetadataCacheReader, @@ -158,7 +158,7 @@ impl Stream for LiquidCacheReader { } impl ParquetFallback { - fn new(config: ParquetFallbackConfig) -> Self { + pub(crate) fn new(config: ParquetFallbackConfig) -> Self { Self { row_group_idx: config.row_group_idx, metadata: config.metadata, @@ -172,7 +172,10 @@ impl ParquetFallback { } } - async fn fetch_batch(&mut self, batch_id: BatchID) -> Result { + pub(crate) async fn fetch_batch( + &mut self, + batch_id: BatchID, + ) -> Result { if self.stream.is_none() || batch_id != self.next_batch_id { self.rebuild_stream(batch_id)?; } @@ -296,43 +299,45 @@ impl LiquidCacheReaderInner { row_filter: &mut Option, selection: Vec, ) -> Result { - let mut input_selection = row_selector_to_boolean_buffer(&selection); + let input_selection = row_selector_to_boolean_buffer(&selection); + + if let Some(snapshot_selection) = self + .cached_row_group + .snapshot_selection(self.current_batch_id) + { + return Ok(boolean_buffer_and_then( + &input_selection, + &snapshot_selection, + )); + } let Some(filter) = row_filter.as_mut() else { return Ok(input_selection); }; - for predicate in filter.predicates_mut() { - if input_selection.count_set_bits() == 0 { - break; - } - - let boolean_array = match self - .cached_row_group - .evaluate_selection_with_predicate( - self.current_batch_id, - &input_selection, - predicate, - ) - .await - { - Some(result) => result?, - None => { - self.evaluate_predicate_after_materialize(&input_selection, predicate) - .await? - } - }; - - let boolean_mask = if boolean_array.null_count() == 0 { - boolean_array.into_parts().0 - } else { - prep_null_mask_filter(&boolean_array).into_parts().0 - }; - - input_selection = boolean_buffer_and_then(&input_selection, &boolean_mask); + if let Some(selection) = apply_predicates( + &self.cached_row_group, + self.current_batch_id, + input_selection.clone(), + filter, + ) + .await? + { + return Ok(selection); } - Ok(input_selection) + self.read_parquet_batch_and_fill_cache(self.current_batch_id) + .await?; + apply_predicates( + &self.cached_row_group, + self.current_batch_id, + input_selection, + filter, + ) + .await? + .ok_or_else(|| { + ArrowError::ComputeError("predicate unavailable after materialization".to_string()) + }) } #[fastrace::trace] @@ -420,9 +425,11 @@ impl LiquidCacheReaderInner { })?; let array = Arc::clone(record_batch.column(col_idx)); - match column.insert(batch_id, array).await { + match column.insert(batch_id, Arc::clone(&array)).await { Ok(()) | Err(InsertArrowArrayError::AlreadyCached) => {} - Err(InsertArrowArrayError::CacheFull) => {} + Err(InsertArrowArrayError::CacheFull) => { + column.insert_snapshot(batch_id, array); + } } } @@ -430,57 +437,6 @@ impl LiquidCacheReaderInner { Ok(record_batch) } - async fn evaluate_predicate_after_materialize( - &mut self, - selection: &BooleanBuffer, - predicate: &mut crate::reader::LiquidPredicate, - ) -> Result { - let record_batch = self - .read_parquet_batch_and_fill_cache(self.current_batch_id) - .await?; - - if let Some(result) = self - .cached_row_group - .evaluate_selection_with_predicate(self.current_batch_id, selection, predicate) - .await - { - return result; - } - - let column_ids = predicate.predicate_column_ids(); - let mut arrays = Vec::with_capacity(column_ids.len()); - let mut fields = Vec::with_capacity(column_ids.len()); - - for column_id in column_ids { - let array = self.parquet_array(&record_batch, column_id)?; - arrays.push(filter_array(array, selection)?); - - let field = self - .cached_row_group - .get_column(column_id as u64) - .ok_or_else(|| { - ArrowError::ComputeError(format!( - "column {column_id} not present in liquid cache" - )) - })? - .field() - .as_ref() - .clone(); - fields.push(field); - } - - let schema = Arc::new(Schema::new(fields)); - let predicate_batch = if arrays.is_empty() { - let options = - RecordBatchOptions::new().with_row_count(Some(selection.count_set_bits())); - RecordBatch::try_new_with_options(schema, arrays, &options)? - } else { - RecordBatch::try_new(schema, arrays)? - }; - - predicate.evaluate(predicate_batch) - } - fn parquet_array( &self, record_batch: &RecordBatch, @@ -501,6 +457,35 @@ impl LiquidCacheReaderInner { } } +pub(crate) async fn apply_predicates( + row_group: &CachedRowGroupRef, + batch_id: BatchID, + mut input_selection: BooleanBuffer, + filter: &mut LiquidRowFilter, +) -> Result, ArrowError> { + for predicate in filter.predicates_mut() { + if input_selection.count_set_bits() == 0 { + break; + } + + let Some(boolean_array) = row_group + .evaluate_selection_with_predicate(batch_id, &input_selection, predicate) + .await + else { + return Ok(None); + }; + let boolean_array = boolean_array?; + let boolean_mask = if boolean_array.null_count() == 0 { + boolean_array.into_parts().0 + } else { + prep_null_mask_filter(&boolean_array).into_parts().0 + }; + input_selection = boolean_buffer_and_then(&input_selection, &boolean_mask); + } + + Ok(Some(input_selection)) +} + fn filter_array(array: ArrayRef, selection: &BooleanBuffer) -> Result { let selection_array = BooleanArray::new(selection.clone(), None); arrow::compute::filter(array.as_ref(), &selection_array) diff --git a/src/datafusion/src/reader/runtime/liquid_predicate.rs b/src/datafusion/src/reader/runtime/liquid_predicate.rs index 2cae21e29..405f59c07 100644 --- a/src/datafusion/src/reader/runtime/liquid_predicate.rs +++ b/src/datafusion/src/reader/runtime/liquid_predicate.rs @@ -46,21 +46,25 @@ fn extract_column_literal(expr: &Arc) -> Option<(&str, Arc() && binary.right().is::() { - return extract_column_literal(binary.left()); + return column_name(binary.left()).map(|name| (name, Arc::clone(expr))); } else if let Some(like_expr) = expr.downcast_ref::() && like_expr.pattern().is::() { - return extract_column_literal(like_expr.expr()); - } else if let Some(cast_expr) = expr.downcast_ref::() { - return extract_column_literal(cast_expr.expr()); - } else if let Some(try_cast_expr) = expr.downcast_ref::() { - return extract_column_literal(try_cast_expr.expr()); - } else if let Some(column) = expr.downcast_ref::() { - return Some((column.name(), Arc::clone(expr))); + return column_name(like_expr.expr()).map(|name| (name, Arc::clone(expr))); } None } +fn column_name(expr: &Arc) -> Option<&str> { + if let Some(cast_expr) = expr.downcast_ref::() { + column_name(cast_expr.expr()) + } else if let Some(try_cast_expr) = expr.downcast_ref::() { + column_name(try_cast_expr.expr()) + } else { + expr.downcast_ref::().map(Column::name) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/datafusion/src/reader/runtime/mod.rs b/src/datafusion/src/reader/runtime/mod.rs index 68be7a68c..b647962fa 100644 --- a/src/datafusion/src/reader/runtime/mod.rs +++ b/src/datafusion/src/reader/runtime/mod.rs @@ -1,6 +1,9 @@ +pub(crate) use liquid_cache_reader::apply_predicates; pub(crate) use liquid_predicate::extract_multi_column_or; -pub(crate) use morsel::{LiquidRowGroupPlanner, build_projection_schema}; -pub(crate) use utils::get_root_column_ids; +pub(crate) use morsel::{ + LiquidRowGroupPlanner, LiquidRowGroupPrefetchContext, build_projection_schema, +}; +pub(crate) use utils::{get_root_column_ids, take_next_batch}; mod liquid_cache_reader; mod liquid_predicate; diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index ee3f611ec..1c6931ffb 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -14,12 +14,14 @@ use parquet::{ }; use crate::{ - cache::CachedFileRef, - reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}, + cache::{CachedFileRef, CachedRowGroupRef, LiquidCacheParquetRef, RowGroupSnapshots}, + reader::plantime::{LiquidFileMetrics, LiquidRowFilter, ParquetMetadataCacheReader}, }; use super::{ - liquid_cache_reader::{LiquidCacheReader, LiquidCacheReaderConfig, ParquetFallbackConfig}, + liquid_cache_reader::{ + LiquidCacheReader, LiquidCacheReaderConfig, ParquetFallback, ParquetFallbackConfig, + }, utils::get_root_column_ids, }; @@ -35,15 +37,13 @@ pub(crate) struct LiquidRowGroupPlanner { pub(crate) projector: Arc, pub(crate) replace_schema: bool, pub(crate) span: Option>, + pub(crate) liquid_cache: LiquidCacheParquetRef, + pub(crate) metrics: LiquidFileMetrics, } impl LiquidRowGroupPlanner { - pub(crate) fn plan( - &self, - row_group_idx: usize, - selection: Option, - ) -> Option { - let metadata = self.metadata.row_group(row_group_idx); + fn cache_details(&self) -> CacheDetails { + let schema_descr = self.metadata.file_metadata().schema_descr(); let mut predicate_projection: Option = None; if let Some(filter) = &self.row_filter { for predicate in filter.predicates() { @@ -55,6 +55,84 @@ impl LiquidRowGroupPlanner { } } } + let mut cache_projection = self.projection.clone(); + if let Some(predicate_projection) = &predicate_projection { + cache_projection.union(predicate_projection); + } + CacheDetails { + projection_column_ids: get_root_column_ids(schema_descr, &self.projection), + cache_column_ids: get_root_column_ids(schema_descr, &cache_projection), + predicate_column_ids: predicate_projection + .as_ref() + .map(|projection| get_root_column_ids(schema_descr, projection)) + .unwrap_or_default(), + cache_projection, + } + } + + fn fallback_config( + &self, + row_group_idx: usize, + details: &CacheDetails, + cache_batch_size: usize, + ) -> ParquetFallbackConfig { + ParquetFallbackConfig { + row_group_idx, + metadata: Arc::clone(&self.metadata), + input: self.input.clone(), + cache_projection: details.cache_projection.clone(), + cache_column_ids: details.cache_column_ids.clone(), + cache_batch_size, + row_count: self.metadata.row_group(row_group_idx).num_rows() as usize, + } + } + + pub(crate) fn estimated_bytes(&self, row_group_idx: usize) -> usize { + let details = self.cache_details(); + self.metadata + .row_group(row_group_idx) + .columns() + .iter() + .enumerate() + .filter(|(idx, _)| details.cache_projection.leaf_included(*idx)) + .map(|(_, column)| column.uncompressed_size() as usize) + .sum() + } + + pub(crate) fn prefetch_context( + &self, + row_group_idx: usize, + snapshots: Arc, + ) -> LiquidRowGroupPrefetchContext { + let details = self.cache_details(); + let cached_row_group = self.cached_file.create_row_group_with_snapshots( + row_group_idx as u64, + details.predicate_column_ids.clone(), + snapshots, + ); + let fallback = ParquetFallback::new(self.fallback_config( + row_group_idx, + &details, + cached_row_group.batch_size(), + )); + LiquidRowGroupPrefetchContext { + cached_row_group, + cache_column_ids: details.cache_column_ids, + predicate_column_ids: details.predicate_column_ids, + projection_column_ids: details.projection_column_ids, + row_filter: self.row_filter.clone(), + fallback, + } + } + + pub(crate) fn plan( + &self, + row_group_idx: usize, + selection: Option, + snapshots: Arc, + ) -> Option { + let metadata = self.metadata.row_group(row_group_idx); + let details = self.cache_details(); let selection = selection .unwrap_or_else(|| vec![RowSelector::select(metadata.num_rows() as usize)].into()); @@ -62,21 +140,13 @@ impl LiquidRowGroupPlanner { return None; } - let mut cache_projection = self.projection.clone(); - if let Some(predicate_projection) = &predicate_projection { - cache_projection.union(predicate_projection); - } - let schema_descr = self.metadata.file_metadata().schema_descr(); - let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); - let predicate_column_ids = predicate_projection - .as_ref() - .map(|projection| get_root_column_ids(schema_descr, projection)) - .unwrap_or_default(); let projection_columns = get_root_column_ids(schema_descr, &self.projection); - let cached_row_group = self - .cached_file - .create_row_group(row_group_idx as u64, predicate_column_ids); + let cached_row_group = self.cached_file.create_row_group_with_snapshots( + row_group_idx as u64, + details.predicate_column_ids.clone(), + snapshots, + ); let cache_batch_size = cached_row_group.batch_size(); Some(LiquidRowGroupMorsel { @@ -87,15 +157,7 @@ impl LiquidRowGroupPlanner { cached_row_group, projection_columns, schema: Arc::clone(&self.stream_schema), - parquet_fallback: ParquetFallbackConfig { - row_group_idx, - metadata: Arc::clone(&self.metadata), - input: self.input.clone(), - cache_projection, - cache_column_ids, - cache_batch_size, - row_count: metadata.num_rows() as usize, - }, + parquet_fallback: self.fallback_config(row_group_idx, &details, cache_batch_size), }, output_schema: Arc::clone(&self.output_schema), projector: Arc::clone(&self.projector), @@ -105,6 +167,22 @@ impl LiquidRowGroupPlanner { } } +struct CacheDetails { + cache_projection: ProjectionMask, + cache_column_ids: Vec, + predicate_column_ids: Vec, + projection_column_ids: Vec, +} + +pub(crate) struct LiquidRowGroupPrefetchContext { + pub(crate) cached_row_group: CachedRowGroupRef, + pub(crate) cache_column_ids: Vec, + pub(crate) predicate_column_ids: Vec, + pub(crate) projection_column_ids: Vec, + pub(crate) row_filter: Option, + pub(crate) fallback: ParquetFallback, +} + pub(crate) fn build_projection_schema( file_schema: &SchemaRef, projection_column_ids: &[usize], diff --git a/src/datafusion/src/reader/runtime/utils.rs b/src/datafusion/src/reader/runtime/utils.rs index 629ffcfd7..ffa8b363c 100644 --- a/src/datafusion/src/reader/runtime/utils.rs +++ b/src/datafusion/src/reader/runtime/utils.rs @@ -27,7 +27,7 @@ pub(crate) fn get_root_column_ids( /// Take the next batch from the selection queue. /// The returning selection will have exactly the batch size, or less if the selection is exhausted. -pub(super) fn take_next_batch( +pub(crate) fn take_next_batch( selection: &mut VecDeque, batch_size: usize, ) -> Option> { From feb0fd97b89d3cfdd9e20edc3c5378a1382c4104 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Thu, 3 Sep 2026 17:27:18 -0400 Subject: [PATCH 07/24] fix high impact bugs (#515) Found by codex audit, those are indeed correctness bugs --- src/datafusion-local/src/tests/mod.rs | 1 + .../src/tests/nested_filter.rs | 98 +++++++++++++++++++ src/datafusion/bench/filter_pushdown.rs | 10 +- src/datafusion/src/cache/mod.rs | 43 ++++++-- src/datafusion/src/cache/stats.rs | 10 +- src/datafusion/src/optimizers/squeeze_hint.rs | 34 ++++++- .../src/reader/plantime/morselizer.rs | 85 ++++++++++++++-- .../src/reader/plantime/row_filter.rs | 15 +-- src/datafusion/src/reader/plantime/source.rs | 4 + .../src/reader/runtime/liquid_cache_reader.rs | 9 +- 10 files changed, 277 insertions(+), 32 deletions(-) create mode 100644 src/datafusion-local/src/tests/nested_filter.rs diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 366669cf1..e89f79b83 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -24,6 +24,7 @@ use datafusion::{ use crate::LiquidCacheLocalBuilder; mod date_optimizer; mod filter_limit; +mod nested_filter; mod squeeze; mod variants; diff --git a/src/datafusion-local/src/tests/nested_filter.rs b/src/datafusion-local/src/tests/nested_filter.rs new file mode 100644 index 000000000..fb51344dc --- /dev/null +++ b/src/datafusion-local/src/tests/nested_filter.rs @@ -0,0 +1,98 @@ +use std::{fs::File, sync::Arc}; + +use arrow::{ + array::{AsArray, Int32Array, StructArray}, + datatypes::{DataType, Field, Fields, Schema}, + record_batch::RecordBatch, +}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig}; +use parquet::arrow::ArrowWriter; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +fn write_people(path: &std::path::Path) { + let person_fields = Fields::from(vec![Field::new("age", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("person", DataType::Struct(person_fields.clone()), false), + Field::new("cohort", DataType::Int32, false), + ])); + let person = StructArray::new( + person_fields, + vec![Arc::new(Int32Array::from(vec![1, 1, 2, 2]))], + None, + ); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(person), + Arc::new(Int32Array::from(vec![10, 20, 20, 30])), + ], + ) + .unwrap(); + + let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +async fn query_rows(sql: &str) -> Vec<(i32, i32)> { + let temp_dir = TempDir::new().unwrap(); + let parquet_path = temp_dir.path().join("people.parquet"); + write_people(&parquet_path); + + let (ctx, _) = LiquidCacheLocalBuilder::new() + .with_cache_dir(temp_dir.path().to_path_buf()) + .build(SessionConfig::new()) + .await + .unwrap(); + ctx.register_parquet( + "people", + parquet_path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + ctx.sql(sql) + .await + .unwrap() + .collect() + .await + .unwrap() + .into_iter() + .flat_map(|batch| { + let ages = batch + .column(0) + .as_primitive::(); + let cohorts = batch + .column(1) + .as_primitive::(); + (0..batch.num_rows()) + .map(|row| (ages.value(row), cohorts.value(row))) + .collect::>() + }) + .collect() +} + +#[tokio::test] +async fn nested_struct_field_filter_keeps_only_matching_rows() { + let rows = query_rows( + "SELECT person['age'], cohort \ + FROM people WHERE person['age'] = 2 ORDER BY cohort", + ) + .await; + + assert_eq!(rows, vec![(2, 20), (2, 30)]); +} + +#[tokio::test] +async fn nested_and_primitive_filters_are_both_applied() { + let rows = query_rows( + "SELECT person['age'], cohort \ + FROM people WHERE person['age'] = 2 AND cohort = 20", + ) + .await; + + assert_eq!(rows, vec![(2, 20)]); +} diff --git a/src/datafusion/bench/filter_pushdown.rs b/src/datafusion/bench/filter_pushdown.rs index 897d8e2e1..2f77f974c 100644 --- a/src/datafusion/bench/filter_pushdown.rs +++ b/src/datafusion/bench/filter_pushdown.rs @@ -15,7 +15,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::metrics; -use liquid_cache_datafusion::cache::{BatchID, LiquidCacheParquet}; +use liquid_cache_datafusion::cache::{BatchID, LiquidCacheParquet, ParquetFileIdentity}; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; use rand::RngExt as _; @@ -52,7 +52,13 @@ fn setup_cache() -> (Arc, tempfile::TempDir) { )); let field = Arc::new(Field::new("test_column", DataType::Int32, false)); let schema = Arc::new(Schema::new(vec![field.clone()])); - let file = cache.register_or_get_file("test_file.parquet".to_string(), schema); + let file = cache.register_or_get_file( + ParquetFileIdentity::new( + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + "test_file.parquet".to_string(), + ), + schema, + ); let row_group = file.create_row_group(0, vec![]); (row_group.get_column(0).unwrap(), tmp_dir) } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 1984ac782..d890f480f 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -9,6 +9,7 @@ use arrow::array::{BooleanArray, RecordBatch}; use arrow::buffer::BooleanBuffer; use arrow_schema::{ArrowError, Field, Schema, SchemaRef}; use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::Column; use liquid_cache::cache::squeeze_policies::SqueezePolicy; @@ -38,6 +39,26 @@ pub use id::{BatchID, ParquetArrayID}; /// [`LiquidParquetSource`](crate::LiquidParquetSource) that opens the file. pub type ColumnSqueezeHints = HashMap>; +/// The identity of a Parquet object within an object store. +/// +/// Object paths are only unique within their object store, so both components +/// are required to keep cached data from different stores isolated. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub struct ParquetFileIdentity { + object_store_url: ObjectStoreUrl, + path: String, +} + +impl ParquetFileIdentity { + /// Create an identity from an object store URL and an object path. + pub fn new(object_store_url: ObjectStoreUrl, path: String) -> Self { + Self { + object_store_url, + path, + } + } +} + /// One column of a row group: (file column index, field, squeeze hint, is-predicate). type CachedColumnSpec = (u64, Arc, Option>, bool); @@ -323,8 +344,8 @@ pub(crate) type CachedFileRef = Arc; /// The main cache structure. #[derive(Debug)] pub struct LiquidCacheParquet { - /// Map file path to file id. - files: Mutex>, + /// Map object-store-qualified file identity to file id. + files: Mutex>, cache_store: Arc, @@ -396,23 +417,23 @@ impl LiquidCacheParquet { /// Register a file in the cache. pub fn register_or_get_file( &self, - file_path: String, + file_identity: ParquetFileIdentity, full_file_schema: SchemaRef, ) -> CachedFileRef { - self.register_or_get_file_with_hints(file_path, full_file_schema, Arc::default()) + self.register_or_get_file_with_hints(file_identity, full_file_schema, Arc::default()) } /// Register a file in the cache, attaching typed squeeze hints derived from /// the query plan (keyed by file-schema column name). pub fn register_or_get_file_with_hints( &self, - file_path: String, + file_identity: ParquetFileIdentity, full_file_schema: SchemaRef, squeeze_hints: Arc, ) -> CachedFileRef { let mut files = self.files.lock().unwrap(); let file_id = *files - .entry(file_path.clone()) + .entry(file_identity) .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); drop(files); @@ -533,7 +554,10 @@ mod tests { Box::new(AlwaysHydrate::new()), ) .await; - let file = cache.register_or_get_file("test".to_string(), schema); + let file = cache.register_or_get_file( + ParquetFileIdentity::new(ObjectStoreUrl::local_filesystem(), "test".to_string()), + schema, + ); file.create_row_group(0, vec![]) } @@ -557,7 +581,10 @@ mod tests { ) .await; cache - .register_or_get_file("test".to_string(), schema) + .register_or_get_file( + ParquetFileIdentity::new(ObjectStoreUrl::local_filesystem(), "test".to_string()), + schema, + ) .create_row_group(0, vec![0, 1]) } diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index 9527f08dd..d1fbcf41f 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -164,7 +164,7 @@ impl LiquidCacheParquet { mod tests { use std::io::Read; - use crate::cache::id::BatchID; + use crate::cache::{ParquetFileIdentity, id::BatchID}; use super::*; use arrow::{ @@ -209,7 +209,13 @@ mod tests { let mut memory_size_sum = 0; for file_no in 0..8 { let file_name = format!("test_{file_no}.parquet"); - let file = cache.register_or_get_file(file_name, schema.clone()); + let file = cache.register_or_get_file( + ParquetFileIdentity::new( + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + file_name, + ), + schema.clone(), + ); for rg in 0..8 { let row_group = file.create_row_group(rg, vec![]); for col in 0..8 { diff --git a/src/datafusion/src/optimizers/squeeze_hint.rs b/src/datafusion/src/optimizers/squeeze_hint.rs index 77c750074..b64cce3a3 100644 --- a/src/datafusion/src/optimizers/squeeze_hint.rs +++ b/src/datafusion/src/optimizers/squeeze_hint.rs @@ -223,11 +223,19 @@ impl HintAnalyzer { let usages = lineage_for_expr(&expr, &child); self.record(&usages); } - for aggr in agg.aggr_expr() { + for (aggr, filter) in agg.aggr_expr().iter().zip(agg.filter_expr()) { for expr in aggr.expressions() { let usages = lineage_for_expr(&expr, &child); self.record(&usages); } + for order_by in aggr.order_bys() { + let usages = lineage_for_expr(&order_by.expr, &child); + self.record(&usages); + } + if let Some(filter) = filter { + let usages = lineage_for_expr(filter, &child); + self.record(&usages); + } } return opaque(plan); } @@ -778,6 +786,30 @@ mod tests { assert_eq!(hints.get("date"), None); } + #[tokio::test] + async fn aggregate_filter_records_raw_column_use() { + let hints = hints_for( + "SELECT AVG(EXTRACT(YEAR FROM date)) \ + FILTER (WHERE date = DATE '2021-01-01') FROM t", + ) + .await; + + // The aggregate argument needs only YEAR, but its filter needs the + // exact date, so squeezing the column to YEAR would change the result. + assert_eq!(hints.get("date"), None); + } + + #[tokio::test] + async fn aggregate_order_by_records_raw_column_use() { + let hints = + hints_for("SELECT FIRST_VALUE(EXTRACT(MONTH FROM date) ORDER BY date DESC) FROM t") + .await; + + // The aggregate value needs only MONTH, but chronological ordering + // needs the exact date. + assert_eq!(hints.get("date"), None); + } + #[tokio::test] async fn substring_search_in_filter() { let hints = hints_for("SELECT date FROM t WHERE url LIKE '%example%'").await; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 6098ed4de..bc43ba4a0 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -40,8 +40,8 @@ use parquet::{ use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; use crate::{ cache::{ - BatchID, ColumnSqueezeHints, InsertArrowArrayError, LiquidCacheParquetRef, PrefetchOutcome, - RowGroupSnapshots, + BatchID, ColumnSqueezeHints, InsertArrowArrayError, LiquidCacheParquetRef, + ParquetFileIdentity, PrefetchOutcome, RowGroupSnapshots, }, reader::{ plantime::row_filter::build_row_filter, @@ -87,7 +87,10 @@ impl Morselizer for LiquidMorselizer { let file_name = partitioned_file.object_meta.location.to_string(); let metrics = LiquidFileMetrics::new(self.partition_index, &file_name, &self.metrics); let metadata_size_hint = partitioned_file.metadata_size_hint; - let file_location = partitioned_file.object_meta.location.to_string(); + let file_identity = ParquetFileIdentity::new( + self.parquet_file_reader_factory.object_store_url().clone(), + partitioned_file.object_meta.location.to_string(), + ); let reader = self.parquet_file_reader_factory.create_liquid_reader( self.partition_index, partitioned_file.clone(), @@ -157,7 +160,7 @@ impl Morselizer for LiquidMorselizer { reorder_filters: self.reorder_filters, liquid_cache: self.liquid_cache.clone(), expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), - file_location, + file_identity, span, squeeze_hints: Arc::clone(&self.squeeze_hints), prefetch: self.prefetch, @@ -207,7 +210,7 @@ struct PreparedLiquidOpen { reorder_filters: bool, liquid_cache: LiquidCacheParquetRef, expr_adapter_factory: Arc, - file_location: String, + file_identity: ParquetFileIdentity, span: Option>, squeeze_hints: Arc, prefetch: bool, @@ -551,7 +554,7 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result>) -> (Vec, Vec) { let mut a = Vec::new(); let mut b = Vec::new(); diff --git a/src/datafusion/src/reader/plantime/row_filter.rs b/src/datafusion/src/reader/plantime/row_filter.rs index d0c97c81a..54c778cec 100644 --- a/src/datafusion/src/reader/plantime/row_filter.rs +++ b/src/datafusion/src/reader/plantime/row_filter.rs @@ -64,7 +64,7 @@ use std::collections::BTreeSet; use std::sync::Arc; use arrow::array::BooleanArray; -use arrow::datatypes::{DataType, Schema}; +use arrow::datatypes::Schema; use arrow::error::{ArrowError, Result as ArrowResult}; use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; @@ -115,7 +115,7 @@ impl LiquidRowFilter { /// /// An expression can be evaluated as a `DatafusionArrowPredicate` if it: /// * Does not reference any projected columns -/// * Does not reference columns with non-primitive types (e.g. structs / lists) +/// * References only columns present in the physical file schema #[derive(Debug, Clone)] pub struct LiquidPredicate { /// the filter expression @@ -299,11 +299,9 @@ impl FilterCandidateBuilder { // a struct that implements TreeNodeRewriter to traverse a PhysicalExpr tree structure to determine // if any column references in the expression would prevent it from being predicate-pushed-down. -// if non_primitive_columns || projected_columns, it can't be pushed down. +// if projected_columns, it can't be pushed down. // can't be reused between calls to `rewrite`; each construction must be used only once. struct PushdownChecker<'schema> { - /// Does the expression require any non-primitive columns (like structs)? - non_primitive_columns: bool, /// Does the expression reference any columns that are not in the file schema? projected_columns: bool, // Indices into the file schema of the columns required to evaluate the expression @@ -314,7 +312,6 @@ struct PushdownChecker<'schema> { impl<'schema> PushdownChecker<'schema> { fn new(file_schema: &'schema Schema) -> Self { Self { - non_primitive_columns: false, projected_columns: false, required_columns: BTreeSet::default(), file_schema, @@ -324,10 +321,6 @@ impl<'schema> PushdownChecker<'schema> { fn check_single_column(&mut self, column_name: &str) -> Option { if let Ok(idx) = self.file_schema.index_of(column_name) { self.required_columns.insert(idx); - if DataType::is_nested(self.file_schema.field(idx).data_type()) { - self.non_primitive_columns = true; - return Some(TreeNodeRecursion::Jump); - } } else { // If the column does not exist in the file schema then it cannot be pushed down. self.projected_columns = true; @@ -339,7 +332,7 @@ impl<'schema> PushdownChecker<'schema> { #[inline] fn prevents_pushdown(&self) -> bool { - self.non_primitive_columns || self.projected_columns + self.projected_columns } } diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 50623c01b..02e3c1d6d 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -52,6 +52,10 @@ impl CachedMetaReaderFactory { Self { store, store_url } } + pub(crate) fn object_store_url(&self) -> &ObjectStoreUrl { + &self.store_url + } + pub(crate) fn create_liquid_reader( &self, partition_index: usize, diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 54e887bc9..063e68de6 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -596,7 +596,14 @@ mod tests { Box::new(AlwaysHydrate::new()), ) .await; - let file = cache.register_or_get_file("test".to_string(), schema.clone()); + let file = cache.register_or_get_file( + crate::cache::ParquetFileIdentity::new( + datafusion::execution::object_store::ObjectStoreUrl::parse("test-runtime:///") + .unwrap(), + "test".to_string(), + ), + schema.clone(), + ); let row_group = file.create_row_group(0, vec![]); let column = row_group.get_column(0).unwrap(); From 3c89a8d8fce704d812b5fbd1d0e47b5be7afcb2a Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Thu, 3 Sep 2026 22:04:12 -0400 Subject: [PATCH 08/24] Remove squeeze/hybrid storage path; replace with eviction-only flow and keep lineage metadata (#516) we will still squeeze, but I'd rather rewrite than improve. --- Cargo.lock | 1 - Cargo.toml | 1 - README.md | 4 +- benchmark/README.md | 12 + benchmark/bench_server.rs | 4 +- benchmark/src/inprocess_runner.rs | 21 +- benchmark/src/lib.rs | 15 +- dev/design/03-squeeze-integer.md | 118 -- dev/design/05-predicate-cache.md | 32 - .../src/components/cache_state_view.rs | 2 +- dev/dev-tools/src/trace/parser.rs | 17 +- dev/dev-tools/src/trace/simulator.rs | 32 +- examples/core.rs | 4 +- examples/datafusion-client-server.rs | 4 +- examples/datafusion-local.rs | 4 +- src/common/src/rpc.rs | 10 +- src/core/Cargo.toml | 10 - src/core/bench/squeeze_date32.rs | 94 -- src/core/src/cache/builders.rs | 42 +- src/core/src/cache/cached_batch.rs | 15 +- src/core/src/cache/core.rs | 317 +--- src/core/src/cache/expressions.rs | 11 +- src/core/src/cache/io_context.rs | 145 +- src/core/src/cache/mod.rs | 13 +- .../src/cache/observer/internal_tracing.rs | 54 +- src/core/src/cache/observer/mod.rs | 31 - src/core/src/cache/observer/stats.rs | 20 - .../src/cache/policies/cache/three_queue.rs | 16 +- src/core/src/cache/policies/eviction.rs | 95 ++ src/core/src/cache/policies/hydration.rs | 200 +-- src/core/src/cache/policies/mod.rs | 6 +- src/core/src/cache/policies/squeeze.rs | 863 ----------- src/core/src/cache/tests/mod.rs | 1 - src/core/src/cache/tests/policies.rs | 6 +- ...he__tests__policies__default_policies.snap | 21 +- ...ests__policies__insert_wont_fit_cache.snap | 9 +- ...ts__squeezed__read_squeezed_date_time.snap | 34 - ...__squeezed__read_squeezed_int64_array.snap | 34 - ..._squeezed__read_squeezed_variant_path.snap | 35 - src/core/src/cache/tests/squeezed.rs | 195 --- src/core/src/cache/transcode.rs | 4 +- src/core/src/cache/utils.rs | 4 +- .../byte_view_array/comparisons.rs | 198 +-- .../byte_view_array/fingerprint.rs | 25 - .../liquid_array/byte_view_array/helpers.rs | 33 +- .../src/liquid_array/byte_view_array/mod.rs | 169 +-- .../byte_view_array/serialization.rs | 6 +- .../src/liquid_array/byte_view_array/tests.rs | 179 +-- src/core/src/liquid_array/decimal_array.rs | 387 +---- src/core/src/liquid_array/float_array.rs | 634 +------- .../liquid_array/hybrid_primitive_array.rs | 1292 ----------------- src/core/src/liquid_array/mod.rs | 175 +-- src/core/src/liquid_array/primitive_array.rs | 169 +-- .../src/liquid_array/raw/bit_pack_array.rs | 1 + src/core/src/liquid_array/raw/fsst_buffer.rs | 90 +- .../src/liquid_array/squeezed_date32_array.rs | 748 ---------- src/core/src/liquid_array/tests.rs | 33 - src/core/src/liquid_array/variant_array.rs | 279 ---- src/core/study/cache_storage.rs | 34 +- src/core/study/squeeze_integer.rs | 516 ------- src/datafusion-client/src/client_exec.rs | 26 +- src/datafusion-client/src/optimizer.rs | 14 +- src/datafusion-local/src/lib.rs | 20 +- .../src/tests/date_optimizer.rs | 16 +- src/datafusion-local/src/tests/mod.rs | 88 +- ...datafusion_local__tests__os_selection.snap | 21 +- ...afusion_local__tests__provide_schema2.snap | 24 +- ...al__tests__provide_schema_with_filter.snap | 7 - ...usion_local__tests__referer_filtering.snap | 15 +- ...ests__single_column_filter_projection.snap | 7 - ..._local__tests__squeeze__basic_squeeze.snap | 74 - ...queeze__squeeze_distinct_search_phase.snap | 32 - ...ocal__tests__squeeze__squeeze_strings.snap | 75 - ...s__squeeze__squeeze_substrings_search.snap | 20 - ...eeze__squeeze_substrings_search_title.snap | 20 - ...on_local__tests__url_prefix_filtering.snap | 11 +- ...al__tests__url_selection_and_ordering.snap | 11 +- src/datafusion-local/src/tests/squeeze.rs | 150 -- src/datafusion-local/src/tests/variants.rs | 20 +- src/datafusion-server/src/lib.rs | 23 +- src/datafusion-server/src/service.rs | 20 +- src/datafusion-server/src/tests/cases.rs | 4 +- src/datafusion-server/src/tests/mod.rs | 25 +- src/datafusion/README.md | 42 +- src/datafusion/bench/filter_pushdown.rs | 4 +- src/datafusion/src/cache/column.rs | 12 +- src/datafusion/src/cache/mod.rs | 48 +- src/datafusion/src/cache/stats.rs | 4 +- src/datafusion/src/io/mod.rs | 22 +- .../{squeeze_hint.rs => lineage.rs} | 54 +- src/datafusion/src/optimizers/mod.rs | 34 +- .../src/reader/plantime/morselizer.rs | 22 +- src/datafusion/src/reader/plantime/source.rs | 20 +- .../src/reader/runtime/liquid_cache_reader.rs | 2 +- 94 files changed, 598 insertions(+), 7913 deletions(-) delete mode 100644 dev/design/03-squeeze-integer.md delete mode 100644 dev/design/05-predicate-cache.md delete mode 100644 src/core/bench/squeeze_date32.rs create mode 100644 src/core/src/cache/policies/eviction.rs delete mode 100644 src/core/src/cache/policies/squeeze.rs delete mode 100644 src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_date_time.snap delete mode 100644 src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_int64_array.snap delete mode 100644 src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_variant_path.snap delete mode 100644 src/core/src/cache/tests/squeezed.rs delete mode 100644 src/core/src/liquid_array/hybrid_primitive_array.rs delete mode 100644 src/core/src/liquid_array/squeezed_date32_array.rs delete mode 100644 src/core/src/liquid_array/variant_array.rs delete mode 100644 src/core/study/squeeze_integer.rs delete mode 100644 src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__basic_squeeze.snap delete mode 100644 src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_distinct_search_phase.snap delete mode 100644 src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap delete mode 100644 src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search.snap delete mode 100644 src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search_title.snap delete mode 100644 src/datafusion-local/src/tests/squeeze.rs rename src/datafusion/src/optimizers/{squeeze_hint.rs => lineage.rs} (95%) diff --git a/Cargo.lock b/Cargo.lock index 76b296915..b440b4e3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4222,7 +4222,6 @@ dependencies = [ "num-traits", "object_store", "parquet", - "parquet-variant-compute", "rand 0.10.2", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 28f79e68e..4f6f0debb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,7 +42,6 @@ parquet = { version = "59.2.0", features = [ "variant_experimental", ] } parquet-variant-json = { version = "59.2.0" } -parquet-variant-compute = { version = "59.2.0" } datafusion = { version = "55.0.0" } datafusion-datasource = { version = "55.0.0" } datafusion-common = { version = "55.0.0" } diff --git a/README.md b/README.md index 6515ed31f..be69b7874 100644 --- a/README.md +++ b/README.md @@ -125,9 +125,7 @@ For performance testing, disable background transcoding: ```rust let (ctx, _) = LiquidCacheLocalBuilder::new() - .with_squeeze_policy(Box::new( - squeeze_policies::Evict, - )) + .with_eviction_policy(Box::new(Evict)) .build(config) .await?; ``` diff --git a/benchmark/README.md b/benchmark/README.md index 056ec543c..2c55789b8 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -118,6 +118,18 @@ cargo run --release --bin in_process -- \ ## In process mode The benchmark uses a JSON manifest file to describe the data tables and queries to run. +### Cache representation ablation + +The core cache benchmark can compare direct Arrow eviction with the Liquid +transcode-and-evict path. Use the same input and memory budget for both runs: + +```bash +cargo bench -p liquid-cache --bench cache_storage -- \ + --parquet ../../examples/nano_hits.parquet --max-memory-mb 1 --mode arrow +cargo bench -p liquid-cache --bench cache_storage -- \ + --parquet ../../examples/nano_hits.parquet --max-memory-mb 1 --mode liquid +``` + ### JSON Format ```json diff --git a/benchmark/bench_server.rs b/benchmark/bench_server.rs index fbccae3c6..8aed63eb1 100644 --- a/benchmark/bench_server.rs +++ b/benchmark/bench_server.rs @@ -67,7 +67,7 @@ async fn main() -> Result<(), Box> { std::process::exit(1); })); } - let squeeze_policy = args.cache_mode.to_squeeze_policy(); + let eviction_policy = args.cache_mode.to_eviction_policy(); // LiquidCache server mode let ctx = LiquidCacheService::context()?; @@ -76,7 +76,7 @@ async fn main() -> Result<(), Box> { max_memory_bytes, args.disk_cache_dir.clone(), Box::new(LiquidPolicy::new()), - squeeze_policy, + eviction_policy, Box::new(NoHydration::new()), ) .await?; diff --git a/benchmark/src/inprocess_runner.rs b/benchmark/src/inprocess_runner.rs index 0aac7d495..21059fb6e 100644 --- a/benchmark/src/inprocess_runner.rs +++ b/benchmark/src/inprocess_runner.rs @@ -9,7 +9,7 @@ use datafusion::parquet::{ }; use datafusion::prelude::{SessionConfig, SessionContext}; use liquid_cache::cache::NoHydration; -use liquid_cache::cache::squeeze_policies::{Evict, TranscodeEvict, TranscodeSqueezeEvict}; +use liquid_cache::cache::{Evict, TranscodeEvict}; use liquid_cache::cache_policies::LiquidPolicy; use liquid_cache_datafusion::{LiquidCacheParquetRef, extract_execution_metrics}; use liquid_cache_datafusion_local::LiquidCacheLocalBuilder; @@ -190,7 +190,6 @@ pub enum InProcessBenchmarkMode { Arrow, #[default] Liquid, - LiquidNoSqueeze, } impl std::str::FromStr for InProcessBenchmarkMode { @@ -202,10 +201,9 @@ impl std::str::FromStr for InProcessBenchmarkMode { "datafusion-default" | "datafusion" => InProcessBenchmarkMode::DataFusionDefault, "arrow" => InProcessBenchmarkMode::Arrow, "liquid" => InProcessBenchmarkMode::Liquid, - "liquid-no-squeeze" => InProcessBenchmarkMode::LiquidNoSqueeze, _ => { return Err(format!( - "Invalid in-process benchmark mode: {s}, must be one of: parquet, datafusion-default, arrow, liquid, liquid-no-squeeze" + "Invalid in-process benchmark mode: {s}, must be one of: parquet, datafusion-default, arrow, liquid" )); } }) @@ -350,7 +348,7 @@ impl InProcessBenchmarkRunner { .with_cache_dir(cache_dir) .with_cache_policy(Box::new(LiquidPolicy::new())) .with_hydration_policy(Box::new(NoHydration::new())) - .with_squeeze_policy(Box::new(Evict)) + .with_eviction_policy(Box::new(Evict)) .build(session_config) .await?; (v.0, Some(v.1)) @@ -361,18 +359,7 @@ impl InProcessBenchmarkRunner { .with_cache_dir(cache_dir) .with_cache_policy(Box::new(LiquidPolicy::new())) .with_hydration_policy(Box::new(NoHydration::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) - .build(session_config) - .await?; - (v.0, Some(v.1)) - } - InProcessBenchmarkMode::LiquidNoSqueeze => { - let v = LiquidCacheLocalBuilder::new() - .with_max_memory_bytes(cache_size) - .with_cache_dir(cache_dir) - .with_cache_policy(Box::new(LiquidPolicy::new())) - .with_hydration_policy(Box::new(NoHydration::new())) - .with_squeeze_policy(Box::new(TranscodeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(session_config) .await?; (v.0, Some(v.1)) diff --git a/benchmark/src/lib.rs b/benchmark/src/lib.rs index 88f28a976..48a1b1cf1 100644 --- a/benchmark/src/lib.rs +++ b/benchmark/src/lib.rs @@ -7,9 +7,7 @@ use datafusion::{error::Result, physical_plan::ExecutionPlan}; use fastrace::Span; use fastrace::future::FutureExt as _; use liquid_cache::cache::CacheStats; -use liquid_cache::cache::squeeze_policies::{ - Evict, SqueezePolicy, TranscodeEvict, TranscodeSqueezeEvict, -}; +use liquid_cache::cache::{Evict, EvictionPolicy, TranscodeEvict}; use liquid_cache_common::rpc::ExecutionMetricsResponse; use liquid_cache_datafusion_server::{ApiResponse, ExecutionStats}; use log::info; @@ -302,15 +300,13 @@ pub enum BenchmarkMode { Arrow, #[default] Liquid, - LiquidNoSqueeze, } impl BenchmarkMode { - pub fn to_squeeze_policy(&self) -> Box { + pub fn to_eviction_policy(&self) -> Box { match self { BenchmarkMode::Arrow => Box::new(Evict), - BenchmarkMode::Liquid => Box::new(TranscodeSqueezeEvict), - BenchmarkMode::LiquidNoSqueeze => Box::new(TranscodeEvict), + BenchmarkMode::Liquid => Box::new(TranscodeEvict), } } } @@ -323,7 +319,6 @@ impl Display for BenchmarkMode { match self { BenchmarkMode::Arrow => "arrow", BenchmarkMode::Liquid => "liquid", - BenchmarkMode::LiquidNoSqueeze => "liquid-no-squeeze", } ) } @@ -336,7 +331,6 @@ impl FromStr for BenchmarkMode { Ok(match s { "arrow" => BenchmarkMode::Arrow, "liquid" => BenchmarkMode::Liquid, - "liquid-no-squeeze" => BenchmarkMode::LiquidNoSqueeze, _ => return Err(format!("Invalid benchmark mode: {s}")), }) } @@ -468,10 +462,9 @@ impl Display for IterationResult { write_border_sep(f, INNER)?; let total_value = format!("{}", cache_stats.total_entries); let stats_value = format!( - "(A:{} L:{} S-L:{} D-L:{} D-A:{})", + "(A:{} L:{} D-L:{} D-A:{})", cache_stats.memory_arrow_entries, cache_stats.memory_liquid_entries, - cache_stats.memory_squeezed_liquid_entries, cache_stats.disk_liquid_entries, cache_stats.disk_arrow_entries ); diff --git a/dev/design/03-squeeze-integer.md b/dev/design/03-squeeze-integer.md deleted file mode 100644 index 62a68f5c3..000000000 --- a/dev/design/03-squeeze-integer.md +++ /dev/null @@ -1,118 +0,0 @@ -# Squeeze Integer - - -This doc describes how to implement `squeeze` for integer arrays. - -Currently, we have LiquidPrimitiveArray, which is a bit-packed array. - -It does not support squeeze, so on an eviction, will be entirely written to disk. - -Instead, we'd like to support squeeze for integer arrays. - -#### Motivating example - -Let's say we have the following integer array: - -``` -[0, 32, 12, 90, 48, 368, 85, 13, 183] -``` - -The bit-width is determined by the max value, which is 368, which requires 9 bits. - -To squeeze it, we squeeze the bit-width to 5 bits, which can now support up to 32. - -We reserve value=32 as the sentinel value, and the array becomes: - -``` -[0, 32, 12, 32, 48, 32, 32, 13, 32] -``` - -Now let's say we want to evaluate bit mask for all values i < 15, consider two cases: -1. if the value is smaller than 32, we know its actual value, which we can evaluate the predicate. -2. if the value is 32, we know its greater or equal to 32, which we know it's false. - -What if we want to evaluate bit mask for all values i > 35? -1. for x < 32, we know the exact value. -2. for x>= 32, we don't know the exact value, so we need to read from disk. - -### Real design - -To squeeze a LiquidPrimitiveArray, it is squeezable if the bit-width is great or equal than 10 (a intuition). -And we always reduce the bit-width by half. - -To properly distinguish the fully in-memory and hybrid squeezed array, we need a LiquidPrimitive and LiquidPrimitive. - -This part of the code is very similar to the LiquidByteViewArray and LiquidByteViewArray implementation. - - -#### Try evaluate predicate -By default, LiquidPrimitive doesn't evaluate predicate, so will return `None`. - -But for LiquidPrimitive, we need to evaluate the predicate. -Using the example above, try to reduce the trip to disk as much as possible. - -The supported predicate is: eq, not eq, lt, lt eq, gt, gt eq. - -#### Get with selection -Get with selection will first check if the selected values are all full value, if so, return the data. -If any of the value is clamped, we need to get data from disk. - - - -## Quantize integer - -A different way to squeeze is to quantize the integer array. - -Let's say we have the following integer array: - -``` -[0, 32, 12, 90, 48, 368, 85, 13, 183] -``` - -The original array requires 9 bits. - -If we want to quantize it to 5 bits, then we'll have 32 buckets, each bucket has 512/32 = 16 values, the bucket range is like: - -``` -[0, 32) -[32, 64) -[64, 96) -... -[336, 352) -[352, 368) -[368, 384) -... -[496, 512) -``` - -The original array becomes: - -``` -[0, 1, 0, 5, 3, 23, 5, 0, 11] -``` - -Let's say we want to find `66 < i < 72`, which belongs to the 3rd bucket. -Then we check if the quantized array contains 2, if not, we can evaluate the predicate, if yes, we need to read from disk. - -### Pros and cons - -Squeeze: -- Pros: can recover the original value without reading from disk. -- Cons: most of the values are clamped, i.e., cutting the range by half will only save 1 bit. If we want to cut 10 bits, we shrink the range by 1024 times -- virtually every value is clamped. - -Quantize: -- Pros: more tight approximation of the original value, good for predicate evaluation. -- Cons: none of the values can be recovered without reading from disk. - - -## Study - -We mainly concerns about the io required for both clamp and quantize. - -To do this, we'll create a study just like `study/cache_storage.rs`, which will: -1. read integer columns from hits.parquet dataset. -2. Identify 10 representative integer filters and their corresponding columns. -2. compress the integer columns using the clamp and quantize approach. -3. Report: 1. arrow data size, 2. clamp size, 3. quantize size. -4. Report: perform the 10 representative integer filters (get_with_predicate) on both clamp and quantize approach over the corresponding column, report the io required. In the meanwhile, check the result is correct. -5. Report: use the bit-mask from above to perform get_with_selection, report the io required. diff --git a/dev/design/05-predicate-cache.md b/dev/design/05-predicate-cache.md deleted file mode 100644 index 76716bf33..000000000 --- a/dev/design/05-predicate-cache.md +++ /dev/null @@ -1,32 +0,0 @@ -Predicate cache is a new approach to cache the cache expression along with its predicate evaluation result. - -Let's consider this query: -```sql -SELECT COUNT(*) FROM table WHERE Url like '%google%'; -``` - -There are a few ways to cache data: -1. Cache the `url` column - data cache -2. Cache the result of the sql, e.g., 15211 - result cache -3. Cache a subquery: `SELECT * FROM table WHERE Url like '%google%'`, next time we rewrite the query to `SELECT COUNT(url) FROM view` - materialized view - -Predicate cache is a new approach: -1. It caches the bit mask of the predicate evaluation result, i.e., pair of (`'url like %google%'`, `[1, 0, 1, ..., 0, 1, 0]`). -2. Next time we see this predicate again, we can directly use the cache bit mask without evaluating the predicate again. - -## Implementation - -To start simple, we only consider predicate cache for string columns (ByteViewArray). - -We will have a new SqueezedArray type: PredicateSqueezedArray, which holds a hashmap of `predicate` -> `bit mask`. - -At squeeze time, we have two choices: -1. squeeze the string array by dropping the fsst buffers, i.e., becoming a `LiquidByteViewArray` -2. squeeze the string array by becoming a `PredicateSqueezedArray`, which evaluates the predicates and stores the bit mask. - -The intuition is that: 1. if the predicate is comparison, we do first, 2. if the predicate is substring search (e.g., like), we do later. - -In order for the squeeze policy to know what expression it previously evaluated, we need to register it at `CachedColumn` creation time. -CachedColumn knows the expression from the lineage analysis in `lineage_opt.rs`, which tells us all the expressions applied to a given column. - -Then in the `eval_predicate` method of LiquidCache, we will check if the predicate is already in the `PredicateSqueezedArray`, if so, we can directly return the bit mask. diff --git a/dev/dev-tools/src/components/cache_state_view.rs b/dev/dev-tools/src/components/cache_state_view.rs index fe8d4a6d9..fb9f7d602 100644 --- a/dev/dev-tools/src/components/cache_state_view.rs +++ b/dev/dev-tools/src/components/cache_state_view.rs @@ -211,7 +211,7 @@ pub fn CacheStateView(simulator: Signal) -> Element { div { class: "flex items-center gap-1 flex-wrap min-h-5", match state.victim_status.get(&entry.entry_id) { Some(VictimStatus::Selected) => rsx!( Badge { label: "victim".to_string(), tone: "warn", class: "".to_string() } ), - Some(VictimStatus::Squeezed) => rsx!( Badge { label: "squeezed".to_string(), tone: "neutral", class: "".to_string() } ), + Some(VictimStatus::Evicted) => rsx!( Badge { label: "evicted".to_string(), tone: "neutral", class: "".to_string() } ), None => rsx! {}, } if let Some(op) = state.current_operations.get(&entry.entry_id) { diff --git a/dev/dev-tools/src/trace/parser.rs b/dev/dev-tools/src/trace/parser.rs index 3512194c6..99f918f51 100644 --- a/dev/dev-tools/src/trace/parser.rs +++ b/dev/dev-tools/src/trace/parser.rs @@ -45,10 +45,10 @@ pub enum TraceEvent { entry: u64, kind: CacheKind, }, - SqueezeBegin { + EvictionBegin { victims: Vec, }, - SqueezeVictim { + EvictionVictim { entry: u64, }, IoWrite { @@ -108,10 +108,10 @@ impl TraceEvent { kind.display_name() ) } - TraceEvent::SqueezeBegin { victims } => { + TraceEvent::EvictionBegin { victims } => { format!("Begin squeeze (victims: {:?})", victims) } - TraceEvent::SqueezeVictim { entry } => { + TraceEvent::EvictionVictim { entry } => { format!("Squeeze victim {}", entry) } TraceEvent::IoWrite { entry, kind, bytes } => { @@ -195,8 +195,8 @@ impl TraceEvent { match self { TraceEvent::InsertSuccess { .. } => "insert_success", TraceEvent::InsertFailed { .. } => "insert_failed", - TraceEvent::SqueezeBegin { .. } => "squeeze_begin", - TraceEvent::SqueezeVictim { .. } => "squeeze_victim", + TraceEvent::EvictionBegin { .. } => "eviction_begin", + TraceEvent::EvictionVictim { .. } => "eviction_victim", TraceEvent::IoWrite { .. } => "io_write", TraceEvent::IoReadSqueezedBacking { .. } => "io_r_squeezed", TraceEvent::IoReadArrow { .. } => "io_read_arrow", @@ -285,13 +285,13 @@ fn parse_event_line(line: &str) -> TraceEvent { .unwrap_or(0), kind: fields.get("kind").map(|s| CacheKind::from_str(s)).unwrap(), }, - Some("squeeze_begin") => TraceEvent::SqueezeBegin { + Some("eviction_begin") => TraceEvent::EvictionBegin { victims: fields .get("victims") .map(|s| parse_victims(s)) .unwrap_or_default(), }, - Some("squeeze_victim") => TraceEvent::SqueezeVictim { + Some("eviction_victim") => TraceEvent::EvictionVictim { entry: fields .get("entry") .and_then(|s| s.parse().ok()) @@ -428,6 +428,7 @@ pub fn parse_trace(input: &str) -> Vec { && !line.starts_with("]") && !line.starts_with("---") && !line.starts_with("source:") + && !line.starts_with("assertion_line:") && !line.starts_with("expression:") }) .map(parse_event_line) diff --git a/dev/dev-tools/src/trace/simulator.rs b/dev/dev-tools/src/trace/simulator.rs index 04695f7fc..ea5f48cb8 100644 --- a/dev/dev-tools/src/trace/simulator.rs +++ b/dev/dev-tools/src/trace/simulator.rs @@ -53,8 +53,8 @@ pub enum EntryOperation { pub enum VictimStatus { /// Entry is selected as victim but not yet processed Selected, - /// Entry has been processed as victim (squeezed) - Squeezed, + /// Entry has been processed as victim (evicted) + Evicted, } /// Represents the complete cache state at a point in time @@ -63,9 +63,9 @@ pub struct CacheState { /// Map of entry_id -> CacheEntry pub entries: BTreeMap, /// Current squeeze victims being processed - pub squeeze_victims: Vec, + pub eviction_victims: Vec, /// Whether we're currently in a squeeze operation - pub in_squeeze: bool, + pub in_eviction: bool, /// Disk I/O statistics pub io_stats: IoStats, /// Delta changes in I/O stats for the current event @@ -82,8 +82,8 @@ impl CacheState { pub fn new() -> Self { Self { entries: BTreeMap::new(), - squeeze_victims: Vec::new(), - in_squeeze: false, + eviction_victims: Vec::new(), + in_eviction: false, io_stats: IoStats::default(), io_stats_delta: IoStatsDelta::default(), current_operations: BTreeMap::new(), @@ -213,9 +213,9 @@ impl CacheSimulator { // Track failed insert as ghost entry self.state.failed_inserts.insert(*entry, kind.clone()); } - TraceEvent::SqueezeBegin { victims } => { - self.state.in_squeeze = true; - self.state.squeeze_victims = victims.clone(); + TraceEvent::EvictionBegin { victims } => { + self.state.in_eviction = true; + self.state.eviction_victims = victims.clone(); // Mark all victims as selected for victim in victims { self.state @@ -223,17 +223,17 @@ impl CacheSimulator { .insert(*victim, VictimStatus::Selected); } } - TraceEvent::SqueezeVictim { entry } => { + TraceEvent::EvictionVictim { entry } => { // Remove from squeeze victims list - self.state.squeeze_victims.retain(|v| v != entry); - // Mark as squeezed + self.state.eviction_victims.retain(|v| v != entry); + // Mark as evicted self.state .victim_status - .insert(*entry, VictimStatus::Squeezed); + .insert(*entry, VictimStatus::Evicted); - // Check if squeeze is complete - if self.state.squeeze_victims.is_empty() { - self.state.in_squeeze = false; + // Check if eviction is complete + if self.state.eviction_victims.is_empty() { + self.state.in_eviction = false; } } TraceEvent::IoWrite { entry, bytes, .. } => { diff --git a/examples/core.rs b/examples/core.rs index d172f601a..1df164c7d 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use arrow::array::UInt64Array; use liquid_cache::cache::{ - AlwaysHydrate, EntryID, LiquidCacheBuilder, LiquidPolicy, TranscodeSqueezeEvict, + AlwaysHydrate, EntryID, LiquidCacheBuilder, LiquidPolicy, TranscodeEvict, }; #[tokio::main] @@ -12,7 +12,7 @@ async fn main() -> Result<(), Box> { .with_max_disk_bytes(1024 * 1024 * 1024 * 10) // 10GB .with_batch_size(8192) .with_cache_policy(Box::new(LiquidPolicy::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_hydration_policy(Box::new(AlwaysHydrate::new())) .build() .await; diff --git a/examples/datafusion-client-server.rs b/examples/datafusion-client-server.rs index 03b5f54b8..ec839399e 100644 --- a/examples/datafusion-client-server.rs +++ b/examples/datafusion-client-server.rs @@ -2,7 +2,7 @@ use arrow_flight::flight_service_server::FlightServiceServer; use clap::{Parser, ValueEnum}; use datafusion::{error::Result, execution::object_store::ObjectStoreUrl, prelude::*}; use liquid_cache_datafusion_client::LiquidCacheClientBuilder; -use liquid_cache_datafusion_local::storage::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache_datafusion_local::storage::cache::TranscodeEvict; use liquid_cache_datafusion_local::storage::cache::{AlwaysHydrate, LiquidPolicy}; use liquid_cache_datafusion_server::LiquidCacheService; use std::path::Path; @@ -57,7 +57,7 @@ async fn run_server() -> std::result::Result<(), Box> { Some(1024 * 1024 * 1024), // max memory size 1GB Some(tempfile::tempdir()?.keep()), // disk cache dir Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await?; diff --git a/examples/datafusion-local.rs b/examples/datafusion-local.rs index 1908ad810..46bff28f9 100644 --- a/examples/datafusion-local.rs +++ b/examples/datafusion-local.rs @@ -1,6 +1,6 @@ use datafusion::prelude::SessionConfig; use liquid_cache_datafusion_local::LiquidCacheLocalBuilder; -use liquid_cache_datafusion_local::storage::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache_datafusion_local::storage::cache::TranscodeEvict; use liquid_cache_datafusion_local::storage::cache_policies::LiquidPolicy; use tempfile::TempDir; @@ -12,7 +12,7 @@ async fn main() -> Result<(), Box> { .with_max_memory_bytes(1024 * 1024 * 1024) // 1GB .with_max_disk_bytes(1024 * 1024 * 1024 * 10) // 10GB .with_cache_dir(temp_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(LiquidPolicy::new())) .build(SessionConfig::new()) .await?; diff --git a/src/common/src/rpc.rs b/src/common/src/rpc.rs index 943c4d4a6..452d4ef56 100644 --- a/src/common/src/rpc.rs +++ b/src/common/src/rpc.rs @@ -65,15 +65,15 @@ impl From for LiquidCacheActions { } } -/// A typed squeeze hint for one file-schema column, shipped alongside a plan. +/// A typed lineage expression for one file-schema column, shipped alongside a plan. /// -/// The cache server cannot re-derive squeeze hints for lineage that lives only +/// The cache server cannot re-derive lineage expressions for lineage that lives only /// in the client-side part of the plan (e.g. a `date_part` projection above the /// pushed-down scan), so the client derives them from the full physical plan and /// ships them here. `hint` is the canonical encoding produced by /// `CacheExpression::to_metadata_value`. #[derive(Clone, PartialEq, ::prost::Message)] -pub struct ColumnSqueezeHint { +pub struct ColumnLineage { /// File-schema column name the hint applies to. #[prost(string, tag = "1")] pub column: ::prost::alloc::string::String, @@ -91,10 +91,10 @@ pub struct RegisterPlanRequest { #[prost(bytes, tag = "2")] pub handle: Bytes, - /// Typed squeeze hints for the (single) scan in this plan fragment, derived + /// Typed lineage expressions for the (single) scan in this plan fragment, derived /// by the client from the full physical plan. #[prost(message, repeated, tag = "3")] - pub squeeze_hints: ::prost::alloc::vec::Vec, + pub lineages: ::prost::alloc::vec::Vec, } impl ProstMessageExt for RegisterPlanRequest { diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml index 1a83b0e2f..b33a6963f 100644 --- a/src/core/Cargo.toml +++ b/src/core/Cargo.toml @@ -30,7 +30,6 @@ congee = { workspace = true } t4 = { workspace = true } log = { workspace = true } parquet = { workspace = true } -parquet-variant-compute = { workspace = true } fastrace = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } @@ -93,15 +92,6 @@ name = "linear_integer_study" path = "study/linear_integer.rs" harness = false -[[bench]] -name = "squeeze_integer_study" -path = "study/squeeze_integer.rs" -harness = false - -[[bench]] -name = "squeeze_date32" -path = "bench/squeeze_date32.rs" -harness = false [[bench]] name = "primitive_encoding" diff --git a/src/core/bench/squeeze_date32.rs b/src/core/bench/squeeze_date32.rs deleted file mode 100644 index 0ead85e6e..000000000 --- a/src/core/bench/squeeze_date32.rs +++ /dev/null @@ -1,94 +0,0 @@ -use arrow::array::{Array, ArrayRef, cast::AsArray}; -use arrow::datatypes::Date32Type; -use clap::Parser; -use datafusion::prelude::*; -use futures::StreamExt; -use liquid_cache::liquid_array::{Date32Field, LiquidPrimitiveArray, SqueezedDate32Array}; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Clone)] -#[command(name = "Squeeze Date32 Study")] -#[command(about = "Compare size of full Date32 vs squeezed YEAR/MONTH/DAY on TPCH lineitem")] -struct CliArgs { - /// Parquet file to read - #[arg( - long, - default_value = "../../benchmark/tpch/data/sf1.0/lineitem.parquet" - )] - parquet: String, - - /// Optional row limit for faster runs - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors - #[arg(long, default_value = "false")] - bench: bool, -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(8192 * 2); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("lineitem", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - let cols = ["l_commitdate", "l_receiptdate", "l_shipdate"]; - - for col in cols { - run_for_column(&ctx, col, args.limit).await; - } -} - -async fn run_for_column(ctx: &SessionContext, col: &str, limit: Option) { - let sql = if let Some(n) = limit { - format!("SELECT {} FROM lineitem LIMIT {n}", col) - } else { - format!("SELECT {} FROM lineitem", col) - }; - let df = ctx.sql(&sql).await.expect("create df"); - let mut stream = df.execute_stream().await.expect("execute stream"); - - let mut total_rows = 0usize; - let mut total_arrow_bytes = 0usize; - let mut total_liquid_bytes = 0usize; - let mut total_year_bytes = 0usize; - let mut total_month_bytes = 0usize; - let mut total_day_bytes = 0usize; - let mut total_dow_bytes = 0usize; - - while let Some(batch_res) = stream.next().await { - let batch = batch_res.expect("stream batch"); - let arr: ArrayRef = batch.column(0).clone(); - assert_eq!(arr.data_type(), &arrow_schema::DataType::Date32); - - total_rows += arr.len(); - total_arrow_bytes += arr.get_array_memory_size(); - - let prim = arr.as_primitive::().clone(); - let liquid = LiquidPrimitiveArray::::from_arrow_array(prim.clone()); - total_liquid_bytes += liquid.get_array_memory_size(); - - let squeezed_year = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Year); - let squeezed_month = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Month); - let squeezed_day = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Day); - let squeezed_dow = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::DayOfWeek); - - total_year_bytes += squeezed_year.get_array_memory_size(); - total_month_bytes += squeezed_month.get_array_memory_size(); - total_day_bytes += squeezed_day.get_array_memory_size(); - total_dow_bytes += squeezed_dow.get_array_memory_size(); - } - - println!( - "Column {col} on {total_rows} rows:\n Arrow(Date32): {total_arrow_bytes} bytes\n Liquid(Date32): {total_liquid_bytes} bytes\n Squeezed YEAR: {total_year_bytes} bytes\n Squeezed MONTH: {total_month_bytes} bytes\n Squeezed DAY: {total_day_bytes} bytes\n Squeezed DOW: {total_dow_bytes} bytes" - ); -} diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 6881f0c46..9bb4a8430 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -9,7 +9,7 @@ use arrow::buffer::BooleanBuffer; use super::cached_batch::CacheEntry; use super::core::LiquidCache; use super::io_context::{DefaultCacheMetadata, EntryMetadata}; -use super::policies::{CachePolicy, HydrationPolicy, SqueezePolicy, TranscodeSqueezeEvict}; +use super::policies::{CachePolicy, EvictionPolicy, HydrationPolicy, TranscodeEvict}; use super::{CacheExpression, CacheFull, EntryID, LiquidExpr, LiquidPolicy}; use crate::sync::Arc; @@ -35,10 +35,10 @@ pub struct LiquidCacheBuilder { max_disk_bytes: usize, cache_policy: Box, hydration_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, metadata: Option>, store: Option, - squeeze_victims_concurrently: bool, + evict_victims_concurrently: bool, } impl Default for LiquidCacheBuilder { @@ -58,10 +58,10 @@ impl LiquidCacheBuilder { max_disk_bytes, cache_policy: Box::new(LiquidPolicy::new()), hydration_policy: Box::new(super::AlwaysHydrate::new()), - squeeze_policy: Box::new(TranscodeSqueezeEvict), + eviction_policy: Box::new(TranscodeEvict), metadata: None, store: None, - squeeze_victims_concurrently: !cfg!(test), + evict_victims_concurrently: !cfg!(test), } } @@ -100,10 +100,10 @@ impl LiquidCacheBuilder { self } - /// Set the squeeze policy for the cache. - /// Default is [TranscodeSqueezeEvict]. - pub fn with_squeeze_policy(mut self, policy: Box) -> Self { - self.squeeze_policy = policy; + /// Set the eviction policy for the cache. + /// Default is [TranscodeEvict]. + pub fn with_eviction_policy(mut self, policy: Box) -> Self { + self.eviction_policy = policy; self } @@ -121,9 +121,9 @@ impl LiquidCacheBuilder { self } - /// Set whether cache victims are squeezed concurrently. - pub fn with_squeeze_victims_concurrently(mut self, enabled: bool) -> Self { - self.squeeze_victims_concurrently = enabled; + /// Set whether cache victims are evicted concurrently. + pub fn with_evict_victims_concurrently(mut self, enabled: bool) -> Self { + self.evict_victims_concurrently = enabled; self } @@ -149,12 +149,12 @@ impl LiquidCacheBuilder { self.batch_size, self.max_memory_bytes, self.max_disk_bytes, - self.squeeze_policy, + self.eviction_policy, self.cache_policy, self.hydration_policy, metadata, store, - self.squeeze_victims_concurrently, + self.evict_victims_concurrently, )) } } @@ -177,7 +177,7 @@ pub struct Insert<'a> { pub(super) entry_id: EntryID, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, - pub(super) squeeze_hint: Option>, + pub(super) lineage: Option>, } impl<'a> Insert<'a> { @@ -187,7 +187,7 @@ impl<'a> Insert<'a> { entry_id, batch, skip_gc: false, - squeeze_hint: None, + lineage: None, } } @@ -197,9 +197,9 @@ impl<'a> Insert<'a> { self } - /// Set a squeeze hint for the entry. - pub fn with_squeeze_hint(mut self, expression: Arc) -> Self { - self.squeeze_hint = Some(expression); + /// Set a lineage expression for the entry. + pub fn with_lineage(mut self, expression: Arc) -> Self { + self.lineage = Some(expression); self } @@ -209,8 +209,8 @@ impl<'a> Insert<'a> { } else { maybe_gc_view_arrays(&self.batch).unwrap_or_else(|| self.batch.clone()) }; - if let Some(squeeze_hint) = self.squeeze_hint { - self.storage.add_squeeze_hint(&self.entry_id, squeeze_hint); + if let Some(lineage) = self.lineage { + self.storage.add_lineage(&self.entry_id, lineage); } let batch = CacheEntry::memory_arrow(batch); self.storage.insert_inner(self.entry_id, batch).await diff --git a/src/core/src/cache/cached_batch.rs b/src/core/src/cache/cached_batch.rs index fda9950a5..d107be5f0 100644 --- a/src/core/src/cache/cached_batch.rs +++ b/src/core/src/cache/cached_batch.rs @@ -5,7 +5,7 @@ use std::{fmt::Display, sync::Arc}; use arrow::array::ArrayRef; use arrow_schema::DataType; -use crate::liquid_array::{LiquidArrayRef, LiquidSqueezedArrayRef}; +use crate::liquid_array::LiquidArrayRef; /// A cached entry storing data in various formats. #[derive(Debug, Clone)] @@ -14,8 +14,6 @@ pub enum CacheEntry { MemoryArrow(ArrayRef), /// Cached batch in memory as liquid array. MemoryLiquid(LiquidArrayRef), - /// Cached batch in memory as squeezed liquid array. - MemorySqueezedLiquid(LiquidSqueezedArrayRef), /// Cached batch on disk as liquid array. DiskLiquid { /// Original Arrow data type. @@ -43,11 +41,6 @@ impl CacheEntry { Self::MemoryLiquid(array) } - /// Construct a cached batch stored as an in-memory squeezed Liquid array. - pub fn memory_squeezed_liquid(array: LiquidSqueezedArrayRef) -> Self { - Self::MemorySqueezedLiquid(array) - } - /// Construct a cached batch stored on disk as Liquid bytes. pub fn disk_liquid(data_type: DataType, disk_bytes: usize) -> Self { Self::DiskLiquid { @@ -69,7 +62,6 @@ impl CacheEntry { match self { Self::MemoryArrow(array) => array.get_array_memory_size(), Self::MemoryLiquid(array) => array.get_array_memory_size(), - Self::MemorySqueezedLiquid(array) => array.get_array_memory_size(), Self::DiskLiquid { .. } | Self::DiskArrow { .. } => 0, } } @@ -79,7 +71,6 @@ impl CacheEntry { match self { Self::MemoryArrow(array) => Arc::strong_count(array), Self::MemoryLiquid(array) => Arc::strong_count(array), - Self::MemorySqueezedLiquid(array) => Arc::strong_count(array), Self::DiskLiquid { .. } | Self::DiskArrow { .. } => 0, } } @@ -90,7 +81,6 @@ impl Display for CacheEntry { match self { Self::MemoryArrow(_) => write!(f, "MemoryArrow"), Self::MemoryLiquid(_) => write!(f, "MemoryLiquid"), - Self::MemorySqueezedLiquid(_) => write!(f, "MemorySqueezedLiquid"), Self::DiskLiquid { .. } => write!(f, "DiskLiquid"), Self::DiskArrow { .. } => write!(f, "DiskArrow"), } @@ -104,8 +94,6 @@ pub enum CachedBatchType { MemoryArrow, /// Cached batch in memory as liquid array. MemoryLiquid, - /// Cached batch in memory as squeezed liquid array. - MemorySqueezedLiquid, /// Cached batch on disk as liquid array. DiskLiquid, /// Cached batch on disk as Arrow array. @@ -117,7 +105,6 @@ impl From<&CacheEntry> for CachedBatchType { match batch { CacheEntry::MemoryArrow(_) => Self::MemoryArrow, CacheEntry::MemoryLiquid(_) => Self::MemoryLiquid, - CacheEntry::MemorySqueezedLiquid(_) => Self::MemorySqueezedLiquid, CacheEntry::DiskLiquid { .. } => Self::DiskLiquid, CacheEntry::DiskArrow { .. } => Self::DiskArrow, } diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 4fc02f9f7..26859aa34 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -15,15 +15,10 @@ use super::{ policies::{CachePolicy, HydrationPolicy, HydrationRequest, MaterializedEntry}, utils::CacheConfig, }; -use crate::cache::DefaultSqueezeIo; -use crate::cache::policies::{SqueezeOutcome, SqueezePolicy}; +use crate::cache::policies::{EvictionOutcome, EvictionPolicy}; use crate::cache::utils::{LiquidCompressorStates, arrow_to_bytes}; use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; use crate::cache::{CacheFull, CacheStats, EventTrace}; -use crate::liquid_array::{ - LiquidSqueezedArrayRef, SqueezeIoHandler, SqueezedBacking, SqueezedDate32Array, - VariantStructSqueezedArray, -}; use crate::sync::Arc; // CacheStats and RuntimeStats moved to stats.rs @@ -55,19 +50,17 @@ pub struct LiquidCache { budget: BudgetAccounting, cache_policy: Box, hydration_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, observer: Arc, metadata: Arc, store: t4::Store, - squeeze_victims_concurrently: bool, + evict_victims_concurrently: bool, } /// Outcome of [`LiquidCache::prefetch`]. pub enum PrefetchResult { /// A memory-form snapshot of the entry (Arrow or Liquid), ready to hand to a reader. Snapshot(Arc), - /// The entry is squeezed; prefetch leaves it alone. - Squeezed, /// The entry is not in the index, or its disk blob is gone. Absent, } @@ -81,13 +74,11 @@ impl LiquidCache { let mut memory_arrow_entries = 0usize; let mut memory_liquid_entries = 0usize; - let mut memory_squeezed_liquid_entries = 0usize; let mut disk_liquid_entries = 0usize; let mut disk_arrow_entries = 0usize; let mut memory_arrow_bytes = 0usize; let mut memory_liquid_bytes = 0usize; - let mut memory_squeezed_liquid_bytes = 0usize; self.index.for_each(|_, batch| match batch { CacheEntry::MemoryArrow(array) => { @@ -98,10 +89,6 @@ impl LiquidCache { memory_liquid_entries += 1; memory_liquid_bytes += array.get_array_memory_size(); } - CacheEntry::MemorySqueezedLiquid(array) => { - memory_squeezed_liquid_entries += 1; - memory_squeezed_liquid_bytes += array.get_array_memory_size(); - } CacheEntry::DiskLiquid { .. } => disk_liquid_entries += 1, CacheEntry::DiskArrow { .. } => disk_arrow_entries += 1, }); @@ -114,12 +101,10 @@ impl LiquidCache { total_entries, memory_arrow_entries, memory_liquid_entries, - memory_squeezed_liquid_entries, disk_liquid_entries, disk_arrow_entries, memory_arrow_bytes, memory_liquid_bytes, - memory_squeezed_liquid_bytes, memory_usage_bytes, disk_usage_bytes, max_memory_bytes: self.config.max_memory_bytes(), @@ -176,7 +161,6 @@ impl LiquidCache { .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_liquid(array))) } - CacheEntry::MemorySqueezedLiquid(_) => PrefetchResult::Squeezed, } } @@ -200,13 +184,6 @@ impl LiquidCache { .await; Some(liquid) } - CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { - SqueezedBacking::Liquid(_) => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - Some(liquid) - } - SqueezedBacking::Arrow(_) => None, - }, CacheEntry::DiskArrow { .. } | CacheEntry::MemoryArrow(_) => None, } } @@ -254,9 +231,9 @@ impl LiquidCache { self.metadata.get_compressor(entry_id) } - /// Add a squeeze hint for an entry. - pub fn add_squeeze_hint(&self, entry_id: &EntryID, expression: Arc) { - self.metadata.add_squeeze_hint(entry_id, expression); + /// Add a lineage expression for an entry. + pub fn add_lineage(&self, entry_id: &EntryID, expression: Arc) { + self.metadata.add_lineage(entry_id, expression); } /// Flush all entries to disk. @@ -301,12 +278,6 @@ impl LiquidCache { Err(CacheFull) => self.drop_memory_entry(entry_id, &batch), } } - CacheEntry::MemorySqueezedLiquid(array) => { - // We don't have to do anything, because it's already on disk - let disk_entry = Self::disk_entry_from_squeezed(array); - self.try_insert(entry_id, disk_entry) - .expect("failed to insert disk entry"); - } CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { // Already on disk, skip } @@ -325,23 +296,17 @@ impl LiquidCache { ) -> Result { match &batch { batch @ CacheEntry::MemoryArrow(_) => { - let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( - self.store.clone(), - entry_id, - self.observer.clone(), - )); - let outcome = self.squeeze_policy.squeeze( + let outcome = self.eviction_policy.evict( batch, self.metadata.get_compressor(&entry_id).as_ref(), None, - &squeeze_io, ); - let SqueezeOutcome::Replace { + let EvictionOutcome::Replace { entry: new_batch, bytes_to_write, } = outcome else { - unreachable!("memory arrow squeeze cannot remove entry"); + unreachable!("memory Arrow eviction cannot remove entry"); }; if let Some(bytes_to_write) = bytes_to_write { self.write_batch_to_disk(entry_id, &new_batch, bytes_to_write) @@ -359,15 +324,6 @@ impl LiquidCache { disk_bytes, )) } - CacheEntry::MemorySqueezedLiquid(squeezed_array) => { - // The full data is already on disk, so we just need to mark ourself as disk entry - let data_type = squeezed_array.original_arrow_data_type(); - let entry = match squeezed_array.disk_backing() { - SqueezedBacking::Liquid(n) => CacheEntry::disk_liquid(data_type, n), - SqueezedBacking::Arrow(n) => CacheEntry::disk_arrow(data_type, n), - }; - Ok(entry) - } CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } => { unreachable!("Unexpected batch in write_in_memory_batch_to_disk") } @@ -400,7 +356,7 @@ impl LiquidCache { batch_to_cache = on_disk_batch; continue; } - self.squeeze_victims(victims).await?; + self.evict_victims(victims).await?; batch_to_cache = not_inserted; crate::utils::yield_now_if_shuttle(); @@ -413,12 +369,12 @@ impl LiquidCache { batch_size: usize, max_memory_bytes: usize, max_disk_bytes: usize, - squeeze_policy: Box, + eviction_policy: Box, cache_policy: Box, hydration_policy: Box, metadata: Arc, store: t4::Store, - squeeze_victims_concurrently: bool, + evict_victims_concurrently: bool, ) -> Self { let config = CacheConfig::new(batch_size, max_memory_bytes, max_disk_bytes); let observer = Arc::new(Observer::new()); @@ -432,11 +388,11 @@ impl LiquidCache { config, cache_policy, hydration_policy, - squeeze_policy, + eviction_policy, observer, metadata, store, - squeeze_victims_concurrently, + evict_victims_concurrently, } } @@ -480,9 +436,7 @@ impl LiquidCache { assert!( matches!( removed.as_ref(), - CacheEntry::MemoryArrow(_) - | CacheEntry::MemoryLiquid(_) - | CacheEntry::MemorySqueezedLiquid(_) + CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) ), "flush should only drop memory entries" ); @@ -529,67 +483,58 @@ impl LiquidCache { } #[fastrace::trace] - async fn squeeze_victims(&self, victims: Vec) -> Result<(), CacheFull> { - self.trace(InternalEvent::SqueezeBegin { + async fn evict_victims(&self, victims: Vec) -> Result<(), CacheFull> { + self.trace(InternalEvent::EvictionBegin { victims: victims.clone(), }); - if self.squeeze_victims_concurrently { + if self.evict_victims_concurrently { let results = futures::stream::iter(victims) - .map(|victim| self.squeeze_victim_inner(victim)) + .map(|victim| self.evict_victim_inner(victim)) .buffer_unordered(usize::MAX) .collect::>() .await; results.into_iter().collect::, _>>()?; } else { for victim in victims { - self.squeeze_victim_inner(victim).await?; + self.evict_victim_inner(victim).await?; } } Ok(()) } - async fn squeeze_victim_inner(&self, to_squeeze: EntryID) -> Result<(), CacheFull> { - let Some(mut to_squeeze_batch) = self.index.get(&to_squeeze) else { + async fn evict_victim_inner(&self, victim: EntryID) -> Result<(), CacheFull> { + let Some(mut victim_entry) = self.index.get(&victim) else { return Ok(()); }; - self.trace(InternalEvent::SqueezeVictim { entry: to_squeeze }); - let compressor = self.metadata.get_compressor(&to_squeeze); - let squeeze_hint_arc = self.metadata.squeeze_hint(&to_squeeze); - let squeeze_hint = squeeze_hint_arc.as_deref(); - let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( - self.store.clone(), - to_squeeze, - self.observer.clone(), - )); - + self.trace(InternalEvent::EvictionVictim { entry: victim }); + let compressor = self.metadata.get_compressor(&victim); + let lineage_arc = self.metadata.lineage(&victim); + let lineage = lineage_arc.as_deref(); loop { - let outcome = self.squeeze_policy.squeeze( - to_squeeze_batch.as_ref(), - compressor.as_ref(), - squeeze_hint, - &squeeze_io, - ); + let outcome = + self.eviction_policy + .evict(victim_entry.as_ref(), compressor.as_ref(), lineage); match outcome { - SqueezeOutcome::Replace { + EvictionOutcome::Replace { entry: new_batch, bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) + self.write_batch_to_disk(victim, &new_batch, bytes_to_write) .await?; } - match self.try_insert(to_squeeze, new_batch) { + match self.try_insert(victim, new_batch) { Ok(()) => { break; } Err(batch) => { - to_squeeze_batch = Arc::new(batch); + victim_entry = Arc::new(batch); } } } - SqueezeOutcome::Remove => { - self.remove_disk_entry(to_squeeze).await; + EvictionOutcome::Remove => { + self.remove_disk_entry(victim).await; break; } } @@ -597,14 +542,6 @@ impl LiquidCache { Ok(()) } - fn disk_entry_from_squeezed(array: &LiquidSqueezedArrayRef) -> CacheEntry { - let data_type = array.original_arrow_data_type(); - match array.disk_backing() { - SqueezedBacking::Liquid(n) => CacheEntry::disk_liquid(data_type, n), - SqueezedBacking::Arrow(n) => CacheEntry::disk_arrow(data_type, n), - } - } - async fn maybe_hydrate( &self, entry_id: &EntryID, @@ -612,13 +549,11 @@ impl LiquidCache { materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, ) { - let compressor = self.metadata.get_compressor(entry_id); if let Some(new_entry) = self.hydration_policy.hydrate(&HydrationRequest { entry_id: *entry_id, cached, materialized, expression, - compressor, }) { let cached_type = CachedBatchType::from(cached); let new_type = CachedBatchType::from(&new_entry); @@ -690,10 +625,6 @@ impl LiquidCache { self.read_disk_array(entry, entry_id, expression, selection) .await } - CacheEntry::MemorySqueezedLiquid(array) => { - self.read_squeezed_array(array, entry_id, expression, selection) - .await - } } } @@ -750,107 +681,6 @@ impl LiquidCache { } } - async fn read_squeezed_array( - &self, - array: &LiquidSqueezedArrayRef, - entry_id: &EntryID, - expression: Option<&CacheExpression>, - selection: Option<&BooleanBuffer>, - ) -> Option { - if let Some(array) = self.try_read_squeezed_date32_array(array, expression, selection) { - self.observer.on_get_squeezed_success(); - self.trace(InternalEvent::ReadSqueezedData { - entry: *entry_id, - expression: expression.unwrap().clone(), - }); - return Some(array); - } - - if let Some(array) = self - .try_read_squeezed_variant_array(array, entry_id, expression, selection) - .await - { - self.observer.on_get_squeezed_success(); - self.trace(InternalEvent::ReadSqueezedData { - entry: *entry_id, - expression: expression.unwrap().clone(), - }); - return Some(array); - } - - // no shortcut, needs to read full data - let out = match selection { - Some(selection) => array.filter(selection).await, - None => array.to_arrow_array().await, - }; - Some(out) - } - - fn try_read_squeezed_date32_array( - &self, - array: &LiquidSqueezedArrayRef, - expression: Option<&CacheExpression>, - selection: Option<&BooleanBuffer>, - ) -> Option { - if let Some(field) = expression.and_then(CacheExpression::as_date32_field) - && let Some(squeezed) = array.as_any().downcast_ref::() - && squeezed.field() == field - { - let component = squeezed.to_component_array(); - self.observer.on_hit_date32_expression(); - if let Some(selection) = selection { - let selection_array = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::filter(&component, &selection_array).ok()?; - return Some(filtered); - } - return Some(component); - } - None - } - - async fn try_read_squeezed_variant_array( - &self, - array: &LiquidSqueezedArrayRef, - entry_id: &EntryID, - expression: Option<&CacheExpression>, - selection: Option<&BooleanBuffer>, - ) -> Option { - let requests = expression.and_then(|expr| expr.variant_requests())?; - let variant_squeezed = array - .as_any() - .downcast_ref::()?; - let all_paths_present = requests - .iter() - .all(|request| variant_squeezed.contains_path(request.path())); - - let full_array = if !all_paths_present { - let batch = CacheEntry::MemorySqueezedLiquid(array.clone()); - self.observer.on_get_squeezed_needs_io(); - let full_array = self.read_disk_arrow_array(entry_id).await?; - self.maybe_hydrate( - entry_id, - &batch, - MaterializedEntry::Arrow(&full_array), - expression, - ) - .await; - full_array - } else { - let requested_paths = requests.iter().map(|r| r.path()); - variant_squeezed - .to_arrow_array_with_paths(requested_paths) - .unwrap() - }; - - match selection { - Some(selection) => { - let selection_array = BooleanArray::new(selection.clone(), None); - arrow::compute::filter(&full_array, &selection_array).ok() - } - None => Some(full_array), - } - } - async fn write_batch_to_disk( &self, entry_id: EntryID, @@ -1010,27 +840,9 @@ impl LiquidCache { }); Some(liquid.try_eval_predicate(predicate, selection)) } - CacheEntry::MemorySqueezedLiquid(array) => { - self.eval_predicate_on_squeezed(array, selection_opt, predicate) - .await - } } } - async fn eval_predicate_on_squeezed( - &self, - array: &LiquidSqueezedArrayRef, - selection_opt: Option<&BooleanBuffer>, - predicate: &LiquidExpr, - ) -> Option { - let mut owned = None; - let selection = selection_opt.unwrap_or_else(|| { - owned = Some(BooleanBuffer::new_set(array.len())); - owned.as_ref().unwrap() - }); - Some(array.try_eval_predicate(predicate, selection).await) - } - fn eval_predicate_on_array(&self, array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", @@ -1054,19 +866,15 @@ impl LiquidCache { mod tests { use super::*; use crate::cache::{ - CacheEntry, CacheExpression, CachePolicy, LiquidCacheBuilder, LiquidPolicy, - TranscodeSqueezeEvict, transcode_liquid_inner, + CacheEntry, CachePolicy, LiquidCacheBuilder, LiquidPolicy, TranscodeEvict, + transcode_liquid_inner, utils::{ LiquidCompressorStates, arrow_to_bytes, create_cache_store, create_test_array, create_test_arrow_array, }, }; - use crate::liquid_array::{ - Date32Field, LiquidPrimitiveArray, LiquidSqueezedArrayRef, SqueezedDate32Array, - }; use crate::sync::thread; - use arrow::array::{Array, ArrayRef, Date32Array, Int32Array}; - use arrow::datatypes::Date32Type; + use arrow::array::{Array, ArrayRef, Int32Array}; use std::future::Future; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1133,43 +941,6 @@ mod tests { assert!(store.index().get(&EntryID::from(999)).is_none()); } - #[tokio::test] - async fn get_arrow_array_with_expression_extracts_year() { - let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; - let entry_id = EntryID::from(42); - - let date_values = Date32Array::from(vec![Some(2), Some(365 + 1), None, Some(365 + 100)]); - let liquid = LiquidPrimitiveArray::::from_arrow_array(date_values.clone()); - let squeezed = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Year); - let squeezed: LiquidSqueezedArrayRef = Arc::new(squeezed); - - store - .insert_inner( - entry_id, - CacheEntry::memory_squeezed_liquid(squeezed.clone()), - ) - .await - .unwrap(); - - let expr = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); - let result = store - .get(&entry_id) - .with_expression_hint(expr) - .read() - .await - .expect("array present"); - - let result = result - .as_any() - .downcast_ref::() - .expect("date32 result"); - assert_eq!(result.len(), 4); - assert_eq!(result.value(0), 0); - assert_eq!(result.value(1), 365); - assert!(result.is_null(2)); - assert_eq!(result.value(3), 365); - } - #[tokio::test] async fn test_cache_advice_strategies() { // Comprehensive test of all three advice types @@ -1270,7 +1041,7 @@ mod tests { // Build a small cache in blocking liquid mode to avoid background tasks let storage = LiquidCacheBuilder::new() .with_max_memory_bytes(10 * 1024 * 1024) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build() .await; @@ -1365,7 +1136,7 @@ mod tests { let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(0) .with_max_disk_bytes(0) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build() .await; let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); @@ -1385,7 +1156,7 @@ mod tests { let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(first_bytes.max(second_bytes)) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(LiquidPolicy::new())) .build() .await; @@ -1414,7 +1185,7 @@ mod tests { let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(disk_bytes) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(LiquidPolicy::new())) .build() .await; @@ -1436,7 +1207,7 @@ mod tests { let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(disk_bytes) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(LiquidPolicy::new())) .build() .await; @@ -1456,7 +1227,7 @@ mod tests { let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(0) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build() .await; let entry_id = EntryID::from(901usize); diff --git a/src/core/src/cache/expressions.rs b/src/core/src/cache/expressions.rs index 64e7d781b..a1deb122c 100644 --- a/src/core/src/cache/expressions.rs +++ b/src/core/src/cache/expressions.rs @@ -36,10 +36,9 @@ impl VariantRequest { /// Experimental expression descriptor for cache lookups. /// -/// A `CacheExpression` is a *squeeze hint*: it tells the cache how a column is -/// consumed by a query so that, under memory pressure, the cache can keep only -/// the part of the column the query actually needs (e.g. a single date -/// component, or a handful of variant paths) instead of evicting it wholesale. +/// A `CacheExpression` records how a query consumes a column. The lineage is +/// retained even when a cache policy does not use it, allowing representation +/// decisions to evolve independently from query analysis. #[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize)] pub enum CacheExpression { /// Extract one or more components (YEAR/MONTH/DAY/DOW) from a `Date32` or @@ -148,7 +147,7 @@ impl CacheExpression { /// expression for exactly one component. /// /// Multi-component extractions return `None`: there is no single-component - /// squeezed representation that satisfies all of them, so the squeeze path + /// partial representation that satisfies all of them, so the representation-selection path /// keeps the column intact rather than dropping a needed component. pub fn as_date32_field(&self) -> Option { match self { @@ -306,7 +305,7 @@ mod tests { let encoded = expr.to_metadata_value(); let decoded = CacheExpression::from_metadata_value(&encoded).unwrap(); assert_eq!(decoded, expr); - // Multi-component extractions do not collapse to a single squeezable field. + // Multi-component extractions do not collapse to a single specialized field. assert_eq!(decoded.as_date32_field(), None); assert_eq!( decoded.date32_fields().unwrap(), diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index a9c8974fe..9c5e0a179 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -1,36 +1,29 @@ -use std::{fmt::Debug, ops::Range}; +use std::fmt::Debug; use ahash::AHashMap; -use bytes::Bytes; -use crate::sync::{Arc, RwLock}; -use crate::{ - cache::{ - CacheExpression, Observer, - observer::InternalEvent, - utils::{EntryID, LiquidCompressorStates}, - }, - liquid_array::SqueezeIoHandler, +use crate::cache::{ + CacheExpression, + utils::{EntryID, LiquidCompressorStates}, }; +use crate::sync::{Arc, RwLock}; /// Per-entry metadata used by the cache. /// /// This trait covers only the metadata side of the cache: where to find a -/// batch's compressor and squeeze hints. All actual byte IO goes through the +/// batch's compressor and lineage expressions. All actual byte IO goes through the /// [`t4::Store`] held by the cache itself. pub trait EntryMetadata: Debug + Send + Sync { - /// Add a squeeze hint for an entry. - fn add_squeeze_hint(&self, _entry_id: &EntryID, _expression: Arc) { + /// Add a lineage expression for an entry. + fn add_lineage(&self, _entry_id: &EntryID, _expression: Arc) { // Do nothing by default } - /// Get the squeeze hint for an entry. + /// Get the lineage expression for an entry. /// If None, the entry will be evicted to disk entirely. - /// If Some, the entry will be squeezed according to the cache expressions previously recorded for this column. - /// For example, if expression is `ExtractDate32 { fields: [Date32Field::Year] }`, - /// the entry will be squeezed to a [crate::liquid_array::SqueezedDate32Array] with the year - /// component (Date32 or Timestamp input). - fn squeeze_hint(&self, _entry_id: &EntryID) -> Option> { + /// The expression records how the column is used by query plans and may inform + /// encoding decisions without discarding any values. + fn lineage(&self, _entry_id: &EntryID) -> Option> { None } @@ -45,12 +38,12 @@ pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { /// A default implementation of [`EntryMetadata`]. /// -/// All entries share a single [`LiquidCompressorStates`] and squeeze hints are +/// All entries share a single [`LiquidCompressorStates`] and lineage expressions are /// stored in a flat map keyed by [`EntryID`]. #[derive(Debug, Default)] pub struct DefaultCacheMetadata { compressor_state: Arc, - squeeze_hints: RwLock>>, + lineages: RwLock>>, } impl DefaultCacheMetadata { @@ -58,19 +51,19 @@ impl DefaultCacheMetadata { pub fn new() -> Self { Self { compressor_state: Arc::new(LiquidCompressorStates::new()), - squeeze_hints: RwLock::new(AHashMap::new()), + lineages: RwLock::new(AHashMap::new()), } } } impl EntryMetadata for DefaultCacheMetadata { - fn add_squeeze_hint(&self, entry_id: &EntryID, expression: Arc) { - let mut guard = self.squeeze_hints.write().unwrap(); + fn add_lineage(&self, entry_id: &EntryID, expression: Arc) { + let mut guard = self.lineages.write().unwrap(); guard.insert(*entry_id, expression); } - fn squeeze_hint(&self, entry_id: &EntryID) -> Option> { - let guard = self.squeeze_hints.read().unwrap(); + fn lineage(&self, entry_id: &EntryID) -> Option> { + let guard = self.lineages.read().unwrap(); guard.get(entry_id).cloned() } @@ -78,103 +71,3 @@ impl EntryMetadata for DefaultCacheMetadata { self.compressor_state.clone() } } - -/// A default implementation of [SqueezeIoHandler] backed by a [`t4::Store`]. -#[derive(Debug)] -pub struct DefaultSqueezeIo { - store: t4::Store, - entry_id: EntryID, - observer: Arc, -} - -impl DefaultSqueezeIo { - /// Create a new instance of [DefaultSqueezeIo]. - pub fn new(store: t4::Store, entry_id: EntryID, observer: Arc) -> Self { - Self { - store, - entry_id, - observer, - } - } -} - -#[async_trait::async_trait] -impl SqueezeIoHandler for DefaultSqueezeIo { - async fn read(&self, range: Option>) -> std::io::Result { - let key = entry_id_to_key(&self.entry_id); - let bytes = match range { - Some(range) => { - let len = range.end - range.start; - self.store - .get_range(&key, range.start, len) - .await - .map_err(|e| std::io::Error::other(e.to_string()))? - } - None => self - .store - .get(&key) - .await - .map_err(|e| std::io::Error::other(e.to_string()))?, - }; - let bytes = Bytes::from(bytes); - self.observer - .record_internal(InternalEvent::IoReadSqueezedBacking { - entry: self.entry_id, - bytes: bytes.len(), - }); - Ok(bytes) - } - - fn tracing_decompress_count(&self, decompress_cnt: usize, total_cnt: usize) { - self.observer - .record_internal(InternalEvent::DecompressSqueezed { - entry: self.entry_id, - decompressed: decompress_cnt, - total: total_cnt, - }); - } - - fn trace_io_saved(&self) { - self.observer.runtime_stats().incr_squeeze_io_saved(); - } -} - -#[cfg(test)] -#[derive(Debug, Default)] -pub(crate) struct TestSqueezeIo { - bytes: std::sync::Mutex>, - reads: std::sync::atomic::AtomicUsize, -} - -#[cfg(test)] -impl TestSqueezeIo { - pub(crate) fn set_bytes(&self, bytes: Bytes) { - *self.bytes.lock().unwrap() = Some(bytes); - } - - pub(crate) fn reads(&self) -> usize { - self.reads.load(std::sync::atomic::Ordering::SeqCst) - } - - pub(crate) fn reset_reads(&self) { - self.reads.store(0, std::sync::atomic::Ordering::SeqCst); - } -} - -#[cfg(test)] -#[async_trait::async_trait] -impl SqueezeIoHandler for TestSqueezeIo { - async fn read(&self, range: Option>) -> std::io::Result { - self.reads.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - let bytes = self - .bytes - .lock() - .unwrap() - .clone() - .expect("test squeeze backing set"); - Ok(match range { - Some(range) => bytes.slice(range.start as usize..range.end as usize), - None => bytes, - }) - } -} diff --git a/src/core/src/cache/mod.rs b/src/core/src/cache/mod.rs index 19ae5dc13..4783322c6 100644 --- a/src/core/src/cache/mod.rs +++ b/src/core/src/cache/mod.rs @@ -17,16 +17,14 @@ pub use builders::{EvaluatePredicate, Get, Insert, LiquidCacheBuilder, default_m pub use cached_batch::{CacheEntry, CachedBatchType}; pub use core::{LiquidCache, PrefetchResult}; pub use expressions::{CacheExpression, VariantRequest}; -#[cfg(test)] -pub(crate) use io_context::TestSqueezeIo; -pub use io_context::{DefaultCacheMetadata, DefaultSqueezeIo, EntryMetadata}; +pub use io_context::{DefaultCacheMetadata, EntryMetadata}; pub use liquid_expr::LiquidExpr; pub use observer::EventTrace; pub use observer::Observer; pub use observer::{CacheStats, RuntimeStats, RuntimeStatsSnapshot}; pub use policies::{ - AlwaysHydrate, CachePolicy, HydrationPolicy, HydrationRequest, LiquidPolicy, MaterializedEntry, - NoHydration, SqueezePolicy, TranscodeSqueezeEvict, + AlwaysHydrate, CachePolicy, Evict, EvictionPolicy, HydrationPolicy, HydrationRequest, + LiquidPolicy, MaterializedEntry, NoHydration, TranscodeEvict, }; pub use transcode::{transcode_liquid_inner, transcode_liquid_inner_with_hint}; pub use utils::{EntryID, LiquidCompressorStates}; @@ -46,10 +44,5 @@ pub mod hydration_policies { pub use super::policies::hydration::*; } -/// Legacy path: re-export squeeze policy types under `cache::squeeze_policies`. -pub mod squeeze_policies { - pub use super::policies::squeeze::*; -} - #[cfg(test)] mod tests; diff --git a/src/core/src/cache/observer/internal_tracing.rs b/src/core/src/cache/observer/internal_tracing.rs index 21c4a5343..26c2a7c27 100644 --- a/src/core/src/cache/observer/internal_tracing.rs +++ b/src/core/src/cache/observer/internal_tracing.rs @@ -13,10 +13,10 @@ pub(crate) enum InternalEvent { entry: EntryID, kind: CachedBatchType, }, - SqueezeBegin { + EvictionBegin { victims: Vec, }, - SqueezeVictim { + EvictionVictim { entry: EntryID, }, IoWrite { @@ -28,10 +28,6 @@ pub(crate) enum InternalEvent { entry: EntryID, bytes: usize, }, - IoReadSqueezedBacking { - entry: EntryID, - bytes: usize, - }, IoReadArrow { entry: EntryID, bytes: usize, @@ -56,18 +52,9 @@ pub(crate) enum InternalEvent { selection: bool, cached: CachedBatchType, }, - ReadSqueezedData { - entry: EntryID, - expression: CacheExpression, - }, TryReadLiquid { entry: EntryID, }, - DecompressSqueezed { - entry: EntryID, - decompressed: usize, - total: usize, - }, } #[derive(Debug)] @@ -106,13 +93,13 @@ impl fmt::Display for InternalEvent { kind ) } - InternalEvent::SqueezeBegin { victims } => { + InternalEvent::EvictionBegin { victims } => { let mut buf = String::new(); fmt_entry_list(&mut buf, victims)?; - write!(f, "event=squeeze_begin victims={}", buf) + write!(f, "event=eviction_begin victims={}", buf) } - InternalEvent::SqueezeVictim { entry } => { - write!(f, "event=squeeze_victim entry={}", usize::from(*entry)) + InternalEvent::EvictionVictim { entry } => { + write!(f, "event=eviction_victim entry={}", usize::from(*entry)) } InternalEvent::IoWrite { entry, kind, bytes } => { write!( @@ -131,14 +118,6 @@ impl fmt::Display for InternalEvent { bytes ) } - InternalEvent::IoReadSqueezedBacking { entry, bytes } => { - write!( - f, - "event=io_read_squeezed_backing entry={} bytes={}", - usize::from(*entry), - bytes - ) - } InternalEvent::IoReadArrow { entry, bytes } => { write!( f, @@ -191,27 +170,6 @@ impl fmt::Display for InternalEvent { InternalEvent::TryReadLiquid { entry } => { write!(f, "event=try_read_liquid entry={}", usize::from(*entry)) } - InternalEvent::ReadSqueezedData { entry, expression } => { - write!( - f, - "event=read_squeezed_data entry={} expression={}", - usize::from(*entry), - expression - ) - } - InternalEvent::DecompressSqueezed { - entry, - decompressed, - total, - } => { - write!( - f, - "event=decompress_squeezed entry={} decompressed={} total={}", - usize::from(*entry), - decompressed, - total - ) - } } } } diff --git a/src/core/src/cache/observer/mod.rs b/src/core/src/cache/observer/mod.rs index 38de5fc45..28419b6d0 100644 --- a/src/core/src/cache/observer/mod.rs +++ b/src/core/src/cache/observer/mod.rs @@ -84,21 +84,6 @@ impl Observer { self.runtime.incr_eval_predicate(); } - #[inline] - pub(crate) fn on_get_squeezed_success(&self) { - self.runtime.incr_get_squeezed_success(); - } - - #[inline] - pub(crate) fn on_get_squeezed_needs_io(&self) { - self.runtime.incr_get_squeezed_needs_io(); - } - - #[inline] - pub(crate) fn on_hit_date32_expression(&self) { - self.runtime.incr_hit_date32_expression(); - } - #[inline] pub(crate) fn on_disk_reservation_failure(&self) { self.runtime.incr_disk_reservation_failures(); @@ -111,18 +96,6 @@ impl Observer { InternalEvent::IoReadArrow { .. } | InternalEvent::IoReadLiquid { .. } => { self.runtime.incr_read_io_count() } - InternalEvent::IoReadSqueezedBacking { .. } => { - self.runtime.incr_read_io_count(); - self.runtime.incr_get_squeezed_needs_io(); - } - InternalEvent::DecompressSqueezed { - decompressed, - total, - .. - } => { - self.runtime - .track_decompress_squeezed_count(decompressed, total); - } _ => {} } @@ -130,8 +103,4 @@ impl Observer { self.event_tracer.record(event); } } - - pub(crate) fn runtime_stats(&self) -> &RuntimeStats { - &self.runtime - } } diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index fa0c3d9ae..6c871583d 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -97,28 +97,12 @@ define_runtime_stats! { (get, "Number of `get` calls issued via `CachedData`.", incr_get), (get_with_selection, "Number of `get_with_selection` calls issued via `CachedData`.", incr_get_with_selection), (eval_predicate, "Number of `eval_predicate` calls issued via `CachedData`.", incr_eval_predicate), - (get_squeezed_success, "Number of Squeezed-Liquid full evaluations finished without IO.", incr_get_squeezed_success), - (get_squeezed_needs_io, "Number of Squeezed-Liquid full paths that required IO.", incr_get_squeezed_needs_io), (try_read_liquid_calls, "Number of `try_read_liquid` calls issued via `CachedData`.", incr_try_read_liquid), - (hit_date32_expression_calls, "Number of `hit_date32_expression` calls.", incr_hit_date32_expression), (read_io_count, "Number of read IO operations.", incr_read_io_count), (write_io_count, "Number of write IO operations.", incr_write_io_count), (disk_evictions, "Number of disk cache entries evicted.", incr_disk_evictions), (disk_reservation_failures, "Number of failed disk budget reservations.", incr_disk_reservation_failures), (eval_predicate_on_liquid_failed, "Number of `eval_predicate` calls that failed on Liquid array.", incr_eval_predicate_on_liquid_failed), - (squeezed_decompressed_count, "Number of decompressed Squeezed-Liquid entries.", __incr_squeezed_decompressed_count), - (squeezed_total_count, "Total number of Squeezed-Liquid entries.", __incr_squeezed_total_count), - (squeeze_io_saved, "Number of io saved by squeezing.", incr_squeeze_io_saved), -} - -impl RuntimeStats { - /// Track the number of decompressed Squeezed-Liquid entries. - pub fn track_decompress_squeezed_count(&self, decompressed: usize, total: usize) { - self.squeezed_decompressed_count - .fetch_add(decompressed as u64, Ordering::Relaxed); - self.squeezed_total_count - .fetch_add(total as u64, Ordering::Relaxed); - } } /// Snapshot of cache statistics. @@ -130,8 +114,6 @@ pub struct CacheStats { pub memory_arrow_entries: usize, /// Number of in-memory Liquid entries. pub memory_liquid_entries: usize, - /// Number of in-memory Squeezed-Liquid entries. - pub memory_squeezed_liquid_entries: usize, /// Number of on-disk Liquid entries. pub disk_liquid_entries: usize, /// Number of on-disk Arrow entries. @@ -140,8 +122,6 @@ pub struct CacheStats { pub memory_arrow_bytes: usize, /// Total size of in-memory Liquid entries in bytes. pub memory_liquid_bytes: usize, - /// Total size of in-memory Squeezed-Liquid entries in bytes. - pub memory_squeezed_liquid_bytes: usize, /// Total memory usage of the cache. pub memory_usage_bytes: usize, /// Total disk usage of the cache. diff --git a/src/core/src/cache/policies/cache/three_queue.rs b/src/core/src/cache/policies/cache/three_queue.rs index 655c19847..f2780e245 100644 --- a/src/core/src/cache/policies/cache/three_queue.rs +++ b/src/core/src/cache/policies/cache/three_queue.rs @@ -11,7 +11,6 @@ use super::doubly_linked_list::{DoublyLinkedList, DoublyLinkedNode, drop_boxed_n enum QueueKind { Arrow, Liquid, - Squeezed, Disk, } @@ -28,7 +27,6 @@ struct LiquidQueueInternalState { map: HashMap, arrow: DoublyLinkedList, liquid: DoublyLinkedList, - squeezed: DoublyLinkedList, disk: DoublyLinkedList, } @@ -37,7 +35,6 @@ impl LiquidQueueInternalState { match queue { QueueKind::Arrow => &mut self.arrow, QueueKind::Liquid => &mut self.liquid, - QueueKind::Squeezed => &mut self.squeezed, QueueKind::Disk => &mut self.disk, } } @@ -81,7 +78,6 @@ impl LiquidQueueInternalState { let list = match queue { QueueKind::Arrow => &mut self.arrow, QueueKind::Liquid => &mut self.liquid, - QueueKind::Squeezed => &mut self.squeezed, QueueKind::Disk => &mut self.disk, }; @@ -117,7 +113,6 @@ impl Drop for LiquidQueueInternalState { match node_ptr.as_ref().data.queue { QueueKind::Arrow => self.arrow.unlink(node_ptr), QueueKind::Liquid => self.liquid.unlink(node_ptr), - QueueKind::Squeezed => self.squeezed.unlink(node_ptr), QueueKind::Disk => self.disk.unlink(node_ptr), } drop_boxed_node(node_ptr); @@ -127,7 +122,6 @@ impl Drop for LiquidQueueInternalState { unsafe { self.arrow.drop_all(); self.liquid.drop_all(); - self.squeezed.drop_all(); self.disk.drop_all(); } } @@ -158,7 +152,6 @@ impl CachePolicy for LiquidPolicy { let target = match batch_type { CachedBatchType::MemoryArrow => QueueKind::Arrow, CachedBatchType::MemoryLiquid => QueueKind::Liquid, - CachedBatchType::MemorySqueezedLiquid => QueueKind::Squeezed, CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow => QueueKind::Disk, }; @@ -184,11 +177,6 @@ impl CachePolicy for LiquidPolicy { continue; } - if let Some(entry) = inner.pop_front(QueueKind::Squeezed) { - victims.push(entry); - continue; - } - break; } @@ -255,15 +243,13 @@ mod tests { let arrow_entry = entry(1); let liquid_entry = entry(2); - let hybrid_entry = entry(3); policy.notify_insert(&liquid_entry, CachedBatchType::MemoryLiquid); - policy.notify_insert(&hybrid_entry, CachedBatchType::MemorySqueezedLiquid); policy.notify_insert(&arrow_entry, CachedBatchType::MemoryArrow); // Request more victims than available to ensure we only get what exists. let victims = policy.find_memory_victim(5); - assert_eq!(victims, vec![arrow_entry, liquid_entry, hybrid_entry]); + assert_eq!(victims, vec![arrow_entry, liquid_entry]); } #[test] diff --git a/src/core/src/cache/policies/eviction.rs b/src/core/src/cache/policies/eviction.rs new file mode 100644 index 000000000..b81e305cc --- /dev/null +++ b/src/core/src/cache/policies/eviction.rs @@ -0,0 +1,95 @@ +//! Policies for moving cache entries to cheaper storage under memory pressure. + +use arrow::array::Array; +use bytes::Bytes; + +use crate::cache::{ + CacheExpression, LiquidCompressorStates, cached_batch::CacheEntry, + transcode_liquid_inner_with_hint, utils::arrow_to_bytes, +}; + +/// The next storage representation selected for a cache entry. +#[derive(Debug, Clone)] +pub enum EvictionOutcome { + /// Replace the cache entry, optionally persisting these bytes first. + Replace { + /// Replacement cache entry. + entry: CacheEntry, + /// Bytes that must be persisted before installing the replacement. + bytes_to_write: Option, + }, + /// Remove an already-on-disk entry. + Remove, +} + +/// Chooses the next representation for an entry under memory pressure. +pub trait EvictionPolicy: std::fmt::Debug + Send + Sync { + /// Move the entry one step toward cheaper storage. + fn evict( + &self, + entry: &CacheEntry, + compressor: &LiquidCompressorStates, + expression: Option<&CacheExpression>, + ) -> EvictionOutcome; +} + +/// Evict memory entries directly to disk. +#[derive(Debug, Default, Clone)] +pub struct Evict; + +impl EvictionPolicy for Evict { + fn evict( + &self, + entry: &CacheEntry, + _compressor: &LiquidCompressorStates, + _expression: Option<&CacheExpression>, + ) -> EvictionOutcome { + persist(entry) + } +} + +/// Transcode Arrow to Liquid before eventually evicting it to disk. +#[derive(Debug, Default, Clone)] +pub struct TranscodeEvict; + +impl EvictionPolicy for TranscodeEvict { + fn evict( + &self, + entry: &CacheEntry, + compressor: &LiquidCompressorStates, + expression: Option<&CacheExpression>, + ) -> EvictionOutcome { + match entry { + CacheEntry::MemoryArrow(array) => { + match transcode_liquid_inner_with_hint(array, compressor, expression) { + Ok(liquid) => EvictionOutcome::Replace { + entry: CacheEntry::memory_liquid(liquid), + bytes_to_write: None, + }, + Err(_) => persist(entry), + } + } + _ => persist(entry), + } + } +} + +fn persist(entry: &CacheEntry) -> EvictionOutcome { + match entry { + CacheEntry::MemoryArrow(array) => { + let bytes = arrow_to_bytes(array).expect("failed to serialize Arrow array"); + EvictionOutcome::Replace { + entry: CacheEntry::disk_arrow(array.data_type().clone(), bytes.len()), + bytes_to_write: Some(bytes), + } + } + CacheEntry::MemoryLiquid(array) => { + let bytes = Bytes::from(array.to_bytes()); + EvictionOutcome::Replace { + entry: CacheEntry::disk_liquid(array.original_arrow_data_type(), bytes.len()), + bytes_to_write: Some(bytes), + } + } + CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } => EvictionOutcome::Remove, + } +} diff --git a/src/core/src/cache/policies/hydration.rs b/src/core/src/cache/policies/hydration.rs index ea96175d6..fbb3fdf2f 100644 --- a/src/core/src/cache/policies/hydration.rs +++ b/src/core/src/cache/policies/hydration.rs @@ -1,21 +1,12 @@ -//! Hydration policies decide whether and how to promote squeezed/on-disk entries back into memory. - -use std::sync::Arc; +//! Policies for promoting on-disk entries back into memory. use arrow::array::ArrayRef; use crate::{ - cache::{ - CacheExpression, LiquidCompressorStates, VariantRequest, cached_batch::CacheEntry, - utils::EntryID, - }, - liquid_array::{ - LiquidArrayRef, LiquidSqueezedArray, LiquidSqueezedArrayRef, VariantStructSqueezedArray, - }, + cache::{CacheExpression, cached_batch::CacheEntry, utils::EntryID}, + liquid_array::LiquidArrayRef, }; -use super::squeeze::try_variant_squeeze; - /// The materialized representation produced by a cache read. #[derive(Debug, Clone)] pub enum MaterializedEntry<'a> { @@ -25,129 +16,56 @@ pub enum MaterializedEntry<'a> { Liquid(&'a LiquidArrayRef), } -/// Request context provided to a [`HydrationPolicy`]. +/// Context for deciding whether to retain a materialized disk entry. #[derive(Debug, Clone)] pub struct HydrationRequest<'a> { /// Cache key being materialized. pub entry_id: EntryID, - /// The cached entry before materialization (e.g., `DiskArrow`). + /// The cached entry before materialization. pub cached: &'a CacheEntry, /// The fully materialized entry produced by the read path. pub materialized: MaterializedEntry<'a>, - /// Optional expression hint associated with the read. + /// Lineage expression associated with the read, when available. pub expression: Option<&'a CacheExpression>, - /// Compressor state used for hydrating into squeezed representations. - pub compressor: Arc, } -/// Decide if a materialized entry should be promoted back into memory. +/// Decide whether a materialized entry should be promoted back into memory. pub trait HydrationPolicy: std::fmt::Debug + Send + Sync { - /// Determine how to hydrate a cache entry that was just materialized. - /// Return a new cache entry to insert if hydration is desired. + /// Return a memory entry when hydration is desired. fn hydrate(&self, request: &HydrationRequest<'_>) -> Option; } -/// Default hydration policy: always keep a materialized cache miss in memory -/// by promoting along the path: disk -> squeezed -> liquid -> arrow. +/// Always retain materialized disk reads in memory. #[derive(Debug, Default, Clone)] pub struct AlwaysHydrate; impl AlwaysHydrate { - /// Create a new [`AlwaysHydrate`] policy. + /// Create a new policy. pub fn new() -> Self { Self } } -fn hydrate_variant_paths( - requests: &[VariantRequest], - materialized: &ArrayRef, - squeezed: &VariantStructSqueezedArray, - compressor: &LiquidCompressorStates, -) -> Option { - let missing_requests: Vec = requests - .iter() - .filter(|request| !squeezed.contains_path(request.path())) - .cloned() - .collect(); - if missing_requests.is_empty() { - return None; - } - - let (new_squeezed, _) = try_variant_squeeze(materialized, &missing_requests, compressor)?; - let new_variant = new_squeezed - .as_any() - .downcast_ref::()?; - - let mut combined_values = squeezed.typed_values(); - combined_values.extend(new_variant.typed_values()); - - let nulls = squeezed.nulls().or_else(|| new_variant.nulls()); - let merged = VariantStructSqueezedArray::new( - combined_values, - nulls, - squeezed.original_arrow_data_type(), - squeezed.disk_backing().disk_bytes(), - ); - Some(CacheEntry::memory_squeezed_liquid( - Arc::new(merged) as LiquidSqueezedArrayRef - )) -} - impl HydrationPolicy for AlwaysHydrate { fn hydrate(&self, request: &HydrationRequest<'_>) -> Option { - match (request.cached, &request.materialized) { - (CacheEntry::DiskArrow { disk_bytes, .. }, MaterializedEntry::Arrow(arr)) => { - if let Some(CacheExpression::VariantGet { requests }) = request.expression - && let Some((squeezed, _bytes)) = - try_variant_squeeze(arr, requests, request.compressor.as_ref()) - { - let variant = squeezed - .as_any() - .downcast_ref::()?; - let squeezed = VariantStructSqueezedArray::new( - variant.typed_values(), - variant.nulls(), - variant.original_arrow_data_type(), - *disk_bytes, - ); - return Some(CacheEntry::memory_squeezed_liquid( - Arc::new(squeezed) as LiquidSqueezedArrayRef - )); - } - Some(CacheEntry::memory_arrow((*arr).clone())) - } - (CacheEntry::DiskLiquid { .. }, MaterializedEntry::Liquid(liq)) => { - Some(CacheEntry::memory_liquid((*liq).clone())) - } - (CacheEntry::MemoryLiquid(_), _) => None, - // When already squeezed/hybrid or liquid in memory, prefer promoting to Arrow if available. - (CacheEntry::MemorySqueezedLiquid(squeezed_entry), MaterializedEntry::Arrow(arr)) => { - if let Some(CacheExpression::VariantGet { requests }) = request.expression - && let Some(squeezed) = squeezed_entry - .as_any() - .downcast_ref::() - && let Some(entry) = - hydrate_variant_paths(requests, arr, squeezed, request.compressor.as_ref()) - { - return Some(entry); - } - Some(CacheEntry::memory_arrow((*arr).clone())) + match (&request.cached, &request.materialized) { + (CacheEntry::DiskArrow { .. }, MaterializedEntry::Arrow(array)) => { + Some(CacheEntry::memory_arrow((*array).clone())) } - (CacheEntry::MemorySqueezedLiquid(_), MaterializedEntry::Liquid(liq)) => { - Some(CacheEntry::memory_liquid((*liq).clone())) + (CacheEntry::DiskLiquid { .. }, MaterializedEntry::Liquid(array)) => { + Some(CacheEntry::memory_liquid((*array).clone())) } _ => None, } } } -/// No hydration policy: never promote a materialized entry back into memory. +/// Never retain materialized disk reads in memory. #[derive(Debug, Default, Clone)] pub struct NoHydration; impl NoHydration { - /// Create a new [`NoHydration`] policy. + /// Create a new policy. pub fn new() -> Self { Self } @@ -158,87 +76,3 @@ impl HydrationPolicy for NoHydration { None } } - -#[cfg(test)] -mod tests { - use super::*; - use parquet_variant_compute::json_to_variant; - - use crate::cache::utils::LiquidCompressorStates; - use arrow::array::StringArray; - use arrow_schema::DataType; - - fn variant_array_from_json(values: &[&str]) -> ArrayRef { - let strings: ArrayRef = Arc::new(StringArray::from_iter_values(values.iter().copied())); - let variant = json_to_variant(&strings).expect("variant array"); - let struct_array = variant.into_inner(); - Arc::new(struct_array) as ArrayRef - } - - #[test] - fn hydrates_disk_arrow_variant_to_squeezed() { - let arr = variant_array_from_json(&[r#"{"name":"Ada","age":30}"#]); - let expr = CacheExpression::variant_get("age", DataType::Int64); - let policy = AlwaysHydrate::new(); - let compressor = Arc::new(LiquidCompressorStates::new()); - let cached_entry = CacheEntry::disk_arrow(arr.data_type().clone(), 1); - - let hydrated = policy.hydrate(&HydrationRequest { - entry_id: EntryID::from(0), - cached: &cached_entry, - materialized: MaterializedEntry::Arrow(&arr), - expression: Some(&expr), - compressor, - }); - - let entry = hydrated.expect("should hydrate"); - match entry { - CacheEntry::MemorySqueezedLiquid(squeezed) => { - let variant = squeezed - .as_any() - .downcast_ref::() - .expect("variant squeezed"); - assert!(variant.contains_path("age")); - } - other => panic!("expected squeezed entry, got {other:?}"), - } - } - - #[test] - fn hydrates_squeezed_variant_adds_missing_path() { - let arr = variant_array_from_json(&[r#"{"name":"Ada","age":30}"#]); - let name_expr = CacheExpression::variant_get("name", DataType::Utf8); - let age_expr = CacheExpression::variant_get("age", DataType::Int64); - let compressor = Arc::new(LiquidCompressorStates::new()); - - // Build an initial squeezed array containing only the "name" path. - let (squeezed, _) = try_variant_squeeze( - &arr, - name_expr.variant_requests().unwrap(), - compressor.as_ref(), - ) - .expect("squeeze name"); - let cached_entry = CacheEntry::memory_squeezed_liquid(squeezed.clone()); - - let policy = AlwaysHydrate::new(); - let hydrated = policy.hydrate(&HydrationRequest { - entry_id: EntryID::from(1), - cached: &cached_entry, - materialized: MaterializedEntry::Arrow(&arr), - expression: Some(&age_expr), - compressor, - }); - - let entry = hydrated.expect("should hydrate"); - let squeezed = match entry { - CacheEntry::MemorySqueezedLiquid(sq) => sq, - other => panic!("expected squeezed entry, got {other:?}"), - }; - let variant = squeezed - .as_any() - .downcast_ref::() - .expect("variant squeezed"); - assert!(variant.contains_path("name")); - assert!(variant.contains_path("age")); - } -} diff --git a/src/core/src/cache/policies/mod.rs b/src/core/src/cache/policies/mod.rs index 4dde5cfe7..624c66b99 100644 --- a/src/core/src/cache/policies/mod.rs +++ b/src/core/src/cache/policies/mod.rs @@ -1,9 +1,9 @@ -//! Policy modules for cache eviction, hydration, and squeezing. +//! Policy modules for cache eviction and hydration. pub mod cache; +pub mod eviction; pub mod hydration; -pub mod squeeze; pub use cache::*; +pub use eviction::*; pub use hydration::*; -pub use squeeze::*; diff --git a/src/core/src/cache/policies/squeeze.rs b/src/core/src/cache/policies/squeeze.rs deleted file mode 100644 index 7941a1feb..000000000 --- a/src/core/src/cache/policies/squeeze.rs +++ /dev/null @@ -1,863 +0,0 @@ -//! Squeeze policies for liquid cache. - -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, StructArray}; -use arrow_schema::DataType; -use bytes::Bytes; -use parquet::variant::VariantPath; -use parquet_variant_compute::{VariantArray, shred_variant, unshred_variant}; - -use crate::cache::{ - CacheExpression, LiquidCompressorStates, VariantRequest, cached_batch::CacheEntry, - transcode_liquid_inner, transcode_liquid_inner_with_hint, utils::arrow_to_bytes, -}; -use crate::liquid_array::{ - LiquidSqueezedArrayRef, SqueezeIoHandler, SqueezedBacking, VariantStructSqueezedArray, -}; -use crate::utils::VariantSchema; - -/// What to do when we need to squeeze an entry? -#[derive(Debug, Clone)] -pub enum SqueezeOutcome { - /// Replace the cache entry, optionally writing bytes to disk first. - Replace { - /// Replacement cache entry. - entry: CacheEntry, - /// Bytes that must be written before inserting the replacement. - bytes_to_write: Option, - }, - /// Remove the entry entirely. - Remove, -} - -/// Policy that chooses the next representation for an entry under memory pressure. -pub trait SqueezePolicy: std::fmt::Debug + Send + Sync { - /// Squeeze the entry. - fn squeeze( - &self, - entry: &CacheEntry, - compressor: &LiquidCompressorStates, - squeeze_hint: Option<&CacheExpression>, - squeeze_io: &Arc, - ) -> SqueezeOutcome; -} - -/// Squeeze the entry to disk. -#[derive(Debug, Default, Clone)] -pub struct Evict; - -impl SqueezePolicy for Evict { - fn squeeze( - &self, - entry: &CacheEntry, - _compressor: &LiquidCompressorStates, - _squeeze_hint: Option<&CacheExpression>, - _squeeze_io: &Arc, - ) -> SqueezeOutcome { - match entry { - CacheEntry::MemoryArrow(array) => { - let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); - SqueezeOutcome::Replace { - entry: CacheEntry::disk_arrow(array.data_type().clone(), bytes.len()), - bytes_to_write: Some(bytes), - } - } - CacheEntry::MemoryLiquid(liquid_array) => { - let disk_data = liquid_array.to_bytes(); - SqueezeOutcome::Replace { - entry: CacheEntry::disk_liquid( - liquid_array.original_arrow_data_type(), - disk_data.len(), - ), - bytes_to_write: Some(Bytes::from(disk_data)), - } - } - CacheEntry::MemorySqueezedLiquid(squeezed_array) => { - let data_type = squeezed_array.original_arrow_data_type(); - let new_entry = match squeezed_array.disk_backing() { - SqueezedBacking::Liquid(n) => CacheEntry::disk_liquid(data_type, n), - SqueezedBacking::Arrow(n) => CacheEntry::disk_arrow(data_type, n), - }; - SqueezeOutcome::Replace { - entry: new_entry, - bytes_to_write: None, - } - } - CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } => SqueezeOutcome::Remove, - } - } -} - -/// Squeeze the entry to liquid memory. -#[derive(Debug, Default, Clone)] -pub struct TranscodeSqueezeEvict; - -impl SqueezePolicy for TranscodeSqueezeEvict { - fn squeeze( - &self, - entry: &CacheEntry, - compressor: &LiquidCompressorStates, - squeeze_hint: Option<&CacheExpression>, - squeeze_io: &Arc, - ) -> SqueezeOutcome { - match entry { - CacheEntry::MemoryArrow(array) => { - if let Some(requests) = - squeeze_hint.and_then(|expression| expression.variant_requests()) - && let Some((squeezed_array, bytes)) = - try_variant_squeeze(array, requests, compressor) - { - return SqueezeOutcome::Replace { - entry: CacheEntry::memory_squeezed_liquid(squeezed_array), - bytes_to_write: Some(bytes), - }; - } - match transcode_liquid_inner_with_hint(array, compressor, squeeze_hint) { - Ok(liquid_array) => SqueezeOutcome::Replace { - entry: CacheEntry::memory_liquid(liquid_array), - bytes_to_write: None, - }, - Err(_) => { - let bytes = - arrow_to_bytes(array).expect("failed to convert arrow to bytes"); - SqueezeOutcome::Replace { - entry: CacheEntry::disk_arrow(array.data_type().clone(), bytes.len()), - bytes_to_write: Some(bytes), - } - } - } - } - CacheEntry::MemoryLiquid(liquid_array) => { - let (squeezed_array, bytes) = - match liquid_array.squeeze(squeeze_io.clone(), squeeze_hint) { - Some(result) => result, - None => { - let bytes = Bytes::from(liquid_array.to_bytes()); - return SqueezeOutcome::Replace { - entry: CacheEntry::disk_liquid( - liquid_array.original_arrow_data_type(), - bytes.len(), - ), - bytes_to_write: Some(bytes), - }; - } - }; - SqueezeOutcome::Replace { - entry: CacheEntry::memory_squeezed_liquid(squeezed_array), - bytes_to_write: Some(bytes), - } - } - CacheEntry::MemorySqueezedLiquid(squeezed_array) => { - let data_type = squeezed_array.original_arrow_data_type(); - let new_entry = match squeezed_array.disk_backing() { - SqueezedBacking::Liquid(n) => CacheEntry::disk_liquid(data_type, n), - SqueezedBacking::Arrow(n) => CacheEntry::disk_arrow(data_type, n), - }; - SqueezeOutcome::Replace { - entry: new_entry, - bytes_to_write: None, - } - } - CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } => SqueezeOutcome::Remove, - } - } -} - -/// Squeeze the entry to liquid memory, but don't convert to squeezed. -#[derive(Debug, Default, Clone)] -pub struct TranscodeEvict; - -impl SqueezePolicy for TranscodeEvict { - fn squeeze( - &self, - entry: &CacheEntry, - compressor: &LiquidCompressorStates, - _squeeze_hint: Option<&CacheExpression>, - _squeeze_io: &Arc, - ) -> SqueezeOutcome { - match entry { - CacheEntry::MemoryArrow(array) => { - match transcode_liquid_inner_with_hint(array, compressor, None) { - Ok(liquid_array) => SqueezeOutcome::Replace { - entry: CacheEntry::memory_liquid(liquid_array), - bytes_to_write: None, - }, - Err(_) => { - let bytes = - arrow_to_bytes(array).expect("failed to convert arrow to bytes"); - SqueezeOutcome::Replace { - entry: CacheEntry::disk_arrow(array.data_type().clone(), bytes.len()), - bytes_to_write: Some(bytes), - } - } - } - } - CacheEntry::MemoryLiquid(liquid_array) => { - let bytes = Bytes::from(liquid_array.to_bytes()); - SqueezeOutcome::Replace { - entry: CacheEntry::disk_liquid( - liquid_array.original_arrow_data_type(), - bytes.len(), - ), - bytes_to_write: Some(bytes), - } - } - CacheEntry::MemorySqueezedLiquid(squeezed_array) => { - let data_type = squeezed_array.original_arrow_data_type(); - let new_entry = match squeezed_array.disk_backing() { - SqueezedBacking::Liquid(n) => CacheEntry::disk_liquid(data_type, n), - SqueezedBacking::Arrow(n) => CacheEntry::disk_arrow(data_type, n), - }; - SqueezeOutcome::Replace { - entry: new_entry, - bytes_to_write: None, - } - } - CacheEntry::DiskLiquid { .. } | CacheEntry::DiskArrow { .. } => SqueezeOutcome::Remove, - } - } -} - -pub(crate) fn try_variant_squeeze( - array: &ArrayRef, - requests: &[VariantRequest], - compressor: &LiquidCompressorStates, -) -> Option<(LiquidSqueezedArrayRef, Bytes)> { - let struct_array = array.as_any().downcast_ref::()?; - let mut variant_array = VariantArray::try_new(struct_array).ok()?; - if variant_array.is_empty() { - return None; - } - - if requests.is_empty() { - return None; - } - - let mut shredded_array: Option = None; - if let Some(shredding_type) = build_shredding_schema(struct_array, requests) - && let Ok(unshredded) = unshred_variant(&variant_array) - && let Ok(shredded) = shred_variant(&unshredded, &shredding_type) - { - let shredded_struct: ArrayRef = Arc::new(shredded.into_inner()); - variant_array = VariantArray::try_new(shredded_struct.as_ref()).ok()?; - shredded_array = Some(shredded_struct); - } - - let typed_root = variant_array.typed_value_column()?; - let typed_root = typed_root.as_any().downcast_ref::()?; - - let mut collected = Vec::new(); - for request in requests { - let path = request.path().trim(); - if path.is_empty() { - continue; - } - let Some(path_struct) = extract_typed_values_for_path(typed_root, path) else { - continue; - }; - let path_struct = path_struct.as_any().downcast_ref::()?; - let Some(typed_values) = path_struct.column_by_name("typed_value") else { - continue; - }; - if typed_values.len() != array.len() { - continue; - } - collected.push((Arc::::from(path.to_string()), typed_values.clone())); - } - - if collected.is_empty() { - return None; - } - - let backing_array = shredded_array.as_ref().unwrap_or(array); - let nulls = variant_array.inner().nulls().cloned(); - let bytes = arrow_to_bytes(backing_array).ok()?; - let mut liquid_values = Vec::with_capacity(collected.len()); - for (path, typed_values) in collected { - let Ok(liquid_array) = transcode_liquid_inner(&typed_values, compressor) else { - return None; - }; - liquid_values.push((path, liquid_array)); - } - let squeezed = VariantStructSqueezedArray::new( - liquid_values, - nulls, - backing_array.data_type().clone(), - bytes.len(), - ); - Some((Arc::new(squeezed) as LiquidSqueezedArrayRef, bytes)) -} - -fn build_shredding_schema( - variant_struct: &StructArray, - requests: &[VariantRequest], -) -> Option { - let typed_field = match variant_struct.data_type() { - DataType::Struct(fields) => fields - .iter() - .find(|child| child.name() == "typed_value") - .cloned(), - _ => None, - }; - - let mut schema = VariantSchema::new(typed_field.as_deref()); - for request in requests { - let path = request.path().trim(); - if path.is_empty() { - continue; - } - schema.insert_path(path, request.data_type()); - } - schema.shredding_type() -} - -fn extract_typed_values_for_path(typed_root: &StructArray, path: &str) -> Option { - let path = VariantPath::try_from(path).ok()?; - if path.is_empty() { - return None; - } - - let mut cursor = typed_root; - for (idx, element) in path.iter().enumerate() { - let field_name = match element { - parquet::variant::VariantPathElement::Field { name } => name.as_ref(), - parquet::variant::VariantPathElement::Index { .. } => return None, - }; - let field = cursor.column_by_name(field_name)?; - if idx == path.len() - 1 { - return Some(field.clone()); - } - let struct_field = field.as_any().downcast_ref::()?; - let typed_value = struct_field.column_by_name("typed_value")?; - cursor = typed_value.as_any().downcast_ref::()?; - } - - None -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::cache::cached_batch::CacheEntry; - use crate::cache::{CacheExpression, io_context::TestSqueezeIo}; - use crate::liquid_array::{LiquidSqueezedArray, SqueezedBacking, VariantStructSqueezedArray}; - use arrow::array::{Array, ArrayRef, Int32Array, StringArray, StructArray}; - use arrow_schema::Fields; - use arrow_schema::{DataType, Field}; - use parquet::variant::VariantPath; - use parquet_variant_compute::{GetOptions, json_to_variant, variant_get}; - use std::collections::BTreeMap; - use std::sync::Arc; - - fn int_array(n: i32) -> ArrayRef { - Arc::new(Int32Array::from_iter_values(0..n)) - } - - fn decode_arrow(bytes: &Bytes) -> ArrayRef { - let cursor = std::io::Cursor::new(bytes.to_vec()); - let mut reader = - arrow::ipc::reader::StreamReader::try_new(cursor, None).expect("arrow stream"); - let batch = reader - .next() - .expect("non-empty stream") - .expect("read stream"); - batch.column(0).clone() - } - - fn into_replace(outcome: SqueezeOutcome) -> (CacheEntry, Option) { - match outcome { - SqueezeOutcome::Replace { - entry, - bytes_to_write, - } => (entry, bytes_to_write), - SqueezeOutcome::Remove => panic!("expected replacement"), - } - } - - fn struct_array() -> ArrayRef { - let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; - let field = Arc::new(Field::new("value", DataType::Int32, true)); - Arc::new(StructArray::from(vec![(field, values)])) - } - - #[test] - fn test_squeeze_to_disk_policy() { - let disk = Evict; - let states = LiquidCompressorStates::new(); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - // MemoryArrow -> DiskArrow + bytes (Arrow IPC) - let arr = int_array(8); - let (new_batch, bytes) = into_replace(disk.squeeze( - &CacheEntry::memory_arrow(arr.clone()), - &states, - None, - &squeeze_io, - )); - let data = new_batch; - match (data, bytes) { - ( - CacheEntry::DiskArrow { - data_type: dt, - disk_bytes, - }, - Some(b), - ) => { - assert_eq!(dt, DataType::Int32); - assert_eq!(disk_bytes, b.len()); - let decoded = decode_arrow(&b); - assert_eq!(decoded.as_ref(), arr.as_ref()); - } - other => panic!("unexpected: {other:?}"), - } - - // MemoryLiquid (strings) -> MemoryHybridLiquid + bytes - let strings = Arc::new(StringArray::from(vec!["a", "b", "a"])) as ArrayRef; - let liquid = transcode_liquid_inner(&strings, &states).unwrap(); - let (new_batch, bytes) = into_replace(disk.squeeze( - &CacheEntry::memory_liquid(liquid.clone()), - &states, - None, - &squeeze_io, - )); - let data = new_batch; - match (data, bytes) { - (CacheEntry::DiskLiquid { disk_bytes, .. }, Some(b)) => { - assert_eq!(disk_bytes, b.len()); - assert!(!b.is_empty()); - } - other => panic!("unexpected: {other:?}"), - } - - let expression = Some(&CacheExpression::PredicateColumn); - // MemorySqueezedLiquid -> DiskLiquid, no extra bytes - let squeezed = match liquid.squeeze(squeeze_io.clone(), expression) { - Some((h, _b)) => h, - None => panic!("squeeze should succeed for byte-view"), - }; - let (new_batch, bytes) = into_replace(disk.squeeze( - &CacheEntry::memory_squeezed_liquid(squeezed), - &states, - expression, - &squeeze_io, - )); - let data = new_batch; - match (data, bytes) { - ( - CacheEntry::DiskLiquid { - data_type: _data_type, - .. - }, - None, - ) => {} - other => panic!("unexpected: {other:?}"), - } - - // Disk* -> remove - let b1 = disk.squeeze( - &CacheEntry::disk_arrow(DataType::Utf8, 1), - &states, - expression, - &squeeze_io, - ); - assert!(matches!(b1, SqueezeOutcome::Remove)); - let b2 = disk.squeeze( - &CacheEntry::disk_liquid(DataType::Utf8, 1), - &states, - expression, - &squeeze_io, - ); - assert!(matches!(b2, SqueezeOutcome::Remove)); - } - - #[test] - fn test_squeeze_to_liquid_policy() { - let to_liquid = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - - // MemoryArrow -> MemoryLiquid, no bytes - let arr = int_array(8); - let (new_batch, bytes) = into_replace(to_liquid.squeeze( - &CacheEntry::memory_arrow(arr.clone()), - &states, - None, - &squeeze_io, - )); - assert!(bytes.is_none()); - match new_batch { - CacheEntry::MemoryLiquid(liq) => { - assert_eq!(liq.to_arrow_array().as_ref(), arr.as_ref()); - } - other => panic!("unexpected: {other:?}"), - } - let expression = Some(&CacheExpression::PredicateColumn); - - // MemoryLiquid (strings) -> MemorySqueezedLiquid + bytes - let strings = Arc::new(StringArray::from(vec!["x", "y", "x"])) as ArrayRef; - let liquid = transcode_liquid_inner(&strings, &states).unwrap(); - let (new_batch, bytes) = into_replace(to_liquid.squeeze( - &CacheEntry::memory_liquid(liquid), - &states, - expression, - &squeeze_io, - )); - match (new_batch, bytes) { - (CacheEntry::MemorySqueezedLiquid(_), Some(b)) => assert!(!b.is_empty()), - other => panic!("unexpected: {other:?}"), - } - - // MemorySqueezedLiquid -> DiskLiquid, no bytes - let strings = Arc::new(StringArray::from(vec!["m", "n"])) as ArrayRef; - let liquid = transcode_liquid_inner(&strings, &states).unwrap(); - let squeezed = liquid.squeeze(squeeze_io.clone(), expression).unwrap().0; - let (new_batch, bytes) = into_replace(to_liquid.squeeze( - &CacheEntry::memory_squeezed_liquid(squeezed), - &states, - expression, - &squeeze_io, - )); - match (new_batch, bytes) { - ( - CacheEntry::DiskLiquid { - data_type: DataType::Utf8, - .. - }, - None, - ) => {} - other => panic!("unexpected: {other:?}"), - } - - // Disk* -> remove - let b1 = to_liquid.squeeze( - &CacheEntry::disk_arrow(DataType::Utf8, 1), - &states, - expression, - &squeeze_io, - ); - assert!(matches!(b1, SqueezeOutcome::Remove)); - let b2 = to_liquid.squeeze( - &CacheEntry::disk_liquid(DataType::Utf8, 1), - &states, - expression, - &squeeze_io, - ); - assert!(matches!(b2, SqueezeOutcome::Remove)); - } - - #[test] - fn transcode_squeeze_struct_falls_back_to_disk_arrow() { - let to_liquid = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - let struct_arr = struct_array(); - let (new_batch, bytes) = into_replace(to_liquid.squeeze( - &CacheEntry::memory_arrow(struct_arr.clone()), - &states, - None, - &squeeze_io, - )); - match (new_batch, bytes) { - ( - CacheEntry::DiskArrow { - data_type: dt, - disk_bytes, - }, - Some(b), - ) => { - assert_eq!(&dt, struct_arr.data_type()); - assert_eq!(disk_bytes, b.len()); - assert_eq!(decode_arrow(&b).as_ref(), struct_arr.as_ref()); - } - other => panic!("expected disk arrow fallback, got {other:?}"), - } - } - - #[test] - fn transcode_evict_struct_falls_back_to_disk_arrow() { - let to_disk = TranscodeEvict; - let states = LiquidCompressorStates::new(); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - let struct_arr = struct_array(); - let (new_batch, bytes) = into_replace(to_disk.squeeze( - &CacheEntry::memory_arrow(struct_arr.clone()), - &states, - None, - &squeeze_io, - )); - match (new_batch, bytes) { - ( - CacheEntry::DiskArrow { - data_type: dt, - disk_bytes, - }, - Some(b), - ) => { - assert_eq!(&dt, struct_arr.data_type()); - assert_eq!(disk_bytes, b.len()); - assert_eq!(decode_arrow(&b).as_ref(), struct_arr.as_ref()); - } - other => panic!("expected disk arrow fallback, got {other:?}"), - } - } - - fn enriched_variant_array(path: &str, data_type: DataType) -> ArrayRef { - enriched_variant_array_with_paths(&[(path, data_type)]) - } - - fn enriched_variant_array_with_paths(entries: &[(&str, DataType)]) -> ArrayRef { - let values: ArrayRef = Arc::new(StringArray::from(vec![ - Some(r#"{"name": "Alice", "age": 30}"#), - Some(r#"{"name": "Bob", "age": 25}"#), - Some(r#"{"name": "Charlie", "age": 35}"#), - ])); - let base_variant = json_to_variant(&values).unwrap(); - let base_arr: ArrayRef = Arc::new(base_variant.inner().clone()); - - let mut typed_structs: BTreeMap = BTreeMap::new(); - - for (path, data_type) in entries.iter() { - let typed_values = variant_get( - &base_arr, - GetOptions::new_with_path( - VariantPath::try_from(*path).expect("variant path should parse"), - ) - .with_as_type(Some(Arc::new(Field::new( - "typed_value", - data_type.clone(), - true, - )))), - ) - .unwrap(); - - typed_structs - .entry(path.to_string()) - .or_insert(Arc::new(StructArray::new( - Fields::from(vec![Arc::new(Field::new( - "typed_value", - data_type.clone(), - true, - ))]), - vec![typed_values.clone()], - None, - ))); - } - - let mut typed_fields: Vec> = Vec::new(); - let mut typed_columns: Vec = Vec::new(); - for (name, tree) in typed_structs { - typed_fields.push(Arc::new(Field::new( - name.as_str(), - tree.data_type().clone(), - true, - ))); - typed_columns.push(tree.clone()); - } - - let typed_struct = Arc::new(StructArray::new( - Fields::from(typed_fields), - typed_columns, - base_variant.inner().nulls().cloned(), - )); - - let inner = base_variant.inner(); - use arrow::array::BinaryViewArray; - Arc::new(StructArray::new( - Fields::from(vec![ - Arc::new(Field::new("metadata", DataType::BinaryView, false)), - Arc::new(Field::new("value", DataType::BinaryView, true)), - Arc::new(Field::new( - "typed_value", - typed_struct.data_type().clone(), - true, - )), - ]), - vec![ - inner - .column_by_name("metadata") - .cloned() - .unwrap_or_else(|| base_variant.metadata_column().clone()), - inner.column_by_name("value").cloned().unwrap_or_else(|| { - Arc::new(BinaryViewArray::from(vec![None::<&[u8]>; inner.len()])) as ArrayRef - }), - typed_struct as ArrayRef, - ], - inner.nulls().cloned(), - )) as ArrayRef - } - - fn assert_variant_squeezed( - squeezed: &LiquidSqueezedArrayRef, - expected_path: &str, - bytes: &Bytes, - ) { - use futures::executor::block_on; - - assert!(!bytes.is_empty()); - assert!(matches!(squeezed.disk_backing(), SqueezedBacking::Arrow(_))); - let struct_squeezed = squeezed - .as_any() - .downcast_ref::() - .expect("squeezed variant struct"); - let arrow_array = block_on(struct_squeezed.to_arrow_array()); - let struct_array = arrow_array - .as_any() - .downcast_ref::() - .expect("variant struct"); - let value_column = struct_array - .column_by_name("value") - .expect("value column present"); - assert_eq!(value_column.len(), value_column.null_count()); - let typed_struct = struct_array - .column_by_name("typed_value") - .expect("typed_value column") - .as_any() - .downcast_ref::() - .expect("typed struct"); - assert!( - extract_typed_values_for_path(typed_struct, expected_path).is_some(), - "typed path {expected_path} missing from squeezed variant" - ); - } - - #[test] - fn test_variant_squeeze_with_hint() { - let policy = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let variant_arr = enriched_variant_array("name", DataType::Utf8); - let hint = CacheExpression::variant_get("name", DataType::Utf8); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - - let (new_batch, bytes) = into_replace(policy.squeeze( - &CacheEntry::memory_arrow(variant_arr), - &states, - Some(&hint), - &squeeze_io, - )); - - match (new_batch, bytes) { - (CacheEntry::MemorySqueezedLiquid(squeezed), Some(b)) => { - assert_variant_squeezed(&squeezed, "name", &b); - } - other => panic!("expected MemorySqueezedLiquid with bytes, got {other:?}"), - } - } - - #[test] - fn test_variant_squeeze_with_int64_path() { - let policy = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let variant_arr = enriched_variant_array("age", DataType::Int64); - let hint = CacheExpression::variant_get("age", DataType::Int64); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - - let (new_batch, bytes) = into_replace(policy.squeeze( - &CacheEntry::memory_arrow(variant_arr), - &states, - Some(&hint), - &squeeze_io, - )); - - match (new_batch, bytes) { - (CacheEntry::MemorySqueezedLiquid(squeezed), Some(b)) => { - assert_variant_squeezed(&squeezed, "age", &b); - } - other => panic!("expected MemorySqueezedLiquid with bytes, got {other:?}"), - } - } - - #[test] - fn test_variant_squeeze_with_multiple_paths_preserves_all_fields() { - let policy = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let variant_arr = enriched_variant_array_with_paths(&[ - ("name", DataType::Utf8), - ("age", DataType::Int64), - ]); - let hint = CacheExpression::variant_get("name", DataType::Utf8); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - - let (new_batch, bytes) = into_replace(policy.squeeze( - &CacheEntry::memory_arrow(variant_arr), - &states, - Some(&hint), - &squeeze_io, - )); - - match (new_batch, bytes) { - (CacheEntry::MemorySqueezedLiquid(squeezed), Some(b)) => { - assert!(!b.is_empty()); - let struct_squeezed = squeezed - .as_any() - .downcast_ref::() - .unwrap(); - let arrow_array = futures::executor::block_on(struct_squeezed.to_arrow_array()); - let struct_array = arrow_array.as_any().downcast_ref::().unwrap(); - let typed_value = struct_array - .column_by_name("typed_value") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert!(typed_value.column_by_name("name").is_some()); - assert!(typed_value.column_by_name("age").is_none()); - } - other => panic!("expected MemorySqueezedLiquid with bytes, got {other:?}"), - } - } - - #[test] - fn test_variant_squeeze_without_hint() { - let policy = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let variant_arr = enriched_variant_array("name", DataType::Utf8); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - - let (new_batch, bytes) = into_replace(policy.squeeze( - &CacheEntry::memory_arrow(variant_arr), - &states, - None, - &squeeze_io, - )); - - match (new_batch, bytes) { - (CacheEntry::DiskArrow { disk_bytes, .. }, Some(b)) => { - assert_eq!(disk_bytes, b.len()); - assert!(!b.is_empty()); - } - (CacheEntry::MemoryLiquid(_), None) => {} - other => panic!("expected DiskArrow with bytes or MemoryLiquid, got {other:?}"), - } - } - - #[test] - fn test_variant_squeeze_skips_when_path_missing() { - let policy = TranscodeSqueezeEvict; - let states = LiquidCompressorStates::new(); - let squeeze_io: Arc = Arc::new(TestSqueezeIo::default()); - let variant_arr = enriched_variant_array("name", DataType::Utf8); - let hint = CacheExpression::variant_get("age", DataType::Int64); - - let (new_batch, bytes) = into_replace(policy.squeeze( - &CacheEntry::memory_arrow(variant_arr.clone()), - &states, - Some(&hint), - &squeeze_io, - )); - - match (new_batch, bytes) { - ( - CacheEntry::DiskArrow { - data_type: dt, - disk_bytes, - }, - Some(b), - ) => { - assert_eq!(dt, variant_arr.data_type().clone()); - assert_eq!(disk_bytes, b.len()); - assert!(!b.is_empty()); - } - other => panic!("expected DiskArrow fallback when path missing, got {other:?}"), - } - } -} diff --git a/src/core/src/cache/tests/mod.rs b/src/core/src/cache/tests/mod.rs index 86beb85c9..1ac026073 100644 --- a/src/core/src/cache/tests/mod.rs +++ b/src/core/src/cache/tests/mod.rs @@ -1,2 +1 @@ mod policies; -mod squeezed; diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index 29c849902..1e26f45a9 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -1,5 +1,5 @@ use crate::cache::{ - AlwaysHydrate, EntryID, LiquidCacheBuilder, LiquidPolicy, TranscodeSqueezeEvict, + AlwaysHydrate, EntryID, LiquidCacheBuilder, LiquidPolicy, TranscodeEvict, utils::create_test_arrow_array, }; @@ -11,7 +11,7 @@ async fn default_policies() { let cache = LiquidCacheBuilder::new() .with_cache_policy(Box::new(LiquidPolicy::new())) .with_hydration_policy(Box::new(AlwaysHydrate::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_max_memory_bytes(capacity) .build() .await; @@ -39,7 +39,7 @@ async fn insert_wont_fit_cache() { let cache = LiquidCacheBuilder::new() .with_cache_policy(Box::new(LiquidPolicy::new())) .with_hydration_policy(Box::new(AlwaysHydrate::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_max_memory_bytes(capacity) .build() .await; diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap index 83ab4e86d..f4dc0a56a 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap @@ -1,33 +1,34 @@ --- source: src/core/src/cache/tests/policies.rs +assertion_line: 31 expression: trace --- EventTrace: [ event=insert_success entry=0 kind=MemoryArrow event=insert_success entry=1 kind=MemoryArrow event=insert_failed entry=2 kind=MemoryArrow -event=squeeze_begin victims=[0,1] -event=squeeze_victim entry=0 +event=eviction_begin victims=[0,1] +event=eviction_victim entry=0 event=insert_success entry=0 kind=MemoryLiquid -event=squeeze_victim entry=1 +event=eviction_victim entry=1 event=insert_success entry=1 kind=MemoryLiquid event=insert_success entry=2 kind=MemoryArrow event=insert_failed entry=3 kind=MemoryArrow -event=squeeze_begin victims=[2,0,1] -event=squeeze_victim entry=2 +event=eviction_begin victims=[2,0,1] +event=eviction_victim entry=2 event=insert_success entry=2 kind=MemoryLiquid -event=squeeze_victim entry=0 +event=eviction_victim entry=0 event=io_write entry=0 kind=DiskLiquid bytes=1320 event=insert_success entry=0 kind=DiskLiquid -event=squeeze_victim entry=1 +event=eviction_victim entry=1 event=io_write entry=1 kind=DiskLiquid bytes=1320 event=insert_success entry=1 kind=DiskLiquid event=insert_success entry=3 kind=MemoryArrow event=insert_failed entry=4 kind=MemoryArrow -event=squeeze_begin victims=[3,2] -event=squeeze_victim entry=3 +event=eviction_begin victims=[3,2] +event=eviction_victim entry=3 event=insert_success entry=3 kind=MemoryLiquid -event=squeeze_victim entry=2 +event=eviction_victim entry=2 event=io_write entry=2 kind=DiskLiquid bytes=1320 event=insert_success entry=2 kind=DiskLiquid event=insert_success entry=4 kind=MemoryArrow diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap index 16dbb27bd..284f51c2b 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap @@ -1,16 +1,17 @@ --- source: src/core/src/cache/tests/policies.rs +assertion_line: 62 expression: trace --- EventTrace: [ event=insert_success entry=0 kind=MemoryArrow event=insert_failed entry=1 kind=MemoryArrow -event=squeeze_begin victims=[0] -event=squeeze_victim entry=0 +event=eviction_begin victims=[0] +event=eviction_victim entry=0 event=insert_success entry=0 kind=MemoryLiquid event=insert_failed entry=1 kind=MemoryArrow -event=squeeze_begin victims=[0] -event=squeeze_victim entry=0 +event=eviction_begin victims=[0] +event=eviction_victim entry=0 event=io_write entry=0 kind=DiskLiquid bytes=1320 event=insert_success entry=0 kind=DiskLiquid event=insert_failed entry=1 kind=MemoryArrow diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_date_time.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_date_time.snap deleted file mode 100644 index dbb4d543a..000000000 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_date_time.snap +++ /dev/null @@ -1,34 +0,0 @@ ---- -source: src/core/src/cache/tests/squeezed.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=0 kind=MemoryArrow -event=insert_success entry=1 kind=MemoryArrow -event=insert_failed entry=2 kind=MemoryArrow -event=squeeze_begin victims=[0,1] -event=squeeze_victim entry=0 -event=insert_success entry=0 kind=MemoryLiquid -event=squeeze_victim entry=1 -event=insert_success entry=1 kind=MemoryLiquid -event=insert_success entry=2 kind=MemoryArrow -event=insert_failed entry=3 kind=MemoryArrow -event=squeeze_begin victims=[2,0,1] -event=squeeze_victim entry=2 -event=insert_success entry=2 kind=MemoryLiquid -event=squeeze_victim entry=0 -event=io_write entry=0 kind=MemorySqueezedLiquid bytes=6184 -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=squeeze_victim entry=1 -event=io_write entry=1 kind=MemorySqueezedLiquid bytes=6184 -event=insert_success entry=1 kind=MemorySqueezedLiquid -event=insert_success entry=3 kind=MemoryArrow -event=read entry=0 selection=false expr=ExtractDate32[Year] cached=MemorySqueezedLiquid -event=read_squeezed_data entry=0 expression=ExtractDate32[Year] -event=read entry=1 selection=false expr=ExtractDate32[Year] cached=MemorySqueezedLiquid -event=read_squeezed_data entry=1 expression=ExtractDate32[Year] -event=read entry=2 selection=false expr=ExtractDate32[Year] cached=MemoryLiquid -event=read entry=3 selection=false expr=ExtractDate32[Year] cached=MemoryArrow -event=read entry=1 selection=false expr=ExtractDate32[Month] cached=MemorySqueezedLiquid -event=io_read_squeezed_backing entry=1 bytes=6184 -] diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_int64_array.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_int64_array.snap deleted file mode 100644 index d0ae4ddbc..000000000 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_int64_array.snap +++ /dev/null @@ -1,34 +0,0 @@ ---- -source: src/core/src/cache/tests/squeezed.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=0 kind=MemoryArrow -event=insert_success entry=1 kind=MemoryArrow -event=insert_failed entry=2 kind=MemoryArrow -event=squeeze_begin victims=[0,1] -event=squeeze_victim entry=0 -event=insert_success entry=0 kind=MemoryLiquid -event=squeeze_victim entry=1 -event=insert_success entry=1 kind=MemoryLiquid -event=insert_success entry=2 kind=MemoryArrow -event=insert_failed entry=3 kind=MemoryArrow -event=squeeze_begin victims=[2,0,1] -event=squeeze_victim entry=2 -event=insert_success entry=2 kind=MemoryLiquid -event=squeeze_victim entry=0 -event=io_write entry=0 kind=MemorySqueezedLiquid bytes=6184 -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=squeeze_victim entry=1 -event=io_write entry=1 kind=DiskLiquid bytes=6184 -event=insert_success entry=1 kind=DiskLiquid -event=insert_success entry=3 kind=MemoryArrow -event=read entry=0 selection=false expr=PredicateColumn cached=MemorySqueezedLiquid -event=io_read_squeezed_backing entry=0 bytes=6184 -event=read entry=1 selection=false expr=PredicateColumn cached=DiskLiquid -event=io_read_liquid entry=1 bytes=6184 -event=hydrate entry=1 cached=DiskLiquid new=MemoryLiquid -event=insert_success entry=1 kind=MemoryLiquid -event=read entry=2 selection=false expr=PredicateColumn cached=MemoryLiquid -event=read entry=3 selection=false expr=PredicateColumn cached=MemoryArrow -] diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_variant_path.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_variant_path.snap deleted file mode 100644 index 18756a3b4..000000000 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__squeezed__read_squeezed_variant_path.snap +++ /dev/null @@ -1,35 +0,0 @@ ---- -source: src/core/src/cache/tests/squeezed.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=0 kind=MemoryArrow -event=insert_failed entry=1 kind=MemoryArrow -event=squeeze_begin victims=[0] -event=squeeze_victim entry=0 -event=io_write entry=0 kind=MemorySqueezedLiquid bytes=7816 -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=insert_success entry=1 kind=MemoryArrow -event=insert_failed entry=2 kind=MemoryArrow -event=squeeze_begin victims=[1,0] -event=squeeze_victim entry=1 -event=io_write entry=1 kind=MemorySqueezedLiquid bytes=7816 -event=insert_success entry=1 kind=MemorySqueezedLiquid -event=squeeze_victim entry=0 -event=insert_success entry=0 kind=DiskArrow -event=insert_success entry=2 kind=MemoryArrow -event=read entry=0 selection=false expr=VariantGet[name:Utf8] cached=DiskArrow -event=io_read_arrow entry=0 bytes=7816 -event=hydrate entry=0 cached=DiskArrow new=MemorySqueezedLiquid -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=read entry=0 selection=false expr=VariantGet[age:Int64] cached=MemorySqueezedLiquid -event=io_read_arrow entry=0 bytes=7816 -event=hydrate entry=0 cached=MemorySqueezedLiquid new=MemorySqueezedLiquid -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=read_squeezed_data entry=0 expression=VariantGet[age:Int64] -event=read entry=1 selection=false expr=VariantGet[address.zipcode:Int64] cached=MemorySqueezedLiquid -event=io_read_arrow entry=1 bytes=7816 -event=hydrate entry=1 cached=MemorySqueezedLiquid new=MemorySqueezedLiquid -event=insert_success entry=1 kind=MemorySqueezedLiquid -event=read_squeezed_data entry=1 expression=VariantGet[address.zipcode:Int64] -] diff --git a/src/core/src/cache/tests/squeezed.rs b/src/core/src/cache/tests/squeezed.rs deleted file mode 100644 index ec46a4507..000000000 --- a/src/core/src/cache/tests/squeezed.rs +++ /dev/null @@ -1,195 +0,0 @@ -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, Date32Array, Int64Array, StringArray}; -use arrow_schema::DataType; -use parquet_variant_compute::json_to_variant; - -use crate::{ - cache::{ - AlwaysHydrate, CacheExpression, EntryID, LiquidCacheBuilder, LiquidPolicy, - TranscodeSqueezeEvict, - }, - liquid_array::Date32Field, -}; - -fn create_date32_array() -> ArrayRef { - let date32_array = Date32Array::from_iter_values(0..4096); - Arc::new(date32_array) -} - -#[tokio::test] -async fn read_squeezed_date_time() { - let temp_dir = tempfile::tempdir().unwrap(); - let array = create_date32_array(); - let array_size = array.get_array_memory_size(); - - let cache = LiquidCacheBuilder::new() - .with_cache_policy(Box::new(LiquidPolicy::new())) - .with_hydration_policy(Box::new(AlwaysHydrate::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) - .with_max_memory_bytes(array_size * 2) - .with_store( - t4::mount(temp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(), - ) - .build() - .await; - - let expression = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); - - for i in 0..4 { - let entry_id = EntryID::from(i); - cache - .insert(entry_id, array.clone()) - .with_squeeze_hint(expression.clone()) - .await - .unwrap(); - } - - for i in 0..4 { - let entry_id = EntryID::from(i); - let array = cache - .get(&entry_id) - .with_expression_hint(expression.clone()) - .await - .unwrap(); - assert_eq!(array.len(), array.len()); - } - cache - .get(&EntryID::from(1)) - .with_expression_hint(Arc::new(CacheExpression::extract_date32( - Date32Field::Month, - ))) - .await - .unwrap(); - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} - -fn create_variant_array() -> ArrayRef { - let mut values = Vec::new(); - for i in 0..64 { - if i % 2 == 0 { - values.push(Some(r#"{"name":"Ada", "address": {"zipcode": 90001}}"#)); - } else { - values.push(Some( - r#"{"name":"Bob", "age": 29, "address": {"city": "New York"}}"#, - )); - } - } - let json_values: ArrayRef = Arc::new(StringArray::from(values)); - let variant = json_to_variant(&json_values).expect("variant from json"); - ArrayRef::from(variant) -} - -#[tokio::test] -async fn read_squeezed_variant_path() { - let temp_dir = tempfile::tempdir().unwrap(); - let variant_array = create_variant_array(); - let array_size = variant_array.get_array_memory_size(); - - let cache = LiquidCacheBuilder::new() - .with_cache_policy(Box::new(LiquidPolicy::new())) - .with_hydration_policy(Box::new(AlwaysHydrate::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) - .with_max_memory_bytes(array_size * 3 / 2) - .with_store( - t4::mount(temp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(), - ) - .build() - .await; - - let name_expr = Arc::new(CacheExpression::variant_get("name", DataType::Utf8)); - let age_expr = Arc::new(CacheExpression::variant_get("age", DataType::Int64)); - let zipcode_expr = Arc::new(CacheExpression::variant_get( - "address.zipcode", - DataType::Int64, - )); - for i in 0..3 { - let entry_id = EntryID::from(i); - cache - .insert(entry_id, variant_array.clone()) - .with_squeeze_hint(name_expr.clone()) - .await - .unwrap(); - } - - let squeezed = cache - .get(&EntryID::from(0)) - .with_expression_hint(name_expr.clone()) - .read() - .await - .unwrap(); - assert_eq!(squeezed.len(), variant_array.len()); - - cache - .get(&EntryID::from(0)) - .with_expression_hint(age_expr.clone()) - .read() - .await - .unwrap(); - cache - .get(&EntryID::from(1)) - .with_expression_hint(zipcode_expr.clone()) - .read() - .await - .unwrap(); - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} - -fn create_int64_array() -> ArrayRef { - let int64_array = Int64Array::from_iter_values(0..4096); - Arc::new(int64_array) -} - -#[tokio::test] -async fn read_squeezed_int64_array() { - let temp_dir = tempfile::tempdir().unwrap(); - let int64_array = create_int64_array(); - let array_size = int64_array.get_array_memory_size(); - - let cache = LiquidCacheBuilder::new() - .with_cache_policy(Box::new(LiquidPolicy::new())) - .with_hydration_policy(Box::new(AlwaysHydrate::new())) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) - .with_max_memory_bytes(array_size * 2) - .with_store( - t4::mount(temp_dir.path().join("liquid_cache.t4")) - .await - .unwrap(), - ) - .build() - .await; - - let expression = Arc::new(CacheExpression::PredicateColumn); - - for i in 0..4 { - let entry_id = EntryID::from(i); - if i % 2 == 0 { - cache - .insert(entry_id, int64_array.clone()) - .with_squeeze_hint(expression.clone()) - .await - .unwrap(); - } else { - cache.insert(entry_id, int64_array.clone()).await.unwrap(); - } - } - - for i in 0..4 { - let entry_id = EntryID::from(i); - let array = cache - .get(&entry_id) - .with_expression_hint(expression.clone()) - .read() - .await - .unwrap(); - assert_eq!(array.len(), int64_array.len()); - } - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} diff --git a/src/core/src/cache/transcode.rs b/src/core/src/cache/transcode.rs index c3a4f07a6..6e862954c 100644 --- a/src/core/src/cache/transcode.rs +++ b/src/core/src/cache/transcode.rs @@ -46,7 +46,7 @@ pub fn transcode_liquid_inner<'a>( pub fn transcode_liquid_inner_with_hint<'a>( array: &'a ArrayRef, state: &LiquidCompressorStates, - squeeze_hint: Option<&CacheExpression>, + lineage: Option<&CacheExpression>, ) -> Result { let data_type = array.data_type(); if data_type.is_primitive() { @@ -162,7 +162,7 @@ pub fn transcode_liquid_inner_with_hint<'a>( } // Handle string/dictionary types. - let build_fingerprints = matches!(squeeze_hint, Some(CacheExpression::SubstringSearch)); + let build_fingerprints = matches!(lineage, Some(CacheExpression::SubstringSearch)); match array.data_type() { DataType::Utf8View => { let options = diff --git a/src/core/src/cache/utils.rs b/src/core/src/cache/utils.rs index a4d253435..4efd4eabe 100644 --- a/src/core/src/cache/utils.rs +++ b/src/core/src/cache/utils.rs @@ -55,14 +55,14 @@ pub(crate) async fn create_cache_store( max_memory_bytes: usize, policy: Box, ) -> Arc { - use crate::cache::{AlwaysHydrate, LiquidCacheBuilder, TranscodeSqueezeEvict}; + use crate::cache::{AlwaysHydrate, LiquidCacheBuilder, TranscodeEvict}; let batch_size = 128; let builder = LiquidCacheBuilder::new() .with_batch_size(batch_size) .with_max_memory_bytes(max_memory_bytes) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_hydration_policy(Box::new(AlwaysHydrate::new())) .with_cache_policy(policy); builder.build().await diff --git a/src/core/src/liquid_array/byte_view_array/comparisons.rs b/src/core/src/liquid_array/byte_view_array/comparisons.rs index 38b5efc7d..22c47f57f 100644 --- a/src/core/src/liquid_array/byte_view_array/comparisons.rs +++ b/src/core/src/liquid_array/byte_view_array/comparisons.rs @@ -14,7 +14,7 @@ use super::LiquidByteViewArray; use super::fingerprint::{StringFingerprint, substring_pattern_bytes}; use crate::liquid_array::byte_view_array::operator::{self, ByteViewOperator}; use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::{DiskBuffer, FsstBacking, PrefixKey}; +use crate::liquid_array::raw::fsst_buffer::{FsstBacking, PrefixKey}; impl LiquidByteViewArray { /// Compare equality with a byte needle @@ -183,133 +183,6 @@ impl LiquidByteViewArray { } } -impl LiquidByteViewArray { - pub(crate) async fn compare_with(&self, needle: &[u8], op: &ByteViewOperator) -> BooleanArray { - match op { - ByteViewOperator::Equality(operator::Equality::Eq) => self.compare_equals(needle).await, - ByteViewOperator::Equality(operator::Equality::NotEq) => { - self.compare_not_equals(needle).await - } - ByteViewOperator::Comparison(op) => self.compare_with_inner(needle, op).await, - ByteViewOperator::SubString(op) => { - let pattern = substring_pattern_bytes(needle).expect("Invalid substring pattern"); - let fingerprints = self - .string_fingerprints - .as_ref() - .expect("Fingerprints not initialized"); - self.compare_like_substring(pattern, *op, fingerprints) - .await - } - } - } - - /// Compare not equals with a byte needle - async fn compare_not_equals(&self, needle: &[u8]) -> BooleanArray { - let result = self.compare_equals(needle).await; - let (values, nulls) = result.into_parts(); - let values = !&values; - BooleanArray::new(values, nulls) - } - - /// Compare equality with a byte needle - pub(super) async fn compare_equals(&self, needle: &[u8]) -> BooleanArray { - let (mut dict_results, ambiguous) = self.compare_equals_with_prefix(needle); - if !ambiguous.is_empty() { - let bytes = self - .fsst_buffer - .squeeze_io() - .read(Some(self.fsst_buffer.disk_range())) - .await - .expect("read squeezed backing"); - let hydrated = LiquidByteViewArray::::from_bytes( - bytes, - self.fsst_buffer.compressor_arc(), - ); - let compressed_needle = compress_needle(hydrated.fsst_buffer.compressor(), needle); - - for &dict_index in ambiguous.iter() { - let compressed_value = hydrated.fsst_buffer.get_compressed_slice(dict_index); - if compressed_value == compressed_needle.as_slice() { - dict_results[dict_index] = true; - } - } - } else { - self.fsst_buffer.squeeze_io().trace_io_saved(); - } - - self.map_dictionary_results_to_array_results(dict_results) - } - - /// Prefix optimization for ordering operations - pub(super) async fn compare_with_inner( - &self, - needle: &[u8], - op: &operator::Comparison, - ) -> BooleanArray { - let (mut dict_results, ambiguous) = self.compare_with_prefix(needle, op); - - // For values needing full comparison, load buffer and decompress - if !ambiguous.is_empty() { - let (values_buffer, offsets_buffer) = - self.fsst_buffer.to_uncompressed_selected(&ambiguous).await; - let binary_array = - unsafe { BinaryArray::new_unchecked(offsets_buffer, values_buffer, None) }; - - for (pos, &dict_index) in ambiguous.iter().enumerate() { - let value_cmp = bytes_cmp_short_auto(binary_array.value(pos), needle); - let result = match (op, value_cmp) { - (operator::Comparison::Lt, std::cmp::Ordering::Less) => true, - (operator::Comparison::Lt, _) => false, - ( - operator::Comparison::LtEq, - std::cmp::Ordering::Less | std::cmp::Ordering::Equal, - ) => true, - (operator::Comparison::LtEq, _) => false, - (operator::Comparison::Gt, std::cmp::Ordering::Greater) => true, - (operator::Comparison::Gt, _) => false, - ( - operator::Comparison::GtEq, - std::cmp::Ordering::Greater | std::cmp::Ordering::Equal, - ) => true, - (operator::Comparison::GtEq, _) => false, - }; - dict_results[dict_index] = result; - } - } else { - self.fsst_buffer.squeeze_io().trace_io_saved(); - } - - self.map_dictionary_results_to_array_results(dict_results) - } - - pub(super) async fn compare_like_substring( - &self, - needle: &[u8], - operator: operator::SubString, - fingerprints: &Arc<[u32]>, - ) -> BooleanArray { - let (dict_results, ambiguous) = compute_fingerprint_candidates(needle, fingerprints); - - let dict_results = if !ambiguous.is_empty() { - let (values_buffer, offsets_buffer) = - self.fsst_buffer.to_uncompressed_selected(&ambiguous).await; - apply_like_match_on_candidates( - dict_results, - ambiguous, - values_buffer, - offsets_buffer, - needle, - operator, - ) - } else { - self.fsst_buffer.squeeze_io().trace_io_saved(); - dict_results - }; - - self.map_dictionary_results_to_array_results(dict_results) - } -} - impl LiquidByteViewArray { /// Return (selected_rows, ambiguous_rows, unique_rows) based on prefix-only comparison. pub fn prefix_compare_counts( @@ -405,65 +278,6 @@ impl LiquidByteViewArray { } // returns a tuple of compare_results and ambiguous indices - fn compare_equals_with_prefix(&self, needle: &[u8]) -> (Vec, Vec) { - let shared_prefix_len = self.shared_prefix.len(); - let num_unique = self.prefix_keys.len(); - if needle.len() < shared_prefix_len || needle[..shared_prefix_len] != self.shared_prefix { - return (vec![false; num_unique], Vec::new()); - } - - let needle_suffix = &needle[shared_prefix_len..]; - let needle_len = needle_suffix.len(); - let prefix_len = PrefixKey::prefix_len(); - - let mut dict_results = vec![false; num_unique]; - let mut ambiguous = Vec::new(); - - for (i, prefix_key) in self.prefix_keys.iter().enumerate().take(num_unique) { - let known_len = if prefix_key.len_byte() == 255 { - None - } else { - Some(prefix_key.len_byte() as usize) - }; - - // 1) Length gate - match known_len { - Some(l) => { - if l != needle_len { - continue; - } - } - None => { - if needle_len < 255 { - continue; - } - } - } - - // 2) Prefix classification - match known_len { - None => { - // Long strings: prefix match => need full comparison. - if prefix_key.prefix7()[..prefix_len] == needle_suffix[..prefix_len] { - ambiguous.push(i); - } - } - Some(l) if l <= prefix_len => { - // Small strings: exact compare on the known length. - if prefix_key.prefix7()[..l] == needle_suffix[..l] { - dict_results[i] = true; - } - } - Some(_l) => { - // Medium strings: prefix match => need full comparison. - if prefix_key.prefix7()[..prefix_len] == needle_suffix[..prefix_len] { - ambiguous.push(i); - } - } - } - } - (dict_results, ambiguous) - } /// Check if shared prefix comparison can short-circuit the entire operation fn compare_with_shared_prefix(&self, needle: &[u8], op: &operator::Comparison) -> Option { @@ -587,16 +401,6 @@ fn bytes_cmp_short(left: &[u8], right: &[u8], len: usize) -> std::cmp::Ordering } } -fn bytes_cmp_short_auto(left: &[u8], right: &[u8]) -> std::cmp::Ordering { - let len = left.len().min(right.len()); - let ordering = bytes_cmp_short(left, right, len); - if ordering == std::cmp::Ordering::Equal { - left.len().cmp(&right.len()) - } else { - ordering - } -} - /// Compute which dictionary entries are candidates for matching based on fingerprints. /// Returns a tuple of (dict_results, ambiguous_indices). fn compute_fingerprint_candidates( diff --git a/src/core/src/liquid_array/byte_view_array/fingerprint.rs b/src/core/src/liquid_array/byte_view_array/fingerprint.rs index 454edd530..bc6309d45 100644 --- a/src/core/src/liquid_array/byte_view_array/fingerprint.rs +++ b/src/core/src/liquid_array/byte_view_array/fingerprint.rs @@ -1,7 +1,3 @@ -use std::sync::Arc; - -use arrow::buffer::{Buffer, OffsetBuffer}; - const FINGERPRINT_BUCKETS: u8 = 32; const FINGERPRINT_MASK: u8 = FINGERPRINT_BUCKETS - 1; @@ -35,27 +31,6 @@ impl StringFingerprint { } } -pub(super) fn build_fingerprints(values: &Buffer, offsets: &OffsetBuffer) -> Arc<[u32]> { - let offsets = offsets.as_ref(); - if offsets.len() < 2 { - return Arc::from([]); - } - - let mut fingerprints = Vec::with_capacity(offsets.len().saturating_sub(1)); - let values = values.as_slice(); - - for window in offsets.windows(2) { - let start = window[0] as usize; - let end = window[1] as usize; - debug_assert!(start <= end); - debug_assert!(end <= values.len()); - let bytes = &values[start..end]; - fingerprints.push(StringFingerprint::from_bytes(bytes).0); - } - - Arc::from(fingerprints.into_boxed_slice()) -} - pub(super) fn substring_pattern_bytes(pattern: &[u8]) -> Option<&[u8]> { if pattern.len() < 2 { return None; diff --git a/src/core/src/liquid_array/byte_view_array/helpers.rs b/src/core/src/liquid_array/byte_view_array/helpers.rs index ff7906d4c..c9209a456 100644 --- a/src/core/src/liquid_array/byte_view_array/helpers.rs +++ b/src/core/src/liquid_array/byte_view_array/helpers.rs @@ -9,7 +9,7 @@ use super::LiquidByteViewArray; use super::operator::{ByteViewExpression, ByteViewOperator}; use crate::liquid_array::byte_view_array::operator::UnsupportedExpression; use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::{DiskBuffer, FsstBacking}; +use crate::liquid_array::raw::fsst_buffer::FsstBacking; pub(super) fn build_dict_selection( keys: &UInt16Array, @@ -91,37 +91,6 @@ pub(super) fn try_eval_predicate_in_memory( Some(array.compare_with(needle, op)) } -pub(super) async fn try_eval_predicate_on_disk( - expr: &Arc, - array: &LiquidByteViewArray, -) -> Option { - let cur_expr = match ByteViewExpression::try_from(expr) { - Ok(expr) => expr, - Err(UnsupportedExpression::Constant(v)) => { - let bool_array = if v { - BooleanBuffer::new_set(array.len()) - } else { - BooleanBuffer::new_unset(array.len()) - }; - return Some(BooleanArray::new(bool_array, array.nulls().cloned())); - } - Err(UnsupportedExpression::Expr) | Err(UnsupportedExpression::Op) => { - return None; - } - }; - - let op = cur_expr.op(); - let needle = cur_expr.literal(); - - if let ByteViewOperator::SubString(_substring_op) = op - && array.string_fingerprints.as_ref().is_none() - { - return None; - } - let result = array.compare_with(needle, op).await; - Some(result) -} - use std::fmt::Display; /// Detailed memory usage of the byte view array diff --git a/src/core/src/liquid_array/byte_view_array/mod.rs b/src/core/src/liquid_array/byte_view_array/mod.rs index 1d484148b..cb06173b2 100644 --- a/src/core/src/liquid_array/byte_view_array/mod.rs +++ b/src/core/src/liquid_array/byte_view_array/mod.rs @@ -7,21 +7,16 @@ use arrow::array::{ use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer}; use arrow::compute::cast; use arrow_schema::DataType; -use bytes::Bytes; use std::any::Any; use std::sync::Arc; #[cfg(test)] use std::cell::Cell; -use crate::cache::{CacheExpression, LiquidExpr}; -use crate::liquid_array::byte_view_array::fingerprint::build_fingerprints; +use crate::cache::LiquidExpr; use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::{DiskBuffer, FsstBacking, PrefixKey}; -use crate::liquid_array::{ - LiquidArray, LiquidDataType, LiquidSqueezedArray, LiquidSqueezedArrayRef, SqueezeIoHandler, - SqueezedBacking, eval_predicate_on_array, -}; +use crate::liquid_array::raw::fsst_buffer::{FsstBacking, PrefixKey}; +use crate::liquid_array::{LiquidArray, LiquidDataType, eval_predicate_on_array}; mod comparisons; mod conversions; @@ -295,41 +290,6 @@ impl LiquidByteViewArray { } } -impl LiquidByteViewArray { - /// Check if the FSST buffer is currently stored on disk - pub fn is_fsst_buffer_on_disk(&self) -> bool { - true - } - - /// Convert to Arrow DictionaryArray - pub async fn to_dict_arrow(&self) -> DictionaryArray { - if self.should_decompress_keyed() { - self.to_dict_arrow_decompress_keyed().await - } else { - self.to_dict_arrow_decompress_all().await - } - } - - async fn to_dict_arrow_decompress_all(&self) -> DictionaryArray { - let (values_buffer, offsets_buffer) = self.fsst_buffer.to_uncompressed().await; - self.to_dict_arrow_inner(self.dictionary_keys.clone(), values_buffer, offsets_buffer) - } - - async fn to_dict_arrow_decompress_keyed(&self) -> DictionaryArray { - let (selected, new_keys) = - helpers::build_dict_selection(&self.dictionary_keys, self.prefix_keys.len()); - let (values_buffer, offsets_buffer) = - self.fsst_buffer.to_uncompressed_selected(&selected).await; - self.to_dict_arrow_inner(new_keys, values_buffer, offsets_buffer) - } - - /// Convert to Arrow array with original type - pub async fn to_arrow_array(&self) -> ArrayRef { - let dict = self.to_dict_arrow().await; - cast(&dict, &self.original_arrow_type.to_arrow_type()).unwrap() - } -} - impl LiquidArray for LiquidByteViewArray { fn as_any(&self) -> &dyn Any { self @@ -362,7 +322,7 @@ impl LiquidArray for LiquidByteViewArray { } fn to_bytes(&self) -> Vec { - self.to_bytes_inner().expect("InMemoryFsstBuffer") + self.to_bytes_inner() } fn original_arrow_data_type(&self) -> DataType { @@ -373,129 +333,8 @@ impl LiquidArray for LiquidByteViewArray { LiquidDataType::ByteViewArray } - fn squeeze( - &self, - io: Arc, - squeeze_hint: Option<&CacheExpression>, - ) -> Option<(LiquidSqueezedArrayRef, Bytes)> { - squeeze_hint?; - - let string_fingerprints = if matches!(squeeze_hint, Some(CacheExpression::SubstringSearch)) - { - self.string_fingerprints.clone().or_else(|| { - let (values_buffer, offsets_buffer) = self.fsst_buffer.to_uncompressed(); - Some(build_fingerprints(&values_buffer, &offsets_buffer)) - }) - } else { - None - }; - - // Serialize full IPC bytes first - let bytes = match self.to_bytes_inner() { - Ok(b) => b, - Err(_) => return None, - }; - - // Build the hybrid (disk-backed FSST) view - let disk_range = 0u64..(bytes.len() as u64); - let compressor = self.fsst_buffer.compressor_arc(); - let disk = DiskBuffer::new( - self.fsst_buffer.uncompressed_bytes(), - io, - disk_range, - compressor, - ); - let hybrid = LiquidByteViewArray:: { - dictionary_keys: self.dictionary_keys.clone(), - prefix_keys: self.prefix_keys.clone(), - fsst_buffer: disk, - original_arrow_type: self.original_arrow_type, - shared_prefix: self.shared_prefix.clone(), - string_fingerprints, - }; - - let bytes = Bytes::from(bytes); - Some((Arc::new(hybrid) as LiquidSqueezedArrayRef, bytes)) - } - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { let filtered = helpers::filter_inner(self, selection); filtered.to_arrow_array() } } - -#[async_trait::async_trait] -impl LiquidSqueezedArray for LiquidByteViewArray { - /// Get the underlying any type. - fn as_any(&self) -> &dyn Any { - self - } - - /// Get the memory size of the Liquid array. - fn get_array_memory_size(&self) -> usize { - self.get_detailed_memory_usage().total() - } - - /// Get the length of the Liquid array. - fn len(&self) -> usize { - self.dictionary_keys.len() - } - - /// Check if the Liquid array is empty. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Convert the Liquid array to an Arrow array. - async fn to_arrow_array(&self) -> ArrayRef { - let bytes = self - .fsst_buffer - .squeeze_io() - .read(Some(self.fsst_buffer.disk_range())) - .await - .expect("read squeezed backing"); - let hydrated = - LiquidByteViewArray::::from_bytes(bytes, self.fsst_buffer.compressor_arc()); - LiquidByteViewArray::::to_arrow_array(&hydrated) - } - - /// Get the logical data type of the Liquid array. - fn data_type(&self) -> LiquidDataType { - LiquidDataType::ByteViewArray - } - - fn original_arrow_data_type(&self) -> DataType { - self.original_arrow_type.to_arrow_type() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Liquid(self.fsst_buffer.disk_range_len()) - } - - /// Filter the Liquid array with a boolean array and return an **arrow array**. - async fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let select_any = selection.count_set_bits() > 0; - if !select_any { - return arrow::array::new_empty_array(&self.original_arrow_data_type()); - } - let filtered = helpers::filter_inner(self, selection); - filtered.to_arrow_array().await - } - - /// Try to evaluate a predicate on the Liquid array with a filter. - /// Returns `Ok(None)` if the predicate is not supported. - /// - /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. - /// The returned boolean mask is nullable if the the original array is nullable. - async fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - // Reuse generic filter path first to reduce input rows if any - let filtered = helpers::filter_inner(self, filter); - if let Some(mask) = - helpers::try_eval_predicate_on_disk(expr.physical_expr(), &filtered).await - { - mask - } else { - eval_predicate_on_array(filtered.to_arrow_array().await, expr) - } - } -} diff --git a/src/core/src/liquid_array/byte_view_array/serialization.rs b/src/core/src/liquid_array/byte_view_array/serialization.rs index f37870a57..297dee2e5 100644 --- a/src/core/src/liquid_array/byte_view_array/serialization.rs +++ b/src/core/src/liquid_array/byte_view_array/serialization.rs @@ -4,12 +4,12 @@ use fsst::Compressor; use std::sync::Arc; use super::{ArrowByteType, LiquidByteViewArray}; +use crate::liquid_array::LiquidDataType; use crate::liquid_array::ipc::LiquidIPCHeader; use crate::liquid_array::raw::BitPackedArray; use crate::liquid_array::raw::fsst_buffer::{ FsstArray, PrefixKey, RawFsstBuffer, decode_compact_offsets, empty_compact_offsets, }; -use crate::liquid_array::{LiquidDataType, SqueezeResult}; // Header for LiquidByteViewArray serialization #[repr(C)] @@ -119,7 +119,7 @@ impl LiquidByteViewArray { | Optional string fingerprints (u32 per entry) | +--------------------------------------------------+ */ - pub(crate) fn to_bytes_inner(&self) -> SqueezeResult> { + pub(crate) fn to_bytes_inner(&self) -> Vec { let header_size = LiquidIPCHeader::size() + ByteViewArrayHeader::size(); let mut result = Vec::with_capacity(header_size + 1024); result.resize(header_size, 0); @@ -216,7 +216,7 @@ impl LiquidByteViewArray { header_slice[0..LiquidIPCHeader::size()].copy_from_slice(&ipc.to_bytes()); header_slice[LiquidIPCHeader::size()..header_size].copy_from_slice(&view_header.to_bytes()); - Ok(result) + result } /// Deserialize a LiquidByteViewArray from bytes. diff --git a/src/core/src/liquid_array/byte_view_array/tests.rs b/src/core/src/liquid_array/byte_view_array/tests.rs index 63267e34d..507912784 100644 --- a/src/core/src/liquid_array/byte_view_array/tests.rs +++ b/src/core/src/liquid_array/byte_view_array/tests.rs @@ -9,12 +9,12 @@ use rand::{RngExt as _, SeedableRng}; use std::sync::Arc; use crate::cache::transcode_liquid_inner_with_hint; -use crate::cache::{CacheExpression, LiquidCompressorStates, TestSqueezeIo}; +use crate::cache::{CacheExpression, LiquidCompressorStates}; use crate::liquid_array::byte_view_array::operator::{ ByteViewOperator, Comparison, Equality, SubString, }; -use crate::liquid_array::raw::fsst_buffer::{DiskBuffer, FsstArray, PrefixKey}; -use crate::liquid_array::{LiquidArray, LiquidDataType, LiquidSqueezedArray}; +use crate::liquid_array::raw::fsst_buffer::{FsstArray, PrefixKey}; +use crate::liquid_array::{LiquidArray, LiquidDataType}; #[test] fn test_dictionary_view_structure() { @@ -38,42 +38,6 @@ fn test_original_arrow_data_type_returns_utf8() { assert_eq!(array.original_arrow_data_type(), DataType::Utf8); } -#[test] -fn test_hybrid_original_arrow_data_type_returns_utf8() { - let input = StringArray::from(vec!["foo", "bar"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let in_memory = LiquidByteViewArray::::from_string_array(&input, compressor); - let (hybrid, _) = in_memory - .squeeze( - Arc::new(TestSqueezeIo::default()), - Some(&CacheExpression::PredicateColumn), - ) - .expect("squeeze should succeed"); - let disk_view = hybrid - .as_any() - .downcast_ref::>() - .expect("should downcast to disk array"); - assert_eq!(disk_view.original_arrow_data_type(), DataType::Utf8); -} - -#[test] -fn test_squeeze_builds_string_fingerprints() { - let input = StringArray::from(vec!["alpha", "beta", "alphabet"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let in_memory = LiquidByteViewArray::::from_string_array(&input, compressor); - let (hybrid, _) = in_memory - .squeeze( - Arc::new(TestSqueezeIo::default()), - Some(&CacheExpression::substring_search()), - ) - .expect("squeeze should succeed"); - let disk_view = hybrid - .as_any() - .downcast_ref::>() - .expect("should downcast to disk array"); - assert!(disk_view.string_fingerprints.is_some()); -} - #[test] fn test_ipc_roundtrip_preserves_string_fingerprints() { let input = StringArray::from(vec!["alpha", "beta", "alphabet"]); @@ -103,42 +67,6 @@ fn test_ipc_roundtrip_preserves_string_fingerprints() { ); } -#[tokio::test] -async fn test_string_fingerprint_skips_disk_read_for_impossible_substring() { - let input = StringArray::from(vec!["alpha", "ALP", "beta", "gamma"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let in_memory = LiquidByteViewArray::::from_string_array(&input, compressor); - - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = in_memory - .squeeze(io.clone(), Some(&CacheExpression::substring_search())) - .expect("squeeze should succeed"); - io.set_bytes(bytes); - - let disk_view = hybrid - .as_any() - .downcast_ref::>() - .expect("should downcast to disk array"); - - let fingerprints = disk_view - .string_fingerprints - .as_ref() - .expect("fingerprints should be present"); - let result = disk_view - .compare_like_substring(b"zzz", SubString::Contains, fingerprints) - .await; - let expected = BooleanArray::from(vec![false, false, false, false]); - assert_eq!(result, expected); - assert_eq!(io.reads(), 0); - - let result = disk_view - .compare_like_substring(b"alp", SubString::Contains, fingerprints) - .await; - let expected = BooleanArray::from(vec![true, false, false, false]); - assert_eq!(result, expected); - assert_eq!(io.reads(), 1); -} - #[test] fn test_ipc_roundtrip_sliced_dictionary_nulls() { let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); @@ -679,78 +607,6 @@ fn test_compare_with_prefix_optimization_utf8_and_binary() { assert_eq!(lte_result, lte_expected); } -fn test_compare_equals(input: StringArray, needle: &[u8], expected: BooleanArray) { - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - let result = liquid_array.compare_equals(needle); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_equals_on_disk() { - let input = StringArray::from(vec![ - Some("apple_orange"), - None, - Some("apple_orange_long_string"), - Some("apple_b"), - Some("apple_oo_long_string"), - Some("apple_b"), - Some("apple"), - ]); - test_compare_equals( - input.clone(), - b"apple", - BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(false), - Some(false), - Some(true), - ]), - ); - test_compare_equals( - input.clone(), - b"", - BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(false), - Some(false), - Some(false), - ]), - ); - test_compare_equals( - input.clone(), - b"apple_b", - BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(true), - Some(false), - Some(true), - Some(false), - ]), - ); - test_compare_equals( - input.clone(), - b"apple_oo_long_string", - BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(true), - Some(false), - Some(false), - ]), - ); -} - #[test] fn test_compare_equals_long_string_len_byte_255() { let common = "prefix_"; @@ -850,35 +706,6 @@ fn test_compare_with_like_fallback() { assert_eq!(result, expected); } -#[tokio::test] -async fn test_compare_equals_on_disk_long_prefix() { - let common = "prefix_"; - let long_len = 260; - let suffix_len = long_len - common.len(); - let long_a = format!("{}{}", common, "a".repeat(suffix_len)); - let long_b = format!("{}{}", common, "b".repeat(suffix_len)); - - let input = StringArray::from(vec![long_a.as_str(), long_b.as_str(), "z"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let in_memory = LiquidByteViewArray::::from_string_array(&input, compressor); - - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = in_memory - .squeeze(io.clone(), Some(&CacheExpression::PredicateColumn)) - .expect("squeeze should succeed"); - io.set_bytes(bytes); - - let disk_view = hybrid - .as_any() - .downcast_ref::>() - .expect("should downcast to disk array"); - - let result = disk_view.compare_equals(long_b.as_bytes()).await; - let expected = BooleanArray::from(vec![false, true, false]); - assert_eq!(result, expected); - assert_eq!(io.reads(), 1); -} - // Benchmark tests for v2 offset compression improvements fn generate_mixed_size_strings(count: usize, seed: u64) -> Vec { let mut rng = rand::rngs::StdRng::seed_from_u64(seed); diff --git a/src/core/src/liquid_array/decimal_array.rs b/src/core/src/liquid_array/decimal_array.rs index 374ff9cc7..6d462ee0b 100644 --- a/src/core/src/liquid_array/decimal_array.rs +++ b/src/core/src/liquid_array/decimal_array.rs @@ -1,29 +1,13 @@ -use std::any::Any; -use std::mem::size_of; -use std::num::NonZero; -use std::sync::Arc; +use bytes::Bytes; +use std::{any::Any, mem::size_of, sync::Arc}; -use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::array::{Array, ArrayRef, PrimitiveArray}; +use arrow::buffer::ScalarBuffer; use arrow::datatypes::{Decimal128Type, Decimal256Type, DecimalType, UInt64Type, i256}; use arrow_schema::DataType; -use bytes::Bytes; -use datafusion_common::ScalarValue; -use datafusion_expr_common::columnar_value::ColumnarValue; -use datafusion_expr_common::operator::Operator as DFOperator; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{ - BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, -}; -use datafusion_physical_expr_common::datum::apply_cmp; use num_traits::ToPrimitive; -use super::{ - LiquidArray, LiquidDataType, LiquidSqueezedArray, LiquidSqueezedArrayRef, NeedsBacking, - Operator, SqueezeIoHandler, SqueezeResult, SqueezedBacking, -}; -use crate::cache::{CacheExpression, LiquidExpr}; -use crate::liquid_array::eval_predicate_on_array; +use super::{LiquidArray, LiquidDataType}; use crate::liquid_array::ipc::{LiquidIPCHeader, get_physical_type_id}; use crate::liquid_array::raw::BitPackedArray; use crate::utils::get_bit_width; @@ -296,344 +280,12 @@ impl LiquidArray for LiquidDecimalArray { fn data_type(&self) -> LiquidDataType { LiquidDataType::Decimal } - - fn squeeze( - &self, - io: Arc, - expression_hint: Option<&CacheExpression>, - ) -> Option<(LiquidSqueezedArrayRef, Bytes)> { - let _expression_hint = expression_hint?; - let full_bytes = Bytes::from(self.to_bytes_inner()); - let disk_range = 0u64..(full_bytes.len() as u64); - - let orig_bw = self.bit_packed.bit_width()?; - if orig_bw.get() < 8 { - return None; - } - - let new_bw_u8 = NonZero::new((orig_bw.get() / 2).max(1)).unwrap(); - let unsigned_array = self.bit_packed.to_primitive(); - let (_dt, values, nulls) = unsigned_array.into_parts(); - - let max_offset = values.iter().copied().max().unwrap_or(0); - let bucket_count_u64 = 1u64 << (new_bw_u8.get() as u64); - let range_size = max_offset.saturating_add(1); - let bucket_width_u64 = (range_size.div_ceil(bucket_count_u64)).max(1); - - let quantized_values: ScalarBuffer = - ScalarBuffer::from_iter(values.iter().map(|&v| { - let mut idx_u64 = v / bucket_width_u64; - if idx_u64 >= bucket_count_u64 { - idx_u64 = bucket_count_u64 - 1; - } - idx_u64 - })); - let quantized_unsigned = PrimitiveArray::::new(quantized_values, nulls); - let quantized_bitpacked = BitPackedArray::from_primitive(quantized_unsigned, new_bw_u8); - - let hybrid = LiquidDecimalQuantizedArray { - quantized: quantized_bitpacked, - reference_value: self.reference_value, - bucket_width: bucket_width_u64, - disk_range, - io, - meta: self.meta, - }; - Some((Arc::new(hybrid) as LiquidSqueezedArrayRef, full_bytes)) - } -} - -#[derive(Debug, Clone)] -pub(crate) struct LiquidDecimalQuantizedArray { - quantized: BitPackedArray, - reference_value: u64, - bucket_width: u64, - disk_range: std::ops::Range, - io: Arc, - meta: DecimalMeta, -} - -impl LiquidDecimalQuantizedArray { - fn len(&self) -> usize { - self.quantized.len() - } - - fn new_from_filtered(&self, filtered: PrimitiveArray) -> Self { - let bit_width = self - .quantized - .bit_width() - .expect("quantized bit width must exist"); - let quantized = BitPackedArray::from_primitive(filtered, bit_width); - Self { - quantized, - reference_value: self.reference_value, - bucket_width: self.bucket_width, - disk_range: self.disk_range.clone(), - io: self.io.clone(), - meta: self.meta, - } - } - - fn filter_inner(&self, selection: &BooleanBuffer) -> Self { - let q_prim: PrimitiveArray = self.quantized.to_primitive(); - let selection = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::kernels::filter::filter(&q_prim, &selection).unwrap(); - let filtered = filtered.as_primitive::().clone(); - self.new_from_filtered(filtered) - } - - async fn hydrate_full_arrow(&self) -> ArrayRef { - let bytes = self - .io - .read(Some(self.disk_range.clone())) - .await - .expect("read squeezed backing"); - let liquid = crate::liquid_array::ipc::read_from_bytes( - bytes, - &crate::liquid_array::ipc::LiquidIPCContext::new(None), - ); - liquid.to_arrow_array() - } - - fn literal_to_u64(&self, literal: &Literal) -> Option { - match literal.value() { - ScalarValue::Decimal128(Some(v), _precision, scale) => { - if *scale != self.meta.scale { - return None; - } - v.to_u64() - } - ScalarValue::Decimal256(Some(v), _precision, scale) => { - if *scale != self.meta.scale { - return None; - } - v.to_u64() - } - _ => None, - } - } - - fn try_eval_predicate_inner( - &self, - op: &Operator, - literal: &Literal, - ) -> SqueezeResult> { - let k = match self.literal_to_u64(literal) { - Some(k) => k, - None => return Ok(None), - }; - - let q_prim = self.quantized.to_primitive(); - let (_dt, values, _nulls) = q_prim.into_parts(); - let nulls_opt = self.quantized.nulls(); - - let mut out_vals: Vec = Vec::with_capacity(values.len()); - - let push_const_for_below = |op: &Operator| -> bool { - match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => false, - Operator::LtEq => false, - Operator::Gt => true, - Operator::GtEq => true, - } - }; - - if k < self.reference_value { - let const_val = push_const_for_below(op); - if let Some(n) = nulls_opt { - for (i, _b) in values.iter().enumerate() { - out_vals.push(n.is_valid(i) && const_val); - } - } else { - out_vals.resize(values.len(), const_val); - } - } else { - let rel = k - self.reference_value; - let bw = self.bucket_width; - let q = rel / bw; - let r = rel % bw; - - let less_side: bool = matches!( - op, - Operator::Eq | Operator::NotEq | Operator::Lt | Operator::LtEq - ); - let greater_side: bool = matches!(op, Operator::NotEq | Operator::Gt | Operator::GtEq); - let on_equal_bucket = |r: u64, bw: u64| -> Option { - match op { - Operator::Eq | Operator::NotEq => None, - Operator::Lt => (r == 0).then_some(false), - Operator::LtEq => (r + 1 == bw).then_some(true), - Operator::Gt => (r + 1 == bw).then_some(false), - Operator::GtEq => (r == 0).then_some(true), - } - }; - - if let Some(n) = nulls_opt { - for (i, &b) in values.iter().enumerate() { - if !n.is_valid(i) { - out_vals.push(false); - continue; - } - let v = if b < q { - less_side - } else if b > q { - greater_side - } else { - match on_equal_bucket(r, bw) { - Some(val) => val, - None => return Err(NeedsBacking), - } - }; - out_vals.push(v); - } - } else { - for &b in values.iter() { - let v = if b < q { - less_side - } else if b > q { - greater_side - } else { - match on_equal_bucket(r, bw) { - Some(val) => val, - None => return Err(NeedsBacking), - } - }; - out_vals.push(v); - } - } - } - - let bool_buf = BooleanBuffer::from_iter(out_vals); - let out = BooleanArray::new(bool_buf, self.quantized.nulls().cloned()); - Ok(Some(out)) - } -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.quantized.get_array_memory_size() + size_of::() + size_of::() - } - - fn len(&self) -> usize { - LiquidDecimalQuantizedArray::len(self) - } - - async fn to_arrow_array(&self) -> ArrayRef { - self.hydrate_full_arrow().await - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Decimal - } - - fn original_arrow_data_type(&self) -> DataType { - self.meta.data_type() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Liquid((self.disk_range.end - self.disk_range.start) as usize) - } - - async fn try_eval_predicate( - &self, - liquid_expr: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - let filtered = self.filter_inner(filter); - - let expr = if let Some(expr) = unwrap_dynamic_filter(liquid_expr.physical_expr()) { - expr - } else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(binary_expr) = expr.downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - if !binary_expr.left().is::() { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - - let Some(literal) = binary_expr.right().downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - - let Some(op) = Operator::from_datafusion(binary_expr.op()) else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - match filtered.try_eval_predicate_inner(&op, literal) { - Ok(Some(mask)) => { - self.io.trace_io_saved(); - return mask; - } - Ok(None) => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - Err(NeedsBacking) => {} - } - - use arrow::array::cast::AsArray; - - let full = self.hydrate_full_arrow().await; - let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); - let filtered_len = filtered_arr.len(); - - let lhs = ColumnarValue::Array(filtered_arr); - let rhs = ColumnarValue::Scalar(literal.value().clone()); - let result = match binary_expr.op() { - DFOperator::NotEq => apply_cmp(DFOperator::NotEq, &lhs, &rhs), - DFOperator::Eq => apply_cmp(DFOperator::Eq, &lhs, &rhs), - DFOperator::Lt => apply_cmp(DFOperator::Lt, &lhs, &rhs), - DFOperator::LtEq => apply_cmp(DFOperator::LtEq, &lhs, &rhs), - DFOperator::Gt => apply_cmp(DFOperator::Gt, &lhs, &rhs), - DFOperator::GtEq => apply_cmp(DFOperator::GtEq, &lhs, &rhs), - _ => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() - } -} - -fn unwrap_dynamic_filter(expr: &Arc) -> Option> { - if let Some(dynamic_filter) = expr.downcast_ref::() { - dynamic_filter.current().ok() - } else { - Some(expr.clone()) - } } #[cfg(test)] mod tests { use super::*; - use crate::cache::{CacheExpression, TestSqueezeIo}; use arrow::array::Decimal128Builder; - use arrow::buffer::BooleanBuffer; - use datafusion_common::ScalarValue; - use datafusion_expr_common::operator::Operator as DFOperator; - use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; - use futures::executor::block_on; - use std::sync::Arc; #[test] fn decimal_u64_roundtrip() { @@ -661,33 +313,4 @@ mod tests { let arrow = decoded.to_arrow_array(); assert_eq!(arrow.as_ref(), &original); } - - #[test] - fn decimal_quantized_predicate_eval() { - let mut builder = Decimal128Builder::new(); - builder.append_value(100_i128); - builder.append_value(200_i128); - builder.append_null(); - builder.append_value(300_i128); - let original = builder.finish().with_precision_and_scale(10, 2).unwrap(); - - let liquid = LiquidDecimalArray::from_decimal_array(&original); - let hint = CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liquid.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - - let mask = BooleanBuffer::new_set(original.len()); - let lit = Arc::new(Literal::new(ScalarValue::Decimal128(Some(100_i128), 10, 2))); - let col = Arc::new(Column::new("col", 0)); - let expr: Arc = Arc::new(BinaryExpr::new(col, DFOperator::GtEq, lit)); - - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = BooleanArray::from(vec![Some(true), Some(true), None, Some(true)]); - assert_eq!(got, expected); - assert_eq!(io.reads(), 0); - } } diff --git a/src/core/src/liquid_array/float_array.rs b/src/core/src/liquid_array/float_array.rs index 683dece88..d3e175bf9 100644 --- a/src/core/src/liquid_array/float_array.rs +++ b/src/core/src/liquid_array/float_array.rs @@ -1,3 +1,4 @@ +use bytes::Bytes; /// /// Acknowledgement: /// The ALP compression implemented in this file is based on the Rust implementation available at https://github.com/spiraldb/alp @@ -5,42 +6,27 @@ use std::{ any::Any, fmt::Debug, - num::NonZero, ops::{Mul, Shl, Shr}, sync::Arc, }; use arrow::{ - array::{ - Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray, - PrimitiveArray, - }, - buffer::{BooleanBuffer, ScalarBuffer}, + array::{Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, PrimitiveArray}, + buffer::ScalarBuffer, datatypes::{ ArrowNativeType, Float32Type, Float64Type, Int32Type, Int64Type, UInt32Type, UInt64Type, }, }; use arrow_schema::DataType; -use datafusion_common::ScalarValue; -use datafusion_expr_common::columnar_value::ColumnarValue; -use datafusion_expr_common::operator::Operator as DFOperator; -use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; -use datafusion_physical_expr_common::datum::apply_cmp; use fastlanes::BitPacking; use num_traits::{AsPrimitive, Float, FromPrimitive}; use super::LiquidDataType; -use crate::cache::LiquidExpr; use crate::liquid_array::LiquidArray; +use crate::liquid_array::ipc::LiquidIPCHeader; use crate::liquid_array::ipc::{PhysicalTypeMarker, get_physical_type_id}; use crate::liquid_array::raw::BitPackedArray; -use crate::liquid_array::{ - LiquidSqueezedArray, LiquidSqueezedArrayRef, NeedsBacking, Operator, SqueezeResult, - SqueezedBacking, eval_predicate_on_array, ipc::LiquidIPCHeader, -}; use crate::utils::get_bit_width; -use crate::{cache::CacheExpression, liquid_array::SqueezeIoHandler}; -use bytes::Bytes; mod private { use arrow::{ @@ -57,13 +43,6 @@ mod private { const NUM_SAMPLES: usize = 1024; // we use FASTLANES to encode array, the sample size needs to be at least 1024 to get a good estimate of the best exponents -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum FloatSqueezePolicy { - /// Quantize values into buckets (good for coarse filtering; requires disk to recover values). - #[default] - Quantize = 0, -} - /// LiquidFloatType is a sealed trait that represents all the float types supported by Liquid. /// Implementors are Float32Type and Float64Type. TODO(): What about Float16Type, decimal types? pub trait LiquidFloatType: @@ -236,7 +215,6 @@ pub struct LiquidFloatArray { patch_indices: Vec, patch_values: Vec, reference_value: ::Native, - squeeze_policy: FloatSqueezePolicy, } impl LiquidFloatArray @@ -267,11 +245,6 @@ where let best_exponents = get_best_exponents::(&arrow_array); encode_arrow_array(&arrow_array, &best_exponents) } - - /// Get current squeeze policy for this array - pub fn squeeze_policy(&self) -> FloatSqueezePolicy { - self.squeeze_policy - } } impl LiquidArray for LiquidFloatArray @@ -334,60 +307,6 @@ where fn to_best_arrow_array(&self) -> ArrayRef { self.to_arrow_array() } - - fn squeeze( - &self, - io: Arc, - _expression_hint: Option<&CacheExpression>, - ) -> Option<(super::LiquidSqueezedArrayRef, bytes::Bytes)> { - let orig_bw = self.bit_packed.bit_width()?; - if orig_bw.get() < 8 { - return None; - } - - // New squeezed bit width is half of the original - let new_bw = orig_bw.get() / 2; - - let full_bytes = Bytes::from(self.to_bytes_inner()); - let disk_range = 0u64..(full_bytes.len() as u64); - - let (_dt, values, nulls) = self.bit_packed.to_primitive().into_parts(); - - match self.squeeze_policy { - FloatSqueezePolicy::Quantize => { - let shift = orig_bw.get() - new_bw; - let quantized_min = self.reference_value.shr(shift); - // let quantized_max = values - let quantized_values: ScalarBuffer< - ::Native, - > = ScalarBuffer::from_iter(values.iter().map(|&v| { - let signed_val: ::Native = v.as_(); - let v_signed = self.reference_value.add_wrapping(signed_val); - let v_quantized: ::Native = - v_signed.shr(shift); - v_quantized.sub_wrapping(quantized_min).as_() - })); - let quantized_array = - PrimitiveArray::<::UnsignedIntType>::new( - quantized_values, - nulls.clone(), - ); - let quantized_bitpacked = - BitPackedArray::from_primitive(quantized_array, NonZero::new(new_bw).unwrap()); - let hybrid = LiquidFloatQuantizedArray:: { - exponent: self.exponent, - quantized: quantized_bitpacked, - reference_value: self.reference_value, - bucket_width: shift, - disk_range, - io, - patch_indices: self.patch_indices.clone(), - patch_values: self.patch_values.clone(), - }; - Some((Arc::new(hybrid) as LiquidSqueezedArrayRef, full_bytes)) - } - } - } } impl LiquidFloatArray @@ -595,7 +514,6 @@ where patch_indices, patch_values, reference_value, - squeeze_policy: FloatSqueezePolicy::Quantize, } } } @@ -625,7 +543,6 @@ fn encode_arrow_array( patch_indices: Vec::new(), patch_values: Vec::new(), reference_value: ::Native::ZERO, - squeeze_policy: FloatSqueezePolicy::Quantize, }; } @@ -708,7 +625,6 @@ fn encode_arrow_array( patch_indices, patch_values, reference_value: min, - squeeze_policy: FloatSqueezePolicy::Quantize, } } @@ -739,324 +655,10 @@ fn get_best_exponents(arrow_array: &PrimitiveArray) -> Ex best_exponents } -#[derive(Debug)] -struct LiquidFloatQuantizedArray { - exponent: Exponents, - quantized: BitPackedArray, - reference_value: ::Native, - bucket_width: u8, // Width of each bucket (in bits) - disk_range: std::ops::Range, - io: Arc, - patch_indices: Vec, - patch_values: Vec, -} - -impl LiquidFloatQuantizedArray -where - T: LiquidFloatType, -{ - #[allow(dead_code)] - fn as_any(&self) -> &dyn Any { - self - } - - #[inline] - fn len(&self) -> usize { - self.quantized.len() - } - - fn new_from_filtered( - &self, - filtered: PrimitiveArray<::UnsignedIntType>, - ) -> Self { - let bit_width = self - .quantized - .bit_width() - .expect("quantized bit width must exist"); - let quantized = BitPackedArray::from_primitive(filtered, bit_width); - Self { - exponent: self.exponent, - quantized, - reference_value: self.reference_value, - bucket_width: self.bucket_width, - io: self.io.clone(), - patch_indices: self.patch_indices.clone(), - patch_values: self.patch_values.clone(), - disk_range: self.disk_range.clone(), - } - } - - fn filter_inner(&self, selection: &BooleanBuffer) -> Self { - let q_prim: PrimitiveArray = self.quantized.to_primitive(); - let selection = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::kernels::filter::filter(&q_prim, &selection).unwrap(); - let filtered = filtered.as_primitive::().clone(); - self.new_from_filtered(filtered) - } - - async fn hydrate_full_arrow(&self) -> ArrayRef { - let bytes = self - .io - .read(Some(self.disk_range.clone())) - .await - .expect("read squeezed backing"); - let liquid = crate::liquid_array::ipc::read_from_bytes( - bytes, - &crate::liquid_array::ipc::LiquidIPCContext::new(None), - ); - liquid.to_arrow_array() - } - - #[inline] - fn handle_eq(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k < lo || k > hi { Some(false) } else { None } - } - - #[inline] - fn handle_neq(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k < lo || k > hi { Some(true) } else { None } - } - - #[inline] - fn handle_lt(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k <= lo { - Some(false) - } else if hi < k { - Some(true) - } else { - None - } - } - - #[inline] - fn handle_lteq(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k < lo { - Some(false) - } else if hi <= k { - Some(true) - } else { - None - } - } - - #[inline] - fn handle_gt(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k < lo { - Some(true) - } else if hi <= k { - Some(false) - } else { - None - } - } - - #[inline] - fn handle_gteq(lo: T::Native, hi: T::Native, k: T::Native) -> Option { - if k <= lo { - Some(true) - } else if hi < k { - Some(false) - } else { - None - } - } - - fn try_eval_predicate_inner( - &self, - op: &Operator, - literal: &Literal, - ) -> SqueezeResult> { - // Extract scalar value as T::Native - let k_opt: Option = match literal.value() { - ScalarValue::Int8(Some(v)) => T::Native::from_i8(*v), - ScalarValue::Int16(Some(v)) => T::Native::from_i16(*v), - ScalarValue::Int32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Int64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::UInt8(Some(v)) => T::Native::from_u8(*v), - ScalarValue::UInt16(Some(v)) => T::Native::from_u16(*v), - ScalarValue::UInt32(Some(v)) => T::Native::from_u32(*v), - ScalarValue::UInt64(Some(v)) => T::Native::from_u64(*v), - ScalarValue::Date32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Date64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::Float32(Some(v)) => T::Native::from_f32(*v), - ScalarValue::Float64(Some(v)) => T::Native::from_f64(*v), - _ => None, - }; - let Some(k) = k_opt else { return Ok(None) }; - let q_prim = self.quantized.to_primitive(); - let (_dt, values, _nulls) = q_prim.into_parts(); - - let mut out_vals: Vec = Vec::with_capacity(values.len()); - let mut next_patch_index = 0; - let mut ignore_patches = false; - if self.patch_indices.is_empty() { - ignore_patches = true; - } - let comp_fn = match op { - Operator::Eq => Self::handle_eq, - Operator::NotEq => Self::handle_neq, - Operator::Lt => Self::handle_lt, - Operator::LtEq => Self::handle_lteq, - Operator::Gt => Self::handle_gt, - Operator::GtEq => Self::handle_gteq, - }; - // TODO(): This might not be very vectorization-friendly right now. Figure out optimizations - for (i, &b) in values.iter().enumerate() { - if let Some(nulls) = self.quantized.nulls() - && !nulls.is_valid(i) - { - out_vals.push(false); - continue; - } - if !ignore_patches && i as u64 == self.patch_indices[next_patch_index] { - next_patch_index += 1; - if next_patch_index == self.patch_indices.len() { - ignore_patches = true; - } - out_vals.push(false); - continue; - } - - let val: ::Native = b.as_(); - let lo = (val << self.bucket_width).add_wrapping(self.reference_value); - let hi = ((val.add_wrapping(1i32.into())) << self.bucket_width) - .add_wrapping(self.reference_value); - let val_lower = T::decode_single(&lo, &self.exponent); - let val_higher = T::decode_single(&hi, &self.exponent); - - let decided = comp_fn(val_lower, val_higher, k); - if let Some(v) = decided { - out_vals.push(v); - } else { - return Err(NeedsBacking); - } - } - - // Handle patches separately - // TODO(): Vectorize this - for (idx, patch_idx) in self.patch_indices.iter().enumerate() { - let patch_value = self.patch_values[idx]; - out_vals[*patch_idx as usize] = match op { - Operator::Eq => patch_value == k, - Operator::NotEq => patch_value != k, - Operator::Lt => patch_value < k, - Operator::LtEq => patch_value <= k, - Operator::Gt => patch_value > k, - Operator::GtEq => patch_value >= k, - } - } - - let bool_buf = arrow::buffer::BooleanBuffer::from_iter(out_vals); - let out = BooleanArray::new(bool_buf, self.quantized.nulls().cloned()); - Ok(Some(out)) - } -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for LiquidFloatQuantizedArray -where - T: LiquidFloatType, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.quantized.get_array_memory_size() - + size_of::() - + self.patch_indices.capacity() * size_of::() - + self.patch_values.capacity() * size_of::() - + size_of::<::Native>() - } - - fn len(&self) -> usize { - LiquidFloatQuantizedArray::::len(self) - } - - async fn to_arrow_array(&self) -> ArrayRef { - self.hydrate_full_arrow().await - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Float - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Liquid((self.disk_range.end - self.disk_range.start) as usize) - } - - async fn try_eval_predicate( - &self, - liquid_expr: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - // Apply selection first to reduce input rows - let filtered = self.filter_inner(filter); - let expr = liquid_expr.physical_expr(); - - if let Some(binary_expr) = expr.downcast_ref::() - && let Some(literal) = binary_expr.right().downcast_ref::() - { - let op = binary_expr.op(); - let supported_op = Operator::from_datafusion(op); - if let Some(supported_op) = supported_op { - match filtered.try_eval_predicate_inner(&supported_op, literal) { - Ok(Some(mask)) => return mask, - Ok(None) => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - Err(NeedsBacking) => {} - } - - // Fallback: hydrate full Arrow and evaluate predicate on filtered rows. - use arrow::array::cast::AsArray; - - let full = self.hydrate_full_arrow().await; - let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); - let filtered_len = filtered_arr.len(); - - let lhs = ColumnarValue::Array(filtered_arr); - let rhs = ColumnarValue::Scalar(literal.value().clone()); - let result = match op { - DFOperator::NotEq => apply_cmp(DFOperator::NotEq, &lhs, &rhs), - DFOperator::Eq => apply_cmp(DFOperator::Eq, &lhs, &rhs), - DFOperator::Lt => apply_cmp(DFOperator::Lt, &lhs, &rhs), - DFOperator::LtEq => apply_cmp(DFOperator::LtEq, &lhs, &rhs), - DFOperator::Gt => apply_cmp(DFOperator::Gt, &lhs, &rhs), - DFOperator::GtEq => apply_cmp(DFOperator::GtEq, &lhs, &rhs), - _ => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - return result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone(); - } - } - let fallback = self.filter(filter).await; - eval_predicate_on_array(fallback, liquid_expr) - } -} - #[cfg(test)] mod tests { - use datafusion_expr_common::operator::Operator; - use datafusion_physical_expr::PhysicalExpr; - use futures::executor::block_on; - use rand::{RngExt as _, SeedableRng as _, distr::uniform::SampleUniform, rngs::StdRng}; - use crate::cache::TestSqueezeIo; + use arrow::buffer::BooleanBuffer; use super::*; @@ -1208,230 +810,4 @@ mod tests { // Run for f64 run_compression_test::("f64", |i| i as f64); } - - // --------- Hybrid (squeeze) tests ---------- - fn make_f_array_with_range( - len: usize, - base_min: T::Native, - range: T::Native, - null_prob: f32, - rng: &mut StdRng, - ) -> PrimitiveArray - where - T: LiquidFloatType, - ::Native: SampleUniform, - PrimitiveArray: From::Native>>>, - { - let mut vals: Vec> = Vec::with_capacity(len); - for _ in 0..len { - if rng.random_bool(null_prob as f64) { - vals.push(None); - } else { - vals.push(Some(rng.random_range(base_min..(base_min + range)))); - } - } - PrimitiveArray::::from(vals) - } - - #[test] - fn hybrid_squeeze_unsqueezable_small_range() { - let mut rng = StdRng::seed_from_u64(0x51_71); - let arr = make_f_array_with_range::(64, 10_000.0, 100.0, 0.1, &mut rng); - let liquid = LiquidFloatArray::::from_arrow_array(arr); - assert!( - liquid - .squeeze(Arc::new(TestSqueezeIo::default()), None) - .is_none() - ); - } - - #[test] - fn hybrid_squeeze_full_read_roundtrip_f32() { - let mut rng = StdRng::seed_from_u64(0x51_72); - let arr = make_f_array_with_range::( - 2000, - -50_000.0, - (1 << 16) as f32, - 0.1, - &mut rng, - ); - let liq = LiquidFloatArray::::from_arrow_array(arr.clone()); - let bytes_baseline = liq.to_bytes(); - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), None).expect("squeezable"); - io.set_bytes(bytes.clone()); - // ensure we can recover the original by hydrating from full bytes - let recovered = LiquidFloatArray::::from_bytes(bytes.clone()); - assert_eq!( - recovered.to_arrow_array().as_primitive::(), - &arr - ); - assert_eq!(bytes_baseline, recovered.to_bytes()); - - let min = arrow::compute::kernels::aggregate::min(&arr).unwrap(); - let mask = BooleanBuffer::from(vec![true; arr.len()]); - let build_expr = |op: Operator, k: f32| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::Float32(Some(k)))); - Arc::new(BinaryExpr::new(lit.clone(), op, lit)) - }; - - // Expect resolvable results without IO - let resolvable_cases: Vec<(Operator, f32, bool)> = vec![ - (Operator::Eq, min - 1.0, false), // eq false everywhere - (Operator::NotEq, min - 1.0, true), // neq true everywhere - (Operator::Lt, min, false), // lt false everywhere - (Operator::LtEq, min - 1.0, false), // lte false everywhere - (Operator::Gt, min - 1.0, true), // gt true everywhere - (Operator::GtEq, min, true), // gte true everywhere - ]; - - for (op, k, expected_const) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(expected_const) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - // Unresolvable for Eq: pick a present value (ensures ambiguous bucket) - let k_present = (0..arr.len()) - .find_map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } - }) - .unwrap(); - let expr_eq_present = build_expr(Operator::Eq, k_present); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i) == k_present) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert!(io.reads() > 0); - assert_eq!(got, expected); - } - - #[test] - fn hybrid_squeeze_full_read_roundtrip_f64() { - let mut rng = StdRng::seed_from_u64(0x51_72); - let arr = make_f_array_with_range::( - 2000, - -50_000.0f64, - (1 << 16) as f64, - 0.1, - &mut rng, - ); - let liq = LiquidFloatArray::::from_arrow_array(arr.clone()); - let bytes_baseline = liq.to_bytes(); - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), None).expect("squeezable"); - io.set_bytes(bytes.clone()); - // ensure we can recover the original by hydrating from full bytes - let recovered = LiquidFloatArray::::from_bytes(bytes.clone()); - assert_eq!( - recovered.to_arrow_array().as_primitive::(), - &arr - ); - assert_eq!(bytes_baseline, recovered.to_bytes()); - - let min = arrow::compute::kernels::aggregate::min(&arr).unwrap(); - let mask = BooleanBuffer::from(vec![true; arr.len()]); - let build_expr = |op: Operator, k: f64| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::Float64(Some(k)))); - Arc::new(BinaryExpr::new(lit.clone(), op, lit)) - }; - - // Expect resolvable results without IO - let resolvable_cases: Vec<(Operator, f64, bool)> = vec![ - (Operator::Eq, min - 1.0, false), // eq false everywhere - (Operator::NotEq, min - 1.0, true), // neq true everywhere - (Operator::Lt, min, false), // lt false everywhere - (Operator::LtEq, min - 1.0, false), // lte false everywhere - (Operator::Gt, min - 1.0, true), // gt true everywhere - (Operator::GtEq, min, true), // gte true everywhere - ]; - - for (op, k, expected_const) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(expected_const) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - // Unresolvable for Eq: pick a present value (ensures ambiguous bucket) - let k_present = (0..arr.len()) - .find_map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } - }) - .unwrap(); - let expr_eq_present = build_expr(Operator::Eq, k_present); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i) == k_present) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert!(io.reads() > 0); - assert_eq!(got, expected); - } } diff --git a/src/core/src/liquid_array/hybrid_primitive_array.rs b/src/core/src/liquid_array/hybrid_primitive_array.rs deleted file mode 100644 index f2d0845bb..000000000 --- a/src/core/src/liquid_array/hybrid_primitive_array.rs +++ /dev/null @@ -1,1292 +0,0 @@ -use std::any::Any; -use std::sync::Arc; - -use arrow::array::ArrowNativeTypeOp; -use arrow::array::{ArrayRef, BooleanArray, PrimitiveArray, cast::AsArray}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; -use arrow::datatypes::{ArrowNativeType, ArrowPrimitiveType}; -use arrow_schema::{DataType, TimeUnit}; -use datafusion_common::ScalarValue; -use datafusion_expr_common::columnar_value::ColumnarValue; -use datafusion_expr_common::operator::Operator as DFOperator; -use datafusion_physical_expr::expressions::{ - BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, -}; -use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; -use datafusion_physical_expr_common::datum::apply_cmp; -use num_traits::{AsPrimitive, FromPrimitive}; - -use crate::cache::LiquidExpr; -use crate::liquid_array::eval_predicate_on_array; -use crate::liquid_array::raw::BitPackedArray; - -use super::primitive_array::LiquidPrimitiveType; -use super::{ - LiquidDataType, LiquidSqueezedArray, NeedsBacking, Operator, PrimitiveKind, SqueezeIoHandler, - SqueezeResult, SqueezedBacking, -}; - -#[derive(Clone, Copy)] -enum PredicateLhs { - Plain, - ToTimestampSeconds, -} - -fn unwrap_dynamic_filter(expr: &Arc) -> Option> { - if let Some(dynamic_filter) = expr.downcast_ref::() { - dynamic_filter.current().ok() - } else { - Some(expr.clone()) - } -} - -fn predicate_lhs_kind(expr: &Arc) -> Option { - if expr.is::() { - return Some(PredicateLhs::Plain); - } - if let Some(func) = expr.downcast_ref::() - && func.name() == "to_timestamp_seconds" - && let [arg] = func.args() - && arg.is::() - { - return Some(PredicateLhs::ToTimestampSeconds); - } - None -} - -fn can_eval_to_timestamp_seconds_direct() -> bool { - matches!( - T::DATA_TYPE, - DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::UInt8 - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::Timestamp(TimeUnit::Second, _) - ) -} - -#[derive(Debug, Clone)] -pub(crate) struct LiquidPrimitiveClampedArray { - pub(crate) squeezed: BitPackedArray, - pub(crate) reference_value: T::Native, - // Range in the on-disk payload needed to reconstruct the full array (we use full bytes) - pub(crate) disk_range: std::ops::Range, - pub(crate) io: Arc, -} - -impl LiquidPrimitiveClampedArray -where - T: LiquidPrimitiveType + PrimitiveKind, -{ - #[inline] - pub(crate) fn len(&self) -> usize { - self.squeezed.len() - } - - pub(crate) fn new_from_filtered( - &self, - filtered: PrimitiveArray<::UnSignedType>, - ) -> Self { - let bit_width = self - .squeezed - .bit_width() - .expect("squeezed bit width must exist"); - let squeezed = BitPackedArray::from_primitive(filtered, bit_width); - Self { - squeezed, - reference_value: self.reference_value, - disk_range: self.disk_range.clone(), - io: self.io.clone(), - } - } - - pub(crate) fn filter_inner(&self, selection: &BooleanBuffer) -> Self { - let unsigned_array: PrimitiveArray = self.squeezed.to_primitive(); - let selection = BooleanArray::new(selection.clone(), None); - let filtered_values = - arrow::compute::kernels::filter::filter(&unsigned_array, &selection).unwrap(); - let filtered_values = filtered_values.as_primitive::().clone(); - self.new_from_filtered(filtered_values) - } - - async fn hydrate_full_arrow(&self) -> ArrayRef { - let bytes = self - .io - .read(Some(self.disk_range.clone())) - .await - .expect("read squeezed backing"); - let liquid = crate::liquid_array::ipc::read_from_bytes( - bytes, - &crate::liquid_array::ipc::LiquidIPCContext::new(None), - ); - liquid.to_arrow_array() - } - - pub(crate) fn to_arrow_known_only(&self) -> Option { - // Convert squeezed to primitive and ensure no sentinel exists. - type U = <::UnSignedType as ArrowPrimitiveType>::Native; - let squeezed_prim = self.squeezed.to_primitive(); - let (_dt, values, nulls) = squeezed_prim.into_parts(); - let bw = self.squeezed.bit_width().expect("bit width").get(); - let sentinel: U = U::::usize_as((1usize << bw) - 1); - - // If any valid value equals sentinel, cannot fully materialize without disk - if let Some(n) = self.squeezed.nulls() { - for (i, v) in values.iter().enumerate() { - if n.is_valid(i) && *v == sentinel { - return None; - } - } - } else if values.contains(&sentinel) { - return None; - } - - // All values are known; reconstruct to full Arrow by adding reference - let ref_u: U = self.reference_value.as_(); - let restored_vals: ScalarBuffer = - ScalarBuffer::from_iter(values.iter().map(|&u| { - let t_val: T::Native = u.add_wrapping(ref_u).as_(); - t_val - })); - let arr = PrimitiveArray::::new(restored_vals, nulls); - Some(Arc::new(arr)) - } - - // Evaluate a simple comparison if fully decidable without disk; otherwise return Err(NeedsBacking) - pub(crate) fn try_eval_predicate_inner( - &self, - op: &Operator, - literal: &Literal, - ) -> SqueezeResult> { - // Extract scalar value as T::Native - let k_opt: Option = match literal.value() { - ScalarValue::Int8(Some(v)) => T::Native::from_i8(*v), - ScalarValue::Int16(Some(v)) => T::Native::from_i16(*v), - ScalarValue::Int32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Int64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::UInt8(Some(v)) => T::Native::from_u8(*v), - ScalarValue::UInt16(Some(v)) => T::Native::from_u16(*v), - ScalarValue::UInt32(Some(v)) => T::Native::from_u32(*v), - ScalarValue::UInt64(Some(v)) => T::Native::from_u64(*v), - ScalarValue::Date32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Date64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::TimestampSecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampMillisecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampMicrosecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampNanosecond(Some(v), _) => T::Native::from_i64(*v), - _ => None, - }; - let Some(k) = k_opt else { return Ok(None) }; - - // Prepare squeezed data and thresholds - type U = <::UnSignedType as ArrowPrimitiveType>::Native; - let squeezed_prim = self.squeezed.to_primitive(); - let (_dt, values, _nulls) = squeezed_prim.into_parts(); - let bw = self.squeezed.bit_width().expect("bit width").get(); - let sentinel: U = U::::usize_as((1usize << bw) - 1); - - // Precompute whether sentinel rows can be resolved under this operator and literal - let is_unsigned = ::IS_UNSIGNED; - let resolves_on_sentinel: bool = if is_unsigned { - let ref_u: U = self.reference_value.as_(); - let k_u: U = k.as_(); - let ref_u64: u64 = num_traits::AsPrimitive::::as_(ref_u); - let sent_u64: u64 = num_traits::AsPrimitive::::as_(sentinel); - let k_u64: u64 = num_traits::AsPrimitive::::as_(k_u); - let sent_abs: u64 = ref_u64 + sent_u64; - match op { - Operator::Eq | Operator::NotEq | Operator::Gt | Operator::LtEq => k_u64 < sent_abs, - Operator::Lt | Operator::GtEq => k_u64 <= sent_abs, - } - } else { - // signed types (including Date32/Date64) - let ref_i: i64 = self.reference_value.as_(); - let k_i: i64 = k.as_(); - let sent_abs: i64 = ref_i + (num_traits::AsPrimitive::::as_(sentinel) as i64); - match op { - Operator::Eq | Operator::NotEq | Operator::Gt | Operator::LtEq => k_i < sent_abs, - Operator::Lt | Operator::GtEq => k_i <= sent_abs, - } - }; - - // Build boolean values in a single pass; if an unresolved sentinel is seen, return IO range - let ref_u: U = self.reference_value.as_(); - let k_t: T::Native = k; - let mut out_vals: Vec = Vec::with_capacity(values.len()); - if let Some(n) = self.squeezed.nulls() { - for (i, &u) in values.iter().enumerate() { - if !n.is_valid(i) { - out_vals.push(false); - continue; - } - if u == sentinel { - if !resolves_on_sentinel { - return Err(NeedsBacking); - } - let b = match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => false, - Operator::LtEq => false, - Operator::Gt => true, - Operator::GtEq => true, - }; - out_vals.push(b); - } else { - let actual: T::Native = u.add_wrapping(ref_u).as_(); - let b = match op { - Operator::Eq => actual == k_t, - Operator::NotEq => actual != k_t, - Operator::Lt => actual < k_t, - Operator::LtEq => actual <= k_t, - Operator::Gt => actual > k_t, - Operator::GtEq => actual >= k_t, - }; - out_vals.push(b); - } - } - } else { - for &u in values.iter() { - if u == sentinel { - if !resolves_on_sentinel { - return Err(NeedsBacking); - } - let b = match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => false, - Operator::LtEq => false, - Operator::Gt => true, - Operator::GtEq => true, - }; - out_vals.push(b); - } else { - let actual: T::Native = u.add_wrapping(ref_u).as_(); - let b = match op { - Operator::Eq => actual == k_t, - Operator::NotEq => actual != k_t, - Operator::Lt => actual < k_t, - Operator::LtEq => actual <= k_t, - Operator::Gt => actual > k_t, - Operator::GtEq => actual >= k_t, - }; - out_vals.push(b); - } - } - } - - let bool_buf = BooleanBuffer::from_iter(out_vals); - let out = BooleanArray::new(bool_buf, self.squeezed.nulls().cloned()); - Ok(Some(out)) - } -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for LiquidPrimitiveClampedArray -where - T: LiquidPrimitiveType, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.squeezed.get_array_memory_size() + std::mem::size_of::() - } - - fn len(&self) -> usize { - LiquidPrimitiveClampedArray::::len(self) - } - - async fn to_arrow_array(&self) -> ArrayRef { - if let Some(arr) = self.to_arrow_known_only() { - return arr; - } - self.hydrate_full_arrow().await - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Integer - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Liquid((self.disk_range.end - self.disk_range.start) as usize) - } - - async fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - if selection.count_set_bits() == 0 { - return arrow::array::new_empty_array(&self.original_arrow_data_type()); - } - let filtered = self.filter_inner(selection); - if let Some(arr) = filtered.to_arrow_known_only() { - return arr; - } - let full = self.hydrate_full_arrow().await; - let selection_array = BooleanArray::new(selection.clone(), None); - arrow::compute::kernels::filter::filter(&full, &selection_array).unwrap() - } - - async fn try_eval_predicate( - &self, - liquid_expr: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - // Apply selection first to reduce input rows - let filtered = self.filter_inner(filter); - - let expr = if let Some(expr) = unwrap_dynamic_filter(liquid_expr.physical_expr()) { - expr - } else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(binary_expr) = expr.downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(lhs_kind) = predicate_lhs_kind(binary_expr.left()) else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(literal) = binary_expr.right().downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - - let op = binary_expr.op(); - let Some(supported_op) = Operator::from_datafusion(op) else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let can_eval_without_cast = match lhs_kind { - PredicateLhs::Plain => true, - PredicateLhs::ToTimestampSeconds => can_eval_to_timestamp_seconds_direct::(), - }; - if can_eval_without_cast { - match filtered.try_eval_predicate_inner(&supported_op, literal) { - Ok(Some(mask)) => { - self.io.trace_io_saved(); - return mask; - } - Ok(None) => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - Err(NeedsBacking) => {} - } - } - - // Fallback: hydrate full Arrow and evaluate predicate on filtered rows. - use arrow::array::cast::AsArray; - - let full = self.hydrate_full_arrow().await; - let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); - let filtered_len = filtered_arr.len(); - let lhs_array = match lhs_kind { - PredicateLhs::Plain => filtered_arr, - PredicateLhs::ToTimestampSeconds => { - let target_type = literal.value().data_type(); - arrow::compute::cast(&filtered_arr, &target_type) - .expect("to_timestamp_seconds cast must succeed") - } - }; - - let lhs = ColumnarValue::Array(lhs_array); - let rhs = ColumnarValue::Scalar(literal.value().clone()); - let result = match op { - DFOperator::NotEq => apply_cmp(DFOperator::NotEq, &lhs, &rhs), - DFOperator::Eq => apply_cmp(DFOperator::Eq, &lhs, &rhs), - DFOperator::Lt => apply_cmp(DFOperator::Lt, &lhs, &rhs), - DFOperator::LtEq => apply_cmp(DFOperator::LtEq, &lhs, &rhs), - DFOperator::Gt => apply_cmp(DFOperator::Gt, &lhs, &rhs), - DFOperator::GtEq => apply_cmp(DFOperator::GtEq, &lhs, &rhs), - _ => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() - } -} - -// Quantized hybrid array: stores bucket indices of value offsets -#[derive(Debug, Clone)] -pub(crate) struct LiquidPrimitiveQuantizedArray { - pub(crate) quantized: BitPackedArray, - pub(crate) reference_value: T::Native, - // bucket width in terms of absolute offset units - pub(crate) bucket_width: u64, - pub(crate) disk_range: std::ops::Range, - pub(crate) io: Arc, -} - -impl LiquidPrimitiveQuantizedArray -where - T: LiquidPrimitiveType + PrimitiveKind, -{ - #[inline] - pub(crate) fn len(&self) -> usize { - self.quantized.len() - } - - pub(crate) fn new_from_filtered( - &self, - filtered: PrimitiveArray<::UnSignedType>, - ) -> Self { - let bit_width = self - .quantized - .bit_width() - .expect("quantized bit width must exist"); - let quantized = BitPackedArray::from_primitive(filtered, bit_width); - Self { - quantized, - reference_value: self.reference_value, - bucket_width: self.bucket_width, - disk_range: self.disk_range.clone(), - io: self.io.clone(), - } - } - - pub(crate) fn filter_inner(&self, selection: &BooleanBuffer) -> Self { - let q_prim: PrimitiveArray = self.quantized.to_primitive(); - let selection = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::kernels::filter::filter(&q_prim, &selection).unwrap(); - let filtered = filtered.as_primitive::().clone(); - self.new_from_filtered(filtered) - } - - async fn hydrate_full_arrow(&self) -> ArrayRef { - let bytes = self - .io - .read(Some(self.disk_range.clone())) - .await - .expect("read squeezed backing"); - let liquid = crate::liquid_array::ipc::read_from_bytes( - bytes, - &crate::liquid_array::ipc::LiquidIPCContext::new(None), - ); - liquid.to_arrow_array() - } - - // Evaluate using bucket interval semantics; return Err if any ambiguous bucket is encountered - pub(crate) fn try_eval_predicate_inner( - &self, - op: &Operator, - literal: &Literal, - ) -> SqueezeResult> { - type U = <::UnSignedType as ArrowPrimitiveType>::Native; - - // Extract scalar value as T::Native - let k_opt: Option = match literal.value() { - ScalarValue::Int8(Some(v)) => T::Native::from_i8(*v), - ScalarValue::Int16(Some(v)) => T::Native::from_i16(*v), - ScalarValue::Int32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Int64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::UInt8(Some(v)) => T::Native::from_u8(*v), - ScalarValue::UInt16(Some(v)) => T::Native::from_u16(*v), - ScalarValue::UInt32(Some(v)) => T::Native::from_u32(*v), - ScalarValue::UInt64(Some(v)) => T::Native::from_u64(*v), - ScalarValue::Date32(Some(v)) => T::Native::from_i32(*v), - ScalarValue::Date64(Some(v)) => T::Native::from_i64(*v), - ScalarValue::TimestampSecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampMillisecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampMicrosecond(Some(v), _) => T::Native::from_i64(*v), - ScalarValue::TimestampNanosecond(Some(v), _) => T::Native::from_i64(*v), - _ => None, - }; - let Some(k) = k_opt else { return Ok(None) }; - - let q_prim = self.quantized.to_primitive(); - let (_dt, values, _nulls) = q_prim.into_parts(); - - let mut out_vals: Vec = Vec::with_capacity(values.len()); - let nulls_opt = self.quantized.nulls(); - - // Common fast-path constants when literal is below the minimum (reference) - let push_const_for_below = |op: &Operator| -> bool { - match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => false, - Operator::LtEq => false, - Operator::Gt => true, - Operator::GtEq => true, - } - }; - - // Minimal signed/unsigned split: only to compute below_ref and relative offset - let (below_ref, rel_opt): (bool, Option) = if T::IS_UNSIGNED { - let ref_u_native: U = self.reference_value.as_(); - let ref_u: u64 = num_traits::AsPrimitive::::as_(ref_u_native); - let k_u_native: U = k.as_(); - let k_u: u64 = num_traits::AsPrimitive::::as_(k_u_native); - if k_u < ref_u { - (true, None) - } else { - (false, Some(k_u - ref_u)) - } - } else { - let ref_i: i64 = self.reference_value.as_(); - let k_i: i64 = k.as_(); - if k_i < ref_i { - (true, None) - } else { - (false, Some((k_i - ref_i) as u64)) - } - }; - - if below_ref { - let const_val = push_const_for_below(op); - if let Some(n) = nulls_opt { - for (i, _b) in values.iter().enumerate() { - out_vals.push(n.is_valid(i) && const_val); - } - } else { - out_vals.resize(values.len(), const_val); - } - } else { - let rel = rel_opt.expect("rel must exist when not below_ref"); - let bw: u64 = self.bucket_width; - debug_assert!(bw > 0, "bucket_width must be > 0"); - let q = rel / bw; // target bucket index for k - let r = rel % bw; // position of k within its bucket - - // Precompute decisions outside the loop - let less_side: bool = match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => true, - Operator::LtEq => true, - Operator::Gt => false, - Operator::GtEq => false, - }; - let greater_side: bool = match op { - Operator::Eq => false, - Operator::NotEq => true, - Operator::Lt => false, - Operator::LtEq => false, - Operator::Gt => true, - Operator::GtEq => true, - }; - let on_equal_bucket = |r: u64, bw: u64| -> Option { - match op { - Operator::Eq | Operator::NotEq => None, - Operator::Lt => { - if r == 0 { - Some(false) - } else { - None - } - } - Operator::LtEq => { - if r + 1 == bw { - Some(true) - } else { - None - } - } - Operator::Gt => { - if r + 1 == bw { - Some(false) - } else { - None - } - } - Operator::GtEq => { - if r == 0 { - Some(true) - } else { - None - } - } - } - }; - - if let Some(n) = nulls_opt { - for (i, &b_native) in values.iter().enumerate() { - if !n.is_valid(i) { - out_vals.push(false); - continue; - } - let b: u64 = num_traits::AsPrimitive::::as_(b_native); - let v = if b < q { - less_side - } else if b > q { - greater_side - } else { - match on_equal_bucket(r, bw) { - Some(val) => val, - None => { - return Err(NeedsBacking); - } - } - }; - out_vals.push(v); - } - } else { - for &b_native in values.iter() { - let b: u64 = num_traits::AsPrimitive::::as_(b_native); - let v = if b < q { - less_side - } else if b > q { - greater_side - } else { - match on_equal_bucket(r, bw) { - Some(val) => val, - None => { - return Err(NeedsBacking); - } - } - }; - out_vals.push(v); - } - } - } - - let bool_buf = BooleanBuffer::from_iter(out_vals); - let out = BooleanArray::new(bool_buf, self.quantized.nulls().cloned()); - Ok(Some(out)) - } -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for LiquidPrimitiveQuantizedArray -where - T: LiquidPrimitiveType + PrimitiveKind, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.quantized.get_array_memory_size() + std::mem::size_of::() - } - - fn len(&self) -> usize { - LiquidPrimitiveQuantizedArray::::len(self) - } - - async fn to_arrow_array(&self) -> ArrayRef { - self.hydrate_full_arrow().await - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Integer - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Liquid((self.disk_range.end - self.disk_range.start) as usize) - } - - async fn try_eval_predicate( - &self, - liquid_expr: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - // Apply selection first to reduce input rows - let filtered = self.filter_inner(filter); - - let expr = if let Some(expr) = unwrap_dynamic_filter(liquid_expr.physical_expr()) { - expr - } else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(binary_expr) = expr.downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(lhs_kind) = predicate_lhs_kind(binary_expr.left()) else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let Some(literal) = binary_expr.right().downcast_ref::() else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - - let op = binary_expr.op(); - let Some(supported_op) = Operator::from_datafusion(op) else { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - }; - let can_eval_without_cast = match lhs_kind { - PredicateLhs::Plain => true, - PredicateLhs::ToTimestampSeconds => can_eval_to_timestamp_seconds_direct::(), - }; - if can_eval_without_cast { - match filtered.try_eval_predicate_inner(&supported_op, literal) { - Ok(Some(mask)) => { - self.io.trace_io_saved(); - return mask; - } - Ok(None) => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - Err(NeedsBacking) => {} - } - } - - // Fallback: hydrate full Arrow and evaluate predicate on filtered rows. - use arrow::array::cast::AsArray; - - let full = self.hydrate_full_arrow().await; - let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); - let filtered_len = filtered_arr.len(); - let lhs_array = match lhs_kind { - PredicateLhs::Plain => filtered_arr, - PredicateLhs::ToTimestampSeconds => { - let target_type = literal.value().data_type(); - arrow::compute::cast(&filtered_arr, &target_type) - .expect("to_timestamp_seconds cast must succeed") - } - }; - - let lhs = ColumnarValue::Array(lhs_array); - let rhs = ColumnarValue::Scalar(literal.value().clone()); - let result = match op { - DFOperator::NotEq => apply_cmp(DFOperator::NotEq, &lhs, &rhs), - DFOperator::Eq => apply_cmp(DFOperator::Eq, &lhs, &rhs), - DFOperator::Lt => apply_cmp(DFOperator::Lt, &lhs, &rhs), - DFOperator::LtEq => apply_cmp(DFOperator::LtEq, &lhs, &rhs), - DFOperator::Gt => apply_cmp(DFOperator::Gt, &lhs, &rhs), - DFOperator::GtEq => apply_cmp(DFOperator::GtEq, &lhs, &rhs), - _ => { - let fallback = self.filter(filter).await; - return eval_predicate_on_array(fallback, liquid_expr); - } - }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::cache::TestSqueezeIo; - use crate::liquid_array::LiquidArray; - use crate::liquid_array::primitive_array::{IntegerSqueezePolicy, LiquidPrimitiveArray}; - use crate::utils::get_bit_width; - use arrow::array::{Array, BooleanArray, PrimitiveArray}; - use arrow::buffer::BooleanBuffer; - use arrow::datatypes::{Int32Type, UInt32Type}; - use datafusion_common::ScalarValue; - use datafusion_expr_common::operator::Operator; - use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; - use futures::executor::block_on; - use rand::rngs::StdRng; - use rand::{RngExt as _, SeedableRng}; - use std::sync::Arc; - - // ---------- Hybrid (squeeze) tests ---------- - - fn make_i32_array_with_range( - len: usize, - base_min: i32, - range: i32, - null_prob: f32, - rng: &mut StdRng, - ) -> PrimitiveArray { - let mut vals: Vec> = Vec::with_capacity(len); - for _ in 0..len { - if rng.random_bool(null_prob as f64) { - vals.push(None); - } else { - let delta = rng.random_range(0..=range); - vals.push(Some(base_min.saturating_add(delta))); - } - } - PrimitiveArray::::from(vals) - } - - fn make_u32_array_with_range( - len: usize, - base_min: u32, - range: u32, - null_prob: f32, - rng: &mut StdRng, - ) -> PrimitiveArray { - let mut vals: Vec> = Vec::with_capacity(len); - for _ in 0..len { - if rng.random_bool(null_prob as f64) { - vals.push(None); - } else { - let delta = rng.random_range(0..=range); - vals.push(Some(base_min.saturating_add(delta))); - } - } - PrimitiveArray::::from(vals) - } - - fn compute_boundary_i32(arr: &PrimitiveArray) -> Option { - // boundary = min + ((1 << (bit_width(range)/2)) - 1) - let min = arrow::compute::kernels::aggregate::min(arr)?; - let max = arrow::compute::kernels::aggregate::max(arr)?; - let range = (max as i64 - min as i64) as u64; - let bw = get_bit_width(range); - let half = (bw.get() / 2) as u32; - let sentinel = if half == 0 { 0 } else { (1u64 << half) - 1 } as i64; - (min as i64 + sentinel).try_into().ok() - } - - fn compute_boundary_u32(arr: &PrimitiveArray) -> Option { - let min = arrow::compute::kernels::aggregate::min(arr)?; - let max = arrow::compute::kernels::aggregate::max(arr)?; - let range = (max as u128 - min as u128) as u64; - let bw = get_bit_width(range); - let half = (bw.get() / 2) as u32; - let sentinel = if half == 0 { 0 } else { (1u128 << half) - 1 } as u128; - let b = (min as u128 + sentinel) as u64 as u32; - Some(b) - } - - #[test] - fn clamp_unsqueezable_small_range() { - // range < 512 -> bit width < 10 => None - let mut rng = StdRng::seed_from_u64(0x51_71); - let arr = make_i32_array_with_range(64, 10_000, 100, 0.1, &mut rng); - let liquid = LiquidPrimitiveArray::::from_arrow_array(arr) - .with_squeeze_policy(IntegerSqueezePolicy::Clamp); - let hint = crate::cache::CacheExpression::PredicateColumn; - assert!( - liquid - .squeeze(Arc::new(TestSqueezeIo::default()), Some(&hint)) - .is_none() - ); - } - - #[test] - fn clamp_squeeze_full_read_roundtrip_i32() { - let mut rng = StdRng::seed_from_u64(0x51_72); - let arr = make_i32_array_with_range(128, -50_000, 1 << 16, 0.1, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Clamp); - let bytes_baseline = liq.to_bytes(); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes.clone()); - // ensure we can recover the original by hydrating from full bytes - let recovered = LiquidPrimitiveArray::::from_bytes(bytes.clone()); - assert_eq!(recovered.to_arrow_array().as_primitive::(), &arr); - assert_eq!(bytes_baseline, recovered.to_bytes()); - - // If we filter to only known values, hybrid can materialize without IO - let boundary = compute_boundary_i32(&arr).unwrap(); - let mask_bits: Vec = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - true - } else { - arr.value(i) < boundary - } - }) - .collect(); - let mask = BooleanBuffer::from_iter(mask_bits.iter().copied()); - io.reset_reads(); - let filtered_arrow = block_on(hybrid.filter(&mask)); - assert_eq!(io.reads(), 0); - - let expected = { - let vals: Vec> = (0..arr.len()) - .zip(mask_bits.iter()) - .filter(|&(_, &keep)| keep) - .map(|(i, &_keep)| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } - }) - .collect(); - PrimitiveArray::::from(vals) - }; - assert_eq!(filtered_arrow.as_primitive::(), &expected); - } - - #[test] - fn clamp_predicate_eval_i32_resolvable_and_unresolvable() { - let mut rng = StdRng::seed_from_u64(0x51_73); - let arr = make_i32_array_with_range(200, -1_000_000, 1 << 16, 0.2, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Clamp); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - - let boundary = compute_boundary_i32(&arr).unwrap(); - // selection mask: random subset - let mask_bits: Vec = (0..arr.len()).map(|_| rng.random()).collect(); - let mask = BooleanBuffer::from_iter(mask_bits.iter().copied()); - - let col = Arc::new(Column::new("col", 0)); - let build_expr = |op: Operator, k: i32| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::Int32(Some(k)))); - Arc::new(BinaryExpr::new(col.clone(), op, lit)) - }; - - // Helper to compute expected boolean array on selected rows - let expected_for = |op: Operator, k: i32| -> BooleanArray { - let vals: Vec> = (0..arr.len()) - .zip(mask_bits.iter()) - .filter(|&(_, &keep)| keep) - .map(|(i, &_keep)| { - if arr.is_null(i) { - None - } else { - let v = arr.value(i); - Some(match op { - Operator::Eq => v == k, - Operator::NotEq => v != k, - Operator::Lt => v < k, - Operator::LtEq => v <= k, - Operator::Gt => v > k, - Operator::GtEq => v >= k, - _ => unreachable!(), - }) - } - }) - .collect(); - BooleanArray::from(vals) - }; - - // Resolvable cases: K strictly less than boundary for Eq,Neq,LtEq,Gt; K <= boundary for Lt,GtEq - let resolvable_cases: Vec<(Operator, i32)> = vec![ - (Operator::Eq, boundary - 1), - (Operator::NotEq, boundary - 1), - (Operator::Lt, boundary), - (Operator::LtEq, boundary - 1), - (Operator::Gt, boundary - 1), - (Operator::GtEq, boundary), - ]; - - for (op, k) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = expected_for(op, k); - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - // Unresolvable: choose constants >= boundary for ops that require disk - let unresolvable_cases: Vec<(Operator, i32)> = vec![ - (Operator::Eq, boundary), - (Operator::NotEq, boundary), - (Operator::Lt, boundary + 1), - (Operator::LtEq, boundary), - (Operator::Gt, boundary + 1), - (Operator::GtEq, boundary + 1), - ]; - for (op, k) in unresolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = expected_for(op, k); - assert!(io.reads() > 0); - assert_eq!(got, expected); - } - } - - #[test] - fn clamp_predicate_eval_u32_resolvable_and_unresolvable() { - let mut rng = StdRng::seed_from_u64(0x51_74); - let arr = make_u32_array_with_range(180, 1_000_000, 1 << 16, 0.15, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Clamp); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - - let boundary = compute_boundary_u32(&arr).unwrap(); - let mask_bits: Vec = (0..arr.len()).map(|_| rng.random()).collect(); - let mask = BooleanBuffer::from_iter(mask_bits.iter().copied()); - - let col = Arc::new(Column::new("col", 0)); - let build_expr = |op: Operator, k: u32| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::UInt32(Some(k)))); - Arc::new(BinaryExpr::new(col.clone(), op, lit)) - }; - - let expected_for = |op: Operator, k: u32| -> BooleanArray { - let vals: Vec> = (0..arr.len()) - .zip(mask_bits.iter()) - .filter(|&(_, &keep)| keep) - .map(|(i, &_keep)| { - if arr.is_null(i) { - None - } else { - let v = arr.value(i); - Some(match op { - Operator::Eq => v == k, - Operator::NotEq => v != k, - Operator::Lt => v < k, - Operator::LtEq => v <= k, - Operator::Gt => v > k, - Operator::GtEq => v >= k, - _ => unreachable!(), - }) - } - }) - .collect(); - BooleanArray::from(vals) - }; - - let resolvable_cases: Vec<(Operator, u32)> = vec![ - (Operator::Eq, boundary - 1), - (Operator::NotEq, boundary - 1), - (Operator::Lt, boundary), - (Operator::LtEq, boundary - 1), - (Operator::Gt, boundary - 1), - (Operator::GtEq, boundary), - ]; - for (op, k) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = expected_for(op, k); - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - let unresolvable_cases: Vec<(Operator, u32)> = vec![ - (Operator::Eq, boundary), - (Operator::NotEq, boundary), - (Operator::Lt, boundary + 1), - (Operator::LtEq, boundary), - (Operator::Gt, boundary + 1), - (Operator::GtEq, boundary + 1), - ]; - for (op, k) in unresolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = expected_for(op, k); - assert!(io.reads() > 0); - assert_eq!(got, expected); - } - } - - #[test] - fn quantize_predicate_eval_u32_resolvable_and_unresolvable() { - let mut rng = StdRng::seed_from_u64(0x51_84); - let arr = make_u32_array_with_range(200, 1_000_000, 1 << 16, 0.2, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Quantize); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - - let min = arrow::compute::kernels::aggregate::min(&arr).unwrap(); - - let mask = BooleanBuffer::from(vec![true; arr.len()]); - let col = Arc::new(Column::new("col", 0)); - let build_expr = |op: Operator, k: u32| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::UInt32(Some(k)))); - Arc::new(BinaryExpr::new(col.clone(), op, lit)) - }; - - // Expect resolvable results without IO - let resolvable_cases: Vec<(Operator, u32, bool)> = vec![ - (Operator::Eq, min.saturating_sub(1), false), // eq false everywhere - (Operator::NotEq, min.saturating_sub(1), true), // neq true everywhere - (Operator::Lt, min, false), // lt false everywhere - (Operator::LtEq, min.saturating_sub(1), false), // lte false everywhere - (Operator::Gt, min.saturating_sub(1), true), // gt true everywhere - (Operator::GtEq, min, true), // gte true everywhere - ]; - for (op, k, expected_const) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(expected_const) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - // Unresolvable for Eq: pick a present value (ensures ambiguous bucket) - let k_present = (0..arr.len()) - .find_map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } - }) - .unwrap(); - let expr_eq_present = build_expr(Operator::Eq, k_present); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i) == k_present) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert!(io.reads() > 0); - assert_eq!(got, expected); - } - - #[test] - fn quantize_predicate_eval_i32_resolvable_and_unresolvable() { - let mut rng = StdRng::seed_from_u64(0x51_85); - let arr = make_i32_array_with_range(220, -1_000_000, 1 << 16, 0.2, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Quantize); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - - let min = arrow::compute::kernels::aggregate::min(&arr).unwrap(); - let mask = BooleanBuffer::from(vec![true; arr.len()]); - let col = Arc::new(Column::new("col", 0)); - let build_expr = |op: Operator, k: i32| -> Arc { - let lit = Arc::new(Literal::new(ScalarValue::Int32(Some(k)))); - Arc::new(BinaryExpr::new(col.clone(), op, lit)) - }; - - let resolvable_cases: Vec<(Operator, i32, bool)> = vec![ - (Operator::Eq, min - 1, false), // eq false everywhere - (Operator::NotEq, min - 1, true), - (Operator::Lt, min, false), - (Operator::LtEq, min - 1, false), - (Operator::Gt, min - 1, true), - (Operator::GtEq, min, true), - ]; - for (op, k, expected_const) in resolvable_cases { - let expr = build_expr(op, k); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(expected_const) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert_eq!(io.reads(), 0); - assert_eq!(got, expected); - } - - // Unresolvable for Eq: pick a present value - let k_present = (0..arr.len()) - .find_map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } - }) - .unwrap(); - let expr_eq_present = build_expr(Operator::Eq, k_present); - io.reset_reads(); - let got = block_on(hybrid.try_eval_predicate( - &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), - &mask, - )); - let expected = { - let vals: Vec> = (0..arr.len()) - .map(|i| { - if arr.is_null(i) { - None - } else { - Some(arr.value(i) == k_present) - } - }) - .collect(); - BooleanArray::from(vals) - }; - assert!(io.reads() > 0); - assert_eq!(got, expected); - } - - #[test] - fn quantize_to_arrow_is_err() { - let mut rng = StdRng::seed_from_u64(0x51_86); - let arr = make_u32_array_with_range(64, 1000, 1 << 12, 0.0, &mut rng); - let liq = LiquidPrimitiveArray::::from_arrow_array(arr.clone()) - .with_squeeze_policy(IntegerSqueezePolicy::Quantize); - let hint = crate::cache::CacheExpression::PredicateColumn; - let io = Arc::new(TestSqueezeIo::default()); - let (hybrid, bytes) = liq.squeeze(io.clone(), Some(&hint)).expect("squeezable"); - io.set_bytes(bytes); - io.reset_reads(); - let materialized = block_on(hybrid.to_arrow_array()); - assert!(io.reads() > 0); - assert_eq!(materialized.as_primitive::(), &arr); - } -} diff --git a/src/core/src/liquid_array/mod.rs b/src/core/src/liquid_array/mod.rs index 7776c7640..637016cb3 100644 --- a/src/core/src/liquid_array/mod.rs +++ b/src/core/src/liquid_array/mod.rs @@ -5,18 +5,15 @@ pub mod byte_view_array; mod decimal_array; mod fix_len_byte_array; mod float_array; -mod hybrid_primitive_array; pub mod ipc; mod linear_integer_array; mod primitive_array; pub mod raw; -mod squeezed_date32_array; #[cfg(test)] mod tests; pub(crate) mod utils; -mod variant_array; -use std::{any::Any, ops::Range, sync::Arc}; +use std::{any::Any, sync::Arc}; use arrow::{ array::{ArrayRef, BooleanArray, cast::AsArray}, @@ -25,8 +22,6 @@ use arrow::{ }; use arrow_schema::{DataType, Field, Schema}; pub use byte_view_array::LiquidByteViewArray; -use bytes::Bytes; -use datafusion_expr_common::operator::Operator as DFOperator; pub use decimal_array::LiquidDecimalArray; pub use fix_len_byte_array::LiquidFixedLenByteArray; pub use float_array::{LiquidFloat32Array, LiquidFloat64Array, LiquidFloatArray}; @@ -35,16 +30,26 @@ pub use linear_integer_array::{ LiquidLinearI16Array, LiquidLinearI32Array, LiquidLinearI64Array, LiquidLinearU8Array, LiquidLinearU16Array, LiquidLinearU32Array, LiquidLinearU64Array, }; -pub use primitive_array::IntegerSqueezePolicy; pub use primitive_array::{ LiquidDate32Array, LiquidDate64Array, LiquidI8Array, LiquidI16Array, LiquidI32Array, LiquidI64Array, LiquidPrimitiveArray, LiquidPrimitiveDeltaArray, LiquidPrimitiveType, LiquidU8Array, LiquidU16Array, LiquidU32Array, LiquidU64Array, }; -pub use squeezed_date32_array::{Date32Field, SqueezedDate32Array}; -pub use variant_array::VariantStructSqueezedArray; -use crate::cache::{CacheExpression, LiquidExpr}; +use crate::cache::LiquidExpr; + +/// A date or timestamp component observed by lineage analysis. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub enum Date32Field { + /// Year component. + Year, + /// Month component. + Month, + /// Day component. + Day, + /// Day of week, where Sunday is zero. + DayOfWeek, +} /// Liquid data type is only logical type #[derive(Debug, Clone, Copy)] @@ -128,141 +133,12 @@ pub trait LiquidArray: std::fmt::Debug + Send + Sync { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } - - /// Squeeze the Liquid array to a `LiquidHybridArrayRef` and a `bytes::Bytes`. - /// Return `None` if the Liquid array cannot be squeezed. - /// - /// This is the bridge from in-memory array to hybrid array. - /// Important: The returned `Bytes` is the data that is stored on disk, it is the same as to_bytes(). - /// - /// Hydrating the hybrid array from the stored bytes should yield the same `LiquidArray`. - fn squeeze( - &self, - _io: Arc, - _expression_hint: Option<&CacheExpression>, - ) -> Option<(LiquidSqueezedArrayRef, bytes::Bytes)> { - None - } } /// A reference to a Liquid array. pub type LiquidArrayRef = Arc; -/// On-disk backing for a squeezed array. -/// -/// Each variant carries the byte length of the persisted backing data, so the -/// cache can release the disk budget when the entry is evicted. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SqueezedBacking { - /// Bytes are stored using the Liquid IPC format. - Liquid(usize), - /// Bytes are stored using Arrow IPC (or another Arrow-compatible encoding). - Arrow(usize), -} - -impl SqueezedBacking { - /// Byte length of the backing data persisted on disk. - pub fn disk_bytes(&self) -> usize { - match self { - Self::Liquid(n) | Self::Arrow(n) => *n, - } - } -} - -/// A reference to a Liquid squeezed array. -pub type LiquidSqueezedArrayRef = Arc; - -/// Signals that the squeezed representation needs to be hydrated from disk. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct NeedsBacking; - -/// Result type for squeezed operations that may require disk hydration. -pub type SqueezeResult = Result; - -enum Operator { - Eq, - NotEq, - Lt, - LtEq, - Gt, - GtEq, -} - -impl Operator { - fn from_datafusion(op: &DFOperator) -> Option { - let op = match op { - DFOperator::Eq => Operator::Eq, - DFOperator::NotEq => Operator::NotEq, - DFOperator::Lt => Operator::Lt, - DFOperator::LtEq => Operator::LtEq, - DFOperator::Gt => Operator::Gt, - DFOperator::GtEq => Operator::GtEq, - _ => return None, - }; - Some(op) - } -} - -/// A Liquid squeezed array is a Liquid array that part of its data is stored on disk. -/// `LiquidSqueezedArray` is more complex than in-memory `LiquidArray` because it needs to handle IO. -#[async_trait::async_trait] -pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync { - /// Get the underlying any type. - fn as_any(&self) -> &dyn Any; - - /// Get the memory size of the Liquid array. - fn get_array_memory_size(&self) -> usize; - - /// Get the length of the Liquid array. - fn len(&self) -> usize; - - /// Check if the Liquid array is empty. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Convert the Liquid array to an Arrow array. - async fn to_arrow_array(&self) -> ArrayRef; - - /// Convert the Liquid array to an Arrow array. - /// Except that it will pick the best encoding for the arrow array. - /// Meaning that it may not obey the data type of the original arrow array. - async fn to_best_arrow_array(&self) -> ArrayRef { - self.to_arrow_array().await - } - - /// Get the logical data type of the Liquid array. - fn data_type(&self) -> LiquidDataType; - - /// Get the original arrow data type of the Liquid squeezed array. - fn original_arrow_data_type(&self) -> DataType; - - /// Filter the Liquid array with a boolean array and return an **arrow array**. - async fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let arrow_array = self.to_arrow_array().await; - let selection = BooleanArray::new(selection.clone(), None); - arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() - } - - /// Evaluate a predicate on the Liquid array with a filter. - /// - /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. - /// The returned boolean mask is nullable if the the original array is nullable. - async fn try_eval_predicate( - &self, - predicate: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - let filtered = self.filter(filter).await; - eval_predicate_on_array(filtered, predicate) - } - - /// Describe how the squeezed array persists its backing bytes on disk, - /// including the byte length of the persisted data. - fn disk_backing(&self) -> SqueezedBacking; -} - -pub(crate) fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { +fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), @@ -279,25 +155,6 @@ pub(crate) fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) - boolean_array.as_boolean().clone() } -/// A trait to read the backing bytes of a squeezed array from disk. -#[async_trait::async_trait] -pub trait SqueezeIoHandler: std::fmt::Debug + Send + Sync { - /// Read the backing bytes of a squeezed array from disk. - async fn read(&self, range: Option>) -> std::io::Result; - - /// Trace the number of decompressions performed. - // TODO: this is ugly. - fn tracing_decompress_count(&self, _decompress_cnt: usize, _total_cnt: usize) { - // Do nothing by default - } - - /// Trace the number of IO saved by squeezing. - // TODO: this is ugly. - fn trace_io_saved(&self) { - // Do nothing by default - } -} - /// Compile-time info about primitive kind (signed vs unsigned) and bounds. /// Implemented for all Liquid-supported primitive integer and date types. pub trait PrimitiveKind { diff --git a/src/core/src/liquid_array/primitive_array.rs b/src/core/src/liquid_array/primitive_array.rs index c651112b5..c1412618e 100644 --- a/src/core/src/liquid_array/primitive_array.rs +++ b/src/core/src/liquid_array/primitive_array.rs @@ -1,3 +1,4 @@ +use bytes::Bytes; use std::any::Any; use std::fmt::{Debug, Display}; use std::sync::Arc; @@ -16,30 +17,12 @@ use fastlanes::BitPacking; use num_traits::{AsPrimitive, FromPrimitive}; use super::LiquidDataType; -use crate::cache::{CacheExpression, LiquidExpr}; -use crate::liquid_array::hybrid_primitive_array::{ - LiquidPrimitiveClampedArray, LiquidPrimitiveQuantizedArray, -}; +use crate::cache::LiquidExpr; use crate::liquid_array::ipc::{LiquidIPCHeader, PhysicalTypeMarker, get_physical_type_id}; use crate::liquid_array::raw::BitPackedArray; -use crate::liquid_array::{ - LiquidArray, LiquidSqueezedArrayRef, PrimitiveKind, SqueezeIoHandler, SqueezedDate32Array, - eval_predicate_on_array, -}; +use crate::liquid_array::{LiquidArray, PrimitiveKind, eval_predicate_on_array}; use crate::utils::get_bit_width; use arrow::datatypes::ArrowNativeType; -use bytes::Bytes; - -/// Squeeze policy for primitive integer arrays. -/// Users can choose whether to clamp or quantize when squeezing. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum IntegerSqueezePolicy { - /// Clamp values above the squeezed range to a sentinel (recoverable for non-clamped rows). - Clamp = 0, - /// Quantize values into buckets (good for coarse filtering; requires disk to recover values). - #[default] - Quantize = 1, -} mod private { pub trait Sealed {} @@ -124,7 +107,6 @@ pub type LiquidDate64Array = LiquidPrimitiveArray; pub struct LiquidPrimitiveArray { bit_packed: BitPackedArray, reference_value: T::Native, - squeeze_policy: IntegerSqueezePolicy, } /// Liquid's primitive array which uses delta encoding for compression @@ -140,9 +122,7 @@ where { /// Get the memory size of the Liquid primitive array. pub fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() - + std::mem::size_of::() - + std::mem::size_of::() + self.bit_packed.get_array_memory_size() + std::mem::size_of::() } /// Get the length of the Liquid primitive array. @@ -164,7 +144,6 @@ where return Self { bit_packed: BitPackedArray::new_null_array(arrow_array.len()), reference_value: T::Native::ZERO, - squeeze_policy: IntegerSqueezePolicy::default(), }; } }; @@ -201,25 +180,8 @@ where Self { bit_packed: bit_packed_array, reference_value: min, - squeeze_policy: IntegerSqueezePolicy::default(), } } - - /// Get the current squeeze policy for this array. - pub fn squeeze_policy(&self) -> IntegerSqueezePolicy { - self.squeeze_policy - } - - /// Set the squeeze policy for this array. - pub fn set_squeeze_policy(&mut self, policy: IntegerSqueezePolicy) { - self.squeeze_policy = policy; - } - - /// Set the squeeze policy, returning self for chaining. - pub fn with_squeeze_policy(mut self, policy: IntegerSqueezePolicy) -> Self { - self.squeeze_policy = policy; - self - } } impl LiquidPrimitiveDeltaArray @@ -385,119 +347,6 @@ where fn data_type(&self) -> LiquidDataType { LiquidDataType::Integer } - - fn squeeze( - &self, - io: Arc, - expression_hint: Option<&CacheExpression>, - ) -> Option<(LiquidSqueezedArrayRef, Bytes)> { - let expression_hint = expression_hint?; - // Full bytes (original format) are what we store to disk - let full_bytes = Bytes::from(self.to_bytes_inner()); - let disk_range = 0u64..(full_bytes.len() as u64); - - if T::DATA_TYPE == DataType::Date32 { - // Special handle for Date32 arrays with component extraction support. - let field = expression_hint.as_date32_field()?; - let squeezed = - SqueezedDate32Array::from_liquid_date32(self, field).with_backing(io, disk_range); - return Some((Arc::new(squeezed) as LiquidSqueezedArrayRef, full_bytes)); - } - if matches!(T::DATA_TYPE, DataType::Timestamp(_, _)) { - let field = expression_hint.as_date32_field()?; - let squeezed = SqueezedDate32Array::from_liquid_timestamp(self, field) - .with_backing(io, disk_range); - return Some((Arc::new(squeezed) as LiquidSqueezedArrayRef, full_bytes)); - } - - // Only squeeze if we have a concrete bit width and it is large enough - let orig_bw = self.bit_packed.bit_width()?; - if orig_bw.get() < 8 { - return None; - } - - // New squeezed bit width is half of the original - let new_bw_u8 = std::num::NonZero::new((orig_bw.get() / 2).max(1)).unwrap(); - - // Decode original unsigned offsets - let unsigned_array = self.bit_packed.to_primitive(); - let (_dt, values, nulls) = unsigned_array.into_parts(); - - match self.squeeze_policy { - IntegerSqueezePolicy::Clamp => { - // Sentinel is the max representable value with new_bw bits - type U = - <::UnSignedType as ArrowPrimitiveType>::Native; - let sentinel: U = U::::usize_as((1usize << new_bw_u8.get()) - 1); - - // Clamp values to the squeezed width; values >= sentinel become sentinel - let squeezed_values: ScalarBuffer> = ScalarBuffer::from_iter( - values - .iter() - .map(|&v| if v >= sentinel { sentinel } else { v }), - ); - let squeezed_unsigned = - PrimitiveArray::<::UnSignedType>::new( - squeezed_values, - nulls, - ); - let squeezed_bitpacked = - BitPackedArray::from_primitive(squeezed_unsigned, new_bw_u8); - - let hybrid = LiquidPrimitiveClampedArray:: { - squeezed: squeezed_bitpacked, - reference_value: self.reference_value, - disk_range, - io: io.clone(), - }; - Some((Arc::new(hybrid) as LiquidSqueezedArrayRef, full_bytes)) - } - IntegerSqueezePolicy::Quantize => { - // Quantize value offsets into buckets of width W. - // Determine actual max offset value. - type U = - <::UnSignedType as ArrowPrimitiveType>::Native; - let max_offset: U = if let Some(m) = values.iter().copied().max() { - m - } else { - U::::ZERO - }; - - // Compute bucket count and width: ceil((max_offset+1)/bucket_count) - let bucket_count_u64 = 1u64 << (new_bw_u8.get() as u64); - let max_off_u64: u64 = num_traits::AsPrimitive::::as_(max_offset); - let range_size = max_off_u64.saturating_add(1); - let bucket_width_u64 = (range_size.div_ceil(bucket_count_u64)).max(1); - - let quantized_values: ScalarBuffer> = - ScalarBuffer::from_iter(values.iter().map(|&v| { - // v / bucket_width, clamped to last bucket - let v_u64: u64 = num_traits::AsPrimitive::::as_(v); - let mut idx_u64 = v_u64 / bucket_width_u64; - if idx_u64 >= bucket_count_u64 { - idx_u64 = bucket_count_u64 - 1; - } - U::::usize_as(idx_u64 as usize) - })); - let quantized_unsigned = - PrimitiveArray::<::UnSignedType>::new( - quantized_values, - nulls, - ); - let quantized_bitpacked = - BitPackedArray::from_primitive(quantized_unsigned, new_bw_u8); - - let hybrid = LiquidPrimitiveQuantizedArray:: { - quantized: quantized_bitpacked, - reference_value: self.reference_value, - bucket_width: bucket_width_u64, - disk_range, - io, - }; - Some((Arc::new(hybrid) as LiquidSqueezedArrayRef, full_bytes)) - } - } - } } impl LiquidArray for LiquidPrimitiveDeltaArray @@ -585,15 +434,6 @@ where fn data_type(&self) -> LiquidDataType { LiquidDataType::Integer } - - fn squeeze( - &self, - _io: Arc, - _expression_hint: Option<&CacheExpression>, - ) -> Option<(crate::liquid_array::LiquidSqueezedArrayRef, bytes::Bytes)> { - // Not implemented for delta arrays - None - } } impl LiquidPrimitiveArray @@ -674,7 +514,6 @@ where Self { bit_packed, reference_value, - squeeze_policy: IntegerSqueezePolicy::default(), } } } diff --git a/src/core/src/liquid_array/raw/bit_pack_array.rs b/src/core/src/liquid_array/raw/bit_pack_array.rs index 8da3339e1..4b8954cde 100644 --- a/src/core/src/liquid_array/raw/bit_pack_array.rs +++ b/src/core/src/liquid_array/raw/bit_pack_array.rs @@ -57,6 +57,7 @@ where self.nulls.as_ref() } + #[cfg(test)] pub(crate) fn bit_width(&self) -> Option> { self.bit_width } diff --git a/src/core/src/liquid_array/raw/fsst_buffer.rs b/src/core/src/liquid_array/raw/fsst_buffer.rs index 0789623ff..5ce63275b 100644 --- a/src/core/src/liquid_array/raw/fsst_buffer.rs +++ b/src/core/src/liquid_array/raw/fsst_buffer.rs @@ -9,11 +9,9 @@ use bytes; use fsst::{Compressor, Decompressor, Symbol}; use std::io::Result; use std::io::{Error, ErrorKind}; -use std::ops::Range; use std::sync::Arc; use crate::liquid_array::fix_len_byte_array::ArrowFixedLenByteArrayType; -use crate::liquid_array::{LiquidByteViewArray, SqueezeIoHandler}; mod sealed { pub trait Sealed {} @@ -21,7 +19,6 @@ mod sealed { /// Raw FSST buffer that stores compressed data using Arrow Buffer. /// Offsets are managed externally as a `u32` slice (including the final sentinel offset). -#[derive(Clone)] pub(crate) struct RawFsstBuffer { values: Buffer, uncompressed_bytes: usize, @@ -616,7 +613,7 @@ impl FsstArray { } /// FSST backing store for `LiquidByteViewArray` (in-memory or disk-only handle). pub trait FsstBacking: std::fmt::Debug + Clone + sealed::Sealed { - /// Get the uncompressed bytes of the FSST buffer (used for sizing / squeeze bookkeeping). + /// Get the uncompressed bytes of the FSST buffer. fn uncompressed_bytes(&self) -> usize; /// Get the in-memory size of the FSST backing (raw bytes + any in-memory indices). @@ -624,7 +621,6 @@ pub trait FsstBacking: std::fmt::Debug + Clone + sealed::Sealed { } impl sealed::Sealed for FsstArray {} -impl sealed::Sealed for DiskBuffer {} impl FsstArray { pub(crate) fn to_uncompressed(&self) -> (Buffer, OffsetBuffer) { @@ -676,90 +672,6 @@ impl FsstBacking for FsstArray { } } -/// Disk buffer for FSST buffer. -#[derive(Clone)] -pub struct DiskBuffer { - uncompressed_bytes: usize, - io: Arc, - disk_range: Range, - compressor: Arc, -} - -impl std::fmt::Debug for DiskBuffer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DiskBuffer") - .field("uncompressed_bytes", &self.uncompressed_bytes) - .field("disk_range", &self.disk_range) - .field("io", &self.io) - .field("compressor", &"") - .finish() - } -} - -impl DiskBuffer { - pub(crate) fn new( - uncompressed_bytes: usize, - io: Arc, - disk_range: Range, - compressor: Arc, - ) -> Self { - Self { - uncompressed_bytes, - io, - disk_range, - compressor, - } - } - - pub(crate) fn squeeze_io(&self) -> &Arc { - &self.io - } - - pub(crate) fn disk_range(&self) -> Range { - self.disk_range.clone() - } - - pub(crate) fn disk_range_len(&self) -> usize { - (self.disk_range.end - self.disk_range.start) as usize - } - - pub(crate) fn compressor_arc(&self) -> Arc { - self.compressor.clone() - } -} - -impl DiskBuffer { - pub(crate) async fn to_uncompressed(&self) -> (Buffer, OffsetBuffer) { - let bytes = self.io.read(Some(self.disk_range.clone())).await.unwrap(); - let byte_view = - LiquidByteViewArray::::from_bytes(bytes, self.compressor.clone()); - byte_view.fsst_buffer.to_uncompressed() - } - - pub(crate) async fn to_uncompressed_selected( - &self, - selected: &[usize], - ) -> (Buffer, OffsetBuffer) { - let bytes = self.io.read(Some(self.disk_range.clone())).await.unwrap(); - let byte_view = - LiquidByteViewArray::::from_bytes(bytes, self.compressor.clone()); - let total_count = byte_view.prefix_keys.len(); - self.io - .tracing_decompress_count(selected.len(), total_count); - byte_view.fsst_buffer.to_uncompressed_selected(selected) - } -} - -impl FsstBacking for DiskBuffer { - fn uncompressed_bytes(&self) -> usize { - self.uncompressed_bytes - } - - fn get_array_memory_size(&self) -> usize { - 0 - } -} - impl CompactOffsets { fn write_residuals(&self, out: &mut Vec) { out.extend_from_slice(&self.header.slope.to_le_bytes()); diff --git a/src/core/src/liquid_array/squeezed_date32_array.rs b/src/core/src/liquid_array/squeezed_date32_array.rs deleted file mode 100644 index e2497e556..000000000 --- a/src/core/src/liquid_array/squeezed_date32_array.rs +++ /dev/null @@ -1,748 +0,0 @@ -use arrow::array::{ - ArrayRef, BooleanArray, PrimitiveArray, - cast::AsArray, - types::{ - TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, - TimestampSecondType, - }, -}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; -use arrow::datatypes::{ArrowPrimitiveType, Date32Type, Int32Type, UInt32Type}; -use arrow_schema::{DataType, TimeUnit}; -use bytes::Bytes; -use num_traits::AsPrimitive; -use std::ops::Range; -use std::sync::Arc; - -use super::LiquidArray; -use super::primitive_array::LiquidPrimitiveArray; -use super::{LiquidDataType, LiquidSqueezedArray, SqueezedBacking}; -use crate::cache::LiquidExpr; -use crate::liquid_array::LiquidPrimitiveType; -use crate::liquid_array::SqueezeIoHandler; -use crate::liquid_array::eval_predicate_on_array; -use crate::liquid_array::raw::BitPackedArray; -use crate::utils::get_bit_width; - -/// Which component to extract from a `Date32`/Timestamp (days since UNIX epoch). -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -pub enum Date32Field { - /// Year component - Year, - /// Month component - Month, - /// Day component - Day, - /// Day-of-week component (Sunday=0). - DayOfWeek, -} - -/// A bit-packed array that stores a single extracted component (YEAR/MONTH/DAY/DOW) -/// from a `Date32`/Timestamp array. -/// -/// Values are stored as unsigned offsets from `reference_value`, using the same -/// bit-packing machinery as primitive arrays. -#[derive(Debug, Clone)] -pub struct SqueezedDate32Array { - field: Date32Field, - bit_packed: BitPackedArray, - /// The minimum extracted value used as reference for offsetting. - reference_value: i32, - original_data_type: DataType, - backing: Option, -} - -#[derive(Debug, Clone)] -struct DiskBacking { - io: Arc, - disk_range: Range, -} - -impl SqueezedDate32Array { - /// Build a squeezed representation (YEAR/MONTH/DAY/DAYOFWEEK) from a `Date32` array. - pub fn from_liquid_date32( - array: &LiquidPrimitiveArray, - field: Date32Field, - ) -> Self { - // Decode the logical Date32 array (i32: days since epoch) from the liquid array. - let arrow_array: PrimitiveArray = - array.to_arrow_array().as_primitive::().clone(); - - let (_dt, values, nulls) = arrow_array.into_parts(); - - // Compute min and max for the extracted component, skipping nulls. - let mut has_value = false; - let mut min_component: i32 = i32::MAX; - let mut max_component: i32 = i32::MIN; - - // Fast path: if all nulls, return a null bit-packed array of the same length. - if let Some(nulls_buf) = &nulls - && nulls_buf.null_count() == values.len() - { - return Self { - field, - bit_packed: BitPackedArray::new_null_array(values.len()), - reference_value: 0, - original_data_type: DataType::Date32, - backing: None, - }; - } - - for (idx, &days) in values.iter().enumerate() { - if let Some(nulls_buf) = &nulls - && nulls_buf.is_null(idx) - { - continue; - } - let comp = component_from_days(field, days); - has_value = true; - if comp < min_component { - min_component = comp; - } - if comp > max_component { - max_component = comp; - } - } - - // If no non-null values found, return an all-null structure (defensive) - if !has_value { - return Self { - field, - bit_packed: BitPackedArray::new_null_array(values.len()), - reference_value: 0, - original_data_type: DataType::Date32, - backing: None, - }; - } - - // Compute bit width from the value range. - let max_offset = (max_component as i64 - min_component as i64) as u64; - let bit_width = get_bit_width(max_offset); - - // Build unsigned offsets for packing; placeholders are fine for nulls. - let offsets: ScalarBuffer<::Native> = - ScalarBuffer::from_iter((0..values.len()).map(|idx| { - if nulls.as_ref().is_some_and(|n| n.is_null(idx)) { - 0u32 - } else { - let comp = component_from_days(field, values[idx]); - (comp - min_component) as u32 - } - })); - - let unsigned_array = PrimitiveArray::::new(offsets, nulls); - let bit_packed = BitPackedArray::from_primitive(unsigned_array, bit_width); - - Self { - field, - bit_packed, - reference_value: min_component, - original_data_type: DataType::Date32, - backing: None, - } - } - - /// Build a squeezed representation (YEAR/MONTH/DAY/DAYOFWEEK) from a timestamp array. - pub fn from_liquid_timestamp( - array: &LiquidPrimitiveArray, - field: Date32Field, - ) -> Self { - let unit = timestamp_unit(&T::DATA_TYPE).expect("timestamp data type"); - let arrow_array: PrimitiveArray = array.to_arrow_array().as_primitive::().clone(); - let (_dt, values, nulls) = arrow_array.into_parts(); - - let mut has_value = false; - let mut min_component: i32 = i32::MAX; - let mut max_component: i32 = i32::MIN; - - if let Some(nulls_buf) = &nulls - && nulls_buf.null_count() == values.len() - { - return Self { - field, - bit_packed: BitPackedArray::new_null_array(values.len()), - reference_value: 0, - original_data_type: T::DATA_TYPE.clone(), - backing: None, - }; - } - - for (idx, &value) in values.iter().enumerate() { - if let Some(nulls_buf) = &nulls - && nulls_buf.is_null(idx) - { - continue; - } - let days = timestamp_to_days_since_epoch(value.as_(), unit); - let comp = component_from_days(field, days); - has_value = true; - if comp < min_component { - min_component = comp; - } - if comp > max_component { - max_component = comp; - } - } - - if !has_value { - return Self { - field, - bit_packed: BitPackedArray::new_null_array(values.len()), - reference_value: 0, - original_data_type: T::DATA_TYPE.clone(), - backing: None, - }; - } - - let max_offset = (max_component as i64 - min_component as i64) as u64; - let bit_width = get_bit_width(max_offset); - - let offsets: ScalarBuffer<::Native> = - ScalarBuffer::from_iter((0..values.len()).map(|idx| { - if nulls.as_ref().is_some_and(|n| n.is_null(idx)) { - 0u32 - } else { - let days = timestamp_to_days_since_epoch(values[idx].as_(), unit); - let comp = component_from_days(field, days); - (comp - min_component) as u32 - } - })); - - let unsigned_array = PrimitiveArray::::new(offsets, nulls); - let bit_packed = BitPackedArray::from_primitive(unsigned_array, bit_width); - - Self { - field, - bit_packed, - reference_value: min_component, - original_data_type: T::DATA_TYPE.clone(), - backing: None, - } - } - - pub(crate) fn with_backing( - mut self, - io: Arc, - disk_range: Range, - ) -> Self { - self.backing = Some(DiskBacking { io, disk_range }); - self - } - - async fn read_backing(&self) -> Bytes { - let backing = self - .backing - .as_ref() - .expect("SqueezedDate32Array backing not set"); - backing - .io - .read(Some(backing.disk_range.clone())) - .await - .expect("read squeezed backing") - } - - /// Length of the array. - pub fn len(&self) -> usize { - self.bit_packed.len() - } - - /// Whether the array has no elements. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Memory size of the bit-packed representation plus reference value. - pub fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() + std::mem::size_of::() - } - - /// The extracted component type. - pub fn field(&self) -> Date32Field { - self.field - } - - /// Convert to an Arrow array shaped like the original input, encoded so that - /// re-applying `date_part` (or any equivalent extraction) recovers the - /// component value originally squeezed. - pub fn to_component_array(&self) -> ArrayRef { - match &self.original_data_type { - DataType::Date32 => Arc::new(self.to_arrow_date32_lossy()) as ArrayRef, - DataType::Timestamp(unit, _) => self.to_arrow_timestamp_lossy(*unit), - _ => Arc::new(self.to_arrow_date32_lossy()) as ArrayRef, - } - } - - /// Convert back to an Arrow `Int32` array representing the extracted component values. - /// Useful for verification or future pushdown logic. - pub fn to_component_date32(&self) -> PrimitiveArray { - let unsigned: PrimitiveArray = self.bit_packed.to_primitive(); - let (_dt, values, nulls) = unsigned.into_parts(); - let ref_v = self.reference_value; - let signed_values: ScalarBuffer<::Native> = - ScalarBuffer::from_iter(values.iter().map(|&v| (v as i32).saturating_add(ref_v))); - PrimitiveArray::::new(signed_values, nulls) - } - - /// Lossy reconstruction to Arrow Timestamp at the requested unit, using the - /// same date mapping as [`Self::to_arrow_date32_lossy`] (midnight UTC of the - /// reconstructed date). - pub fn to_arrow_timestamp_lossy(&self, unit: TimeUnit) -> ArrayRef { - let date_array = self.to_arrow_date32_lossy(); - let (_dt, day_values, nulls) = date_array.into_parts(); - let ticks_per_day: i64 = match unit { - TimeUnit::Second => 86_400, - TimeUnit::Millisecond => 86_400_000, - TimeUnit::Microsecond => 86_400_000_000, - TimeUnit::Nanosecond => 86_400_000_000_000, - }; - let tick_values: ScalarBuffer = - ScalarBuffer::from_iter(day_values.iter().map(|&d| (d as i64) * ticks_per_day)); - match unit { - TimeUnit::Second => Arc::new(PrimitiveArray::::new( - tick_values, - nulls, - )), - TimeUnit::Millisecond => Arc::new(PrimitiveArray::::new( - tick_values, - nulls, - )), - TimeUnit::Microsecond => Arc::new(PrimitiveArray::::new( - tick_values, - nulls, - )), - TimeUnit::Nanosecond => Arc::new(PrimitiveArray::::new( - tick_values, - nulls, - )), - } - } - - /// Lossy reconstruction to Arrow `Date32` (days since epoch). - /// - /// Mapping used: - /// - Year: (year, 1, 1) - /// - Month: (1970, month, 1) - /// - Day: (1970, 1, day) - /// - DayOfWeek: (1970, 1, 4 + dow) where 1970-01-04 is Sunday - pub fn to_arrow_date32_lossy(&self) -> PrimitiveArray { - let unsigned: PrimitiveArray = self.bit_packed.to_primitive(); - let (_dt, values, nulls) = unsigned.into_parts(); - - let ref_v = self.reference_value; - let days_values: ScalarBuffer<::Native> = - ScalarBuffer::from_iter(values.iter().enumerate().map(|(i, &off)| { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - 0i32 - } else { - match self.field { - Date32Field::Year => { - let y = ref_v + off as i32; - ymd_to_epoch_days(y, 1, 1) - } - Date32Field::Month => { - let m = (ref_v + off as i32) as u32; - ymd_to_epoch_days(1970, m, 1) - } - Date32Field::Day => { - let d = (ref_v + off as i32) as u32; - ymd_to_epoch_days(1970, 1, d) - } - Date32Field::DayOfWeek => { - let dow = ref_v + off as i32; - ymd_to_epoch_days(1970, 1, 4).saturating_add(dow) - } - } - } - })); - - PrimitiveArray::::new(days_values, nulls) - } -} - -/// Convert days since UNIX epoch (1970-01-01) to (year, month, day) in the -/// proleptic Gregorian calendar using a branchless integer algorithm. -fn ymd_from_epoch_days(days_since_epoch: i32) -> (i32, u32, u32) { - // Based on Howard Hinnant's civil_from_days algorithm. - let z = days_since_epoch as i64 + 719_468; // shift to civil epoch - let era = if z >= 0 { - z / 146_097 - } else { - (z - 146_096) / 146_097 - }; - let doe = z - era * 146_097; // [0, 146096] - let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365; // [0, 399] - let mut y = yoe + era * 400; - let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); // [0, 365] - let mp = (5 * doy + 2) / 153; // [0, 11] - let d = (doy - (153 * mp + 2) / 5) + 1; // [1, 31] - let m = mp + if mp < 10 { 3 } else { -9 }; // [1, 12] - if m <= 2 { - y += 1; - } - (y as i32, m as u32, d as u32) -} - -fn component_from_days(field: Date32Field, days: i32) -> i32 { - let (year, month, day) = ymd_from_epoch_days(days); - match field { - Date32Field::Year => year, - Date32Field::Month => month as i32, - Date32Field::Day => day as i32, - Date32Field::DayOfWeek => day_of_week_sunday0(days), - } -} - -fn day_of_week_sunday0(days_since_epoch: i32) -> i32 { - (days_since_epoch + 4).rem_euclid(7) -} - -fn timestamp_unit(data_type: &DataType) -> Option { - match data_type { - DataType::Timestamp(unit, _) => Some(*unit), - _ => None, - } -} - -fn timestamp_to_days_since_epoch(value: i64, unit: TimeUnit) -> i32 { - let ticks_per_day = match unit { - TimeUnit::Second => 86_400, - TimeUnit::Millisecond => 86_400_000, - TimeUnit::Microsecond => 86_400_000_000, - TimeUnit::Nanosecond => 86_400_000_000_000, - }; - (value.div_euclid(ticks_per_day)) as i32 -} - -/// Convert a date (year, month, day) in proleptic Gregorian calendar to -/// days since UNIX epoch (1970-01-01). -fn ymd_to_epoch_days(year: i32, month: u32, day: u32) -> i32 { - // Based on Howard Hinnant's civil_to_days algorithm. - let y = year as i64 - if month <= 2 { 1 } else { 0 }; - let era = if y >= 0 { y / 400 } else { (y - 399) / 400 }; - let yoe = y - era * 400; // [0, 399] - let m = month as i64; - let d = day as i64; - let mp = m + if m > 2 { -3 } else { 9 }; // Mar=0..Jan=10,Feb=11 - let doy = (153 * mp + 2) / 5 + d - 1; // [0, 365] - let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; // [0, 146096] - (era * 146_097 + doe - 719_468) as i32 -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for SqueezedDate32Array { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.get_array_memory_size() - } - - fn len(&self) -> usize { - self.len() - } - - async fn to_arrow_array(&self) -> ArrayRef { - let bytes = self.read_backing().await; - let liquid = crate::liquid_array::ipc::read_from_bytes( - bytes, - &crate::liquid_array::ipc::LiquidIPCContext::new(None), - ); - liquid.to_arrow_array() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Integer - } - - fn original_arrow_data_type(&self) -> DataType { - self.original_data_type.clone() - } - - fn disk_backing(&self) -> SqueezedBacking { - let backing = self - .backing - .as_ref() - .expect("SqueezedDate32Array backing not set"); - SqueezedBacking::Liquid((backing.disk_range.end - backing.disk_range.start) as usize) - } - - async fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - if selection.count_set_bits() == 0 { - return arrow::array::new_empty_array(&self.original_arrow_data_type()); - } - let full = self.to_arrow_array().await; - let selection_array = BooleanArray::new(selection.clone(), None); - arrow::compute::filter(&full, &selection_array).unwrap() - } - - async fn try_eval_predicate( - &self, - predicate: &LiquidExpr, - filter: &BooleanBuffer, - ) -> BooleanArray { - let filtered = self.filter(filter).await; - eval_predicate_on_array(filtered, predicate) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::types::TimestampMicrosecondType; - use arrow::array::{Array, PrimitiveArray}; - use std::sync::Arc; - - fn dates(vals: &[Option]) -> PrimitiveArray { - PrimitiveArray::::from(vals.to_vec()) - } - - fn assert_prim_eq(a: PrimitiveArray, b: PrimitiveArray) { - let a_ref: arrow::array::ArrayRef = Arc::new(a); - let b_ref: arrow::array::ArrayRef = Arc::new(b); - assert_eq!(a_ref.as_ref(), b_ref.as_ref()); - } - - fn extract(field: Date32Field, input: Vec>) -> PrimitiveArray { - let arr = dates(&input); - let liquid = LiquidPrimitiveArray::::from_arrow_array(arr); - let squeezed = SqueezedDate32Array::from_liquid_date32(&liquid, field); - squeezed.to_component_date32() - } - - fn lossy(field: Date32Field, input: Vec>) -> PrimitiveArray { - let arr = dates(&input); - let liquid = LiquidPrimitiveArray::::from_arrow_array(arr); - let squeezed = SqueezedDate32Array::from_liquid_date32(&liquid, field); - squeezed.to_arrow_date32_lossy() - } - - #[test] - fn test_extraction_correctness() { - // YEAR - let input = vec![ - Some(-1), - Some(0), - Some(ymd_to_epoch_days(1971, 7, 15)), - None, - ]; - let expected = - PrimitiveArray::::from(vec![Some(1969), Some(1970), Some(1971), None]); - assert_prim_eq(extract(Date32Field::Year, input), expected); - - // MONTH - let input = vec![ - Some(ymd_to_epoch_days(1970, 1, 31)), - Some(ymd_to_epoch_days(1970, 2, 1)), - Some(ymd_to_epoch_days(1970, 12, 31)), - None, - ]; - let expected = PrimitiveArray::::from(vec![Some(1), Some(2), Some(12), None]); - assert_prim_eq(extract(Date32Field::Month, input), expected); - - // DAY - let input = vec![ - Some(ymd_to_epoch_days(1970, 1, 1)), - Some(ymd_to_epoch_days(1970, 1, 31)), - Some(ymd_to_epoch_days(1970, 2, 1)), - None, - ]; - let expected = PrimitiveArray::::from(vec![Some(1), Some(31), Some(1), None]); - assert_prim_eq(extract(Date32Field::Day, input), expected); - - // DAYOFWEEK (Sunday=0) - let input = vec![ - Some(ymd_to_epoch_days(1970, 1, 4)), - Some(ymd_to_epoch_days(1970, 1, 5)), - Some(ymd_to_epoch_days(1970, 1, 10)), - None, - ]; - let expected = PrimitiveArray::::from(vec![Some(0), Some(1), Some(6), None]); - assert_prim_eq(extract(Date32Field::DayOfWeek, input), expected); - } - - #[test] - fn test_lossy_reconstruction_mapping() { - // YEAR → (y,1,1) - let input = vec![ - Some(ymd_to_epoch_days(1999, 12, 31)), - Some(ymd_to_epoch_days(2000, 6, 1)), - None, - ]; - let expected = PrimitiveArray::::from(vec![ - Some(ymd_to_epoch_days(1999, 1, 1)), - Some(ymd_to_epoch_days(2000, 1, 1)), - None, - ]); - assert_prim_eq(lossy(Date32Field::Year, input), expected); - - // MONTH → (1970,m,1) - let input = vec![ - Some(ymd_to_epoch_days(1980, 3, 14)), - Some(ymd_to_epoch_days(1977, 12, 5)), - None, - ]; - let expected = PrimitiveArray::::from(vec![ - Some(ymd_to_epoch_days(1970, 3, 1)), - Some(ymd_to_epoch_days(1970, 12, 1)), - None, - ]); - assert_prim_eq(lossy(Date32Field::Month, input), expected); - - // DAY → (1970,1,d) - let input = vec![ - Some(ymd_to_epoch_days(1980, 3, 14)), - Some(ymd_to_epoch_days(1977, 12, 5)), - None, - ]; - let expected = PrimitiveArray::::from(vec![ - Some(ymd_to_epoch_days(1970, 1, 14)), - Some(ymd_to_epoch_days(1970, 1, 5)), - None, - ]); - assert_prim_eq(lossy(Date32Field::Day, input), expected); - - // DAYOFWEEK → (1970,1,4 + dow) - let input = vec![ - Some(ymd_to_epoch_days(2020, 5, 17)), - Some(ymd_to_epoch_days(2020, 5, 18)), - None, - ]; - let expected = PrimitiveArray::::from(vec![ - Some(ymd_to_epoch_days(1970, 1, 4)), - Some(ymd_to_epoch_days(1970, 1, 5)), - None, - ]); - assert_prim_eq(lossy(Date32Field::DayOfWeek, input), expected); - } - - #[test] - fn test_roundtrip_idempotence() { - let input = vec![ - Some(ymd_to_epoch_days(1969, 12, 31)), - Some(ymd_to_epoch_days(1970, 1, 1)), - Some(ymd_to_epoch_days(1970, 1, 31)), - Some(ymd_to_epoch_days(1970, 2, 1)), - Some(ymd_to_epoch_days(1971, 7, 15)), - None, - ]; - - for &field in &[ - Date32Field::Year, - Date32Field::Month, - Date32Field::Day, - Date32Field::DayOfWeek, - ] { - let comp1 = extract(field, input.clone()); - let lossy_dt = lossy(field, input.clone()); - let liquid2 = LiquidPrimitiveArray::::from_arrow_array(lossy_dt); - let comp2 = - SqueezedDate32Array::from_liquid_date32(&liquid2, field).to_component_date32(); - assert_prim_eq(comp1, comp2); - } - } - - #[test] - fn test_all_nulls_behavior() { - let input = vec![None, None, None]; - - for &field in &[ - Date32Field::Year, - Date32Field::Month, - Date32Field::Day, - Date32Field::DayOfWeek, - ] { - let comp = extract(field, input.clone()); - let expected_comp = PrimitiveArray::::from(vec![None, None, None]); - assert_prim_eq(comp, expected_comp); - - let lossy_dt = lossy(field, input.clone()); - let expected_dt = PrimitiveArray::::from(vec![None, None, None]); - assert_prim_eq(lossy_dt, expected_dt); - } - } - - /// `to_component_array` is consumed by [`crate::cache::core::LiquidCache::try_read_squeezed_date32_array`] - /// as the SQL fast path. The query plan still runs `date_part` on the returned array, so the - /// values must round-trip through `component_from_days`: feeding a returned Date32 day-value - /// back into `component_from_days(field, days)` must recover the original component. - /// - /// Before the encoding fix, the Year case returned `Date32(year_int)` (e.g. year 1970 became - /// Date32 day-1970 = 1975-05-24), so re-extracting the year gave 1975 instead of 1970. - #[test] - fn to_component_array_date32_round_trips_through_extract() { - let inputs: Vec> = vec![ - Some(ymd_to_epoch_days(1970, 1, 1)), - Some(ymd_to_epoch_days(1971, 7, 15)), - Some(ymd_to_epoch_days(1999, 12, 31)), - Some(ymd_to_epoch_days(2024, 2, 29)), - Some(ymd_to_epoch_days(4709, 11, 24)), - None, - ]; - let expected_components: Vec> = inputs - .iter() - .map(|opt| opt.map(|d| component_from_days(Date32Field::Year, d))) - .collect(); - - let arr = dates(&inputs); - let liquid = LiquidPrimitiveArray::::from_arrow_array(arr); - let squeezed = SqueezedDate32Array::from_liquid_date32(&liquid, Date32Field::Year); - let component = squeezed - .to_component_array() - .as_any() - .downcast_ref::>() - .expect("date32 component array") - .clone(); - - for (idx, expected) in expected_components.iter().enumerate() { - match expected { - Some(year) => { - assert!(!component.is_null(idx), "row {idx} unexpectedly null"); - let recovered = component_from_days(Date32Field::Year, component.value(idx)); - assert_eq!( - recovered, *year, - "row {idx}: extracting Year from to_component_array output recovered {recovered}, expected {year}", - ); - } - None => assert!(component.is_null(idx), "row {idx} should be null"), - } - } - } - - #[test] - fn test_timestamp_extraction() { - // Two Microsecond timestamps at 2021-01-01 00:00:00 UTC and 2022-01-01 00:00:00 UTC. - let input = vec![ - Some(1_609_459_200_000_000), - Some(1_640_995_200_000_000), - None, - ]; - let arr = PrimitiveArray::::from(input); - let liquid = LiquidPrimitiveArray::::from_arrow_array(arr); - let squeezed = SqueezedDate32Array::from_liquid_timestamp(&liquid, Date32Field::Year); - let component = squeezed.to_component_array(); - let out = component - .as_any() - .downcast_ref::>() - .expect("timestamp array"); - - // to_component_array returns Timestamps that round-trip through `date_part`: - // year 2021 maps to (2021,1,1) at midnight UTC. - let micros_per_day: i64 = 86_400_000_000; - assert_eq!( - out.value(0), - ymd_to_epoch_days(2021, 1, 1) as i64 * micros_per_day, - ); - assert_eq!( - out.value(1), - ymd_to_epoch_days(2022, 1, 1) as i64 * micros_per_day, - ); - assert!(out.is_null(2)); - - // Direct integer view is still available via `to_component_date32`. - let int_view = squeezed.to_component_date32(); - assert_eq!(int_view.value(0), 2021); - assert_eq!(int_view.value(1), 2022); - assert!(int_view.is_null(2)); - } -} diff --git a/src/core/src/liquid_array/tests.rs b/src/core/src/liquid_array/tests.rs index de3475737..5fc575625 100644 --- a/src/core/src/liquid_array/tests.rs +++ b/src/core/src/liquid_array/tests.rs @@ -14,7 +14,6 @@ mod byte_view_tests { use rand::SeedableRng; use rand::prelude::*; - use crate::cache::{CacheExpression, TestSqueezeIo}; use crate::liquid_array::raw::FsstArray; use crate::liquid_array::{LiquidArray, LiquidByteViewArray}; @@ -197,36 +196,4 @@ mod byte_view_tests { ]); assert_eq!(result, expected); } - - #[test] - fn squeeze_and_full_read_roundtrip() { - let input = StringArray::from(vec![ - Some("hello"), - Some("world"), - Some("hello"), - None, - Some("byteview"), - ]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let decode_compressor = compressor.clone(); - let liquid = LiquidByteViewArray::::from_string_array(&input, compressor); - - let baseline = liquid.to_bytes(); - let Some((_hybrid, bytes)) = liquid.squeeze( - Arc::new(TestSqueezeIo::default()), - Some(&CacheExpression::PredicateColumn), - ) else { - panic!("squeeze should succeed"); - }; - - let restored = crate::liquid_array::ipc::read_from_bytes( - bytes.clone(), - &crate::liquid_array::ipc::LiquidIPCContext::new(Some(decode_compressor)), - ); - - let a1 = LiquidArray::to_arrow_array(&liquid); - let a2 = restored.to_arrow_array(); - assert_eq!(a1.as_ref(), a2.as_ref()); - assert_eq!(baseline, restored.to_bytes()); - } } diff --git a/src/core/src/liquid_array/variant_array.rs b/src/core/src/liquid_array/variant_array.rs deleted file mode 100644 index aaae6f38a..000000000 --- a/src/core/src/liquid_array/variant_array.rs +++ /dev/null @@ -1,279 +0,0 @@ -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, BinaryViewArray, StructArray}; -use arrow::buffer::NullBuffer; -use arrow_schema::{DataType, Field, Fields}; - -use crate::liquid_array::{ - LiquidArrayRef, LiquidDataType, LiquidSqueezedArray, NeedsBacking, SqueezedBacking, -}; -use ahash::AHashMap; - -/// Squeezed representation for variant arrays that contain multiple typed fields. -#[derive(Debug)] -pub struct VariantStructSqueezedArray { - values: AHashMap, LiquidArrayRef>, - len: usize, - nulls: Option, - original_arrow_type: DataType, - disk_backing_size: usize, -} - -impl VariantStructSqueezedArray { - /// Create a squeezed representation that keeps only the typed variant columns resident. - pub fn new( - values: Vec<(Arc, LiquidArrayRef)>, - nulls: Option, - original_arrow_type: DataType, - disk_backing_size: usize, - ) -> Self { - let len = values.first().map(|(_, array)| array.len()).unwrap_or(0); - let mut map = AHashMap::with_capacity(values.len()); - for (path, array) in values { - debug_assert_eq!(array.len(), len, "variant paths must share length"); - map.insert(path, array); - } - Self { - values: map, - len, - nulls, - original_arrow_type, - disk_backing_size, - } - } - - fn build_root_struct(&self) -> StructArray { - let metadata = Arc::new(BinaryViewArray::from(vec![b"" as &[u8]; self.len])) as ArrayRef; - let value_placeholder = Arc::new(BinaryViewArray::new_null(self.len)) as ArrayRef; - let typed_struct = self.build_typed_struct(); - - let metadata_field = Arc::new(Field::new("metadata", DataType::BinaryView, false)); - let value_field = Arc::new(Field::new("value", DataType::BinaryView, true)); - let typed_field = Arc::new(Field::new( - "typed_value", - typed_struct.data_type().clone(), - true, - )); - - StructArray::new( - Fields::from(vec![metadata_field, value_field, typed_field]), - vec![metadata, value_placeholder, typed_struct as ArrayRef], - self.nulls.clone(), - ) - } - - fn build_typed_struct(&self) -> Arc { - let mut root = VariantTreeNode::new(self.len); - for (path, array) in &self.values { - let segments: Vec<&str> = path - .split('.') - .filter(|segment| !segment.is_empty()) - .collect(); - if segments.is_empty() { - continue; - } - root.insert(&segments, array.to_arrow_array()); - } - root.into_struct_array() - } - - /// Returns true if the squeezed contains the provided variant path. - pub fn contains_path(&self, path: &str) -> bool { - self.values.contains_key(path) - } - - /// Build an Arrow array that includes only the provided variant paths. - /// If `paths` is empty or none match, it falls back to the full array. - pub fn to_arrow_array_with_paths<'a>( - &self, - paths: impl IntoIterator, - ) -> Result { - let mut filtered: Vec<(Arc, LiquidArrayRef)> = Vec::new(); - for path in paths.into_iter() { - if let Some(array) = self.values.get(path) { - filtered.push((Arc::from(path.to_string()), array.clone())); - } - } - - if filtered.is_empty() { - return Ok(Arc::new(self.build_root_struct()) as ArrayRef); - } - - let filtered = VariantStructSqueezedArray::new( - filtered, - self.nulls.clone(), - self.original_arrow_type.clone(), - self.disk_backing_size, - ); - Ok(Arc::new(filtered.build_root_struct()) as ArrayRef) - } - - /// Clone the stored typed values keyed by variant path. - pub fn typed_values(&self) -> Vec<(Arc, LiquidArrayRef)> { - self.values - .iter() - .map(|(path, array)| (path.clone(), array.clone())) - .collect() - } - - /// Null buffer shared by all stored paths, if present. - pub fn nulls(&self) -> Option { - self.nulls.clone() - } -} - -#[async_trait::async_trait] -impl LiquidSqueezedArray for VariantStructSqueezedArray { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.values - .values() - .map(|array| array.get_array_memory_size()) - .sum() - } - - fn len(&self) -> usize { - self.len - } - - async fn to_arrow_array(&self) -> ArrayRef { - Arc::new(self.build_root_struct()) as ArrayRef - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::ByteViewArray - } - - fn original_arrow_data_type(&self) -> DataType { - self.original_arrow_type.clone() - } - - fn disk_backing(&self) -> SqueezedBacking { - SqueezedBacking::Arrow(self.disk_backing_size) - } -} - -#[derive(Default)] -struct VariantTreeNode { - len: usize, - leaf: Option, - children: AHashMap, -} - -impl VariantTreeNode { - fn new(len: usize) -> Self { - Self { - len, - leaf: None, - children: AHashMap::new(), - } - } - - fn insert(&mut self, segments: &[&str], values: ArrayRef) { - if segments.is_empty() { - self.leaf = Some(values); - return; - } - let (head, tail) = segments.split_first().unwrap(); - self.children - .entry(head.to_string()) - .or_insert_with(|| VariantTreeNode::new(self.len)) - .insert(tail, values); - } - - fn into_struct_array(self) -> Arc { - let mut fields = Vec::with_capacity(self.children.len()); - let mut arrays = Vec::with_capacity(self.children.len()); - let mut entries: Vec<_> = self.children.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - for (name, child) in entries { - let field_array = child.into_field_array(); - fields.push(Arc::new(Field::new( - name.as_str(), - field_array.data_type().clone(), - false, - ))); - arrays.push(field_array); - } - Arc::new(StructArray::new(Fields::from(fields), arrays, None)) - } - - fn into_field_array(self) -> ArrayRef { - let len = self.len; - if self.children.is_empty() { - let values = self.leaf.expect("variant leaf value present"); - wrap_typed_value(len, values) - } else { - let typed_struct = self.into_struct_array() as ArrayRef; - wrap_typed_value(len, typed_struct) - } - } -} - -fn wrap_typed_value(len: usize, values: ArrayRef) -> ArrayRef { - let placeholder = Arc::new(BinaryViewArray::new_null(len)) as ArrayRef; - Arc::new(StructArray::new( - Fields::from(vec![ - Arc::new(Field::new("value", DataType::BinaryView, true)), - Arc::new(Field::new("typed_value", values.data_type().clone(), true)), - ]), - vec![placeholder, values], - None, - )) as ArrayRef -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{Int64Array, StringArray}; - use arrow_schema::DataType; - - use crate::liquid_array::{LiquidByteViewArray, LiquidPrimitiveArray, raw::FsstArray}; - - #[test] - fn to_arrow_array_with_paths_prunes_extra_fields() { - // Build squeezed variant with two typed paths: did (utf8) and time_us (int64). - let did_arrow = StringArray::from(vec![Some("d")]); - let (_comp, did_liquid) = LiquidByteViewArray::::train_from_arrow(&did_arrow); - let did_liquid: LiquidArrayRef = Arc::new(did_liquid); - - let time_arrow = Int64Array::from(vec![1_i64]); - let time_liquid = - LiquidPrimitiveArray::::from_arrow_array(time_arrow); - let time_liquid: LiquidArrayRef = Arc::new(time_liquid); - - let squeezed = VariantStructSqueezedArray::new( - vec![ - (Arc::from("did"), did_liquid), - (Arc::from("time_us"), time_liquid), - ], - None, - DataType::Struct(Fields::from(Vec::>::new())), - 0, - ); - - // Request only time_us; did should be pruned from typed_value. - let array = squeezed - .to_arrow_array_with_paths(["time_us"]) - .expect("arrow array"); - let root = array - .as_any() - .downcast_ref::() - .expect("struct root"); - let typed_value = root - .column_by_name("typed_value") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let field_names: Vec<_> = typed_value - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(field_names, vec!["time_us"]); - } -} diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index 0c69f4482..fb3eea97f 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -5,7 +5,7 @@ use std::time::Instant; use arrow::array::ArrayRef; use arrow::buffer::BooleanBuffer; use arrow_schema::DataType; -use clap::Parser; +use clap::{Parser, ValueEnum}; use datafusion::logical_expr::Operator; use datafusion::prelude::*; use datafusion::scalar::ScalarValue; @@ -15,7 +15,8 @@ use liquid_cache::cache::LiquidCache; use liquid_cache::cache::LiquidCacheBuilder; use liquid_cache::cache::LiquidExpr; use liquid_cache::cache::LiquidPolicy; -use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache::cache::TranscodeEvict; +use liquid_cache::cache::{Evict, EvictionPolicy}; #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; @@ -35,22 +36,41 @@ struct CliArgs { /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors #[arg(long, default_value = "false")] bench: bool, + + /// Representation used while evicting memory entries. + #[arg(long, value_enum, default_value_t = CacheMode::Liquid)] + mode: CacheMode, + + /// Memory budget in MiB. + #[arg(long, default_value_t = 500)] + max_memory_mb: usize, +} + +#[derive(Debug, Default, Clone, Copy, ValueEnum)] +enum CacheMode { + Arrow, + #[default] + Liquid, } fn main() { let args = CliArgs::parse(); - // 1) Build cache storage with FILO and a small budget (500 MB) + // 1) Build cache storage with the requested representation and budget. let cache_dir = args .cache_dir .clone() .unwrap_or_else(|| tempfile::tempdir().unwrap().keep()); let store_path = cache_dir.join("liquid_cache.t4"); let store = tokio_test::block_on(t4::mount(&store_path)).expect("failed to mount t4 store"); + let eviction_policy: Box = match args.mode { + CacheMode::Arrow => Box::new(Evict), + CacheMode::Liquid => Box::new(TranscodeEvict), + }; let storage = tokio_test::block_on(async { LiquidCacheBuilder::new() - .with_max_memory_bytes(500 * 1024 * 1024) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_max_memory_bytes(args.max_memory_mb * 1024 * 1024) + .with_eviction_policy(eviction_policy) .with_cache_policy(Box::new(LiquidPolicy::new())) .with_store(store) .build() @@ -72,7 +92,7 @@ fn main() { let pred_expr: Arc = Arc::new(BinaryExpr::new( Arc::new(Column::new("col", 0)), Operator::Eq, - Arc::new(Literal::new(ScalarValue::Utf8View(Some(String::new())))), + Arc::new(Literal::new(ScalarValue::Utf8(Some(String::new())))), )); // 4) Scan all entries with selection=all-true and time get_with_predicate @@ -85,7 +105,7 @@ fn main() { let selection = BooleanBuffer::new_set(len); let Some(liquid_expr) = LiquidExpr::try_new( Arc::clone(&pred_expr), - &DataType::Utf8View, + &DataType::Utf8, Some(&liquid_cache::cache::CacheExpression::PredicateColumn), ) else { continue; diff --git a/src/core/study/squeeze_integer.rs b/src/core/study/squeeze_integer.rs deleted file mode 100644 index cfdb89cc1..000000000 --- a/src/core/study/squeeze_integer.rs +++ /dev/null @@ -1,516 +0,0 @@ -use std::ops::Range; -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, BooleanArray, PrimitiveArray, cast::AsArray}; -use arrow::buffer::BooleanBuffer; -use arrow::datatypes::DataType; -use bytes::Bytes; -use clap::Parser; -use datafusion::prelude::*; -use datafusion::scalar::ScalarValue; -use futures::StreamExt; -use liquid_cache::cache::{CacheExpression, LiquidExpr}; -use liquid_cache::liquid_array::{ - IntegerSqueezePolicy, LiquidArray, LiquidPrimitiveArray, LiquidPrimitiveType, - LiquidSqueezedArray, SqueezeIoHandler, -}; -use std::sync::Mutex; -use std::sync::atomic::{AtomicUsize, Ordering}; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Default, Clone)] -#[command(name = "Squeeze Integer Study")] -#[command(about = "Clamp vs Quantize squeeze on representative integer filters from ClickBench")] -struct CliArgs { - /// Parquet file to read - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Optional row limit for each column (useful for faster runs) - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors - #[arg(long, default_value = "false")] - bench: bool, -} - -#[derive(Debug, Clone)] -struct FilterCase { - column: String, - op: datafusion::logical_expr::Operator, - scalar: ScalarValue, -} - -#[derive(Default, Debug, Clone)] -struct Stats { - rows: usize, - arrow_bytes: usize, - liquid_bytes: usize, - clamp_mem_bytes: usize, - clamp_disk_bytes: usize, - quant_mem_bytes: usize, - quant_disk_bytes: usize, - // IO incurred (bytes) for try_eval_predicate - clamp_pred_io_bytes: usize, - quant_pred_io_bytes: usize, - // IO incurred (bytes) for get_with_selection (filter_to_arrow) - clamp_select_io_bytes: usize, - quant_select_io_bytes: usize, - // number of predicate cases executed on this column - pred_cases: usize, -} - -impl Stats { - fn add(&mut self, other: &Stats) { - self.rows += other.rows; - self.arrow_bytes += other.arrow_bytes; - self.liquid_bytes += other.liquid_bytes; - self.clamp_mem_bytes += other.clamp_mem_bytes; - self.clamp_disk_bytes += other.clamp_disk_bytes; - self.quant_mem_bytes += other.quant_mem_bytes; - self.quant_disk_bytes += other.quant_disk_bytes; - self.clamp_pred_io_bytes += other.clamp_pred_io_bytes; - self.quant_pred_io_bytes += other.quant_pred_io_bytes; - self.clamp_select_io_bytes += other.clamp_select_io_bytes; - self.quant_select_io_bytes += other.quant_select_io_bytes; - self.pred_cases += other.pred_cases; - } -} - -// Hardcoded representative integer filters based on ClickBench queries -fn representative_integer_filters() -> Vec { - use datafusion::logical_expr::Operator as Op; - vec![ - // SELECT COUNT(*) FROM hits WHERE "AdvEngineID" <> 0; - FilterCase { - column: "AdvEngineID".to_string(), - op: Op::NotEq, - scalar: ScalarValue::Int64(Some(0)), - }, - // SELECT "UserID" FROM hits WHERE "UserID" = 435090932899640449; - FilterCase { - column: "UserID".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(435_090_932_899_640_449)), - }, - // WHERE "CounterID" = 62 - FilterCase { - column: "CounterID".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(62)), - }, - // WHERE "IsRefresh" = 0 - FilterCase { - column: "IsRefresh".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(0)), - }, - // WHERE "DontCountHits" = 0 - FilterCase { - column: "DontCountHits".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(0)), - }, - // WHERE "IsLink" <> 0 - FilterCase { - column: "IsLink".to_string(), - op: Op::NotEq, - scalar: ScalarValue::Int64(Some(0)), - }, - // WHERE "IsDownload" = 0 - FilterCase { - column: "IsDownload".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(0)), - }, - // WHERE "TraficSourceID" IN (-1, 6) - FilterCase { - column: "TraficSourceID".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(-1)), - }, - FilterCase { - column: "TraficSourceID".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(6)), - }, - // WHERE "RefererHash" = 3594120000172545465 - FilterCase { - column: "RefererHash".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(3_594_120_000_172_545_465)), - }, - // WHERE "URLHash" = 2868770270353813622 - FilterCase { - column: "URLHash".to_string(), - op: Op::Eq, - scalar: ScalarValue::Int64(Some(2_868_770_270_353_813_622)), - }, - ] -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(8192 * 2); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - // Hardcoded representative integer filters inspired by ClickBench queries - let cases = representative_integer_filters(); - - println!("Squeeze Integer Study over {} case(s)", cases.len()); - - // Stream the requested columns; for simplicity scan per-case so we only pull needed column - // and accumulate stats per column then sum totals. - let mut grand = Stats::default(); - for case in &cases { - let stats = run_case(&ctx, case, args.limit).await; - println!( - "Case on column '{}', op '{:?}', scalar {:?}:\n rows: {}\n sizes (bytes) -> arrow: {}, liquid: {}, clamp: {} (mem: {}, disk: {}), quant: {} (mem: {}, disk: {})\n io (bytes) -> pred: clamp {}, quant {}; select: clamp {}, quant {}", - case.column, - case.op, - case.scalar, - stats.rows, - stats.arrow_bytes, - stats.liquid_bytes, - stats.clamp_mem_bytes + stats.clamp_disk_bytes, - stats.clamp_mem_bytes, - stats.clamp_disk_bytes, - stats.quant_mem_bytes + stats.quant_disk_bytes, - stats.quant_mem_bytes, - stats.quant_disk_bytes, - stats.clamp_pred_io_bytes, - stats.quant_pred_io_bytes, - stats.clamp_select_io_bytes, - stats.quant_select_io_bytes - ); - grand.add(&stats); - } - - println!( - "TOTAL\n rows: {}\n sizes (bytes) -> arrow: {}, liquid: {}, clamp: {} (mem: {}, disk: {}), quant: {} (mem: {}, disk: {})\n io (bytes) -> pred: clamp {}, quant {}; select: clamp {}, quant {}", - grand.rows, - grand.arrow_bytes, - grand.liquid_bytes, - grand.clamp_mem_bytes + grand.clamp_disk_bytes, - grand.clamp_mem_bytes, - grand.clamp_disk_bytes, - grand.quant_mem_bytes + grand.quant_disk_bytes, - grand.quant_mem_bytes, - grand.quant_disk_bytes, - grand.clamp_pred_io_bytes, - grand.quant_pred_io_bytes, - grand.clamp_select_io_bytes, - grand.quant_select_io_bytes - ); -} - -// --- end of helpers removed after switching to fixed cases --- - -async fn run_case(ctx: &SessionContext, case: &FilterCase, limit: Option) -> Stats { - let sql = if let Some(n) = limit { - format!("SELECT \"{}\" FROM \"hits\" LIMIT {n}", case.column) - } else { - format!("SELECT \"{}\" FROM \"hits\"", case.column) - }; - let df = ctx.sql(&sql).await.expect("create df"); - let mut stream = df.execute_stream().await.expect("execute stream"); - - let mut stats = Stats::default(); - while let Some(batch_res) = stream.next().await { - let batch = batch_res.expect("stream batch"); - let array: ArrayRef = batch.column(0).clone(); - stats.rows += array.len(); - stats.arrow_bytes += array.get_array_memory_size(); - - // Dispatch by datatype - match array.data_type() { - DataType::Int8 => run_for_array::(&array, case, &mut stats), - DataType::Int16 => { - run_for_array::(&array, case, &mut stats) - } - DataType::Int32 => { - run_for_array::(&array, case, &mut stats) - } - DataType::Int64 => { - run_for_array::(&array, case, &mut stats) - } - DataType::UInt8 => { - run_for_array::(&array, case, &mut stats) - } - DataType::UInt16 => { - run_for_array::(&array, case, &mut stats) - } - DataType::UInt32 => { - run_for_array::(&array, case, &mut stats) - } - DataType::UInt64 => { - run_for_array::(&array, case, &mut stats) - } - DataType::Date32 => { - run_for_array::(&array, case, &mut stats) - } - DataType::Date64 => { - run_for_array::(&array, case, &mut stats) - } - _ => {} - } - } - - stats -} - -#[derive(Debug, Default)] -struct InMemorySqueezeIo { - bytes: Mutex>, - bytes_read: AtomicUsize, -} - -impl InMemorySqueezeIo { - fn set_bytes(&self, bytes: Bytes) { - *self.bytes.lock().unwrap() = Some(bytes); - } - - fn bytes(&self) -> Bytes { - self.bytes - .lock() - .unwrap() - .clone() - .expect("in-memory squeeze bytes set") - } - - fn reset_bytes_read(&self) { - self.bytes_read.store(0, Ordering::SeqCst); - } - - fn bytes_read(&self) -> usize { - self.bytes_read.load(Ordering::SeqCst) - } -} - -#[async_trait::async_trait] -impl SqueezeIoHandler for InMemorySqueezeIo { - async fn read(&self, range: Option>) -> std::io::Result { - let bytes = self.bytes(); - let out = match range { - Some(range) => bytes.slice(range.start as usize..range.end as usize), - None => bytes, - }; - self.bytes_read.fetch_add(out.len(), Ordering::SeqCst); - Ok(out) - } -} - -fn run_for_array(array: &ArrayRef, case: &FilterCase, stats: &mut Stats) -where - ::Native: num_traits::cast::AsPrimitive - + num_traits::FromPrimitive - + num_traits::bounds::Bounded, -{ - let prim = array.as_primitive::().clone(); - - // Build Liquid primitive array (unsqueezed) and track its size - let liquid = LiquidPrimitiveArray::::from_arrow_array(prim.clone()); - stats.liquid_bytes += liquid.get_array_memory_size(); - - let hint = CacheExpression::PredicateColumn; - - // Build Liquid primitive array and squeeze with Clamp - let mut lp = LiquidPrimitiveArray::::from_arrow_array(prim.clone()); - let clamp_io = Arc::new(InMemorySqueezeIo::default()); - let clamp_hybrid_and_bytes = { - lp.set_squeeze_policy(IntegerSqueezePolicy::Clamp); - lp.squeeze(clamp_io.clone(), Some(&hint)) - }; - - // Build Quantize - let mut lq = LiquidPrimitiveArray::::from_arrow_array(prim.clone()); - let quant_io = Arc::new(InMemorySqueezeIo::default()); - let quant_hybrid_and_bytes = { - lq.set_squeeze_policy(IntegerSqueezePolicy::Quantize); - lq.squeeze(quant_io.clone(), Some(&hint)) - }; - - // Size accounting (for squeezable ones) - if let Some((h, bytes)) = clamp_hybrid_and_bytes.as_ref() { - clamp_io.set_bytes(bytes.clone()); - stats.clamp_mem_bytes += h.get_array_memory_size(); - stats.clamp_disk_bytes += bytes.len(); - } - if let Some((h, bytes)) = quant_hybrid_and_bytes.as_ref() { - quant_io.set_bytes(bytes.clone()); - stats.quant_mem_bytes += h.get_array_memory_size(); - stats.quant_disk_bytes += bytes.len(); - } - - // Build predicate expr: Column op Literal(scalar) - use datafusion::physical_plan::expressions::{BinaryExpr, Column, Literal}; - let expr: std::sync::Arc = - std::sync::Arc::new(BinaryExpr::new( - std::sync::Arc::new(Column::new("col", 0)), - case.op, - std::sync::Arc::new(Literal::new(case.scalar.clone())), - )); - - let all_true = BooleanBuffer::new_set(prim.len()); - - // Evaluate predicate on clamp - if let Some((hy, _full_bytes)) = clamp_hybrid_and_bytes.clone() { - let (mask, pred_io_bytes) = - try_eval_or_fetch::(&*hy, clamp_io.as_ref(), &expr, &all_true); - stats.clamp_pred_io_bytes += pred_io_bytes; - let sel = bool_array_to_selection(&mask); - // Expected selection result from Arrow - let expected_filtered = filter_expected::(&prim, &case.op, &case.scalar); - // Try get with selection from hybrid - let sel_io = get_with_selection(&*hy, clamp_io.as_ref(), &sel, expected_filtered.as_ref()); - stats.clamp_select_io_bytes += sel_io; - } - - // Evaluate predicate on quantized - if let Some((hy, _full_bytes)) = quant_hybrid_and_bytes.clone() { - let (mask, pred_io_bytes) = - try_eval_or_fetch::(&*hy, quant_io.as_ref(), &expr, &all_true); - stats.quant_pred_io_bytes += pred_io_bytes; - let sel = bool_array_to_selection(&mask); - // Expected selection result from Arrow - let expected_filtered = filter_expected::(&prim, &case.op, &case.scalar); - let sel_io = get_with_selection(&*hy, quant_io.as_ref(), &sel, expected_filtered.as_ref()); - stats.quant_select_io_bytes += sel_io; - } - - stats.pred_cases += 1; -} - -fn try_eval_or_fetch( - hybrid: &dyn LiquidSqueezedArray, - io: &InMemorySqueezeIo, - expr: &std::sync::Arc, - filter: &BooleanBuffer, -) -> (BooleanArray, usize) { - io.reset_bytes_read(); - let maybe_expr = LiquidExpr::try_new(expr.clone(), &hybrid.original_arrow_data_type(), None); - if let Some(liquid_expr) = maybe_expr { - let mask = futures::executor::block_on(hybrid.try_eval_predicate(&liquid_expr, filter)); - return (mask, io.bytes_read()); - } - // Not supported in hybrid form: materialize from full bytes and compute via Arrow. - // Count this as a full backing read for apples-to-apples IO accounting. - let full_bytes = io.bytes(); - let liq = LiquidPrimitiveArray::::from_bytes(full_bytes.clone()); - let arr = liq.to_arrow_array(); - let mask = eval_on_arrow(&arr, expr); - (mask, full_bytes.len()) -} - -fn get_with_selection( - hybrid: &dyn LiquidSqueezedArray, - io: &InMemorySqueezeIo, - selection: &BooleanBuffer, - expected: &dyn Array, -) -> usize { - io.reset_bytes_read(); - let arr = futures::executor::block_on(hybrid.filter(selection)); - assert_eq!(arr.as_ref(), expected); - io.bytes_read() -} - -fn eval_on_arrow( - array: &ArrayRef, - expr: &std::sync::Arc, -) -> BooleanArray { - use arrow::compute::cast; - use datafusion::logical_expr::ColumnarValue; - use datafusion::physical_expr_common::datum::apply_cmp; - use datafusion::physical_plan::expressions::{BinaryExpr, Literal}; - - if let Some(be) = expr.downcast_ref::() - && let Some(lit) = be.right().downcast_ref::() - { - let target_dt = scalar_data_type(lit.value()).unwrap_or_else(|| array.data_type().clone()); - let lhs_arr = if &target_dt == array.data_type() { - array.clone() - } else { - cast(array, &target_dt).expect("cast lhs for comparison") - }; - let lhs = ColumnarValue::Array(lhs_arr); - let rhs = ColumnarValue::Scalar(lit.value().clone()); - let res = match be.op() { - datafusion::logical_expr::Operator::Eq => { - apply_cmp(datafusion::logical_expr::Operator::Eq, &lhs, &rhs) - } - datafusion::logical_expr::Operator::NotEq => { - apply_cmp(datafusion::logical_expr::Operator::NotEq, &lhs, &rhs) - } - datafusion::logical_expr::Operator::Lt => { - apply_cmp(datafusion::logical_expr::Operator::Lt, &lhs, &rhs) - } - datafusion::logical_expr::Operator::LtEq => { - apply_cmp(datafusion::logical_expr::Operator::LtEq, &lhs, &rhs) - } - datafusion::logical_expr::Operator::Gt => { - apply_cmp(datafusion::logical_expr::Operator::Gt, &lhs, &rhs) - } - datafusion::logical_expr::Operator::GtEq => { - apply_cmp(datafusion::logical_expr::Operator::GtEq, &lhs, &rhs) - } - _ => panic!("unsupported operator"), - } - .expect("cmp ok"); - let arr = res.into_array(array.len()).unwrap(); - arr.as_boolean().clone() - } else { - panic!("unexpected expression kind for numeric predicate") - } -} - -fn bool_array_to_selection(mask: &BooleanArray) -> BooleanBuffer { - // selection must be non-nullable; treat nulls as false - let iter = (0..mask.len()).map(|i| mask.is_valid(i) && mask.value(i)); - BooleanBuffer::from_iter(iter) -} - -fn filter_expected( - prim: &PrimitiveArray, - op: &datafusion::logical_expr::Operator, - scalar: &ScalarValue, -) -> ArrayRef { - use datafusion::physical_plan::expressions::{BinaryExpr, Column, Literal}; - // Build the same predicate expr and evaluate mask on Arrow (eval_on_arrow may cast internally). - let arr: ArrayRef = std::sync::Arc::new(prim.clone()); - let expr: std::sync::Arc = - std::sync::Arc::new(BinaryExpr::new( - std::sync::Arc::new(Column::new("col", 0)), - *op, - std::sync::Arc::new(Literal::new(scalar.clone())), - )); - let mask = eval_on_arrow(&arr, &expr); - // Apply mask to the original-typed array to keep dtype identical to hybrid’s result - arrow::compute::kernels::filter::filter(&arr, &mask).unwrap() -} - -fn scalar_data_type(sv: &ScalarValue) -> Option { - Some(match sv { - ScalarValue::Int8(_) => DataType::Int8, - ScalarValue::Int16(_) => DataType::Int16, - ScalarValue::Int32(_) => DataType::Int32, - ScalarValue::Int64(_) => DataType::Int64, - ScalarValue::UInt8(_) => DataType::UInt8, - ScalarValue::UInt16(_) => DataType::UInt16, - ScalarValue::UInt32(_) => DataType::UInt32, - ScalarValue::UInt64(_) => DataType::UInt64, - ScalarValue::Date32(_) => DataType::Date32, - ScalarValue::Date64(_) => DataType::Date64, - _ => return None, - }) -} diff --git a/src/datafusion-client/src/client_exec.rs b/src/datafusion-client/src/client_exec.rs index 13ee28f55..b9f751bf5 100644 --- a/src/datafusion-client/src/client_exec.rs +++ b/src/datafusion-client/src/client_exec.rs @@ -40,10 +40,10 @@ use fastrace::future::FutureExt; use fastrace::prelude::*; use futures::{Stream, TryStreamExt, future::BoxFuture, ready}; use liquid_cache_common::rpc::{ - ColumnSqueezeHint, FetchResults, LiquidCacheActions, RegisterObjectStoreRequest, + ColumnLineage, FetchResults, LiquidCacheActions, RegisterObjectStoreRequest, RegisterPlanRequest, }; -use liquid_cache_datafusion::cache::ColumnSqueezeHints; +use liquid_cache_datafusion::cache::ColumnLineages; use tonic::Request; use uuid::Uuid; @@ -66,9 +66,9 @@ pub struct LiquidCacheClientExec { uuid: Uuid, plan_registered: Arc, properties: Arc, - /// Typed squeeze hints for the scan in `remote_plan`, derived by the client + /// Typed lineage expressions for the scan in `remote_plan`, derived by the client /// from the full physical plan and shipped to the cache server. - squeeze_hints: ColumnSqueezeHints, + lineages: ColumnLineages, } impl std::fmt::Debug for LiquidCacheClientExec { @@ -91,7 +91,7 @@ impl LiquidCacheClientExec { remote_plan: Arc, cache_server: String, object_stores: Vec<(ObjectStoreUrl, HashMap)>, - squeeze_hints: ColumnSqueezeHints, + lineages: ColumnLineages, ) -> Self { let properties = Self::plan_properties(&remote_plan); let uuid = Uuid::new_v4(); @@ -103,7 +103,7 @@ impl LiquidCacheClientExec { uuid, metrics: ExecutionPlanMetricsSet::new(), properties, - squeeze_hints, + lineages, } } @@ -171,7 +171,7 @@ impl ExecutionPlan for LiquidCacheClientExec { metrics: self.metrics.clone(), uuid: self.uuid, properties, - squeeze_hints: self.squeeze_hints.clone(), + lineages: self.lineages.clone(), })) } @@ -215,7 +215,7 @@ impl ExecutionPlan for LiquidCacheClientExec { self.uuid, partition, self.object_stores.clone(), - squeeze_hints_to_wire(&self.squeeze_hints), + lineages_to_wire(&self.lineages), ); Ok(Box::pin(FlightStream::new( Some(Box::pin(stream)), @@ -272,11 +272,11 @@ impl ExecutionPlan for LiquidCacheClientExec { } } -/// Convert typed squeeze hints into their wire form (canonical string encoding). -fn squeeze_hints_to_wire(hints: &ColumnSqueezeHints) -> Vec { +/// Convert typed lineage expressions into their wire form (canonical string encoding). +fn lineages_to_wire(hints: &ColumnLineages) -> Vec { hints .iter() - .map(|(column, expr)| ColumnSqueezeHint { + .map(|(column, expr)| ColumnLineage { column: column.clone(), hint: expr.to_metadata_value(), }) @@ -290,7 +290,7 @@ async fn flight_stream( handle: Uuid, partition: usize, object_stores: Vec<(ObjectStoreUrl, HashMap)>, - squeeze_hints: Vec, + lineages: Vec, ) -> Result { // Materialized scalar-subquery results are embedded in scan predicates as // `ScalarSubqueryExpr`, which cannot be serialized on its own and is @@ -330,7 +330,7 @@ async fn flight_stream( let action = LiquidCacheActions::RegisterPlan(RegisterPlanRequest { plan: plan_bytes.to_vec(), handle: handle.into_bytes().to_vec().into(), - squeeze_hints: squeeze_hints.clone(), + lineages: lineages.clone(), }) .into(); client diff --git a/src/datafusion-client/src/optimizer.rs b/src/datafusion-client/src/optimizer.rs index ad75324ba..992cdc4ef 100644 --- a/src/datafusion-client/src/optimizer.rs +++ b/src/datafusion-client/src/optimizer.rs @@ -12,7 +12,7 @@ use datafusion::{ physical_plan::{ExecutionPlan, execution_plan::replace_children_if_necessary}, }; -use liquid_cache_datafusion::optimizers::SqueezeHintMap; +use liquid_cache_datafusion::optimizers::LineageHints; use crate::client_exec::LiquidCacheClientExec; @@ -47,7 +47,7 @@ impl PushdownOptimizer { fn optimize_plan( &self, plan: Arc, - hints: &SqueezeHintMap, + hints: &LineageHints, ) -> Result> { // If this node is already a LiquidCacheClientExec, return it as is if plan.is::() { @@ -58,16 +58,16 @@ impl PushdownOptimizer { if let Some(candidate) = find_pushdown_candidate(&plan) { // If the current node is the one to be pushed down, wrap it if Arc::ptr_eq(&plan, &candidate) { - // The fragment is single-scan; collect that scan's squeeze + // The fragment is single-scan; collect that scan's lineage // hints (derived from the full plan, which includes the // client-side projections that won't be shipped) so the server // can apply them when it rebuilds the LiquidParquetSource. - let squeeze_hints = hints.for_fragment(&plan); + let lineages = hints.for_fragment(&plan); return Ok(Arc::new(LiquidCacheClientExec::new( plan, self.cache_server.clone(), self.object_stores.clone(), - squeeze_hints, + lineages, ))); } } @@ -148,11 +148,11 @@ impl PhysicalOptimizerRule for PushdownOptimizer { plan: Arc, _config: &ConfigOptions, ) -> Result> { - // Derive squeeze hints from the full physical plan up front: the + // Derive lineage expressions from the full physical plan up front: the // lineage that justifies a hint (e.g. a `date_part` projection) often // lives above the node we push down, so it must be captured before the // plan is split into fragments. - let hints = SqueezeHintMap::analyze(&plan); + let hints = LineageHints::analyze(&plan); self.optimize_plan(plan, &hints) } diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 32e2f0d58..bfccbaa13 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -10,8 +10,8 @@ use std::sync::Arc; use datafusion::logical_expr::ScalarUDF; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion::{common::config::ConfigNonZeroUsize, error::Result}; -use liquid_cache::cache::squeeze_policies::{SqueezePolicy, TranscodeSqueezeEvict}; use liquid_cache::cache::{AlwaysHydrate, HydrationPolicy, default_max_memory_bytes}; +use liquid_cache::cache::{EvictionPolicy, TranscodeEvict}; use liquid_cache::cache_policies::{CachePolicy, LiquidPolicy}; use liquid_cache_datafusion::optimizers::LocalModeOptimizer; use liquid_cache_datafusion::{ @@ -65,8 +65,8 @@ pub struct LiquidCacheLocalBuilder { cache_dir: PathBuf, /// Cache policy cache_policy: Box, - /// Squeeze policy - squeeze_policy: Box, + /// Eviction policy + eviction_policy: Box, /// Hydration policy hydration_policy: Box, prefetch: bool, @@ -83,7 +83,7 @@ impl Default for LiquidCacheLocalBuilder { max_disk_bytes, cache_dir: std::env::temp_dir(), cache_policy: Box::new(LiquidPolicy::new()), - squeeze_policy: Box::new(TranscodeSqueezeEvict), + eviction_policy: Box::new(TranscodeEvict), hydration_policy: Box::new(AlwaysHydrate::new()), prefetch: true, span: fastrace::Span::enter_with_local_parent("liquid_cache_datafusion_local_builder"), @@ -123,9 +123,9 @@ impl LiquidCacheLocalBuilder { self } - /// Set squeeze policy - pub fn with_squeeze_policy(mut self, squeeze_policy: Box) -> Self { - self.squeeze_policy = squeeze_policy; + /// Set eviction policy + pub fn with_eviction_policy(mut self, eviction_policy: Box) -> Self { + self.eviction_policy = eviction_policy; self } @@ -179,19 +179,19 @@ impl LiquidCacheLocalBuilder { self.max_disk_bytes, store, self.cache_policy, - self.squeeze_policy, + self.eviction_policy, self.hydration_policy, ) .await; #[cfg(test)] - let cache = LiquidCacheParquet::new_with_squeeze_victim_concurrency( + let cache = LiquidCacheParquet::new_with_eviction_concurrency( self.batch_size, self.max_memory_bytes, self.max_disk_bytes, store, self.cache_policy, - self.squeeze_policy, + self.eviction_policy, self.hydration_policy, false, ) diff --git a/src/datafusion-local/src/tests/date_optimizer.rs b/src/datafusion-local/src/tests/date_optimizer.rs index 2d5fb9986..4167efa4f 100644 --- a/src/datafusion-local/src/tests/date_optimizer.rs +++ b/src/datafusion-local/src/tests/date_optimizer.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use arrow::array::{Date32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use datafusion::parquet::arrow::ArrowWriter; -use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache::cache::TranscodeEvict; use tempfile::TempDir; use crate::tests::CacheStatsSummary; @@ -39,7 +39,7 @@ async fn general_test(sql: &str) -> CacheStatsSummary { let lc_builder = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 1024) .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(liquid_cache::cache_policies::LiquidPolicy::new())); let mut config = SessionConfig::new(); config.options_mut().execution.target_partitions = 1; @@ -96,27 +96,23 @@ async fn general_test(sql: &str) -> CacheStatsSummary { #[tokio::test] async fn test_date_extraction() { let sql = r#"select AVG(EXTRACT(YEAR from date_a)) as year from test_table"#; - let stats = general_test(sql).await; - assert_eq!(stats.stats.runtime.hit_date32_expression_calls, 86); + general_test(sql).await; } #[tokio::test] async fn date_extraction_month() { let sql = r#"select AVG(EXTRACT(MONTH from date_a)) as month from test_table"#; - let stats = general_test(sql).await; - assert_eq!(stats.stats.runtime.hit_date32_expression_calls, 78); + general_test(sql).await; } #[tokio::test] async fn date_extraction_day() { let sql = r#"select AVG(EXTRACT(DAY from date_a)) as day from test_table"#; - let stats = general_test(sql).await; - assert_eq!(stats.stats.runtime.hit_date32_expression_calls, 86); + general_test(sql).await; } #[tokio::test] async fn test_date_extraction_case2() { let sql = r#"select AVG(EXTRACT(YEAR from date_a) + 1) as year, (SELECT MAX(EXTRACT(YEAR from date_a)) FROM test_table) as max_year from test_table"#; - let stats = general_test(sql).await; - assert_eq!(stats.stats.runtime.hit_date32_expression_calls, 172); // we know this. + general_test(sql).await; } diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index e89f79b83..834e82e22 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -1,9 +1,6 @@ use arrow_schema::{DataType, Field, Schema}; use liquid_cache::{ - cache::{ - CacheStats, - squeeze_policies::{Evict, SqueezePolicy, TranscodeEvict, TranscodeSqueezeEvict}, - }, + cache::{CacheStats, Evict, EvictionPolicy, TranscodeEvict}, cache_policies::LiquidPolicy, }; use liquid_cache_datafusion::LiquidCacheParquetRef; @@ -25,7 +22,6 @@ use crate::LiquidCacheLocalBuilder; mod date_optimizer; mod filter_limit; mod nested_filter; -mod squeeze; mod variants; const TEST_FILE: &str = "../../examples/nano_hits.parquet"; @@ -83,11 +79,6 @@ impl fmt::Display for CacheStatsSummary { "entries.memory.liquid: {}", self.stats.memory_liquid_entries )?; - writeln!( - f, - "entries.memory.squeezed_liquid: {}", - self.stats.memory_squeezed_liquid_entries - )?; writeln!(f, "entries.disk.liquid: {}", self.stats.disk_liquid_entries)?; writeln!(f, "entries.disk.arrow: {}", self.stats.disk_arrow_entries)?; writeln!(f, "usage.memory_bytes: {}", self.stats.memory_usage_bytes)?; @@ -98,7 +89,7 @@ impl fmt::Display for CacheStatsSummary { } async fn create_session_context_with_liquid_cache( - squeeze_policy: Box, + eviction_policy: Box, cache_size_bytes: usize, cache_dir: &Path, ) -> Result<(SessionContext, LiquidCacheParquetRef)> { @@ -112,7 +103,7 @@ async fn create_session_context_with_liquid_cache( .with_prefetch(false) .with_max_memory_bytes(cache_size_bytes) .with_cache_dir(cache_dir.to_path_buf()) - .with_squeeze_policy(squeeze_policy) + .with_eviction_policy(eviction_policy) .with_cache_policy(Box::new(LiquidPolicy::new())) .build(config) .await?; @@ -181,12 +172,12 @@ async fn prefetch_matches_lazy_io() { async fn run_sql_with_cache( sql: &str, - squeeze_policy: Box, + eviction_policy: Box, cache_size_bytes: usize, cache_dir: &Path, ) -> QueryOutcome { let (ctx, cache) = - create_session_context_with_liquid_cache(squeeze_policy, cache_size_bytes, cache_dir) + create_session_context_with_liquid_cache(eviction_policy, cache_size_bytes, cache_dir) .await .unwrap(); @@ -217,14 +208,11 @@ async fn test_runner(sql: &str, reference: &str, cache_dir: &Path) { let cache_sizes = [10 * 1024, 1024 * 1024, usize::MAX]; // 10KB, 1MB, unlimited for cache_size in cache_sizes { - let squeeze_policies: Vec> = vec![ - Box::new(TranscodeSqueezeEvict), - Box::new(Evict), - Box::new(TranscodeEvict), - ]; - for squeeze_policy in squeeze_policies { + let eviction_policies: Vec> = + vec![Box::new(TranscodeEvict), Box::new(Evict)]; + for eviction_policy in eviction_policies { let QueryOutcome { values, .. } = - run_sql_with_cache(sql, squeeze_policy, cache_size, cache_dir).await; + run_sql_with_cache(sql, eviction_policy, cache_size, cache_dir).await; assert_eq!( values, reference, "Results differ, cache_size: {cache_size}" @@ -242,13 +230,7 @@ async fn test_url_prefix_filtering() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 1024, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 1024, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -271,13 +253,7 @@ async fn test_url_selection_and_ordering() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 300, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 300, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -300,13 +276,7 @@ async fn test_os_selection() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 1024, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 1024, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -330,13 +300,7 @@ async fn test_referer_filtering() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 1024, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 1024, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -360,13 +324,7 @@ async fn test_single_column_filter_projection() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 1024, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 1024, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -393,7 +351,7 @@ async fn test_provide_schema2() { .with_prefetch(false) .with_cache_dir(cache_dir.path().to_path_buf()) .with_max_memory_bytes(1024 * 1024) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(config) .await .unwrap(); @@ -468,7 +426,7 @@ async fn test_provide_schema2() { #[cfg(target_arch = "aarch64")] let snapshot = snapshot .replace("usage.memory_bytes: 999980", "usage.memory_bytes: 1000915") - .replace("usage.memory_bytes: 1035369", "usage.memory_bytes: 1036304"); + .replace("usage.memory_bytes: 1035368", "usage.memory_bytes: 1036303"); insta::assert_snapshot!(snapshot); } @@ -482,13 +440,7 @@ async fn test_provide_schema_with_filter() { values, plan, stats, - } = run_sql_with_cache( - sql, - Box::new(TranscodeSqueezeEvict), - 1024 * 1024, - cache_dir.path(), - ) - .await; + } = run_sql_with_cache(sql, Box::new(TranscodeEvict), 1024 * 1024, cache_dir.path()).await; assert!(stats.has_cache_hits()); assert!(stats.entries_reused()); @@ -501,7 +453,7 @@ async fn test_provide_schema_with_filter() { )); let (ctx, _) = LiquidCacheLocalBuilder::new() - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -545,7 +497,7 @@ async fn test_repartitioned_file_scan_cache_correctness() { let reference = run_sql_with_cache( sql, - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), 1024 * 1024, reference_cache_dir.path(), ) @@ -559,7 +511,7 @@ async fn test_repartitioned_file_scan_cache_correctness() { let (ctx, cache) = LiquidCacheLocalBuilder::new() .with_max_memory_bytes(1024 * 1024) .with_cache_dir(parallel_cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .with_cache_policy(Box::new(LiquidPolicy::new())) .build(config) .await diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap index fe7b869c7..a2b47c338 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap @@ -36,26 +36,19 @@ values: stats: entries.total: 8 entries.after_first_run: 8 -entries.memory.arrow: 3 -entries.memory.liquid: 4 -entries.memory.squeezed_liquid: 1 -entries.disk.liquid: 0 +entries.memory.arrow: 0 +entries.memory.liquid: 7 +entries.disk.liquid: 1 entries.disk.arrow: 0 -usage.memory_bytes: 954618 -usage.disk_bytes: 164536 +usage.memory_bytes: 329550 +usage.disk_bytes: 468740 RuntimeStatsSnapshot: get: 3 get_with_selection: 3 eval_predicate: 4 - get_squeezed_success: 0 - get_squeezed_needs_io: 1 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 - read_io_count: 3 - write_io_count: 0 + read_io_count: 5 + write_io_count: 3 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 2141 - squeezed_total_count: 2164 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap index c127d3fe1..d6e659e63 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap @@ -1,6 +1,5 @@ --- source: src/datafusion-local/src/tests/mod.rs -assertion_line: 428 expression: snapshot --- query[0]: SELECT * from default where log like '%hhj%' order by _timestamp @@ -24,7 +23,6 @@ entries.total: 8 entries.after_first_run: 8 entries.memory.arrow: 5 entries.memory.liquid: 1 -entries.memory.squeezed_liquid: 0 entries.disk.liquid: 2 entries.disk.arrow: 0 usage.memory_bytes: 1000915 @@ -33,18 +31,12 @@ RuntimeStatsSnapshot: get: 4 get_with_selection: 4 eval_predicate: 2 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 query[1]: SELECT date_bin(interval '10 second', to_timestamp_micros(_timestamp), to_timestamp('2001-01-01T00:00:00')) AS zo_sql_key, count(*) AS zo_sql_num from default WHERE log like '%hhj%' or message like '%hhj%' GROUP BY zo_sql_key ORDER BY zo_sql_key @@ -180,7 +172,6 @@ entries.total: 8 entries.after_first_run: 8 entries.memory.arrow: 5 entries.memory.liquid: 1 -entries.memory.squeezed_liquid: 0 entries.disk.liquid: 2 entries.disk.arrow: 0 usage.memory_bytes: 1000915 @@ -189,18 +180,12 @@ RuntimeStatsSnapshot: get: 5 get_with_selection: 5 eval_predicate: 0 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 3 - hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 query[2]: SELECT _timestamp, kubernetes_namespace_name from default order by _timestamp desc limit 100 @@ -226,24 +211,17 @@ entries.total: 8 entries.after_first_run: 8 entries.memory.arrow: 5 entries.memory.liquid: 3 -entries.memory.squeezed_liquid: 0 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 1036304 +usage.memory_bytes: 1036303 usage.disk_bytes: 35000 RuntimeStatsSnapshot: get: 4 get_with_selection: 4 eval_predicate: 0 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap index 0b214cc49..6ff89dd3e 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap @@ -46,7 +46,6 @@ entries.total: 12 entries.after_first_run: 12 entries.memory.arrow: 12 entries.memory.liquid: 0 -entries.memory.squeezed_liquid: 0 entries.disk.liquid: 0 entries.disk.arrow: 0 usage.memory_bytes: 574612 @@ -55,15 +54,9 @@ RuntimeStatsSnapshot: get: 13 get_with_selection: 13 eval_predicate: 4 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap index df9e16133..567317870 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap @@ -35,26 +35,19 @@ values: stats: entries.total: 8 entries.after_first_run: 8 -entries.memory.arrow: 1 -entries.memory.liquid: 7 -entries.memory.squeezed_liquid: 0 +entries.memory.arrow: 2 +entries.memory.liquid: 6 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 917947 +usage.memory_bytes: 951300 usage.disk_bytes: 877216 RuntimeStatsSnapshot: get: 2 get_with_selection: 2 eval_predicate: 8 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 - read_io_count: 5 + read_io_count: 4 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap index 5b5807c29..5973251b1 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap @@ -23,7 +23,6 @@ entries.total: 4 entries.after_first_run: 4 entries.memory.arrow: 4 entries.memory.liquid: 0 -entries.memory.squeezed_liquid: 0 entries.disk.liquid: 0 entries.disk.arrow: 0 usage.memory_bytes: 262528 @@ -32,15 +31,9 @@ RuntimeStatsSnapshot: get: 1 get_with_selection: 1 eval_predicate: 4 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 0 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__basic_squeeze.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__basic_squeeze.snap deleted file mode 100644 index 45d651a98..000000000 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__basic_squeeze.snap +++ /dev/null @@ -1,74 +0,0 @@ ---- -source: src/datafusion-local/src/tests/squeeze.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=0 kind=MemoryArrow -event=insert_failed entry=262144 kind=MemoryArrow -event=squeeze_begin victims=[0] -event=squeeze_victim entry=0 -event=insert_success entry=0 kind=MemoryLiquid -event=insert_success entry=262144 kind=MemoryArrow -event=eval_predicate entry=0 selection=true cached=MemoryLiquid -event=read entry=262144 selection=true expr=None cached=MemoryArrow -event=insert_failed entry=1 kind=MemoryArrow -event=squeeze_begin victims=[262144,0] -event=squeeze_victim entry=262144 -event=insert_success entry=262144 kind=MemoryLiquid -event=squeeze_victim entry=0 -event=io_write entry=0 kind=MemorySqueezedLiquid bytes=63528 -event=insert_success entry=0 kind=MemorySqueezedLiquid -event=insert_success entry=1 kind=MemoryArrow -event=insert_failed entry=262145 kind=MemoryArrow -event=squeeze_begin victims=[1,262144,0] -event=squeeze_victim entry=1 -event=insert_success entry=1 kind=MemoryLiquid -event=squeeze_victim entry=262144 -event=io_write entry=262144 kind=DiskLiquid bytes=17448 -event=insert_success entry=262144 kind=DiskLiquid -event=squeeze_victim entry=0 -event=insert_success entry=0 kind=DiskLiquid -event=insert_success entry=262145 kind=MemoryArrow -event=eval_predicate entry=1 selection=true cached=MemoryLiquid -event=read entry=262145 selection=true expr=None cached=MemoryArrow -event=insert_failed entry=2 kind=MemoryArrow -event=squeeze_begin victims=[262145,1] -event=squeeze_victim entry=262145 -event=insert_success entry=262145 kind=MemoryLiquid -event=squeeze_victim entry=1 -event=io_write entry=1 kind=MemorySqueezedLiquid bytes=63528 -event=insert_success entry=1 kind=MemorySqueezedLiquid -event=insert_success entry=2 kind=MemoryArrow -event=insert_failed entry=262146 kind=MemoryArrow -event=squeeze_begin victims=[2,262145,1] -event=squeeze_victim entry=2 -event=insert_success entry=2 kind=MemoryLiquid -event=squeeze_victim entry=262145 -event=io_write entry=262145 kind=DiskLiquid bytes=17448 -event=insert_success entry=262145 kind=DiskLiquid -event=squeeze_victim entry=1 -event=insert_success entry=1 kind=DiskLiquid -event=insert_success entry=262146 kind=MemoryArrow -event=eval_predicate entry=2 selection=true cached=MemoryLiquid -event=read entry=262146 selection=true expr=None cached=MemoryArrow -event=insert_failed entry=4294967296 kind=MemoryArrow -event=squeeze_begin victims=[262146,2] -event=squeeze_victim entry=262146 -event=insert_success entry=262146 kind=MemoryLiquid -event=squeeze_victim entry=2 -event=io_write entry=2 kind=MemorySqueezedLiquid bytes=63528 -event=insert_success entry=2 kind=MemorySqueezedLiquid -event=insert_success entry=4294967296 kind=MemoryArrow -event=insert_failed entry=4295229440 kind=MemoryArrow -event=squeeze_begin victims=[4294967296,262146,2] -event=squeeze_victim entry=4294967296 -event=insert_success entry=4294967296 kind=MemoryLiquid -event=squeeze_victim entry=262146 -event=io_write entry=262146 kind=DiskLiquid bytes=17448 -event=insert_success entry=262146 kind=DiskLiquid -event=squeeze_victim entry=2 -event=insert_success entry=2 kind=DiskLiquid -event=insert_success entry=4295229440 kind=MemoryArrow -event=eval_predicate entry=4294967296 selection=true cached=MemoryLiquid -event=read entry=4295229440 selection=true expr=None cached=MemoryArrow -] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_distinct_search_phase.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_distinct_search_phase.snap deleted file mode 100644 index fa1468071..000000000 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_distinct_search_phase.snap +++ /dev/null @@ -1,32 +0,0 @@ ---- -source: src/datafusion-local/src/tests/squeeze.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=2555904 kind=MemoryArrow -event=eval_predicate entry=2555904 selection=true cached=MemoryArrow -event=read entry=2555904 selection=true expr=None cached=MemoryArrow -event=insert_success entry=2555905 kind=MemoryArrow -event=eval_predicate entry=2555905 selection=true cached=MemoryArrow -event=read entry=2555905 selection=true expr=None cached=MemoryArrow -event=insert_failed entry=2555906 kind=MemoryArrow -event=squeeze_begin victims=[2555904,2555905] -event=squeeze_victim entry=2555904 -event=insert_success entry=2555904 kind=MemoryLiquid -event=squeeze_victim entry=2555905 -event=insert_success entry=2555905 kind=MemoryLiquid -event=insert_success entry=2555906 kind=MemoryArrow -event=eval_predicate entry=2555906 selection=true cached=MemoryArrow -event=read entry=2555906 selection=true expr=None cached=MemoryArrow -event=insert_success entry=4297523200 kind=MemoryArrow -event=eval_predicate entry=4297523200 selection=true cached=MemoryArrow -event=read entry=4297523200 selection=true expr=None cached=MemoryArrow -event=eval_predicate entry=2555904 selection=true cached=MemoryLiquid -event=read entry=2555904 selection=true expr=None cached=MemoryLiquid -event=eval_predicate entry=2555905 selection=true cached=MemoryLiquid -event=read entry=2555905 selection=true expr=None cached=MemoryLiquid -event=eval_predicate entry=2555906 selection=true cached=MemoryArrow -event=read entry=2555906 selection=true expr=None cached=MemoryArrow -event=eval_predicate entry=4297523200 selection=true cached=MemoryArrow -event=read entry=4297523200 selection=true expr=None cached=MemoryArrow -] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap deleted file mode 100644 index 2943782b2..000000000 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap +++ /dev/null @@ -1,75 +0,0 @@ ---- -source: src/datafusion-local/src/tests/squeeze.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=851968 kind=MemoryArrow -event=insert_failed entry=917504 kind=MemoryArrow -event=squeeze_begin victims=[851968] -event=squeeze_victim entry=851968 -event=insert_success entry=851968 kind=MemoryLiquid -event=insert_failed entry=917504 kind=MemoryArrow -event=squeeze_begin victims=[851968] -event=squeeze_victim entry=851968 -event=io_write entry=851968 kind=DiskLiquid bytes=139416 -event=insert_success entry=851968 kind=DiskLiquid -event=insert_failed entry=917504 kind=MemoryArrow -event=insert_success entry=917504 kind=MemoryLiquid -event=eval_predicate entry=917504 selection=true cached=MemoryLiquid -event=read entry=851968 selection=true expr=None cached=DiskLiquid -event=io_read_liquid entry=851968 bytes=139416 -event=hydrate entry=851968 cached=DiskLiquid new=MemoryLiquid -event=insert_success entry=851968 kind=MemoryLiquid -event=insert_success entry=851969 kind=MemoryArrow -event=insert_failed entry=917505 kind=MemoryArrow -event=squeeze_begin victims=[851969,917504,851968] -event=squeeze_victim entry=851969 -event=insert_success entry=851969 kind=MemoryLiquid -event=squeeze_victim entry=917504 -event=io_write entry=917504 kind=MemorySqueezedLiquid bytes=136440 -event=insert_success entry=917504 kind=MemorySqueezedLiquid -event=squeeze_victim entry=851968 -event=io_write entry=851968 kind=DiskLiquid bytes=139416 -event=insert_success entry=851968 kind=DiskLiquid -event=insert_success entry=917505 kind=MemoryArrow -event=eval_predicate entry=917505 selection=true cached=MemoryArrow -event=read entry=851969 selection=true expr=None cached=MemoryLiquid -event=insert_failed entry=851970 kind=MemoryArrow -event=squeeze_begin victims=[917505,851969,917504] -event=squeeze_victim entry=917505 -event=insert_success entry=917505 kind=MemoryLiquid -event=squeeze_victim entry=851969 -event=io_write entry=851969 kind=DiskLiquid bytes=139376 -event=insert_success entry=851969 kind=DiskLiquid -event=squeeze_victim entry=917504 -event=insert_success entry=917504 kind=DiskLiquid -event=insert_success entry=851970 kind=MemoryArrow -event=insert_failed entry=917506 kind=MemoryArrow -event=squeeze_begin victims=[851970,917505] -event=squeeze_victim entry=851970 -event=insert_success entry=851970 kind=MemoryLiquid -event=squeeze_victim entry=917505 -event=io_write entry=917505 kind=MemorySqueezedLiquid bytes=141576 -event=insert_success entry=917505 kind=MemorySqueezedLiquid -event=insert_failed entry=917506 kind=MemoryArrow -event=squeeze_begin victims=[851970,917505] -event=squeeze_victim entry=851970 -event=io_write entry=851970 kind=DiskLiquid bytes=146184 -event=insert_success entry=851970 kind=DiskLiquid -event=squeeze_victim entry=917505 -event=insert_success entry=917505 kind=DiskLiquid -event=insert_success entry=917506 kind=MemoryArrow -event=eval_predicate entry=917506 selection=true cached=MemoryArrow -event=read entry=851970 selection=true expr=None cached=DiskLiquid -event=io_read_liquid entry=851970 bytes=146184 -event=hydrate entry=851970 cached=DiskLiquid new=MemoryLiquid -event=insert_failed entry=851970 kind=MemoryLiquid -event=squeeze_begin victims=[917506] -event=squeeze_victim entry=917506 -event=insert_success entry=917506 kind=MemoryLiquid -event=insert_success entry=851970 kind=MemoryLiquid -event=insert_success entry=4295819264 kind=MemoryArrow -event=insert_success entry=4295884800 kind=MemoryArrow -event=eval_predicate entry=4295884800 selection=true cached=MemoryArrow -event=read entry=4295819264 selection=true expr=None cached=MemoryArrow -] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search.snap deleted file mode 100644 index 9a8975b9e..000000000 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search.snap +++ /dev/null @@ -1,20 +0,0 @@ ---- -source: src/datafusion-local/src/tests/squeeze.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=2555904 kind=MemoryArrow -event=eval_predicate entry=2555904 selection=true cached=MemoryArrow -event=insert_success entry=2555905 kind=MemoryArrow -event=eval_predicate entry=2555905 selection=true cached=MemoryArrow -event=insert_failed entry=2555906 kind=MemoryArrow -event=squeeze_begin victims=[2555904,2555905] -event=squeeze_victim entry=2555904 -event=insert_success entry=2555904 kind=MemoryLiquid -event=squeeze_victim entry=2555905 -event=insert_success entry=2555905 kind=MemoryLiquid -event=insert_success entry=2555906 kind=MemoryArrow -event=eval_predicate entry=2555906 selection=true cached=MemoryArrow -event=insert_success entry=4297523200 kind=MemoryArrow -event=eval_predicate entry=4297523200 selection=true cached=MemoryArrow -] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search_title.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search_title.snap deleted file mode 100644 index 12ee8a2cd..000000000 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_substrings_search_title.snap +++ /dev/null @@ -1,20 +0,0 @@ ---- -source: src/datafusion-local/src/tests/squeeze.rs -expression: trace ---- -EventTrace: [ -event=insert_success entry=131072 kind=MemoryArrow -event=eval_predicate entry=131072 selection=true cached=MemoryArrow -event=insert_success entry=131073 kind=MemoryArrow -event=eval_predicate entry=131073 selection=true cached=MemoryArrow -event=insert_failed entry=131074 kind=MemoryArrow -event=squeeze_begin victims=[131072,131073] -event=squeeze_victim entry=131072 -event=insert_success entry=131072 kind=MemoryLiquid -event=squeeze_victim entry=131073 -event=insert_success entry=131073 kind=MemoryLiquid -event=insert_success entry=131074 kind=MemoryArrow -event=eval_predicate entry=131074 selection=true cached=MemoryArrow -event=insert_success entry=4295098368 kind=MemoryArrow -event=eval_predicate entry=4295098368 selection=true cached=MemoryArrow -] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap index e83ef4900..7539741d9 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap @@ -52,25 +52,18 @@ stats: entries.total: 4 entries.after_first_run: 4 entries.memory.arrow: 2 -entries.memory.liquid: 1 -entries.memory.squeezed_liquid: 1 +entries.memory.liquid: 2 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 912473 +usage.memory_bytes: 1018212 usage.disk_bytes: 139416 RuntimeStatsSnapshot: get: 4 get_with_selection: 4 eval_predicate: 0 - get_squeezed_success: 0 - get_squeezed_needs_io: 1 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 1 write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap index 7b6d7f6dd..3e480e87a 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap @@ -38,24 +38,17 @@ entries.total: 4 entries.after_first_run: 4 entries.memory.arrow: 0 entries.memory.liquid: 2 -entries.memory.squeezed_liquid: 1 -entries.disk.liquid: 1 +entries.disk.liquid: 2 entries.disk.arrow: 0 -usage.memory_bytes: 189553 +usage.memory_bytes: 147293 usage.disk_bytes: 706252 RuntimeStatsSnapshot: get: 3 get_with_selection: 3 eval_predicate: 4 - get_squeezed_success: 0 - get_squeezed_needs_io: 0 try_read_liquid_calls: 0 - hit_date32_expression_calls: 0 read_io_count: 4 write_io_count: 4 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 - squeezed_decompressed_count: 0 - squeezed_total_count: 0 - squeeze_io_saved: 0 diff --git a/src/datafusion-local/src/tests/squeeze.rs b/src/datafusion-local/src/tests/squeeze.rs deleted file mode 100644 index f2c40dfc6..000000000 --- a/src/datafusion-local/src/tests/squeeze.rs +++ /dev/null @@ -1,150 +0,0 @@ -use arrow::{array::AsArray, datatypes::Int64Type, util::pretty::pretty_format_batches}; -use datafusion::prelude::SessionConfig; -use tempfile::TempDir; - -use crate::LiquidCacheLocalBuilder; - -const TEST_FILE: &str = "../../examples/nano_hits.parquet"; - -fn squeeze_test_config() -> SessionConfig { - SessionConfig::new().with_repartition_file_scans(false) -} - -#[tokio::test] -async fn basic_squeeze() { - let cache_dir = TempDir::new().unwrap(); - let (ctx, cache) = LiquidCacheLocalBuilder::new() - .with_prefetch(false) - .with_max_memory_bytes(1024 * 128) - .with_cache_dir(cache_dir.path().to_path_buf()) - .build(squeeze_test_config()) - .await - .unwrap(); - ctx.register_parquet("hits", TEST_FILE, Default::default()) - .await - .unwrap(); - - let plan = ctx - .sql("SELECT COUNT(DISTINCT(\"EventTime\")) FROM hits WHERE \"WatchID\" <> 0") - .await - .unwrap(); - let result = plan.collect().await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!( - result[0].column(0).as_primitive::().value(0), - 12385 - ); - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} - -#[tokio::test] -async fn squeeze_strings() { - let cache_dir = TempDir::new().unwrap(); - let (ctx, cache) = LiquidCacheLocalBuilder::new() - .with_prefetch(false) - .with_max_memory_bytes(1024 * 1024) - .with_cache_dir(cache_dir.path().to_path_buf()) - .build(squeeze_test_config()) - .await - .unwrap(); - ctx.register_parquet("hits", TEST_FILE, Default::default()) - .await - .unwrap(); - - let plan = ctx - .sql("SELECT COUNT(DISTINCT(\"URL\")) FROM hits WHERE \"Referer\" <> '0'") - .await - .unwrap(); - let result = plan.collect().await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!( - result[0].column(0).as_primitive::().value(0), - 5639 - ); - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} - -#[tokio::test] -async fn squeeze_substrings_search() { - let cache_dir = TempDir::new().unwrap(); - let (ctx, cache) = LiquidCacheLocalBuilder::new() - .with_prefetch(false) - .with_max_memory_bytes(1024 * 256) - .with_cache_dir(cache_dir.path().to_path_buf()) - .build(squeeze_test_config()) - .await - .unwrap(); - ctx.register_parquet("hits", TEST_FILE, Default::default()) - .await - .unwrap(); - - let plan = ctx - .sql("SELECT COUNT(*) FROM hits WHERE \"SearchPhrase\" LIKE '%abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789%'") - .await - .unwrap(); - let result = plan.collect().await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0].column(0).as_primitive::().value(0), 0); - let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); -} - -#[tokio::test] -async fn squeeze_substrings_search_title() { - let cache_dir = TempDir::new().unwrap(); - let (ctx, cache) = LiquidCacheLocalBuilder::new() - .with_prefetch(false) - .with_max_memory_bytes(1024 * 1024 * 4) - .with_cache_dir(cache_dir.path().to_path_buf()) - .build(squeeze_test_config()) - .await - .unwrap(); - ctx.register_parquet("hits", TEST_FILE, Default::default()) - .await - .unwrap(); - - let plan = ctx - .sql("SELECT COUNT(*) FROM hits WHERE \"Title\" LIKE '%Cosplay%'") - .await - .unwrap(); - let result = plan.collect().await.unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0].column(0).as_primitive::().value(0), 6); - let trace = cache.consume_event_trace(); - println!("{:?}", cache.storage().stats()); - insta::assert_snapshot!(trace); -} - -#[tokio::test] -async fn squeeze_distinct_search_phase() { - let cache_dir = TempDir::new().unwrap(); - let (ctx, cache) = LiquidCacheLocalBuilder::new() - .with_prefetch(false) - .with_max_memory_bytes(1024 * 256) - .with_cache_dir(cache_dir.path().to_path_buf()) - .build(squeeze_test_config()) - .await - .unwrap(); - ctx.register_parquet("hits", TEST_FILE, Default::default()) - .await - .unwrap(); - - let plan = ctx - .sql("SELECT DISTINCT(\"SearchPhrase\") FROM hits ORDER BY \"SearchPhrase\" LIMIT 10") - .await - .unwrap(); - let result = plan.collect().await.unwrap(); - println!("{}", pretty_format_batches(result.as_ref()).unwrap()); - assert_eq!(result.len(), 1); - ctx.sql("SELECT DISTINCT(\"SearchPhrase\") FROM hits ORDER BY \"SearchPhrase\" LIMIT 10") - .await - .unwrap() - .collect() - .await - .unwrap(); - let trace = cache.consume_event_trace(); - println!("{:?}", cache.storage().stats()); - insta::assert_snapshot!(trace); -} diff --git a/src/datafusion-local/src/tests/variants.rs b/src/datafusion-local/src/tests/variants.rs index aa550ab51..cd76a88d9 100644 --- a/src/datafusion-local/src/tests/variants.rs +++ b/src/datafusion-local/src/tests/variants.rs @@ -10,7 +10,7 @@ use arrow::{ }; use arrow_schema::{Field, Fields, Schema}; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; -use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache::cache::TranscodeEvict; use parquet::{ arrow::ArrowWriter, variant::{VariantArray, VariantType, json_to_variant}, @@ -135,7 +135,7 @@ async fn test_variant_parquet_naive_read() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -172,7 +172,7 @@ async fn test_variant_transcoding_falls_back_to_disk_arrow() { .with_batch_size(1) .with_max_memory_bytes(64) .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -236,7 +236,7 @@ async fn test_variant_get() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -268,7 +268,7 @@ async fn test_variant_predicate() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -300,7 +300,7 @@ async fn test_variant_get_fails_when_value_field_not_nullable() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -367,7 +367,7 @@ fn write_large_variant_parquet_file(dir: &Path, num_rows: usize) -> PathBuf { } #[tokio::test] -async fn test_large_variant_squeeze() { +async fn test_large_variant_under_memory_pressure() { let cache_dir = TempDir::new().unwrap(); let parquet_dir = TempDir::new().unwrap(); let num_rows = 1_000; @@ -377,7 +377,7 @@ async fn test_large_variant_squeeze() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) .with_max_memory_bytes(1024) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -417,7 +417,7 @@ async fn variant_multi_queries() { let (ctx, _cache) = LiquidCacheLocalBuilder::new() .with_cache_dir(cache_dir.path().to_path_buf()) .with_max_memory_bytes(1024) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); @@ -467,7 +467,7 @@ async fn variant_multi_queries_complex() { .with_cache_dir(cache_dir.path().to_path_buf()) .with_max_memory_bytes(1024 * 600) .with_batch_size(8) - .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_eviction_policy(Box::new(TranscodeEvict)) .build(SessionConfig::new()) .await .unwrap(); diff --git a/src/datafusion-server/src/lib.rs b/src/datafusion-server/src/lib.rs index a12b6dd4d..fcaf0a772 100644 --- a/src/datafusion-server/src/lib.rs +++ b/src/datafusion-server/src/lib.rs @@ -38,7 +38,7 @@ use fastrace::prelude::SpanContext; use futures::{Stream, TryStreamExt}; use liquid_cache::cache::CacheExpression; use liquid_cache_common::rpc::{FetchResults, LiquidCacheActions}; -use liquid_cache_datafusion::cache::{ColumnSqueezeHints, LiquidCacheParquetRef}; +use liquid_cache_datafusion::cache::{ColumnLineages, LiquidCacheParquetRef}; use log::info; use prost::bytes::Bytes; use service::LiquidCacheServiceInner; @@ -59,8 +59,7 @@ pub use errors::{ pub use liquid_cache as storage; use liquid_cache::{ cache::{ - AlwaysHydrate, HydrationPolicy, - squeeze_policies::{SqueezePolicy, TranscodeSqueezeEvict}, + AlwaysHydrate, HydrationPolicy, {EvictionPolicy, TranscodeEvict}, }, cache_policies::{CachePolicy, LiquidPolicy}, }; @@ -79,7 +78,7 @@ mod tests; /// use arrow_flight::flight_service_server::FlightServiceServer; /// use datafusion::prelude::SessionContext; /// use liquid_cache_datafusion_server::LiquidCacheService; -/// use liquid_cache_datafusion_server::storage::cache::squeeze_policies::TranscodeSqueezeEvict; +/// use liquid_cache_datafusion_server::storage::cache::TranscodeEvict; /// use liquid_cache_datafusion_server::storage::cache::AlwaysHydrate; /// use liquid_cache_datafusion_server::storage::cache_policies::LiquidPolicy; /// use tonic::transport::Server; @@ -90,7 +89,7 @@ mod tests; /// None, /// None, /// Box::new(LiquidPolicy::new()), -/// Box::new(TranscodeSqueezeEvict), +/// Box::new(TranscodeEvict), /// Box::new(AlwaysHydrate::new()), /// ) /// .await @@ -115,7 +114,7 @@ impl LiquidCacheService { None, None, Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await @@ -133,7 +132,7 @@ impl LiquidCacheService { max_memory_bytes: Option, disk_cache_dir: Option, cache_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, hydration_policy: Box, ) -> anyhow::Result { let disk_cache_dir = match disk_cache_dir { @@ -150,7 +149,7 @@ impl LiquidCacheService { max_memory_bytes, disk_cache_dir, cache_policy, - squeeze_policy, + eviction_policy, hydration_policy, ) .await, @@ -253,13 +252,13 @@ impl LiquidCacheService { LiquidCacheActions::RegisterPlan(cmd) => { let plan = physical_plan_from_bytes(&cmd.plan, &self.inner.get_ctx().task_ctx())?; let handle = Uuid::from_bytes_ref(cmd.handle.as_ref().try_into()?); - let mut squeeze_hints = ColumnSqueezeHints::default(); - for hint in &cmd.squeeze_hints { + let mut lineages = ColumnLineages::default(); + for hint in &cmd.lineages { if let Some(expr) = CacheExpression::from_metadata_value(&hint.hint) { - squeeze_hints.insert(hint.column.clone(), Arc::new(expr)); + lineages.insert(hint.column.clone(), Arc::new(expr)); } } - self.inner.register_plan(*handle, plan, squeeze_hints); + self.inner.register_plan(*handle, plan, lineages); let output = futures::stream::iter(vec![Ok(arrow_flight::Result { body: Bytes::default(), })]); diff --git a/src/datafusion-server/src/service.rs b/src/datafusion-server/src/service.rs index 06015fa3b..04c671243 100644 --- a/src/datafusion-server/src/service.rs +++ b/src/datafusion-server/src/service.rs @@ -5,13 +5,13 @@ use datafusion::{ physical_plan::{ExecutionPlan, display::DisplayableExecutionPlan}, prelude::SessionContext, }; -use liquid_cache::{ByteCache, cache::squeeze_policies::SqueezePolicy}; +use liquid_cache::{ByteCache, cache::EvictionPolicy}; use liquid_cache::{cache::HydrationPolicy, cache_policies::CachePolicy}; use liquid_cache_common::rpc::ExecutionMetricsResponse; use liquid_cache_datafusion::{ - cache::{ColumnSqueezeHints, LiquidCacheParquet, LiquidCacheParquetRef}, + cache::{ColumnLineages, LiquidCacheParquet, LiquidCacheParquetRef}, extract_execution_metrics, - optimizers::{SqueezeHintMap, rewrite_data_source_plan_with_hints}, + optimizers::{LineageHints, rewrite_data_source_plan_with_hints}, }; use log::{debug, info}; use object_store::ObjectStore; @@ -49,7 +49,7 @@ impl LiquidCacheServiceInner { max_memory_bytes: Option, disk_cache_dir: PathBuf, cache_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, hydration_policy: Box, ) -> Self { let batch_size = default_ctx.state().config().batch_size(); @@ -68,7 +68,7 @@ impl LiquidCacheServiceInner { usize::MAX, store, cache_policy, - squeeze_policy, + eviction_policy, hydration_policy, ) .await, @@ -151,15 +151,15 @@ impl LiquidCacheServiceInner { &self, handle: Uuid, plan: Arc, - squeeze_hints: ColumnSqueezeHints, + lineages: ColumnLineages, ) { let cache = self.cache(); // Hints the server can derive from the fragment itself (e.g. a // `date_part` inside a pushed-down partial aggregate). The client ships // hints for lineage that only exists in the client-side part of the // plan; those take precedence on conflict. - let mut hints = SqueezeHintMap::analyze(&plan).for_fragment(&plan); - hints.extend(squeeze_hints); + let mut hints = LineageHints::analyze(&plan).for_fragment(&plan); + hints.extend(lineages); self.execution_plans.write().unwrap().insert( handle, ExecutionPlanEntry::new(rewrite_data_source_plan_with_hints(plan, cache, &hints)), @@ -228,7 +228,7 @@ impl LiquidCacheServiceInner { mod tests { use super::*; use liquid_cache::{ - cache::{AlwaysHydrate, squeeze_policies::TranscodeSqueezeEvict}, + cache::{AlwaysHydrate, TranscodeEvict}, cache_policies::LiquidPolicy, }; #[tokio::test] @@ -239,7 +239,7 @@ mod tests { None, temp_dir.path().to_path_buf(), Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await; diff --git a/src/datafusion-server/src/tests/cases.rs b/src/datafusion-server/src/tests/cases.rs index 856ec9e47..86a898ef3 100644 --- a/src/datafusion-server/src/tests/cases.rs +++ b/src/datafusion-server/src/tests/cases.rs @@ -2,7 +2,7 @@ use std::fs::File; use std::path::{Path, PathBuf}; use std::sync::Arc; -use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache::cache::TranscodeEvict; use crate::tests::run_sql; @@ -34,7 +34,7 @@ async fn test_parquet_with_page_index() { let result = run_sql( "SELECT * FROM hits WHERE id = 0", - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), 1000, file_path, ) diff --git a/src/datafusion-server/src/tests/mod.rs b/src/datafusion-server/src/tests/mod.rs index 56210b51f..fac8ea380 100644 --- a/src/datafusion-server/src/tests/mod.rs +++ b/src/datafusion-server/src/tests/mod.rs @@ -6,10 +6,7 @@ use datafusion::{ prelude::SessionContext, }; use liquid_cache::{ - cache::{ - AlwaysHydrate, - squeeze_policies::{Evict, SqueezePolicy, TranscodeEvict, TranscodeSqueezeEvict}, - }, + cache::{AlwaysHydrate, Evict, EvictionPolicy, TranscodeEvict}, cache_policies::LiquidPolicy, }; use uuid::Uuid; @@ -28,7 +25,7 @@ async fn get_physical_plan(sql: &str, ctx: &SessionContext) -> Arc, + eviction_policy: Box, cache_size_bytes: usize, file_path: &str, ) -> String { @@ -42,7 +39,7 @@ async fn run_sql( Some(cache_size_bytes), tmp_dir.path().to_path_buf(), Box::new(LiquidPolicy::new()), - squeeze_policy, + eviction_policy, Box::new(AlwaysHydrate::new()), ) .await; @@ -69,11 +66,7 @@ async fn test_runner(sql: &str, reference: &str) { let sizes = [10, 573960, usize::MAX]; for size in sizes { - let policies: [Box; 3] = [ - Box::new(TranscodeSqueezeEvict), - Box::new(Evict), - Box::new(TranscodeEvict), - ]; + let policies: [Box; 2] = [Box::new(TranscodeEvict), Box::new(Evict)]; for policy in policies { let result = run_sql(sql, policy, size, TEST_FILE).await; assert_eq!(result, reference); @@ -84,7 +77,7 @@ async fn test_runner(sql: &str, reference: &str) { #[tokio::test] async fn test_url_prefix() { let sql = r#"select COUNT(*) from hits where "URL" like 'https://%'"#; - let reference = run_sql(sql, Box::new(TranscodeSqueezeEvict), 573960, TEST_FILE).await; + let reference = run_sql(sql, Box::new(TranscodeEvict), 573960, TEST_FILE).await; insta::assert_snapshot!(reference); test_runner(sql, &reference).await; } @@ -92,7 +85,7 @@ async fn test_url_prefix() { #[tokio::test] async fn test_url() { let sql = r#"select "URL" from hits where "URL" like '%tours%' order by "URL" desc"#; - let reference = run_sql(sql, Box::new(TranscodeSqueezeEvict), 573960, TEST_FILE).await; + let reference = run_sql(sql, Box::new(TranscodeEvict), 573960, TEST_FILE).await; insta::assert_snapshot!(reference); test_runner(sql, &reference).await; } @@ -100,7 +93,7 @@ async fn test_url() { #[tokio::test] async fn test_os() { let sql = r#"select "OS" from hits where "URL" like '%tours%' order by "OS" desc"#; - let reference = run_sql(sql, Box::new(TranscodeSqueezeEvict), 573960, TEST_FILE).await; + let reference = run_sql(sql, Box::new(TranscodeEvict), 573960, TEST_FILE).await; insta::assert_snapshot!(reference); test_runner(sql, &reference).await; } @@ -108,7 +101,7 @@ async fn test_os() { #[tokio::test] async fn test_referer() { let sql = r#"select "Referer" from hits where "Referer" <> '' AND "URL" like '%tours%' order by "Referer" desc"#; - let reference = run_sql(sql, Box::new(TranscodeSqueezeEvict), 573960, TEST_FILE).await; + let reference = run_sql(sql, Box::new(TranscodeEvict), 573960, TEST_FILE).await; insta::assert_snapshot!(reference); test_runner(sql, &reference).await; } @@ -116,7 +109,7 @@ async fn test_referer() { #[tokio::test] async fn test_min_max() { let sql = r#"select min("Referer"), max("Referer") from hits where "Referer" <> '' AND "URL" like '%tours%'"#; - let reference = run_sql(sql, Box::new(TranscodeSqueezeEvict), 573960, TEST_FILE).await; + let reference = run_sql(sql, Box::new(TranscodeEvict), 573960, TEST_FILE).await; insta::assert_snapshot!(reference); test_runner(sql, &reference).await; } diff --git a/src/datafusion/README.md b/src/datafusion/README.md index 028c1610f..f2f0a4e39 100644 --- a/src/datafusion/README.md +++ b/src/datafusion/README.md @@ -2,37 +2,11 @@ Parquet reader with liquid array caching and optimized data formats. -## Squeeze-hint (lineage) pushdown - -A *squeeze hint* tells the cache how a column is used, so under memory pressure it -keeps only what the query needs instead of the whole column — e.g. just the `YEAR` -of a date read via `EXTRACT(YEAR FROM d)`, the paths of a `variant_get`, or a -substring fingerprint for `LIKE '%foo%'`. - -A physical optimizer rule drives the whole flow: - -```text - physical plan - | - v - LocalModeOptimizer (physical optimizer rule) - | 1. analyze the plan -> one CacheExpression per scan column - | 2. find the parquet scan (ParquetSource) - | 3. replace it with a LiquidParquetSource carrying the hints - v - LiquidParquetSource - | on open, passes the hints down as squeeze hints - v - liquid cache - | squeezes under memory pressure - v - keeps only the hinted form (e.g. YEAR); full data stays on disk -``` - -The analysis is conservative: a column used in a way the analyzer doesn't model -gets no hint, so the cache never drops data a query still needs. - -**Flight mode** is the same flow split across the wire: the server only sees the -pushed-down fragment (which may lack the lineage), so the client derives the hints -from the full plan and ships them with the plan, and the server attaches them to -the `LiquidParquetSource` it builds. +## Lineage expression pushdown + +LiquidCache analyzes physical plans to record how each file column is consumed, +including date extraction, variant paths, predicates, and substring searches. +These expressions continue to flow through local and distributed execution even +though the cache currently retains complete Arrow or Liquid arrays. Keeping the +analysis separate preserves the information needed for future representation +work without enabling partial-data storage today. diff --git a/src/datafusion/bench/filter_pushdown.rs b/src/datafusion/bench/filter_pushdown.rs index 2f77f974c..5eff8fd18 100644 --- a/src/datafusion/bench/filter_pushdown.rs +++ b/src/datafusion/bench/filter_pushdown.rs @@ -1,7 +1,7 @@ use arrow::buffer::BooleanBuffer; use divan::Bencher; use liquid_cache::cache::AlwaysHydrate; -use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; +use liquid_cache::cache::TranscodeEvict; use liquid_cache::cache_policies::LiquidPolicy; use liquid_cache_datafusion::cache::CachedColumn; use liquid_cache_datafusion::{FilterCandidateBuilder, LiquidPredicate}; @@ -47,7 +47,7 @@ fn setup_cache() -> (Arc, tempfile::TempDir) { usize::MAX, store, Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), )); let field = Arc::new(Field::new("test_column", DataType::Int32, false)); diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index fc02ccebc..c27c4ff8b 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -41,7 +41,6 @@ pub enum InsertArrowArrayError { pub(crate) enum PrefetchOutcome { Snapshotted, AlreadySnapshotted, - Squeezed, Missing, } @@ -60,21 +59,21 @@ impl CachedColumn { is_predicate_column: bool, snapshots: Arc, ) -> Self { - // Register the column's squeeze hint. Squeeze hints are column-scoped; + // Register the column's lineage expression. Lineage expressions are column-scoped; // `ParquetCacheMetadata` keys them by column (the batch id is masked // off), so registering once on any batch covers every batch. // // The read-path `expression` is the typed lineage hint derived from the // plan (date/variant/substring). A pure predicate column carries no such - // hint, but still registers `PredicateColumn` to guide squeezing — it is + // hint, but still registers `PredicateColumn` to describe its usage — it is // deliberately *not* stored as `expression`, since it does not change how // the column is materialized on read. - let squeeze_hint = expression + let lineage = expression .clone() .or_else(|| is_predicate_column.then(|| Arc::new(CacheExpression::PredicateColumn))); - if let Some(hint) = squeeze_hint { + if let Some(hint) = lineage { let hint_entry_id = column_access_path.entry_id(BatchID::from_raw(0)).into(); - cache_store.add_squeeze_hint(&hint_entry_id, hint); + cache_store.add_lineage(&hint_entry_id, hint); } Self { field, @@ -246,7 +245,6 @@ impl CachedColumn { self.snapshots.insert(entry_id, entry); PrefetchOutcome::Snapshotted } - PrefetchResult::Squeezed => PrefetchOutcome::Squeezed, PrefetchResult::Absent => PrefetchOutcome::Missing, } } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index d890f480f..fed837f6f 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -12,7 +12,7 @@ use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::Column; -use liquid_cache::cache::squeeze_policies::SqueezePolicy; +use liquid_cache::cache::EvictionPolicy; use liquid_cache::cache::{ CacheEntry, CacheExpression, CachePolicy, EntryID, EventTrace, HydrationPolicy, LiquidCache, LiquidCacheBuilder, @@ -32,12 +32,12 @@ pub(crate) use column::{InsertArrowArrayError, PrefetchOutcome}; pub(crate) use id::ColumnAccessPath; pub use id::{BatchID, ParquetArrayID}; -/// Typed squeeze hints for a single file, keyed by file-schema column name. +/// Typed lineage expressions for a single file, keyed by file-schema column name. /// -/// Produced by the physical squeeze-hint analyzer (local mode) or shipped from +/// Produced by the physical lineage analyzer (local mode) or shipped from /// the client (Flight mode), and attached to the /// [`LiquidParquetSource`](crate::LiquidParquetSource) that opens the file. -pub type ColumnSqueezeHints = HashMap>; +pub type ColumnLineages = HashMap>; /// The identity of a Parquet object within an object store. /// @@ -59,7 +59,7 @@ impl ParquetFileIdentity { } } -/// One column of a row group: (file column index, field, squeeze hint, is-predicate). +/// One column of a row group: (file column index, field, lineage expression, is-predicate). type CachedColumnSpec = (u64, Arc, Option>, bool); #[derive(Default, Debug)] @@ -268,7 +268,7 @@ pub struct CachedFile { cache_store: Arc, file_id: u64, file_schema: SchemaRef, - squeeze_hints: Arc, + lineages: Arc, } impl CachedFile { @@ -276,13 +276,13 @@ impl CachedFile { cache_store: Arc, file_id: u64, file_schema: SchemaRef, - squeeze_hints: Arc, + lineages: Arc, ) -> Self { Self { cache_store, file_id, file_schema, - squeeze_hints, + lineages, } } @@ -308,7 +308,7 @@ impl CachedFile { .enumerate() .map(|(idx, field)| { let is_predicate_column = predicate_column_ids.contains(&idx); - let expression = self.squeeze_hints.get(field.name()).cloned(); + let expression = self.lineages.get(field.name()).cloned(); ( idx as u64, Arc::clone(field), @@ -363,34 +363,34 @@ impl LiquidCacheParquet { max_disk_bytes: usize, store: t4::Store, cache_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, hydration_policy: Box, ) -> Self { - Self::new_with_squeeze_victim_concurrency( + Self::new_with_eviction_concurrency( batch_size, max_memory_bytes, max_disk_bytes, store, cache_policy, - squeeze_policy, + eviction_policy, hydration_policy, !cfg!(test), ) .await } - /// Create a new cache for parquet files with explicit victim squeeze concurrency. + /// Create a new cache for parquet files with explicit victim eviction concurrency. #[doc(hidden)] #[allow(clippy::too_many_arguments)] - pub async fn new_with_squeeze_victim_concurrency( + pub async fn new_with_eviction_concurrency( batch_size: usize, max_memory_bytes: usize, max_disk_bytes: usize, store: t4::Store, cache_policy: Box, - squeeze_policy: Box, + eviction_policy: Box, hydration_policy: Box, - squeeze_victims_concurrently: bool, + evict_victims_concurrently: bool, ) -> Self { assert!(batch_size.is_power_of_two()); let metadata = Arc::new(ParquetCacheMetadata::new()); @@ -398,12 +398,12 @@ impl LiquidCacheParquet { .with_batch_size(batch_size) .with_max_memory_bytes(max_memory_bytes) .with_max_disk_bytes(max_disk_bytes) - .with_squeeze_policy(squeeze_policy) + .with_eviction_policy(eviction_policy) .with_cache_policy(cache_policy) .with_hydration_policy(hydration_policy) .with_metadata(metadata) .with_store(store) - .with_squeeze_victims_concurrently(squeeze_victims_concurrently) + .with_evict_victims_concurrently(evict_victims_concurrently) .build() .await; @@ -423,13 +423,13 @@ impl LiquidCacheParquet { self.register_or_get_file_with_hints(file_identity, full_file_schema, Arc::default()) } - /// Register a file in the cache, attaching typed squeeze hints derived from + /// Register a file in the cache, attaching typed lineage expressions derived from /// the query plan (keyed by file-schema column name). pub fn register_or_get_file_with_hints( &self, file_identity: ParquetFileIdentity, full_file_schema: SchemaRef, - squeeze_hints: Arc, + lineages: Arc, ) -> CachedFileRef { let mut files = self.files.lock().unwrap(); let file_id = *files @@ -441,7 +441,7 @@ impl LiquidCacheParquet { self.cache_store.clone(), file_id, full_file_schema, - squeeze_hints, + lineages, )) } @@ -533,7 +533,7 @@ mod tests { use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; use datafusion::physical_plan::expressions::Column; use liquid_cache::cache::AlwaysHydrate; - use liquid_cache::cache::squeeze_policies::TranscodeSqueezeEvict; + use liquid_cache::cache::TranscodeEvict; use liquid_cache::cache_policies::LiquidPolicy; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; @@ -550,7 +550,7 @@ mod tests { usize::MAX, store, Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await; @@ -576,7 +576,7 @@ mod tests { usize::MAX, store, Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await; diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index d1fbcf41f..ac4b6e734 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -128,14 +128,12 @@ impl LiquidCacheParquet { let row_count = match cached_batch { CacheEntry::MemoryArrow(array) => Some(array.len() as u64), CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), - CacheEntry::MemorySqueezedLiquid(array) => Some(array.len() as u64), CacheEntry::DiskLiquid { .. } => None, CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count }; let cache_type = match cached_batch { CacheEntry::MemoryArrow(_) => "InMemory", CacheEntry::MemoryLiquid(_) => "LiquidMemory", - CacheEntry::MemorySqueezedLiquid(_) => "LiquidSqueezed", CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", CacheEntry::DiskArrow { .. } => "OnDiskArrow", }; @@ -173,7 +171,7 @@ mod tests { }; use bytes::Bytes; use liquid_cache::{ - cache::{AlwaysHydrate, squeeze_policies::Evict}, + cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, }; use parquet::arrow::arrow_reader::ParquetRecordBatchReader; diff --git a/src/datafusion/src/io/mod.rs b/src/datafusion/src/io/mod.rs index 0982379fc..f88a72d2d 100644 --- a/src/datafusion/src/io/mod.rs +++ b/src/datafusion/src/io/mod.rs @@ -56,14 +56,14 @@ impl ColumnExpressionTracker { } impl EntryMetadata for ParquetCacheMetadata { - fn add_squeeze_hint(&self, entry_id: &EntryID, expression: Arc) { + fn add_lineage(&self, entry_id: &EntryID, expression: Arc) { let column_path = ColumnAccessPath::from(ParquetArrayID::from(*entry_id)); let mut guard = self.expression_hints.write().unwrap(); let expression_tracker = guard.entry(column_path).or_default(); expression_tracker.record(expression.clone()); } - fn squeeze_hint(&self, entry_id: &EntryID) -> Option> { + fn lineage(&self, entry_id: &EntryID) -> Option> { let column_path = ColumnAccessPath::from(ParquetArrayID::from(*entry_id)); let guard = self.expression_hints.read().unwrap(); guard @@ -96,31 +96,31 @@ mod tests { } #[test] - fn squeeze_hint_tracks_majority() { + fn lineage_tracks_majority() { let meta = make_meta(); let e = entry(1, 2, 3); let month = Arc::new(CacheExpression::extract_date32(Date32Field::Month)); let year = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); - meta.add_squeeze_hint(&e, month.clone()); - meta.add_squeeze_hint(&e, month.clone()); - meta.add_squeeze_hint(&e, year.clone()); + meta.add_lineage(&e, month.clone()); + meta.add_lineage(&e, month.clone()); + meta.add_lineage(&e, year.clone()); - let majority = meta.squeeze_hint(&e).expect("hint"); + let majority = meta.lineage(&e).expect("hint"); assert_eq!(majority, month); } #[test] - fn squeeze_hint_prefers_recent_on_tie() { + fn lineage_prefers_recent_on_tie() { let meta = make_meta(); let e = entry(9, 9, 9); let year = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); let day = Arc::new(CacheExpression::extract_date32(Date32Field::Day)); - meta.add_squeeze_hint(&e, year.clone()); - meta.add_squeeze_hint(&e, day.clone()); + meta.add_lineage(&e, year.clone()); + meta.add_lineage(&e, day.clone()); - let majority = meta.squeeze_hint(&e).expect("hint"); + let majority = meta.lineage(&e).expect("hint"); assert_eq!(majority, day); } } diff --git a/src/datafusion/src/optimizers/squeeze_hint.rs b/src/datafusion/src/optimizers/lineage.rs similarity index 95% rename from src/datafusion/src/optimizers/squeeze_hint.rs rename to src/datafusion/src/optimizers/lineage.rs index b64cce3a3..2413c6f9c 100644 --- a/src/datafusion/src/optimizers/squeeze_hint.rs +++ b/src/datafusion/src/optimizers/lineage.rs @@ -1,11 +1,11 @@ -//! Physical-plan lineage analysis for squeeze hints. +//! Physical-plan lineage analysis for lineage expressions. //! //! This replaces the old logical [`LineageOptimizer`](super) + global //! `Arc::as_ptr` registry + field-metadata-string machinery. We analyze the //! *physical* plan directly: for every parquet scan we look at how each of its //! output columns is consumed by the operators above it (and by the scan's own //! pushed-down projection/filter), and we derive a typed [`CacheExpression`] -//! per file column describing the cheapest faithful squeeze. +//! per file column describing the safe cache expression. //! //! Working on the physical plan buys two things over the previous logical //! approach: @@ -49,7 +49,7 @@ use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMerge use liquid_cache::cache::CacheExpression; use liquid_cache::liquid_array::Date32Field; -use crate::cache::ColumnSqueezeHints; +use crate::cache::ColumnLineages; /// Stable identity of a node within a single analysis pass: the data address of /// its `Arc`. Valid only for the lifetime of one [`HintAnalyzer::analyze`] call @@ -93,7 +93,7 @@ struct ColumnStats { usages: Vec>, } -/// Analyzes a physical plan and produces, per parquet scan, the typed squeeze +/// Analyzes a physical plan and produces, per parquet scan, the typed lineage /// hint for each of its file columns. #[derive(Default)] pub(crate) struct HintAnalyzer { @@ -103,18 +103,18 @@ pub(crate) struct HintAnalyzer { scan_columns: HashMap>, } -/// Squeeze hints derived from a physical plan, keyed by the analyzed scan +/// Lineage expressions derived from a physical plan, keyed by the analyzed scan /// nodes. Valid only for the plan it was analyzed from (keyed by `Arc` identity). /// /// Local mode consumes this directly during the parquet-scan rewrite; the Flight /// client uses [`Self::for_fragment`] to extract the hints for the single scan /// inside each pushed-down fragment so they can be shipped to the cache server. -pub struct SqueezeHintMap { - per_scan: HashMap, +pub struct LineageHints { + per_scan: HashMap, } -impl SqueezeHintMap { - /// Analyze a physical plan and derive per-scan squeeze hints. +impl LineageHints { + /// Analyze a physical plan and derive per-scan lineage expressions. pub fn analyze(plan: &std::sync::Arc) -> Self { Self { per_scan: HintAnalyzer::analyze(plan), @@ -130,8 +130,8 @@ impl SqueezeHintMap { /// /// Pushed-down fragments are single-scan, so this returns that scan's hints; /// `fragment` must be a node from the same plan this map was analyzed from. - pub fn for_fragment(&self, fragment: &std::sync::Arc) -> ColumnSqueezeHints { - let mut merged = ColumnSqueezeHints::default(); + pub fn for_fragment(&self, fragment: &std::sync::Arc) -> ColumnLineages { + let mut merged = ColumnLineages::default(); fragment .apply(|node| { if let Some(hints) = self.per_scan.get(&node_ptr(node)) { @@ -147,11 +147,11 @@ impl SqueezeHintMap { } impl HintAnalyzer { - /// Analyze `plan` and return, keyed by scan node pointer, the squeeze hints + /// Analyze `plan` and return, keyed by scan node pointer, the lineage expressions /// for that scan's file columns. pub(crate) fn analyze( plan: &std::sync::Arc, - ) -> HashMap { + ) -> HashMap { let mut analyzer = HintAnalyzer::default(); let root = analyzer.visit(plan); // Columns that escape the top of the analyzed plan (returned to the @@ -328,7 +328,7 @@ impl HintAnalyzer { // A pushed-down filter consumes columns directly at the scan (with // filter pushdown enabled, `WHERE col LIKE '%x%'` lives here rather than // in a FilterExec above). Record those usages so substring searches are - // detected and columns used in other predicates are not wrongly squeezed. + // detected and columns used in other predicates are not given a partial representation. if let Some(predicate) = parquet.filter() { let usages = lineage_for_expr(&predicate, &base); self.record(&usages); @@ -396,8 +396,8 @@ impl HintAnalyzer { opaque(plan) } - fn finish(self) -> HashMap { - let mut per_scan: HashMap = HashMap::new(); + fn finish(self) -> HashMap { + let mut per_scan: HashMap = HashMap::new(); for ((scan, col), stats) in &self.stats { let Some(columns) = self.scan_columns.get(scan) else { @@ -503,7 +503,7 @@ fn propagate_other(expr: &std::sync::Arc, input: &LineageMap) combined } -/// Decide the squeeze hint for one file column from its observed op chains. +/// Decide the lineage expression for one file column from its observed op chains. fn derive_hint(data_type: &DataType, usages: &[Vec]) -> Option { if usages.is_empty() { return None; @@ -562,7 +562,7 @@ fn derive_variant(usages: &[Vec]) -> Option { None => { seen.insert(path.clone(), data_type.clone()); // A variant_get without an explicit type hint cannot be - // squeezed to a typed column; only record typed paths. + // represented as a typed column; only record typed paths. if let Some(dt) = data_type { requests.push((path.clone(), dt.clone())); } @@ -570,7 +570,7 @@ fn derive_variant(usages: &[Vec]) -> Option { } } // Raw passthrough of a variant column does not invalidate the hint: - // the squeezed representation keeps a disk backing for full reads. + // the partial representation keeps a disk backing for full reads. None => continue, _ => return None, } @@ -659,17 +659,17 @@ fn is_date_part_type(data_type: &DataType) -> bool { /// liquid-cache-backed equivalent. Returns `None` for non-parquet nodes. pub(crate) type ScanConverter<'a> = dyn FnMut( &std::sync::Arc, - ColumnSqueezeHints, + ColumnLineages, ) -> Option> + 'a; -/// Rewrite every parquet scan in `plan`, attaching the squeeze hints derived for +/// Rewrite every parquet scan in `plan`, attaching the lineage expressions derived for /// it. `hints` resolves a scan's node pointer to its hints; scans absent from -/// the map get [`ColumnSqueezeHints::default`]. +/// the map get [`ColumnLineages::default`]. pub(crate) fn rewrite_with_hints( plan: std::sync::Arc, convert: &mut ScanConverter<'_>, - hints: &HashMap, + hints: &HashMap, ) -> std::sync::Arc { plan.transform_up(|node| { let ptr = node_ptr(&node); @@ -727,7 +727,7 @@ mod tests { writer.close().unwrap(); } - async fn hints_for(sql: &str) -> ColumnSqueezeHints { + async fn hints_for(sql: &str) -> ColumnLineages { let mut config = SessionConfig::new(); // Mirror liquid cache: predicates are pushed into the parquet scan. config.options_mut().execution.parquet.pushdown_filters = true; @@ -741,7 +741,7 @@ mod tests { let df = ctx.sql(sql).await.unwrap(); let plan = df.create_physical_plan().await.unwrap(); - let map = SqueezeHintMap::analyze(&plan); + let map = LineageHints::analyze(&plan); // Single-table queries: one scan, so the merged fragment hints are it. map.for_fragment(&plan) } @@ -781,7 +781,7 @@ mod tests { #[tokio::test] async fn mixed_raw_and_extract_gets_no_hint() { - // `date` escapes raw in the projection, so it cannot be squeezed. + // `date` escapes raw in the projection, so it must remain available in full. let hints = hints_for("SELECT date, EXTRACT(YEAR FROM date) AS y FROM t").await; assert_eq!(hints.get("date"), None); } @@ -795,7 +795,7 @@ mod tests { .await; // The aggregate argument needs only YEAR, but its filter needs the - // exact date, so squeezing the column to YEAR would change the result. + // exact date, so retaining only YEAR would change the result. assert_eq!(hints.get("date"), None); } diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 3d3ae3dcb..7e30acde2 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -1,6 +1,6 @@ //! Optimizers for the Parquet module -mod squeeze_hint; +mod lineage; use std::sync::Arc; @@ -13,16 +13,16 @@ use datafusion::{ physical_plan::ExecutionPlan, }; -pub(crate) use squeeze_hint::HintAnalyzer; -pub use squeeze_hint::SqueezeHintMap; +pub(crate) use lineage::HintAnalyzer; +pub use lineage::LineageHints; -use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHints}; +use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnLineages}; /// Physical optimizer rule for local mode liquid cache. /// /// Rewrites `DataSourceExec` parquet scans to use [`LiquidParquetSource`], and -/// in the same pass derives typed squeeze hints from the full physical plan -/// (via the squeeze-hint analyzer) and attaches each scan's hints to its source. +/// in the same pass derives typed lineage expressions from the full physical plan +/// (via the lineage analyzer) and attaches each scan's hints to its source. #[derive(Debug)] pub struct LocalModeOptimizer { cache: LiquidCacheParquetRef, @@ -59,14 +59,10 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { let analysis = HintAnalyzer::analyze(&plan); let cache = self.cache.clone(); let prefetch = self.prefetch; - let mut convert = |node: &Arc, hints: ColumnSqueezeHints| { + let mut convert = |node: &Arc, hints: ColumnLineages| { convert_parquet_scan(node, &cache, hints, prefetch) }; - Ok(squeeze_hint::rewrite_with_hints( - plan, - &mut convert, - &analysis, - )) + Ok(lineage::rewrite_with_hints(plan, &mut convert, &analysis)) } fn name(&self) -> &str { @@ -87,7 +83,7 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { pub fn rewrite_data_source_plan_with_hints( plan: Arc, cache: &LiquidCacheParquetRef, - hints: &ColumnSqueezeHints, + hints: &ColumnLineages, ) -> Arc { plan.transform_up( |node| match convert_parquet_scan(&node, cache, hints.clone(), true) { @@ -103,12 +99,12 @@ pub fn rewrite_data_source_plan_with_hints( .data } -/// Rewrite the data source plan to use liquid cache (no squeeze hints). +/// Rewrite the data source plan to use liquid cache (no lineage expressions). pub fn rewrite_data_source_plan( plan: Arc, cache: &LiquidCacheParquetRef, ) -> Arc { - rewrite_data_source_plan_with_hints(plan, cache, &ColumnSqueezeHints::default()) + rewrite_data_source_plan_with_hints(plan, cache, &ColumnLineages::default()) } /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent @@ -116,7 +112,7 @@ pub fn rewrite_data_source_plan( fn convert_parquet_scan( node: &Arc, cache: &LiquidCacheParquetRef, - hints: ColumnSqueezeHints, + hints: ColumnLineages, prefetch: bool, ) -> Option> { let data_source_exec = node.downcast_ref::()?; @@ -125,7 +121,7 @@ fn convert_parquet_scan( let new_source = LiquidParquetSource::from_parquet_source(parquet_source.clone(), cache.clone()) - .with_squeeze_hints(Arc::new(hints)) + .with_lineages(Arc::new(hints)) .with_prefetch(prefetch); let mut new_config = file_scan_config.clone(); @@ -151,7 +147,7 @@ mod tests { prelude::SessionContext, }; use liquid_cache::{ - cache::{AlwaysHydrate, squeeze_policies::TranscodeSqueezeEvict}, + cache::{AlwaysHydrate, TranscodeEvict}, cache_policies::LiquidPolicy, }; use parquet::{arrow::ArrowWriter, file::properties::WriterProperties}; @@ -169,7 +165,7 @@ mod tests { usize::MAX, store, Box::new(LiquidPolicy::new()), - Box::new(TranscodeSqueezeEvict), + Box::new(TranscodeEvict), Box::new(AlwaysHydrate::new()), ) .await, diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index bc43ba4a0..6334d39f2 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -40,8 +40,8 @@ use parquet::{ use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; use crate::{ cache::{ - BatchID, ColumnSqueezeHints, InsertArrowArrayError, LiquidCacheParquetRef, - ParquetFileIdentity, PrefetchOutcome, RowGroupSnapshots, + BatchID, ColumnLineages, InsertArrowArrayError, LiquidCacheParquetRef, ParquetFileIdentity, + PrefetchOutcome, RowGroupSnapshots, }, reader::{ plantime::row_filter::build_row_filter, @@ -67,7 +67,7 @@ pub(crate) struct LiquidMorselizer { pub(crate) liquid_cache: LiquidCacheParquetRef, pub(crate) expr_adapter_factory: Arc, pub(crate) span: Option>, - pub(crate) squeeze_hints: Arc, + pub(crate) lineages: Arc, pub(crate) prefetch: bool, } @@ -162,7 +162,7 @@ impl Morselizer for LiquidMorselizer { expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), file_identity, span, - squeeze_hints: Arc::clone(&self.squeeze_hints), + lineages: Arc::clone(&self.lineages), prefetch: self.prefetch, })), })) @@ -212,7 +212,7 @@ struct PreparedLiquidOpen { expr_adapter_factory: Arc, file_identity: ParquetFileIdentity, span: Option>, - squeeze_hints: Arc, + lineages: Arc, prefetch: bool, } @@ -556,7 +556,7 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result summary.any_snapshotted = true, PrefetchOutcome::Missing => summary.any_missing = true, - PrefetchOutcome::AlreadySnapshotted | PrefetchOutcome::Squeezed => {} + PrefetchOutcome::AlreadySnapshotted => {} } } summary @@ -935,7 +935,7 @@ mod tests { }; use futures::StreamExt; use liquid_cache::{ - cache::{AlwaysHydrate, squeeze_policies::Evict}, + cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, }; use object_store::local::LocalFileSystem; @@ -1109,7 +1109,7 @@ mod tests { liquid_cache: cache.clone(), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), span: None, - squeeze_hints: Arc::default(), + lineages: Arc::default(), prefetch: true, }; let cached_file = cache.register_or_get_file( @@ -1239,7 +1239,7 @@ mod tests { liquid_cache: Arc::clone(&cache), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), span: None, - squeeze_hints: Arc::default(), + lineages: Arc::default(), prefetch: true, }; let morselizer_b = LiquidMorselizer { @@ -1257,7 +1257,7 @@ mod tests { liquid_cache: cache, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), span: None, - squeeze_hints: Arc::default(), + lineages: Arc::default(), prefetch: true, }; diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 02e3c1d6d..5f0c764e1 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -1,5 +1,5 @@ use super::LiquidMorselizer; -use crate::cache::{ColumnSqueezeHints, LiquidCacheParquetRef}; +use crate::cache::{ColumnLineages, LiquidCacheParquetRef}; use ahash::{HashMap, HashMapExt}; use bytes::Bytes; use datafusion::{ @@ -199,7 +199,7 @@ pub struct LiquidParquetSource { projection: ProjectionExprs, table_schema: TableSchema, span: Option>, - squeeze_hints: Arc, + lineages: Arc, prefetch: bool, } @@ -224,11 +224,11 @@ impl LiquidParquetSource { } } - /// Attach typed squeeze hints (keyed by file-schema column name) derived + /// Attach typed lineage expressions (keyed by file-schema column name) derived /// from the query plan. These flow to the cache when the file is opened. - pub fn with_squeeze_hints(&self, squeeze_hints: Arc) -> Self { + pub fn with_lineages(&self, lineages: Arc) -> Self { Self { - squeeze_hints, + lineages, ..self.clone() } } @@ -239,9 +239,9 @@ impl LiquidParquetSource { self } - /// The typed squeeze hints currently attached to this source. - pub fn squeeze_hints(&self) -> &Arc { - &self.squeeze_hints + /// The typed lineage expressions currently attached to this source. + pub fn lineages(&self) -> &Arc { + &self.lineages } /// Set predicate information. @@ -271,7 +271,7 @@ impl LiquidParquetSource { metrics: source.metrics().clone(), predicate: None, span: None, - squeeze_hints: Arc::default(), + lineages: Arc::default(), prefetch: true, }; @@ -334,7 +334,7 @@ impl FileSource for LiquidParquetSource { reorder_filters: self.reorder_filters(), expr_adapter_factory, span: execution_span.map(Arc::new), - squeeze_hints: Arc::clone(&self.squeeze_hints), + lineages: Arc::clone(&self.lineages), prefetch: self.prefetch, })) } diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 063e68de6..37babe4ad 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -511,7 +511,7 @@ mod tests { scalar::ScalarValue, }; use futures::{StreamExt, pin_mut}; - use liquid_cache::cache::{AlwaysHydrate, squeeze_policies::Evict}; + use liquid_cache::cache::{AlwaysHydrate, Evict}; use liquid_cache::cache_policies::LiquidPolicy; use object_store::local::LocalFileSystem; use parquet::arrow::{ From bd2a4e3fdddb457dc87e1176fcc16429f6072bd1 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Fri, 4 Sep 2026 09:54:05 -0400 Subject: [PATCH 09/24] Switch to cache vortex array (#517) I think vortex is well maintained and this switch actually gets slightly performance gain. We'll add back some liquid cache secret sauce very soon --- Cargo.lock | 690 ++++++++++- Cargo.toml | 21 +- benchmark/src/lib.rs | 6 +- dev/design/00-fsst-view.md | 224 ---- fuzz/Cargo.toml | 8 +- fuzz/README.md | 8 +- fuzz/fuzz_targets/fsst_view.rs | 117 -- fuzz/fuzz_targets/liquid_array.rs | 75 ++ src/core/Cargo.toml | 80 +- src/core/README.md | 3 +- src/core/bench/bitpacking.rs | 81 -- src/core/bench/byte_view_compare.rs | 108 -- src/core/bench/fsstarray.rs | 97 -- src/core/bench/liquid_float_array.rs | 76 -- src/core/bench/primitive_encoding.rs | 89 -- src/core/src/cache/core.rs | 45 +- src/core/src/cache/expressions.rs | 38 +- src/core/src/cache/io_context.rs | 17 +- src/core/src/cache/liquid_expr.rs | 69 +- src/core/src/cache/mod.rs | 6 +- src/core/src/cache/policies/eviction.rs | 47 +- src/core/src/cache/tests/policies.rs | 108 +- ...he__tests__policies__default_policies.snap | 13 +- ...ests__policies__insert_wont_fit_cache.snap | 9 +- src/core/src/cache/transcode.rs | 437 ------- src/core/src/cache/utils.rs | 46 +- src/core/src/liquid_array/array.rs | 819 +++++++++++++ .../byte_view_array/comparisons.rs | 457 ------- .../byte_view_array/conversions.rs | 392 ------ .../byte_view_array/fingerprint.rs | 49 - .../liquid_array/byte_view_array/helpers.rs | 147 --- .../src/liquid_array/byte_view_array/mod.rs | 340 ------ .../liquid_array/byte_view_array/operator.rs | 174 --- .../byte_view_array/serialization.rs | 326 ----- .../src/liquid_array/byte_view_array/tests.rs | 857 ------------- src/core/src/liquid_array/decimal_array.rs | 316 ----- .../src/liquid_array/fix_len_byte_array.rs | 599 --------- src/core/src/liquid_array/float_array.rs | 813 ------------- src/core/src/liquid_array/ipc.rs | 690 ----------- .../src/liquid_array/linear_integer_array.rs | 727 ----------- src/core/src/liquid_array/mod.rs | 197 +-- src/core/src/liquid_array/primitive_array.rs | 879 -------------- .../src/liquid_array/raw/bit_pack_array.rs | 584 --------- src/core/src/liquid_array/raw/fsst_buffer.rs | 1068 ----------------- src/core/src/liquid_array/raw/mod.rs | 8 - src/core/src/liquid_array/tests.rs | 199 --- src/core/src/liquid_array/utils.rs | 32 - src/core/src/utils/mod.rs | 183 --- src/core/study/cache_storage.rs | 7 +- .../study/filter_selectivity_ambiguity.rs | 321 ----- src/core/study/fsst_selectivity.rs | 244 ---- src/core/study/fsst_view.rs | 765 ------------ src/core/study/linear_integer.rs | 216 ---- src/core/study/prefix_differentiability.rs | 202 ---- src/core/study/string-fingerprint.rs | 665 ---------- src/datafusion-local/src/tests/mod.rs | 71 +- ...datafusion_local__tests__os_selection.snap | 14 +- ...afusion_local__tests__provide_schema2.snap | 12 +- ...al__tests__provide_schema_with_filter.snap | 4 +- ...usion_local__tests__referer_filtering.snap | 4 +- ...ests__single_column_filter_projection.snap | 4 +- ...on_local__tests__url_prefix_filtering.snap | 10 +- ...al__tests__url_selection_and_ordering.snap | 12 +- src/datafusion/README.md | 4 +- src/datafusion/src/cache/column.rs | 3 +- src/datafusion/src/io/mod.rs | 14 +- src/datafusion/src/optimizers/lineage.rs | 82 +- 67 files changed, 1884 insertions(+), 13144 deletions(-) delete mode 100644 dev/design/00-fsst-view.md delete mode 100644 fuzz/fuzz_targets/fsst_view.rs create mode 100644 fuzz/fuzz_targets/liquid_array.rs delete mode 100644 src/core/bench/bitpacking.rs delete mode 100644 src/core/bench/byte_view_compare.rs delete mode 100644 src/core/bench/fsstarray.rs delete mode 100644 src/core/bench/liquid_float_array.rs delete mode 100644 src/core/bench/primitive_encoding.rs delete mode 100644 src/core/src/cache/transcode.rs create mode 100644 src/core/src/liquid_array/array.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/comparisons.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/conversions.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/fingerprint.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/helpers.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/mod.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/operator.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/serialization.rs delete mode 100644 src/core/src/liquid_array/byte_view_array/tests.rs delete mode 100644 src/core/src/liquid_array/decimal_array.rs delete mode 100644 src/core/src/liquid_array/fix_len_byte_array.rs delete mode 100644 src/core/src/liquid_array/float_array.rs delete mode 100644 src/core/src/liquid_array/ipc.rs delete mode 100644 src/core/src/liquid_array/linear_integer_array.rs delete mode 100644 src/core/src/liquid_array/primitive_array.rs delete mode 100644 src/core/src/liquid_array/raw/bit_pack_array.rs delete mode 100644 src/core/src/liquid_array/raw/fsst_buffer.rs delete mode 100644 src/core/src/liquid_array/raw/mod.rs delete mode 100644 src/core/src/liquid_array/tests.rs delete mode 100644 src/core/src/liquid_array/utils.rs delete mode 100644 src/core/study/filter_selectivity_ambiguity.rs delete mode 100644 src/core/study/fsst_selectivity.rs delete mode 100644 src/core/study/fsst_view.rs delete mode 100644 src/core/study/linear_integer.rs delete mode 100644 src/core/study/prefix_differentiability.rs delete mode 100644 src/core/study/string-fingerprint.rs diff --git a/Cargo.lock b/Cargo.lock index b440b4e3f..8df28fc69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -70,6 +70,17 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "alp" +version = "0.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41ceb950ed9f82837662e3f3c470403987f32e3a25e5b741c8f17b3b0cca394d" +dependencies = [ + "fastlanes 0.6.1", + "itertools 0.15.0", + "num-traits", +] + [[package]] name = "android_system_properties" version = "0.1.6" @@ -153,6 +164,21 @@ dependencies = [ "derive_arbitrary", ] +[[package]] +name = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arcref" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28f6098a1e8ab66ff91324cce8fea6643101882cf7d09c85acdb1485ecf61e29" + [[package]] name = "arrayvec" version = "0.7.8" @@ -208,6 +234,7 @@ dependencies = [ "chrono-tz", "half", "hashbrown 0.17.1", + "libc", "num-complex", "num-integer", "num-traits", @@ -375,6 +402,7 @@ version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" dependencies = [ + "bitflags 2.13.1", "serde", "serde_core", "serde_json", @@ -435,6 +463,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -938,7 +977,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2922,6 +2961,26 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "enum-iterator" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4549325971814bda7a44061bf3fe7e487d447cba01e4220a4b454d630d7a016" +dependencies = [ + "enum-iterator-derive", +] + +[[package]] +name = "enum-iterator-derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "685adfa4d6f3d765a26bc5dbc936577de9abf756c1feeb3089b01dd395034842" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "enumset" version = "1.1.14" @@ -2976,7 +3035,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2989,6 +3048,26 @@ dependencies = [ "serde", ] +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "examples" version = "0.0.0" @@ -3008,6 +3087,35 @@ dependencies = [ "url", ] +[[package]] +name = "ext-trait" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d772df1c1a777963712fb68e014235e80863d6a91a85c4e06ba2d16243a310e5" +dependencies = [ + "ext-trait-proc_macros", +] + +[[package]] +name = "ext-trait-proc_macros" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ab7934152eaf26aa5aa9f7371408ad5af4c31357073c9e84c3b9d7f11ad639a" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "extension-traits" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a296e5a895621edf9fa8329c83aa1cb69a964643e36cf54d8d7a69b789089537" +dependencies = [ + "ext-trait", +] + [[package]] name = "fastant" version = "0.1.11" @@ -3018,6 +3126,18 @@ dependencies = [ "web-time", ] +[[package]] +name = "fastlanes" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a516528425408429de7661e7101f7e7ad54f7d3fab859a3a6de239f508d17e2" +dependencies = [ + "const_for", + "num-traits", + "pastey", + "seq-macro", +] + [[package]] name = "fastlanes" version = "0.7.0" @@ -3426,6 +3546,10 @@ name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", + "allocator-api2", +] [[package]] name = "hashbrown" @@ -3565,6 +3689,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humansize" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb51c9a029ddc91b07a787f1d86b53ccfa49b0e86688c946ebe8d3555685dd7" +dependencies = [ + "libm", +] + [[package]] name = "humantime" version = "2.4.0" @@ -3837,7 +3970,9 @@ checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" dependencies = [ "console", "once_cell", + "regex", "similar", + "strip-ansi-escapes", "tempfile", ] @@ -3875,7 +4010,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -4051,6 +4186,16 @@ version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4933f3f57a8e9d9da04db23fb153356ecaf00cbd14aee46279c33dc80925c37" +[[package]] +name = "lasso" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e14eda50a3494b3bf7b9ce51c52434a761e383d7238ce1dd5dcec2fbc13e9fb" +dependencies = [ + "dashmap", + "hashbrown 0.14.5", +] + [[package]] name = "lazy-js-bundle" version = "0.7.10" @@ -4063,6 +4208,31 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "lending-iterator" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc07588c853b50689205fb5c00498aa681d89828e0ce8cbd965ebc7a5d8ae260" +dependencies = [ + "extension-traits", + "lending-iterator-proc_macros", + "macro_rules_attribute", + "never-say-never", + "nougat", + "polonius-the-crab", +] + +[[package]] +name = "lending-iterator-proc_macros" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5445dd1c0deb1e97b8a16561d17fc686ca83e8411128fb036e9668a72d51b1d" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "lexical-core" version = "1.0.6" @@ -4209,20 +4379,14 @@ dependencies = [ "datafusion-common", "datafusion-expr-common", "datafusion-physical-expr", - "datafusion-physical-expr-common", - "divan", - "fastlanes", "fastrace", - "fsst-rs", "futures", "insta", "liquid-cache-common", "log", "mimalloc", - "num-traits", "object_store", "parquet", - "rand 0.10.2", "serde", "serde_json", "shuttle", @@ -4232,6 +4396,24 @@ dependencies = [ "tokio", "tokio-test", "tracing-subscriber", + "vortex-alp", + "vortex-array", + "vortex-arrow", + "vortex-btrblocks", + "vortex-buffer", + "vortex-bytebool", + "vortex-datetime-parts", + "vortex-decimal-byte-parts", + "vortex-error", + "vortex-fastlanes", + "vortex-fsst", + "vortex-ipc", + "vortex-mask", + "vortex-runend", + "vortex-sequence", + "vortex-session", + "vortex-sparse", + "vortex-zigzag", ] [[package]] @@ -4399,6 +4581,10 @@ version = "0.1.13" dependencies = [ "arbitrary", "arrow", + "bytes", + "datafusion-common", + "datafusion-expr-common", + "datafusion-physical-expr", "libfuzzer-sys", "liquid-cache", ] @@ -4548,6 +4734,22 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "macro_rules_attribute" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf0c9b980bf4f3a37fd7b1c066941dd1b1d0152ce6ee6e8fe8c49b9f6810d862" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d" + [[package]] name = "manganis" version = "0.7.10" @@ -4762,6 +4964,12 @@ dependencies = [ "jni-sys 0.3.1", ] +[[package]] +name = "never-say-never" +version = "6.6.666" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf5a574dadd7941adeaa71823ecba5e28331b8313fb2e1c6a5c7e5981ea53ad6" + [[package]] name = "nix" version = "0.26.4" @@ -4782,6 +4990,27 @@ dependencies = [ "memchr", ] +[[package]] +name = "nougat" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97b57b9ced431322f054fc673f1d3c7fa52d80efd9df74ad2fc759f044742510" +dependencies = [ + "macro_rules_attribute", + "nougat-proc_macros", +] + +[[package]] +name = "nougat-proc_macros" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c84f77a45e99a2f9b492695d99e1c23844619caa5f3e57647cffacad773ca257" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -4797,7 +5026,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -5098,6 +5327,12 @@ version = "4.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c45bb4a6ae1280ec0803b1ef9d3455eb50f01efbbe1447ab020f1d54fba9d8" +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "parking_lot" version = "0.12.5" @@ -5205,6 +5440,12 @@ dependencies = [ "uuid", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pastey" version = "0.2.3" @@ -5368,6 +5609,12 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" +[[package]] +name = "polonius-the-crab" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a69ee997a6282f8462abf1e0d8c38c965e968799e912b3bed8c9e8a28c2f9f" + [[package]] name = "portable-atomic" version = "1.15.0" @@ -5595,7 +5842,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -5934,7 +6181,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -6472,12 +6719,27 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "str_stack" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f446288b699d66d0fd2e30d1cfe7869194312524b3b9252594868ed26ef056a" +[[package]] +name = "strip-ansi-escapes" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8f8038e7e7969abb3f1b7c2a811225e9296da208539e0f79c5251d6cac0025" +dependencies = [ + "vte", +] + [[package]] name = "strsim" version = "0.11.1" @@ -6541,6 +6803,17 @@ dependencies = [ "symbolic-common", ] +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.119" @@ -6657,10 +6930,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -6670,7 +6943,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -7411,6 +7684,373 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "vortex-alp" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "alp", + "itertools 0.14.0", + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-array" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "arc-swap", + "arcref", + "arrow-buffer", + "async-lock", + "bytes", + "cfg-if", + "enum-iterator", + "flatbuffers", + "futures", + "half", + "humansize", + "inventory", + "itertools 0.14.0", + "jiff", + "memchr", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "pin-project-lite", + "prost", + "rand 0.10.2", + "regex", + "regex-syntax", + "rustc-hash 2.1.3", + "simdutf8", + "smallvec", + "static_assertions", + "tracing", + "uuid", + "vortex-array-macros", + "vortex-buffer", + "vortex-compute", + "vortex-error", + "vortex-flatbuffers", + "vortex-mask", + "vortex-proto", + "vortex-session", + "vortex-utils", +] + +[[package]] +name = "vortex-array-macros" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + +[[package]] +name = "vortex-arrow" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "arrow-select", + "itertools 0.14.0", + "num-traits", + "simdutf8", + "tracing", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-runend", + "vortex-session", +] + +[[package]] +name = "vortex-btrblocks" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "itertools 0.14.0", + "rand 0.10.2", + "vortex-alp", + "vortex-array", + "vortex-buffer", + "vortex-compressor", + "vortex-datetime-parts", + "vortex-decimal-byte-parts", + "vortex-error", + "vortex-fastlanes", + "vortex-fsst", + "vortex-runend", + "vortex-sequence", + "vortex-sparse", + "vortex-utils", + "vortex-zigzag", +] + +[[package]] +name = "vortex-buffer" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "allocator-api2", + "arrow-buffer", + "bitvec", + "bytes", + "itertools 0.14.0", + "simdutf8", + "vortex-error", +] + +[[package]] +name = "vortex-bytebool" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "num-traits", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-session", +] + +[[package]] +name = "vortex-compressor" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "itertools 0.14.0", + "num-traits", + "parking_lot", + "rand 0.10.2", + "rustc-hash 2.1.3", + "tracing", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-utils", +] + +[[package]] +name = "vortex-compute" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "vortex-buffer", +] + +[[package]] +name = "vortex-datetime-parts" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-decimal-byte-parts" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-error" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "arrow-schema", + "flatbuffers", + "jiff", + "prost", +] + +[[package]] +name = "vortex-fastlanes" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "fastlanes 0.7.0", + "itertools 0.14.0", + "lending-iterator", + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-flatbuffers" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "flatbuffers", + "vortex-buffer", + "vortex-error", +] + +[[package]] +name = "vortex-fsst" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "fsst-rs", + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-ipc" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "bytes", + "flatbuffers", + "futures", + "itertools 0.14.0", + "pin-project-lite", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-flatbuffers", + "vortex-session", +] + +[[package]] +name = "vortex-mask" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "itertools 0.14.0", + "vortex-buffer", + "vortex-error", +] + +[[package]] +name = "vortex-proto" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "prost", + "prost-types", +] + +[[package]] +name = "vortex-runend" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "itertools 0.14.0", + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-sequence" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "num-traits", + "prost", + "smallvec", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-proto", + "vortex-session", +] + +[[package]] +name = "vortex-session" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "arc-swap", + "lasso", + "parking_lot", + "vortex-error", + "vortex-utils", +] + +[[package]] +name = "vortex-sparse" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "itertools 0.14.0", + "num-traits", + "prost", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", +] + +[[package]] +name = "vortex-utils" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "hashbrown 0.17.1", +] + +[[package]] +name = "vortex-zigzag" +version = "0.1.0" +source = "git+https://github.com/vortex-data/vortex?rev=265b7053ac80c3ee9ddccf1fb66f85effa902c62#265b7053ac80c3ee9ddccf1fb66f85effa902c62" +dependencies = [ + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-mask", + "vortex-session", + "zigzag", +] + [[package]] name = "vstd" version = "0.0.0-2026-08-30-0159" @@ -7422,6 +8062,15 @@ dependencies = [ "verus_state_machines_macros", ] +[[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "memchr", +] + [[package]] name = "walkdir" version = "2.5.0" @@ -7597,7 +8246,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -8032,6 +8681,15 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "zigzag" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70b40401a28d86ce16a330b863b86fd7dbee4d7c940587ab09ab8c019f9e3fdf" +dependencies = [ + "num-traits", +] + [[package]] name = "zlib-rs" version = "0.6.7" diff --git a/Cargo.toml b/Cargo.toml index 4f6f0debb..0e789f42d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,7 +47,6 @@ datafusion-datasource = { version = "55.0.0" } datafusion-common = { version = "55.0.0" } datafusion-expr-common = { version = "55.0.0" } datafusion-physical-expr = { version = "55.0.0" } -datafusion-physical-expr-common = { version = "55.0.0" } datafusion-proto = { version = "55.0.0" } async-trait = "0.1.89" futures = { version = "0.3.32", default-features = false, features = ["std"] } @@ -66,8 +65,26 @@ uuid = { version = "1.23.3", features = ["v4"] } fastrace = "0.7" fastrace-tonic = "0.2" congee = "0.4.1" -insta = "1.47.2" +insta = { version = "1.47.2", features = ["filters"] } t4 = "0.1.9" +vortex-array = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-arrow = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-btrblocks = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-ipc = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-buffer = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-mask = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-session = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-error = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-fastlanes = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-fsst = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-alp = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-runend = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-sparse = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-zigzag = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-datetime-parts = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-decimal-byte-parts = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-bytebool = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } +vortex-sequence = { git = "https://github.com/vortex-data/vortex", rev = "265b7053ac80c3ee9ddccf1fb66f85effa902c62", default-features = false } [profile.dev.package] insta.opt-level = 3 diff --git a/benchmark/src/lib.rs b/benchmark/src/lib.rs index 48a1b1cf1..a72185b02 100644 --- a/benchmark/src/lib.rs +++ b/benchmark/src/lib.rs @@ -331,7 +331,11 @@ impl FromStr for BenchmarkMode { Ok(match s { "arrow" => BenchmarkMode::Arrow, "liquid" => BenchmarkMode::Liquid, - _ => return Err(format!("Invalid benchmark mode: {s}")), + _ => { + return Err(format!( + "Invalid benchmark mode: {s}, must be one of: arrow, liquid" + )); + } }) } } diff --git a/dev/design/00-fsst-view.md b/dev/design/00-fsst-view.md deleted file mode 100644 index a34cd6219..000000000 --- a/dev/design/00-fsst-view.md +++ /dev/null @@ -1,224 +0,0 @@ -# FSSTView in LiquidCache - -## Current string representation - -``` -Dictionary(u16) FSSTArray(BinaryArray) - ┌────────┐ ┌─────────────────────────────────────┐ - │ ┌────┐ │ │ Offset(i32)Nulls FSST Buffer │ - │ │ 12 │ │ │ ┌─────┐ ┌─────┐ ┌──────────────┐ │ - │ └────┘ │ │ │ │ │ │ │ │ │ - │ ┌────┐ │ │ │ │ │ │ │ │ │ - │ │ 37 │ │ │ │ │ │ │ │ │ │ - │ └────┘ │ │ │ │ │ │ │ │ │ - │ ┌────┐ │ │ │ │ │ │ │ │ │ - │ │ 42 │ │ │ │ │ │ │ │ │ │ - │ └────┘ │ │ │ │ │ │ │ │ │ - │ ┌────┐ │ │ │ │ │ │ │ │ │ - │ │ 17 │ │ │ │ │ │ │ │ │ │ - │ └────┘ │ │ │ │ │ │ │ │ │ - └────────┘ │ └─────┘ └─────┘ └──────────────┘ │ - └─────────────────────────────────────┘ -``` -## New string representation - -``` - keys - Nulls (u16) OffsetView(u64) │ FSST Buffer - ┌──┐ ┌────┐ ┌──────────────────┐ │ ┌─────────────────────┐ - │ │ │┌──┐│ │┌──────┐┌────────┐│ │ │ │ - │ │ ││12││ ││offset││Prefix ││ │ │ │ - │ │ │└──┘│ │└──────┘└────────┘│ │ │ │ - │ │ │┌──┐│ │┌──────┐┌────────┐│ │ │ │ - │ │ ││37││ ││offset││Prefix ││ │ │ │ - │ │ │└──┘│ │└──────┘└────────┘│ │ │ │ - │ │ │┌──┐│ │┌──────┐┌────────┐│ │ │ │ - │ │ ││42││ ││offset││Prefix ││ │ │ │ - │ │ │└──┘│ │└──────┘└────────┘│ │ │ │ - │ │ │┌──┐│ └──────────────────┘ │ │ │ - │ │ ││17││ ┌──────────────────┐ │ │ │ - │ │ │└──┘│ │Shared prefix │ │ │ │ - └──┘ └────┘ └──────────────────┘ │ └─────────────────────┘ - │ - In-memory│Disk -``` - -TLDR: -1. keys, offset and nulls are stored in memory. -2. FSST buffer is stored on disk. - -Design decisions: -1. The strings in the FSST buffer are unique, i.e., if two strings are different, their dictionary keys are different, and vice versa. -2. There's only one FSST buffer, this avoids the need to track buffer ids as in StringView representation in Arrow. -3. Shared prefix is the prefix that is shared across all strings in the array. -4. Offset refers to the string offset in the FSST buffer, it use Arrow's offset buffer. -5. Nulls refers to the null bit of the DictionaryView. -6. Keys are stored as u16, this is the index to the OffsetView. -7. The OffsetView has 12 bytes, with 4 bytes of offset and 8 bytes of prefix. -8. Everything but FSST buffer is stored in memory. - -For example if the array is: -- "hello" -- "hello world" -- "hello rust program" - -Then the shared prefix is "hello", the prefix of the offset view are: -- "" (empty string) -- " world" -- " rust pr" - - - -Questions: -1. should we bit-pack the offsets? -2. should we merge the null bits into DictionaryView? Maybe read this paper: https://dl.acm.org/doi/pdf/10.1145/3662010.3663452 -3. Should we extract a common prefix of the dictionary? - - -## Design notes - -### Prefix are uncompressed -The whole purpose of the prefix is to skip the symbol table thing. - -### Prefix is 6 bytes -It can be 14 bytes, but that's probably too much. - -### DictionaryView is different than StringView -Each view in StringView has 16 bytes, this is too large. - -FSSTView shrinks it to 8 bytes by: -1. remove buffer id -2. remove offset and len by storing dictionary offset in a separate field, and use dictionary index to get the offset and len. - -### Use prefix to skip decompression - -When comparing FSSTView with a string needle, we can skip the decompression by using the prefix: first check if the 6-byte prefix is enough to determine the result, if not, then decompress the string for comparison. - -Design discussion: -- Sometimes it's faster to decompress the entire array and then do the comparison. But when? - -### Use prefix to skip disk io - -Currently, each string will have a inlined 8-byte prefix, along with its offset to the compressed FSST buffer. - -Let's say if we have a needle "hello", and the prefix is: -| h | e | l | l | o | | | | - -Can we skip the disk IO by using the prefix? The answer is no, because we don't know the length of the string -- we don't know whether the stored string is "hello" or just "hello" with some zero bytes. -(this is less of a problem if we don't allow \0 in the middle of a string, but realistically, \0 is also a valid character.) - -To address this, we need to record the length of the string. Normally, this means an extra 4-byte for every string. As is done in StringView. - -Instead, we only borrow one byte from the prefix to store the length. -Nuance: how to handle long strings where the length is greater than 255? -Answer: we don't. If the length is greater or equal to 255, we simply says "we don't know", and a disk IO is required. Our study shows that 99% of the real world string have length less than 100, so we can confidently determine the length for the most of the time. - - -### FSST buffer contains full strings -Although the OffsetView contains the prefix (both shared and non-shared), the FSST buffer contains the full strings. -This allows faster conversion to arrow StringViewArray, because we don't need to prepend the prefix to the decompressed strings. -After all, we don't need FSST buffer to be short, because they are on disk anyway. - -### Efficient sort -- Sorting fsst view requires first sorting the dictionary, then use the dictionary rank to sort the keys. -- When sorting the dictionary, we should use the prefix to delay the decompression/loading from disk. -- Unlike `compare_with`, if we ever need to decompress one string from array, we simply decompress the entire array. this makes the sort simpler to implement and potentially faster (without needing to track the decompressed strings). - -## Engineering details - -### Fuzz test - -1. Use cargo-fuzz to test the FSSTView implementation. -2. We use [structured fuzzing](https://rust-fuzz.github.io/book/cargo-fuzz/structure-aware-fuzzing.html#structure-aware-fuzzing) to generate array of strings and 10 `compare_with` operations. -3. We test the following functions: -- Roundtrip from and to arrow StringArray. -- The `compare_with` function. We test that our `compare_with` function is equivalent to the Arrow's equivalent function. - -### Evict to disk - -FSSTView can be evicted to disk, but we only evict the FSST buffer and keep the OffsetView in memory, this allows most of the time to avoid decompression and IO. - -To do this, we need to change the `fsst_buffer` to be an enum, with two variants: -1. `InMemory(FsstArray)` -2. `OnDisk(PathBuf)` - -We will need to add two functions: -1. `evict_to_disk`: evict the FSST buffer to disk, and keep the OffsetView in memory. the enum will change from `InMemory` to `OnDisk`. -2. `load_from_disk`: load the FSST buffer from disk, and keep the OffsetView in memory. the enum will change from `OnDisk` to `InMemory`. - -The above two functions will need to be thread-safe, so a `std::sync::RwLock` is needed. - -When we need to read from `fsst_buffer`: -1. If it's `InMemory`, we can read from it directly. -2. If it's `OnDisk`, we read it from disk, do the work, and drop the in-memory data, i.e., **no promotion policy**. - - -## Performance evaluation - -All the benchmark below should be self-sufficient, i.e., the benchmark should be able to run without any external dependencies, without any external setup. Just cargo run and it should work. - -### Encode and decode performance - -(1) convert arrow StringViewArray to FSSTView, (2) convert arrow StringViewArray to baseline dictionary-based array. -(3) convert arrow IPC format and compress it with Snappy/Zstd/LZ4. - -Compare: 1. encode time, 2. encode size, 3. decode time (decode to arrow StringViewArray). - -Workload 1: [fineweb dataset](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu). -- the `id` column. -- the `date` column. -- the `url` column. -- the `file_path` column. - -(We don't use the `text` column because it's too large for the benchmark purpose.) -The fineweb dataset link: -- https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/blob/main/data/CC-MAIN-2025-26/000_00000.parquet -- https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/blob/main/data/CC-MAIN-2025-26/000_00001.parquet -- https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/blob/main/data/CC-MAIN-2025-26/000_00002.parquet -- https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/blob/main/data/CC-MAIN-2025-26/000_00003.parquet -- https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/blob/main/data/CC-MAIN-2025-26/000_00004.parquet - -Workload 2: [ClickBench dataset](https://datasets.clickhouse.com/hits_compatible/athena/hits.parquet). -- the `title` column. -- the `url` column. -- the `search_phrase` column. - -Setup: -- everything is in memory, no IO yet. - -Download phase: -- to load data, we simply register the parquet file to datafusion, and use sql (e.g., `SELECT url, title, search_phrase FROM parquet_file`) to read the columns we care about. -- once we read the record batch for the first time, we save it to tmp disk using arrow IPC format to avoid re-downloading the data. -- if the data is already downloaded, we simply read it from disk. - - - -### Sort performance - -This is to exercise the effectiveness of the prefix. - -Same workload as above, but we compare the performance of sorting the array. -Instead of sorting the entire column, we sort the each of the batch independently, each batch is 8192*2 rows. - -It turns out that fsst-view doesn't help sort performance, because: -- Efficient sort should first perform on the dictionary, which will not use the prefix. -- Even if for array that are all unique, we still need to decompress the array as long as there's one string that can't be resolved by the prefix. Decompressing individual strings and track+cache the decompressed strings can be slower than decompressing the entire array in one shot. -- But we probably can still use prefix to skip the comparison of the dictionary. - -### Find needle performance - -Randomly pick one string from the array, and find it across the entire column. - -This exercise both the effectiveness of the prefix, and the effectiveness of evaluating on encoded data. - -### IO performance - -We need to implement a cache abstraction, where the cache size is 1%, 10%, 30%, and 100% of the total size. - -For arrow StringViewArray and existing dictionary-based array, we stop inserting to cache when the cache is full, and write data to disk. -For FSSTView, we initially insert the entire column to cache, and when cache is full, we evict some of the previously inserted FSST buffer to disk to make room for the new data, which only keeps the OffsetView in memory. - -Then we compare the performance of the following operations: -1. Sorting the column. -2. Finding a needle in the column. -3. Convert to arrow StringViewArray. diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 1ae1a25ce..eec656163 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -12,10 +12,14 @@ libfuzzer-sys = "0.4" arbitrary = { version = "1.4", features = ["derive"] } arrow = { workspace = true } liquid-cache = { workspace = true } +bytes = { workspace = true } +datafusion-common = { workspace = true } +datafusion-expr-common = { workspace = true } +datafusion-physical-expr = { workspace = true } [[bin]] -name = "fsst_view_fuzz" -path = "fuzz_targets/fsst_view.rs" +name = "liquid_array" +path = "fuzz_targets/liquid_array.rs" test = false doc = false bench = false diff --git a/fuzz/README.md b/fuzz/README.md index 4fa668f2e..1e665921c 100644 --- a/fuzz/README.md +++ b/fuzz/README.md @@ -11,18 +11,18 @@ cargo install cargo-fuzz ## Run ```bash -cargo fuzz run fsst_view -- -jobs=12 +cargo fuzz run liquid_array -- -jobs=12 ``` ## Coverage ```bash -cargo fuzz coverage fsst_view +cargo fuzz coverage liquid_array ``` ```bash -llvm-cov show target/x86_64-unknown-linux-gnu/coverage/x86_64-unknown-linux-gnu/release/fsst_view \ - --instr-profile fuzz/coverage/fsst_view/coverage.profdata \ +llvm-cov show target/x86_64-unknown-linux-gnu/coverage/x86_64-unknown-linux-gnu/release/liquid_array \ + --instr-profile fuzz/coverage/liquid_array/coverage.profdata \ --format html \ --ignore-filename-regex "\.cargo" \ > index.html diff --git a/fuzz/fuzz_targets/fsst_view.rs b/fuzz/fuzz_targets/fsst_view.rs deleted file mode 100644 index 30efba42a..000000000 --- a/fuzz/fuzz_targets/fsst_view.rs +++ /dev/null @@ -1,117 +0,0 @@ -#![no_main] - -use arbitrary::{Arbitrary, Unstructured}; -use arrow::array::cast::AsArray; -use arrow::array::{Array, StringArray}; -use arrow::compute::kernels::cmp; -use arrow::error::ArrowError; -use libfuzzer_sys::fuzz_target; -use liquid_cache::liquid_array::LiquidByteViewArray; -use liquid_cache::liquid_array::byte_view_array::{ByteViewOperator, Comparison, Equality}; -use liquid_cache::liquid_array::raw::FsstArray; -#[derive(Debug, Clone, Arbitrary)] -struct FuzzInput { - strings: Vec>, - compare_operations: [CompareOperation; 5], -} - -#[derive(Debug, Clone, Arbitrary)] -struct CompareOperation { - needle: String, - operator: FuzzOperator, -} - -#[derive(Debug, Clone, Arbitrary)] -enum FuzzOperator { - Eq, - NotEq, - Lt, - LtEq, - Gt, - GtEq, -} - -impl FuzzOperator { - fn to_byte_view_operator(&self) -> ByteViewOperator { - match self { - FuzzOperator::Eq => ByteViewOperator::Equality(Equality::Eq), - FuzzOperator::NotEq => ByteViewOperator::Equality(Equality::NotEq), - FuzzOperator::Lt => ByteViewOperator::Comparison(Comparison::Lt), - FuzzOperator::LtEq => ByteViewOperator::Comparison(Comparison::LtEq), - FuzzOperator::Gt => ByteViewOperator::Comparison(Comparison::Gt), - FuzzOperator::GtEq => ByteViewOperator::Comparison(Comparison::GtEq), - } - } -} - -fuzz_target!(|data: &[u8]| { - let mut u = Unstructured::new(data); - let input = match FuzzInput::arbitrary(&mut u) { - Ok(input) => input, - Err(_) => return, - }; - - if input.strings.is_empty() { - return; - } - - // Test roundtrip from and to arrow StringArray - let (liquid_array, original_array) = test_roundtrip(&input.strings); - - // Test compare_with function equivalence - test_compare_with(&liquid_array, &original_array, &input.compare_operations); -}); - -fn test_roundtrip(strings: &[Option]) -> (LiquidByteViewArray, StringArray) { - let original_array = StringArray::from(strings.to_vec()); - - // Train compressor and create LiquidByteViewArray - let compressor = LiquidByteViewArray::::train_compressor(original_array.iter()); - let liquid_array = - LiquidByteViewArray::::from_string_array(&original_array, compressor); - - // Convert back to StringArray - let roundtrip_array = liquid_array.to_arrow_array(); - let roundtrip_string_array = roundtrip_array.as_string::(); - - assert_eq!(&original_array, roundtrip_string_array); - - (liquid_array, original_array) -} - -fn test_compare_with( - liquid_array: &LiquidByteViewArray, - arrow_array: &StringArray, - operations: &[CompareOperation], -) { - for op in operations { - let needle_bytes = op.needle.as_bytes(); - let operator = op.operator.to_byte_view_operator(); - - // Get expected result from Arrow operations - let arrow_result = compute_arrow_comparison(arrow_array, &op.needle, &op.operator); - - // Get result from LiquidByteViewArray - let liquid_result = liquid_array.compare_with(needle_bytes, &operator); - assert_eq!(arrow_result.unwrap(), liquid_result); - } -} - -fn compute_arrow_comparison( - array: &StringArray, - needle: &str, - operator: &FuzzOperator, -) -> Result { - let needle_array = StringArray::from(vec![needle; array.len()]); - - let result = match operator { - FuzzOperator::Eq => cmp::eq(array, &needle_array)?, - FuzzOperator::NotEq => cmp::neq(array, &needle_array)?, - FuzzOperator::Lt => cmp::lt(array, &needle_array)?, - FuzzOperator::LtEq => cmp::lt_eq(array, &needle_array)?, - FuzzOperator::Gt => cmp::gt(array, &needle_array)?, - FuzzOperator::GtEq => cmp::gt_eq(array, &needle_array)?, - }; - - Ok(result) -} diff --git a/fuzz/fuzz_targets/liquid_array.rs b/fuzz/fuzz_targets/liquid_array.rs new file mode 100644 index 000000000..d36a75deb --- /dev/null +++ b/fuzz/fuzz_targets/liquid_array.rs @@ -0,0 +1,75 @@ +#![no_main] + +use std::sync::Arc; + +use arbitrary::Arbitrary; +use arrow::array::{Array, ArrayRef, StringViewArray}; +use arrow::buffer::BooleanBuffer; +use bytes::Bytes; +use datafusion_common::ScalarValue; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, LikeExpr, Literal}; +use libfuzzer_sys::fuzz_target; +use liquid_cache::cache::LiquidExpr; +use liquid_cache::liquid_array::{LiquidArray, eval_predicate_on_array}; + +#[derive(Arbitrary, Debug)] +enum FuzzOperator { + Eq, + NotEq, + Lt, + LtEq, + Gt, + GtEq, + Like, +} + +#[derive(Arbitrary, Debug)] +struct Input { + values: Vec>, + predicates: Vec<(String, FuzzOperator)>, +} + +fuzz_target!(|input: Input| { + let arrow: ArrayRef = Arc::new(StringViewArray::from_iter( + input.values.iter().map(|value| value.as_deref()), + )); + let liquid = LiquidArray::from_arrow_array(&arrow).unwrap(); + assert_eq!(liquid.to_arrow_array().as_ref(), arrow.as_ref()); + + let selection = BooleanBuffer::new_set(arrow.len()); + for (needle, operator) in input.predicates.iter().take(8) { + let column: Arc = Arc::new(Column::new("c", 0)); + let needle = match operator { + FuzzOperator::Like => format!("%{needle}%"), + _ => needle.clone(), + }; + let literal: Arc = + Arc::new(Literal::new(ScalarValue::Utf8View(Some(needle)))); + let expr: Arc = match operator { + FuzzOperator::Like => Arc::new(LikeExpr::new(false, false, column, literal)), + operator => Arc::new(BinaryExpr::new(column, binary_operator(operator), literal)), + }; + let predicate = LiquidExpr::try_new(expr, arrow.data_type()).unwrap(); + assert_eq!( + liquid.try_eval_predicate(&predicate, &selection), + eval_predicate_on_array(arrow.clone(), &predicate) + ); + } + + let decoded = LiquidArray::from_bytes(Bytes::from(liquid.to_bytes())); + assert_eq!(decoded.to_arrow_array().as_ref(), arrow.as_ref()); +}); + +fn binary_operator(operator: &FuzzOperator) -> Operator { + match operator { + FuzzOperator::Eq => Operator::Eq, + FuzzOperator::NotEq => Operator::NotEq, + FuzzOperator::Lt => Operator::Lt, + FuzzOperator::LtEq => Operator::LtEq, + FuzzOperator::Gt => Operator::Gt, + FuzzOperator::GtEq => Operator::GtEq, + FuzzOperator::Like => unreachable!(), + } +} diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml index b33a6963f..b8e038cee 100644 --- a/src/core/Cargo.toml +++ b/src/core/Cargo.toml @@ -18,12 +18,8 @@ liquid-cache-common = { workspace = true } datafusion-common = { workspace = true } datafusion-expr-common = { workspace = true } datafusion-physical-expr = { workspace = true } -datafusion-physical-expr-common = { workspace = true } arrow = { workspace = true } arrow-schema = { workspace = true } -fastlanes = "0.7.0" -num-traits = "0.2.19" -fsst-rs = "0.6.0" ahash = { workspace = true } tempfile = { workspace = true } congee = { workspace = true } @@ -34,81 +30,41 @@ fastrace = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } sysinfo = { version = "0.39.6", default-features = false, features = ["system"] } +vortex-array = { workspace = true } +vortex-arrow = { workspace = true } +vortex-btrblocks = { workspace = true } +vortex-ipc = { workspace = true } +vortex-buffer = { workspace = true } +vortex-mask = { workspace = true } +vortex-session = { workspace = true } +vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } +vortex-fsst = { workspace = true } +vortex-alp = { workspace = true } +vortex-runend = { workspace = true } +vortex-sparse = { workspace = true } +vortex-zigzag = { workspace = true } +vortex-datetime-parts = { workspace = true } +vortex-decimal-byte-parts = { workspace = true } +vortex-bytebool = { workspace = true } +vortex-sequence = { workspace = true } [dev-dependencies] tempfile = { workspace = true } shuttle = "0.9.1" tokio-test = "0.4" tracing-subscriber = "0.3.23" -rand = "0.10.1" serde.workspace = true serde_json = { workspace = true } mimalloc = "0.1.52" clap = { version = "4.6.1", features = ["derive"] } -divan = "0.1.21" insta = { workspace = true } datafusion = { workspace = true } [features] shuttle = ["t4/shuttle"] -[[bench]] -name = "fsst_view_study" -path = "study/fsst_view.rs" -harness = false - -[[bench]] -name = "fsst_selectivity" -path = "study/fsst_selectivity.rs" -harness = false - -[[bench]] -name = "fsstarray" -path = "bench/fsstarray.rs" -harness = false - -[[bench]] -name = "byte_view_compare" -path = "bench/byte_view_compare.rs" -harness = false - -[[bench]] -name = "bitpacking" -path = "bench/bitpacking.rs" -harness = false - -[[bench]] -name = "liquid_float_array" -harness = false -path = "bench/liquid_float_array.rs" - [[bench]] name = "cache_storage" path = "study/cache_storage.rs" harness = false - -[[bench]] -name = "linear_integer_study" -path = "study/linear_integer.rs" -harness = false - - -[[bench]] -name = "primitive_encoding" -path = "bench/primitive_encoding.rs" -harness = false - -[[bench]] -name = "string_fingerprint_study" -path = "study/string-fingerprint.rs" -harness = false - -[[bench]] -name = "filter_selectivity_ambiguity" -path = "study/filter_selectivity_ambiguity.rs" -harness = false - -[[bench]] -name = "prefix_differentiability" -path = "study/prefix_differentiability.rs" -harness = false diff --git a/src/core/README.md b/src/core/README.md index 1338bc2e6..e0e752125 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -1,6 +1,6 @@ # liquid-cache -Storage layer providing byte caching and liquid array data structures. +Storage layer providing byte caching with a Vortex-backed encoded cache tier. This library provides one way to insert into the cache and three ways to read from it: - read as Arrow array @@ -88,7 +88,6 @@ let expr: Arc = Arc::new(BinaryExpr::new( let liquid_expr = liquid_cache::cache::LiquidExpr::try_new( expr, &DataType::Utf8, - Some(&liquid_cache::cache::CacheExpression::PredicateColumn), ) .unwrap(); diff --git a/src/core/bench/bitpacking.rs b/src/core/bench/bitpacking.rs deleted file mode 100644 index 9e6e9ef77..000000000 --- a/src/core/bench/bitpacking.rs +++ /dev/null @@ -1,81 +0,0 @@ -use arrow::buffer::BooleanBuffer; -use divan::Bencher; - -use std::num::NonZero; - -use arrow::array::PrimitiveArray; -use liquid_cache::liquid_array::raw::BitPackedArray; -use liquid_cache::liquid_array::{LiquidArray, LiquidPrimitiveArray}; -use rand::RngExt as _; - -const ARRAY_SIZES: [usize; 4] = [8192, 16384, 32768, 65536]; -const BIT_WIDTHS: [u8; 6] = [1, 3, 7, 11, 19, 27]; - -fn create_random_vec(array_size: usize, bit_width: u8) -> Vec { - let max_value = (1u32 << bit_width) - 1; - let mut rng = rand::rng(); - let values: Vec = (0..array_size) - .map(|_| rng.random_range(0..=max_value)) - .collect(); - values -} - -fn create_selection_array(array_size: usize, selectivity: f64) -> BooleanBuffer { - let mut rng = rand::rng(); - let values: Vec = (0..array_size) - .map(|_| rng.random::() < selectivity) - .collect(); - BooleanBuffer::from(values) -} - -#[divan::bench(args = BIT_WIDTHS, consts = ARRAY_SIZES)] -fn from_primitive_benchmark(bencher: Bencher, bit_width: u8) { - use arrow::datatypes::UInt32Type; - - let values: Vec = create_random_vec(SIZE, bit_width); - let array = PrimitiveArray::::from(values); - let bit_width = NonZero::new(bit_width).unwrap(); - - bencher - .with_inputs(|| array.clone()) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(|array| { - std::hint::black_box(BitPackedArray::from_primitive(array, bit_width)) - }); -} - -#[divan::bench(args = BIT_WIDTHS, consts = ARRAY_SIZES)] -fn to_primitive_benchmark(bencher: Bencher, bit_width: u8) { - use arrow::datatypes::UInt32Type; - - let values: Vec = create_random_vec(SIZE, bit_width); - let array = PrimitiveArray::::from(values); - let bit_width = NonZero::new(bit_width).unwrap(); - let bit_packed = BitPackedArray::from_primitive(array, bit_width); - - bencher - .with_inputs(|| bit_packed.clone()) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(|bit_packed| std::hint::black_box(bit_packed.to_primitive())); -} - -#[divan::bench(args = [(0.01, 1), (0.01, 3), (0.01, 7), (0.01, 11), (0.01, 19), (0.01, 27), (0.1, 1), (0.1, 3), (0.1, 7), (0.1, 11), (0.1, 19), (0.1, 27), (0.3, 1), (0.3, 3), (0.3, 7), (0.3, 11), (0.3, 19), (0.3, 27), (0.7, 1), (0.7, 3), (0.7, 7), (0.7, 11), (0.7, 19), (0.7, 27), (0.9, 1), (0.9, 3), (0.9, 7), (0.9, 11), (0.9, 19), (0.9, 27)], consts = ARRAY_SIZES)] -fn filter_benchmark(bencher: Bencher, (selectivity, bit_width): (f64, u8)) { - use arrow::datatypes::UInt32Type; - - let values: Vec = create_random_vec(SIZE, bit_width); - let array = PrimitiveArray::::from(values); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - let selection = create_selection_array(SIZE, selectivity); - - bencher - .with_inputs(|| (&liquid_array, &selection)) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(|(liquid_array, selection)| { - std::hint::black_box(liquid_array.filter(selection)) - }); -} - -fn main() { - divan::main(); -} diff --git a/src/core/bench/byte_view_compare.rs b/src/core/bench/byte_view_compare.rs deleted file mode 100644 index f9e608885..000000000 --- a/src/core/bench/byte_view_compare.rs +++ /dev/null @@ -1,108 +0,0 @@ -use divan::Bencher; -use rand::SeedableRng; -use rand::seq::SliceRandom; -use std::fmt::Write; -use std::sync::Arc; - -extern crate arrow; - -use arrow::array::{DictionaryArray, StringArray, UInt16Array}; -use liquid_cache::liquid_array::LiquidByteViewArray; -use liquid_cache::liquid_array::byte_view_array::{ByteViewOperator, Comparison, Equality}; -use liquid_cache::liquid_array::raw::FsstArray; - -const ROW_COUNT: usize = 10_000; -const STRING_LEN: usize = 16; -// Keep in sync with PrefixKey::prefix_len(). -const PREFIX_LEN: usize = 7; -const SUFFIX_LEN: usize = STRING_LEN - PREFIX_LEN; - -const DECIDABLE_PCTS: [u8; 6] = [0, 5, 20, 50, 80, 100]; - -const NEEDLE_PREFIX: &str = "aaaaaaa"; -const LOWER_PREFIX: &str = "0000000"; -const HIGHER_PREFIX: &str = "zzzzzzz"; -const BREAKER_PREFIX: &str = "bbbbbbb"; - -fn make_value(prefix: &str, idx: usize) -> String { - debug_assert_eq!(prefix.len(), PREFIX_LEN); - let mut value = String::with_capacity(STRING_LEN); - value.push_str(prefix); - write!(&mut value, "{:0width$}", idx, width = SUFFIX_LEN).expect("format suffix"); - debug_assert_eq!(value.len(), STRING_LEN); - value -} - -fn build_case(decidable_pct: u8) -> (LiquidByteViewArray, Vec) { - assert!(decidable_pct <= 100); - assert!(ROW_COUNT <= u16::MAX as usize); - - let decidable_rows = ROW_COUNT * decidable_pct as usize / 100; - let ambiguous_rows = ROW_COUNT - decidable_rows; - - let mut values = Vec::with_capacity(ROW_COUNT + 1); - - for idx in 0..ambiguous_rows { - values.push(make_value(NEEDLE_PREFIX, idx)); - } - - let lower_count = decidable_rows / 2; - let higher_count = decidable_rows - lower_count; - - for idx in 0..lower_count { - values.push(make_value(LOWER_PREFIX, idx)); - } - for idx in 0..higher_count { - values.push(make_value(HIGHER_PREFIX, idx)); - } - - // Keep shared_prefix empty even when all rows share the needle prefix. - values.push(make_value(BREAKER_PREFIX, 0)); - - let dict_values = Arc::new(StringArray::from(values)); - let mut keys: Vec = (0..ROW_COUNT).map(|idx| idx as u16).collect(); - let mut rng = rand::rngs::StdRng::seed_from_u64(0x9e37_79b9 ^ decidable_pct as u64); - keys.shuffle(&mut rng); - let keys = UInt16Array::from(keys); - let dict = DictionaryArray::new(keys, dict_values.clone()); - - let compressor = LiquidByteViewArray::::train_compressor(dict_values.iter()); - // Safety: dictionary values are unique, and we intentionally include an unused breaker entry. - let array = - unsafe { LiquidByteViewArray::::from_unique_dict_array(&dict, compressor) }; - - let needle = make_value(NEEDLE_PREFIX, 0).into_bytes(); - (array, needle) -} - -#[divan::bench(args = DECIDABLE_PCTS)] -fn byte_view_eq_prefix_decidable(bencher: Bencher, decidable_pct: u8) { - let (array, needle) = build_case(decidable_pct); - - bencher - .with_inputs(|| (array.clone(), needle.clone())) - .input_counter(|_| divan::counter::BytesCount::new(ROW_COUNT * STRING_LEN)) - .bench_values(|(array, needle)| { - std::hint::black_box( - array.compare_with(&needle, &ByteViewOperator::Equality(Equality::Eq)), - ) - }); -} - -#[divan::bench(args = DECIDABLE_PCTS)] -fn byte_view_lt_prefix_decidable(bencher: Bencher, decidable_pct: u8) { - let (array, needle) = build_case(decidable_pct); - - bencher - .with_inputs(|| (array.clone(), needle.clone())) - .input_counter(|_| divan::counter::BytesCount::new(ROW_COUNT * STRING_LEN)) - .bench_values(|(array, needle)| { - std::hint::black_box( - array.compare_with(&needle, &ByteViewOperator::Comparison(Comparison::Lt)), - ) - }); -} - -fn main() { - divan::main(); -} diff --git a/src/core/bench/fsstarray.rs b/src/core/bench/fsstarray.rs deleted file mode 100644 index 750816f08..000000000 --- a/src/core/bench/fsstarray.rs +++ /dev/null @@ -1,97 +0,0 @@ -use divan::Bencher; -use std::sync::Arc; - -extern crate arrow; - -use arrow::{ - array::{Array, StringArray, StringBuilder}, - datatypes::Utf8Type, -}; -use liquid_cache::liquid_array::raw::FsstArray; -use std::fs; - -const CHUNK_SIZE: [usize; 5] = [12, 32, 64, 128, 256]; - -fn create_string_arrays_from_file() -> Vec<(usize, StringArray)> { - const TEST_FILE_PATH: &str = "../../README.md"; - const LICENSE_FILE_PATH: &str = "../../LICENSE"; - - let readme = fs::read_to_string(TEST_FILE_PATH).expect("Failed to read file"); - let license = fs::read_to_string(LICENSE_FILE_PATH).expect("Failed to read file"); - let content = format!("{readme}\n\n{license}"); - - let mut result = Vec::new(); - - let chars: Vec = content.chars().collect(); - - for &chunk_size in &CHUNK_SIZE { - let mut builder = StringBuilder::new(); - for chunk in chars.chunks(chunk_size) { - let chunk_str: String = chunk.iter().collect(); - builder.append_value(chunk_str); - } - result.push((chunk_size, builder.finish())); - } - - result -} - -#[divan::bench(args = CHUNK_SIZE)] -fn compressor_benchmark(bencher: Bencher, chunk_size: usize) { - let string_arrays = create_string_arrays_from_file(); - let (_, string_array) = string_arrays - .into_iter() - .find(|(s, _)| *s == chunk_size) - .unwrap(); - let total_size = chunk_size * string_array.len(); - - bencher - .with_inputs(|| string_array.clone()) - .input_counter(move |_| divan::counter::BytesCount::new(total_size)) - .bench_values(|string_array| { - let input = string_array.iter().flat_map(|s| s.map(|a| a.as_bytes())); - FsstArray::train_compressor(input) - }); -} - -#[divan::bench(args = CHUNK_SIZE)] -fn from_byte_array_with_compressor_benchmark(bencher: Bencher, chunk_size: usize) { - let string_arrays = create_string_arrays_from_file(); - let (_, string_array) = string_arrays - .into_iter() - .find(|(s, _)| *s == chunk_size) - .unwrap(); - let compressor = - FsstArray::train_compressor(string_array.iter().flat_map(|s| s.map(|s| s.as_bytes()))); - let uncompressed_size = chunk_size * string_array.len(); - - bencher - .with_inputs(|| (string_array.clone(), Arc::new(compressor.clone()))) - .input_counter(move |_| divan::counter::BytesCount::new(uncompressed_size)) - .bench_values(|(string_array, compressor)| { - FsstArray::from_byte_array_with_compressor(&string_array, compressor) - }); -} - -#[divan::bench(args = CHUNK_SIZE)] -fn to_arrow_byte_array_benchmark(bencher: Bencher, chunk_size: usize) { - let string_arrays = create_string_arrays_from_file(); - let (_, string_array) = string_arrays - .into_iter() - .find(|(s, _)| *s == chunk_size) - .unwrap(); - let compressor = - FsstArray::train_compressor(string_array.iter().flat_map(|s| s.map(|s| s.as_bytes()))); - let fsst_values = - FsstArray::from_byte_array_with_compressor(&string_array, Arc::new(compressor)); - let total_size = chunk_size * string_array.len(); - - bencher - .with_inputs(|| fsst_values.clone()) - .input_counter(move |_| divan::counter::BytesCount::new(total_size)) - .bench_values(|fsst_values| fsst_values.to_arrow_byte_array::()); -} - -fn main() { - divan::main(); -} diff --git a/src/core/bench/liquid_float_array.rs b/src/core/bench/liquid_float_array.rs deleted file mode 100644 index 72668dff9..000000000 --- a/src/core/bench/liquid_float_array.rs +++ /dev/null @@ -1,76 +0,0 @@ -use datafusion::arrow::{ - array::PrimitiveArray, - buffer::ScalarBuffer, - datatypes::{Float32Type, Float64Type}, -}; -use divan::Bencher; -use liquid_cache::liquid_array::{LiquidArray, LiquidFloatArray}; -use rand::RngExt as _; - -const BENCH_SIZES: [usize; 3] = [8192, 16384, 24576]; - -#[divan::bench(consts = BENCH_SIZES)] -fn float32_liquid_encode(bencher: Bencher) { - bencher - .with_inputs(|| { - let mut rng = rand::rng(); - let mut array: Vec = vec![]; - for _ in 0..SIZE { - array.push(rng.random_range(-1.3e3..1.3e3)); - } - PrimitiveArray::new(ScalarBuffer::from(array), None) - }) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(LiquidFloatArray::::from_arrow_array); -} - -#[divan::bench(consts = BENCH_SIZES)] -fn float64_liquid_encode(bencher: Bencher) { - bencher - .with_inputs(|| { - let mut rng = rand::rng(); - let mut array: Vec = vec![]; - for _ in 0..SIZE { - array.push(rng.random_range(-1.3e3..1.3e3)); - } - PrimitiveArray::new(ScalarBuffer::from(array), None) - }) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(LiquidFloatArray::::from_arrow_array); -} - -#[divan::bench(consts = BENCH_SIZES)] -fn float32_liquid_decode(bencher: Bencher) { - bencher - .with_inputs(|| { - let mut rng = rand::rng(); - let mut array: Vec = vec![]; - for _ in 0..SIZE { - array.push(rng.random_range(-1.3e3..1.3e3)); - } - let arrow_array = PrimitiveArray::::new(ScalarBuffer::from(array), None); - LiquidFloatArray::::from_arrow_array(arrow_array) - }) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(|liquid_array| liquid_array.to_arrow_array()); -} - -#[divan::bench(consts = BENCH_SIZES)] -fn float64_liquid_decode(bencher: Bencher) { - bencher - .with_inputs(|| { - let mut rng = rand::rng(); - let mut array: Vec = vec![]; - for _ in 0..SIZE { - array.push(rng.random_range(-1.3e3..1.3e3)); - } - let arrow_array = PrimitiveArray::::new(ScalarBuffer::from(array), None); - LiquidFloatArray::::from_arrow_array(arrow_array) - }) - .input_counter(|_| divan::counter::BytesCount::new(SIZE * std::mem::size_of::())) - .bench_values(|liquid_array| liquid_array.to_arrow_array()); -} - -fn main() { - divan::main(); -} diff --git a/src/core/bench/primitive_encoding.rs b/src/core/bench/primitive_encoding.rs deleted file mode 100644 index 5bfdb2b36..000000000 --- a/src/core/bench/primitive_encoding.rs +++ /dev/null @@ -1,89 +0,0 @@ -use datafusion::arrow::{array::PrimitiveArray, buffer::ScalarBuffer, datatypes::Int32Type}; -use liquid_cache::liquid_array::{LiquidArray, LiquidPrimitiveArray, LiquidPrimitiveDeltaArray}; -use rand::{RngExt as _, SeedableRng, rngs::StdRng}; -use std::time::Instant; - -fn main() { - let size = 262144; - let data_size_bytes = size * std::mem::size_of::(); - let mut rng = StdRng::seed_from_u64(42); - - // Test both sequential and random data - let datasets = vec![ - ( - "Sequential", - (0..size).map(|x| x as i32).collect::>(), - ), - ( - "Random", - (0..size) - .map(|_| rng.random_range(0..1_000_000)) - .collect::>(), - ), - ]; - - for (name, data) in datasets { - println!("{} Data 1MB ({} integers)", name, size); - println!("{}", "=".repeat(50)); - - let array = PrimitiveArray::::new(ScalarBuffer::from(data.clone()), None); - - // Memory comparison - println!("Memory Consumption:"); - let regular = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); - - let regular_mem_kb = regular.get_array_memory_size() as f64 / 1024.0; - let delta_mem_kb = delta.get_array_memory_size() as f64 / 1024.0; - let memory_savings = ((regular_mem_kb - delta_mem_kb) / regular_mem_kb) * 100.0; - - println!(" Regular: {:.1} KB", regular_mem_kb); - println!(" Delta: {:.1} KB", delta_mem_kb); - println!(" Savings: {:.1}%", memory_savings); - println!(); - - // Encoding speed comparison - println!("Encoding Speed:"); - - let encode_start = Instant::now(); - let regular_encoded = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let regular_encode_time = encode_start.elapsed(); - let regular_encode_speed = - data_size_bytes as f64 / regular_encode_time.as_secs_f64() / 1_000_000.0; - - let encode_start = Instant::now(); - let delta_encoded = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); - let delta_encode_time = encode_start.elapsed(); - let delta_encode_speed = - data_size_bytes as f64 / delta_encode_time.as_secs_f64() / 1_000_000.0; - - let encode_ratio = regular_encode_speed / delta_encode_speed; - - println!(" Regular: {:.1} MB/s", regular_encode_speed); - println!(" Delta: {:.1} MB/s", delta_encode_speed); - println!(" Ratio: {:.1}x slower", encode_ratio); - println!(); - - // Decoding speed comparison - println!("Decoding Speed:"); - - let decode_start = Instant::now(); - let _ = regular_encoded.to_arrow_array(); - let regular_decode_time = decode_start.elapsed(); - let regular_decode_speed = - data_size_bytes as f64 / regular_decode_time.as_secs_f64() / 1_000_000.0; - - let decode_start = Instant::now(); - let _ = delta_encoded.to_arrow_array(); - let delta_decode_time = decode_start.elapsed(); - let delta_decode_speed = - data_size_bytes as f64 / delta_decode_time.as_secs_f64() / 1_000_000.0; - - let decode_ratio = regular_decode_speed / delta_decode_speed; - - println!(" Regular: {:.1} MB/s", regular_decode_speed); - println!(" Delta: {:.1} MB/s", delta_decode_speed); - println!(" Ratio: {:.1}x slower", decode_ratio); - println!(); - } -} diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 26859aa34..7c36dd11a 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -16,7 +16,7 @@ use super::{ utils::CacheConfig, }; use crate::cache::policies::{EvictionOutcome, EvictionPolicy}; -use crate::cache::utils::{LiquidCompressorStates, arrow_to_bytes}; +use crate::cache::utils::arrow_to_bytes; use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; use crate::cache::{CacheFull, CacheStats, EventTrace}; use crate::sync::Arc; @@ -226,11 +226,6 @@ impl LiquidCache { &self.observer } - /// Get the compressor states of the cache. - pub fn compressor_states(&self, entry_id: &EntryID) -> Arc { - self.metadata.get_compressor(entry_id) - } - /// Add a lineage expression for an entry. pub fn add_lineage(&self, entry_id: &EntryID, expression: Arc) { self.metadata.add_lineage(entry_id, expression); @@ -296,11 +291,7 @@ impl LiquidCache { ) -> Result { match &batch { batch @ CacheEntry::MemoryArrow(_) => { - let outcome = self.eviction_policy.evict( - batch, - self.metadata.get_compressor(&entry_id).as_ref(), - None, - ); + let outcome = self.eviction_policy.evict(batch, None); let EvictionOutcome::Replace { entry: new_batch, bytes_to_write, @@ -507,13 +498,11 @@ impl LiquidCache { return Ok(()); }; self.trace(InternalEvent::EvictionVictim { entry: victim }); - let compressor = self.metadata.get_compressor(&victim); - let lineage_arc = self.metadata.lineage(&victim); - let lineage = lineage_arc.as_deref(); loop { - let outcome = - self.eviction_policy - .evict(victim_entry.as_ref(), compressor.as_ref(), lineage); + let outcome = self.eviction_policy.evict( + victim_entry.as_ref(), + self.metadata.lineage(&victim).as_deref(), + ); match outcome { EvictionOutcome::Replace { @@ -744,15 +733,9 @@ impl LiquidCache { entry: *entry_id, bytes: bytes.len(), }); - let compressor_states = self.metadata.get_compressor(entry_id); - let compressor = compressor_states.fsst_compressor(); - - Some( - (crate::liquid_array::ipc::read_from_bytes( - Bytes::from(bytes), - &crate::liquid_array::ipc::LiquidIPCContext::new(compressor), - )) as _, - ) + Some(Arc::new(crate::liquid_array::LiquidArray::from_bytes( + Bytes::from(bytes), + ))) } pub(crate) async fn eval_predicate_internal( @@ -867,11 +850,7 @@ mod tests { use super::*; use crate::cache::{ CacheEntry, CachePolicy, LiquidCacheBuilder, LiquidPolicy, TranscodeEvict, - transcode_liquid_inner, - utils::{ - LiquidCompressorStates, arrow_to_bytes, create_cache_store, create_test_array, - create_test_arrow_array, - }, + utils::{arrow_to_bytes, create_cache_store, create_test_array, create_test_arrow_array}, }; use crate::sync::thread; use arrow::array::{Array, ArrayRef, Int32Array}; @@ -1110,8 +1089,8 @@ mod tests { let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; let entry_id = EntryID::from(322usize); let arrow_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); - let compressor = LiquidCompressorStates::new(); - let liquid = transcode_liquid_inner(&arrow_array, &compressor).unwrap(); + let liquid = + Arc::new(crate::liquid_array::LiquidArray::from_arrow_array(&arrow_array).unwrap()); store .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) diff --git a/src/core/src/cache/expressions.rs b/src/core/src/cache/expressions.rs index a1deb122c..44ad51dd7 100644 --- a/src/core/src/cache/expressions.rs +++ b/src/core/src/cache/expressions.rs @@ -5,7 +5,18 @@ use std::sync::Arc; use arrow_schema::DataType; -use crate::liquid_array::Date32Field; +/// A date or timestamp component observed by lineage analysis. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub enum Date32Field { + /// Year component. + Year, + /// Month component. + Month, + /// Day component. + Day, + /// Day of week, where Sunday is zero. + DayOfWeek, +} /// A typed variant path requested by a query. #[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize)] @@ -58,8 +69,6 @@ pub enum CacheExpression { }, /// A column used for predicate evaluation. PredicateColumn, - /// A column used primarily for substring search (LIKE '%foo%'). - SubstringSearch, } impl std::fmt::Display for CacheExpression { @@ -80,9 +89,6 @@ impl std::fmt::Display for CacheExpression { Self::PredicateColumn => { write!(f, "PredicateColumn") } - Self::SubstringSearch => { - write!(f, "SubstringSearch") - } } } } @@ -138,11 +144,6 @@ impl CacheExpression { } } - /// Build a substring-search expression hint. - pub fn substring_search() -> Self { - Self::SubstringSearch - } - /// Return the requested `Date32` component when this is an extract /// expression for exactly one component. /// @@ -236,7 +237,6 @@ enum CacheExprDto { Date { fields: Vec }, Variant { requests: Vec }, Predicate, - Substring, } #[derive(serde::Serialize, serde::Deserialize)] @@ -261,7 +261,6 @@ impl From<&CacheExpression> for CacheExprDto { .collect(), }, CacheExpression::PredicateColumn => CacheExprDto::Predicate, - CacheExpression::SubstringSearch => CacheExprDto::Substring, } } } @@ -284,7 +283,6 @@ impl CacheExpression { }) } CacheExprDto::Predicate => Some(CacheExpression::PredicateColumn), - CacheExprDto::Substring => Some(CacheExpression::SubstringSearch), } } } @@ -334,14 +332,10 @@ mod tests { } #[test] - fn predicate_and_substring_roundtrip() { - for expr in [ - CacheExpression::PredicateColumn, - CacheExpression::substring_search(), - ] { - let decoded = CacheExpression::from_metadata_value(&expr.to_metadata_value()).unwrap(); - assert_eq!(decoded, expr); - } + fn predicate_roundtrip() { + let expr = CacheExpression::PredicateColumn; + let decoded = CacheExpression::from_metadata_value(&expr.to_metadata_value()).unwrap(); + assert_eq!(decoded, expr); } #[test] diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index 9c5e0a179..7a7f71d3d 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -2,10 +2,7 @@ use std::fmt::Debug; use ahash::AHashMap; -use crate::cache::{ - CacheExpression, - utils::{EntryID, LiquidCompressorStates}, -}; +use crate::cache::{CacheExpression, utils::EntryID}; use crate::sync::{Arc, RwLock}; /// Per-entry metadata used by the cache. @@ -26,9 +23,6 @@ pub trait EntryMetadata: Debug + Send + Sync { fn lineage(&self, _entry_id: &EntryID) -> Option> { None } - - /// Get the compressor for an entry. - fn get_compressor(&self, entry_id: &EntryID) -> Arc; } /// Convert an [`EntryID`] to a t4 key (8-byte little-endian representation). @@ -38,11 +32,9 @@ pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { /// A default implementation of [`EntryMetadata`]. /// -/// All entries share a single [`LiquidCompressorStates`] and lineage expressions are -/// stored in a flat map keyed by [`EntryID`]. +/// Lineage expressions are stored in a flat map keyed by [`EntryID`]. #[derive(Debug, Default)] pub struct DefaultCacheMetadata { - compressor_state: Arc, lineages: RwLock>>, } @@ -50,7 +42,6 @@ impl DefaultCacheMetadata { /// Create a new instance of [`DefaultCacheMetadata`]. pub fn new() -> Self { Self { - compressor_state: Arc::new(LiquidCompressorStates::new()), lineages: RwLock::new(AHashMap::new()), } } @@ -66,8 +57,4 @@ impl EntryMetadata for DefaultCacheMetadata { let guard = self.lineages.read().unwrap(); guard.get(entry_id).cloned() } - - fn get_compressor(&self, _entry_id: &EntryID) -> Arc { - self.compressor_state.clone() - } } diff --git a/src/core/src/cache/liquid_expr.rs b/src/core/src/cache/liquid_expr.rs index a317181c6..84b72d2fd 100644 --- a/src/core/src/cache/liquid_expr.rs +++ b/src/core/src/cache/liquid_expr.rs @@ -6,7 +6,6 @@ use datafusion_physical_expr::expressions::{ }; use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; -use crate::cache::CacheExpression; use crate::sync::Arc; use crate::utils::get_bytes_needle; @@ -28,14 +27,10 @@ impl LiquidExpr { /// Validate and wrap a physical expression for LiquidCache predicate evaluation. /// /// Returns `None` when the expression shape or operator is unsupported for the - /// provided column type and expression hint. - pub fn try_new( - expr: Arc, - data_type: &DataType, - expression_hint: Option<&CacheExpression>, - ) -> Option { + /// provided column type. + pub fn try_new(expr: Arc, data_type: &DataType) -> Option { let normalized = unwrap_dynamic_filter(&expr)?; - if supports_expr(&normalized, data_type, expression_hint) { + if supports_expr(&normalized, data_type) { Some(Self { expr: normalized }) } else { None @@ -46,11 +41,6 @@ impl LiquidExpr { pub fn physical_expr(&self) -> &Arc { &self.expr } - - #[cfg(test)] - pub(crate) fn new_unchecked(expr: Arc) -> Self { - Self { expr } - } } fn unwrap_dynamic_filter(expr: &Arc) -> Option> { @@ -61,17 +51,13 @@ fn unwrap_dynamic_filter(expr: &Arc) -> Option, - data_type: &DataType, - expression_hint: Option<&CacheExpression>, -) -> bool { +fn supports_expr(expr: &Arc, data_type: &DataType) -> bool { if let Some(binary) = expr.downcast_ref::() { - return supports_binary_expr(binary, data_type, expression_hint); + return supports_binary_expr(binary, data_type); } if let Some(like_expr) = expr.downcast_ref::() { - return supports_like_expr(like_expr, data_type, expression_hint); + return supports_like_expr(like_expr, data_type); } if let Some(literal) = expr.downcast_ref::() { @@ -81,11 +67,7 @@ fn supports_expr( false } -fn supports_binary_expr( - binary: &BinaryExpr, - data_type: &DataType, - expression_hint: Option<&CacheExpression>, -) -> bool { +fn supports_binary_expr(binary: &BinaryExpr, data_type: &DataType) -> bool { let Some(literal) = binary.right().downcast_ref::() else { return false; }; @@ -104,7 +86,6 @@ fn supports_binary_expr( | Operator::GtEq => get_bytes_needle(literal.value()).is_some(), Operator::LikeMatch | Operator::NotLikeMatch => { get_bytes_needle(literal.value()).is_some() - && is_substring_hint_enabled(expression_hint) } _ => false, } @@ -123,15 +104,11 @@ fn supports_binary_expr( } } -fn supports_like_expr( - like_expr: &LikeExpr, - data_type: &DataType, - expression_hint: Option<&CacheExpression>, -) -> bool { +fn supports_like_expr(like_expr: &LikeExpr, data_type: &DataType) -> bool { if !is_byte_like(data_type) || like_expr.case_insensitive() { return false; } - if !is_column_like(like_expr.expr()) || !is_substring_hint_enabled(expression_hint) { + if !is_column_like(like_expr.expr()) { return false; } like_expr @@ -141,10 +118,6 @@ fn supports_like_expr( .is_some() } -fn is_substring_hint_enabled(expression_hint: Option<&CacheExpression>) -> bool { - matches!(expression_hint, Some(CacheExpression::SubstringSearch)) -} - fn is_column_like(expr: &Arc) -> bool { if expr.downcast_ref::().is_some() { return true; @@ -209,35 +182,19 @@ mod tests { Operator::Eq, Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), )); - let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8, None); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8); assert!(liquid_expr.is_some()); } #[test] - fn rejects_byte_like_without_substring_hint() { - let expr: Arc = Arc::new(LikeExpr::new( - false, - false, - Arc::new(Column::new("c", 0)), - Arc::new(Literal::new(ScalarValue::Utf8(Some("%abc%".to_string())))), - )); - let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8, None); - assert!(liquid_expr.is_none()); - } - - #[test] - fn accepts_byte_like_with_substring_hint() { + fn accepts_byte_like_like() { let expr: Arc = Arc::new(LikeExpr::new( false, false, Arc::new(Column::new("c", 0)), Arc::new(Literal::new(ScalarValue::Utf8(Some("%abc%".to_string())))), )); - let liquid_expr = LiquidExpr::try_new( - expr, - &DataType::Utf8, - Some(&CacheExpression::SubstringSearch), - ); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8); assert!(liquid_expr.is_some()); } @@ -248,7 +205,7 @@ mod tests { Operator::Gt, Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), )); - let liquid_expr = LiquidExpr::try_new(expr, &DataType::Int32, None); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Int32); assert!(liquid_expr.is_some()); } } diff --git a/src/core/src/cache/mod.rs b/src/core/src/cache/mod.rs index 4783322c6..14c20f5e5 100644 --- a/src/core/src/cache/mod.rs +++ b/src/core/src/cache/mod.rs @@ -10,13 +10,12 @@ mod io_context; mod liquid_expr; mod observer; pub mod policies; -mod transcode; mod utils; pub use builders::{EvaluatePredicate, Get, Insert, LiquidCacheBuilder, default_max_memory_bytes}; pub use cached_batch::{CacheEntry, CachedBatchType}; pub use core::{LiquidCache, PrefetchResult}; -pub use expressions::{CacheExpression, VariantRequest}; +pub use expressions::{CacheExpression, Date32Field, VariantRequest}; pub use io_context::{DefaultCacheMetadata, EntryMetadata}; pub use liquid_expr::LiquidExpr; pub use observer::EventTrace; @@ -26,8 +25,7 @@ pub use policies::{ AlwaysHydrate, CachePolicy, Evict, EvictionPolicy, HydrationPolicy, HydrationRequest, LiquidPolicy, MaterializedEntry, NoHydration, TranscodeEvict, }; -pub use transcode::{transcode_liquid_inner, transcode_liquid_inner_with_hint}; -pub use utils::{EntryID, LiquidCompressorStates}; +pub use utils::EntryID; /// The cache could not reserve enough disk budget for a write. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/core/src/cache/policies/eviction.rs b/src/core/src/cache/policies/eviction.rs index b81e305cc..d64191e01 100644 --- a/src/core/src/cache/policies/eviction.rs +++ b/src/core/src/cache/policies/eviction.rs @@ -1,12 +1,12 @@ //! Policies for moving cache entries to cheaper storage under memory pressure. +use std::sync::Arc; + use arrow::array::Array; use bytes::Bytes; -use crate::cache::{ - CacheExpression, LiquidCompressorStates, cached_batch::CacheEntry, - transcode_liquid_inner_with_hint, utils::arrow_to_bytes, -}; +use crate::cache::{CacheExpression, cached_batch::CacheEntry, utils::arrow_to_bytes}; +use crate::liquid_array::LiquidArray; /// The next storage representation selected for a cache entry. #[derive(Debug, Clone)] @@ -24,13 +24,8 @@ pub enum EvictionOutcome { /// Chooses the next representation for an entry under memory pressure. pub trait EvictionPolicy: std::fmt::Debug + Send + Sync { - /// Move the entry one step toward cheaper storage. - fn evict( - &self, - entry: &CacheEntry, - compressor: &LiquidCompressorStates, - expression: Option<&CacheExpression>, - ) -> EvictionOutcome; + /// Move the entry one step toward cheaper storage; lineage is available for encoding decisions. + fn evict(&self, entry: &CacheEntry, lineage: Option<&CacheExpression>) -> EvictionOutcome; } /// Evict memory entries directly to disk. @@ -38,12 +33,7 @@ pub trait EvictionPolicy: std::fmt::Debug + Send + Sync { pub struct Evict; impl EvictionPolicy for Evict { - fn evict( - &self, - entry: &CacheEntry, - _compressor: &LiquidCompressorStates, - _expression: Option<&CacheExpression>, - ) -> EvictionOutcome { + fn evict(&self, entry: &CacheEntry, _lineage: Option<&CacheExpression>) -> EvictionOutcome { persist(entry) } } @@ -53,22 +43,15 @@ impl EvictionPolicy for Evict { pub struct TranscodeEvict; impl EvictionPolicy for TranscodeEvict { - fn evict( - &self, - entry: &CacheEntry, - compressor: &LiquidCompressorStates, - expression: Option<&CacheExpression>, - ) -> EvictionOutcome { + fn evict(&self, entry: &CacheEntry, _lineage: Option<&CacheExpression>) -> EvictionOutcome { match entry { - CacheEntry::MemoryArrow(array) => { - match transcode_liquid_inner_with_hint(array, compressor, expression) { - Ok(liquid) => EvictionOutcome::Replace { - entry: CacheEntry::memory_liquid(liquid), - bytes_to_write: None, - }, - Err(_) => persist(entry), - } - } + CacheEntry::MemoryArrow(array) => match LiquidArray::from_arrow_array(array) { + Ok(liquid) => EvictionOutcome::Replace { + entry: CacheEntry::memory_liquid(Arc::new(liquid)), + bytes_to_write: None, + }, + Err(_) => persist(entry), + }, _ => persist(entry), } } diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index 1e26f45a9..ee3681a45 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -1,7 +1,15 @@ use crate::cache::{ - AlwaysHydrate, EntryID, LiquidCacheBuilder, LiquidPolicy, TranscodeEvict, - utils::create_test_arrow_array, + AlwaysHydrate, CachedBatchType, EntryID, LiquidCacheBuilder, LiquidExpr, LiquidPolicy, + TranscodeEvict, utils::create_test_arrow_array, }; +use arrow::array::{Array, ArrayRef, Int64Array, StringViewArray}; +use arrow::buffer::BooleanBuffer; +use arrow_schema::DataType; +use datafusion_common::ScalarValue; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; +use std::sync::Arc; #[tokio::test] async fn default_policies() { @@ -28,7 +36,9 @@ async fn default_policies() { } let trace = cache.consume_event_trace(); - insta::assert_snapshot!(trace); + insta::with_settings!({ filters => vec![(r"bytes=\d+", "bytes=[bytes]")] }, { + insta::assert_snapshot!(trace); + }); } #[tokio::test] @@ -59,5 +69,95 @@ async fn insert_wont_fit_cache() { let trace = cache.consume_event_trace(); let json_trace = serde_json::to_string(&trace).unwrap(); println!("{}", json_trace); - insta::assert_snapshot!(trace); + insta::with_settings!({ filters => vec![(r"bytes=\d+", "bytes=[bytes]")] }, { + insta::assert_snapshot!(trace); + }); +} + +#[tokio::test] +async fn liquid_eviction_reads_memory_and_disk() { + let integers: ArrayRef = Arc::new(Int64Array::from_iter( + (0..2_048).map(|value| (value % 13 != 0).then_some(value)), + )); + let strings: ArrayRef = + Arc::new(StringViewArray::from_iter((0..2_048).map(|value| { + (value % 11 != 0).then(|| format!("value-{value}")) + }))); + let capacity = integers.get_array_memory_size() + strings.get_array_memory_size(); + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(capacity) + .with_eviction_policy(Box::new(TranscodeEvict)) + .with_hydration_policy(Box::new(AlwaysHydrate::new())) + .build() + .await; + + for (id, array) in [ + integers.clone(), + strings.clone(), + integers.clone(), + strings.clone(), + ] + .into_iter() + .enumerate() + { + cache.insert(EntryID::from(id), array).await.unwrap(); + } + + let mut states = Vec::new(); + cache.for_each_entry(|_, entry| states.push(CachedBatchType::from(entry))); + assert!( + states + .iter() + .any(|state| *state != CachedBatchType::MemoryArrow) + ); + + let id = EntryID::from(0); + assert_eq!( + cache.get(&id).read().await.unwrap().as_ref(), + integers.as_ref() + ); + + let selection = BooleanBuffer::from_iter((0..integers.len()).map(|index| index % 3 == 0)); + let selected = cache + .get(&id) + .with_selection(&selection) + .read() + .await + .unwrap(); + let expected = arrow::compute::filter( + &integers, + &arrow::array::BooleanArray::new(selection.clone(), None), + ) + .unwrap(); + assert_eq!(selected.as_ref(), expected.as_ref()); + + let physical: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 0)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int64(Some(1_000)))), + )); + let predicate = LiquidExpr::try_new(physical, &DataType::Int64).unwrap(); + let actual = cache.eval_predicate(&id, &predicate).read().await.unwrap(); + let expected = arrow::array::BooleanArray::from_iter( + integers + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|value| value.map(|value| value >= 1_000)), + ); + assert_eq!(actual, expected); + + cache.flush_all_to_disk().await.unwrap(); + let mut disk_states = Vec::new(); + cache.for_each_entry(|_, entry| disk_states.push(CachedBatchType::from(entry))); + assert!(disk_states.iter().all(|state| matches!( + state, + CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow + ))); + assert!(disk_states.contains(&CachedBatchType::DiskLiquid)); + assert_eq!( + cache.get(&id).read().await.unwrap().as_ref(), + integers.as_ref() + ); } diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap index f4dc0a56a..fc702d26f 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap @@ -1,6 +1,5 @@ --- source: src/core/src/cache/tests/policies.rs -assertion_line: 31 expression: trace --- EventTrace: [ @@ -18,10 +17,10 @@ event=eviction_begin victims=[2,0,1] event=eviction_victim entry=2 event=insert_success entry=2 kind=MemoryLiquid event=eviction_victim entry=0 -event=io_write entry=0 kind=DiskLiquid bytes=1320 +event=io_write entry=0 kind=DiskLiquid bytes=[bytes] event=insert_success entry=0 kind=DiskLiquid event=eviction_victim entry=1 -event=io_write entry=1 kind=DiskLiquid bytes=1320 +event=io_write entry=1 kind=DiskLiquid bytes=[bytes] event=insert_success entry=1 kind=DiskLiquid event=insert_success entry=3 kind=MemoryArrow event=insert_failed entry=4 kind=MemoryArrow @@ -29,19 +28,19 @@ event=eviction_begin victims=[3,2] event=eviction_victim entry=3 event=insert_success entry=3 kind=MemoryLiquid event=eviction_victim entry=2 -event=io_write entry=2 kind=DiskLiquid bytes=1320 +event=io_write entry=2 kind=DiskLiquid bytes=[bytes] event=insert_success entry=2 kind=DiskLiquid event=insert_success entry=4 kind=MemoryArrow event=read entry=0 selection=false expr=None cached=DiskLiquid -event=io_read_liquid entry=0 bytes=1320 +event=io_read_liquid entry=0 bytes=[bytes] event=hydrate entry=0 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=0 kind=MemoryLiquid event=read entry=1 selection=false expr=None cached=DiskLiquid -event=io_read_liquid entry=1 bytes=1320 +event=io_read_liquid entry=1 bytes=[bytes] event=hydrate entry=1 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=1 kind=MemoryLiquid event=read entry=2 selection=false expr=None cached=DiskLiquid -event=io_read_liquid entry=2 bytes=1320 +event=io_read_liquid entry=2 bytes=[bytes] event=hydrate entry=2 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=2 kind=MemoryLiquid event=read entry=3 selection=false expr=None cached=MemoryLiquid diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap index 284f51c2b..14eedaf85 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__insert_wont_fit_cache.snap @@ -1,6 +1,5 @@ --- source: src/core/src/cache/tests/policies.rs -assertion_line: 62 expression: trace --- EventTrace: [ @@ -12,16 +11,16 @@ event=insert_success entry=0 kind=MemoryLiquid event=insert_failed entry=1 kind=MemoryArrow event=eviction_begin victims=[0] event=eviction_victim entry=0 -event=io_write entry=0 kind=DiskLiquid bytes=1320 +event=io_write entry=0 kind=DiskLiquid bytes=[bytes] event=insert_success entry=0 kind=DiskLiquid event=insert_failed entry=1 kind=MemoryArrow event=insert_failed entry=1 kind=MemoryLiquid -event=io_write entry=1 kind=MemoryLiquid bytes=34600 +event=io_write entry=1 kind=MemoryLiquid bytes=[bytes] event=insert_success entry=1 kind=DiskLiquid event=read entry=1 selection=false expr=None cached=DiskLiquid -event=io_read_liquid entry=1 bytes=34600 +event=io_read_liquid entry=1 bytes=[bytes] event=hydrate entry=1 cached=DiskLiquid new=MemoryLiquid event=insert_failed entry=1 kind=MemoryLiquid -event=io_write entry=1 kind=MemoryLiquid bytes=34600 +event=io_write entry=1 kind=MemoryLiquid bytes=[bytes] event=insert_success entry=1 kind=DiskLiquid ] diff --git a/src/core/src/cache/transcode.rs b/src/core/src/cache/transcode.rs deleted file mode 100644 index 6e862954c..000000000 --- a/src/core/src/cache/transcode.rs +++ /dev/null @@ -1,437 +0,0 @@ -use std::sync::Arc; - -use arrow::array::types::*; -use arrow::array::{ArrayRef, AsArray}; -use arrow_schema::{DataType, TimeUnit}; - -use crate::liquid_array::byte_view_array::ByteViewBuildOptions; -use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::{ - LiquidArrayRef, LiquidByteViewArray, LiquidDecimalArray, LiquidFixedLenByteArray, - LiquidFloatArray, LiquidPrimitiveArray, -}; - -use super::{CacheExpression, utils::LiquidCompressorStates}; - -fn with_fsst_compressor_or_train( - state: &LiquidCompressorStates, - use_compressor: impl FnOnce(Arc) -> LiquidArrayRef, - train: impl FnOnce() -> (Arc, LiquidArrayRef), -) -> LiquidArrayRef { - if let Some(compressor) = state.fsst_compressor() { - return use_compressor(compressor); - } - - let mut compressors = state.fsst_compressor_raw().write().unwrap(); - if let Some(compressor) = compressors.as_ref() { - return use_compressor(compressor.clone()); - } - - let (compressor, liquid_array) = train(); - *compressors = Some(compressor); - liquid_array -} - -/// This method is used to transcode an arrow array into a liquid array. -/// -/// Returns the transcoded liquid array if successful, otherwise returns the original arrow array. -pub fn transcode_liquid_inner<'a>( - array: &'a ArrayRef, - state: &LiquidCompressorStates, -) -> Result { - transcode_liquid_inner_with_hint(array, state, None) -} - -/// Transcode with an optional hint to precompute metadata (e.g., substring fingerprints). -pub fn transcode_liquid_inner_with_hint<'a>( - array: &'a ArrayRef, - state: &LiquidCompressorStates, - lineage: Option<&CacheExpression>, -) -> Result { - let data_type = array.data_type(); - if data_type.is_primitive() { - // For primitive types, perform the transcoding. - let liquid_array: LiquidArrayRef = match data_type { - DataType::Int8 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Int16 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Int32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Int64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::UInt8 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::UInt16 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::UInt32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::UInt64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Date32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Date64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Timestamp(TimeUnit::Second, None) => Arc::new(LiquidPrimitiveArray::< - TimestampSecondType, - >::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Timestamp(TimeUnit::Millisecond, None) => Arc::new(LiquidPrimitiveArray::< - TimestampMillisecondType, - >::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Timestamp(TimeUnit::Microsecond, None) => Arc::new(LiquidPrimitiveArray::< - TimestampMicrosecondType, - >::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Timestamp(TimeUnit::Nanosecond, None) => Arc::new(LiquidPrimitiveArray::< - TimestampNanosecondType, - >::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Timestamp(_, Some(_)) => { - log::warn!("unsupported timestamp type with timezone {data_type:?}"); - return Err(array); - } - DataType::Float32 => Arc::new(LiquidFloatArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Float64 => Arc::new(LiquidFloatArray::::from_arrow_array( - array.as_primitive::().clone(), - )), - DataType::Decimal128(_, _) => { - let decimals = array.as_primitive::(); - if LiquidDecimalArray::fits_u64(decimals) { - return Ok(Arc::new(LiquidDecimalArray::from_decimal_array(decimals))); - } - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidFixedLenByteArray::from_decimal_array( - decimals, compressor, - )) - }, - || { - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(decimals); - (compressor, Arc::new(liquid_array)) - }, - ); - return Ok(liquid_array); - } - DataType::Decimal256(_, _) => { - let decimals = array.as_primitive::(); - if LiquidDecimalArray::fits_u64(decimals) { - return Ok(Arc::new(LiquidDecimalArray::from_decimal_array(decimals))); - } - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidFixedLenByteArray::from_decimal_array( - decimals, compressor, - )) - }, - || { - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(decimals); - (compressor, Arc::new(liquid_array)) - }, - ); - return Ok(liquid_array); - } - _ => { - // For unsupported primitive types, leave the value unchanged. - log::warn!("unsupported primitive type {data_type:?}"); - return Err(array); - } - }; - return Ok(liquid_array); - } - - // Handle string/dictionary types. - let build_fingerprints = matches!(lineage, Some(CacheExpression::SubstringSearch)); - match array.data_type() { - DataType::Utf8View => { - let options = - ByteViewBuildOptions::for_data_type(array.data_type(), build_fingerprints); - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidByteViewArray::::from_view_array_inner( - array.as_string_view(), - compressor, - options, - )) - }, - || { - let (compressor, compressed) = - LiquidByteViewArray::::train_from_string_view_inner( - array.as_string_view(), - options, - ); - (compressor, Arc::new(compressed)) - }, - ); - Ok(liquid_array) - } - DataType::BinaryView => { - let options = - ByteViewBuildOptions::for_data_type(array.data_type(), build_fingerprints); - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidByteViewArray::::from_view_array_inner( - array.as_binary_view(), - compressor, - options, - )) - }, - || { - let (compressor, compressed) = - LiquidByteViewArray::::train_from_binary_view_inner( - array.as_binary_view(), - options, - ); - (compressor, Arc::new(compressed)) - }, - ); - Ok(liquid_array) - } - DataType::Utf8 => { - let options = - ByteViewBuildOptions::for_data_type(array.data_type(), build_fingerprints); - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidByteViewArray::::from_byte_array_inner( - array.as_string::(), - compressor, - options, - )) - }, - || { - let (compressor, compressed) = - LiquidByteViewArray::::train_from_arrow_inner( - array.as_string::(), - options, - ); - (compressor, Arc::new(compressed)) - }, - ); - Ok(liquid_array) - } - DataType::Binary => { - let options = - ByteViewBuildOptions::for_data_type(array.data_type(), build_fingerprints); - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| { - Arc::new(LiquidByteViewArray::::from_byte_array_inner( - array.as_binary::(), - compressor, - options, - )) - }, - || { - let (compressor, compressed) = - LiquidByteViewArray::::train_from_arrow_inner( - array.as_binary::(), - options, - ); - (compressor, Arc::new(compressed)) - }, - ); - Ok(liquid_array) - } - DataType::Dictionary(_, _) => { - if let Some(dict_array) = array.as_dictionary_opt::() { - let options = - ByteViewBuildOptions::for_data_type(array.data_type(), build_fingerprints); - let liquid_array = with_fsst_compressor_or_train( - state, - |compressor| unsafe { - Arc::new( - LiquidByteViewArray::::from_unique_dict_array_with_options( - dict_array, compressor, options, - ), - ) - }, - || { - let (compressor, liquid_array) = - LiquidByteViewArray::::train_from_arrow_dict_inner( - dict_array, options, - ); - (compressor, Arc::new(liquid_array)) - }, - ); - return Ok(liquid_array); - } - log::warn!("unsupported data type {:?}", array.data_type()); - Err(array) - } - _ => { - log::debug!("unsupported data type {:?}", array.data_type()); - Err(array) - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{ - ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, DictionaryArray, Float32Array, - Float64Array, Int32Array, Int64Array, StringArray, TimestampMicrosecondArray, UInt16Array, - }; - use arrow::datatypes::UInt16Type; - - const TEST_ARRAY_SIZE: usize = 8192; - - fn assert_transcode(original: &ArrayRef, transcoded: &LiquidArrayRef) { - assert!( - transcoded.get_array_memory_size() < original.get_array_memory_size(), - "transcoded size: {}, original size: {}", - transcoded.get_array_memory_size(), - original.get_array_memory_size() - ); - let back_to_arrow = transcoded.to_arrow_array(); - assert_eq!(original, &back_to_arrow); - } - - #[test] - fn test_transcode_int32() { - let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..TEST_ARRAY_SIZE as i32)); - let state = LiquidCompressorStates::new(); - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_int64() { - let array: ArrayRef = Arc::new(Int64Array::from_iter_values(0..TEST_ARRAY_SIZE as i64)); - let state = LiquidCompressorStates::new(); - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_float32() { - let array: ArrayRef = Arc::new(Float32Array::from_iter_values( - (0..TEST_ARRAY_SIZE).map(|i| i as f32), - )); - let state = LiquidCompressorStates::new(); - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_float64() { - let array: ArrayRef = Arc::new(Float64Array::from_iter_values( - (0..TEST_ARRAY_SIZE).map(|i| i as f64), - )); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_timestamp_microsecond() { - let array: ArrayRef = Arc::new(TimestampMicrosecondArray::from_iter_values( - (0..TEST_ARRAY_SIZE).map(|i| (i as i64) * 1_000), - )); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_string() { - let array: ArrayRef = Arc::new(StringArray::from_iter_values( - (0..TEST_ARRAY_SIZE).map(|i| format!("test_string_{i}")), - )); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_binary_view() { - let array: ArrayRef = Arc::new(BinaryViewArray::from_iter_values( - (0..TEST_ARRAY_SIZE).map(|i| format!("test_binary_{i}").into_bytes()), - )); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_dictionary_uft8() { - // Create a dictionary with many repeated values - let values = - StringArray::from_iter_values((0..100).map(|i| format!("value__longer_values_{i}"))); - let keys: Vec = (0..TEST_ARRAY_SIZE).map(|i| (i % 100) as u16).collect(); - - let dict_array = - DictionaryArray::::try_new(UInt16Array::from(keys), Arc::new(values)) - .unwrap(); - - let array: ArrayRef = Arc::new(dict_array); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_dictionary_binary() { - // Create a dictionary with binary values and many repeated values - let values = BinaryArray::from_iter_values( - (0..100).map(|i| format!("binary_value_{i}").into_bytes()), - ); - let keys: Vec = (0..TEST_ARRAY_SIZE).map(|i| (i % 100) as u16).collect(); - - let dict_array = - DictionaryArray::::try_new(UInt16Array::from(keys), Arc::new(values)) - .unwrap(); - - let array: ArrayRef = Arc::new(dict_array); - let state = LiquidCompressorStates::new(); - - let transcoded = transcode_liquid_inner(&array, &state).unwrap(); - assert_transcode(&array, &transcoded); - } - - #[test] - fn test_transcode_unsupported_type() { - // Create a boolean array which is not supported by the transcoder - let values: Vec = (0..TEST_ARRAY_SIZE).map(|i| i.is_multiple_of(2)).collect(); - let array: ArrayRef = Arc::new(BooleanArray::from(values)); - let state = LiquidCompressorStates::new(); - - // Try to transcode and expect an error - let result = transcode_liquid_inner(&array, &state); - - // Verify it returns Err with the original array - assert!(result.is_err()); - if let Err(original) = result { - assert_eq!(&array, original); - } - } -} diff --git a/src/core/src/cache/utils.rs b/src/core/src/cache/utils.rs index 4efd4eabe..1f31f3a0a 100644 --- a/src/core/src/cache/utils.rs +++ b/src/core/src/cache/utils.rs @@ -1,6 +1,7 @@ #[cfg(test)] use crate::cache::cached_batch::CacheEntry; -use crate::sync::{Arc, RwLock}; +#[cfg(test)] +use crate::sync::Arc; use arrow::array::ArrayRef; use arrow_schema::ArrowError; use bytes::Bytes; @@ -86,49 +87,6 @@ impl From for usize { } } -/// States for liquid compressor. -pub struct LiquidCompressorStates { - fsst_compressor: RwLock>>, -} - -impl std::fmt::Debug for LiquidCompressorStates { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "EtcCompressorStates") - } -} - -impl Default for LiquidCompressorStates { - fn default() -> Self { - Self::new() - } -} - -impl LiquidCompressorStates { - /// Create a new instance of LiquidCompressorStates. - pub fn new() -> Self { - Self { - fsst_compressor: RwLock::new(None), - } - } - - /// Create a new instance of LiquidCompressorStates with an fsst compressor. - pub fn new_with_fsst_compressor(fsst_compressor: Arc) -> Self { - Self { - fsst_compressor: RwLock::new(Some(fsst_compressor)), - } - } - - /// Get the fsst compressor. - pub fn fsst_compressor(&self) -> Option> { - self.fsst_compressor.read().unwrap().clone() - } - - /// Get the fsst compressor . - pub fn fsst_compressor_raw(&self) -> &RwLock>> { - &self.fsst_compressor - } -} - pub(crate) fn arrow_to_bytes(array: &ArrayRef) -> Result { use arrow::array::RecordBatch; use arrow::ipc::writer::StreamWriter; diff --git a/src/core/src/liquid_array/array.rs b/src/core/src/liquid_array/array.rs new file mode 100644 index 000000000..ac4fb7175 --- /dev/null +++ b/src/core/src/liquid_array/array.rs @@ -0,0 +1,819 @@ +use std::mem::size_of; +use std::sync::LazyLock; + +use arrow::array::{Array, ArrayRef, BooleanArray}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow_schema::{DataType, Field}; +use bytes::Bytes; +use datafusion_common::ScalarValue; +use datafusion_expr_common::operator::Operator as DataFusionOperator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column, LikeExpr, Literal}; +use vortex_array::arrays::ConstantArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::fns::like::{Like, LikeOptions}; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::{IntoArray, VortexSessionExecute}; +use vortex_arrow::ArrowSessionExt; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_ipc::messages::{BufMessageReader, DecoderMessage, EncoderMessage, MessageEncoder}; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::cache::LiquidExpr; + +use super::eval_predicate_on_array; + +const MAGIC: u32 = 0x4C51_4441; // "LQDA" for LiQuid Data Array +const VERSION: u16 = 1; +const HEADER_SIZE: usize = 16; + +static SESSION: LazyLock = LazyLock::new(session); +static COMPRESSOR: LazyLock = LazyLock::new(BtrBlocksCompressor::default); + +fn session() -> VortexSession { + let session = vortex_array::array_session(); + vortex_arrow::initialize(&session); + vortex_fastlanes::initialize(&session); + vortex_fsst::initialize(&session); + vortex_alp::initialize(&session); + vortex_runend::initialize(&session); + vortex_sparse::initialize(&session); + vortex_zigzag::initialize(&session); + vortex_datetime_parts::initialize(&session); + vortex_decimal_byte_parts::initialize(&session); + vortex_bytebool::initialize(&session); + vortex_sequence::initialize(&session); + session +} + +/// A Liquid array backed by an in-memory compressed Vortex array. +#[derive(Debug)] +pub struct LiquidArray { + array: vortex_array::ArrayRef, + arrow_type: DataType, +} + +impl LiquidArray { + /// Transcode a supported Arrow array into Vortex. + pub fn from_arrow_array(array: &ArrayRef) -> Result { + if !is_supported(array.data_type()) { + return Err(array); + } + + let arrow_type = array.data_type().clone(); + let imported = match import_array(array) { + Ok(imported) => imported, + Err(error) => { + log::warn!("failed to import {arrow_type:?} into Vortex: {error}"); + return Err(array); + } + }; + let mut ctx = SESSION.create_execution_ctx(); + let array = match COMPRESSOR.compress(&imported, &mut ctx) { + Ok(compressed) => compressed, + Err(error) => { + log::warn!("failed to compress {arrow_type:?} with Vortex: {error}"); + return Err(array); + } + }; + Ok(Self { array, arrow_type }) + } + + /// Deserialize a Vortex-backed Liquid array. + pub fn from_bytes(bytes: Bytes) -> Self { + validate_header(&bytes); + let header_size = HEADER_SIZE; + let type_len = u32::from_le_bytes( + bytes[header_size..header_size + 4] + .try_into() + .expect("Vortex Arrow type length"), + ) as usize; + let type_start = header_size + 4; + let type_end = type_start + type_len; + let arrow_type = serde_json::from_slice(&bytes[type_start..type_end]) + .expect("valid serialized Arrow data type"); + let ipc_start = type_end.next_multiple_of(8); + let mut reader = BufMessageReader::new(bytes.slice(ipc_start..)); + let dtype = match reader + .next() + .expect("Vortex dtype message") + .expect("valid Vortex IPC") + { + DecoderMessage::DType(dtype) => { + DType::from_flatbuffer(dtype, &SESSION).expect("valid Vortex dtype") + } + message => panic!("expected Vortex dtype message, got {message:?}"), + }; + let array = match reader + .next() + .expect("Vortex array message") + .expect("valid Vortex IPC") + { + DecoderMessage::Array((array, read_ctx, row_count)) => array + .decode(&dtype, row_count, &read_ctx, &SESSION) + .expect("valid Vortex array"), + message => panic!("expected Vortex array message, got {message:?}"), + }; + Self { array, arrow_type } + } + + fn export(&self, array: vortex_array::ArrayRef) -> ArrayRef { + let field = Field::new("", self.arrow_type.clone(), array.dtype().is_nullable()); + let mut ctx = SESSION.create_execution_ctx(); + SESSION + .arrow() + .execute_arrow(array, Some(&field), &mut ctx) + .expect("Vortex array must export to its import Arrow type") + } + + fn filtered_vortex(&self, filter: &BooleanBuffer) -> vortex_array::ArrayRef { + assert_eq!(filter.len(), self.len(), "filter length must match array"); + let mask = Mask::from_buffer(BitBuffer::from(filter.clone())); + self.array.filter(mask).expect("Vortex filter must succeed") + } + + fn try_vortex_predicate( + &self, + array: &vortex_array::ArrayRef, + predicate: &LiquidExpr, + ) -> Option { + let expr = predicate.physical_expr(); + if let Some(binary) = expr.downcast_ref::() { + binary.left().downcast_ref::()?; + let literal = binary.right().downcast_ref::()?; + return match binary.op() { + DataFusionOperator::Eq + | DataFusionOperator::NotEq + | DataFusionOperator::Lt + | DataFusionOperator::LtEq + | DataFusionOperator::Gt + | DataFusionOperator::GtEq => { + let op = vortex_operator(binary.op())?; + let constant = self.constant(literal.value(), array.len())?; + array.binary(constant, op).ok() + } + DataFusionOperator::LikeMatch | DataFusionOperator::NotLikeMatch => self.try_like( + array, + literal.value(), + LikeOptions { + negated: matches!(binary.op(), DataFusionOperator::NotLikeMatch), + case_insensitive: false, + }, + ), + _ => None, + }; + } + + if let Some(like) = expr.downcast_ref::() { + if like.case_insensitive() || like.expr().downcast_ref::().is_none() { + return None; + } + let pattern = like.pattern().downcast_ref::()?; + return self.try_like( + array, + pattern.value(), + LikeOptions { + negated: like.negated(), + case_insensitive: like.case_insensitive(), + }, + ); + } + None + } + + fn constant(&self, value: &ScalarValue, len: usize) -> Option { + let value = value.cast_to(&self.arrow_type).ok()?; + let arrow = value.to_array_of_size(1).ok()?; + let vortex = SESSION + .arrow() + .from_arrow_array(arrow, value.is_null()) + .ok()?; + let mut ctx = SESSION.create_execution_ctx(); + let scalar = vortex.execute_scalar(0, &mut ctx).ok()?; + Some(ConstantArray::new(scalar, len).into_array()) + } + + fn try_like( + &self, + array: &vortex_array::ArrayRef, + pattern: &ScalarValue, + options: LikeOptions, + ) -> Option { + let pattern = self.constant(pattern, array.len())?; + Some( + Like::try_new(array.clone(), pattern, options) + .ok()? + .into_array(), + ) + } + + fn try_eval_vortex_predicate( + &self, + array: &vortex_array::ArrayRef, + predicate: &LiquidExpr, + ) -> Option { + let result = self.try_vortex_predicate(array, predicate)?; + let field = Field::new("", DataType::Boolean, result.dtype().is_nullable()); + let mut ctx = SESSION.create_execution_ctx(); + let result = SESSION + .arrow() + .execute_arrow(result, Some(&field), &mut ctx) + .ok()?; + let result = result.as_any().downcast_ref::()?; + let source_validity = array + .validity() + .ok()? + .execute_mask(array.len(), &mut ctx) + .ok()?; + let source_nulls = vortex_arrow::to_null_buffer(source_validity); + // Vortex compare on extension dtypes returns non-null false for null input rows. + // Liquid's contract is null where the input is null (see Timestamp = literal). + let nulls = NullBuffer::union(result.nulls(), source_nulls.as_ref()); + Some(BooleanArray::new(result.values().clone(), nulls)) + } +} + +impl LiquidArray { + /// Get the memory size of the Liquid array. + pub fn get_array_memory_size(&self) -> usize { + self.array.nbytes() as usize + size_of::() + } + + /// Get the length of the Liquid array. + pub fn len(&self) -> usize { + self.array.len() + } + + /// Check whether the Liquid array is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Convert the Liquid array to an Arrow array. + pub fn to_arrow_array(&self) -> ArrayRef { + self.export(self.array.clone()) + } + + /// Get the original Arrow data type. + pub fn original_arrow_data_type(&self) -> DataType { + self.arrow_type.clone() + } + + /// Serialize the Liquid array. + pub fn to_bytes(&self) -> Vec { + let arrow_type = serde_json::to_vec(&self.arrow_type).expect("Arrow data type serializes"); + let mut bytes = Vec::new(); + bytes.extend_from_slice(&header_bytes()); + bytes.extend_from_slice(&(arrow_type.len() as u32).to_le_bytes()); + bytes.extend_from_slice(&arrow_type); + bytes.resize(bytes.len().next_multiple_of(8), 0); + + let mut encoder = MessageEncoder::new(SESSION.clone()); + for message in [ + EncoderMessage::DType(self.array.dtype()), + EncoderMessage::Array(&self.array), + ] { + for buffer in encoder.encode(message).expect("Vortex IPC serialization") { + bytes.extend_from_slice(&buffer); + } + } + bytes + } + + /// Filter the Liquid array and return an Arrow array. + pub fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { + self.export(self.filtered_vortex(selection)) + } + + /// Evaluate a predicate on selected rows. + pub fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> BooleanArray { + let filtered = self.filtered_vortex(filter); + if let Some(value) = predicate + .physical_expr() + .downcast_ref::() + .and_then(|literal| match literal.value() { + ScalarValue::Boolean(value) => *value, + _ => None, + }) + { + let values = if value { + BooleanBuffer::new_set(filtered.len()) + } else { + BooleanBuffer::new_unset(filtered.len()) + }; + let mut ctx = SESSION.create_execution_ctx(); + let validity = filtered + .validity() + .and_then(|validity| validity.execute_mask(filtered.len(), &mut ctx)) + .expect("Vortex array validity must execute"); + return BooleanArray::new(values, vortex_arrow::to_null_buffer(validity)); + } + let Some(result) = self.try_eval_vortex_predicate(&filtered, predicate) else { + return eval_predicate_on_array(self.export(filtered), predicate); + }; + result + } +} + +fn header_bytes() -> [u8; HEADER_SIZE] { + let mut bytes = [0; HEADER_SIZE]; + bytes[0..4].copy_from_slice(&MAGIC.to_le_bytes()); + bytes[4..6].copy_from_slice(&VERSION.to_le_bytes()); + bytes +} + +fn validate_header(bytes: &[u8]) { + assert!( + bytes.len() >= HEADER_SIZE, + "value too small for Liquid array header, expected at least {HEADER_SIZE} bytes, got {}", + bytes.len() + ); + let magic = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); + let version = u16::from_le_bytes(bytes[4..6].try_into().unwrap()); + assert_eq!(magic, MAGIC, "Invalid Liquid array magic number"); + assert_eq!(version, VERSION, "Unsupported Liquid array version"); +} + +fn import_array(array: &ArrayRef) -> VortexResult { + SESSION + .arrow() + .from_arrow_array(array.clone(), array.nulls().is_some()) +} + +fn is_supported(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, None) + | DataType::Float32 + | DataType::Float64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Utf8 + | DataType::Utf8View + | DataType::Binary + | DataType::BinaryView + ) || matches!( + data_type, + DataType::Dictionary(key, value) + if key.as_ref() == &DataType::UInt16 + && matches!(value.as_ref(), DataType::Utf8 | DataType::Binary) + ) +} + +fn vortex_operator(op: &DataFusionOperator) -> Option { + match op { + DataFusionOperator::Eq => Some(Operator::Eq), + DataFusionOperator::NotEq => Some(Operator::NotEq), + DataFusionOperator::Lt => Some(Operator::Lt), + DataFusionOperator::LtEq => Some(Operator::Lte), + DataFusionOperator::Gt => Some(Operator::Gt), + DataFusionOperator::GtEq => Some(Operator::Gte), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{ + BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Decimal128Array, DictionaryArray, + Float64Array, Int32Array, Int64Array, StringArray, StringViewArray, StructArray, + TimestampMicrosecondArray, UInt16Array, + }; + use arrow::datatypes::UInt16Type; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::CastExpr; + + use super::*; + + fn roundtrip(array: ArrayRef) { + let vortex = LiquidArray::from_arrow_array(&array).unwrap(); + assert_eq!(vortex.to_arrow_array().as_ref(), array.as_ref()); + } + + #[test] + fn roundtrips_supported_arrow_types() { + let decimal = Decimal128Array::from(vec![Some(12_345), None, Some(-9_876)]) + .with_precision_and_scale(15, 3) + .unwrap(); + let dictionary = DictionaryArray::::from_iter([ + Some("red"), + None, + Some("blue"), + Some("red"), + ]); + let arrays: Vec = vec![ + Arc::new(Int32Array::from_iter_values(0..1_100)), + Arc::new(Int64Array::from(vec![Some(-1), None, Some(7)])), + Arc::new(UInt16Array::from(vec![0, 1, u16::MAX])), + Arc::new(Date32Array::from(vec![Some(0), None, Some(20_000)])), + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(1_000), + None, + Some(9_000), + ])), + Arc::new(Float64Array::from(vec![Some(1.5), None, Some(-0.0)])), + Arc::new(decimal), + Arc::new(StringArray::from(vec![Some("alpha"), None, Some("")])), + Arc::new(StringViewArray::from(vec![ + Some(""), + None, + Some("a longer string value"), + ])), + Arc::new(BinaryArray::from(vec![ + Some(b"valid utf8".as_ref()), + None, + Some(b"".as_ref()), + ])), + Arc::new(BinaryViewArray::from_iter_values([ + vec![0xff, 0x00], + vec![], + vec![0x80], + ])), + Arc::new(dictionary), + ]; + for array in arrays { + roundtrip(array); + } + } + + #[test] + fn roundtrips_empty_array() { + roundtrip(Arc::new(StringViewArray::from(Vec::>::new()))); + } + + #[test] + fn rejects_unsupported_types() { + let boolean: ArrayRef = Arc::new(BooleanArray::from(vec![true, false])); + assert!(LiquidArray::from_arrow_array(&boolean).is_err()); + + let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let structure: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("x", DataType::Int32, false)), + values, + )])); + assert!(LiquidArray::from_arrow_array(&structure).is_err()); + } + + #[test] + fn filters_like_arrow() { + let array: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3), Some(4)])); + let selection = BooleanBuffer::from(vec![true, false, true, false]); + let vortex = LiquidArray::from_arrow_array(&array).unwrap(); + let expected = + arrow::compute::filter(&array, &BooleanArray::new(selection.clone(), None)).unwrap(); + assert_eq!(vortex.filter(&selection).as_ref(), expected.as_ref()); + } + + fn comparison_expr( + op: DataFusionOperator, + value: ScalarValue, + data_type: &DataType, + ) -> LiquidExpr { + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 0)), + op, + Arc::new(Literal::new(value)), + )); + LiquidExpr::try_new(expr, data_type).unwrap() + } + + fn like_expr(pattern: ScalarValue, negated: bool, data_type: &DataType) -> LiquidExpr { + let expr: Arc = Arc::new(LikeExpr::new( + negated, + false, + Arc::new(Column::new("c", 0)), + Arc::new(Literal::new(pattern)), + )); + LiquidExpr::try_new(expr, data_type).unwrap() + } + + fn expected_predicate( + array: &ArrayRef, + expr: &LiquidExpr, + selection: &BooleanBuffer, + ) -> BooleanArray { + let filtered = + arrow::compute::filter(array, &BooleanArray::new(selection.clone(), None)).unwrap(); + eval_predicate_on_array(filtered, expr) + } + + fn assert_predicate_matches(name: &str, array: &ArrayRef, expr: &LiquidExpr) { + let selection = BooleanBuffer::new_set(array.len()); + let vortex = LiquidArray::from_arrow_array(array).unwrap(); + let filtered = vortex.filtered_vortex(&selection); + if vortex.try_eval_vortex_predicate(&filtered, expr).is_none() { + eprintln!("Vortex predicate fell back to Arrow: {name}"); + } + assert_eq!( + vortex.try_eval_predicate(expr, &selection), + expected_predicate(array, expr, &selection), + "predicate case {name}" + ); + } + + #[test] + fn evaluates_comparisons_with_vortex() { + let integers: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(-2), + None, + Some(0), + Some(3), + Some(9), + ])); + let strings: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("x"), + None, + Some("prefix-x"), + Some("z"), + Some("x-suffix"), + ])); + let selections = [ + BooleanBuffer::new_set(5), + BooleanBuffer::from(vec![true, false, true, false, true]), + ]; + let ops = [ + DataFusionOperator::Eq, + DataFusionOperator::NotEq, + DataFusionOperator::Lt, + DataFusionOperator::LtEq, + DataFusionOperator::Gt, + DataFusionOperator::GtEq, + ]; + + for selection in selections { + let vortex = LiquidArray::from_arrow_array(&integers).unwrap(); + for op in ops { + for literal in [-1, 3] { + let expr = comparison_expr( + op, + ScalarValue::Int64(Some(literal)), + integers.data_type(), + ); + assert_eq!( + vortex.try_eval_predicate(&expr, &selection), + expected_predicate(&integers, &expr, &selection) + ); + } + } + + let vortex = LiquidArray::from_arrow_array(&strings).unwrap(); + for op in ops { + for literal in ["x", "z"] { + let expr = comparison_expr( + op, + ScalarValue::Utf8View(Some(literal.into())), + strings.data_type(), + ); + assert_eq!( + vortex.try_eval_predicate(&expr, &selection), + expected_predicate(&strings, &expr, &selection) + ); + } + } + } + } + + #[test] + fn evaluates_tricky_predicates() { + let dictionary: ArrayRef = Arc::new(DictionaryArray::::from_iter([ + Some("x"), + None, + Some("prefix-x"), + Some("z"), + ])); + for op in [ + DataFusionOperator::Eq, + DataFusionOperator::NotEq, + DataFusionOperator::Lt, + ] { + let name = format!("Dictionary(UInt16, Utf8) {op} Utf8"); + let expr = comparison_expr( + op, + ScalarValue::Utf8(Some("x".into())), + dictionary.data_type(), + ); + assert_predicate_matches(&name, &dictionary, &expr); + } + let expr = like_expr( + ScalarValue::Utf8(Some("%x%".into())), + false, + dictionary.data_type(), + ); + assert_predicate_matches("Dictionary(UInt16, Utf8) LIKE Utf8", &dictionary, &expr); + + let binary: ArrayRef = Arc::new(BinaryArray::from(vec![ + Some(b"ax".as_slice()), + None, + Some(b"by".as_slice()), + Some(b"cz".as_slice()), + ])); + for op in [DataFusionOperator::Eq, DataFusionOperator::Lt] { + let name = format!("Binary {op} Binary"); + let expr = comparison_expr( + op, + ScalarValue::Binary(Some(b"by".to_vec())), + binary.data_type(), + ); + assert_predicate_matches(&name, &binary, &expr); + } + + let binary_view: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some(b"\xff\x00".as_slice()), + None, + Some(b"\x80".as_slice()), + Some(b"valid".as_slice()), + ])); + let expr = comparison_expr( + DataFusionOperator::Eq, + ScalarValue::BinaryView(Some(b"\x80".to_vec())), + binary_view.data_type(), + ); + assert_predicate_matches("invalid BinaryView = BinaryView", &binary_view, &expr); + + let dates: ArrayRef = Arc::new(Date32Array::from(vec![Some(0), None, Some(20_000)])); + for (op, value) in [ + (DataFusionOperator::Gt, 10_000), + (DataFusionOperator::Eq, 20_000), + ] { + let name = format!("Date32 {op} Date32"); + let expr = comparison_expr(op, ScalarValue::Date32(Some(value)), dates.data_type()); + assert_predicate_matches(&name, &dates, &expr); + } + + let timestamps: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![ + Some(1_000), + None, + Some(9_000), + ])); + for (op, value) in [ + (DataFusionOperator::Gt, 5_000), + (DataFusionOperator::Eq, 9_000), + ] { + let name = format!("Timestamp(Microsecond) {op} Timestamp(Microsecond)"); + let expr = comparison_expr( + op, + ScalarValue::TimestampMicrosecond(Some(value), None), + timestamps.data_type(), + ); + assert_predicate_matches(&name, ×tamps, &expr); + } + + let decimals: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(1_000), None, Some(12_345)]) + .with_precision_and_scale(15, 3) + .unwrap(), + ); + let expr = comparison_expr( + DataFusionOperator::GtEq, + ScalarValue::Decimal128(Some(5_000), 15, 3), + decimals.data_type(), + ); + assert_predicate_matches("Decimal128(15, 3) >= Decimal128", &decimals, &expr); + + let integers: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); + let physical: Arc = Arc::new(BinaryExpr::new( + Arc::new(CastExpr::new( + Arc::new(Column::new("c", 0)), + DataType::Int32, + None, + )), + DataFusionOperator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(3)))), + )); + let expr = LiquidExpr::try_new(physical, integers.data_type()).unwrap(); + assert_predicate_matches("Int64 = Int32", &integers, &expr); + + let strings: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("alpha"), + None, + Some("beta"), + ])); + let physical: Arc = Arc::new(BinaryExpr::new( + Arc::new(CastExpr::new( + Arc::new(Column::new("c", 0)), + DataType::Utf8, + None, + )), + DataFusionOperator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("beta".into())))), + )); + let expr = LiquidExpr::try_new(physical, strings.data_type()).unwrap(); + assert_predicate_matches("Utf8View = Utf8", &strings, &expr); + } + + #[test] + fn evaluates_like_with_vortex() { + let array: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("x"), + None, + Some("prefix-x"), + Some("x-suffix"), + Some("none"), + ])); + let vortex = LiquidArray::from_arrow_array(&array).unwrap(); + for (pattern, negated) in [("%x%", false), ("x%", false), ("%x", false), ("%x%", true)] { + let expr = like_expr( + ScalarValue::Utf8View(Some(pattern.into())), + negated, + &DataType::Utf8View, + ); + let selection = BooleanBuffer::new_set(array.len()); + assert_eq!( + vortex.try_eval_predicate(&expr, &selection), + expected_predicate(&array, &expr, &selection) + ); + } + } + + #[test] + fn boolean_literal_preserves_string_nulls() { + let strings = (0..3_000) + .map(|index| (index % 11 != 0).then_some("a nullable string value")) + .collect::>(); + let array: ArrayRef = Arc::new(StringViewArray::from(strings)); + let physical: Arc = + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + let predicate = LiquidExpr::try_new(physical, array.data_type()).unwrap(); + let vortex = LiquidArray::from_arrow_array(&array).unwrap(); + assert_eq!( + vortex.try_eval_predicate(&predicate, &BooleanBuffer::new_set(array.len())), + BooleanArray::new(BooleanBuffer::new_set(array.len()), array.nulls().cloned()) + ); + } + + #[test] + fn roundtrips_dictionary_export_branches() { + let identical_values = vec!["same"; 64]; + let identical: ArrayRef = Arc::new(DictionaryArray::::from_iter( + identical_values.iter().map(|value| Some(*value)), + )); + + let unique_values = (0..256) + .map(|index| format!("unique-{index}")) + .collect::>(); + let unique: ArrayRef = Arc::new(DictionaryArray::::from_iter( + unique_values.iter().map(|value| Some(value.as_str())), + )); + + let long_values = (0..1_500) + .map(|index| format!("value-{}", index % 7)) + .collect::>(); + let long: ArrayRef = Arc::new(DictionaryArray::::from_iter( + long_values.iter().map(|value| Some(value.as_str())), + )); + + for array in [identical, unique, long] { + roundtrip(array); + } + } + + #[test] + fn ipc_roundtrips() { + let arrays: Vec = vec![ + Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])), + Arc::new(StringViewArray::from(vec![ + Some("one"), + None, + Some("three"), + ])), + ]; + for array in arrays { + let vortex = LiquidArray::from_arrow_array(&array).unwrap(); + let decoded = LiquidArray::from_bytes(Bytes::from(vortex.to_bytes())); + assert_eq!(decoded.to_arrow_array().as_ref(), array.as_ref()); + } + } + + #[test] + #[should_panic(expected = "Invalid Liquid array magic number")] + fn rejects_bad_ipc_magic() { + LiquidArray::from_bytes(Bytes::from_static(&[0; HEADER_SIZE])); + } + + #[test] + #[should_panic(expected = "Unsupported Liquid array version")] + fn rejects_bad_ipc_version() { + let mut header = header_bytes(); + header[4..6].copy_from_slice(&(VERSION + 1).to_le_bytes()); + LiquidArray::from_bytes(Bytes::copy_from_slice(&header)); + } +} diff --git a/src/core/src/liquid_array/byte_view_array/comparisons.rs b/src/core/src/liquid_array/byte_view_array/comparisons.rs deleted file mode 100644 index 22c47f57f..000000000 --- a/src/core/src/liquid_array/byte_view_array/comparisons.rs +++ /dev/null @@ -1,457 +0,0 @@ -use arrow::array::{Array, DictionaryArray}; -use arrow::array::{BinaryArray, BooleanArray, BooleanBufferBuilder, StringArray, cast::AsArray}; -use arrow::datatypes::UInt16Type; -use arrow_schema::DataType; -use datafusion_common::ScalarValue; -use datafusion_expr_common::columnar_value::ColumnarValue; -use datafusion_expr_common::operator::Operator; -use datafusion_physical_expr_common::datum::apply_cmp; -use fsst::Compressor; -use std::sync::Arc; -use std::vec; - -use super::LiquidByteViewArray; -use super::fingerprint::{StringFingerprint, substring_pattern_bytes}; -use crate::liquid_array::byte_view_array::operator::{self, ByteViewOperator}; -use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::{FsstBacking, PrefixKey}; - -impl LiquidByteViewArray { - /// Compare equality with a byte needle - pub(super) fn compare_equals(&self, needle: &[u8]) -> BooleanArray { - let shared_prefix_len = self.shared_prefix.len(); - let num_unique = self.prefix_keys.len(); - if needle.len() < shared_prefix_len || needle[..shared_prefix_len] != self.shared_prefix { - return self.map_dictionary_results_to_array_results(vec![false; num_unique]); - } - - let needle_suffix = &needle[shared_prefix_len..]; - let needle_len = needle_suffix.len(); - let prefix_len = PrefixKey::prefix_len(); - let mut dict_results = vec![false; num_unique]; - - if needle_len <= prefix_len { - for (i, prefix_key) in self.prefix_keys.iter().enumerate().take(num_unique) { - let known_len = if prefix_key.len_byte() == 255 { - None - } else { - Some(prefix_key.len_byte() as usize) - }; - if let Some(l) = known_len - && l == needle_len - && prefix_key.prefix7()[..l] == needle_suffix[..l] - { - dict_results[i] = true; - } - } - - return self.map_dictionary_results_to_array_results(dict_results); - } - - let compressed_needle = compress_needle(self.fsst_buffer.compressor(), needle); - - for (i, prefix_key) in self.prefix_keys.iter().enumerate().take(num_unique) { - let known_len = if prefix_key.len_byte() == 255 { - None - } else { - Some(prefix_key.len_byte() as usize) - }; - - match known_len { - Some(l) => { - if l != needle_len { - continue; - } - } - None => { - if needle_len < 255 { - continue; - } - } - } - - if prefix_key.prefix7()[..prefix_len] == needle_suffix[..prefix_len] { - let compressed_value = self.fsst_buffer.get_compressed_slice(i); - if compressed_value == compressed_needle.as_slice() { - dict_results[i] = true; - } - } - } - - self.map_dictionary_results_to_array_results(dict_results) - } - - /// Compare not equals with a byte needle - fn compare_not_equals(&self, needle: &[u8]) -> BooleanArray { - let result = self.compare_equals(needle); - let (values, nulls) = result.into_parts(); - let values = !&values; - BooleanArray::new(values, nulls) - } - - /// Compare with prefix optimization and fallback to Arrow operations - pub fn compare_with(&self, needle: &[u8], op: &ByteViewOperator) -> BooleanArray { - match op { - ByteViewOperator::Comparison(cmp) => self.compare_with_inner(needle, cmp), - ByteViewOperator::Equality(operator::Equality::Eq) => self.compare_equals(needle), - ByteViewOperator::Equality(operator::Equality::NotEq) => { - self.compare_not_equals(needle) - } - ByteViewOperator::SubString(op) => { - if let Some(fingerprints) = self.string_fingerprints.as_ref() { - let pattern = - substring_pattern_bytes(needle).expect("Invalid substring pattern"); - self.compare_like_substring(pattern, *op, fingerprints) - } else { - let fallback = ByteViewOperator::SubString(*op); - self.compare_with_arrow_fallback(needle, &fallback) - } - } - } - } - - /// Prefix optimization for ordering operations - pub(super) fn compare_with_inner( - &self, - needle: &[u8], - op: &operator::Comparison, - ) -> BooleanArray { - let (mut dict_results, ambiguous) = self.compare_with_prefix(needle, op); - - // For values needing full comparison, load buffer and decompress - if !ambiguous.is_empty() { - let (values_buffer, offsets_buffer) = - self.fsst_buffer.to_uncompressed_selected(&ambiguous); - let binary_array = - unsafe { BinaryArray::new_unchecked(offsets_buffer, values_buffer, None) }; - - for (pos, &dict_index) in ambiguous.iter().enumerate() { - let value_cmp = binary_array.value(pos).cmp(needle); - let result = match (op, value_cmp) { - (operator::Comparison::Lt, std::cmp::Ordering::Less) => true, - (operator::Comparison::Lt, _) => false, - ( - operator::Comparison::LtEq, - std::cmp::Ordering::Less | std::cmp::Ordering::Equal, - ) => true, - (operator::Comparison::LtEq, _) => false, - (operator::Comparison::Gt, std::cmp::Ordering::Greater) => true, - (operator::Comparison::Gt, _) => false, - ( - operator::Comparison::GtEq, - std::cmp::Ordering::Greater | std::cmp::Ordering::Equal, - ) => true, - (operator::Comparison::GtEq, _) => false, - }; - dict_results[dict_index] = result; - } - } - - self.map_dictionary_results_to_array_results(dict_results) - } - - /// Fallback to Arrow operations for unsupported operations - fn compare_with_arrow_fallback(&self, needle: &[u8], op: &ByteViewOperator) -> BooleanArray { - let dict_array = self.to_dict_arrow(); - compare_with_arrow_inner(dict_array, needle, op) - } - - pub(super) fn compare_like_substring( - &self, - needle: &[u8], - operator: operator::SubString, - fingerprints: &Arc<[u32]>, - ) -> BooleanArray { - let (dict_results, ambiguous) = compute_fingerprint_candidates(needle, fingerprints); - - let dict_results = if !ambiguous.is_empty() { - let (values_buffer, offsets_buffer) = - self.fsst_buffer.to_uncompressed_selected(&ambiguous); - apply_like_match_on_candidates( - dict_results, - ambiguous, - values_buffer, - offsets_buffer, - needle, - operator, - ) - } else { - dict_results - }; - - self.map_dictionary_results_to_array_results(dict_results) - } -} - -impl LiquidByteViewArray { - /// Return (selected_rows, ambiguous_rows, unique_rows) based on prefix-only comparison. - pub fn prefix_compare_counts( - &self, - needle: &[u8], - op: &operator::Comparison, - ) -> (usize, usize, usize) { - let (dict_results, ambiguous) = self.compare_with_prefix(needle, op); - let selected_rows = dict_results.iter().filter(|&x| *x).count(); - (selected_rows, ambiguous.len(), self.dictionary_keys.len()) - } - - fn map_dictionary_results_to_array_results(&self, dict_results: Vec) -> BooleanArray { - let len = self.dictionary_keys.len(); - let mut builder = BooleanBufferBuilder::new(len); - builder.advance(len); - for index in 0..len { - if !self.dictionary_keys.is_valid(index) { - continue; - } - - let dict_index = self.dictionary_keys.value(index) as usize; - debug_assert!(dict_index < dict_results.len()); - if dict_results.get(dict_index).copied().unwrap_or(false) { - builder.set_bit(index, true); - } - } - - let values = builder.finish(); - if let Some(nulls) = self.nulls() { - BooleanArray::new(values, Some(nulls.clone())) - } else { - BooleanArray::new(values, None) - } - } - - // returns a tuple of compare_results and ambiguous indices - #[inline(never)] - pub(super) fn compare_with_prefix( - &self, - needle: &[u8], - op: &operator::Comparison, - ) -> (Vec, Vec) { - // Try to short-circuit based on shared prefix comparison - if let Some(result) = self.compare_with_shared_prefix(needle, op) { - return (vec![result; self.dictionary_keys.len()], Vec::new()); - } - - let needle_suffix = &needle[self.shared_prefix.len()..]; - let num_unique = self.prefix_keys.len(); - let mut dict_results = vec![false; num_unique]; - let mut ambiguous = Vec::new(); - - let cmp_len = needle_suffix.len().min(PrefixKey::prefix_len()); - if cmp_len == 0 { - for (i, prefix_key) in self.prefix_keys.iter().enumerate() { - let is_empty_suffix = prefix_key.len_byte() == 0; - dict_results[i] = match op { - operator::Comparison::Lt => false, - operator::Comparison::LtEq => is_empty_suffix, - operator::Comparison::Gt => !is_empty_suffix, - operator::Comparison::GtEq => true, - }; - } - return (dict_results, ambiguous); - } - - for (i, prefix_key) in self.prefix_keys.iter().enumerate() { - let ordering = bytes_cmp_short(prefix_key.prefix7(), needle_suffix, cmp_len); - match ordering { - std::cmp::Ordering::Less => match op { - operator::Comparison::Lt | operator::Comparison::LtEq => { - dict_results[i] = true; - } - operator::Comparison::Gt | operator::Comparison::GtEq => { - dict_results[i] = false; - } - }, - std::cmp::Ordering::Greater => match op { - operator::Comparison::Lt | operator::Comparison::LtEq => { - dict_results[i] = false; - } - operator::Comparison::Gt | operator::Comparison::GtEq => { - dict_results[i] = true; - } - }, - std::cmp::Ordering::Equal => { - ambiguous.push(i); - } - } - } - (dict_results, ambiguous) - } - - // returns a tuple of compare_results and ambiguous indices - - /// Check if shared prefix comparison can short-circuit the entire operation - fn compare_with_shared_prefix(&self, needle: &[u8], op: &operator::Comparison) -> Option { - let shared_prefix_len = self.shared_prefix.len(); - - let needle_shared_len = std::cmp::min(needle.len(), shared_prefix_len); - let shared_cmp = self.shared_prefix[..needle_shared_len].cmp(&needle[..needle_shared_len]); - match (op, shared_cmp) { - (operator::Comparison::Lt | operator::Comparison::LtEq, std::cmp::Ordering::Less) => { - Some(true) - } - ( - operator::Comparison::Lt | operator::Comparison::LtEq, - std::cmp::Ordering::Greater, - ) => Some(false), - ( - operator::Comparison::Gt | operator::Comparison::GtEq, - std::cmp::Ordering::Greater, - ) => Some(true), - (operator::Comparison::Gt | operator::Comparison::GtEq, std::cmp::Ordering::Less) => { - Some(false) - } - (_, std::cmp::Ordering::Equal) => { - if needle.len() < shared_prefix_len { - match op { - operator::Comparison::Gt | operator::Comparison::GtEq => Some(true), - operator::Comparison::Lt => Some(false), - operator::Comparison::LtEq => Some(false), - } - } else { - None - } - } - } - } -} - -fn compare_with_arrow_inner( - dict_array: DictionaryArray, - needle: &[u8], - op: &ByteViewOperator, -) -> BooleanArray { - let needle_scalar = match dict_array.values().data_type() { - DataType::Utf8 => ScalarValue::Utf8(Some( - std::str::from_utf8(needle) - .expect("utf8 needle") - .to_string(), - )), - DataType::Utf8View => ScalarValue::Utf8View(Some( - std::str::from_utf8(needle) - .expect("utf8 needle") - .to_string(), - )), - DataType::LargeUtf8 => ScalarValue::LargeUtf8(Some( - std::str::from_utf8(needle) - .expect("utf8 needle") - .to_string(), - )), - DataType::Binary => ScalarValue::Binary(Some(needle.to_vec())), - DataType::BinaryView => ScalarValue::BinaryView(Some(needle.to_vec())), - DataType::LargeBinary => ScalarValue::LargeBinary(Some(needle.to_vec())), - _ => ScalarValue::Binary(Some(needle.to_vec())), - }; - let lhs = ColumnarValue::Array(Arc::new(dict_array)); - let rhs = ColumnarValue::Scalar(needle_scalar); - let op = Operator::from(op); - let result = apply_cmp(op, &lhs, &rhs); - - match result.expect("ArrowError") { - ColumnarValue::Array(arr) => arr.as_boolean().clone(), - ColumnarValue::Scalar(_) => unreachable!(), - } -} - -fn compress_needle(compressor: &Compressor, needle: &[u8]) -> Vec { - let mut compressed = Vec::with_capacity(needle.len().saturating_mul(2)); - // SAFETY: the largest compressed size is all escapes == 2 * plaintext_len. - unsafe { - let len = compressor.compress_into(needle, compressed.spare_capacity_mut()); - compressed.set_len(len); - } - compressed -} - -fn bytes_cmp_const(left: &[u8; N], right: &[u8; N]) -> std::cmp::Ordering { - left.cmp(right) -} - -fn bytes_cmp_short(left: &[u8], right: &[u8], len: usize) -> std::cmp::Ordering { - match len { - 0 => std::cmp::Ordering::Equal, - 1 => bytes_cmp_const::<1>( - &left[..1].try_into().unwrap(), - &right[..1].try_into().unwrap(), - ), - 2 => bytes_cmp_const::<2>( - &left[..2].try_into().unwrap(), - &right[..2].try_into().unwrap(), - ), - 3 => bytes_cmp_const::<3>( - &left[..3].try_into().unwrap(), - &right[..3].try_into().unwrap(), - ), - 4 => bytes_cmp_const::<4>( - &left[..4].try_into().unwrap(), - &right[..4].try_into().unwrap(), - ), - 5 => bytes_cmp_const::<5>( - &left[..5].try_into().unwrap(), - &right[..5].try_into().unwrap(), - ), - 6 => bytes_cmp_const::<6>( - &left[..6].try_into().unwrap(), - &right[..6].try_into().unwrap(), - ), - 7 => bytes_cmp_const::<7>( - &left[..7].try_into().unwrap(), - &right[..7].try_into().unwrap(), - ), - _ => left[..len].cmp(&right[..len]), - } -} - -/// Compute which dictionary entries are candidates for matching based on fingerprints. -/// Returns a tuple of (dict_results, ambiguous_indices). -fn compute_fingerprint_candidates( - needle: &[u8], - fingerprints: &Arc<[u32]>, -) -> (Vec, Vec) { - let needle_fp = StringFingerprint::from_bytes(needle); - let dict_results = vec![false; fingerprints.len()]; - let mut ambiguous = Vec::new(); - - for (index, &bits) in fingerprints.iter().enumerate() { - if StringFingerprint::from_bits(bits).might_contain(needle_fp) { - ambiguous.push(index); - } - } - - (dict_results, ambiguous) -} - -/// Apply LIKE match operation on candidate dictionary entries. -/// Returns updated dict_results with matches marked as true. -fn apply_like_match_on_candidates( - mut dict_results: Vec, - ambiguous: Vec, - values_buffer: arrow::buffer::Buffer, - offsets_buffer: arrow::buffer::OffsetBuffer, - needle: &[u8], - operator: operator::SubString, -) -> Vec { - // Safety: the offsets and values are valid because they are from fsst buffer, which already checked utf-8. - let values = unsafe { StringArray::new_unchecked(offsets_buffer, values_buffer, None) }; - let pattern = std::str::from_utf8(needle).ok().unwrap(); - let pattern = format!("%{}%", pattern); - - let lhs = ColumnarValue::Array(Arc::new(values)); - let rhs = ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern))); - let result = apply_cmp(Operator::LikeMatch, &lhs, &rhs).unwrap(); - let result = result.into_array(ambiguous.len()).unwrap(); - let matches = result.as_boolean(); - - for (pos, &dict_index) in ambiguous.iter().enumerate() { - if !matches.is_null(pos) && matches.value(pos) { - dict_results[dict_index] = true; - } - } - - if operator == operator::SubString::NotContains { - for value in &mut dict_results { - *value = !*value; - } - } - - dict_results -} diff --git a/src/core/src/liquid_array/byte_view_array/conversions.rs b/src/core/src/liquid_array/byte_view_array/conversions.rs deleted file mode 100644 index 4fd79b0a3..000000000 --- a/src/core/src/liquid_array/byte_view_array/conversions.rs +++ /dev/null @@ -1,392 +0,0 @@ -use arrow::array::{ - Array, ArrayAccessor, ArrayIter, BinaryArray, BinaryViewArray, DictionaryArray, - GenericByteArray, StringArray, StringViewArray, UInt16Array, cast::AsArray, types::UInt16Type, -}; -use arrow::datatypes::ByteArrayType; -use arrow_schema::DataType; -use fsst::Compressor; -use std::sync::Arc; - -use super::{ArrowByteType, ByteViewBuildOptions, LiquidByteViewArray}; -use crate::liquid_array::byte_view_array::fingerprint::StringFingerprint; -use crate::liquid_array::raw::fsst_buffer::{ - FsstArray, FsstBacking, PrefixKey, RawFsstBuffer, train_compressor, -}; -use crate::utils::CheckedDictionaryArray; - -impl LiquidByteViewArray { - /// Create a LiquidByteViewArray from an Arrow StringViewArray - pub fn from_string_view_array( - array: &StringViewArray, - compressor: Arc, - ) -> LiquidByteViewArray { - Self::from_view_array_inner( - array, - compressor, - ByteViewBuildOptions::new(ArrowByteType::Utf8View), - ) - } - - /// Create a LiquidByteViewArray from an Arrow BinaryViewArray - pub fn from_binary_view_array( - array: &BinaryViewArray, - compressor: Arc, - ) -> LiquidByteViewArray { - Self::from_view_array_inner( - array, - compressor, - ByteViewBuildOptions::new(ArrowByteType::BinaryView), - ) - } - - /// Create a LiquidByteViewArray from an Arrow StringArray - pub fn from_string_array( - array: &StringArray, - compressor: Arc, - ) -> LiquidByteViewArray { - Self::from_byte_array_inner( - array, - compressor, - ByteViewBuildOptions::new(ArrowByteType::Utf8), - ) - } - - /// Create a LiquidByteViewArray from an Arrow BinaryArray - pub fn from_binary_array( - array: &BinaryArray, - compressor: Arc, - ) -> LiquidByteViewArray { - Self::from_byte_array_inner( - array, - compressor, - ByteViewBuildOptions::new(ArrowByteType::Binary), - ) - } - - /// Train a compressor from an Arrow StringViewArray - pub fn train_from_string_view( - array: &StringViewArray, - ) -> (Arc, LiquidByteViewArray) { - Self::train_from_string_view_inner( - array, - ByteViewBuildOptions::new(ArrowByteType::Utf8View), - ) - } - - /// Train a compressor from an Arrow BinaryViewArray - pub fn train_from_binary_view( - array: &BinaryViewArray, - ) -> (Arc, LiquidByteViewArray) { - Self::train_from_binary_view_inner( - array, - ByteViewBuildOptions::new(ArrowByteType::BinaryView), - ) - } - - /// Train a compressor from an Arrow ByteArray. - pub fn train_from_arrow( - array: &GenericByteArray, - ) -> (Arc, LiquidByteViewArray) { - Self::train_from_arrow_inner( - array, - ByteViewBuildOptions::new(ArrowByteType::from_arrow_type(&T::DATA_TYPE)), - ) - } - - /// Only used when the dictionary is read from a trusted parquet reader, - /// which reads a trusted parquet file, written by a trusted writer. - /// - /// # Safety - /// The caller must ensure that the values in the dictionary are unique. - pub unsafe fn from_unique_dict_array( - array: &DictionaryArray, - compressor: Arc, - ) -> LiquidByteViewArray { - let arrow_type = ArrowByteType::from_arrow_type(array.values().data_type()); - Self::from_dict_array_inner( - unsafe { CheckedDictionaryArray::new_unchecked_i_know_what_i_am_doing(array) }, - compressor, - ByteViewBuildOptions::new(arrow_type), - ) - } - - pub(crate) unsafe fn from_unique_dict_array_with_options( - array: &DictionaryArray, - compressor: Arc, - options: ByteViewBuildOptions, - ) -> LiquidByteViewArray { - Self::from_dict_array_inner( - unsafe { CheckedDictionaryArray::new_unchecked_i_know_what_i_am_doing(array) }, - compressor, - options, - ) - } - - /// Train a compressor from an Arrow DictionaryArray. - pub fn train_from_arrow_dict( - array: &DictionaryArray, - ) -> (Arc, LiquidByteViewArray) { - let options = if array.values().data_type() == &DataType::Utf8 { - ByteViewBuildOptions::new(ArrowByteType::Dict16Utf8) - } else if array.values().data_type() == &DataType::Binary { - ByteViewBuildOptions::new(ArrowByteType::Dict16Binary) - } else { - panic!("Unsupported dictionary type: {:?}", array.data_type()) - }; - Self::train_from_arrow_dict_inner(array, options) - } - - pub(crate) fn train_from_string_view_inner( - array: &StringViewArray, - options: ByteViewBuildOptions, - ) -> (Arc, LiquidByteViewArray) { - let compressor = Self::train_compressor(array.iter()); - ( - compressor.clone(), - Self::from_view_array_inner(array, compressor, options), - ) - } - - pub(crate) fn train_from_binary_view_inner( - array: &BinaryViewArray, - options: ByteViewBuildOptions, - ) -> (Arc, LiquidByteViewArray) { - let compressor = Self::train_compressor_bytes(array.iter()); - ( - compressor.clone(), - Self::from_view_array_inner(array, compressor, options), - ) - } - - pub(crate) fn train_from_arrow_inner( - array: &GenericByteArray, - options: ByteViewBuildOptions, - ) -> (Arc, LiquidByteViewArray) { - let dict = CheckedDictionaryArray::from_byte_array::(array); - let value_type = dict.as_ref().values().data_type(); - - let compressor = if value_type == &DataType::Utf8 { - Self::train_compressor(dict.as_ref().values().as_string::().iter()) - } else { - Self::train_compressor_bytes(dict.as_ref().values().as_binary::().iter()) - }; - ( - compressor.clone(), - Self::from_dict_array_inner(dict, compressor, options), - ) - } - - pub(crate) fn train_from_arrow_dict_inner( - array: &DictionaryArray, - options: ByteViewBuildOptions, - ) -> (Arc, LiquidByteViewArray) { - if array.values().data_type() == &DataType::Utf8 { - let values = array.values().as_string::(); - - let compressor = Self::train_compressor(values.iter()); - ( - compressor.clone(), - Self::from_dict_array_inner( - CheckedDictionaryArray::new_checked(array), - compressor, - options, - ), - ) - } else if array.values().data_type() == &DataType::Binary { - let values = array.values().as_binary::(); - let compressor = Self::train_compressor_bytes(values.iter()); - ( - compressor.clone(), - Self::from_dict_array_inner( - CheckedDictionaryArray::new_checked(array), - compressor, - options, - ), - ) - } else { - panic!("Unsupported dictionary type: {:?}", array.data_type()) - } - } - - /// Train a compressor from an iterator of strings - pub fn train_compressor<'a, T: ArrayAccessor>( - array: ArrayIter, - ) -> Arc { - Arc::new(train_compressor( - array.filter_map(|s| s.as_ref().map(|s| s.as_bytes())), - )) - } - - /// Train a compressor from an iterator of byte arrays - pub fn train_compressor_bytes<'a, T: ArrayAccessor>( - array: ArrayIter, - ) -> Arc { - Arc::new(train_compressor( - array.filter_map(|s| s.as_ref().map(|s| *s)), - )) - } - - /// Generic implementation for view arrays (StringViewArray and BinaryViewArray) - pub(crate) fn from_view_array_inner( - array: &T, - compressor: Arc, - options: ByteViewBuildOptions, - ) -> LiquidByteViewArray - where - T: Array + 'static, - { - // Convert view array to CheckedDictionaryArray using existing infrastructure - let dict = if let Some(string_view) = array.as_any().downcast_ref::() { - CheckedDictionaryArray::from_string_view_array(string_view) - } else if let Some(binary_view) = array.as_any().downcast_ref::() { - CheckedDictionaryArray::from_binary_view_array(binary_view) - } else { - panic!("Unsupported view array type") - }; - - Self::from_dict_array_inner(dict, compressor, options) - } - - pub(crate) fn from_byte_array_inner( - array: &GenericByteArray, - compressor: Arc, - options: ByteViewBuildOptions, - ) -> LiquidByteViewArray { - let dict = CheckedDictionaryArray::from_byte_array::(array); - Self::from_dict_array_inner(dict, compressor, options) - } - - /// Core implementation that converts a CheckedDictionaryArray to LiquidByteViewArray - fn from_dict_array_inner( - dict: CheckedDictionaryArray, - compressor: Arc, - options: ByteViewBuildOptions, - ) -> LiquidByteViewArray { - let (keys, values) = dict.as_ref().clone().into_parts(); - let arrow_type = options.arrow_type; - - // Calculate shared prefix directly from values array without intermediate allocations - let shared_prefix = if values.is_empty() { - Vec::new() - } else { - // Get first value as initial candidate for shared prefix - let first_value_bytes = if let Some(string_values) = values.as_string_opt::() { - string_values.value(0).as_bytes() - } else if let Some(binary_values) = values.as_binary_opt::() { - binary_values.value(0) - } else { - panic!("Unsupported dictionary value type") - }; - - let mut shared_prefix = first_value_bytes.to_vec(); - - // Compare with remaining values and truncate shared prefix - for i in 1..values.len() { - let value_bytes = if let Some(string_values) = values.as_string_opt::() { - string_values.value(i).as_bytes() - } else if let Some(binary_values) = values.as_binary_opt::() { - binary_values.value(i) - } else { - panic!("Unsupported dictionary value type") - }; - - let common_len = shared_prefix - .iter() - .zip(value_bytes.iter()) - .take_while(|(a, b)| a == b) - .count(); - shared_prefix.truncate(common_len); - - // Early exit if no common prefix - if shared_prefix.is_empty() { - break; - } - } - - shared_prefix - }; - - let shared_prefix_len = shared_prefix.len(); - - // Prefix keys - one per unique value in dictionary. - let mut prefix_keys = Vec::with_capacity(values.len()); - let mut fingerprints = options - .build_fingerprints - .then(|| Vec::with_capacity(values.len())); - - let mut compress_buffer = Vec::with_capacity(1024 * 1024 * 2); - - // Create the raw buffer and get the byte offsets - let (raw_fsst_buffer, byte_offsets) = - if let Some(string_values) = values.as_string_opt::() { - RawFsstBuffer::from_byte_slices( - string_values.iter().map(|s| s.map(|s| s.as_bytes())), - compressor.clone(), - &mut compress_buffer, - ) - } else if let Some(binary_values) = values.as_binary_opt::() { - RawFsstBuffer::from_byte_slices( - binary_values.iter(), - compressor.clone(), - &mut compress_buffer, - ) - } else { - panic!("Unsupported dictionary value type") - }; - - for i in 0..values.len() { - let value_bytes = if let Some(string_values) = values.as_string_opt::() { - string_values.value(i).as_bytes() - } else if let Some(binary_values) = values.as_binary_opt::() { - binary_values.value(i) - } else { - panic!("Unsupported dictionary value type") - }; - - let remaining_bytes = if shared_prefix_len < value_bytes.len() { - &value_bytes[shared_prefix_len..] - } else { - &[] - }; - - prefix_keys.push(PrefixKey::new(remaining_bytes)); - if let Some(ref mut fingerprints) = fingerprints { - fingerprints.push(StringFingerprint::from_bytes(value_bytes).bits()); - } - } - - assert_eq!(values.len(), byte_offsets.len() - 1); - - let prefix_keys: Arc<[PrefixKey]> = prefix_keys.into(); - - let mut array = LiquidByteViewArray::from_parts( - keys, - prefix_keys, - FsstArray::from_byte_offsets(Arc::new(raw_fsst_buffer), &byte_offsets, compressor), - arrow_type, - shared_prefix, - ); - if let Some(fingerprints) = fingerprints { - array.string_fingerprints = Some(Arc::from(fingerprints.into_boxed_slice())); - } - array - } - - /// Create LiquidByteViewArray from parts - pub(super) fn from_parts( - dictionary_keys: UInt16Array, - prefix_keys: Arc<[PrefixKey]>, - fsst_buffer: B, - original_arrow_type: ArrowByteType, - shared_prefix: Vec, - ) -> Self { - Self { - dictionary_keys, - prefix_keys, - fsst_buffer, - original_arrow_type, - shared_prefix, - string_fingerprints: None, - } - } -} diff --git a/src/core/src/liquid_array/byte_view_array/fingerprint.rs b/src/core/src/liquid_array/byte_view_array/fingerprint.rs deleted file mode 100644 index bc6309d45..000000000 --- a/src/core/src/liquid_array/byte_view_array/fingerprint.rs +++ /dev/null @@ -1,49 +0,0 @@ -const FINGERPRINT_BUCKETS: u8 = 32; -const FINGERPRINT_MASK: u8 = FINGERPRINT_BUCKETS - 1; - -// 32-bit bucketed fingerprint for a string's byte set. -#[derive(Clone, Copy, Debug)] -pub(super) struct StringFingerprint(u32); - -impl StringFingerprint { - // Construct directly from a precomputed 32-bit mask. - pub(super) fn from_bits(bits: u32) -> Self { - Self(bits) - } - - // Map each byte into a bucket and set its bit (round-robin over 32 buckets). - pub(super) fn from_bytes(bytes: &[u8]) -> Self { - let mut bits = 0u32; - for &byte in bytes { - let bucket = (byte & FINGERPRINT_MASK) as u32; - bits |= 1u32 << bucket; - } - Self(bits) - } - - pub(super) fn bits(self) -> u32 { - self.0 - } - - // Returns false only when a substring cannot be present. - pub(super) fn might_contain(self, needle: Self) -> bool { - (self.0 & needle.0) == needle.0 - } -} - -pub(super) fn substring_pattern_bytes(pattern: &[u8]) -> Option<&[u8]> { - if pattern.len() < 2 { - return None; - } - if pattern[0] != b'%' || pattern[pattern.len() - 1] != b'%' { - return None; - } - let inner = &pattern[1..pattern.len() - 1]; - if inner.is_empty() { - return None; - } - if inner.iter().any(|b| *b == b'%' || *b == b'_') { - return None; - } - Some(inner) -} diff --git a/src/core/src/liquid_array/byte_view_array/helpers.rs b/src/core/src/liquid_array/byte_view_array/helpers.rs deleted file mode 100644 index c9209a456..000000000 --- a/src/core/src/liquid_array/byte_view_array/helpers.rs +++ /dev/null @@ -1,147 +0,0 @@ -use arrow::array::{ - BooleanArray, BooleanBufferBuilder, UInt16Array, cast::AsArray, types::UInt16Type, -}; -use arrow::buffer::BooleanBuffer; -use datafusion_physical_expr::PhysicalExpr; -use std::sync::Arc; - -use super::LiquidByteViewArray; -use super::operator::{ByteViewExpression, ByteViewOperator}; -use crate::liquid_array::byte_view_array::operator::UnsupportedExpression; -use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::FsstBacking; - -pub(super) fn build_dict_selection( - keys: &UInt16Array, - dict_len: usize, -) -> (Vec, UInt16Array) { - let mut hit_mask = BooleanBufferBuilder::new(dict_len); - hit_mask.advance(dict_len); - for value in keys.iter().flatten() { - hit_mask.set_bit(value as usize, true); - } - let hit_mask = hit_mask.finish(); - let selected_cnt = hit_mask.count_set_bits(); - - let mut key_map = vec![u16::MAX; dict_len]; - let mut selected = Vec::with_capacity(selected_cnt); - let mut remapped: u16 = 0; - for (index, selected_flag) in hit_mask.iter().enumerate() { - if selected_flag { - key_map[index] = remapped; - selected.push(index); - remapped += 1; - } - } - - let new_keys = UInt16Array::from_iter( - keys.iter() - .map(|value| value.map(|value| key_map[value as usize])), - ); - (selected, new_keys) -} - -pub(super) fn filter_inner( - array: &LiquidByteViewArray, - filter: &BooleanBuffer, -) -> LiquidByteViewArray { - // Only filter the dictionary keys, not the offsets! - // Offset views reference unique values in FSST buffer and should remain unchanged - - // Filter the dictionary keys using Arrow's built-in filter functionality - let filter = BooleanArray::new(filter.clone(), None); - let filtered_keys = arrow::compute::filter(&array.dictionary_keys, &filter).unwrap(); - let filtered_keys = filtered_keys.as_primitive::().clone(); - - LiquidByteViewArray { - dictionary_keys: filtered_keys, - prefix_keys: array.prefix_keys.clone(), - fsst_buffer: array.fsst_buffer.clone(), - original_arrow_type: array.original_arrow_type, - shared_prefix: array.shared_prefix.clone(), - string_fingerprints: array.string_fingerprints.clone(), - } -} - -pub(super) fn try_eval_predicate_in_memory( - expr: &Arc, - array: &LiquidByteViewArray, -) -> Option { - let expr = match ByteViewExpression::try_from(expr) { - Ok(expr) => expr, - Err(UnsupportedExpression::Constant(v)) => { - let bool_array = if v { - BooleanBuffer::new_set(array.len()) - } else { - BooleanBuffer::new_unset(array.len()) - }; - return Some(BooleanArray::new(bool_array, array.nulls().cloned())); - } - Err(UnsupportedExpression::Expr) | Err(UnsupportedExpression::Op) => { - return None; - } - }; - let op = expr.op(); - let needle = expr.literal(); - if let ByteViewOperator::SubString(_substring_op) = op - && array.string_fingerprints.as_ref().is_none() - { - return None; - } - Some(array.compare_with(needle, op)) -} - -use std::fmt::Display; - -/// Detailed memory usage of the byte view array -pub struct ByteViewArrayMemoryUsage { - /// Memory usage of the dictionary key - pub dictionary_key: usize, - /// Memory usage of the prefix keys - pub prefix_keys: usize, - /// Memory usage of the raw FSST buffer - pub fsst_buffer: usize, - /// Memory usage of the shared prefix - pub shared_prefix: usize, - /// Memory usage of the string fingerprints - pub string_fingerprints: usize, - /// Memory usage of the struct size - pub struct_size: usize, -} - -impl Display for ByteViewArrayMemoryUsage { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ByteViewArrayMemoryUsage") - .field("dictionary_key", &self.dictionary_key) - .field("prefix_keys", &self.prefix_keys) - .field("fsst_buffer", &self.fsst_buffer) - .field("shared_prefix", &self.shared_prefix) - .field("string_fingerprints", &self.string_fingerprints) - .field("struct_size", &self.struct_size) - .field("total", &self.total()) - .finish() - } -} - -impl ByteViewArrayMemoryUsage { - /// Get the total memory usage of the byte view array - pub fn total(&self) -> usize { - self.dictionary_key - + self.prefix_keys - + self.fsst_buffer - + self.shared_prefix - + self.string_fingerprints - + self.struct_size - } -} - -impl std::ops::AddAssign for ByteViewArrayMemoryUsage { - fn add_assign(&mut self, other: Self) { - self.dictionary_key += other.dictionary_key; - self.prefix_keys += other.prefix_keys; - self.fsst_buffer += other.fsst_buffer; - self.shared_prefix += other.shared_prefix; - self.string_fingerprints += other.string_fingerprints; - self.struct_size += other.struct_size; - } -} diff --git a/src/core/src/liquid_array/byte_view_array/mod.rs b/src/core/src/liquid_array/byte_view_array/mod.rs deleted file mode 100644 index cb06173b2..000000000 --- a/src/core/src/liquid_array/byte_view_array/mod.rs +++ /dev/null @@ -1,340 +0,0 @@ -//! LiquidByteViewArray - -use arrow::array::BooleanArray; -use arrow::array::{ - Array, ArrayRef, BinaryArray, DictionaryArray, StringArray, UInt16Array, types::UInt16Type, -}; -use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer}; -use arrow::compute::cast; -use arrow_schema::DataType; -use std::any::Any; -use std::sync::Arc; - -#[cfg(test)] -use std::cell::Cell; - -use crate::cache::LiquidExpr; -use crate::liquid_array::raw::FsstArray; -use crate::liquid_array::raw::fsst_buffer::{FsstBacking, PrefixKey}; -use crate::liquid_array::{LiquidArray, LiquidDataType, eval_predicate_on_array}; - -mod comparisons; -mod conversions; -mod fingerprint; -mod helpers; -mod operator; -mod serialization; - -#[cfg(test)] -mod tests; - -pub use helpers::ByteViewArrayMemoryUsage; -pub use operator::{ByteViewOperator, Comparison, Equality, SubString}; - -#[cfg(test)] -thread_local! { - static DISK_READ_COUNTER: Cell = const { Cell::new(0)}; - static FULL_DATA_COMPARISON_COUNTER: Cell = const { Cell::new(0)}; -} - -#[cfg(test)] -fn get_disk_read_counter() -> usize { - DISK_READ_COUNTER.with(|counter| counter.get()) -} - -#[cfg(test)] -fn reset_disk_read_counter() { - DISK_READ_COUNTER.with(|counter| counter.set(0)); -} - -/// An array that stores strings using the FSST format with compact offsets: -/// - Dictionary keys with 2-byte keys stored in memory -/// - Compact offsets with variable-size residuals (1, 2, or 4 bytes) stored in memory -/// - Per-value prefix keys (7-byte prefix + len) stored in memory -/// - FSST buffer can be stored in memory or on disk -/// -/// # Initialization -/// -/// The recommended way to create a `LiquidByteViewArray` is using the `from_*_array` constructors -/// which build a compact (offset + prefix key) representation directly from Arrow inputs. -/// -/// ```rust,ignore -/// let liquid_array = LiquidByteViewArray::from_string_array(&input, compressor); -/// ``` -/// -/// Data access flow: -/// 1. Use dictionary key to index into compact offsets buffer -/// 2. Reconstruct actual offset from linear regression (predicted + residual) -/// 3. Use prefix keys for quick comparisons to avoid decompression when possible -/// 4. Decompress bytes from FSST buffer to get the full value when needed -#[derive(Clone)] -pub struct LiquidByteViewArray { - /// Dictionary keys (u16) - one per array element, using Arrow's UInt16Array for zero-copy - pub(super) dictionary_keys: UInt16Array, - /// Per-value prefix keys (prefix7 + len metadata). - pub(super) prefix_keys: Arc<[PrefixKey]>, - /// FSST-compressed buffer (can be in memory or on disk) - pub(super) fsst_buffer: B, - /// Used to convert back to the original arrow type - pub(super) original_arrow_type: ArrowByteType, - /// Shared prefix across all strings in the array - pub(super) shared_prefix: Vec, - /// Optional per-dictionary string fingerprints (32 bins). - pub(super) string_fingerprints: Option>, -} - -#[derive(Clone, Copy, Debug)] -pub(crate) struct ByteViewBuildOptions { - pub(super) arrow_type: ArrowByteType, - pub(super) build_fingerprints: bool, -} - -impl ByteViewBuildOptions { - pub(crate) fn new(arrow_type: ArrowByteType) -> Self { - Self { - arrow_type, - build_fingerprints: false, - } - } - - pub(crate) fn for_data_type(data_type: &DataType, build_fingerprints: bool) -> Self { - Self { - arrow_type: ArrowByteType::from_arrow_type(data_type), - build_fingerprints, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq)] -#[repr(u16)] -pub(crate) enum ArrowByteType { - Utf8 = 0, - Utf8View = 1, - Dict16Binary = 2, - Dict16Utf8 = 3, - Binary = 4, - BinaryView = 5, -} - -impl From for ArrowByteType { - fn from(value: u16) -> Self { - match value { - 0 => ArrowByteType::Utf8, - 1 => ArrowByteType::Utf8View, - 2 => ArrowByteType::Dict16Binary, - 3 => ArrowByteType::Dict16Utf8, - 4 => ArrowByteType::Binary, - 5 => ArrowByteType::BinaryView, - _ => panic!("Invalid arrow byte type: {value}"), - } - } -} - -impl ArrowByteType { - pub fn to_arrow_type(self) -> DataType { - match self { - ArrowByteType::Utf8 => DataType::Utf8, - ArrowByteType::Utf8View => DataType::Utf8View, - ArrowByteType::Dict16Binary => { - DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Binary)) - } - ArrowByteType::Dict16Utf8 => { - DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) - } - ArrowByteType::Binary => DataType::Binary, - ArrowByteType::BinaryView => DataType::BinaryView, - } - } - - pub fn from_arrow_type(ty: &DataType) -> Self { - match ty { - DataType::Utf8 => ArrowByteType::Utf8, - DataType::Utf8View => ArrowByteType::Utf8View, - DataType::Binary => ArrowByteType::Binary, - DataType::BinaryView => ArrowByteType::BinaryView, - DataType::Dictionary(_, _) => { - if ty - == &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Binary)) - { - ArrowByteType::Dict16Binary - } else if ty - == &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) - { - ArrowByteType::Dict16Utf8 - } else { - panic!("Unsupported arrow type: {ty:?}") - } - } - _ => panic!("Unsupported arrow type: {ty:?}"), - } - } -} - -impl std::fmt::Debug for LiquidByteViewArray { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("LiquidByteViewArray") - .field("dictionary_keys", &self.dictionary_keys) - .field("prefix_keys", &self.prefix_keys) - .field("fsst_buffer", &self.fsst_buffer) - .field("original_arrow_type", &self.original_arrow_type) - .field("shared_prefix", &self.shared_prefix) - .field("string_fingerprints", &self.string_fingerprints) - .finish() - } -} - -impl LiquidByteViewArray { - /// Convert to Arrow DictionaryArray - fn to_dict_arrow_inner( - &self, - keys_array: UInt16Array, - values_buffer: Buffer, - offsets_buffer: OffsetBuffer, - ) -> DictionaryArray { - let values = if self.original_arrow_type == ArrowByteType::Utf8 - || self.original_arrow_type == ArrowByteType::Utf8View - || self.original_arrow_type == ArrowByteType::Dict16Utf8 - { - let string_array = - unsafe { StringArray::new_unchecked(offsets_buffer, values_buffer, None) }; - Arc::new(string_array) as ArrayRef - } else { - let binary_array = - unsafe { BinaryArray::new_unchecked(offsets_buffer, values_buffer, None) }; - Arc::new(binary_array) as ArrayRef - }; - - unsafe { DictionaryArray::::new_unchecked(keys_array, values) } - } - - fn should_decompress_keyed(&self) -> bool { - self.dictionary_keys.len() < 2048 || self.dictionary_keys.len() < self.prefix_keys.len() - } - - /// Get the nulls buffer - pub fn nulls(&self) -> Option<&NullBuffer> { - self.dictionary_keys.nulls() - } - - /// Get detailed memory usage of the byte view array - pub fn get_detailed_memory_usage(&self) -> ByteViewArrayMemoryUsage { - let fingerprint_bytes = self - .string_fingerprints - .as_ref() - .map(|fingerprints| fingerprints.len() * std::mem::size_of::()) - .unwrap_or(0); - ByteViewArrayMemoryUsage { - dictionary_key: self.dictionary_keys.get_array_memory_size(), - prefix_keys: self.prefix_keys.len() * std::mem::size_of::(), - fsst_buffer: self.fsst_buffer.get_array_memory_size(), - shared_prefix: self.shared_prefix.len(), - string_fingerprints: fingerprint_bytes, - struct_size: std::mem::size_of::(), - } - } - - /// Get the length of the array - pub fn len(&self) -> usize { - self.dictionary_keys.len() - } - - /// Is the array empty? - pub fn is_empty(&self) -> bool { - self.dictionary_keys.is_empty() - } - - /// Get disk read count for testing - #[cfg(test)] - pub fn get_disk_read_count(&self) -> usize { - get_disk_read_counter() - } - - /// Reset disk read count for testing - #[cfg(test)] - pub fn reset_disk_read_count(&self) { - reset_disk_read_counter() - } -} - -impl LiquidByteViewArray { - /// Convert to Arrow DictionaryArray - pub fn to_dict_arrow(&self) -> DictionaryArray { - if self.should_decompress_keyed() { - self.to_dict_arrow_decompress_keyed() - } else { - self.to_dict_arrow_decompress_all() - } - } - - fn to_dict_arrow_decompress_all(&self) -> DictionaryArray { - let (values_buffer, offsets_buffer) = self.fsst_buffer.to_uncompressed(); - self.to_dict_arrow_inner(self.dictionary_keys.clone(), values_buffer, offsets_buffer) - } - - fn to_dict_arrow_decompress_keyed(&self) -> DictionaryArray { - let (selected, new_keys) = - helpers::build_dict_selection(&self.dictionary_keys, self.prefix_keys.len()); - let (values_buffer, offsets_buffer) = self.fsst_buffer.to_uncompressed_selected(&selected); - self.to_dict_arrow_inner(new_keys, values_buffer, offsets_buffer) - } - - /// Convert to Arrow array with original type - pub fn to_arrow_array(&self) -> ArrayRef { - let dict = self.to_dict_arrow(); - cast(&dict, &self.original_arrow_type.to_arrow_type()).unwrap() - } - - /// Check if the FSST buffer is currently stored on disk - pub fn is_fsst_buffer_on_disk(&self) -> bool { - false - } -} - -impl LiquidArray for LiquidByteViewArray { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.get_detailed_memory_usage().total() - } - - fn len(&self) -> usize { - self.dictionary_keys.len() - } - - #[inline] - fn to_arrow_array(&self) -> ArrayRef { - let dict = self.to_arrow_array(); - Arc::new(dict) - } - - fn to_best_arrow_array(&self) -> ArrayRef { - let dict = self.to_dict_arrow(); - Arc::new(dict) - } - - fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - let filtered = helpers::filter_inner(self, filter); - - helpers::try_eval_predicate_in_memory(expr.physical_expr(), &filtered) - .unwrap_or_else(|| eval_predicate_on_array(filtered.to_arrow_array(), expr)) - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn original_arrow_data_type(&self) -> DataType { - self.original_arrow_type.to_arrow_type() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::ByteViewArray - } - - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let filtered = helpers::filter_inner(self, selection); - filtered.to_arrow_array() - } -} diff --git a/src/core/src/liquid_array/byte_view_array/operator.rs b/src/core/src/liquid_array/byte_view_array/operator.rs deleted file mode 100644 index 1035e4b6e..000000000 --- a/src/core/src/liquid_array/byte_view_array/operator.rs +++ /dev/null @@ -1,174 +0,0 @@ -use std::sync::Arc; - -use datafusion_common::ScalarValue; -use datafusion_expr_common::operator::Operator; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{ - BinaryExpr, DynamicFilterPhysicalExpr, LikeExpr, Literal, -}; - -use crate::utils::get_bytes_needle; - -/// Supported ordering comparisons for byte views. -#[derive(Debug)] -pub enum Comparison { - /// Less-than. - Lt, - /// Greater-than. - Gt, - /// Less-than or equal. - LtEq, - /// Greater-than or equal. - GtEq, -} - -/// Supported equality comparisons for byte views. -#[derive(Debug)] -pub enum Equality { - /// Equal. - Eq, - /// Not equal. - NotEq, -} - -#[derive(Debug, PartialEq, Eq, Copy, Clone)] -/// Supported substring predicate kinds. -pub enum SubString { - /// Contains a substring. - Contains, - /// Does not contain a substring. - NotContains, -} - -/// Supported operators for byte view predicates. -#[derive(Debug)] -pub enum ByteViewOperator { - /// Ordering comparison. - Comparison(Comparison), - /// Equality comparison. - Equality(Equality), - /// Substring predicate. - SubString(SubString), -} - -impl ByteViewOperator { - fn from_like_expr(like: &LikeExpr) -> Result { - match (like.negated(), like.case_insensitive()) { - (false, false) => Ok(ByteViewOperator::SubString(SubString::Contains)), - (true, false) => Ok(ByteViewOperator::SubString(SubString::NotContains)), - _ => Err(UnsupportedOperator), - } - } -} - -pub struct UnsupportedOperator; - -impl TryFrom<&Operator> for ByteViewOperator { - type Error = UnsupportedOperator; - - fn try_from(operator: &Operator) -> Result { - match operator { - Operator::Eq => Ok(ByteViewOperator::Equality(Equality::Eq)), - Operator::NotEq => Ok(ByteViewOperator::Equality(Equality::NotEq)), - Operator::Lt => Ok(ByteViewOperator::Comparison(Comparison::Lt)), - Operator::Gt => Ok(ByteViewOperator::Comparison(Comparison::Gt)), - Operator::LtEq => Ok(ByteViewOperator::Comparison(Comparison::LtEq)), - Operator::GtEq => Ok(ByteViewOperator::Comparison(Comparison::GtEq)), - Operator::LikeMatch => Ok(ByteViewOperator::SubString(SubString::Contains)), - Operator::NotLikeMatch => Ok(ByteViewOperator::SubString(SubString::NotContains)), - _ => Err(UnsupportedOperator), - } - } -} - -impl From<&ByteViewOperator> for Operator { - fn from(byte_view_operator: &ByteViewOperator) -> Self { - match byte_view_operator { - ByteViewOperator::Comparison(comparison) => match comparison { - Comparison::Lt => Operator::Lt, - Comparison::Gt => Operator::Gt, - Comparison::LtEq => Operator::LtEq, - Comparison::GtEq => Operator::GtEq, - }, - ByteViewOperator::Equality(equality) => match equality { - Equality::Eq => Operator::Eq, - Equality::NotEq => Operator::NotEq, - }, - ByteViewOperator::SubString(substring) => match substring { - SubString::Contains => Operator::LikeMatch, - SubString::NotContains => Operator::NotLikeMatch, - }, - } - } -} - -#[derive(Debug)] -pub(super) struct ByteViewExpression { - op: ByteViewOperator, - literal: Vec, -} - -pub(super) enum UnsupportedExpression { - Op, - Expr, - // This is frequently the case with dynamic filters - Constant(bool), -} - -impl From for UnsupportedExpression { - fn from(_op: UnsupportedOperator) -> Self { - UnsupportedExpression::Op - } -} - -impl ByteViewExpression { - pub(super) fn op(&self) -> &ByteViewOperator { - &self.op - } - - pub(super) fn literal(&self) -> &[u8] { - &self.literal - } -} - -impl TryFrom<&Arc> for ByteViewExpression { - type Error = UnsupportedExpression; - fn try_from(expr: &Arc) -> Result { - let expr = if let Some(dynamic_filter) = expr.downcast_ref::() { - dynamic_filter.current().unwrap() - } else { - expr.clone() - }; - - if let Some(literal) = expr.downcast_ref::() - && let ScalarValue::Boolean(Some(v)) = literal.value() - { - return Err(UnsupportedExpression::Constant(*v)); - } - - if let Some(binary_expr) = expr.downcast_ref::() { - if let Some(literal) = binary_expr.right().downcast_ref::() { - let op = binary_expr.op(); - let byte_view_operator = ByteViewOperator::try_from(op)?; - let literal = - get_bytes_needle(literal.value()).ok_or(UnsupportedExpression::Expr)?; - return Ok(ByteViewExpression { - op: byte_view_operator, - literal, - }); - } - } - // Handle like expressions - else if let Some(like_expr) = expr.downcast_ref::() - && let Some(literal) = like_expr.pattern().downcast_ref::() - { - let byte_view_operator = ByteViewOperator::from_like_expr(like_expr)?; - let literal = get_bytes_needle(literal.value()).ok_or(UnsupportedExpression::Expr)?; - return Ok(ByteViewExpression { - op: byte_view_operator, - literal, - }); - } - Err(UnsupportedExpression::Expr) - } -} diff --git a/src/core/src/liquid_array/byte_view_array/serialization.rs b/src/core/src/liquid_array/byte_view_array/serialization.rs deleted file mode 100644 index 297dee2e5..000000000 --- a/src/core/src/liquid_array/byte_view_array/serialization.rs +++ /dev/null @@ -1,326 +0,0 @@ -use arrow::array::types::UInt16Type; -use bytes::Bytes; -use fsst::Compressor; -use std::sync::Arc; - -use super::{ArrowByteType, LiquidByteViewArray}; -use crate::liquid_array::LiquidDataType; -use crate::liquid_array::ipc::LiquidIPCHeader; -use crate::liquid_array::raw::BitPackedArray; -use crate::liquid_array::raw::fsst_buffer::{ - FsstArray, PrefixKey, RawFsstBuffer, decode_compact_offsets, empty_compact_offsets, -}; - -// Header for LiquidByteViewArray serialization -#[repr(C)] -pub(super) struct ByteViewArrayHeader { - pub(super) keys_size: u32, - pub(super) compact_offsets_size: u32, - pub(super) shared_prefix_size: u32, - pub(super) fsst_raw_size: u32, - pub(super) fingerprint_size: u32, -} - -impl ByteViewArrayHeader { - pub(super) const fn size() -> usize { - const _: () = - assert!(std::mem::size_of::() == ByteViewArrayHeader::size()); - 20 - } - - pub(super) fn to_bytes(&self) -> [u8; Self::size()] { - let mut bytes = [0u8; Self::size()]; - bytes[0..4].copy_from_slice(&self.keys_size.to_le_bytes()); - bytes[4..8].copy_from_slice(&self.compact_offsets_size.to_le_bytes()); - bytes[8..12].copy_from_slice(&self.shared_prefix_size.to_le_bytes()); - bytes[12..16].copy_from_slice(&self.fsst_raw_size.to_le_bytes()); - bytes[16..20].copy_from_slice(&self.fingerprint_size.to_le_bytes()); - bytes - } - - pub(super) fn from_bytes(bytes: &[u8]) -> Self { - if bytes.len() < Self::size() { - panic!( - "value too small for ByteViewArrayHeader, expected at least {} bytes, got {}", - Self::size(), - bytes.len() - ); - } - let keys_size = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); - let compact_offsets_size = u32::from_le_bytes(bytes[4..8].try_into().unwrap()); - let shared_prefix_size = u32::from_le_bytes(bytes[8..12].try_into().unwrap()); - let fsst_raw_size = u32::from_le_bytes(bytes[12..16].try_into().unwrap()); - let fingerprint_size = u32::from_le_bytes(bytes[16..20].try_into().unwrap()); - Self { - keys_size, - compact_offsets_size, - shared_prefix_size, - fsst_raw_size, - fingerprint_size, - } - } -} - -pub(super) fn align_up_8(len: usize) -> usize { - (len + 7) & !7 -} - -fn decode_prefix_keys(bytes: &[u8]) -> Arc<[PrefixKey]> { - let entry_size = std::mem::size_of::(); - if !bytes.len().is_multiple_of(entry_size) { - panic!("Invalid prefix keys size"); - } - if bytes.is_empty() { - return Arc::<[PrefixKey]>::from([]); - } - let mut keys = Vec::with_capacity(bytes.len() / entry_size); - for chunk in bytes.chunks_exact(entry_size) { - let mut prefix7 = [0u8; 7]; - prefix7.copy_from_slice(&chunk[..7]); - let len = chunk[7]; - keys.push(PrefixKey::from_parts(prefix7, len)); - } - keys.into() -} - -impl LiquidByteViewArray { - /* - Serialized LiquidByteViewArray Memory Layout: - - +--------------------------------------------------+ - | LiquidIPCHeader (16 bytes) | - +--------------------------------------------------+ - | ByteViewArrayHeader (20 bytes) | // keys_size, compact_offsets_size, shared_prefix_size, fsst_size, fingerprint_size - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | RawFsstBuffer bytes | - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | BitPackedArray Data (dictionary_keys) | - +--------------------------------------------------+ - | [BitPackedArray Header & Values] | - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | Compact offsets bytes (header + residuals) | - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | Prefix keys bytes (prefix7 + len) | - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | Shared prefix bytes | - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | - +--------------------------------------------------+ - | Optional string fingerprints (u32 per entry) | - +--------------------------------------------------+ - */ - pub(crate) fn to_bytes_inner(&self) -> Vec { - let header_size = LiquidIPCHeader::size() + ByteViewArrayHeader::size(); - let mut result = Vec::with_capacity(header_size + 1024); - result.resize(header_size, 0); - - // A) Align and serialize RawFsstBuffer first (near the start) - while !result.len().is_multiple_of(8) { - result.push(0); - } - let fsst_start = result.len(); - let fsst_raw_bytes = self.fsst_buffer.raw_to_bytes(); - result.extend_from_slice(&fsst_raw_bytes); - let fsst_raw_size = result.len() - fsst_start; - - // B) Alignment before keys - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // C) Serialize dictionary keys - let keys_start = result.len(); - { - use std::num::NonZero; - let bit_packed = BitPackedArray::::from_primitive( - self.dictionary_keys.clone(), - NonZero::new(16).unwrap(), - ); - bit_packed.to_bytes(&mut result); - } - let keys_size = result.len() - keys_start; - - // D) Alignment before compact offsets - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // E) Serialize compact offsets (header + residuals) - let offsets_start = result.len(); - self.fsst_buffer.write_compact_offsets(&mut result); - let compact_offsets_size = result.len() - offsets_start; - - // F) Alignment before prefix keys - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // G) Serialize prefix keys (prefix7 + len) - for prefix in self.prefix_keys.iter() { - result.extend_from_slice(prefix.prefix7()); - result.push(prefix.len_byte()); - } - - // H) Alignment before shared prefix - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // I) Serialize shared prefix - let prefix_start = result.len(); - result.extend_from_slice(&self.shared_prefix); - let shared_prefix_size = result.len() - prefix_start; - - // J) Alignment before fingerprints - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // K) Serialize string fingerprints (u32 per entry) - if let Some(fingerprints) = self.string_fingerprints.as_ref() { - for &fingerprint in fingerprints.iter() { - result.extend_from_slice(&fingerprint.to_le_bytes()); - } - } - - // Prepare headers - let ipc = LiquidIPCHeader::new( - LiquidDataType::ByteViewArray as u16, - self.original_arrow_type as u16, - ); - let view_header = ByteViewArrayHeader { - keys_size: keys_size as u32, - compact_offsets_size: compact_offsets_size as u32, - shared_prefix_size: shared_prefix_size as u32, - fsst_raw_size: fsst_raw_size as u32, - fingerprint_size: (self - .string_fingerprints - .as_ref() - .map(|fingerprints| fingerprints.len()) - .unwrap_or(0) - * std::mem::size_of::()) as u32, - }; - - // Write headers into reserved space at start - let header_slice = &mut result[0..header_size]; - header_slice[0..LiquidIPCHeader::size()].copy_from_slice(&ipc.to_bytes()); - header_slice[LiquidIPCHeader::size()..header_size].copy_from_slice(&view_header.to_bytes()); - - result - } - - /// Deserialize a LiquidByteViewArray from bytes. - pub fn from_bytes(bytes: Bytes, compressor: Arc) -> LiquidByteViewArray { - // 0) Read IPC header and our view header - let ipc = LiquidIPCHeader::from_bytes(&bytes); - let original_arrow_type = ArrowByteType::from(ipc.physical_type_id); - let header_size = LiquidIPCHeader::size() + ByteViewArrayHeader::size(); - let view_header = - ByteViewArrayHeader::from_bytes(&bytes[LiquidIPCHeader::size()..header_size]); - - let mut cursor = header_size; - - // A) Align and read FSST raw buffer first - cursor = align_up_8(cursor); - let fsst_end = cursor + view_header.fsst_raw_size as usize; - if fsst_end > bytes.len() { - panic!("FSST raw buffer extends beyond input buffer"); - } - let fsst_raw = bytes.slice(cursor..fsst_end); - let raw_buffer = RawFsstBuffer::from_bytes(fsst_raw); - cursor = fsst_end; - - // B) Align and read keys - cursor = align_up_8(cursor); - let keys_end = cursor + view_header.keys_size as usize; - if keys_end > bytes.len() { - panic!("Keys data extends beyond input buffer"); - } - let keys_data = bytes.slice(cursor..keys_end); - let bit_packed = BitPackedArray::::from_bytes(keys_data); - let dictionary_keys = bit_packed.to_primitive(); - cursor = keys_end; - - // C) Align and read compact offsets - cursor = align_up_8(cursor); - let offsets_end = cursor + view_header.compact_offsets_size as usize; - if offsets_end > bytes.len() { - panic!("Compact offsets data extends beyond input buffer"); - } - - // Deserialize compact offsets. - let compact_offsets = if view_header.compact_offsets_size > 0 { - let chunk = bytes.slice(cursor..offsets_end); - decode_compact_offsets(chunk.as_ref()) - } else { - empty_compact_offsets() - }; - cursor = offsets_end; - - // D) Align and read prefix keys - cursor = align_up_8(cursor); - let prefix_count = compact_offsets.len().saturating_sub(1); - let prefix_keys_size = prefix_count * std::mem::size_of::(); - let prefix_keys_end = cursor + prefix_keys_size; - if prefix_keys_end > bytes.len() { - panic!("Prefix keys data extends beyond input buffer"); - } - let prefix_keys = if prefix_keys_size > 0 { - decode_prefix_keys(&bytes[cursor..prefix_keys_end]) - } else { - Arc::<[PrefixKey]>::from([]) - }; - cursor = prefix_keys_end; - - // E) Align and read shared prefix - cursor = align_up_8(cursor); - let prefix_end = cursor + view_header.shared_prefix_size as usize; - if prefix_end > bytes.len() { - panic!("Shared prefix data extends beyond input buffer"); - } - let shared_prefix = bytes[cursor..prefix_end].to_vec(); - cursor = prefix_end; - - // F) String fingerprints - cursor = align_up_8(cursor); - let fingerprint_end = cursor + view_header.fingerprint_size as usize; - if fingerprint_end > bytes.len() { - panic!("Fingerprint data extends beyond input buffer"); - } - let string_fingerprints = if view_header.fingerprint_size == 0 { - None - } else { - if !(view_header.fingerprint_size as usize).is_multiple_of(std::mem::size_of::()) { - panic!("Invalid fingerprint data size"); - } - let expected = prefix_count * std::mem::size_of::(); - if view_header.fingerprint_size as usize != expected { - panic!("Fingerprint data size does not match dictionary size"); - } - let mut fingerprints = Vec::with_capacity(view_header.fingerprint_size as usize / 4); - for chunk in bytes[cursor..fingerprint_end].as_chunks::<4>().0 { - fingerprints.push(u32::from_le_bytes(*chunk)); - } - Some(Arc::from(fingerprints.into_boxed_slice())) - }; - - LiquidByteViewArray { - dictionary_keys, - prefix_keys, - fsst_buffer: FsstArray::new(Arc::new(raw_buffer), compact_offsets, compressor), - original_arrow_type, - shared_prefix, - string_fingerprints, - } - } -} diff --git a/src/core/src/liquid_array/byte_view_array/tests.rs b/src/core/src/liquid_array/byte_view_array/tests.rs deleted file mode 100644 index 507912784..000000000 --- a/src/core/src/liquid_array/byte_view_array/tests.rs +++ /dev/null @@ -1,857 +0,0 @@ -use super::*; -use arrow::array::{ - Array, ArrayRef, BooleanArray, DictionaryArray, StringArray, UInt16Array, cast::AsArray, - types::UInt16Type, -}; -use arrow::buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; -use arrow_schema::DataType; -use rand::{RngExt as _, SeedableRng}; -use std::sync::Arc; - -use crate::cache::transcode_liquid_inner_with_hint; -use crate::cache::{CacheExpression, LiquidCompressorStates}; -use crate::liquid_array::byte_view_array::operator::{ - ByteViewOperator, Comparison, Equality, SubString, -}; -use crate::liquid_array::raw::fsst_buffer::{FsstArray, PrefixKey}; -use crate::liquid_array::{LiquidArray, LiquidDataType}; - -#[test] -fn test_dictionary_view_structure() { - // Test PrefixKey structure - let prefix_key = PrefixKey::from_parts([1, 2, 3, 4, 5, 6, 7], 7); - assert_eq!(prefix_key.prefix7(), &[1, 2, 3, 4, 5, 6, 7]); - assert_eq!(prefix_key.len_byte(), 7); - - // Test UInt16Array creation (dictionary keys are now stored directly in UInt16Array) - let keys = UInt16Array::from(vec![42, 100, 255]); - assert_eq!(keys.value(0), 42); - assert_eq!(keys.value(1), 100); - assert_eq!(keys.value(2), 255); -} - -#[test] -fn test_original_arrow_data_type_returns_utf8() { - let input = StringArray::from(vec!["foo", "bar"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let array = LiquidByteViewArray::::from_string_array(&input, compressor); - assert_eq!(array.original_arrow_data_type(), DataType::Utf8); -} - -#[test] -fn test_ipc_roundtrip_preserves_string_fingerprints() { - let input = StringArray::from(vec!["alpha", "beta", "alphabet"]); - let array: ArrayRef = Arc::new(input); - let state = LiquidCompressorStates::new(); - let liquid = transcode_liquid_inner_with_hint( - &array, - &state, - Some(&CacheExpression::substring_search()), - ) - .expect("transcode should succeed"); - let byte_view = liquid - .as_any() - .downcast_ref::>() - .expect("expected byte view array"); - assert!(byte_view.string_fingerprints.is_some()); - - let bytes = byte_view.to_bytes(); - let decoded = LiquidByteViewArray::::from_bytes( - bytes.into(), - byte_view.fsst_buffer.compressor_arc(), - ); - assert!(decoded.string_fingerprints.is_some()); - assert_eq!( - byte_view.string_fingerprints.as_ref().unwrap().as_ref(), - decoded.string_fingerprints.as_ref().unwrap().as_ref() - ); -} - -#[test] -fn test_ipc_roundtrip_sliced_dictionary_nulls() { - let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); - let keys = UInt16Array::from(vec![ - Some(0u16), - None, - Some(2), - Some(1), - None, - Some(3), - Some(0), - Some(2), - Some(1), - ]); - let dict = DictionaryArray::::new(keys, values); - - // Slice to create a non-zero offset (and therefore a non-zero null bitmap bit offset). - let sliced = dict.slice(1, 7); - - let compressor = LiquidByteViewArray::::train_compressor( - sliced.values().as_string::().iter(), - ); - let original = unsafe { - LiquidByteViewArray::::from_unique_dict_array(&sliced, compressor.clone()) - }; - - let before = original.to_arrow_array(); - let bytes = original.to_bytes(); - let decoded = LiquidByteViewArray::::from_bytes(bytes.into(), compressor); - let after = decoded.to_arrow_array(); - - assert_eq!(before.as_ref(), after.as_ref()); -} - -#[test] -fn test_prefix_extraction() { - let input = StringArray::from(vec!["hello", "world", "test"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // With no shared prefix, the prefix keys should be the original strings (truncated to 7 bytes) - assert_eq!(liquid_array.shared_prefix, Vec::::new()); - assert_eq!(liquid_array.prefix_keys[0].prefix7(), b"hello\0\0"); - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"world\0\0"); - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"test\0\0\0"); -} - -#[test] -fn test_shared_prefix_functionality() { - // Test case with shared prefix - let input = StringArray::from(vec![ - "hello_world", - "hello_rust", - "hello_test", - "hello_code", - ]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Should extract "hello_" as shared prefix - assert_eq!(liquid_array.shared_prefix, b"hello_"); - - // Offset view prefixes (7 bytes) and lengths - assert_eq!(liquid_array.prefix_keys[0].prefix7(), b"world\0\0"); - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"rust\0\0\0"); - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"test\0\0\0"); - assert_eq!(liquid_array.prefix_keys[3].prefix7(), b"code\0\0\0"); - - // Test roundtrip - should reconstruct original strings correctly - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test comparison with shared prefix optimization - let result = liquid_array.compare_equals(b"hello_rust"); - let expected = BooleanArray::from(vec![false, true, false, false]); - assert_eq!(result, expected); - - // Test comparison that doesn't match shared prefix - let result = liquid_array.compare_equals(b"goodbye_world"); - let expected = BooleanArray::from(vec![false, false, false, false]); - assert_eq!(result, expected); - - // Test partial shared prefix match - let result = liquid_array.compare_equals(b"hello_"); - let expected = BooleanArray::from(vec![false, false, false, false]); - assert_eq!(result, expected); -} - -#[test] -fn test_shared_prefix_with_short_strings() { - // Test case: short strings that fit entirely in the 6-byte prefix - let input = StringArray::from(vec!["abc", "abcde", "abcdef", "abcdefg"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Should extract "abc" as shared prefix - assert_eq!(liquid_array.shared_prefix, b"abc"); - - // Offset view prefixes should be the remaining parts after shared prefix (7 bytes) - assert_eq!(liquid_array.prefix_keys[0].prefix7(), &[0u8; 7]); // empty after "abc" - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"de\0\0\0\0\0"); // "de" after "abc" - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"def\0\0\0\0"); // "def" after "abc" - assert_eq!(liquid_array.prefix_keys[3].prefix7(), b"defg\0\0\0"); // "defg" after "abc" - - // Test roundtrip - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test equality comparisons with short strings - let result = liquid_array.compare_equals(b"abc"); - let expected = BooleanArray::from(vec![true, false, false, false]); - assert_eq!(result, expected); - - let result = liquid_array.compare_equals(b"abcde"); - let expected = BooleanArray::from(vec![false, true, false, false]); - assert_eq!(result, expected); - - // Test ordering comparisons that can be resolved by shared prefix - let result = liquid_array.compare_with(b"ab", &ByteViewOperator::Comparison(Comparison::Gt)); - let expected = BooleanArray::from(vec![true, true, true, true]); // All start with "abc" > "ab" - assert_eq!(result, expected); - - let result = liquid_array.compare_with(b"abcd", &ByteViewOperator::Comparison(Comparison::Lt)); - let expected = BooleanArray::from(vec![true, false, false, false]); // Only "abc" < "abcd" - assert_eq!(result, expected); -} - -#[test] -fn test_shared_prefix_contains_complete_strings() { - // Test case: shared prefix completely contains some strings - let input = StringArray::from(vec!["data", "database", "data_entry", "data_", "datatype"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Should extract "data" as shared prefix - assert_eq!(liquid_array.shared_prefix, b"data"); - - // Offset view prefixes should be the remaining parts (7 bytes) - assert_eq!(liquid_array.prefix_keys[0].prefix7(), &[0u8; 7]); // "data" - empty remainder - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"base\0\0\0"); // "database" - "base" remainder - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"_entry\0"); // "data_entry" - "_entry" remainder - assert_eq!(liquid_array.prefix_keys[3].prefix7(), b"_\0\0\0\0\0\0"); // "data_" - "_" remainder - assert_eq!(liquid_array.prefix_keys[4].prefix7(), b"type\0\0\0"); // "datatype" - "type" remainder - - // Test roundtrip - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test equality with exact shared prefix - let result = liquid_array.compare_equals(b"data"); - let expected = BooleanArray::from(vec![true, false, false, false, false]); - assert_eq!(result, expected); - - // Test comparisons where shared prefix helps - let result = liquid_array.compare_with(b"dat", &ByteViewOperator::Comparison(Comparison::Gt)); - let expected = BooleanArray::from(vec![true, true, true, true, true]); // All > "dat" - assert_eq!(result, expected); - - let result = liquid_array.compare_with(b"datab", &ByteViewOperator::Comparison(Comparison::Lt)); - let expected = BooleanArray::from(vec![true, false, true, true, false]); // "data", "data_entry", and "data_" < "datab" - assert_eq!(result, expected); - - // Test comparison with needle shorter than shared prefix - let result = liquid_array.compare_with(b"da", &ByteViewOperator::Comparison(Comparison::Gt)); - let expected = BooleanArray::from(vec![true, true, true, true, true]); // All > "da" - assert_eq!(result, expected); - - // Test comparison with needle equal to shared prefix - let result = - liquid_array.compare_with(b"data", &ByteViewOperator::Comparison(Comparison::GtEq)); - let expected = BooleanArray::from(vec![true, true, true, true, true]); // All >= "data" - assert_eq!(result, expected); - - let result = liquid_array.compare_with(b"data", &ByteViewOperator::Comparison(Comparison::Gt)); - let expected = BooleanArray::from(vec![false, true, true, true, true]); // All except exact "data" > "data" - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_large_value_no_panic() { - // Exercise the ordering-compare slow path which must resize its scratch buffer - // for large values instead of panicking inside `fsst::Decompressor::decompress_into`. - let big = "aaaaaaa".to_string() + &"b".repeat(2 * 1024 * 1024 + 128); - let input = StringArray::from(vec![big.as_str()]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = liquid_array.compare_with( - big.as_bytes(), - &ByteViewOperator::Comparison(Comparison::LtEq), - ); - assert_eq!(result.len(), 1); - assert!(result.value(0)); -} - -#[test] -fn test_shared_prefix_corner_case() { - let input = StringArray::from(vec!["data", "database", "data_entry", "data_", "datatype"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - let result = - liquid_array.compare_with(b"data", &ByteViewOperator::Comparison(Comparison::GtEq)); - let expected = BooleanArray::from(vec![true, true, true, true, true]); // All >= "data" - assert_eq!(result, expected); -} - -#[test] -fn test_shared_prefix_edge_cases() { - // Test case 1: All strings are the same (full shared prefix) - let input = StringArray::from(vec!["identical", "identical", "identical"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - assert_eq!(liquid_array.shared_prefix, b"identical"); - // All prefix keys should be empty - for i in 0..liquid_array.prefix_keys.len() { - assert_eq!(liquid_array.prefix_keys[i].prefix7(), &[0u8; 7]); - } - - // Test roundtrip - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test case 2: One string is a prefix of others - let input = StringArray::from(vec!["hello", "hello_world", "hello_test"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - assert_eq!(liquid_array.shared_prefix, b"hello"); - assert_eq!(liquid_array.prefix_keys[0].prefix7(), &[0u8; 7]); // empty after "hello" - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"_world\0"); - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"_test\0\0"); - - // Test roundtrip - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test case 3: Empty string in array (should limit shared prefix) - let input = StringArray::from(vec!["", "hello", "hello_world"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - assert_eq!(liquid_array.shared_prefix, Vec::::new()); // empty shared prefix - assert_eq!(liquid_array.prefix_keys[0].prefix7(), &[0u8; 7]); - assert_eq!(liquid_array.prefix_keys[1].prefix7(), b"hello\0\0"); - assert_eq!(liquid_array.prefix_keys[2].prefix7(), b"hello_w"); // "hello_world" truncated to 7 bytes - - // Test roundtrip - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); -} - -#[test] -fn test_memory_layout() { - let input = StringArray::from(vec!["hello", "world", "test"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Verify memory layout components - assert_eq!(liquid_array.dictionary_keys.len(), 3); - assert_eq!(liquid_array.fsst_buffer.offsets_len(), 4); - assert!(liquid_array.nulls().is_none()); - let _first = liquid_array.fsst_buffer.get_compressed_slice(0); -} - -fn check_filter_result(input: &StringArray, filter: BooleanBuffer) { - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(input, compressor); - let output = liquid_array.filter(&filter); - let expected = { - let selection = BooleanArray::new(filter.clone(), None); - let arrow_filtered = arrow::compute::filter(&input, &selection).unwrap(); - arrow_filtered.as_string::().clone() - }; - assert_eq!(output.as_ref(), &expected); -} - -#[test] -fn test_filter_functionality() { - let input = StringArray::from(vec![ - Some("hello"), - Some("test"), - None, - Some("test"), - None, - Some("test"), - Some("rust"), - ]); - let mut seeded_rng = rand::rngs::StdRng::seed_from_u64(42); - for _i in 0..100 { - let filter = - BooleanBuffer::from_iter((0..input.len()).map(|_| seeded_rng.random::())); - check_filter_result(&input, filter); - } -} - -#[test] -fn test_memory_efficiency() { - let input = StringArray::from(vec!["hello", "world", "hello", "world", "hello"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Verify that dictionary views store unique values efficiently - assert_eq!(liquid_array.dictionary_keys.len(), 5); - - // Verify that FSST buffer contains unique values - let dict = liquid_array.to_dict_arrow(); - assert_eq!(dict.values().len(), 2); // Only "hello" and "world" -} - -#[test] -fn test_to_best_arrow_array() { - let input = StringArray::from(vec!["hello", "world", "test"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let best_array = liquid_array.to_best_arrow_array(); - let dict_array = best_array.as_dictionary::(); - - // Should return dictionary array as the best encoding - assert_eq!(dict_array.len(), 3); - assert_eq!(dict_array.values().len(), 3); // Three unique values -} - -#[test] -fn test_data_type() { - let input = StringArray::from(vec!["hello", "world"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Just verify we can get the data type without errors - let data_type = liquid_array.data_type(); - assert!(matches!(data_type, LiquidDataType::ByteViewArray)); -} - -#[test] -fn test_compare_with_prefix_optimization_fast_path() { - // Test case 1: Prefix comparison can decide most results without decompression - // Uses strings with distinct prefixes to test the fast path - let input = StringArray::from(vec![ - "apple123", // prefix: "apple\0" - "banana456", // prefix: "banana" - "cherry789", // prefix: "cherry" - "apple999", // prefix: "apple\0" (same as first) - "zebra000", // prefix: "zebra\0" - ]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Test Lt with needle "car" (prefix: "car\0\0\0") - // Expected: "apple123" < "car" => true, "banana456" < "car" => true, others false - let result = liquid_array.compare_with_inner(b"car", &Comparison::Lt); - let expected = BooleanArray::from(vec![true, true, false, true, false]); - assert_eq!(result, expected); - - // Test Gt with needle "dog" (prefix: "dog\0\0\0") - // Expected: only "zebra000" > "dog" => true - let result = liquid_array.compare_with_inner(b"dog", &Comparison::Gt); - let expected = BooleanArray::from(vec![false, false, false, false, true]); - assert_eq!(result, expected); - - // Test GtEq with needle "apple" (prefix: "apple\0") - // Expected: all except "apple123" and "apple999" need decompression, others by prefix - let result = liquid_array.compare_with_inner(b"apple", &Comparison::GtEq); - let expected = BooleanArray::from(vec![true, true, true, true, true]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_prefix_optimization_decompression_path() { - // Test case 2: Cases where prefix comparison is inconclusive and requires decompression - // Uses strings with identical prefixes but different suffixes - let input = StringArray::from(vec![ - "prefix_aaa", // prefix: "prefix" - "prefix_bbb", // prefix: "prefix" (same prefix) - "prefix_ccc", // prefix: "prefix" (same prefix) - "prefix_abc", // prefix: "prefix" (same prefix) - "different", // prefix: "differ" (different prefix) - ]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Test Lt with needle "prefix_b" - this will require decompression for prefix matches - // Expected: "prefix_aaa" < "prefix_b" => true, "prefix_bbb" < "prefix_b" => false, etc. - let result = liquid_array.compare_with_inner(b"prefix_b", &Comparison::Lt); - let expected = BooleanArray::from(vec![true, false, false, true, true]); - assert_eq!(result, expected); - - // Test LtEq with needle "prefix_bbb" - exact match case with decompression - let result = liquid_array.compare_with_inner(b"prefix_bbb", &Comparison::LtEq); - let expected = BooleanArray::from(vec![true, true, false, true, true]); - assert_eq!(result, expected); - - // Test Gt with needle "prefix_abc" - requires decompression for prefix matches - let result = liquid_array.compare_with_inner(b"prefix_abc", &Comparison::Gt); - let expected = BooleanArray::from(vec![false, true, true, false, false]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_prefix_optimization_edge_cases_and_nulls() { - // Test case 3: Edge cases including nulls, empty strings, and boundary conditions - let input = StringArray::from(vec![ - Some(""), // Empty string (prefix: "\0\0\0\0\0\0") - None, // Null value - Some("a"), // Single character (prefix: "a\0\0\0\0\0") - Some("abcdef"), // Exactly 6 chars (prefix: "abcdef") - Some("abcdefghij"), // Longer than 6 chars (prefix: "abcdef") - Some("abcdeg"), // Differs at position 5 (prefix: "abcdeg") - ]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Test Lt with empty string needle - should test null handling - let result = liquid_array.compare_with_inner(b"", &Comparison::Lt); - let expected = BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(false), - Some(false), - ]); - assert_eq!(result, expected); - - // Test Gt with needle "abcdef" - tests exact prefix match requiring decompression - let result = liquid_array.compare_with_inner(b"abcdef", &Comparison::Gt); - let expected = BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(true), - Some(true), - ]); - assert_eq!(result, expected); - - // Test LtEq with needle "b" - tests single character comparisons - let result = liquid_array.compare_with_inner(b"b", &Comparison::LtEq); - let expected = BooleanArray::from(vec![ - Some(true), - None, - Some(true), - Some(true), - Some(true), - Some(true), - ]); - assert_eq!(result, expected); - - // Test GtEq with needle "abcdeg" - tests decompression when prefix exactly matches needle prefix - let result = liquid_array.compare_with_inner(b"abcdeg", &Comparison::GtEq); - // b"" >= b"abcdeg" => false - // null => null - // b"a" >= b"abcdeg" => false - // b"abcdef" >= b"abcdeg" => false (because b"abcdef" < b"abcdeg" since 'f' < 'g') - // b"abcdefghij" >= b"abcdeg" => false (because b"abcdefghij" < b"abcdeg" since 'f' < 'g') - // b"abcdeg" >= b"abcdeg" => true (exact match) - let expected = BooleanArray::from(vec![ - Some(false), - None, - Some(false), - Some(false), - Some(false), - Some(true), - ]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_prefix_empty_suffix() { - let input = StringArray::from(vec!["x", "x1"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = liquid_array.compare_with_inner(b"x", &Comparison::LtEq); - let expected = BooleanArray::from(vec![true, false]); - assert_eq!(result, expected); - - let result = liquid_array.compare_with_inner(b"x", &Comparison::Gt); - let expected = BooleanArray::from(vec![false, true]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_prefix_optimization_utf8_and_binary() { - // Test case 4: UTF-8 encoded strings and binary data comparisons - // This demonstrates the advantage of byte-level comparison - let input = StringArray::from(vec![ - "café", // UTF-8: [99, 97, 102, 195, 169] - "naïve", // UTF-8: [110, 97, 195, 175, 118, 101] - "résumé", // UTF-8: [114, 195, 169, 115, 117, 109, 195, 169] - "hello", // ASCII: [104, 101, 108, 108, 111] - "世界", // UTF-8: [228, 184, 150, 231, 149, 140] - ]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Test Lt with UTF-8 needle "naïve" (UTF-8: [110, 97, 195, 175, 118, 101]) - // Expected: "café" < "naïve" => true (99 < 110), "hello" < "naïve" => true, others false - let naive_bytes = "naïve".as_bytes(); // [110, 97, 195, 175, 118, 101] - let result = liquid_array.compare_with_inner(naive_bytes, &Comparison::Lt); - let expected = BooleanArray::from(vec![true, false, false, true, false]); - assert_eq!(result, expected); - - // Test Gt with UTF-8 needle "café" (UTF-8: [99, 97, 102, 195, 169]) - // Expected: strings with first byte > 99 should be true - let cafe_bytes = "café".as_bytes(); // [99, 97, 102, 195, 169] - let result = liquid_array.compare_with_inner(cafe_bytes, &Comparison::Gt); - let expected = BooleanArray::from(vec![false, true, true, true, true]); - assert_eq!(result, expected); - - // Test LtEq with Chinese characters "世界" (UTF-8: [228, 184, 150, 231, 149, 140]) - // Expected: only strings with first byte <= 228 should be true, but since 228 is quite high, - // most Latin characters will be true - let world_bytes = "世界".as_bytes(); // [228, 184, 150, 231, 149, 140] - let result = liquid_array.compare_with_inner(world_bytes, &Comparison::LtEq); - let expected = BooleanArray::from(vec![true, true, true, true, true]); - assert_eq!(result, expected); - - // Test exact equality with "résumé" using GtEq and LtEq to verify byte-level precision - let resume_bytes = "résumé".as_bytes(); // [114, 195, 169, 115, 117, 109, 195, 169] - let gte_result = liquid_array.compare_with_inner(resume_bytes, &Comparison::GtEq); - let lte_result = liquid_array.compare_with_inner(resume_bytes, &Comparison::LtEq); - - // Check GtEq and LtEq results separately - // GtEq: "café"(99) >= "résumé"(114) => false, "naïve"(110) >= "résumé"(114) => false, - // "résumé"(114) >= "résumé"(114) => true, "hello"(104) >= "résumé"(114) => false, - // "世界"(228) >= "résumé"(114) => true - let gte_expected = BooleanArray::from(vec![false, false, true, false, true]); - // LtEq: all strings with first byte <= 114 should be true, "世界"(228) should be false - let lte_expected = BooleanArray::from(vec![true, true, true, true, false]); - assert_eq!(gte_result, gte_expected); - assert_eq!(lte_result, lte_expected); -} - -#[test] -fn test_compare_equals_long_string_len_byte_255() { - let common = "prefix_"; - let long_len = 260; - let suffix_len = long_len - common.len(); - let long_a = format!("{}{}", common, "a".repeat(suffix_len)); - let long_b = format!("{}{}", common, "b".repeat(suffix_len)); - - let input = StringArray::from(vec![long_a.as_str(), long_b.as_str(), "z"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = liquid_array.compare_equals(long_a.as_bytes()); - let expected = BooleanArray::from(vec![true, false, false]); - assert_eq!(result, expected); - - let shorter = format!("{}{}", common, "a".repeat(200)); - let result = liquid_array.compare_equals(shorter.as_bytes()); - let expected = BooleanArray::from(vec![false, false, false]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_not_equals_preserves_nulls() { - let input = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("alpha")]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = liquid_array.compare_with(b"alpha", &ByteViewOperator::Equality(Equality::NotEq)); - let expected = BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_equals_ignores_raw_key_value_in_null_slot() { - let values: ArrayRef = Arc::new(StringArray::from(vec!["alpha", "beta"])); - let keys = UInt16Array::new( - ScalarBuffer::from(vec![0u16, u16::MAX, 1u16]), - Some(NullBuffer::from(BooleanBuffer::from(vec![ - true, false, true, - ]))), - ); - let dict = DictionaryArray::::new(keys, values); - - let compressor = - LiquidByteViewArray::::train_compressor(dict.values().as_string::().iter()); - let liquid_array = - unsafe { LiquidByteViewArray::::from_unique_dict_array(&dict, compressor) }; - - let result = liquid_array.compare_equals(b"alpha"); - let expected = BooleanArray::from(vec![Some(true), None, Some(false)]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_shared_prefix_shorter_needle_lt() { - let input = StringArray::from(vec!["hello_world", "hello_rust"]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = liquid_array.compare_with(b"hell", &ByteViewOperator::Comparison(Comparison::Lt)); - let expected = BooleanArray::from(vec![false, false]); - assert_eq!(result, expected); - - let result = - liquid_array.compare_with(b"hell", &ByteViewOperator::Comparison(Comparison::LtEq)); - let expected = BooleanArray::from(vec![false, false]); - assert_eq!(result, expected); -} - -#[test] -fn test_compare_with_like_fallback() { - let input = StringArray::from(vec![ - Some("Alpha"), - Some("alphabet"), - Some("beta"), - None, - Some("ALPHA"), - ]); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - let result = - liquid_array.compare_with(b"Al%", &ByteViewOperator::SubString(SubString::Contains)); - let expected = BooleanArray::from(vec![ - Some(true), - Some(false), - Some(false), - None, - Some(false), - ]); - assert_eq!(result, expected); - - let result = - liquid_array.compare_with(b"Al%", &ByteViewOperator::SubString(SubString::NotContains)); - let expected = BooleanArray::from(vec![Some(false), Some(true), Some(true), None, Some(true)]); - assert_eq!(result, expected); -} - -// Benchmark tests for v2 offset compression improvements -fn generate_mixed_size_strings(count: usize, seed: u64) -> Vec { - let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - let mut strings = Vec::with_capacity(count); - - for _ in 0..count { - let size_type = rng.random_range(0..4); - let string = match size_type { - 0 => { - // Very short strings (1-3 chars) - stress test prefix optimization - let len = rng.random_range(1..=3); - (0..len) - .map(|_| rng.random_range(b'a'..=b'z') as char) - .collect() - } - 1 => { - // Medium strings (50-200 chars) - test offset compression - let len = rng.random_range(50..=200); - (0..len) - .map(|_| rng.random_range(b'a'..=b'z') as char) - .collect() - } - 2 => { - // Long strings (1000-5000 chars) - stress test linear regression - let len = rng.random_range(1000..=5000); - (0..len) - .map(|_| rng.random_range(b'a'..=b'z') as char) - .collect() - } - _ => { - // Very long strings (10k+ chars) - edge case for offset compression - let len = rng.random_range(10000..=50000); - "x".repeat(len) - } - }; - strings.push(string); - } - - strings -} - -fn generate_zipf_strings(count: usize, base_strings: &[&str], seed: u64) -> Vec { - let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - let mut strings = Vec::with_capacity(count); - - // Simple Zipf-like distribution: first few strings are much more common - for _ in 0..count { - let zipf_choice = rng.random_range(0..100); - let base_idx = if zipf_choice < 50 { - 0 // 50% chance of first string - } else if zipf_choice < 75 { - 1 // 25% chance of second string - } else if zipf_choice < 87 { - 2 // 12% chance of third string - } else { - rng.random_range(3..base_strings.len()) // remaining strings split rest - }; - - let base = base_strings[base_idx]; - - // Add variations to create realistic patterns - let variation = rng.random_range(0..4); - let string = match variation { - 0 => base.to_string(), // Exact duplicate - 1 => format!("{}_{}", base, rng.random_range(1000..9999)), // Common suffix - 2 => format!("{}/{}", base, rng.random_range(100..999)), // Path-like - _ => format!("prefix_{}", base), // Common prefix - }; - strings.push(string); - } - - strings -} - -#[test] -fn test_mixed_size_offset_views() { - let strings = generate_mixed_size_strings(16384, 42); - let input = StringArray::from(strings.clone()); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Verify correctness - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); -} - -#[test] -fn test_zipf_offset_views() { - // Real-world string patterns - let base_patterns = &[ - "error", "warning", "info", "debug", "user", "admin", "guest", "GET", "POST", "PUT", - "DELETE", "success", "failure", "pending", "/api/v1", "/api/v2", "/health", - ]; - - let strings = generate_zipf_strings(16384, base_patterns, 123); - let input = StringArray::from(strings.clone()); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Verify correctness - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - let offset_bytes = liquid_array.fsst_buffer.offset_bytes(); - - assert!( - offset_bytes <= 2, - "Zipf patterns with short strings should use 1 or 2 byte compact offsets, got {} bytes", - offset_bytes - ); -} - -#[test] -fn test_offset_stress() { - let mut strings = Vec::with_capacity(16384); - - // Create strings with problematic offset patterns - for i in 0..16384 { - let string = match i % 8 { - 0 => "a".to_string(), // tiny - 1 => "x".repeat(1000 + (i % 100)), // variable medium - 2 => "b".to_string(), // tiny - 3 => "y".repeat(5000 + (i % 1000)), // variable large - 4 => "c".to_string(), // tiny - 5 => "medium".repeat(50 + (i % 20)), // variable medium - 6 => "huge".repeat(2000 + (i % 500)), // variable huge - _ => format!("string_{}", i), // varied length based on number - }; - strings.push(string); - } - - let input = StringArray::from(strings); - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let liquid_array = LiquidByteViewArray::::from_string_array(&input, compressor); - - // Verify correctness - let output = liquid_array.to_arrow_array(); - assert_eq!(&input, output.as_string::()); - - // Test offset compression handles the stress case - let offsets = liquid_array.fsst_buffer.offsets(); - - // Verify offsets are monotonic - for i in 1..offsets.len() { - assert!(offsets[i] >= offsets[i - 1], "Offsets should be monotonic"); - } -} diff --git a/src/core/src/liquid_array/decimal_array.rs b/src/core/src/liquid_array/decimal_array.rs deleted file mode 100644 index 6d462ee0b..000000000 --- a/src/core/src/liquid_array/decimal_array.rs +++ /dev/null @@ -1,316 +0,0 @@ -use bytes::Bytes; -use std::{any::Any, mem::size_of, sync::Arc}; - -use arrow::array::{Array, ArrayRef, PrimitiveArray}; -use arrow::buffer::ScalarBuffer; -use arrow::datatypes::{Decimal128Type, Decimal256Type, DecimalType, UInt64Type, i256}; -use arrow_schema::DataType; -use num_traits::ToPrimitive; - -use super::{LiquidArray, LiquidDataType}; -use crate::liquid_array::ipc::{LiquidIPCHeader, get_physical_type_id}; -use crate::liquid_array::raw::BitPackedArray; -use crate::utils::get_bit_width; - -#[derive(Debug, Clone, Copy)] -struct DecimalMeta { - precision: u8, - scale: i8, - is_256: bool, -} - -impl DecimalMeta { - fn from_data_type(data_type: &DataType) -> Self { - match data_type { - DataType::Decimal128(precision, scale) => Self { - precision: *precision, - scale: *scale, - is_256: false, - }, - DataType::Decimal256(precision, scale) => Self { - precision: *precision, - scale: *scale, - is_256: true, - }, - _ => panic!("unsupported decimal data type: {data_type:?}"), - } - } - - fn data_type(&self) -> DataType { - if self.is_256 { - DataType::Decimal256(self.precision, self.scale) - } else { - DataType::Decimal128(self.precision, self.scale) - } - } - - fn arrow_code(&self) -> u8 { - if self.is_256 { 1 } else { 0 } - } -} - -#[repr(C)] -struct DecimalArrayHeader { - arrow_type: u8, // 0 for Decimal128, 1 for Decimal256 - precision: u8, - scale: i8, - __padding: u8, - __reserved: u32, -} - -impl DecimalArrayHeader { - const fn size() -> usize { - 8 - } - - fn from_meta(meta: DecimalMeta) -> Self { - Self { - arrow_type: meta.arrow_code(), - precision: meta.precision, - scale: meta.scale, - __padding: 0, - __reserved: 0, - } - } - - fn to_bytes(&self) -> [u8; Self::size()] { - let mut bytes = [0; Self::size()]; - bytes[0] = self.arrow_type; - bytes[1] = self.precision; - bytes[2] = self.scale as u8; - bytes - } - - fn from_bytes(bytes: &[u8]) -> Self { - if bytes.len() < Self::size() { - panic!( - "value too small for DecimalArrayHeader, expected at least {} bytes, got {}", - Self::size(), - bytes.len() - ); - } - Self { - arrow_type: bytes[0], - precision: bytes[1], - scale: bytes[2] as i8, - __padding: 0, - __reserved: 0, - } - } -} - -/// Liquid decimal array stored as a compressed u64 primitive. -#[derive(Debug)] -pub struct LiquidDecimalArray { - meta: DecimalMeta, - bit_packed: BitPackedArray, - reference_value: u64, -} - -impl LiquidDecimalArray { - pub(crate) fn fits_u64(array: &PrimitiveArray) -> bool - where - T::Native: ToPrimitive, - { - array.iter().flatten().all(|v| v.to_u64().is_some()) - } - - pub(crate) fn from_decimal_array(array: &PrimitiveArray) -> Self - where - T::Native: ToPrimitive, - { - debug_assert!(Self::fits_u64(array)); - let meta = DecimalMeta::from_data_type(array.data_type()); - if array.null_count() == array.len() { - return Self { - meta, - bit_packed: BitPackedArray::new_null_array(array.len()), - reference_value: 0, - }; - } - - let nulls = array.nulls().cloned(); - let mut min = u64::MAX; - let mut max = 0u64; - let values: Vec = array - .iter() - .map(|v| match v { - Some(v) => { - let value = v.to_u64().expect("decimal fits u64"); - if value < min { - min = value; - } - if value > max { - max = value; - } - value - } - None => 0, - }) - .collect(); - - let bit_width = get_bit_width(max - min); - let offsets = ScalarBuffer::from_iter(values.iter().map(|v| v.saturating_sub(min))); - let unsigned_array = PrimitiveArray::::new(offsets, nulls); - let bit_packed = BitPackedArray::from_primitive(unsigned_array, bit_width); - - Self { - meta, - bit_packed, - reference_value: min, - } - } - - fn bit_pack_starting_loc() -> usize { - let header_size = LiquidIPCHeader::size() + DecimalArrayHeader::size(); - (header_size + size_of::() + 7) & !7 - } - - fn to_u64_array(&self) -> PrimitiveArray { - let unsigned_array = self.bit_packed.to_primitive(); - let (_data_type, values, _nulls) = unsigned_array.into_parts(); - let nulls = self.bit_packed.nulls(); - let values = if self.reference_value != 0 { - let reference_value = self.reference_value; - ScalarBuffer::from_iter(values.iter().map(|v| v.wrapping_add(reference_value))) - } else { - values - }; - PrimitiveArray::::new(values, nulls.cloned()) - } - - pub(crate) fn to_bytes_inner(&self) -> Vec { - let header_size = LiquidIPCHeader::size() + DecimalArrayHeader::size(); - let mut result = Vec::with_capacity(Self::bit_pack_starting_loc() + 256); - result.resize(header_size, 0); - - let logical_type_id = LiquidDataType::Decimal as u16; - let physical_type_id = get_physical_type_id::(); - let ipc_header = LiquidIPCHeader::new(logical_type_id, physical_type_id); - result[0..LiquidIPCHeader::size()].copy_from_slice(&ipc_header.to_bytes()); - - let decimal_header = DecimalArrayHeader::from_meta(self.meta); - result[LiquidIPCHeader::size()..header_size].copy_from_slice(&decimal_header.to_bytes()); - - result.extend_from_slice(&self.reference_value.to_le_bytes()); - while result.len() < Self::bit_pack_starting_loc() { - result.push(0); - } - self.bit_packed.to_bytes(&mut result); - result - } - - pub(crate) fn from_bytes(bytes: Bytes) -> Self { - let header_size = LiquidIPCHeader::size() + DecimalArrayHeader::size(); - let header = LiquidIPCHeader::from_bytes(&bytes); - - assert_eq!(header.logical_type_id, LiquidDataType::Decimal as u16); - assert_eq!( - header.physical_type_id, - get_physical_type_id::() - ); - - let decimal_header = - DecimalArrayHeader::from_bytes(&bytes[LiquidIPCHeader::size()..header_size]); - let meta = DecimalMeta { - precision: decimal_header.precision, - scale: decimal_header.scale, - is_256: match decimal_header.arrow_type { - 0 => false, - 1 => true, - _ => panic!( - "unsupported decimal type code: {}", - decimal_header.arrow_type - ), - }, - }; - - let ref_start = header_size; - let ref_end = ref_start + size_of::(); - let reference_value = u64::from_le_bytes(bytes[ref_start..ref_end].try_into().unwrap()); - - let bit_packed_data = bytes.slice(Self::bit_pack_starting_loc()..); - let bit_packed = BitPackedArray::::from_bytes(bit_packed_data); - - Self { - meta, - bit_packed, - reference_value, - } - } -} - -impl LiquidArray for LiquidDecimalArray { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() + size_of::() + size_of::() - } - - fn len(&self) -> usize { - self.bit_packed.len() - } - - fn to_arrow_array(&self) -> ArrayRef { - let u64_array = self.to_u64_array(); - let (_data_type, values, nulls) = u64_array.into_parts(); - let data_type = self.meta.data_type(); - if self.meta.is_256 { - let values_i256 = - ScalarBuffer::from_iter(values.iter().map(|v| i256::from_i128(*v as i128))); - let array = PrimitiveArray::::new(values_i256, nulls); - Arc::new(array.with_data_type(data_type)) - } else { - let values_i128 = ScalarBuffer::from_iter(values.iter().map(|v| *v as i128)); - let array = PrimitiveArray::::new(values_i128, nulls); - Arc::new(array.with_data_type(data_type)) - } - } - - fn original_arrow_data_type(&self) -> DataType { - self.meta.data_type() - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Decimal - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::Decimal128Builder; - - #[test] - fn decimal_u64_roundtrip() { - let mut builder = Decimal128Builder::new(); - builder.append_value(100_i128); - builder.append_null(); - builder.append_value(250_i128); - let original = builder.finish().with_precision_and_scale(10, 2).unwrap(); - - let liquid = LiquidDecimalArray::from_decimal_array(&original); - let arrow = liquid.to_arrow_array(); - assert_eq!(arrow.as_ref(), &original); - } - - #[test] - fn decimal_u64_ipc_roundtrip() { - let mut builder = Decimal128Builder::new(); - builder.append_value(12345_i128); - builder.append_value(67890_i128); - let original = builder.finish().with_precision_and_scale(12, 3).unwrap(); - - let liquid = LiquidDecimalArray::from_decimal_array(&original); - let bytes = liquid.to_bytes(); - let decoded = LiquidDecimalArray::from_bytes(bytes.into()); - let arrow = decoded.to_arrow_array(); - assert_eq!(arrow.as_ref(), &original); - } -} diff --git a/src/core/src/liquid_array/fix_len_byte_array.rs b/src/core/src/liquid_array/fix_len_byte_array.rs deleted file mode 100644 index b82dd0043..000000000 --- a/src/core/src/liquid_array/fix_len_byte_array.rs +++ /dev/null @@ -1,599 +0,0 @@ -use std::{any::Any, sync::Arc}; - -use ahash::HashMap; -use arrow::{ - array::{ - Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, PrimitiveArray, - UInt16Array, - }, - compute::kernels::cast, - datatypes::{Decimal128Type, Decimal256Type, DecimalType, UInt16Type}, -}; -use arrow_schema::DataType; -use bytes::Bytes; -use fsst::Compressor; - -use crate::utils::CheckedDictionaryArray; - -use super::{ - LiquidArray, LiquidDataType, - raw::{BitPackedArray, FsstArray}, -}; -use crate::liquid_array::ipc::LiquidIPCHeader; - -/// A fixed length byte array. -#[derive(Debug)] -pub struct LiquidFixedLenByteArray { - arrow_type: ArrowFixedLenByteArrayType, - keys: BitPackedArray, - values: FsstArray, -} - -#[derive(Debug, Clone)] -pub enum ArrowFixedLenByteArrayType { - Decimal128(u8, i8), - Decimal256(u8, i8), -} - -impl From<&DataType> for ArrowFixedLenByteArrayType { - fn from(value: &DataType) -> Self { - match value { - DataType::Decimal128(precision, scale) => { - ArrowFixedLenByteArrayType::Decimal128(*precision, *scale) - } - DataType::Decimal256(precision, scale) => { - ArrowFixedLenByteArrayType::Decimal256(*precision, *scale) - } - _ => panic!("Unsupported arrow type: {value:?}"), - } - } -} - -impl From<&ArrowFixedLenByteArrayType> for DataType { - fn from(value: &ArrowFixedLenByteArrayType) -> Self { - match value { - ArrowFixedLenByteArrayType::Decimal128(precision, scale) => { - DataType::Decimal128(*precision, *scale) - } - ArrowFixedLenByteArrayType::Decimal256(precision, scale) => { - DataType::Decimal256(*precision, *scale) - } - } - } -} - -impl ArrowFixedLenByteArrayType { - pub fn value_width(&self) -> usize { - match self { - ArrowFixedLenByteArrayType::Decimal128(_, _) => Decimal128Type::BYTE_LENGTH, - ArrowFixedLenByteArrayType::Decimal256(_, _) => Decimal256Type::BYTE_LENGTH, - } - } -} - -impl LiquidArray for LiquidFixedLenByteArray { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.keys.get_array_memory_size() + self.values.get_array_memory_size() - } - - fn len(&self) -> usize { - self.keys.len() - } - - fn to_arrow_array(&self) -> ArrayRef { - if self.keys.len() < 2048 || self.keys.len() < self.values.len() { - // Use keyed decompression for smaller arrays - self.to_arrow_array_decompress_keyed() - } else { - // Use full decompression for larger arrays - self.to_arrow_array_decompress_all() - } - } - - fn to_best_arrow_array(&self) -> ArrayRef { - self.to_arrow_array() - } - - fn original_arrow_data_type(&self) -> DataType { - DataType::from(&self.arrow_type) - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::FixedLenByteArray - } -} - -// Specialized header for fixed-length byte arrays -#[repr(C)] -struct FixedLenByteArrayHeader { - key_size: u32, - value_size: u32, - arrow_type: u8, // 0 for Decimal128, 1 for Decimal256 - precision: u8, - scale: i8, - __padding: u8, -} - -impl FixedLenByteArrayHeader { - const fn size() -> usize { - 12 - } - - fn to_bytes(&self) -> [u8; Self::size()] { - let mut bytes = [0; Self::size()]; - bytes[0..4].copy_from_slice(&self.key_size.to_le_bytes()); - bytes[4..8].copy_from_slice(&self.value_size.to_le_bytes()); - bytes[8] = self.arrow_type; - bytes[9] = self.precision; - bytes[10] = self.scale as u8; - bytes - } - - fn from_bytes(bytes: &[u8]) -> Self { - if bytes.len() < Self::size() { - panic!( - "value too small for FixedLenByteArrayHeader, expected at least {} bytes, got {}", - Self::size(), - bytes.len() - ); - } - let key_size = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); - let value_size = u32::from_le_bytes(bytes[4..8].try_into().unwrap()); - let arrow_type = bytes[8]; - let precision = bytes[9]; - let scale = bytes[10] as i8; - Self { - key_size, - value_size, - arrow_type, - precision, - scale, - __padding: 0, - } - } -} - -impl LiquidFixedLenByteArray { - pub(crate) fn to_bytes_inner(&self) -> Vec { - // Create a buffer for the final output data, starting with the header - let header_size = LiquidIPCHeader::size() + FixedLenByteArrayHeader::size(); - let mut result = Vec::with_capacity(header_size + 1024); // Pre-allocate a reasonable size - - result.resize(header_size, 0); - - // Serialize the BitPackedArray (keys) - let keys_start = result.len(); - self.keys().to_bytes(&mut result); - let keys_size = result.len() - keys_start; - - // Add padding to ensure FsstArray starts at an 8-byte aligned position - while !result.len().is_multiple_of(8) { - result.push(0); - } - - // Serialize the FsstArray (values) - let values_start = result.len(); - self.values().to_bytes(&mut result); - let values_size = result.len() - values_start; - - // Go back and fill in the header - let ipc_header = LiquidIPCHeader::new(LiquidDataType::FixedLenByteArray as u16, 0); - let header = &mut result[0..header_size]; - header[0..LiquidIPCHeader::size()].copy_from_slice(&ipc_header.to_bytes()); - - // Map the ArrowFixedLenByteArrayType to our header - let (arrow_type, precision, scale) = match self.arrow_type() { - ArrowFixedLenByteArrayType::Decimal128(p, s) => (0, *p, *s), - ArrowFixedLenByteArrayType::Decimal256(p, s) => (1, *p, *s), - }; - - let fixed_len_byte_array_header = FixedLenByteArrayHeader { - key_size: keys_size as u32, - value_size: values_size as u32, - arrow_type, - precision, - scale, - __padding: 0, - }; - header[LiquidIPCHeader::size()..header_size] - .copy_from_slice(&fixed_len_byte_array_header.to_bytes()); - - result - } - - /// Deserialize a LiquidFixedLenByteArray from bytes, using zero-copy where possible. - pub fn from_bytes(bytes: Bytes, compressor: Arc) -> Self { - let header_size = LiquidIPCHeader::size() + FixedLenByteArrayHeader::size(); - let header = LiquidIPCHeader::from_bytes(&bytes); - - // Verify the logical type - assert_eq!( - header.logical_type_id, - LiquidDataType::FixedLenByteArray as u16 - ); - - let fixed_len_header = - FixedLenByteArrayHeader::from_bytes(&bytes[LiquidIPCHeader::size()..header_size]); - - // Parse arrow type based on the header - let arrow_type = match fixed_len_header.arrow_type { - 0 => ArrowFixedLenByteArrayType::Decimal128( - fixed_len_header.precision, - fixed_len_header.scale, - ), - 1 => ArrowFixedLenByteArrayType::Decimal256( - fixed_len_header.precision, - fixed_len_header.scale, - ), - _ => panic!( - "Unsupported arrow type code: {}", - fixed_len_header.arrow_type - ), - }; - - // Calculate offsets - let keys_size = fixed_len_header.key_size as usize; - let values_size = fixed_len_header.value_size as usize; - - let keys_start = header_size; - let keys_end = keys_start + keys_size; - - if keys_end > bytes.len() { - panic!("Keys data extends beyond input buffer"); - } - - // Ensure values data starts at 8-byte aligned position - let values_start = (keys_end + 7) & !7; // Round up to next 8-byte boundary - let values_end = values_start + values_size; - - if values_end > bytes.len() { - panic!("Values data extends beyond input buffer"); - } - - // Extract and deserialize components - let keys_data = bytes.slice(keys_start..keys_end); - let keys = BitPackedArray::::from_bytes(keys_data); - - let values_data = bytes.slice(values_start..values_end); - let values = FsstArray::from_bytes(values_data, compressor); - - Self::from_parts(arrow_type, keys, values) - } -} - -impl LiquidFixedLenByteArray { - /// Create a new fixed length byte array from a decimal array. - pub fn from_decimal_array( - array: &PrimitiveArray, - compressor: Arc, - ) -> Self { - let dict = CheckedDictionaryArray::from_decimal_array(array); - Self::from_dict_array_inner( - dict, - compressor, - ArrowFixedLenByteArrayType::from(array.data_type()), - ) - } - - /// Train a new fixed length byte array from a decimal array. - pub fn train_from_decimal_array( - array: &PrimitiveArray, - ) -> (Arc, Self) { - let value_width = array.data_type().primitive_width().unwrap(); - let value_buffer = array.values().inner().chunks(value_width); - let compressor = FsstArray::train_compressor(value_buffer); - let compressor = Arc::new(compressor); - let liquid_array = Self::from_decimal_array(array, compressor.clone()); - (compressor, liquid_array) - } - - fn from_dict_array_inner( - array: CheckedDictionaryArray, - compressor: Arc, - arrow_type: ArrowFixedLenByteArrayType, - ) -> Self { - let bit_width_for_key = array.bit_width_for_key(); - let (keys, values) = array.into_inner().into_parts(); - let bit_packed_array = BitPackedArray::from_primitive(keys, bit_width_for_key); - - let fsst_values = match arrow_type { - ArrowFixedLenByteArrayType::Decimal128(_, _) => { - let values = values.as_primitive::(); - FsstArray::from_decimal128_array_with_compressor(values, compressor) - } - ArrowFixedLenByteArrayType::Decimal256(_, _) => { - let values = values.as_primitive::(); - FsstArray::from_decimal256_array_with_compressor(values, compressor) - } - }; - Self { - arrow_type, - keys: bit_packed_array, - values: fsst_values, - } - } - - /// Convert to arrow array by decompressing all values - fn to_arrow_array_decompress_all(&self) -> ArrayRef { - match self.arrow_type { - ArrowFixedLenByteArrayType::Decimal128(precision, scale) => { - let array = self.values.to_decimal128_array(&self.arrow_type); - let keys = self.keys.to_primitive(); - let dict = - unsafe { DictionaryArray::::new_unchecked(keys, Arc::new(array)) }; - cast(&dict, &DataType::Decimal128(precision, scale)).unwrap() - } - ArrowFixedLenByteArrayType::Decimal256(precision, scale) => { - let array = self.values.to_decimal256_array(&self.arrow_type); - let keys = self.keys.to_primitive(); - let dict = - unsafe { DictionaryArray::::new_unchecked(keys, Arc::new(array)) }; - cast(&dict, &DataType::Decimal256(precision, scale)).unwrap() - } - } - } - - /// Convert to arrow array by only decompressing values referenced by keys - fn to_arrow_array_decompress_keyed(&self) -> ArrayRef { - let primitive_key = self.keys.to_primitive(); - let mut hit_mask = BooleanBufferBuilder::new(self.values.len()); - hit_mask.advance(self.values.len()); - for v in primitive_key.iter().flatten() { - hit_mask.set_bit(v as usize, true); - } - let hit_mask = hit_mask.finish(); - let selected_cnt = hit_mask.count_set_bits(); - - let mut key_map = - HashMap::with_capacity_and_hasher(selected_cnt, ahash::RandomState::new()); - let mut offset = 0; - for (i, select) in hit_mask.iter().enumerate() { - if select { - key_map.insert(i, offset); - offset += 1; - } - } - let new_keys = UInt16Array::from_iter( - primitive_key - .iter() - .map(|v| v.map(|v| key_map[&(v as usize)])), - ); - - let decompressed_values = self.decompress_keyed_values(&hit_mask); - let dict = - unsafe { DictionaryArray::::new_unchecked(new_keys, decompressed_values) }; - - match self.arrow_type { - ArrowFixedLenByteArrayType::Decimal128(precision, scale) => { - cast(&dict, &DataType::Decimal128(precision, scale)).unwrap() - } - ArrowFixedLenByteArrayType::Decimal256(precision, scale) => { - cast(&dict, &DataType::Decimal256(precision, scale)).unwrap() - } - } - } - - /// Decompress only the values that are selected by the hit mask - fn decompress_keyed_values(&self, hit_mask: &arrow::buffer::BooleanBuffer) -> ArrayRef { - let value_width = self.arrow_type.value_width(); - let selected_cnt = hit_mask.count_set_bits(); - assert_eq!(hit_mask.len(), self.values.len()); - let selected: Vec = hit_mask - .iter() - .enumerate() - .filter_map(|(i, select)| select.then_some(i)) - .collect(); - - let (value_buffer, offsets) = self.values.to_uncompressed_selected(&selected); - - debug_assert_eq!(offsets.len(), selected_cnt + 1); - debug_assert_eq!(value_buffer.len(), selected_cnt * value_width); - - match self.arrow_type { - ArrowFixedLenByteArrayType::Decimal128(precision, scale) => { - let array_data = - arrow::array::ArrayDataBuilder::new(DataType::Decimal128(precision, scale)) - .len(selected_cnt) - .add_buffer(value_buffer) - .build() - .unwrap(); - Arc::new(arrow::array::Decimal128Array::from(array_data)) - } - ArrowFixedLenByteArrayType::Decimal256(precision, scale) => { - let array_data = - arrow::array::ArrayDataBuilder::new(DataType::Decimal256(precision, scale)) - .len(selected_cnt) - .add_buffer(value_buffer) - .build() - .unwrap(); - Arc::new(arrow::array::Decimal256Array::from(array_data)) - } - } - } - - pub(crate) fn from_parts( - arrow_type: ArrowFixedLenByteArrayType, - keys: BitPackedArray, - values: FsstArray, - ) -> Self { - Self { - arrow_type, - keys, - values, - } - } - - pub(super) fn values(&self) -> &FsstArray { - &self.values - } - - pub(super) fn keys(&self) -> &BitPackedArray { - &self.keys - } - - pub(super) fn arrow_type(&self) -> &ArrowFixedLenByteArrayType { - &self.arrow_type - } -} - -#[cfg(test)] -mod tests { - use crate::liquid_array::utils::gen_test_decimal_array; - - use super::*; - use arrow_schema::DataType; - - fn test_decimal_roundtrip(data_type: DataType) { - let original_array = gen_test_decimal_array::(data_type); - let (_compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&original_array); - - let arrow_array = liquid_array.to_arrow_array(); - let roundtrip_array = arrow_array.as_primitive::(); - - assert_eq!(original_array.len(), roundtrip_array.len()); - - for i in 0..original_array.len() { - assert_eq!(original_array.is_null(i), roundtrip_array.is_null(i)); - if !original_array.is_null(i) { - assert_eq!(original_array.value(i), roundtrip_array.value(i)); - } - } - } - - #[test] - fn test_original_arrow_data_type_returns_decimal128() { - let data_type = DataType::Decimal128(15, 3); - let original_array = gen_test_decimal_array::(data_type); - let (_compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&original_array); - - assert_eq!( - liquid_array.original_arrow_data_type(), - DataType::Decimal128(15, 3) - ); - } - - #[test] - fn test_decimal128_roundtrip() { - test_decimal_roundtrip::(DataType::Decimal128(15, 3)); - } - - #[test] - fn test_decimal256_roundtrip() { - test_decimal_roundtrip::(DataType::Decimal256(38, 6)); - } - - fn test_decimal_filter_operation(data_type: DataType) { - let original_array = gen_test_decimal_array::(data_type); - let (_compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&original_array); - - let mut filter_builder = arrow::array::BooleanBuilder::new(); - for i in 0..liquid_array.len() { - filter_builder.append_value(i.is_multiple_of(2)); - } - let filter = filter_builder.finish(); - let (filter, _null) = filter.into_parts(); - let arrow_filtered = liquid_array.filter(&filter); - let arrow_typed = arrow_filtered.as_primitive::(); - - assert_eq!(arrow_filtered.len(), original_array.len() / 2); - - for (i, val) in arrow_typed.iter().enumerate() { - if original_array.is_null(i * 2) { - assert!(arrow_typed.is_null(i)); - } else { - assert_eq!(val.unwrap(), original_array.value(i * 2)); - } - } - } - - #[test] - fn test_decimal128_filter_operation() { - test_decimal_filter_operation::(DataType::Decimal128(12, 2)); - } - - #[test] - fn test_decimal256_filter_operation() { - test_decimal_filter_operation::(DataType::Decimal256(38, 4)); - } - - #[test] - fn test_keyed_decompression_optimization() { - // Create a larger decimal array to test the optimization logic - let mut builder = arrow::array::Decimal128Builder::new(); - - // Create 10 distinct values - for i in 0..10 { - builder.append_value(i as i128 * 1000); - } - let distinct_values = builder.finish().with_precision_and_scale(15, 3).unwrap(); - - let (_compressor, mut liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&distinct_values); - - // Create a small keys array that only references a few values - // This should trigger the keyed decompression path (keys.len() < 2048) - let small_keys = UInt16Array::from(vec![0, 2, 4, 2, 0]); // Only references indices 0, 2, 4 - liquid_array.keys = - BitPackedArray::from_primitive(small_keys, std::num::NonZero::new(3).unwrap()); - - // Test both decompress_all and decompress_keyed should give the same result - let result_all = liquid_array.to_arrow_array_decompress_all(); - let result_keyed = liquid_array.to_arrow_array_decompress_keyed(); - - // Both should be equal - assert_eq!( - result_all.as_primitive::().values(), - result_keyed.as_primitive::().values() - ); - - // Verify the actual values are correct - let expected_values = vec![0, 2000, 4000, 2000, 0]; // i * 1000 for i in [0, 2, 4, 2, 0] - let actual_values: Vec = result_keyed - .as_primitive::() - .values() - .iter() - .copied() - .collect(); - assert_eq!(expected_values, actual_values); - } - - #[test] - fn test_large_array_uses_full_decompression() { - // Test that large arrays (>= 2048) use full decompression - let distinct_values = gen_test_decimal_array::(DataType::Decimal128(15, 3)); - let (_compressor, mut liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&distinct_values); - - // Create a large keys array - let large_keys: Vec = (0..3000) - .map(|i| (i % distinct_values.len()) as u16) - .collect(); - let large_keys = UInt16Array::from(large_keys); - liquid_array.keys = BitPackedArray::from_primitive( - large_keys, - std::num::NonZero::new(4).unwrap(), // Adjust bit width as needed - ); - - // This should use the full decompression path since keys.len() >= 2048 - let result = liquid_array.to_arrow_array(); - assert_eq!(result.len(), 3000); - - // Verify the result is valid by checking it matches decompress_all - let result_all = liquid_array.to_arrow_array_decompress_all(); - assert_eq!( - result.as_primitive::().values(), - result_all.as_primitive::().values() - ); - } -} diff --git a/src/core/src/liquid_array/float_array.rs b/src/core/src/liquid_array/float_array.rs deleted file mode 100644 index d3e175bf9..000000000 --- a/src/core/src/liquid_array/float_array.rs +++ /dev/null @@ -1,813 +0,0 @@ -use bytes::Bytes; -/// -/// Acknowledgement: -/// The ALP compression implemented in this file is based on the Rust implementation available at https://github.com/spiraldb/alp -/// -use std::{ - any::Any, - fmt::Debug, - ops::{Mul, Shl, Shr}, - sync::Arc, -}; - -use arrow::{ - array::{Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, PrimitiveArray}, - buffer::ScalarBuffer, - datatypes::{ - ArrowNativeType, Float32Type, Float64Type, Int32Type, Int64Type, UInt32Type, UInt64Type, - }, -}; -use arrow_schema::DataType; -use fastlanes::BitPacking; -use num_traits::{AsPrimitive, Float, FromPrimitive}; - -use super::LiquidDataType; -use crate::liquid_array::LiquidArray; -use crate::liquid_array::ipc::LiquidIPCHeader; -use crate::liquid_array::ipc::{PhysicalTypeMarker, get_physical_type_id}; -use crate::liquid_array::raw::BitPackedArray; -use crate::utils::get_bit_width; - -mod private { - use arrow::{ - array::ArrowNumericType, - datatypes::{Float32Type, Float64Type}, - }; - use num_traits::AsPrimitive; - - pub trait Sealed: ArrowNumericType + AsPrimitive> {} - - impl Sealed for Float32Type {} - impl Sealed for Float64Type {} -} - -const NUM_SAMPLES: usize = 1024; // we use FASTLANES to encode array, the sample size needs to be at least 1024 to get a good estimate of the best exponents - -/// LiquidFloatType is a sealed trait that represents all the float types supported by Liquid. -/// Implementors are Float32Type and Float64Type. TODO(): What about Float16Type, decimal types? -pub trait LiquidFloatType: - ArrowPrimitiveType< - Native: AsPrimitive< - ::Native // Native must be convertible to the Native type of Self::UnSignedType - > - + AsPrimitive<::Native> - + FromPrimitive - + AsPrimitive<::Native> - + Mul<::Native> - + Float // required for decode_single and encode_single_unchecked - > - + private::Sealed - + Debug - + PhysicalTypeMarker -{ - type UnsignedIntType: - ArrowPrimitiveType< - Native: BitPacking + - AsPrimitive<::Native> - + AsPrimitive<::Native> - + AsPrimitive - > - + Debug; - type SignedIntType: - ArrowPrimitiveType< - Native: AsPrimitive<::Native> - + AsPrimitive<::Native> - + Ord - + Shr::Native> - + Shl::Native> - + From - > - + Debug + Sync + Send; - - const SWEET: ::Native; - const MAX_EXPONENT: u8; - const FRACTIONAL_BITS: u8; - const F10: &'static [::Native]; - const IF10: &'static [::Native]; - - #[inline] - fn fast_round(val: ::Native) -> ::Native { - ((val + Self::SWEET) - Self::SWEET).as_() - } - - #[inline] - fn encode_single_unchecked(val: &::Native, exp: &Exponents) -> ::Native { - Self::fast_round(*val * Self::F10[exp.e as usize] * Self::IF10[exp.f as usize]) - } - - #[inline] - fn decode_single(val: &::Native, exp: &Exponents) -> ::Native { - let decoded_float: ::Native = (*val).as_(); - decoded_float * Self::F10[exp.f as usize] * Self::IF10[exp.e as usize] - } - -} - -impl LiquidFloatType for Float32Type { - type UnsignedIntType = UInt32Type; - type SignedIntType = Int32Type; - const FRACTIONAL_BITS: u8 = 23; - const MAX_EXPONENT: u8 = 10; - const SWEET: ::Native = (1 << Self::FRACTIONAL_BITS) - as ::Native - + (1 << (Self::FRACTIONAL_BITS - 1)) as ::Native; - const F10: &'static [::Native] = &[ - 1.0, - 10.0, - 100.0, - 1000.0, - 10000.0, - 100000.0, - 1000000.0, - 10000000.0, - 100000000.0, - 1000000000.0, - 10000000000.0, // 10^10 - ]; - const IF10: &'static [::Native] = &[ - 1.0, - 0.1, - 0.01, - 0.001, - 0.0001, - 0.00001, - 0.000001, - 0.0000001, - 0.00000001, - 0.000000001, - 0.0000000001, // 10^-10 - ]; -} - -impl LiquidFloatType for Float64Type { - type UnsignedIntType = UInt64Type; - type SignedIntType = Int64Type; - const FRACTIONAL_BITS: u8 = 52; - const MAX_EXPONENT: u8 = 18; - const SWEET: ::Native = (1u64 << Self::FRACTIONAL_BITS) - as ::Native - + (1u64 << (Self::FRACTIONAL_BITS - 1)) as ::Native; - const F10: &'static [::Native] = &[ - 1.0, - 10.0, - 100.0, - 1000.0, - 10000.0, - 100000.0, - 1000000.0, - 10000000.0, - 100000000.0, - 1000000000.0, - 10000000000.0, - 100000000000.0, - 1000000000000.0, - 10000000000000.0, - 100000000000000.0, - 1000000000000000.0, - 10000000000000000.0, - 100000000000000000.0, - 1000000000000000000.0, - 10000000000000000000.0, - 100000000000000000000.0, - 1000000000000000000000.0, - 10000000000000000000000.0, - 100000000000000000000000.0, // 10^23 - ]; - - const IF10: &'static [::Native] = &[ - 1.0, - 0.1, - 0.01, - 0.001, - 0.0001, - 0.00001, - 0.000001, - 0.0000001, - 0.00000001, - 0.000000001, - 0.0000000001, - 0.00000000001, - 0.000000000001, - 0.0000000000001, - 0.00000000000001, - 0.000000000000001, - 0.0000000000000001, - 0.00000000000000001, - 0.000000000000000001, - 0.0000000000000000001, - 0.00000000000000000001, - 0.000000000000000000001, - 0.0000000000000000000001, - 0.00000000000000000000001, // 10^-23 - ]; -} - -/// Liquid's single-precision floating point array -pub type LiquidFloat32Array = LiquidFloatArray; -/// Liquid's double precision floating point array -pub type LiquidFloat64Array = LiquidFloatArray; - -/// An array that stores floats in ALP -#[derive(Debug, Clone)] -pub struct LiquidFloatArray { - exponent: Exponents, - bit_packed: BitPackedArray, - patch_indices: Vec, - patch_values: Vec, - reference_value: ::Native, -} - -impl LiquidFloatArray -where - T: LiquidFloatType, -{ - /// Check if the Liquid float array is empty. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Get the length of the Liquid float array. - pub fn len(&self) -> usize { - self.bit_packed.len() - } - - /// Get the memory size of the Liquid primitive array. - pub fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() - + size_of::() - + self.patch_indices.capacity() * size_of::() - + self.patch_values.capacity() * size_of::() - + size_of::<::Native>() - } - - /// Create a Liquid primitive array from an Arrow float array. - pub fn from_arrow_array(arrow_array: arrow::array::PrimitiveArray) -> LiquidFloatArray { - let best_exponents = get_best_exponents::(&arrow_array); - encode_arrow_array(&arrow_array, &best_exponents) - } -} - -impl LiquidArray for LiquidFloatArray -where - T: LiquidFloatType, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn get_array_memory_size(&self) -> usize { - self.get_array_memory_size() - } - - fn len(&self) -> usize { - self.len() - } - - #[inline] - fn to_arrow_array(&self) -> ArrayRef { - let unsigned_array = self.bit_packed.to_primitive(); - let (_data_type, values, _nulls) = unsigned_array.into_parts(); - let nulls = self.bit_packed.nulls(); - // TODO(): Check if we should align vectors to cache line boundary - let mut decoded_values = Vec::from_iter(values.iter().map(|v| { - let mut val: ::Native = (*v).as_(); - val = val.add_wrapping(self.reference_value); - T::decode_single(&val, &self.exponent) - })); - - // Patch values - if !self.patch_indices.is_empty() { - for i in 0..self.patch_indices.len() { - decoded_values[self.patch_indices[i].as_usize()] = self.patch_values[i]; - } - } - - Arc::new(PrimitiveArray::::new( - ScalarBuffer::<::Native>::from(decoded_values), - nulls.cloned(), - )) - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Float - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn is_empty(&self) -> bool { - self.len() == 0 - } - - fn to_best_arrow_array(&self) -> ArrayRef { - self.to_arrow_array() - } -} - -impl LiquidFloatArray -where - T: LiquidFloatType, -{ - /* - Serialized LiquidFloatArray Memory Layout: - +--------------------------------------------------+ - | LiquidIPCHeader (16 bytes) | - +--------------------------------------------------+ - - +--------------------------------------------------+ - | reference_value | - | (size_of:: bytes) | // The reference value (e.g. minimum value) - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | // Padding to ensure 8-byte alignment - +--------------------------------------------------+ - - +--------------------------------------------------+ - | Exponents | - +--------------------------------------------------+ - | e (1 byte) | - +--------------------------------------------------+ - | f (1 byte) | - +--------------------------------------------------+ - | Padding (6 bytes) | - +--------------------------------------------------+ - - +--------------------------------------------------+ - | Patch Data | - +--------------------------------------------------+ - | patch_length (8 bytes) | - +--------------------------------------------------+ - | patch_indices (8 * patch_length btyes) | - +--------------------------------------------------+ - | patch_values (length *size_of:: btyes;| - | 8-byte aligned) | - +--------------------------------------------------+ - - +--------------------------------------------------+ - | BitPackedArray Data | - +--------------------------------------------------+ - | [BitPackedArray Header & Bit-Packed Values] | // Written by self.bit_packed.to_bytes() - +--------------------------------------------------+ - */ - pub(crate) fn to_bytes_inner(&self) -> Vec { - // Determine type ID based on the type - let physical_type_id = get_physical_type_id::(); - let logical_type_id = LiquidDataType::Float as u16; - let header = LiquidIPCHeader::new(logical_type_id, physical_type_id); - - let mut result = Vec::with_capacity(256); // Pre-allocate a reasonable size - - // Write header - result.extend_from_slice(&header.to_bytes()); - - // Write reference value - let ref_value_bytes = unsafe { - std::slice::from_raw_parts( - &self.reference_value as *const ::Native - as *const u8, - std::mem::size_of::<::Native>(), - ) - }; - result.extend_from_slice(ref_value_bytes); - - let exponents_starting_loc = (result.len() + 7) & !7; - // Insert padding before exponents start - while result.len() < exponents_starting_loc { - result.push(0); - } - - let exponent_e_bytes = - unsafe { std::slice::from_raw_parts(&self.exponent.e as *const u8, 1) }; - let exponent_f_bytes = - unsafe { std::slice::from_raw_parts(&self.exponent.f as *const u8, 1) }; - // Write exponents and padding - result.extend_from_slice(exponent_e_bytes); - result.extend_from_slice(exponent_f_bytes); - for _i in 0..6 { - result.push(0); - } - - // Number of bytes occupied by usize is target-dependent; use u64 instead - let patch_length = self.patch_indices.len() as u64; - - let patch_length_bytes = unsafe { - std::slice::from_raw_parts( - &patch_length as *const u64 as *const u8, - std::mem::size_of::(), - ) - }; - - // Write the patch length - result.extend_from_slice(patch_length_bytes); - - if !self.patch_indices.is_empty() { - let patch_indices_bytes = unsafe { - std::slice::from_raw_parts( - self.patch_indices.as_ptr() as *const u8, - std::mem::size_of::() * self.patch_indices.len(), - ) - }; - - // Write the patch indices - result.extend_from_slice(patch_indices_bytes); - - // Write the patch values - let patch_values_bytes = unsafe { - std::slice::from_raw_parts( - self.patch_values.as_ptr() as *const u8, - std::mem::size_of::() * self.patch_indices.len(), - ) - }; - result.extend_from_slice(patch_values_bytes); - } - let padding = ((result.len() + 7) & !7) - result.len(); - - // Add padding before writing bit-packed array - for _i in 0..padding { - result.push(0); - } - - // Serialize bit-packed values - self.bit_packed.to_bytes(&mut result); - - result - } - - /// Deserialize a LiquidFloatArray from bytes, using zero-copy where possible. - pub fn from_bytes(bytes: Bytes) -> Self { - let header = LiquidIPCHeader::from_bytes(&bytes); - - // Verify the type id - let physical_id = header.physical_type_id; - assert_eq!(physical_id, get_physical_type_id::()); - let logical_id = header.logical_type_id; - assert_eq!(logical_id, LiquidDataType::Float as u16); - - // Get the reference value - let ref_value_ptr = &bytes[LiquidIPCHeader::size()]; - let reference_value = unsafe { - (ref_value_ptr as *const u8 as *const ::Native) - .read_unaligned() - }; - - // Read exponents (e, f) & skip padding - let mut next = ((LiquidIPCHeader::size() - + std::mem::size_of::<::Native>()) - + 7) - & !7; - - // Read exponent fields (1 byte each) and skip 6 padding bytes - let exponent_e = bytes[next]; - let exponent_f = bytes[next + 1]; - next += 8; - - // Read patch length (8 bytes) - let mut patch_length = 0u64; - patch_length |= bytes[next] as u64; - patch_length |= (bytes[next + 1] as u64) << 8; - patch_length |= (bytes[next + 2] as u64) << 16; - patch_length |= (bytes[next + 3] as u64) << 24; - patch_length |= (bytes[next + 4] as u64) << 32; - patch_length |= (bytes[next + 5] as u64) << 40; - patch_length |= (bytes[next + 6] as u64) << 48; - patch_length |= (bytes[next + 7] as u64) << 56; - next += 8; - - // Read patch indices - let mut patch_indices = Vec::new(); - let mut patch_values = Vec::new(); - if patch_length > 0 { - let count = patch_length as usize; - let idx_bytes = count * std::mem::size_of::(); - let val_bytes = count * std::mem::size_of::(); - - let indices_slice = bytes.slice(next..next + idx_bytes); - next += idx_bytes; - patch_indices = unsafe { - let ptr = indices_slice.as_ptr() as *const u64; - std::slice::from_raw_parts(ptr, count).to_vec() - }; - - let values_slice = bytes.slice(next..next + val_bytes); - next += val_bytes; - patch_values = unsafe { - let ptr = values_slice.as_ptr() as *const T::Native; - std::slice::from_raw_parts(ptr, count).to_vec() - }; - } - - // Align up to 8 bytes for bit-packed array - next = (next + 7) & !7; - - let bit_packed = BitPackedArray::::from_bytes(bytes.slice(next..)); - - Self { - exponent: Exponents { - e: exponent_e, - f: exponent_f, - }, - bit_packed, - patch_indices, - patch_values, - reference_value, - } - } -} - -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct Exponents { - pub(crate) e: u8, - pub(crate) f: u8, -} - -fn encode_arrow_array( - arrow_array: &PrimitiveArray, - exp: &Exponents, // fill_value: &mut Option<::Native> -) -> LiquidFloatArray { - let mut patch_indices: Vec = Vec::new(); - let mut patch_values: Vec = Vec::new(); - let mut patch_count: usize = 0; - let mut fill_value: Option<::Native> = None; - let values = arrow_array.values(); - let nulls = arrow_array.nulls(); - - // All values are null - if arrow_array.null_count() == arrow_array.len() { - return LiquidFloatArray:: { - bit_packed: BitPackedArray::new_null_array(arrow_array.len()), - exponent: Exponents { e: 0, f: 0 }, - patch_indices: Vec::new(), - patch_values: Vec::new(), - reference_value: ::Native::ZERO, - }; - } - - let mut encoded_values = Vec::with_capacity(arrow_array.len()); - for v in values.iter() { - let encoded = T::encode_single_unchecked(&v.as_(), exp); - let decoded = T::decode_single(&encoded, exp); - // TODO(): Check if this is a bitwise comparison - let neq = !decoded.eq(&v.as_()) as usize; - patch_count += neq; - encoded_values.push(encoded); - } - - if patch_count > 0 { - patch_indices.resize_with(patch_count + 1, Default::default); - patch_values.resize_with(patch_count + 1, Default::default); - let mut patch_index: usize = 0; - - for i in 0..encoded_values.len() { - let decoded = T::decode_single(&encoded_values[i], exp); - patch_indices[patch_index] = i.as_(); - patch_values[patch_index] = arrow_array.value(i).as_(); - patch_index += !(decoded.eq(&values[i].as_())) as usize; - } - assert_eq!(patch_index, patch_count); - unsafe { - patch_indices.set_len(patch_count); - patch_values.set_len(patch_count); - } - } - - // find the first successfully encoded value (i.e., not patched) - // this is our fill value for missing values - if patch_count > 0 && patch_count < arrow_array.len() { - for i in 0..encoded_values.len() { - if i >= patch_indices.len() || patch_indices[i] != i as u64 { - fill_value = encoded_values.get(i).copied(); - break; - } - } - } - - // replace the patched values in the encoded array with the fill value - // for better downstream compression - if let Some(fill_value) = fill_value { - // handle the edge case where the first N >= 1 chunks are all patches - for patch_idx in &patch_indices { - encoded_values[*patch_idx as usize] = fill_value; - } - } - - let min = *encoded_values - .iter() - .min() - .expect("`encoded_values` shouldn't be all nulls"); - let max = *encoded_values - .iter() - .max() - .expect("`encoded_values` shouldn't be all nulls"); - let sub: ::Native = max.sub_wrapping(min).as_(); - - let unsigned_encoded_values = encoded_values - .iter() - .map(|v| { - let k: ::Native = v.sub_wrapping(min).as_(); - k - }) - .collect::>(); - let encoded_output = PrimitiveArray::<::UnsignedIntType>::new( - ScalarBuffer::from(unsigned_encoded_values), - nulls.cloned(), - ); - - let bit_width = get_bit_width(sub.as_()); - let bit_packed_array = BitPackedArray::from_primitive(encoded_output, bit_width); - - LiquidFloatArray:: { - bit_packed: bit_packed_array, - exponent: *exp, - patch_indices, - patch_values, - reference_value: min, - } -} - -fn get_best_exponents(arrow_array: &PrimitiveArray) -> Exponents { - let mut best_exponents = Exponents { e: 0, f: 0 }; - let mut min_encoded_size: usize = usize::MAX; - - let sample_arrow_array: Option> = - (arrow_array.len() > NUM_SAMPLES).then(|| { - arrow_array - .iter() - .step_by(arrow_array.len() / NUM_SAMPLES) - .filter(|s| s.is_some()) - .collect() - }); - - for e in 0..T::MAX_EXPONENT { - for f in 0..e { - let exp = Exponents { e, f }; - let liquid_array = - encode_arrow_array(sample_arrow_array.as_ref().unwrap_or(arrow_array), &exp); - if liquid_array.get_array_memory_size() < min_encoded_size { - best_exponents = exp; - min_encoded_size = liquid_array.get_array_memory_size(); - } - } - } - best_exponents -} - -#[cfg(test)] -mod tests { - - use arrow::buffer::BooleanBuffer; - - use super::*; - - macro_rules! test_roundtrip { - ($test_name: ident, $type:ty, $values: expr) => { - #[test] - fn $test_name() { - let original: Vec::Native>> = $values; - let array = PrimitiveArray::<$type>::from(original.clone()); - - // Convert to Liquid array and back - let liquid_array = LiquidFloatArray::<$type>::from_arrow_array(array.clone()); - let result_array = liquid_array.to_arrow_array(); - let bytes_array = - LiquidFloatArray::<$type>::from_bytes(liquid_array.to_bytes().into()); - - assert_eq!(result_array.as_ref(), &array); - assert_eq!(bytes_array.to_arrow_array().as_ref(), &array); - } - }; - } - - // Test cases for Float32 - test_roundtrip!( - test_float32_roundtrip_basic, - Float32Type, - vec![Some(-1.0), Some(1.0), Some(0.0)] - ); - - test_roundtrip!( - test_float32_roundtrip_with_nones, - Float32Type, - vec![Some(-1.0), Some(1.0), Some(0.0), None] - ); - - test_roundtrip!( - test_float32_roundtrip_all_nones, - Float32Type, - vec![None, None, None, None] - ); - - test_roundtrip!(test_float32_roundtrip_empty, Float32Type, vec![]); - - // Test cases for Float64 - test_roundtrip!( - test_float64_roundtrip_basic, - Float64Type, - vec![Some(-1.0), Some(1.0), Some(0.0)] - ); - - test_roundtrip!( - test_float64_roundtrip_with_nones, - Float64Type, - vec![Some(-1.0), Some(1.0), Some(0.0), None] - ); - - test_roundtrip!( - test_float64_roundtrip_all_nones, - Float64Type, - vec![None, None, None, None] - ); - - test_roundtrip!(test_float64_roundtrip_empty, Float64Type, vec![]); - - // Tests with ilters - #[test] - fn test_filter_basic() { - // Create original array with some values - let original = vec![Some(1.0), Some(2.1), Some(3.2), None, Some(5.5)]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidFloatArray::::from_arrow_array(array); - - // Create selection mask: keep indices 0, 2, and 4 - let selection = BooleanBuffer::from(vec![true, false, true, false, true]); - - // Apply filter - let result_array = liquid_array.filter(&selection); - - // Expected result after filtering - let expected = PrimitiveArray::::from(vec![Some(1.0), Some(3.2), Some(5.5)]); - - assert_eq!(result_array.as_ref(), &expected); - } - - #[test] - fn test_original_arrow_data_type_returns_float32() { - let array = PrimitiveArray::::from(vec![Some(1.0), Some(2.5)]); - let liquid = LiquidFloatArray::::from_arrow_array(array); - assert_eq!(liquid.original_arrow_data_type(), DataType::Float32); - } - - #[test] - fn test_filter_all_nulls() { - // Create array with all nulls - let original = vec![None, None, None, None]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidFloatArray::::from_arrow_array(array); - - // Keep first and last elements - let selection = BooleanBuffer::from(vec![true, false, false, true]); - - let result_array = liquid_array.filter(&selection); - - let expected = PrimitiveArray::::from(vec![None, None]); - - assert_eq!(result_array.as_ref(), &expected); - } - - #[test] - fn test_filter_empty_result() { - let original = vec![Some(1.0), Some(2.1), Some(3.3)]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidFloatArray::::from_arrow_array(array); - - // Filter out all elements - let selection = BooleanBuffer::from(vec![false, false, false]); - - let result_array = liquid_array.filter(&selection); - - assert_eq!(result_array.len(), 0); - } - - #[test] - fn test_compression_f32_f64() { - fn run_compression_test( - type_name: &str, - data_fn: impl Fn(usize) -> T::Native, - ) { - let original: Vec = (0..2000).map(data_fn).collect(); - let array = PrimitiveArray::::from_iter_values(original); - let uncompressed_size = array.get_array_memory_size(); - - let liquid_array = LiquidFloatArray::::from_arrow_array(array); - let compressed_size = liquid_array.get_array_memory_size(); - - println!( - "Type: {type_name}, uncompressed_size: {uncompressed_size}, compressed_size: {compressed_size}" - ); - // Assert that compression actually reduced the size - assert!( - compressed_size < uncompressed_size, - "{type_name} compression failed to reduce size" - ); - } - - // Run for f32 - run_compression_test::("f32", |i| i as f32); - - // Run for f64 - run_compression_test::("f64", |i| i as f64); - } -} diff --git a/src/core/src/liquid_array/ipc.rs b/src/core/src/liquid_array/ipc.rs deleted file mode 100644 index 2b8682402..000000000 --- a/src/core/src/liquid_array/ipc.rs +++ /dev/null @@ -1,690 +0,0 @@ -//! IPC for liquid array. - -use std::mem::size_of; -use std::sync::Arc; - -use arrow::array::ArrowPrimitiveType; -use arrow::datatypes::{ - Date32Type, Date64Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, - TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, - TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, -}; -use bytes::Bytes; -use fsst::Compressor; - -use crate::liquid_array::LiquidByteViewArray; -use crate::liquid_array::LiquidDecimalArray; -use crate::liquid_array::LiquidPrimitiveArray; -use crate::liquid_array::raw::FsstArray; - -use super::linear_integer_array::LiquidLinearArray; -use super::{LiquidArrayRef, LiquidDataType, LiquidFixedLenByteArray, LiquidFloatArray}; - -const MAGIC: u32 = 0x4C51_4441; // "LQDA" for LiQuid Data Array -const VERSION: u16 = 1; - -macro_rules! primitive_physical_type_entries { - ($macro:ident) => { - $macro!([ - (Int8, Int8Type, 0, Integer), - (Int16, Int16Type, 1, Integer), - (Int32, Int32Type, 2, Integer), - (Int64, Int64Type, 3, Integer), - (UInt8, UInt8Type, 4, Integer), - (UInt16, UInt16Type, 5, Integer), - (UInt32, UInt32Type, 6, Integer), - (UInt64, UInt64Type, 7, Integer), - (Float32, Float32Type, 8, Float), - (Float64, Float64Type, 9, Float), - (Date32, Date32Type, 10, Integer), - (Date64, Date64Type, 11, Integer), - (TimestampSecond, TimestampSecondType, 12, Integer), - (TimestampMillisecond, TimestampMillisecondType, 13, Integer), - (TimestampMicrosecond, TimestampMicrosecondType, 14, Integer), - (TimestampNanosecond, TimestampNanosecondType, 15, Integer) - ]); - }; -} - -macro_rules! physical_type_integer_body { - (Integer, $arrow_ty:ty, $bytes:expr, $self:expr) => { - Arc::new(LiquidPrimitiveArray::<$arrow_ty>::from_bytes($bytes)) as LiquidArrayRef - }; - (Float, $arrow_ty:ty, $bytes:expr, $self:expr) => { - panic!( - "Physical type {:?} cannot be decoded as an integer array", - $self - ) - }; -} - -macro_rules! physical_type_linear_body { - (Integer, $arrow_ty:ty, $bytes:expr, $self:expr) => { - Arc::new(LiquidLinearArray::<$arrow_ty>::from_bytes($bytes)) as LiquidArrayRef - }; - (Float, $arrow_ty:ty, $bytes:expr, $self:expr) => { - panic!( - "Physical type {:?} cannot be decoded as a linear integer array", - $self - ) - }; -} - -macro_rules! physical_type_float_body { - (Float, $arrow_ty:ty, $bytes:expr, $self:expr) => { - Arc::new(LiquidFloatArray::<$arrow_ty>::from_bytes($bytes)) as LiquidArrayRef - }; - (Integer, $arrow_ty:ty, $bytes:expr, $self:expr) => { - panic!( - "Physical type {:?} cannot be decoded as a float array", - $self - ) - }; -} - -macro_rules! define_physical_types { - ( [ $(($variant:ident, $arrow_ty:ty, $id:expr, $category:ident)),+ $(,)? ] ) => { - /// Physical primitive types supported by Liquid IPC. - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - #[allow(missing_docs)] - #[repr(u16)] - pub enum PrimitivePhysicalType { - $( $variant = $id, )+ - } - - /// Marker trait implemented for Arrow primitive types that have a Liquid physical ID. - pub trait PhysicalTypeMarker: ArrowPrimitiveType { - /// The physical type associated with the Arrow primitive. - const PHYSICAL_TYPE: PrimitivePhysicalType; - } - - $(impl PhysicalTypeMarker for $arrow_ty { - const PHYSICAL_TYPE: PrimitivePhysicalType = PrimitivePhysicalType::$variant; - })+ - - impl PrimitivePhysicalType { - fn from_arrow_type() -> PrimitivePhysicalType - where - T: ArrowPrimitiveType + PhysicalTypeMarker, - { - T::PHYSICAL_TYPE - } - - fn deserialize_integer(self, bytes: Bytes) -> LiquidArrayRef { - match self { - $( PrimitivePhysicalType::$variant => { - physical_type_integer_body!($category, $arrow_ty, bytes, self) - }, )+ - } - } - - fn deserialize_linear_integer(self, bytes: Bytes) -> LiquidArrayRef { - match self { - $( PrimitivePhysicalType::$variant => { - physical_type_linear_body!($category, $arrow_ty, bytes, self) - }, )+ - } - } - - fn deserialize_float(self, bytes: Bytes) -> LiquidArrayRef { - match self { - $( PrimitivePhysicalType::$variant => { - physical_type_float_body!($category, $arrow_ty, bytes, self) - }, )+ - } - } - } - - impl TryFrom for PrimitivePhysicalType { - type Error = u16; - - fn try_from(value: u16) -> Result { - match value { - $( $id => Ok(PrimitivePhysicalType::$variant), )+ - _ => Err(value), - } - } - } - }; -} - -primitive_physical_type_entries!(define_physical_types); - -fn expect_physical_type(id: u16, label: &str) -> PrimitivePhysicalType { - PrimitivePhysicalType::try_from(id) - .unwrap_or_else(|value| panic!("Unsupported {label} physical type: {value}")) -} - -/* - +--------------------------------------------------+ - | LiquidIPCHeader (16 bytes) | - +--------------------------------------------------+ - | MAGIC (4 bytes) | // Offset 0..3: "LQDA" magic number (0x4C51_4441) - +--------------------------------------------------+ - | VERSION (2 bytes) | // Offset 4..5: Version (currently 1) - +--------------------------------------------------+ - | logical_type_id (2 bytes) | // Offset 6..7: Logical type identifier (e.g. Integer) - +--------------------------------------------------+ - | physical_type_id (2 bytes) | // Offset 8..9: Physical type identifier for T - +--------------------------------------------------+ - | __padding (6 bytes) | // Offset 10..15: Padding to ensure 16 byte header - +--------------------------------------------------+ -*/ -#[repr(C)] -pub(super) struct LiquidIPCHeader { - pub(super) magic: [u8; 4], - pub(super) version: u16, - pub(super) logical_type_id: u16, - pub(super) physical_type_id: u16, - pub(super) __padding: [u8; 6], -} - -const _: () = assert!(size_of::() == LiquidIPCHeader::size()); - -impl LiquidIPCHeader { - pub(super) const fn size() -> usize { - 16 - } - - pub(super) fn new(logical_type_id: u16, physical_type_id: u16) -> Self { - Self { - magic: MAGIC.to_le_bytes(), - version: VERSION, - logical_type_id, - physical_type_id, - __padding: [0; 6], - } - } - - pub(super) fn to_bytes(&self) -> [u8; Self::size()] { - let mut bytes = [0; Self::size()]; - bytes[0..4].copy_from_slice(&self.magic); - bytes[4..6].copy_from_slice(&self.version.to_le_bytes()); - bytes[6..8].copy_from_slice(&self.logical_type_id.to_le_bytes()); - bytes[8..10].copy_from_slice(&self.physical_type_id.to_le_bytes()); - bytes - } - - pub(super) fn from_bytes(bytes: &[u8]) -> Self { - if bytes.len() < Self::size() { - panic!( - "value too small for LiquidIPCHeader, expected at least {} bytes, got {}", - Self::size(), - bytes.len() - ); - } - let magic = bytes[0..4].try_into().unwrap(); - let version = u16::from_le_bytes(bytes[4..6].try_into().unwrap()); - let logical_type_id = u16::from_le_bytes(bytes[6..8].try_into().unwrap()); - let physical_type_id = u16::from_le_bytes(bytes[8..10].try_into().unwrap()); - - if magic != MAGIC.to_le_bytes() { - panic!("Invalid magic number"); - } - if version != VERSION { - panic!("Unsupported version"); - } - - Self { - magic, - version, - logical_type_id, - physical_type_id, - __padding: [0; 6], - } - } -} - -/// Context for liquid IPC. -pub struct LiquidIPCContext { - compressor: Option>, -} - -impl LiquidIPCContext { - /// Create a new instance of LiquidIPCContext. - pub fn new(compressor: Option>) -> Self { - Self { compressor } - } -} - -/// Read a liquid array from bytes. -pub fn read_from_bytes(bytes: Bytes, context: &LiquidIPCContext) -> LiquidArrayRef { - let header = LiquidIPCHeader::from_bytes(&bytes); - let logical_type = LiquidDataType::from(header.logical_type_id); - match logical_type { - LiquidDataType::Integer => { - let physical_type = expect_physical_type(header.physical_type_id, "integer"); - physical_type.deserialize_integer(bytes) - } - LiquidDataType::ByteViewArray => { - let compressor = context.compressor.as_ref().expect("Expected a compressor"); - Arc::new(LiquidByteViewArray::::from_bytes( - bytes, - compressor.clone(), - )) - } - LiquidDataType::Float => { - let physical_type = expect_physical_type(header.physical_type_id, "float"); - physical_type.deserialize_float(bytes) - } - LiquidDataType::FixedLenByteArray => { - let compressor = context.compressor.as_ref().expect("Expected a compressor"); - Arc::new(LiquidFixedLenByteArray::from_bytes( - bytes, - compressor.clone(), - )) - } - LiquidDataType::LinearInteger => { - let physical_type = expect_physical_type(header.physical_type_id, "linear-integer"); - physical_type.deserialize_linear_integer(bytes) - } - LiquidDataType::Decimal => Arc::new(LiquidDecimalArray::from_bytes(bytes)), - } -} - -pub(super) fn get_physical_type_id() -> u16 -where - T: ArrowPrimitiveType + PhysicalTypeMarker, -{ - PrimitivePhysicalType::from_arrow_type::() as u16 -} - -#[cfg(test)] -mod tests { - use arrow::{ - array::{AsArray, BinaryViewArray, PrimitiveArray, StringArray}, - datatypes::{ - Decimal128Type, Decimal256Type, DecimalType, Int32Type, TimestampMicrosecondType, - TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, i256, - }, - }; - use arrow_schema::DataType; - - use crate::liquid_array::raw::FsstArray; - use crate::liquid_array::{LiquidArray, utils::gen_test_decimal_array}; - - use super::*; - - #[test] - fn test_to_bytes() { - // Create a simple array - let original: Vec> = vec![Some(10), Some(20), Some(30), None, Some(50)]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - - // Serialize to bytes - let bytes = liquid_array.to_bytes_inner(); - - // Basic validation - let header = LiquidIPCHeader::from_bytes(&bytes); - assert_eq!( - header.magic, - MAGIC.to_le_bytes(), - "Magic number should be LQDA" - ); - assert_eq!(header.version, VERSION, "Version should be 1"); - assert_eq!( - header.physical_type_id, 2, - "Type ID for Int32Type should be 2" - ); - assert_eq!( - header.logical_type_id, - LiquidDataType::Integer as u16, - "Logical type ID should be 1" - ); - - // Check that the total size makes sense (we can't predict the exact size without knowing bit_width) - assert!( - bytes.len() > 100, - "Serialized data should have a reasonable size" - ); - } - - #[test] - fn test_roundtrip_bytes() { - let original: Vec> = vec![Some(10), Some(20), Some(30), None, Some(50)]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - - let deserialized_array = LiquidPrimitiveArray::::from_bytes(bytes); - - let result_array = deserialized_array.to_arrow_array(); - - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_roundtrip_edge_cases() { - // Test various edge cases - - // 1. All nulls array - let all_nulls: Vec> = vec![None; 1000]; - let array = PrimitiveArray::::from(all_nulls); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // 2. No nulls array - let no_nulls: Vec> = (0..1000).map(Some).collect(); - let array = PrimitiveArray::::from(no_nulls); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // 3. Single value array - let single_value: Vec> = vec![Some(42)]; - let array = PrimitiveArray::::from(single_value); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // 4. Empty array - let empty: Vec> = vec![]; - let array = PrimitiveArray::::from(empty); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // 5. Large array with very sparse nulls - let sparse_nulls: Vec> = (0..10_000) - .map(|i| { - if i == 1000 || i == 5000 || i == 9000 { - None - } else { - Some(i) - } - }) - .collect(); - let array = PrimitiveArray::::from(sparse_nulls); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - } - - #[test] - fn test_roundtrip_multiple_data_types() { - use arrow::datatypes::{Int16Type, UInt32Type, UInt64Type}; - - // Test with Int16Type - let i16_values: Vec> = (0..2000) - .map(|i| { - if i % 11 == 0 { - None - } else { - Some((i % 300 - 150) as i16) - } - }) - .collect(); - let array = PrimitiveArray::::from(i16_values); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // Test with UInt32Type - let u32_values: Vec> = (0..2000) - .map(|i| { - if i % 13 == 0 { - None - } else { - Some(i as u32 * 10000) - } - }) - .collect(); - let array = PrimitiveArray::::from(u32_values); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - - // Test with UInt64Type - let u64_values: Vec> = (0..2000) - .map(|i| { - if i % 17 == 0 { - None - } else { - Some(u64::MAX - (i as u64 * 1000000)) - } - }) - .collect(); - let array = PrimitiveArray::::from(u64_values); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidPrimitiveArray::::from_bytes(bytes); - let result = deserialized.to_arrow_array(); - assert_eq!(result.as_ref(), &array); - } - - #[test] - fn test_date_types_ipc_roundtrip() { - // Test Date32Type - let date32_array = PrimitiveArray::::from(vec![Some(18628), None, Some(0)]); - let liquid_array = - LiquidPrimitiveArray::::from_arrow_array(date32_array.clone()); - let bytes = Bytes::from(liquid_array.to_bytes()); - let context = LiquidIPCContext::new(None); - let deserialized = read_from_bytes(bytes, &context); - assert_eq!(deserialized.to_arrow_array().as_ref(), &date32_array); - - // Test Date64Type - let date64_array = - PrimitiveArray::::from(vec![Some(1609459200000), None, Some(0)]); - let liquid_array = - LiquidPrimitiveArray::::from_arrow_array(date64_array.clone()); - let bytes = Bytes::from(liquid_array.to_bytes()); - let context = LiquidIPCContext::new(None); - let deserialized = read_from_bytes(bytes, &context); - assert_eq!(deserialized.to_arrow_array().as_ref(), &date64_array); - } - - #[test] - fn test_ipc_roundtrip_utf8_byte_view() { - let input = StringArray::from(vec![ - Some("hello"), - Some("world"), - None, - Some("liquid"), - Some("byte"), - Some("array"), - Some("hello"), - ]); - - // LiquidByteViewArray - let compressor_bv = LiquidByteViewArray::::train_compressor(input.iter()); - let original_bv = - LiquidByteViewArray::::from_string_array(&input, compressor_bv.clone()); - let bytes_bv = Bytes::from(original_bv.to_bytes()); - let deserialized_bv = LiquidByteViewArray::::from_bytes(bytes_bv, compressor_bv); - let output_bv = deserialized_bv.to_arrow_array(); - assert_eq!(output_bv.as_string::(), &input); - } - - #[test] - fn test_ipc_roundtrip_binaryview_byte_view() { - let input = BinaryViewArray::from(vec![ - Some(b"hello".as_slice()), - Some(b"world".as_slice()), - Some(b"hello".as_slice()), - Some(b"rust\x00".as_slice()), - None, - Some(b"This is a very long string that should be compressed well"), - Some(b""), - Some(b"This is a very long string that should be compressed well"), - ]); - - // LiquidByteViewArray via BinaryView - let (compressor_bv, original_bv) = - LiquidByteViewArray::::train_from_binary_view(&input); - let bytes_bv = Bytes::from(original_bv.to_bytes()); - let deserialized_bv = LiquidByteViewArray::::from_bytes(bytes_bv, compressor_bv); - let output_bv = deserialized_bv.to_arrow_array(); - assert_eq!(output_bv.as_binary_view(), &input); - } - - #[test] - fn test_float32_array_roundtrip() { - let arr = PrimitiveArray::::from(vec![ - Some(-1.3e7), - Some(1.9), - Some(6.6e4), - None, - Some(9.1e-5), - ]); - let original = LiquidFloatArray::::from_arrow_array(arr.clone()); - let serialized = Bytes::from(original.to_bytes_inner()); - let deserialized = LiquidFloatArray::::from_bytes(serialized).to_arrow_array(); - assert_eq!(deserialized.as_ref(), &arr); - } - - #[test] - fn test_float64_array_roundtrip() { - let arr = PrimitiveArray::::from(vec![ - Some(-1.3e7), - Some(1.9), - Some(6.6e4), - None, - Some(9.1e-5), - ]); - let original = LiquidFloatArray::::from_arrow_array(arr.clone()); - let serialized = Bytes::from(original.to_bytes_inner()); - let deserialized = LiquidFloatArray::::from_bytes(serialized).to_arrow_array(); - assert_eq!(deserialized.as_ref(), &arr); - } - - fn test_decimal_roundtrip(data_type: DataType) { - let original_array = gen_test_decimal_array::(data_type); - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&original_array); - - let bytes = liquid_array.to_bytes_inner(); - let bytes = Bytes::from(bytes); - let deserialized = LiquidFixedLenByteArray::from_bytes(bytes, compressor); - let deserialized_arrow = deserialized.to_arrow_array(); - assert_eq!(deserialized_arrow.as_ref(), &original_array); - } - - #[test] - fn test_decimal128_array_roundtrip() { - test_decimal_roundtrip::(DataType::Decimal128(10, 2)); - } - - #[test] - fn test_decimal256_array_roundtrip() { - test_decimal_roundtrip::(DataType::Decimal256(38, 6)); - } - - #[test] - fn test_fixed_len_byte_array_ipc_roundtrip() { - // Test both Decimal128 and Decimal256 through the full IPC pipeline - - let decimal128_array = - gen_test_decimal_array::(DataType::Decimal128(15, 3)); - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&decimal128_array); - - let bytes = liquid_array.to_bytes(); - let bytes = Bytes::from(bytes); - - let context = LiquidIPCContext::new(Some(compressor.clone())); - let deserialized_ref = read_from_bytes(bytes, &context); - assert!(matches!( - deserialized_ref.data_type(), - LiquidDataType::FixedLenByteArray - )); - let result_arrow = deserialized_ref.to_arrow_array(); - assert_eq!(result_arrow.as_ref(), &decimal128_array); - - // Test Decimal256 - let decimal256_array = - gen_test_decimal_array::(DataType::Decimal256(38, 6)); - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&decimal256_array); - - let bytes = liquid_array.to_bytes(); - let bytes = Bytes::from(bytes); - - let context = LiquidIPCContext::new(Some(compressor.clone())); - let deserialized_ref = read_from_bytes(bytes, &context); - - assert!(matches!( - deserialized_ref.data_type(), - LiquidDataType::FixedLenByteArray - )); - - let result_arrow = deserialized_ref.to_arrow_array(); - assert_eq!(result_arrow.as_ref(), &decimal256_array); - } - - #[test] - fn test_fixed_len_byte_array_ipc_edge_cases() { - // Test edge cases with FixedLenByteArray IPC - - let mut builder = arrow::array::Decimal128Builder::new(); - builder.append_value(123456789_i128); - builder.append_null(); - builder.append_value(-987654321_i128); - builder.append_null(); - builder.append_value(0_i128); - let array_with_nulls = builder.finish().with_precision_and_scale(15, 3).unwrap(); - - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&array_with_nulls); - - let bytes = liquid_array.to_bytes(); - let bytes = Bytes::from(bytes); - - let context = LiquidIPCContext::new(Some(compressor)); - let deserialized_ref = read_from_bytes(bytes, &context); - let result_arrow = deserialized_ref.to_arrow_array(); - - assert_eq!(result_arrow.as_ref(), &array_with_nulls); - - // Test with single value - let mut builder = arrow::array::Decimal256Builder::new(); - builder.append_value(i256::from_i128(42_i128)); - let single_value_array = builder.finish().with_precision_and_scale(38, 6).unwrap(); - - let (compressor, liquid_array) = - LiquidFixedLenByteArray::train_from_decimal_array(&single_value_array); - - let bytes = liquid_array.to_bytes(); - let bytes = Bytes::from(bytes); - - let context = LiquidIPCContext::new(Some(compressor)); - let deserialized_ref = read_from_bytes(bytes, &context); - let result_arrow = deserialized_ref.to_arrow_array(); - - assert_eq!(result_arrow.as_ref(), &single_value_array); - } - - #[test] - fn test_timestamp_physical_type_ids() { - assert_eq!(get_physical_type_id::(), 12); - assert_eq!(get_physical_type_id::(), 13); - assert_eq!(get_physical_type_id::(), 14); - assert_eq!(get_physical_type_id::(), 15); - } -} diff --git a/src/core/src/liquid_array/linear_integer_array.rs b/src/core/src/liquid_array/linear_integer_array.rs deleted file mode 100644 index 4a8bc14ed..000000000 --- a/src/core/src/liquid_array/linear_integer_array.rs +++ /dev/null @@ -1,727 +0,0 @@ -use std::any::Any; -use std::fmt::Debug; -use std::marker::PhantomData; -use std::sync::Arc; - -use super::PrimitiveKind; -use super::{LiquidArray, LiquidDataType, LiquidPrimitiveType}; -use crate::cache::LiquidExpr; -use crate::liquid_array::LiquidPrimitiveArray; -use crate::liquid_array::eval_predicate_on_array; -use crate::liquid_array::ipc::{LiquidIPCHeader, get_physical_type_id}; -use arrow::array::{ - Array, ArrayRef, ArrowPrimitiveType, BooleanArray, PrimitiveArray, - cast::AsArray, - types::{ - Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, - UInt32Type, UInt64Type, - }, -}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; -use arrow::compute::kernels::filter; -use arrow_schema::DataType; -use bytes::Bytes; -use num_traits::{AsPrimitive, Bounded, FromPrimitive}; - -/// A linear-model based integer array, **only use it when you know the array is monotonic** and **you don't care about encoding speed!**. -/// -/// Under the hood, it uses a linear model to predict the values and store the residuals: -/// value\[i\] = intercept + round(slope * i) + residual\[i\] -/// -/// Where `intercept` and `slope` are computed using a L-infinity linear fit, **this is time-consuming!**. -/// -/// This array is only recommended if you know the array follows a linear model, e.g., kinetic values, offsets, etc. -/// -/// Examples of not recommended use cases: random values like ids, categorical values, etc. -#[derive(Debug)] -pub struct LiquidLinearArray -where - T::Native: AsPrimitive + FromPrimitive + Bounded, -{ - // Signed residuals, bit-packed as a Liquid primitive array of i64. - residuals: LiquidPrimitiveArray, - // Intercept term stored as f64 for simpler math/IO. - intercept: f64, - // Slope term of the linear model. - slope: f64, - // Keep the logical type parameter. - _phantom: PhantomData, -} - -/// Backward-compatible alias for i32. -pub type LiquidLinearI32Array = LiquidLinearArray; -/// Linear-model array for `i8`. -pub type LiquidLinearI8Array = LiquidLinearArray; -/// Linear-model array for `i16`. -pub type LiquidLinearI16Array = LiquidLinearArray; -/// Linear-model array for `i64`. -pub type LiquidLinearI64Array = LiquidLinearArray; -/// Linear-model array for `u8`. -pub type LiquidLinearU8Array = LiquidLinearArray; -/// Linear-model array for `u16`. -pub type LiquidLinearU16Array = LiquidLinearArray; -/// Linear-model array for `u32`. -pub type LiquidLinearU32Array = LiquidLinearArray; -/// Linear-model array for `u64`. -pub type LiquidLinearU64Array = LiquidLinearArray; -/// Linear-model array for `Date32` (days since epoch). -pub type LiquidLinearDate32Array = LiquidLinearArray; -/// Linear-model array for `Date64` (ms since epoch). -pub type LiquidLinearDate64Array = LiquidLinearArray; - -impl LiquidLinearArray -where - T: LiquidPrimitiveType, - T::Native: AsPrimitive + FromPrimitive + Bounded, -{ - /// Build from an Arrow `PrimitiveArray` by training a linear model - /// using a fast L-infinity fit (Option 3) and storing residuals. - pub fn from_arrow_array(arrow_array: PrimitiveArray) -> Self { - let len = arrow_array.len(); - - // All nulls - if arrow_array.null_count() == len { - // All nulls - let res = PrimitiveArray::::new_null(len); - return Self { - residuals: LiquidPrimitiveArray::::from_arrow_array(res), - intercept: 0.0, // arbitrary, unused since all nulls - slope: 0.0, - _phantom: PhantomData, - }; - } - - // Prepare compact non-null buffers for fast fitting (avoid iterator Option cost). - let (nn_values, nn_indices) = collect_non_null_f64_and_indices::(&arrow_array); - - // Option 3 parameters: L-infinity (Chebyshev) regression. - let (mut intercept, mut slope) = fit_linf(&nn_values, &nn_indices); - - // Compute residuals in one unified loop; also track ranges for fallback decision. - let mut residuals: Vec = Vec::with_capacity(len); - let is_unsigned = ::IS_UNSIGNED; - let vals = arrow_array.values(); - let nulls_opt = arrow_array.nulls(); - - // Original value range - let mut orig_min_u64 = u64::MAX; - let mut orig_max_u64 = 0u64; - let mut orig_min_i64 = i64::MAX; - let mut orig_max_i64 = i64::MIN; - // Residual range - let mut res_min = i64::MAX; - let mut res_max = i64::MIN; - - if is_unsigned { - let max_u64: u64 = ::MAX_U64; - for i in 0..len { - let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); - if valid { - type U = - <::UnSignedType as ArrowPrimitiveType>::Native; - let v_u: U = vals[i].as_(); - let v_u64: u64 = v_u.as_(); - if v_u64 < orig_min_u64 { - orig_min_u64 = v_u64; - } - if v_u64 > orig_max_u64 { - orig_max_u64 = v_u64; - } - let pr = slope * (i as f64) + intercept; - let p = predict_u64_saturated(pr, max_u64); - let (pos, mag) = if v_u64 >= p { - (true, v_u64 - p) - } else { - (false, p - v_u64) - }; - let m = (mag & (i64::MAX as u64)) as i64; - let r = if pos { m } else { -m }; - if r < res_min { - res_min = r; - } - if r > res_max { - res_max = r; - } - residuals.push(r); - } else { - residuals.push(0); - } - } - } else { - let (min_i64, max_i64): (i64, i64) = - (::MIN_I64, ::MAX_I64); - for i in 0..len { - let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); - if valid { - let v_i64: i64 = vals[i].as_(); - if v_i64 < orig_min_i64 { - orig_min_i64 = v_i64; - } - if v_i64 > orig_max_i64 { - orig_max_i64 = v_i64; - } - let pr = slope * (i as f64) + intercept; - let p = predict_i64_saturated(pr, min_i64, max_i64); - let r = v_i64 - p; - if r < res_min { - res_min = r; - } - if r > res_max { - res_max = r; - } - residuals.push(r); - } else { - residuals.push(0); - } - } - } - - // Fallback: ensure residual range is strictly smaller than original range - let res_width: u128 = (res_max as i128 - res_min as i128) as u128; - let orig_width: u128 = if is_unsigned { - (orig_max_u64 as u128).saturating_sub(orig_min_u64 as u128) - } else { - (orig_max_i64 as i128 - orig_min_i64 as i128) as u128 - }; - if res_width >= orig_width { - // Rebuild residuals with zero model - intercept = 0.0; - slope = 0.0; - residuals.clear(); - if is_unsigned { - for i in 0..len { - let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); - if valid { - type U = <::UnSignedType as ArrowPrimitiveType>::Native; - let v_u: U = vals[i].as_(); - let v_u64: u64 = v_u.as_(); - let r = (v_u64 & (i64::MAX as u64)) as i64; - residuals.push(r); - } else { - residuals.push(0); - } - } - } else { - for i in 0..len { - let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); - if valid { - let v_i64: i64 = vals[i].as_(); - residuals.push(v_i64); - } else { - residuals.push(0); - } - } - } - } - let residuals_buf: ScalarBuffer = ScalarBuffer::from(residuals); - let nulls = arrow_array.nulls().cloned(); - let res_prim = PrimitiveArray::::new(residuals_buf, nulls); - let residuals = LiquidPrimitiveArray::::from_arrow_array(res_prim); - - Self { - residuals, - intercept, - slope, - _phantom: PhantomData, - } - } - - fn len(&self) -> usize { - self.residuals.len() - } - - fn residual_starting_loc() -> usize { - // Header + intercept(native) + slope(f64), aligned to 8 bytes boundary - let header_size = - LiquidIPCHeader::size() + std::mem::size_of::() + std::mem::size_of::(); - (header_size + 7) & !7 - } - - fn to_bytes_inner(&self) -> Vec { - let header = LiquidIPCHeader::new( - LiquidDataType::LinearInteger as u16, - get_physical_type_id::(), - ); - let start = Self::residual_starting_loc(); - let mut out = Vec::with_capacity(start + 256); - - // Header - out.extend_from_slice(&header.to_bytes()); - // Model params (f64 intercept, then f64 slope) - out.extend_from_slice(&self.intercept.to_le_bytes()); - out.extend_from_slice(&self.slope.to_le_bytes()); - - while out.len() < start { - out.push(0); - } - - // Encode residuals (already LiquidPrimitiveArray) - out.extend_from_slice(&self.residuals.to_bytes_inner()); - out - } - - /// Decode a `LiquidLinearArray` from bytes. - pub fn from_bytes(bytes: Bytes) -> Self { - let _hdr = LiquidIPCHeader::from_bytes(&bytes); - - // Read intercept as f64 - let intercept_off = LiquidIPCHeader::size(); - let intercept = - f64::from_le_bytes(bytes[intercept_off..intercept_off + 8].try_into().unwrap()); - - // Read slope - let slope_off = intercept_off + std::mem::size_of::(); - let slope = f64::from_le_bytes(bytes[slope_off..slope_off + 8].try_into().unwrap()); - - // Decode residuals - let start = Self::residual_starting_loc(); - let res_bytes = bytes.slice(start..); - let residuals = LiquidPrimitiveArray::::from_bytes(res_bytes); - - Self { - residuals, - intercept, - slope, - _phantom: PhantomData, - } - } -} - -impl LiquidArray for LiquidLinearArray -where - T: LiquidPrimitiveType, - T::Native: AsPrimitive + FromPrimitive + Bounded, -{ - fn as_any(&self) -> &dyn Any { - self - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn get_array_memory_size(&self) -> usize { - self.residuals.get_array_memory_size() - + std::mem::size_of::() // intercept - + std::mem::size_of::() // slope - } - - fn len(&self) -> usize { - self.len() - } - - fn to_arrow_array(&self) -> ArrayRef { - let arr = self.residuals.to_arrow_array(); - let (_dt, residuals, nulls) = arr.as_primitive::().clone().into_parts(); - - // Reconstruct final values: predicted(i) +/- |residual_i| - let mut final_values = Vec::::with_capacity(self.len()); - let is_unsigned = ::IS_UNSIGNED; - if is_unsigned { - let max_u64: u64 = ::MAX_U64; - for (i, &e) in residuals.iter().enumerate() { - let pr = self.slope * (i as f64) + self.intercept; - let p = predict_u64_saturated(pr, max_u64); - let mag = e.unsigned_abs(); - let sum = if e >= 0 { - p.saturating_add(mag) - } else { - p.saturating_sub(mag) - }; - final_values.push(T::Native::from_u64(sum).unwrap()); - } - } else { - let (min_i64, max_i64): (i64, i64) = - (::MIN_I64, ::MAX_I64); - for (i, &e) in residuals.iter().enumerate() { - let pr = self.slope * (i as f64) + self.intercept; - let p = predict_i64_saturated(pr, min_i64, max_i64); - let sum = p.saturating_add(e); - final_values.push(T::Native::from_i64(sum).unwrap()); - } - } - - let values_buf: ScalarBuffer = ScalarBuffer::from(final_values); - Arc::new(PrimitiveArray::::new(values_buf, nulls)) - } - - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let arr = self.to_arrow_array(); - let selection = BooleanArray::new(selection.clone(), None); - filter::filter(&arr, &selection).unwrap() - } - - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - let arr = self.filter(filter); - eval_predicate_on_array(arr, predicate) - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::LinearInteger - } -} - -#[inline] -fn predict_u64_saturated(pred: f64, max_u64: u64) -> u64 { - if !pred.is_finite() || pred <= 0.0 { - 0 - } else if pred >= max_u64 as f64 { - max_u64 - } else { - pred.round() as u64 - } -} - -#[inline] -fn predict_i64_saturated(pred: f64, min_i64: i64, max_i64: i64) -> i64 { - if !pred.is_finite() { - 0 - } else if pred <= min_i64 as f64 { - min_i64 - } else if pred >= max_i64 as f64 { - max_i64 - } else { - pred.round() as i64 - } -} - -/// L-infinity linear fit for y\[i\] ≈ intercept + slope * i (before rounding), -/// minimizing the maximum absolute error over all non-null points. -/// -/// Approach: minimize R(m) = max_i (y_i - m i) - min_i (y_i - m i), which is -/// convex in m. Given m, the best intercept is b = (max_i s_i + min_i s_i)/2 -/// where s_i = y_i - m i. We find m by a few rounds of bisection using the -/// subgradient sign derived from the argmax/argmin indices. O(n) per iteration. -fn fit_linf(values: &[f64], idxs: &[u32]) -> (f64, f64) { - let n = values.len(); - assert_eq!(values.len(), idxs.len()); - if n == 0 { - return (0.0, 0.0); - } - if n == 1 { - return (values[0], 0.0); - } - - let mut slope_min = f64::INFINITY; - let mut slope_max = f64::NEG_INFINITY; - for k in 1..n { - let di = (idxs[k] - idxs[k - 1]) as f64; - if di > 0.0 { - let dv = values[k] - values[k - 1]; - let s = dv / di; - if s < slope_min { - slope_min = s; - } - if s > slope_max { - slope_max = s; - } - } - } - if !slope_min.is_finite() || !slope_max.is_finite() { - slope_min = 0.0; - slope_max = 0.0; - } - - let mut lo = slope_min.min(slope_max); - let mut hi = slope_min.max(slope_max); - if (hi - lo).abs() < 1e-12 { - let pad = if hi.abs() < 1.0 { 1.0 } else { hi.abs() * 1e-6 }; - lo -= pad; - hi += pad; - } - - #[inline] - fn range_stats(values: &[f64], idxs: &[u32], m: f64) -> (f64, u32, f64, u32) { - let mut min_s = f64::INFINITY; - let mut max_s = f64::NEG_INFINITY; - let mut i_min = 0u32; - let mut i_max = 0u32; - for k in 0..values.len() { - let i = idxs[k] as f64; - let s = values[k] - m * i; - if s < min_s { - min_s = s; - i_min = idxs[k]; - } - if s > max_s { - max_s = s; - i_max = idxs[k]; - } - } - (min_s, i_min, max_s, i_max) - } - - const MAX_ITERS: usize = 8; - for _ in 0..MAX_ITERS { - let m = 0.5 * (lo + hi); - let (_min_s, i_min, _max_s, i_max) = range_stats(values, idxs, m); - let g = (i_min as i64) - (i_max as i64); - if g > 0 { - hi = m; - } else if g < 0 { - lo = m; - } else { - lo = m; - hi = m; - break; - } - if (hi - lo).abs() < 1e-12 { - break; - } - } - - let m = 0.5 * (lo + hi); - let (min_s, _i_min, max_s, _i_max) = range_stats(values, idxs, m); - let b = 0.5 * (max_s + min_s); - (b, m) -} - -#[inline] -fn collect_non_null_f64_and_indices(arr: &PrimitiveArray) -> (Vec, Vec) -where - T: LiquidPrimitiveType, - T::Native: AsPrimitive, -{ - let nn = arr.len() - arr.null_count(); - let mut values = Vec::with_capacity(nn); - let mut idxs = Vec::with_capacity(nn); - let vals = arr.values(); - if arr.null_count() == 0 { - for (i, v) in vals.iter().enumerate() { - values.push(v.as_()); - idxs.push(i as u32); - } - } else { - let nulls = arr.nulls().unwrap(); - for (i, v) in vals.iter().enumerate() { - if nulls.is_valid(i) { - values.push(v.as_()); - idxs.push(i as u32); - } - } - } - (values, idxs) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn roundtrip_eq(values: Vec>) { - let arr = PrimitiveArray::::from(values.clone()); - let linear = LiquidLinearI32Array::from_arrow_array(arr.clone()); - let decoded = linear.to_arrow_array(); - assert_eq!(decoded.as_ref(), &arr); - - let bytes = Bytes::from(linear.to_bytes()); - let decoded = LiquidLinearI32Array::from_bytes(bytes); - let round = decoded.to_arrow_array(); - assert_eq!(round.as_ref(), &arr); - } - - macro_rules! roundtrip_eq_t { - ($T:ty, $values:expr) => {{ - let arr = PrimitiveArray::<$T>::from(($values).clone()); - let linear = LiquidLinearArray::<$T>::from_arrow_array(arr.clone()); - let decoded = linear.to_arrow_array(); - assert_eq!(decoded.as_ref(), &arr); - - let bytes = Bytes::from(linear.to_bytes()); - let decoded = LiquidLinearArray::<$T>::from_bytes(bytes); - let round = decoded.to_arrow_array(); - assert_eq!(round.as_ref(), &arr); - }}; - } - - #[test] - fn test_roundtrip_basic() { - // Non-monotonic values to ensure we don't rely on simple increasing sequences - roundtrip_eq(vec![ - Some(10), - Some(15), - Some(14), - Some(20), - Some(18), - Some(25), - Some(24), - ]); - } - - #[test] - fn test_roundtrip_with_nulls() { - roundtrip_eq(vec![Some(10), None, Some(30), None, Some(50), Some(70)]); - } - - #[test] - fn test_all_nulls() { - roundtrip_eq(vec![None, None, None, None]); - } - - #[test] - fn test_single_value() { - roundtrip_eq(vec![Some(42)]); - } - - #[test] - fn test_empty() { - roundtrip_eq(vec![]); - } - - #[test] - fn test_negative_values() { - roundtrip_eq(vec![ - Some(-100), - Some(-50), - Some(0), - Some(50), - Some(25), - None, - Some(-25), - ]); - } - - #[test] - fn test_filter_basic() { - let original: Vec> = vec![Some(1), Some(2), Some(3), None, Some(5), Some(8)]; - let arr = PrimitiveArray::::from(original.clone()); - let linear = LiquidLinearI32Array::from_arrow_array(arr); - let selection = BooleanBuffer::from(vec![true, false, true, false, true, false]); - let result = linear.filter(&selection); - let expected = PrimitiveArray::::from(vec![Some(1), Some(3), Some(5)]); - assert_eq!(result.as_ref(), &expected); - } - - #[test] - fn test_original_arrow_data_type_returns_int32() { - let arr = PrimitiveArray::::from(vec![Some(1), Some(2)]); - let linear = LiquidLinearI32Array::from_arrow_array(arr); - assert_eq!(linear.original_arrow_data_type(), DataType::Int32); - } - - #[test] - fn test_roundtrip_i8() { - roundtrip_eq_t!(Int8Type, vec![Some(-10), Some(0), Some(10), None, Some(20)]); - } - - #[test] - fn test_roundtrip_i16() { - roundtrip_eq_t!( - Int16Type, - vec![Some(-1000), Some(0), Some(1000), None, Some(2000)] - ); - } - - #[test] - fn test_roundtrip_i64() { - roundtrip_eq_t!( - Int64Type, - vec![ - Some(-10_000_000_000), - Some(0), - Some(10_000_000_000), - None, - Some(20_000_000_000), - ] - ); - } - - #[test] - fn test_roundtrip_u8() { - roundtrip_eq_t!( - UInt8Type, - vec![Some(0), Some(10), Some(200), None, Some(255)] - ); - } - - #[test] - fn test_roundtrip_u16() { - roundtrip_eq_t!( - UInt16Type, - vec![Some(0), Some(1000), Some(60000), None, Some(500)] - ); - } - - #[test] - fn test_roundtrip_u32() { - roundtrip_eq_t!( - UInt32Type, - vec![ - Some(0), - Some(1_000_000), - Some(3_000_000_000), - None, - Some(123_456_789), - ] - ); - } - - #[test] - fn test_roundtrip_u64() { - roundtrip_eq_t!( - UInt64Type, - vec![ - Some(0), - Some(10_000_000_000), - Some(9_000_000_000_000_000_000u64), - None, - Some(42), - ] - ); - } - - #[test] - fn test_roundtrip_date32() { - roundtrip_eq_t!( - Date32Type, - vec![Some(-365), Some(0), Some(365), None, Some(18262)] - ); - } - - #[test] - fn test_roundtrip_date64() { - roundtrip_eq_t!( - Date64Type, - vec![ - Some(-86_400_000), - Some(0), - Some(86_400_000), - None, - Some(1_000_000_000_000), - ] - ); - } - - #[test] - fn test_compression() { - let original = (0..1_000_000).step_by(100).collect::>(); - - let original = PrimitiveArray::::from_iter_values(original); - let arrow_size = original.get_array_memory_size(); - - let liquid_linear = LiquidLinearI32Array::from_arrow_array(original.clone()); - let liquid_linear_size = liquid_linear.get_array_memory_size(); - - let liquid_primitive = - LiquidPrimitiveArray::::from_arrow_array(original.clone()); - let liquid_primitive_size = liquid_primitive.get_array_memory_size(); - - println!( - "arrow_size: {arrow_size}, liquid_linear_size: {liquid_linear_size}, liquid_primitive_size: {liquid_primitive_size}", - ); - - assert!(liquid_linear_size < arrow_size); - assert!(liquid_primitive_size < arrow_size); - assert!(liquid_linear_size < liquid_primitive_size); - - let original: ArrayRef = Arc::new(original); - assert_eq!(original.as_ref(), liquid_linear.to_arrow_array().as_ref()); - assert_eq!( - original.as_ref(), - liquid_primitive.to_arrow_array().as_ref() - ); - } -} diff --git a/src/core/src/liquid_array/mod.rs b/src/core/src/liquid_array/mod.rs index 637016cb3..4a11d1ed9 100644 --- a/src/core/src/liquid_array/mod.rs +++ b/src/core/src/liquid_array/mod.rs @@ -1,144 +1,24 @@ -//! LiquidArray is the core data structure of LiquidCache. -//! You should not use this module directly. -//! Instead, use `liquid_cache_datafusion_server` or `liquid_cache_datafusion_client` to interact with LiquidCache. -pub mod byte_view_array; -mod decimal_array; -mod fix_len_byte_array; -mod float_array; -pub mod ipc; -mod linear_integer_array; -mod primitive_array; -pub mod raw; -#[cfg(test)] -mod tests; -pub(crate) mod utils; +//! Vortex-backed encoded array used by LiquidCache. -use std::{any::Any, sync::Arc}; +mod array; + +use std::sync::Arc; use arrow::{ array::{ArrayRef, BooleanArray, cast::AsArray}, - buffer::BooleanBuffer, record_batch::RecordBatch, }; -use arrow_schema::{DataType, Field, Schema}; -pub use byte_view_array::LiquidByteViewArray; -pub use decimal_array::LiquidDecimalArray; -pub use fix_len_byte_array::LiquidFixedLenByteArray; -pub use float_array::{LiquidFloat32Array, LiquidFloat64Array, LiquidFloatArray}; -pub use linear_integer_array::{ - LiquidLinearArray, LiquidLinearDate32Array, LiquidLinearDate64Array, LiquidLinearI8Array, - LiquidLinearI16Array, LiquidLinearI32Array, LiquidLinearI64Array, LiquidLinearU8Array, - LiquidLinearU16Array, LiquidLinearU32Array, LiquidLinearU64Array, -}; -pub use primitive_array::{ - LiquidDate32Array, LiquidDate64Array, LiquidI8Array, LiquidI16Array, LiquidI32Array, - LiquidI64Array, LiquidPrimitiveArray, LiquidPrimitiveDeltaArray, LiquidPrimitiveType, - LiquidU8Array, LiquidU16Array, LiquidU32Array, LiquidU64Array, -}; +use arrow_schema::{Field, Schema}; use crate::cache::LiquidExpr; -/// A date or timestamp component observed by lineage analysis. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -pub enum Date32Field { - /// Year component. - Year, - /// Month component. - Month, - /// Day component. - Day, - /// Day of week, where Sunday is zero. - DayOfWeek, -} - -/// Liquid data type is only logical type -#[derive(Debug, Clone, Copy)] -#[repr(u16)] -pub enum LiquidDataType { - /// A byte-view array (dictionary + FSST raw + views). - ByteViewArray = 4, - /// An integer. - Integer = 1, - /// A float. - Float = 2, - /// A fixed length byte array. - FixedLenByteArray = 3, - /// A decimal encoded as a primitive u64 array. - Decimal = 6, - /// A linear-model based integer (signed residuals + model params). - LinearInteger = 5, -} - -impl From for LiquidDataType { - fn from(value: u16) -> Self { - match value { - 4 => LiquidDataType::ByteViewArray, - 1 => LiquidDataType::Integer, - 2 => LiquidDataType::Float, - 3 => LiquidDataType::FixedLenByteArray, - 5 => LiquidDataType::LinearInteger, - 6 => LiquidDataType::Decimal, - _ => panic!("Invalid liquid data type: {value}"), - } - } -} - -/// A Liquid array. -pub trait LiquidArray: std::fmt::Debug + Send + Sync { - /// Get the underlying any type. - fn as_any(&self) -> &dyn Any; - - /// Get the memory size of the Liquid array. - fn get_array_memory_size(&self) -> usize; - - /// Get the length of the Liquid array. - fn len(&self) -> usize; - - /// Check if the Liquid array is empty. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Convert the Liquid array to an Arrow array. - fn to_arrow_array(&self) -> ArrayRef; - - /// Convert the Liquid array to an Arrow array. - /// Except that it will pick the best encoding for the arrow array. - /// Meaning that it may not obey the data type of the original arrow array. - fn to_best_arrow_array(&self) -> ArrayRef { - self.to_arrow_array() - } - - /// Get the logical data type of the Liquid array. - fn data_type(&self) -> LiquidDataType; - - /// Get the original arrow data type of the Liquid array. - fn original_arrow_data_type(&self) -> DataType; - - /// Serialize the Liquid array to a byte array. - fn to_bytes(&self) -> Vec; - - /// Filter the Liquid array with a boolean array and return an **arrow array**. - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let arrow_array = self.to_arrow_array(); - let selection = BooleanArray::new(selection.clone(), None); - arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() - } - - /// Evaluate a predicate on the Liquid array with a filter. - /// - /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. - /// The returned boolean mask is nullable if the the original array is nullable. - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - let filtered = self.filter(filter); - eval_predicate_on_array(filtered, predicate) - } -} +pub use array::LiquidArray; /// A reference to a Liquid array. -pub type LiquidArrayRef = Arc; +pub type LiquidArrayRef = Arc; -fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { +/// Evaluate a validated Liquid predicate on an Arrow array. +pub fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), @@ -154,62 +34,3 @@ fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanAr .expect("predicate output must be an array"); boolean_array.as_boolean().clone() } - -/// Compile-time info about primitive kind (signed vs unsigned) and bounds. -/// Implemented for all Liquid-supported primitive integer and date types. -pub trait PrimitiveKind { - /// Whether the logical type is unsigned (true for u8/u16/u32/u64). - const IS_UNSIGNED: bool; - /// Maximum representable value as u64 for unsigned types (unused for signed). - const MAX_U64: u64; - /// Minimum representable value as i64 for signed/date types (unused for unsigned). - const MIN_I64: i64; - /// Maximum representable value as i64 for signed/date types (unused for unsigned). - const MAX_I64: i64; -} - -macro_rules! impl_unsigned_kind { - ($t:ty, $max:expr) => { - impl PrimitiveKind for $t { - const IS_UNSIGNED: bool = true; - const MAX_U64: u64 = $max as u64; - const MIN_I64: i64 = 0; // unused - const MAX_I64: i64 = 0; // unused - } - }; -} - -macro_rules! impl_signed_kind { - ($t:ty, $min:expr, $max:expr) => { - impl PrimitiveKind for $t { - const IS_UNSIGNED: bool = false; - const MAX_U64: u64 = 0; // unused - const MIN_I64: i64 = $min as i64; - const MAX_I64: i64 = $max as i64; - } - }; -} - -use arrow::datatypes::{ - Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, TimestampMicrosecondType, - TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, - UInt32Type, UInt64Type, -}; - -impl_unsigned_kind!(UInt8Type, u8::MAX); -impl_unsigned_kind!(UInt16Type, u16::MAX); -impl_unsigned_kind!(UInt32Type, u32::MAX); -impl_unsigned_kind!(UInt64Type, u64::MAX); - -impl_signed_kind!(Int8Type, i8::MIN, i8::MAX); -impl_signed_kind!(Int16Type, i16::MIN, i16::MAX); -impl_signed_kind!(Int32Type, i32::MIN, i32::MAX); -impl_signed_kind!(Int64Type, i64::MIN, i64::MAX); - -// Dates are logically signed in Arrow (Date32: i32 days, Date64: i64 ms) -impl_signed_kind!(Date32Type, i32::MIN, i32::MAX); -impl_signed_kind!(Date64Type, i64::MIN, i64::MAX); -impl_signed_kind!(TimestampSecondType, i64::MIN, i64::MAX); -impl_signed_kind!(TimestampMillisecondType, i64::MIN, i64::MAX); -impl_signed_kind!(TimestampMicrosecondType, i64::MIN, i64::MAX); -impl_signed_kind!(TimestampNanosecondType, i64::MIN, i64::MAX); diff --git a/src/core/src/liquid_array/primitive_array.rs b/src/core/src/liquid_array/primitive_array.rs deleted file mode 100644 index c1412618e..000000000 --- a/src/core/src/liquid_array/primitive_array.rs +++ /dev/null @@ -1,879 +0,0 @@ -use bytes::Bytes; -use std::any::Any; -use std::fmt::{Debug, Display}; -use std::sync::Arc; - -use arrow::array::{ - ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray, PrimitiveArray, - types::{ - Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, - TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, - TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, - }, -}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; -use arrow_schema::DataType; -use fastlanes::BitPacking; -use num_traits::{AsPrimitive, FromPrimitive}; - -use super::LiquidDataType; -use crate::cache::LiquidExpr; -use crate::liquid_array::ipc::{LiquidIPCHeader, PhysicalTypeMarker, get_physical_type_id}; -use crate::liquid_array::raw::BitPackedArray; -use crate::liquid_array::{LiquidArray, PrimitiveKind, eval_predicate_on_array}; -use crate::utils::get_bit_width; -use arrow::datatypes::ArrowNativeType; - -mod private { - pub trait Sealed {} -} - -/// LiquidPrimitiveType is a sealed trait that represents the primitive types supported by Liquid. -/// Implemented for all supported integer, date, and timestamp Arrow primitive types. -/// -/// I have to admit this trait is super complicated. -/// Luckily users never have to worry about it, they can just use the types that are already implemented. -/// We could have implemented this as a macro, but macro is ugly. -/// Type is spec, code is proof. -pub trait LiquidPrimitiveType: - ArrowPrimitiveType< - Native: AsPrimitive<::Native> - + AsPrimitive - + FromPrimitive - + Display, - > + Debug - + Send - + Sync - + private::Sealed - + PrimitiveKind - + PhysicalTypeMarker -{ - /// The unsigned type that can be used to represent the signed type. - type UnSignedType: ArrowPrimitiveType + AsPrimitive + BitPacking> - + Debug; -} - -macro_rules! impl_has_unsigned_type { - ($($signed:ty => $unsigned:ty),*) => { - $( - impl private::Sealed for $signed {} - impl LiquidPrimitiveType for $signed { - type UnSignedType = $unsigned; - } - )* - } -} - -impl_has_unsigned_type! { - Int32Type => UInt32Type, - Int64Type => UInt64Type, - Int16Type => UInt16Type, - Int8Type => UInt8Type, - UInt32Type => UInt32Type, - UInt64Type => UInt64Type, - UInt16Type => UInt16Type, - UInt8Type => UInt8Type, - Date64Type => UInt64Type, - Date32Type => UInt32Type, - TimestampSecondType => UInt64Type, - TimestampMillisecondType => UInt64Type, - TimestampMicrosecondType => UInt64Type, - TimestampNanosecondType => UInt64Type -} - -/// Liquid's unsigned 8-bit integer array. -pub type LiquidU8Array = LiquidPrimitiveArray; -/// Liquid's unsigned 16-bit integer array. -pub type LiquidU16Array = LiquidPrimitiveArray; -/// Liquid's unsigned 32-bit integer array. -pub type LiquidU32Array = LiquidPrimitiveArray; -/// Liquid's unsigned 64-bit integer array. -pub type LiquidU64Array = LiquidPrimitiveArray; -/// Liquid's signed 8-bit integer array. -pub type LiquidI8Array = LiquidPrimitiveArray; -/// Liquid's signed 16-bit integer array. -pub type LiquidI16Array = LiquidPrimitiveArray; -/// Liquid's signed 32-bit integer array. -pub type LiquidI32Array = LiquidPrimitiveArray; -/// Liquid's signed 64-bit integer array. -pub type LiquidI64Array = LiquidPrimitiveArray; -/// Liquid's 32-bit date array. -pub type LiquidDate32Array = LiquidPrimitiveArray; -/// Liquid's 64-bit date array. -pub type LiquidDate64Array = LiquidPrimitiveArray; - -/// Liquid's primitive array -#[derive(Debug)] -pub struct LiquidPrimitiveArray { - bit_packed: BitPackedArray, - reference_value: T::Native, -} - -/// Liquid's primitive array which uses delta encoding for compression -#[derive(Debug, Clone)] -pub struct LiquidPrimitiveDeltaArray { - bit_packed: BitPackedArray, - reference_value: T::Native, -} - -impl LiquidPrimitiveArray -where - T: LiquidPrimitiveType, -{ - /// Get the memory size of the Liquid primitive array. - pub fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() + std::mem::size_of::() - } - - /// Get the length of the Liquid primitive array. - pub fn len(&self) -> usize { - self.bit_packed.len() - } - - /// Check if the Liquid primitive array is empty. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Create a Liquid primitive array from an Arrow primitive array. - pub fn from_arrow_array(arrow_array: PrimitiveArray) -> LiquidPrimitiveArray { - let min = match arrow::compute::kernels::aggregate::min(&arrow_array) { - Some(v) => v, - None => { - // entire array is null - return Self { - bit_packed: BitPackedArray::new_null_array(arrow_array.len()), - reference_value: T::Native::ZERO, - }; - } - }; - let max = arrow::compute::kernels::aggregate::max(&arrow_array).unwrap(); - - // be careful of overflow: - // Want: 127i8 - (-128i8) -> 255u64, - // but we get -1i8 - // (-1i8) as u8 as u64 -> 255u64 - let sub = max.sub_wrapping(min) as ::Native; - let sub: <::UnSignedType as ArrowPrimitiveType>::Native = - sub.as_(); - let bit_width = get_bit_width(sub.as_()); - - let (_data_type, values, nulls) = arrow_array.clone().into_parts(); - let values = if min != T::Native::ZERO { - ScalarBuffer::from_iter(values.iter().map(|v| { - let k: <::UnSignedType as ArrowPrimitiveType>::Native = - v.sub_wrapping(min).as_(); - k - })) - } else { - #[allow(clippy::missing_transmute_annotations)] - unsafe { - std::mem::transmute(values) - } - }; - - let unsigned_array = - PrimitiveArray::<::UnSignedType>::new(values, nulls); - - let bit_packed_array = BitPackedArray::from_primitive(unsigned_array, bit_width); - - Self { - bit_packed: bit_packed_array, - reference_value: min, - } - } -} - -impl LiquidPrimitiveDeltaArray -where - T: LiquidPrimitiveType, -{ - /// Get the memory size of the Liquid primitive delta array. - pub fn get_array_memory_size(&self) -> usize { - self.bit_packed.get_array_memory_size() + std::mem::size_of::() - } - - /// Get the length of the Liquid primitive delta array. - pub fn len(&self) -> usize { - self.bit_packed.len() - } - - /// Check if the Liquid primitive delta array is empty. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Create a Liquid primitive delta array from an Arrow primitive array. - pub fn from_arrow_array(arrow_array: PrimitiveArray) -> LiquidPrimitiveDeltaArray { - use arrow::array::Array; - - let len = arrow_array.len(); - // check if entire array is already null - if arrow_array.null_count() == len { - return Self { - bit_packed: BitPackedArray::new_null_array(len), - reference_value: T::Native::ZERO, - }; - } - - let (_dt, values, nulls) = arrow_array.clone().into_parts(); - let vals: Vec = values.to_vec(); - - type UnsignedNative = - <::UnSignedType as ArrowPrimitiveType>::Native; - let mut out: Vec> = Vec::with_capacity(len); - let mut max_value: UnsignedNative = UnsignedNative::::ZERO; - let mut anchor: T::Native = T::Native::ZERO; - - if let Some(_nb) = &nulls { - // Nulls present: write 0 for nulls; prev will the last prev non-null value - let nb = nulls.as_ref().unwrap(); - let mut have_prev = false; - let mut prev: T::Native = T::Native::ZERO; - - for (i, &cur) in vals.iter().enumerate() { - if !nb.is_valid(i) { - out.push(UnsignedNative::::ZERO); - continue; - } - if !have_prev { - anchor = cur; - prev = cur; - have_prev = true; - out.push(UnsignedNative::::ZERO); - continue; - } - let delta: T::Native = cur.sub_wrapping(prev); - // zig zag encoding - let delta_i64: i64 = delta.as_(); - let zigzag: u64 = ((delta_i64 << 1) ^ (delta_i64 >> 63)) as u64; - let delta_unsigned: UnsignedNative = - UnsignedNative::::usize_as(zigzag as usize); - if delta_unsigned > max_value { - max_value = delta_unsigned; - } - out.push(delta_unsigned); - prev = cur; - } - } else { - // No nulls: first value is anchor, remainder are deltas with their previous values - anchor = vals[0]; - let mut prev: T::Native = anchor; - out.push(UnsignedNative::::ZERO); // anchor will have a difference of 0 - for &cur in vals.iter().skip(1) { - let delta: T::Native = cur.sub_wrapping(prev); - // zig zag encoding - let delta_i64: i64 = delta.as_(); - let zigzag: u64 = ((delta_i64 << 1) ^ (delta_i64 >> 63)) as u64; - let delta_unsigned: UnsignedNative = - UnsignedNative::::usize_as(zigzag as usize); - if delta_unsigned > max_value { - max_value = delta_unsigned; - } - out.push(delta_unsigned); - prev = cur; - } - } - - let bit_width = get_bit_width(max_value.as_()); - let values = ScalarBuffer::from_iter(out); - let unsigned_array = - PrimitiveArray::<::UnSignedType>::new(values, nulls); - let bit_packed_array = BitPackedArray::from_primitive(unsigned_array, bit_width); - - Self { - bit_packed: bit_packed_array, - reference_value: anchor, - } - } -} - -impl LiquidArray for LiquidPrimitiveArray -where - T: LiquidPrimitiveType + super::PrimitiveKind, -{ - fn get_array_memory_size(&self) -> usize { - self.get_array_memory_size() - } - - fn len(&self) -> usize { - self.len() - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn as_any(&self) -> &dyn Any { - self - } - - #[inline] - fn to_arrow_array(&self) -> ArrayRef { - let unsigned_array = self.bit_packed.to_primitive(); - let (_data_type, values, _nulls) = unsigned_array.into_parts(); - let nulls = self.bit_packed.nulls(); - let values = if self.reference_value != T::Native::ZERO { - let reference_v = self.reference_value.as_(); - ScalarBuffer::from_iter(values.iter().map(|v| { - let k: ::Native = (*v).add_wrapping(reference_v).as_(); - k - })) - } else { - #[allow(clippy::missing_transmute_annotations)] - unsafe { - std::mem::transmute(values) - } - }; - - Arc::new(PrimitiveArray::::new(values, nulls.cloned())) - } - - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let arrow_array = self.to_arrow_array(); - let selection = BooleanArray::new(selection.clone(), None); - arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() - } - - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - let filtered = self.filter(filter); - eval_predicate_on_array(filtered, predicate) - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Integer - } -} - -impl LiquidArray for LiquidPrimitiveDeltaArray -where - T: LiquidPrimitiveType + super::PrimitiveKind, -{ - fn get_array_memory_size(&self) -> usize { - self.get_array_memory_size() - } - - fn len(&self) -> usize { - self.len() - } - - fn original_arrow_data_type(&self) -> DataType { - T::DATA_TYPE.clone() - } - - fn as_any(&self) -> &dyn Any { - self - } - - #[inline] - fn to_arrow_array(&self) -> ArrayRef { - // Reconstruct original values from deltas - let unsigned_array = self.bit_packed.to_primitive(); - let (_data_type, delta_values, _nulls) = unsigned_array.into_parts(); - let nulls = self.bit_packed.nulls(); - - // Reconstruct original values by applying deltas - let mut reconstructed = Vec::with_capacity(delta_values.len()); - let mut current_value = self.reference_value; // anchor - - if let Some(nulls) = nulls { - let mut have_prev = false; - for (i, &delta_unsigned) in delta_values.iter().enumerate() { - if !nulls.is_valid(i) { - reconstructed.push(T::Native::ZERO); // Will be masked out by nulls - continue; - } - if !have_prev { - // First non-null value is the anchor - reconstructed.push(current_value); - have_prev = true; - } else { - // Apply delta to get next value - let zigzag: u64 = delta_unsigned.as_(); - let delta_i64 = (zigzag >> 1) as i64 ^ -((zigzag & 1) as i64); - let delta: T::Native = T::Native::from_i64(delta_i64).unwrap(); - current_value = current_value.add_wrapping(delta); - reconstructed.push(current_value); - } - } - } else { - // No nulls case - reconstructed.push(current_value); // First value is anchor - for &delta_unsigned in delta_values.iter().skip(1) { - let zigzag: u64 = delta_unsigned.as_(); - let delta_i64 = (zigzag >> 1) as i64 ^ -((zigzag & 1) as i64); - let delta: T::Native = T::Native::from_i64(delta_i64).unwrap(); - current_value = current_value.add_wrapping(delta); - reconstructed.push(current_value); - } - } - - let values = ScalarBuffer::from_iter(reconstructed); - Arc::new(PrimitiveArray::::new(values, nulls.cloned())) - } - - fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { - let arrow_array = self.to_arrow_array(); - let selection = BooleanArray::new(selection.clone(), None); - arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() - } - - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { - let filtered = self.filter(filter); - eval_predicate_on_array(filtered, predicate) - } - - fn to_bytes(&self) -> Vec { - self.to_bytes_inner() - } - - fn data_type(&self) -> LiquidDataType { - LiquidDataType::Integer - } -} - -impl LiquidPrimitiveArray -where - T: LiquidPrimitiveType, -{ - fn bit_pack_starting_loc() -> usize { - let header_size = LiquidIPCHeader::size() + std::mem::size_of::(); - (header_size + 7) & !7 - } - - /* - Serialized LiquidPrimitiveArray Memory Layout: - +--------------------------------------------------+ - | LiquidIPCHeader (16 bytes) | - +--------------------------------------------------+ - - +--------------------------------------------------+ - | reference_value (size_of:: bytes) | // The reference value (e.g. minimum value) - +--------------------------------------------------+ - | Padding (to 8-byte alignment) | // Padding to ensure 8-byte alignment - +--------------------------------------------------+ - - +--------------------------------------------------+ - | BitPackedArray Data | - +--------------------------------------------------+ - | [BitPackedArray Header & Bit-Packed Values] | // Written by self.bit_packed.to_bytes() - +--------------------------------------------------+ - */ - pub(crate) fn to_bytes_inner(&self) -> Vec { - // Determine type ID based on the type - let physical_type_id = get_physical_type_id::(); - let logical_type_id = super::LiquidDataType::Integer as u16; - let header = LiquidIPCHeader::new(logical_type_id, physical_type_id); - - let bit_pack_starting_loc = Self::bit_pack_starting_loc(); - let mut result = Vec::with_capacity(bit_pack_starting_loc + 256); // Pre-allocate a reasonable size - - // Write header - result.extend_from_slice(&header.to_bytes()); - - // Write reference value - let ref_value_bytes = unsafe { - std::slice::from_raw_parts( - &self.reference_value as *const T::Native as *const u8, - std::mem::size_of::(), - ) - }; - result.extend_from_slice(ref_value_bytes); - while result.len() < bit_pack_starting_loc { - result.push(0); - } - - // Let BitPackedArray write the rest of the data - self.bit_packed.to_bytes(&mut result); - - result - } - - /// Deserialize a LiquidPrimitiveArray from bytes - pub fn from_bytes(bytes: Bytes) -> Self { - let header = LiquidIPCHeader::from_bytes(&bytes); - - let physical_id = header.physical_type_id; - assert_eq!(physical_id, get_physical_type_id::()); - let logical_id = header.logical_type_id; - assert_eq!(logical_id, super::LiquidDataType::Integer as u16); - - // Get the reference value - let ref_value_ptr = &bytes[LiquidIPCHeader::size()]; - let reference_value = - unsafe { (ref_value_ptr as *const u8 as *const T::Native).read_unaligned() }; - - // Skip ahead to the BitPackedArray data - let bit_packed_data = bytes.slice(Self::bit_pack_starting_loc()..); - let bit_packed = BitPackedArray::::from_bytes(bit_packed_data); - - Self { - bit_packed, - reference_value, - } - } -} - -impl LiquidPrimitiveDeltaArray -where - T: LiquidPrimitiveType, -{ - fn bit_pack_starting_loc() -> usize { - let header_size = LiquidIPCHeader::size() + std::mem::size_of::(); - (header_size + 7) & !7 - } - - pub(crate) fn to_bytes_inner(&self) -> Vec { - // Determine type ID based on the type - let physical_type_id = get_physical_type_id::(); - let logical_type_id = 1u16; // Delta encoding type ID - let header = LiquidIPCHeader::new(logical_type_id, physical_type_id); - - let bit_pack_starting_loc = Self::bit_pack_starting_loc(); - let mut result = Vec::with_capacity(bit_pack_starting_loc + 256); - - // Write header - result.extend_from_slice(&header.to_bytes()); - - // Write anchor value (reference_value) - let ref_value_bytes = unsafe { - std::slice::from_raw_parts( - &self.reference_value as *const T::Native as *const u8, - std::mem::size_of::(), - ) - }; - result.extend_from_slice(ref_value_bytes); - while result.len() < bit_pack_starting_loc { - result.push(0); - } - - // Let BitPackedArray write the rest of the data - self.bit_packed.to_bytes(&mut result); - - result - } - - /// Deserialize a LiquidPrimitiveDeltaArray from bytes - pub fn from_bytes(bytes: Bytes) -> Self { - let header = LiquidIPCHeader::from_bytes(&bytes); - - let physical_id = header.physical_type_id; - assert_eq!(physical_id, get_physical_type_id::()); - let logical_id = header.logical_type_id; - assert_eq!(logical_id, 1u16); // Delta encoding type ID - - // Get the anchor value - let ref_value_ptr = &bytes[LiquidIPCHeader::size()]; - let reference_value = - unsafe { (ref_value_ptr as *const u8 as *const T::Native).read_unaligned() }; - - // Skip ahead to the BitPackedArray data - let bit_packed_data = bytes.slice(Self::bit_pack_starting_loc()..); - let bit_packed = BitPackedArray::::from_bytes(bit_packed_data); - - Self { - bit_packed, - reference_value, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::Array; - - macro_rules! test_roundtrip { - ($test_name:ident, $type:ty, $values:expr) => { - #[test] - fn $test_name() { - // Create the original array - let original: Vec::Native>> = $values; - let array = PrimitiveArray::<$type>::from(original.clone()); - - // Convert to Liquid array and back - let liquid_array = LiquidPrimitiveArray::<$type>::from_arrow_array(array.clone()); - let result_array = liquid_array.to_arrow_array(); - let bytes_array = - LiquidPrimitiveArray::<$type>::from_bytes(liquid_array.to_bytes().into()); - - assert_eq!(result_array.as_ref(), &array); - assert_eq!(bytes_array.to_arrow_array().as_ref(), &array); - } - }; - } - - // Test cases for Int8Type - test_roundtrip!( - test_int8_roundtrip_basic, - Int8Type, - vec![Some(1), Some(2), Some(3), None, Some(5)] - ); - test_roundtrip!( - test_int8_roundtrip_negative, - Int8Type, - vec![Some(-128), Some(-64), Some(0), Some(63), Some(127)] - ); - - // Test cases for Int16Type - test_roundtrip!( - test_int16_roundtrip_basic, - Int16Type, - vec![Some(1), Some(2), Some(3), None, Some(5)] - ); - test_roundtrip!( - test_int16_roundtrip_negative, - Int16Type, - vec![ - Some(-32768), - Some(-16384), - Some(0), - Some(16383), - Some(32767) - ] - ); - - // Test cases for Int32Type - test_roundtrip!( - test_int32_roundtrip_basic, - Int32Type, - vec![Some(1), Some(2), Some(3), None, Some(5)] - ); - test_roundtrip!( - test_int32_roundtrip_negative, - Int32Type, - vec![ - Some(-2147483648), - Some(-1073741824), - Some(0), - Some(1073741823), - Some(2147483647) - ] - ); - - // Test cases for Int64Type - test_roundtrip!( - test_int64_roundtrip_basic, - Int64Type, - vec![Some(1), Some(2), Some(3), None, Some(5)] - ); - test_roundtrip!( - test_int64_roundtrip_negative, - Int64Type, - vec![ - Some(-9223372036854775808), - Some(-4611686018427387904), - Some(0), - Some(4611686018427387903), - Some(9223372036854775807) - ] - ); - - // Test cases for unsigned types - test_roundtrip!( - test_uint8_roundtrip, - UInt8Type, - vec![Some(0), Some(128), Some(255), None, Some(64)] - ); - test_roundtrip!( - test_uint16_roundtrip, - UInt16Type, - vec![Some(0), Some(32768), Some(65535), None, Some(16384)] - ); - test_roundtrip!( - test_uint32_roundtrip, - UInt32Type, - vec![ - Some(0), - Some(2147483648), - Some(4294967295), - None, - Some(1073741824) - ] - ); - test_roundtrip!( - test_uint64_roundtrip, - UInt64Type, - vec![ - Some(0), - Some(9223372036854775808), - Some(18446744073709551615), - None, - Some(4611686018427387904) - ] - ); - - test_roundtrip!( - test_date32_roundtrip, - Date32Type, - vec![Some(-365), Some(0), Some(365), None, Some(18262)] - ); - - test_roundtrip!( - test_date64_roundtrip, - Date64Type, - vec![Some(-365), Some(0), Some(365), None, Some(18262)] - ); - - // Edge cases - #[test] - fn test_all_nulls() { - let original: Vec> = vec![None, None, None]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - let result_array = liquid_array.to_arrow_array(); - - assert_eq!(result_array.len(), original.len()); - assert_eq!(result_array.null_count(), original.len()); - } - - #[test] - fn test_all_nulls_filter() { - let original: Vec> = vec![None, None, None]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - let result_array = liquid_array.filter(&BooleanBuffer::from(vec![true, false, true])); - - assert_eq!(result_array.len(), 2); - assert_eq!(result_array.null_count(), 2); - } - - #[test] - fn test_zero_reference_value() { - let original: Vec> = vec![Some(0), Some(1), Some(2), None, Some(4)]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let result_array = liquid_array.to_arrow_array(); - - assert_eq!(liquid_array.reference_value, 0); - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_single_value() { - let original: Vec> = vec![Some(42)]; - let array = PrimitiveArray::::from(original.clone()); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let result_array = liquid_array.to_arrow_array(); - - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_filter_basic() { - // Create original array with some values - let original = vec![Some(1), Some(2), Some(3), None, Some(5)]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - - // Create selection mask: keep indices 0, 2, and 4 - let selection = BooleanBuffer::from(vec![true, false, true, false, true]); - - // Apply filter - let result_array = liquid_array.filter(&selection); - - // Expected result after filtering - let expected = PrimitiveArray::::from(vec![Some(1), Some(3), Some(5)]); - - assert_eq!(result_array.as_ref(), &expected); - } - - #[test] - fn test_original_arrow_data_type_returns_int32() { - let array = PrimitiveArray::::from(vec![Some(1), Some(2)]); - let liquid = LiquidPrimitiveArray::::from_arrow_array(array); - assert_eq!(liquid.original_arrow_data_type(), DataType::Int32); - } - - #[test] - fn test_filter_all_nulls() { - // Create array with all nulls - let original = vec![None, None, None, None]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - - // Keep first and last elements - let selection = BooleanBuffer::from(vec![true, false, false, true]); - - let result_array = liquid_array.filter(&selection); - - let expected = PrimitiveArray::::from(vec![None, None]); - - assert_eq!(result_array.as_ref(), &expected); - } - - #[test] - fn test_filter_empty_result() { - let original = vec![Some(1), Some(2), Some(3)]; - let array = PrimitiveArray::::from(original); - let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); - - // Filter out all elements - let selection = BooleanBuffer::from(vec![false, false, false]); - - let result_array = liquid_array.filter(&selection); - - assert_eq!(result_array.len(), 0); - } - - #[test] - fn test_delta_encoding_basic_roundtrip() { - let original = vec![Some(1), Some(3), Some(6), Some(10), Some(15)]; - let array = PrimitiveArray::::from(original.clone()); - - let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); - let result_array = liquid_delta.to_arrow_array(); - - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_delta_encoding_with_nulls() { - let original = vec![Some(1), None, Some(4), Some(7), None, Some(12)]; - let array = PrimitiveArray::::from(original.clone()); - - let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); - let result_array = liquid_delta.to_arrow_array(); - - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_delta_encoding_serialization() { - let original = vec![Some(1), Some(3), Some(6), Some(10), Some(15)]; - let array = PrimitiveArray::::from(original.clone()); - - let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); - let bytes = liquid_delta.to_bytes(); - let reconstructed = LiquidPrimitiveDeltaArray::::from_bytes(bytes.into()); - let result_array = reconstructed.to_arrow_array(); - - assert_eq!(result_array.as_ref(), &array); - } - - #[test] - fn test_memory_comparison_sequential_data() { - // Sequential data: delta encoding performs better - let sequential_data: Vec> = (0..1000).map(Some).collect(); - let array = PrimitiveArray::::from(sequential_data); - - let liquid_regular = LiquidPrimitiveArray::::from_arrow_array(array.clone()); - let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array); - - let regular_size = liquid_regular.get_array_memory_size(); - let delta_size = liquid_delta.get_array_memory_size(); - - println!( - "Sequential data - Regular: {} bytes, Delta: {} bytes", - regular_size, delta_size - ); - assert!( - delta_size <= regular_size, - "Delta encoding should be more efficient for sequential data" - ); - } -} diff --git a/src/core/src/liquid_array/raw/bit_pack_array.rs b/src/core/src/liquid_array/raw/bit_pack_array.rs deleted file mode 100644 index 4b8954cde..000000000 --- a/src/core/src/liquid_array/raw/bit_pack_array.rs +++ /dev/null @@ -1,584 +0,0 @@ -use std::mem::size_of; -use std::num::NonZero; - -use arrow::array::{ArrowPrimitiveType, PrimitiveArray}; -use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer}; -use arrow::datatypes::ArrowNativeType; -use bytes; -use fastlanes::BitPacking; - -/// A bit-packed array. -#[derive(Debug)] -pub struct BitPackedArray -where - T::Native: BitPacking, -{ - packed_values: ScalarBuffer, - nulls: Option, - bit_width: Option>, // if None, the array is entirely null - original_len: usize, -} - -/// Implement Clone for any T that implements ArrowPrimitiveType and BitPacking -/// This allows us to clone it without requiring T to implement Clone -impl Clone for BitPackedArray -where - T::Native: BitPacking, -{ - fn clone(&self) -> Self { - Self { - packed_values: self.packed_values.clone(), - nulls: self.nulls.clone(), - bit_width: self.bit_width, - original_len: self.original_len, - } - } -} - -impl BitPackedArray -where - T::Native: BitPacking, -{ - /// Creates a new null array with the given length. - pub fn new_null_array(len: usize) -> Self { - Self { - packed_values: vec![T::Native::usize_as(0); len].into(), - nulls: Some(NullBuffer::new_null(len)), - bit_width: None, - original_len: len, - } - } - - pub(crate) fn len(&self) -> usize { - self.original_len - } - - pub(crate) fn nulls(&self) -> Option<&NullBuffer> { - self.nulls.as_ref() - } - - #[cfg(test)] - pub(crate) fn bit_width(&self) -> Option> { - self.bit_width - } - - /// Returns true if the array is nullable. - #[cfg(test)] - fn is_nullable(&self) -> bool { - self.nulls.is_some() - } - - /// Creates a new bit-packed array from a primitive array and a bit width. - pub fn from_primitive(array: PrimitiveArray, bit_width: NonZero) -> Self { - let original_len = array.len(); - let (_data_type, values, nulls) = array.into_parts(); - - let bit_width_usize = bit_width.get() as usize; - let num_chunks = original_len.div_ceil(1024); - let num_full_chunks = original_len / 1024; - let packed_len = (1024 * bit_width_usize).div_ceil(size_of::() * 8); - - let mut output = Vec::::with_capacity(num_chunks * packed_len); - - (0..num_full_chunks).for_each(|i| { - let start_elem = i * 1024; - - output.reserve(packed_len); - let output_len = output.len(); - unsafe { - output.set_len(output_len + packed_len); - BitPacking::unchecked_pack( - bit_width_usize, - &values[start_elem..][..1024], - &mut output[output_len..][..packed_len], - ); - } - }); - - if num_chunks != num_full_chunks { - let last_chunk_size = values.len() % 1024; - let mut last_chunk = vec![T::Native::default(); 1024]; - last_chunk[..last_chunk_size] - .copy_from_slice(&values[values.len() - last_chunk_size..]); - - output.reserve(packed_len); - let output_len = output.len(); - unsafe { - output.set_len(output_len + packed_len); - BitPacking::unchecked_pack( - bit_width_usize, - &last_chunk, - &mut output[output_len..][..packed_len], - ); - } - } - - let buffer = Buffer::from(output); - let scalar_buffer = ScalarBuffer::new(buffer, 0, num_chunks * packed_len); - - Self { - packed_values: scalar_buffer, - nulls, - bit_width: Some(bit_width), - original_len, - } - } - - /// Converts the bit-packed array to a primitive array. - pub fn to_primitive(&self) -> PrimitiveArray { - // Special case for all nulls, don't unpack - let bit_width = if let Some(bit_width) = self.bit_width { - bit_width.get() as usize - } else { - return PrimitiveArray::::new_null(self.original_len); - }; - let packed = self.packed_values.as_ref(); - let length = self.original_len; - let offset = 0; - - let num_chunks = (offset + length).div_ceil(1024); - let elements_per_chunk = (1024 * bit_width).div_ceil(size_of::() * 8); - - let mut output = Vec::::with_capacity(num_chunks * 1024 - offset); - - let first_full_chunk = if offset != 0 { - let chunk: &[T::Native] = &packed[0..elements_per_chunk]; - let mut decoded = vec![T::Native::default(); 1024]; - unsafe { BitPacking::unchecked_unpack(bit_width, chunk, &mut decoded) }; - output.extend_from_slice(&decoded[offset..]); - 1 - } else { - 0 - }; - - (first_full_chunk..num_chunks).for_each(|i| { - let chunk: &[T::Native] = &packed[i * elements_per_chunk..][0..elements_per_chunk]; - unsafe { - let output_len = output.len(); - output.set_len(output_len + 1024); - BitPacking::unchecked_unpack(bit_width, chunk, &mut output[output_len..][..1024]); - } - }); - - output.truncate(length); - if output.len() < 1024 { - output.shrink_to_fit(); - } - - let nulls = self.nulls.clone(); - PrimitiveArray::::new(ScalarBuffer::from(output), nulls) - } - - /// Returns the memory size of the bit-packed array. - pub fn get_array_memory_size(&self) -> usize { - std::mem::size_of::() - + self.packed_values.inner().capacity() - + self - .nulls - .as_ref() - .map_or(0, |nulls| nulls.buffer().capacity()) - } - - /* - Memory Layout (serialized): - - +-----------------------------+ // Header (16 bytes total) - | original_len (4 bytes) | // Offset 0 - 3: Array length (u32) - +-----------------------------+ // - | bit_width (1 byte) | // Offset 4: Bit width (u8) - +-----------------------------+ // - | has_nulls (1 byte) | // Offset 5: Null flag (1 if nulls present) - +-----------------------------+ // - | nulls_len (4 bytes) | // Offset 6 - 9: Length of nulls buffer (u32) - +-----------------------------+ // - | values_len (4 bytes) | // Offset 10 - 13: Length of values buffer (u32) - +-----------------------------+ // - | padding (2 bytes) | // Offset 14 - 15: Padding to ensure 16-byte header - +-----------------------------+ - - [If has_nulls == 1] - +-----------------------------+ // Nulls Buffer - | nulls data (nulls_len bytes)| // Offset 16 - (16 + nulls_len - 1) - +-----------------------------+ - - +-----------------------------+ - | Padding for 8-byte alignment| // Ensure values buffer is 8-byte aligned - +-----------------------------+ - - +-----------------------------+ // Values Buffer (bit-packed data) - | values data (values_len) | // Starts at the 8-byte aligned offset - +-----------------------------+ - */ - /// Serializes the bit-packed array to a byte buffer. - pub fn to_bytes(&self, buffer: &mut Vec) { - let has_nulls = self.nulls.is_some() as u8; - - let nulls_sliced; - let nulls_bytes = if has_nulls == 1 { - let nulls = self.nulls.as_ref().unwrap(); - if nulls.offset() == 0 { - nulls.buffer().as_slice() - } else { - nulls_sliced = Some(nulls.inner().sliced()); - nulls_sliced.as_ref().unwrap().as_slice() - } - } else { - &[] - }; - - let values_bytes = self.packed_values.inner().as_slice(); - - let header_size = 16; - - let values_offset_base = header_size + if has_nulls == 1 { nulls_bytes.len() } else { 0 }; - let values_offset = (values_offset_base + 7) & !7; - - let total_size = values_offset + values_bytes.len(); - buffer.reserve(total_size); - - let start_offset = buffer.len(); - - buffer.extend_from_slice(&(self.original_len as u32).to_le_bytes()); - buffer.push(self.bit_width.map_or(0, |bit_width| bit_width.get())); - buffer.push(has_nulls); - buffer.extend_from_slice(&(nulls_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&(values_bytes.len() as u32).to_le_bytes()); - buffer.extend_from_slice(&[0, 0]); - - if has_nulls == 1 { - buffer.extend_from_slice(nulls_bytes); - } - - while (buffer.len() - start_offset) < values_offset { - buffer.push(0); - } - - buffer.extend_from_slice(values_bytes); - } - - /// Deserializes a bit-packed array from a byte buffer. - pub fn from_bytes(bytes: bytes::Bytes) -> Self - where - T::Native: BitPacking, - { - use std::mem::size_of; - - if bytes.len() < 16 { - panic!("Input buffer too small for header"); - } - - // Read header fields - let original_len = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize; - let bit_width = bytes[4]; - let has_nulls = bytes[5] != 0; - let nulls_len = u32::from_le_bytes(bytes[6..10].try_into().unwrap()) as usize; - let values_len = u32::from_le_bytes(bytes[10..14].try_into().unwrap()) as usize; - - // Calculate offsets - let header_size = 16; - let nulls_offset = if has_nulls { header_size } else { 0 }; - let values_offset_base = header_size + if has_nulls { nulls_len } else { 0 }; - let values_offset = (values_offset_base + 7) & !7; // 8-byte aligned - - if values_len == 0 { - // if empty array, return a new null array - return Self::new_null_array(original_len); - } - - // Validate offsets and lengths - if has_nulls { - if nulls_offset == 0 || nulls_len == 0 { - panic!("Array has nulls but null buffer is missing"); - } - if nulls_offset + nulls_len > bytes.len() { - panic!("Null buffer extends beyond input buffer"); - } - } - - if values_offset == 0 || values_len == 0 { - panic!("Values buffer is required"); - } - if values_offset + values_len > bytes.len() { - panic!("Values buffer extends beyond input buffer"); - } - - // Create the nulls buffer if present - let nulls = if has_nulls { - // Create a buffer view into the nulls section - let nulls_slice = bytes.slice(nulls_offset..nulls_offset + nulls_len); - let nulls_buffer = Buffer::from(nulls_slice); - let boolean_buffer = BooleanBuffer::new(nulls_buffer, 0, original_len); - Some(NullBuffer::from(boolean_buffer)) - } else { - None - }; - - let values_slice = bytes.slice(values_offset..values_offset + values_len); - let values_buffer = Buffer::from(values_slice); - - let element_size = size_of::(); - let packed_len = values_len / element_size; - - let packed_values = ScalarBuffer::::new(values_buffer, 0, packed_len); - - if nulls.is_some() && nulls.as_ref().unwrap().null_count() == original_len { - return Self::new_null_array(original_len); - } - - Self { - packed_values, - nulls, - bit_width: Some(NonZero::new(bit_width).unwrap()), - original_len, - } - } -} - -#[allow(dead_code)] -fn best_arrow_primitive_width(bit_width: NonZero) -> usize { - match bit_width.get() { - 0..=8 => 8, - 9..=16 => 16, - 17..=32 => 32, - 33..=64 => 64, - _ => panic!("Unsupported bit width: {}", bit_width.get()), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::{ - array::Array, - datatypes::{UInt16Type, UInt32Type}, - }; - - #[test] - fn test_bit_pack_roundtrip() { - // Test with a full chunk (1024 elements) - let values: Vec = (0..1024).collect(); - - let array = PrimitiveArray::::from(values); - let before_size = array.get_array_memory_size(); - let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); - let after_size = bit_packed.get_array_memory_size(); - println!("before: {before_size}, after: {after_size}"); - let unpacked = bit_packed.to_primitive(); - - assert_eq!(unpacked.len(), 1024); - for i in 0..1024 { - assert_eq!(unpacked.value(i), i as u32); - } - } - - #[test] - fn test_bit_pack_partial_chunk() { - // Test with a partial chunk (500 elements) - let values: Vec = (0..500).collect(); - let array = PrimitiveArray::::from(values); - let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); - let unpacked = bit_packed.to_primitive(); - - assert_eq!(unpacked.len(), 500); - for i in 0..500 { - assert_eq!(unpacked.value(i), i as u32); - } - } - - #[test] - fn test_bit_pack_multiple_chunks() { - // Test with multiple chunks (2048 elements = 2 full chunks) - let values: Vec = (0..2048).collect(); - let array = PrimitiveArray::::from(values); - let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(11).unwrap()); - let unpacked = bit_packed.to_primitive(); - - assert_eq!(unpacked.len(), 2048); - for i in 0..2048 { - assert_eq!(unpacked.value(i), i as u32); - } - } - - #[test] - fn test_bit_pack_with_nulls() { - let values: Vec> = (0..1000) - .map(|i| if i % 2 == 0 { Some(i as u32) } else { None }) - .collect(); - let array = PrimitiveArray::::from(values); - let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); - let unpacked = bit_packed.to_primitive(); - - assert_eq!(unpacked.len(), 1000); - for i in 0..1000_usize { - if i.is_multiple_of(2) { - assert_eq!(unpacked.value(i), i as u32); - } else { - assert!(unpacked.is_null(i)); - } - } - } - - #[test] - fn test_different_bit_widths() { - // Test with different bit widths - let values: Vec = (0..100).map(|i| i * 2).collect(); - let array = PrimitiveArray::::from(values); - - for bit_width in [8, 16, 24, 32] { - let bit_packed = - BitPackedArray::from_primitive(array.clone(), NonZero::new(bit_width).unwrap()); - let unpacked = bit_packed.to_primitive(); - - assert_eq!(unpacked.len(), 100); - for i in 0..100 { - assert_eq!(unpacked.value(i), i as u32 * 2); - } - } - } - - #[test] - fn test_to_bytes_from_bytes_roundtrip() { - // Create a test array with some values - let values: Vec = (0..100).collect(); - let array = PrimitiveArray::::from(values); - let bit_width = NonZero::new(10).unwrap(); - let original = BitPackedArray::from_primitive(array, bit_width); - - // Serialize to bytes - let mut buffer = Vec::new(); - original.to_bytes(&mut buffer); - - // Make sure we have some reasonable amount of data - assert!(!buffer.is_empty()); - assert!(buffer.len() > 16); // At least header size - - // Deserialize back using from_bytes - let bytes = bytes::Bytes::from(buffer); - let deserialized = BitPackedArray::::from_bytes(bytes); - - // Verify the deserialized data matches the original - assert_eq!(deserialized.bit_width(), original.bit_width()); - assert_eq!(deserialized.len(), original.len()); - assert_eq!(deserialized.is_nullable(), original.is_nullable()); - - // Convert to primitive arrays and compare values - let original_primitive = original.to_primitive(); - let deserialized_primitive = deserialized.to_primitive(); - - assert_eq!(original_primitive.len(), deserialized_primitive.len()); - for i in 0..original_primitive.len() { - assert_eq!(original_primitive.value(i), deserialized_primitive.value(i)); - } - } - - #[test] - fn test_to_bytes_from_bytes_with_nulls() { - // Create a test array with some nulls - let values: Vec> = (0..100) - .map(|i: u32| if i.is_multiple_of(3) { None } else { Some(i) }) - .collect(); - let array = PrimitiveArray::::from(values); - let bit_width = NonZero::new(10).unwrap(); - let original = BitPackedArray::from_primitive(array, bit_width); - - // Serialize to bytes - let mut buffer = Vec::new(); - original.to_bytes(&mut buffer); - - // Deserialize back - let bytes = bytes::Bytes::from(buffer); - let deserialized = BitPackedArray::::from_bytes(bytes); - - // Verify the deserialized data matches the original - assert_eq!(deserialized.bit_width(), original.bit_width()); - assert_eq!(deserialized.len(), original.len()); - assert_eq!(deserialized.is_nullable(), original.is_nullable()); - - // Convert to primitive arrays and compare values including nulls - let original_primitive = original.to_primitive(); - let deserialized_primitive = deserialized.to_primitive(); - - assert_eq!(original_primitive.len(), deserialized_primitive.len()); - for i in 0..original_primitive.len() { - assert_eq!( - original_primitive.is_null(i), - deserialized_primitive.is_null(i) - ); - if !original_primitive.is_null(i) { - assert_eq!(original_primitive.value(i), deserialized_primitive.value(i)); - } - } - } - - #[test] - fn test_to_bytes_from_bytes_with_nulls_and_offset() { - let values: Vec> = (0..32) - .map(|i| if i % 3 == 0 { None } else { Some(i as u16) }) - .collect(); - let array = PrimitiveArray::::from(values); - - // Slice to create a non-zero offset (and therefore a non-zero null bitmap bit offset). - let sliced = array.slice(1, 23); - - let bit_width = NonZero::new(16).unwrap(); - let original = BitPackedArray::from_primitive(sliced.clone(), bit_width); - - let mut buffer = Vec::new(); - original.to_bytes(&mut buffer); - let deserialized = BitPackedArray::::from_bytes(buffer.into()); - - let roundtripped = deserialized.to_primitive(); - assert_eq!(roundtripped, sliced); - } - - #[test] - fn test_memory_size_calculation() { - use super::*; - use arrow::buffer::{Buffer, NullBuffer, ScalarBuffer}; - use arrow::datatypes::UInt32Type; - - let scalar_buffer = ScalarBuffer::::new(Buffer::from(vec![0; 1024]), 0, 1024); - - // --- Test without nulls --- - let bit_packed_no_nulls = BitPackedArray:: { - packed_values: scalar_buffer.clone(), - nulls: None, - bit_width: Some(NonZero::new(10).unwrap()), - original_len: 1024, - }; - - let expected_size_no_nulls = - size_of::>() + scalar_buffer.inner().capacity(); - assert_eq!( - bit_packed_no_nulls.get_array_memory_size(), - expected_size_no_nulls, - "Memory size mismatch without nulls" - ); - - // --- Test with nulls --- - // Create dummy null buffer - let null_buffer = NullBuffer::new_null(1024); - let nulls = Some(null_buffer); - - let bit_packed_with_nulls = BitPackedArray:: { - packed_values: scalar_buffer.clone(), - nulls: nulls.clone(), // Clone the Option - bit_width: Some(NonZero::new(10).unwrap()), - original_len: 1024, - }; - - // Calculate expected size including null buffer - // Note: Arrow's Buffer might allocate slightly more than null_bitmap_len_bytes - // We use the actual buffer capacity for a more precise comparison - let actual_null_buffer_size = nulls.as_ref().map_or(0, |nb| nb.buffer().capacity()); - let expected_size_with_nulls = size_of::>() - + scalar_buffer.inner().capacity() - + actual_null_buffer_size; - - assert_eq!( - bit_packed_with_nulls.get_array_memory_size(), - expected_size_with_nulls, - "Memory size mismatch with nulls" - ); - } -} diff --git a/src/core/src/liquid_array/raw/fsst_buffer.rs b/src/core/src/liquid_array/raw/fsst_buffer.rs deleted file mode 100644 index 5ce63275b..000000000 --- a/src/core/src/liquid_array/raw/fsst_buffer.rs +++ /dev/null @@ -1,1068 +0,0 @@ -use arrow::{ - array::{ - ArrayDataBuilder, Decimal128Array, Decimal256Array, GenericByteArray, OffsetBufferBuilder, - }, - buffer::{Buffer, OffsetBuffer}, - datatypes::ByteArrayType, -}; -use bytes; -use fsst::{Compressor, Decompressor, Symbol}; -use std::io::Result; -use std::io::{Error, ErrorKind}; -use std::sync::Arc; - -use crate::liquid_array::fix_len_byte_array::ArrowFixedLenByteArrayType; - -mod sealed { - pub trait Sealed {} -} - -/// Raw FSST buffer that stores compressed data using Arrow Buffer. -/// Offsets are managed externally as a `u32` slice (including the final sentinel offset). -pub(crate) struct RawFsstBuffer { - values: Buffer, - uncompressed_bytes: usize, -} - -impl std::fmt::Debug for RawFsstBuffer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("RawFsstBuffer") - .field("values_len", &self.values.len()) - .field("uncompressed_bytes", &self.uncompressed_bytes) - .finish() - } -} - -impl RawFsstBuffer { - pub(crate) fn from_parts(values: Buffer, uncompressed_bytes: usize) -> Self { - Self { - values, - uncompressed_bytes, - } - } - - /// Create RawFsstBuffer from an iterator of byte slices. - /// Returns the buffer and a vector of byte offsets (including the final sentinel). - pub(crate) fn from_byte_slices( - iter: I, - compressor: Arc, - compress_buffer: &mut Vec, - ) -> (Self, Vec) - where - I: Iterator>, - T: AsRef<[u8]>, - { - let mut values_buffer = Vec::new(); - let mut offsets = Vec::new(); - let mut uncompressed_len = 0; - - offsets.push(0u32); - for item in iter { - if let Some(bytes) = item { - let bytes = bytes.as_ref(); - uncompressed_len += bytes.len(); - - compress_buffer.clear(); - // `fsst::Compressor::compress_into` requires capacity for the worst-case expansion - // (all bytes escaped) which is `2 * plaintext_len`. - compress_buffer.reserve(bytes.len().saturating_mul(2)); - unsafe { - let len = compressor.compress_into(bytes, compress_buffer.spare_capacity_mut()); - compress_buffer.set_len(len); - } - - values_buffer.extend_from_slice(compress_buffer); - } - offsets.push(values_buffer.len() as u32); - } - - values_buffer.shrink_to_fit(); - let values_buffer = Buffer::from(values_buffer); - let raw_buffer = Self::from_parts(values_buffer, uncompressed_len); - - (raw_buffer, offsets) - } - - pub(crate) fn to_uncompressed( - &self, - decompressor: &Decompressor<'_>, - offsets: &[u32], - ) -> (Buffer, OffsetBuffer) { - let mut value_buffer: Vec = Vec::with_capacity(self.uncompressed_bytes + 8); - let num_values = offsets.len().saturating_sub(1); - let mut out_offsets: OffsetBufferBuilder = OffsetBufferBuilder::new(num_values); - - for i in 0..num_values { - let start_offset = offsets[i]; - let end_offset = offsets[i + 1]; - - if start_offset != end_offset { - let compressed_slice = self.get_compressed_slice(start_offset, end_offset); - let decompressed_len = decompressor - .decompress_into(compressed_slice, value_buffer.spare_capacity_mut()); - - let new_len = value_buffer.len() + decompressed_len; - debug_assert!(new_len <= value_buffer.capacity()); - unsafe { - value_buffer.set_len(new_len); - } - out_offsets.push_length(decompressed_len); - } else { - out_offsets.push_length(0); - } - } - - let buffer = Buffer::from(value_buffer); - (buffer, out_offsets.finish()) - } - - /// Get compressed data slice using byte offsets. - pub(crate) fn get_compressed_slice(&self, start_offset: u32, end_offset: u32) -> &[u8] { - let start = start_offset as usize; - let end = end_offset as usize; - debug_assert!(end <= self.values.len(), "Offset out of bounds"); - debug_assert!(start <= end, "Invalid offset range"); - &self.values.as_slice()[start..end] - } - - pub(crate) fn values_len(&self) -> usize { - self.values.len() - } - - pub(crate) fn get_memory_size(&self) -> usize { - self.values.len() + std::mem::size_of::() - } - - pub(crate) fn to_bytes(&self) -> Vec { - let mut buffer = Vec::with_capacity(self.values.len() + 12); - buffer.extend_from_slice(&(self.uncompressed_bytes as u64).to_le_bytes()); - buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); - buffer.extend_from_slice(self.values.as_slice()); - buffer - } - - pub(crate) fn uncompressed_bytes(&self) -> usize { - self.uncompressed_bytes - } - - pub(crate) fn from_bytes(bytes: bytes::Bytes) -> Self { - let uncompressed_bytes = u64::from_le_bytes(bytes[0..8].try_into().unwrap()) as usize; - let values_len = u32::from_le_bytes(bytes[8..12].try_into().unwrap()) as usize; - let values = bytes.slice(12..12 + values_len); - let values = Buffer::from(values); - Self::from_parts(values, uncompressed_bytes) - } -} - -/// PrefixKey stores a small suffix fingerprint (prefix bytes + length metadata). -#[derive(Debug, Clone, Copy)] -#[repr(C)] -pub(crate) struct PrefixKey { - prefix7: [u8; 7], - /// Suffix length in bytes (after shared prefix), or 255 if >= 255 / unknown. - len: u8, -} - -impl PrefixKey { - pub(crate) const fn prefix_len() -> usize { - 7 - } - - /// Construct from the full suffix bytes (after shared prefix). - /// Embeds up to `prefix_len()` bytes into `prefix7` and stores length (or 255 if >=255). - pub(crate) fn new(suffix_bytes: &[u8]) -> Self { - let mut prefix7 = [0u8; 7]; - let copy_len = std::cmp::min(Self::prefix_len(), suffix_bytes.len()); - if copy_len > 0 { - prefix7[..copy_len].copy_from_slice(&suffix_bytes[..copy_len]); - } - let len = if suffix_bytes.len() >= 255 { - 255u8 - } else { - suffix_bytes.len() as u8 - }; - Self { prefix7, len } - } - - /// Construct directly from stored parts (used by deserialization only) - pub(crate) fn from_parts(prefix7: [u8; 7], len: u8) -> Self { - Self { prefix7, len } - } - - #[inline] - pub(crate) fn prefix7(&self) -> &[u8; 7] { - &self.prefix7 - } - - #[inline] - pub(crate) fn len_byte(&self) -> u8 { - self.len - } - - #[cfg(test)] - pub(crate) fn known_suffix_len(&self) -> Option { - if self.len == 255 { - None - } else { - Some(self.len as usize) - } - } -} - -const _: () = if std::mem::size_of::() != 8 { - panic!("PrefixKey must be 8 bytes") -}; - -#[derive(Debug, Clone, Copy)] -struct CompactOffsetHeader { - slope: i32, - intercept: i32, - offset_bytes: u8, // 1, 2, or 4 bytes per residual -} - -#[derive(Debug, Clone)] -enum OffsetResiduals { - One(Arc<[i8]>), - Two(Arc<[i16]>), - Four(Arc<[i32]>), -} - -impl OffsetResiduals { - fn len(&self) -> usize { - match self { - Self::One(values) => values.len(), - Self::Two(values) => values.len(), - Self::Four(values) => values.len(), - } - } - - #[cfg(test)] - fn bytes_per(&self) -> usize { - match self { - Self::One(_) => 1, - Self::Two(_) => 2, - Self::Four(_) => 4, - } - } - - fn get_i32(&self, index: usize) -> i32 { - match self { - Self::One(values) => values[index] as i32, - Self::Two(values) => values[index] as i32, - Self::Four(values) => values[index], - } - } -} - -/// Compact offset index for FSST dictionary values (includes the final sentinel offset). -#[derive(Debug, Clone)] -pub(crate) struct CompactOffsets { - header: CompactOffsetHeader, - residuals: OffsetResiduals, -} - -// Proper least-squares linear regression -fn fit_line(offsets: &[u32]) -> (i32, i32) { - let n = offsets.len(); - if n <= 1 { - return (0, offsets.first().copied().unwrap_or(0) as i32); - } - - let n_f64 = n as f64; - - // Sum of indices: 0 + 1 + 2 + ... + (n-1) = n*(n-1)/2 - let sum_x = (n * (n - 1) / 2) as f64; - - // Sum of offsets - let sum_y: f64 = offsets.iter().map(|&o| o as f64).sum(); - - // Sum of (index * offset) - let sum_xy: f64 = offsets - .iter() - .enumerate() - .map(|(i, &o)| i as f64 * o as f64) - .sum(); - - // Sum of index squared: 0² + 1² + 2² + ... + (n-1)² = n*(n-1)*(2n-1)/6 - let sum_x_sq = (n * (n - 1) * (2 * n - 1) / 6) as f64; - - // Least squares formulas - let slope = (n_f64 * sum_xy - sum_x * sum_y) / (n_f64 * sum_x_sq - sum_x * sum_x); - let intercept = (sum_y - slope * sum_x) / n_f64; - - (slope.round() as i32, intercept.round() as i32) -} - -impl CompactOffsets { - pub(crate) fn from_offsets(offsets: &[u32]) -> Self { - if offsets.is_empty() { - return Self { - header: CompactOffsetHeader { - slope: 0, - intercept: 0, - offset_bytes: 1, - }, - residuals: OffsetResiduals::One(Arc::new([])), - }; - } - - let (slope, intercept) = fit_line(offsets); - - let mut offset_residuals: Vec = Vec::with_capacity(offsets.len()); - let mut min_residual = i32::MAX; - let mut max_residual = i32::MIN; - for (index, &offset) in offsets.iter().enumerate() { - let predicted = slope * index as i32 + intercept; - let residual = offset as i32 - predicted; - offset_residuals.push(residual); - min_residual = min_residual.min(residual); - max_residual = max_residual.max(residual); - } - - let offset_bytes = if min_residual >= i8::MIN as i32 && max_residual <= i8::MAX as i32 { - 1 - } else if min_residual >= i16::MIN as i32 && max_residual <= i16::MAX as i32 { - 2 - } else { - 4 - }; - - let residuals = match offset_bytes { - 1 => OffsetResiduals::One( - offset_residuals - .iter() - .map(|&r| r as i8) - .collect::>() - .into(), - ), - 2 => OffsetResiduals::Two( - offset_residuals - .iter() - .map(|&r| r as i16) - .collect::>() - .into(), - ), - 4 => OffsetResiduals::Four(offset_residuals.into()), - _ => unreachable!("offset_bytes must be 1, 2, or 4"), - }; - - Self { - header: CompactOffsetHeader { - slope, - intercept, - offset_bytes, - }, - residuals, - } - } - - pub(crate) fn len(&self) -> usize { - self.residuals.len() - } - - pub(crate) fn get_offset(&self, index: usize) -> u32 { - let predicted = self.header.slope * index as i32 + self.header.intercept; - (predicted + self.residuals.get_i32(index)) as u32 - } - - pub(crate) fn offsets(&self) -> Vec { - (0..self.len()).map(|i| self.get_offset(i)).collect() - } - - pub(crate) fn memory_usage(&self) -> usize { - let header_size = std::mem::size_of::(); - let residuals_size = match &self.residuals { - OffsetResiduals::One(values) => values.len() * std::mem::size_of::(), - OffsetResiduals::Two(values) => values.len() * std::mem::size_of::(), - OffsetResiduals::Four(values) => values.len() * std::mem::size_of::(), - }; - header_size + residuals_size - } -} - -pub(crate) fn empty_compact_offsets() -> CompactOffsets { - CompactOffsets::from_offsets(&[]) -} - -const SYMBOL_SIZE_BYTES: usize = std::mem::size_of::(); - -pub(crate) fn train_compressor<'a, I>(iter: I) -> Compressor -where - I: Iterator, -{ - let strings: Vec<&[u8]> = iter.collect(); - fsst::Compressor::train(&strings) -} - -/// In-memory FSST dictionary buffer that bundles compressed bytes, compact offsets, and the -/// compressor needed to (de)compress values. -#[derive(Clone)] -pub struct FsstArray { - compressor: Arc, - raw: Arc, - compact_offsets: CompactOffsets, -} - -impl std::fmt::Debug for FsstArray { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("FsstBuffer") - .field("raw", &self.raw) - .field("compact_offsets", &"") - .field("compressor", &"") - .finish() - } -} - -impl FsstArray { - pub(crate) fn new( - raw: Arc, - compact_offsets: CompactOffsets, - compressor: Arc, - ) -> Self { - Self { - compressor, - raw, - compact_offsets, - } - } - - pub(crate) fn from_byte_offsets( - raw: Arc, - byte_offsets: &[u32], - compressor: Arc, - ) -> Self { - Self::new(raw, CompactOffsets::from_offsets(byte_offsets), compressor) - } - - pub(crate) fn raw_to_bytes(&self) -> Vec { - self.raw.to_bytes() - } - - pub(crate) fn write_compact_offsets(&self, out: &mut Vec) { - self.compact_offsets.write_residuals(out) - } - - /// Trains a compressor on a sequence of strings. - pub fn train_compressor<'a>(input: impl Iterator) -> Compressor { - train_compressor(input) - } - - /// Creates a new FSST buffer from a GenericByteArray and a compressor. - pub fn from_byte_array_with_compressor( - input: &GenericByteArray, - compressor: Arc, - ) -> Self { - let iter = input.iter(); - let mut compress_buffer = Vec::with_capacity(2 * 1024 * 1024); - let (raw, offsets) = - RawFsstBuffer::from_byte_slices(iter, compressor.clone(), &mut compress_buffer); - Self::from_byte_offsets(Arc::new(raw), &offsets, compressor) - } - - /// Creates a new FSST buffer from a Decimal128Array and a compressor. - pub fn from_decimal128_array_with_compressor( - array: &Decimal128Array, - compressor: Arc, - ) -> Self { - let iter = array.iter().map(|v| v.map(|v| v.to_le_bytes())); - let mut compress_buffer = Vec::with_capacity(64); - let (raw, offsets) = - RawFsstBuffer::from_byte_slices(iter, compressor.clone(), &mut compress_buffer); - Self::from_byte_offsets(Arc::new(raw), &offsets, compressor) - } - - /// Creates a new FSST buffer from a Decimal256Array and a compressor. - pub fn from_decimal256_array_with_compressor( - array: &Decimal256Array, - compressor: Arc, - ) -> Self { - let iter = array.iter().map(|v| v.map(|v| v.to_le_bytes())); - let mut compress_buffer = Vec::with_capacity(128); - let (raw, offsets) = - RawFsstBuffer::from_byte_slices(iter, compressor.clone(), &mut compress_buffer); - Self::from_byte_offsets(Arc::new(raw), &offsets, compressor) - } - - /// Returns the uncompressed byte size of this buffer. - pub fn uncompressed_bytes(&self) -> usize { - ::uncompressed_bytes(self) - } - - /// Returns the in-memory size of this buffer. - pub fn get_array_memory_size(&self) -> usize { - ::get_array_memory_size(self) - } - - /// Returns the number of values in this buffer. - #[allow(clippy::len_without_is_empty)] - pub fn len(&self) -> usize { - self.compact_offsets.len().saturating_sub(1) - } - - /// Returns a decompressor for this buffer. - pub fn decompressor(&self) -> Decompressor<'_> { - self.compressor.decompressor() - } - - /// Returns a reference to the compressor. - pub fn compressor(&self) -> &Compressor { - &self.compressor - } - - /// Returns a clone of the shared compressor. - pub fn compressor_arc(&self) -> Arc { - self.compressor.clone() - } - - /// Serializes this FSST buffer (raw bytes + compact offsets) to `out`. - pub fn to_bytes(&self, out: &mut Vec) { - out.extend_from_slice(&self.raw.to_bytes()); - self.compact_offsets.write_residuals(out); - } - - /// Deserializes a FSST buffer from the `to_bytes()` format. - pub fn from_bytes(bytes: bytes::Bytes, compressor: Arc) -> Self { - if bytes.len() < 12 { - panic!("Input buffer too small for RawFsstBuffer header"); - } - - let raw_values_len = u32::from_le_bytes(bytes[8..12].try_into().unwrap()) as usize; - let raw_len = 12 + raw_values_len; - if raw_len > bytes.len() { - panic!("RawFsstBuffer extends beyond input buffer"); - } - - let raw = RawFsstBuffer::from_bytes(bytes.slice(0..raw_len)); - let compact = decode_compact_offsets(&bytes[raw_len..]); - - if compact.len() > 0 { - let last = compact.get_offset(compact.len().saturating_sub(1)) as usize; - debug_assert_eq!( - last, - raw.values_len(), - "offsets must end at raw values length" - ); - } - - Self::new(Arc::new(raw), compact, compressor) - } - - /// Decompress all values into an Arrow byte array. - pub fn to_arrow_byte_array>(&self) -> GenericByteArray { - let (value_buffer, offsets) = self.to_uncompressed(); - unsafe { GenericByteArray::::new_unchecked(offsets, value_buffer, None) } - } - - fn decompress_as_fixed_size_binary(&self, value_width: usize) -> Vec { - let decompressor = self.compressor.decompressor(); - let mut value_buffer: Vec = Vec::with_capacity(self.len() * value_width + 8); - - for i in 0..self.len() { - let compressed = self.get_compressed_slice(i); - let required = decompressor.max_decompression_capacity(compressed) + 8; - value_buffer.reserve(required); - let len = decompressor.decompress_into(compressed, value_buffer.spare_capacity_mut()); - debug_assert!(len == value_width); - let new_len = value_buffer.len() + len; - unsafe { - value_buffer.set_len(new_len); - } - } - value_buffer - } - - fn to_decimal_array_inner(&self, data_type: &ArrowFixedLenByteArrayType) -> Buffer { - let value_width = data_type.value_width(); - Buffer::from(self.decompress_as_fixed_size_binary(value_width)) - } - - /// Converts this FSST buffer to a Decimal128Array. - pub fn to_decimal128_array(&self, data_type: &ArrowFixedLenByteArrayType) -> Decimal128Array { - let value_buffer = self.to_decimal_array_inner(data_type); - let array_builder = ArrayDataBuilder::new(data_type.into()) - .len(self.len()) - .add_buffer(value_buffer); - let array_data = unsafe { array_builder.build_unchecked() }; - Decimal128Array::from(array_data) - } - - /// Converts this FSST buffer to a Decimal256Array. - pub fn to_decimal256_array(&self, data_type: &ArrowFixedLenByteArrayType) -> Decimal256Array { - let value_buffer = self.to_decimal_array_inner(data_type); - let array_builder = ArrayDataBuilder::new(data_type.into()) - .len(self.len()) - .add_buffer(value_buffer); - let array_data = unsafe { array_builder.build_unchecked() }; - Decimal256Array::from(array_data) - } - - #[cfg(test)] - pub(crate) fn offsets_len(&self) -> usize { - self.compact_offsets.len() - } - - #[cfg(test)] - pub(crate) fn offset_bytes(&self) -> u8 { - self.compact_offsets.header.offset_bytes - } - - #[cfg(test)] - pub(crate) fn offsets(&self) -> Vec { - self.compact_offsets.offsets() - } -} -/// FSST backing store for `LiquidByteViewArray` (in-memory or disk-only handle). -pub trait FsstBacking: std::fmt::Debug + Clone + sealed::Sealed { - /// Get the uncompressed bytes of the FSST buffer. - fn uncompressed_bytes(&self) -> usize; - - /// Get the in-memory size of the FSST backing (raw bytes + any in-memory indices). - fn get_array_memory_size(&self) -> usize; -} - -impl sealed::Sealed for FsstArray {} - -impl FsstArray { - pub(crate) fn to_uncompressed(&self) -> (Buffer, OffsetBuffer) { - let offsets = self.compact_offsets.offsets(); - self.raw - .to_uncompressed(&self.compressor.decompressor(), &offsets) - } - - pub(crate) fn get_compressed_slice(&self, dict_index: usize) -> &[u8] { - let start_offset = self.compact_offsets.get_offset(dict_index); - let end_offset = self.compact_offsets.get_offset(dict_index + 1); - self.raw.get_compressed_slice(start_offset, end_offset) - } - - /// Decompress the selected values into a buffer. - pub fn to_uncompressed_selected(&self, selected: &[usize]) -> (Buffer, OffsetBuffer) { - let decompressor = self.compressor.decompressor(); - let mut value_buffer: Vec = Vec::with_capacity(self.uncompressed_bytes() + 8); - let mut out_offsets: OffsetBufferBuilder = OffsetBufferBuilder::new(selected.len()); - - for &dict_index in selected { - let start_offset = self.compact_offsets.get_offset(dict_index); - let end_offset = self.compact_offsets.get_offset(dict_index + 1); - - let compressed_value = self.raw.get_compressed_slice(start_offset, end_offset); - let decompressed_len = - decompressor.decompress_into(compressed_value, value_buffer.spare_capacity_mut()); - let new_len = value_buffer.len() + decompressed_len; - debug_assert!(new_len <= value_buffer.capacity()); - unsafe { - value_buffer.set_len(new_len); - } - out_offsets.push_length(decompressed_len); - } - - (Buffer::from(value_buffer), out_offsets.finish()) - } -} - -impl FsstBacking for FsstArray { - fn uncompressed_bytes(&self) -> usize { - self.raw.uncompressed_bytes() - } - - fn get_array_memory_size(&self) -> usize { - self.raw.get_memory_size() - + self.compact_offsets.memory_usage() - + std::mem::size_of::() - } -} - -impl CompactOffsets { - fn write_residuals(&self, out: &mut Vec) { - out.extend_from_slice(&self.header.slope.to_le_bytes()); - out.extend_from_slice(&self.header.intercept.to_le_bytes()); - out.push(self.header.offset_bytes); - - match &self.residuals { - OffsetResiduals::One(residuals) => { - out.extend(residuals.iter().map(|r| *r as u8)); - } - OffsetResiduals::Two(residuals) => { - for r in residuals.iter() { - out.extend_from_slice(&r.to_le_bytes()); - } - } - OffsetResiduals::Four(residuals) => { - for r in residuals.iter() { - out.extend_from_slice(&r.to_le_bytes()); - } - } - } - } -} - -pub(crate) fn decode_compact_offsets(bytes: &[u8]) -> CompactOffsets { - if bytes.len() < 9 { - panic!("CompactOffsets requires at least 9 bytes for header"); - } - - let slope = i32::from_le_bytes(bytes[0..4].try_into().unwrap()); - let intercept = i32::from_le_bytes(bytes[4..8].try_into().unwrap()); - let offset_bytes = bytes[8] as usize; - if !matches!(offset_bytes, 1 | 2 | 4) { - panic!("Invalid offset_bytes value: {}", offset_bytes); - } - - let header = CompactOffsetHeader { - slope, - intercept, - offset_bytes: offset_bytes as u8, - }; - - let payload = &bytes[9..]; - if !payload.len().is_multiple_of(offset_bytes) { - panic!("Invalid payload size for CompactOffsets"); - } - let count = payload.len() / offset_bytes; - - match offset_bytes { - 1 => { - let residuals: Arc<[i8]> = payload.iter().map(|b| *b as i8).collect::>().into(); - CompactOffsets { - header, - residuals: OffsetResiduals::One(residuals), - } - } - 2 => { - let mut residuals = Vec::with_capacity(count); - for i in 0..count { - let base = i * 2; - residuals.push(i16::from_le_bytes( - payload[base..base + 2].try_into().unwrap(), - )); - } - CompactOffsets { - header, - residuals: OffsetResiduals::Two(residuals.into()), - } - } - 4 => { - let mut residuals = Vec::with_capacity(count); - for i in 0..count { - let base = i * 4; - residuals.push(i32::from_le_bytes( - payload[base..base + 4].try_into().unwrap(), - )); - } - CompactOffsets { - header, - residuals: OffsetResiduals::Four(residuals.into()), - } - } - _ => unreachable!("validated offset_bytes"), - } -} - -/// Saves symbol table from the compressor to a buffer. -/// -/// Format: -/// 1. The first byte is the length of the symbol table as a u8. -/// 2. The next bytes are the lengths of each symbol as u8. -/// 3. The next bytes are the symbols as u64. -pub fn save_symbol_table(compressor: Arc, buffer: &mut Vec) -> Result<()> { - let symbols = compressor.symbol_table(); - let symbols_lengths = compressor.symbol_lengths(); - - if symbols.len() != symbols_lengths.len() { - return Err(Error::new( - ErrorKind::InvalidInput, - "Symbol table and symbol lengths have different lengths", - )); - } - - if symbols.len() > u8::MAX as usize { - return Err(Error::new( - ErrorKind::InvalidInput, - "Symbol table too large", - )); - } - - buffer.push(symbols.len() as u8); - - for &len in symbols_lengths.iter() { - buffer.push(len); - } - - for sym in symbols.iter() { - buffer.extend_from_slice(&sym.to_u64().to_le_bytes()); - } - - Ok(()) -} - -/// Loads symbol table from a buffer saved by `save_symbol_table`. -pub fn load_symbol_table(data: bytes::Bytes) -> Result { - if data.is_empty() { - return Err(Error::new(ErrorKind::InvalidInput, "Empty symbol table")); - } - - let symbol_count = data[0] as usize; - let lengths_start = 1; - let lengths_end = lengths_start + symbol_count; - if lengths_end > data.len() { - return Err(Error::new( - ErrorKind::InvalidInput, - "Buffer too small for symbol lengths", - )); - } - - let lengths = &data[lengths_start..lengths_end]; - let symbols_start = lengths_end; - let symbols_end = symbols_start + symbol_count * SYMBOL_SIZE_BYTES; - if symbols_end > data.len() { - return Err(Error::new( - ErrorKind::InvalidInput, - "Buffer too small for symbols", - )); - } - - let mut symbols = Vec::with_capacity(symbol_count); - for i in 0..symbol_count { - let base = symbols_start + i * SYMBOL_SIZE_BYTES; - let bytes: [u8; SYMBOL_SIZE_BYTES] = - data[base..base + SYMBOL_SIZE_BYTES].try_into().unwrap(); - symbols.push(Symbol::from_slice(&bytes)); - } - - Ok(fsst::Compressor::rebuild_from(symbols, lengths)) -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::{ - array::{Array, Decimal128Builder, StringBuilder}, - datatypes::DataType, - }; - - #[test] - fn test_compact_offset_view_round_trip() { - // Test 1: Small offsets (should use OneByte variant) - let small_offsets = vec![100u32, 105, 110, 115]; - test_round_trip(&small_offsets, "small offsets"); - - // Test 2: Medium offsets (should use TwoBytes variant) - let medium_offsets = vec![1000u32, 2000, 3000, 3500]; - test_round_trip(&medium_offsets, "medium offsets"); - - // Test 3: Large offsets (should use FourBytes variant) - let large_offsets = vec![100000u32, 200000, 300000, 310000]; - test_round_trip(&large_offsets, "large offsets"); - - // Test 4: Mixed scenario with varying prefix lengths - let mixed_offsets = vec![1000u32, 1010, 1020, 1030, 1040, 1050]; - test_round_trip(&mixed_offsets, "mixed scenarios"); - - // Test 5: Edge case - empty values (single sentinel offset) - let empty_offsets: Vec = vec![0]; - test_round_trip(&empty_offsets, "empty values"); - - // Test 6: Single value (one prefix, two offsets) - let single_offset = vec![42u32, 50]; - test_round_trip(&single_offset, "single offset"); - } - - fn test_round_trip(offsets: &[u32], test_name: &str) { - let compact_offsets = CompactOffsets::from_offsets(offsets); - - assert_eq!( - offsets.len(), - compact_offsets.len(), - "Length mismatch in {}", - test_name - ); - for (i, offset) in offsets.iter().enumerate() { - assert_eq!( - compact_offsets.get_offset(i), - *offset, - "Offset mismatch at index {} in {}", - i, - test_name - ); - } - - let mut bytes = Vec::new(); - compact_offsets.write_residuals(&mut bytes); - let reconstructed = decode_compact_offsets(&bytes); - - assert_eq!( - offsets.len(), - reconstructed.len(), - "Reconstructed length mismatch in {}", - test_name - ); - for (i, o) in offsets.iter().enumerate() { - assert_eq!(*o, reconstructed.get_offset(i)); - } - } - - #[test] - fn test_compact_offset_view_memory_efficiency() { - // test that compaction actually saves memory - let offsets = vec![1000u32, 1010, 1020, 1030, 1040]; - - let original_size = offsets.len() * std::mem::size_of::(); - let compact_offsets = CompactOffsets::from_offsets(&offsets); - let compact_size = compact_offsets.memory_usage(); - - // for this test case, we should see some savings due to using smaller residuals - assert!( - compact_size <= original_size, - "Compact representation should not be larger" - ); - } - - #[test] - fn test_compact_offset_view_struct_methods() { - let key = PrefixKey::from_parts([1, 2, 3, 4, 5, 6, 7], 15); - assert_eq!(key.prefix7(), &[1, 2, 3, 4, 5, 6, 7]); - assert_eq!(key.len_byte(), 15); - assert_eq!(key.known_suffix_len(), Some(15)); - - let unknown = PrefixKey::from_parts([7, 6, 5, 4, 3, 2, 1], 255); - assert_eq!(unknown.len_byte(), 255); - assert_eq!(unknown.known_suffix_len(), None); - - let r1 = OffsetResiduals::One(vec![-42i8, 7].into()); - assert_eq!(r1.bytes_per(), 1); - assert_eq!(r1.get_i32(0), -42); - - let r2 = OffsetResiduals::Two(vec![12345i16].into()); - assert_eq!(r2.bytes_per(), 2); - assert_eq!(r2.get_i32(0), 12345); - - let r4 = OffsetResiduals::Four(vec![-1000000i32].into()); - assert_eq!(r4.bytes_per(), 4); - assert_eq!(r4.get_i32(0), -1000000); - - assert_eq!(PrefixKey::prefix_len(), 7); - } - - #[test] - fn test_compact_offset_view_group_from_bytes_errors() { - // Test with insufficient bytes for header - let short_bytes = vec![1, 2, 3]; // only 3 bytes, need at least 9 - let result = std::panic::catch_unwind(|| decode_compact_offsets(&short_bytes)); - assert!(result.is_err(), "Should panic with insufficient bytes"); - - // Test with invalid offset_bytes value - let mut invalid_header = vec![0; 9]; - invalid_header[8] = 3; // invalid offset_bytes (should be 1, 2, or 4) - let result = std::panic::catch_unwind(|| decode_compact_offsets(&invalid_header)); - assert!(result.is_err(), "Should panic with invalid offset_bytes"); - - // Test with misaligned residual data for TwoBytes variant - let mut misaligned_two_bytes = vec![0; 9 + 1]; // header + incomplete residual - misaligned_two_bytes[8] = 2; // offset_bytes = 2 - let result = std::panic::catch_unwind(|| decode_compact_offsets(&misaligned_two_bytes)); - assert!( - result.is_err(), - "Should panic with misaligned TwoBytes residuals" - ); - - // Test with misaligned residual data for FourBytes variant - let mut misaligned_four_bytes = vec![0; 9 + 2]; // header + incomplete residual - misaligned_four_bytes[8] = 4; // offset_bytes = 4 - let result = std::panic::catch_unwind(|| decode_compact_offsets(&misaligned_four_bytes)); - assert!( - result.is_err(), - "Should panic with misaligned FourBytes residuals" - ); - } - - #[test] - fn test_compact_offset_view_group_from_bytes_valid() { - // Test OneByte variant roundtrip - let offsets = vec![100u32, 101, 105]; - let original = CompactOffsets::from_offsets(&offsets); - - let mut bytes = Vec::new(); - original.write_residuals(&mut bytes); - let reconstructed = decode_compact_offsets(&bytes); - - // Verify they match - assert_eq!(offsets.len(), reconstructed.len()); - for (i, o) in offsets.iter().enumerate() { - assert_eq!(*o, reconstructed.get_offset(i)); - } - } - - #[test] - fn test_fsst_buffer_bytes_roundtrip() { - let mut builder = StringBuilder::new(); - for i in 0..1000 { - builder.append_value(format!("test string value {i}")); - } - let original = builder.finish(); - - let compressor = - FsstArray::train_compressor(original.iter().flat_map(|s| s.map(|s| s.as_bytes()))); - let compressor_arc = Arc::new(compressor); - let original_fsst = - FsstArray::from_byte_array_with_compressor(&original, compressor_arc.clone()); - - let mut buffer = Vec::new(); - original_fsst.to_bytes(&mut buffer); - - let bytes = bytes::Bytes::from(buffer); - let deserialized = FsstArray::from_bytes(bytes, compressor_arc); - - let original_strings = original_fsst.to_arrow_byte_array::(); - let deserialized_strings = deserialized.to_arrow_byte_array::(); - assert_eq!(original_strings.len(), deserialized_strings.len()); - for (orig, deser) in original_strings.iter().zip(deserialized_strings.iter()) { - assert_eq!(orig, deser); - } - } - - #[test] - fn test_decimal_compression_smoke() { - let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(10, 2)); - for i in 0..4096 { - builder.append_value(i128::from_le_bytes([(i % 16) as u8; 16])); - } - let original = builder.finish(); - let original_size = original.get_array_memory_size(); - - let values = original - .iter() - .filter_map(|v| v.map(|v| v.to_le_bytes())) - .collect::>(); - let compressor = FsstArray::train_compressor(values.iter().map(|b| b.as_slice())); - let compressor_arc = Arc::new(compressor); - - let fsst = FsstArray::from_decimal128_array_with_compressor(&original, compressor_arc); - let compressed_size = fsst.get_array_memory_size(); - assert!(compressed_size < original_size); - } - - #[test] - fn test_save_and_load_symbol_table_roundtrip() { - let mut builder = StringBuilder::new(); - for i in 0..1000 { - builder.append_value(format!("hello world {i}")); - } - let original = builder.finish(); - - let compressor = - FsstArray::train_compressor(original.iter().flat_map(|s| s.map(|s| s.as_bytes()))); - let compressor_arc = Arc::new(compressor); - - let mut bytes = Vec::new(); - save_symbol_table(compressor_arc.clone(), &mut bytes).unwrap(); - let reloaded = load_symbol_table(bytes::Bytes::from(bytes)).unwrap(); - - let fsst_original = FsstArray::from_byte_array_with_compressor(&original, compressor_arc); - let fsst_reloaded = - FsstArray::from_byte_array_with_compressor(&original, Arc::new(reloaded)); - - let a = fsst_original.to_arrow_byte_array::(); - let b = fsst_reloaded.to_arrow_byte_array::(); - assert_eq!(a, b); - } -} diff --git a/src/core/src/liquid_array/raw/mod.rs b/src/core/src/liquid_array/raw/mod.rs deleted file mode 100644 index e92a96644..000000000 --- a/src/core/src/liquid_array/raw/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -//! Low level array primitives. -//! You should not use this module directly. -//! Instead, use `liquid_cache_datafusion_server` or `liquid_cache_datafusion_client` to interact with LiquidCache. -pub(super) mod bit_pack_array; -/// FSST dictionary backing used by byte-view arrays. -pub mod fsst_buffer; -pub use bit_pack_array::BitPackedArray; -pub use fsst_buffer::FsstArray; diff --git a/src/core/src/liquid_array/tests.rs b/src/core/src/liquid_array/tests.rs deleted file mode 100644 index 5fc575625..000000000 --- a/src/core/src/liquid_array/tests.rs +++ /dev/null @@ -1,199 +0,0 @@ -#[cfg(test)] -mod byte_view_tests { - use std::sync::Arc; - - use arrow::array::{ - Array, AsArray, BinaryViewArray, BooleanArray, DictionaryArray, StringArray, - }; - use arrow::buffer::BooleanBuffer; - use arrow_schema::DataType; - use datafusion_common::ScalarValue; - use datafusion_expr_common::operator::Operator; - use datafusion_physical_expr::PhysicalExpr; - use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; - use rand::SeedableRng; - use rand::prelude::*; - - use crate::liquid_array::raw::FsstArray; - use crate::liquid_array::{LiquidArray, LiquidByteViewArray}; - - fn make_byte_view(input: &StringArray) -> Arc { - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - Arc::new(LiquidByteViewArray::::from_string_array( - input, compressor, - )) - } - - fn gen_random_string(rng: &mut StdRng, max_len: usize) -> String { - let len = rng.random_range(0..=max_len); - let mut out = String::new(); - for _ in 0..len { - let ch = (rng.random_range(0x20u8..=0x7Eu8)) as char; - out.push(ch); - } - out - } - - fn gen_vec_opt_string( - rng: &mut StdRng, - max_items: usize, - max_len: usize, - ) -> Vec> { - let n = rng.random_range(0..=max_items); - (0..n) - .map(|_| { - if rng.random_bool(0.2) { - None - } else { - Some(gen_random_string(rng, max_len)) - } - }) - .collect() - } - - fn gen_vec_opt_bytes( - rng: &mut StdRng, - max_items: usize, - max_len: usize, - ) -> Vec>> { - let n = rng.random_range(0..=max_items); - (0..n) - .map(|_| { - if rng.random_bool(0.2) { - None - } else { - let m = rng.random_range(0..=max_len); - let mut v = vec![0u8; m]; - rng.fill_bytes(&mut v); - Some(v) - } - }) - .collect() - } - - #[test] - fn randomized_utf8_roundtrip() { - for seed in 0..50u64 { - let mut rng = StdRng::seed_from_u64(0xC0FFEE + seed); - let vals = gen_vec_opt_string(&mut rng, 64, 64); - let input = StringArray::from(vals); - let liquid = make_byte_view(&input); - assert_eq!(liquid.to_arrow_array().as_string::(), &input); - } - } - - #[test] - fn randomized_binaryview_roundtrip() { - for seed in 0..50u64 { - let mut rng = StdRng::seed_from_u64(0xB1A5E + seed); - let vals = gen_vec_opt_bytes(&mut rng, 64, 64); - let opt_slices: Vec> = vals.iter().map(|o| o.as_deref()).collect(); - let input = BinaryViewArray::from(opt_slices); - let (_compressor, original) = - LiquidByteViewArray::::train_from_binary_view(&input); - let output = original.to_arrow_array(); - assert_eq!(output.as_binary_view(), &input); - } - } - - #[test] - fn to_dict_arrow_preserves_value_type() { - let input_str = StringArray::from(vec!["hello", "world", "test"]); - let (_c, bv) = LiquidByteViewArray::::train_from_arrow(&input_str); - let dict = bv.to_dict_arrow(); - assert_eq!(dict.values().data_type(), &DataType::Utf8); - - let input_bin = arrow::compute::cast(&input_str, &DataType::Binary) - .unwrap() - .as_binary::() - .clone(); - let (_c, bv) = LiquidByteViewArray::::train_from_arrow(&input_bin); - let dict = bv.to_dict_arrow(); - assert_eq!(dict.values().data_type(), &DataType::Binary); - - let dict_array: DictionaryArray = - DictionaryArray::from_iter(input_str.iter()); - let (_c, bv) = LiquidByteViewArray::::train_from_arrow_dict(&dict_array); - let dict = bv.to_dict_arrow(); - assert_eq!(dict.values().data_type(), &DataType::Utf8); - } - - #[test] - fn to_bytes_and_from_bytes_roundtrip() { - let input = StringArray::from(vec![ - Some("a"), - None, - Some("b"), - Some("a"), - Some("longer text"), - Some(""), - ]); - - let compressor = LiquidByteViewArray::::train_compressor(input.iter()); - let original = - LiquidByteViewArray::::from_string_array(&input, compressor.clone()); - let bytes = original.to_bytes(); - let decoded = LiquidByteViewArray::::from_bytes(bytes.into(), compressor); - let output = decoded.to_arrow_array(); - assert_eq!(output.as_string::(), &input); - } - - #[test] - fn filter_even_indices() { - let input = StringArray::from(vec![ - Some("x"), - Some("y"), - None, - Some("z"), - Some("x"), - Some("y"), - Some("z"), - ]); - let mask = BooleanBuffer::from_iter((0..input.len()).map(|i| i.is_multiple_of(2))); - let liquid = make_byte_view(&input); - let filtered = liquid.filter(&mask).as_string::().clone(); - - let expected_vals: Vec> = (0..input.len()) - .filter(|i| i.is_multiple_of(2)) - .map(|i| { - if input.is_null(i) { - None - } else { - Some(input.value(i)) - } - }) - .collect(); - assert_eq!(filtered, StringArray::from(expected_vals)); - } - - #[test] - fn predicate_eq() { - let input = StringArray::from(vec![ - Some("hello"), - None, - Some("world"), - Some("hello"), - Some(""), - Some("rust"), - ]); - let mask = BooleanBuffer::new_set(input.len()); - - let lit: Arc = - Arc::new(Literal::new(ScalarValue::Utf8(Some("hello".to_string())))); - let col: Arc = Arc::new(Column::new("c", 0)); - let expr: Arc = Arc::new(BinaryExpr::new(col, Operator::Eq, lit)); - - let liquid = make_byte_view(&input); - let result = - liquid.try_eval_predicate(&crate::cache::LiquidExpr::new_unchecked(expr), &mask); - let expected = BooleanArray::from(vec![ - Some(true), - None, - Some(false), - Some(true), - Some(false), - Some(false), - ]); - assert_eq!(result, expected); - } -} diff --git a/src/core/src/liquid_array/utils.rs b/src/core/src/liquid_array/utils.rs deleted file mode 100644 index 19b7f13ab..000000000 --- a/src/core/src/liquid_array/utils.rs +++ /dev/null @@ -1,32 +0,0 @@ -#[cfg(test)] -pub(crate) fn gen_test_decimal_array( - data_type: arrow_schema::DataType, -) -> arrow::array::PrimitiveArray { - use arrow::{ - array::{AsArray, Int64Builder}, - compute::kernels::cast, - }; - - let mut builder = Int64Builder::new(); - for i in 0..4096i64 { - if i % 97 == 0 { - builder.append_null(); - } else { - let value = if i % 5 == 0 { - i * 1000 + 123 - } else if i % 3 == 0 { - 42 - } else if i % 7 == 0 { - i * 1_000_000 + 456789 - } else { - i * 100 + 42 - }; - builder.append_value(value); - } - } - let array = builder.finish(); - cast(&array, &data_type) - .unwrap() - .as_primitive::() - .clone() -} diff --git a/src/core/src/utils/mod.rs b/src/core/src/utils/mod.rs index 4d16e8154..6193b438a 100644 --- a/src/core/src/utils/mod.rs +++ b/src/core/src/utils/mod.rs @@ -1,15 +1,5 @@ //! Utility functions for the storage module. -use std::num::NonZero; - -use arrow::{ - array::{ - ArrayAccessor, ArrayIter, BinaryViewArray, DictionaryArray, GenericByteArray, - GenericByteDictionaryBuilder, PrimitiveArray, PrimitiveDictionaryBuilder, StringViewArray, - }, - datatypes::{BinaryType, ByteArrayType, DecimalType, UInt16Type, Utf8Type}, -}; -use arrow_schema::DataType; use datafusion_common::ScalarValue; pub(crate) mod byte_cache; mod variant_schema; @@ -18,19 +8,6 @@ mod variant_utils; pub use variant_schema::VariantSchema; pub use variant_utils::typed_struct_contains_path; -/// Get the bit width for a given max value. -/// Returns 1 if the max value is 0. -/// Returns 64 - max_value.leading_zeros() as u8 otherwise. -pub(crate) fn get_bit_width(max_value: u64) -> NonZero { - if max_value == 0 { - // todo: here we actually should return 0, as we should just use constant encoding. - // but that's not implemented yet. - NonZero::new(1).unwrap() - } else { - NonZero::new(64 - max_value.leading_zeros() as u8).unwrap() - } -} - pub(crate) fn get_bytes_needle(value: &ScalarValue) -> Option> { match value { ScalarValue::Utf8(Some(v)) => Some(v.as_bytes().to_vec()), @@ -45,114 +22,6 @@ pub(crate) fn get_bytes_needle(value: &ScalarValue) -> Option> { } } -/// A wrapper around `DictionaryArray` that ensures the values are unique. -/// This is because we leverage the fact that the values are unique in the dictionary to short cut the -/// comparison process, i.e., return the index on first match. -/// If the values are not unique, we are screwed. -pub(crate) struct CheckedDictionaryArray { - val: DictionaryArray, -} - -impl CheckedDictionaryArray { - pub fn new_checked(array: &DictionaryArray) -> Self { - gc_dictionary_array(array) - } - - pub fn from_byte_array(array: &GenericByteArray) -> Self { - let iter = array.iter(); - byte_array_to_dict_array::(iter) - } - - pub fn from_string_view_array(array: &StringViewArray) -> Self { - let iter = array.iter(); - byte_array_to_dict_array::(iter) - } - - pub fn from_binary_view_array(array: &BinaryViewArray) -> Self { - let iter = array.iter(); - byte_array_to_dict_array::(iter) - } - - pub fn from_decimal_array(array: &PrimitiveArray) -> Self { - decimal_array_to_dict_array(array) - } - - /// # Safety - /// The caller must ensure that the values in the dictionary are unique. - pub unsafe fn new_unchecked_i_know_what_i_am_doing( - array: &DictionaryArray, - ) -> Self { - #[cfg(debug_assertions)] - { - let gc_ed = gc_dictionary_array(array).val; - assert_eq!( - gc_ed.values().len(), - array.values().len(), - "the input dictionary values are not unique" - ); - } - Self { val: array.clone() } - } - - pub fn into_inner(self) -> DictionaryArray { - self.val - } - - pub fn as_ref(&self) -> &DictionaryArray { - &self.val - } - - pub fn bit_width_for_key(&self) -> NonZero { - let distinct_count = self.as_ref().values().len(); - get_bit_width(distinct_count as u64) - } -} - -fn gc_dictionary_array(array: &DictionaryArray) -> CheckedDictionaryArray { - let value_type = array.values().data_type(); - if let DataType::Binary = value_type { - let typed = array - .downcast_dict::>() - .unwrap(); - let iter = typed.into_iter(); - byte_array_to_dict_array::(iter) - } else if let DataType::Utf8 = value_type { - let typed = array.downcast_dict::>().unwrap(); - let iter = typed.into_iter(); - byte_array_to_dict_array::(iter) - } else { - unreachable!("Unsupported dictionary type: {:?}", value_type); - } -} - -fn decimal_array_to_dict_array( - array: &PrimitiveArray, -) -> CheckedDictionaryArray { - let iter = array.iter(); - let mut builder = - PrimitiveDictionaryBuilder::::with_capacity(array.len(), array.len()); - for s in iter { - builder.append_option(s); - } - let dict = builder.finish(); - CheckedDictionaryArray { val: dict } -} - -fn byte_array_to_dict_array<'a, T: ByteArrayType, I: ArrayAccessor>( - input: ArrayIter, -) -> CheckedDictionaryArray { - let mut builder = GenericByteDictionaryBuilder::::with_capacity( - input.size_hint().0, - input.size_hint().0, - input.size_hint().0, - ); - for s in input { - builder.append_option(s); - } - let dict = builder.finish(); - CheckedDictionaryArray { val: dict } -} - pub(crate) fn yield_now_if_shuttle() { #[cfg(all(feature = "shuttle", test))] shuttle::thread::yield_now(); @@ -186,55 +55,3 @@ pub(crate) fn shuttle_replay(test: impl Fn() + Send + Sync + 'static, schedule: .try_init(); shuttle::replay(test, schedule); } - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{BinaryArray, DictionaryArray}; - use std::sync::Arc; - - fn create_test_dictionary(values: Vec<&[u8]>) -> DictionaryArray { - let binary_array = BinaryArray::from_iter_values(values); - DictionaryArray::new(vec![0u16, 1, 2, 3].into(), Arc::new(binary_array)) - } - - #[test] - fn test_gc_behavior() { - // Test duplicate removal - let dup_dict = create_test_dictionary(vec![b"a", b"a", b"b", b"b"]); - let checked = CheckedDictionaryArray::new_checked(&dup_dict); - let dict_values = checked.as_ref().values(); - assert_eq!(dict_values.len(), 2); - assert_eq!( - dict_values - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - b"a" - ); - assert_eq!( - dict_values - .as_any() - .downcast_ref::() - .unwrap() - .value(1), - b"b" - ); - - // Test already unique values - let unique_dict = create_test_dictionary(vec![b"a", b"b", b"c", b"d"]); - let checked_unique = CheckedDictionaryArray::new_checked(&unique_dict); - assert_eq!(checked_unique.as_ref().values().len(), 4); - } - - #[test] - #[cfg(debug_assertions)] - #[should_panic(expected = "the input dictionary values are not unique")] - fn test_unchecked_duplicates_panic() { - let dup_dict = create_test_dictionary(vec![b"a", b"a", b"b", b"b"]); - unsafe { - CheckedDictionaryArray::new_unchecked_i_know_what_i_am_doing(&dup_dict); - } - } -} diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index fb3eea97f..18e01d5c3 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -103,11 +103,8 @@ fn main() { for (i, id) in ids.iter().enumerate() { let len = lens[i]; let selection = BooleanBuffer::new_set(len); - let Some(liquid_expr) = LiquidExpr::try_new( - Arc::clone(&pred_expr), - &DataType::Utf8, - Some(&liquid_cache::cache::CacheExpression::PredicateColumn), - ) else { + let Some(liquid_expr) = LiquidExpr::try_new(Arc::clone(&pred_expr), &DataType::Utf8) + else { continue; }; if storage diff --git a/src/core/study/filter_selectivity_ambiguity.rs b/src/core/study/filter_selectivity_ambiguity.rs deleted file mode 100644 index 53f281ac9..000000000 --- a/src/core/study/filter_selectivity_ambiguity.rs +++ /dev/null @@ -1,321 +0,0 @@ -use arrow::array::{Array, StringArray, cast::AsArray}; -use arrow::record_batch::RecordBatch; -use arrow_schema::DataType; -use clap::Parser; -use datafusion::prelude::*; -use futures::StreamExt; -use liquid_cache::liquid_array::LiquidByteViewArray; -use liquid_cache::liquid_array::byte_view_array::Comparison; -use liquid_cache::liquid_array::raw::FsstArray; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Clone)] -#[command(name = "SearchPhrase Filter Selectivity Study")] -#[command(about = "Compute prefix selectivity and ambiguity for ClickBench SearchPhrase filters")] -struct CliArgs { - /// Parquet file to read. - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Column to evaluate. - #[arg(long, default_value = "SearchPhrase")] - column: String, - - /// Parquet batch size (rows per RecordBatch). - #[arg(long, default_value_t = 8192 * 2)] - batch_size: usize, - - /// Optional row limit (useful for faster runs). - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors. - #[arg(long, default_value = "false")] - bench: bool, -} - -#[derive(Debug, Clone, Copy)] -struct FilterSpec { - generation: u64, - count: usize, - literal: &'static str, -} - -#[derive(Debug)] -struct FilterQuery { - spec: FilterSpec, - op: Comparison, - needle: Vec, - selected_rows: usize, - ambiguous_rows: usize, - distinct_rows: usize, -} - -struct ScanConfig<'a> { - column: &'a str, - limit: Option, -} - -const FILTERS: &[FilterSpec] = &[ - FilterSpec { - generation: 2, - count: 2, - literal: "in-grid madonnasekret@yandex спб", - }, - FilterSpec { - generation: 3, - count: 3, - literal: "erection пермь курском звучка штильники скривода моряков адлера", - }, - FilterSpec { - generation: 4, - count: 36, - literal: "0б1 купить билето.одноклавович и сотряд", - }, - FilterSpec { - generation: 5, - count: 3, - literal: "0б1 купить бамбарды", - }, - FilterSpec { - generation: 6, - count: 1, - literal: "0986 года на скрыть", - }, - FilterSpec { - generation: 7, - count: 1, - literal: "03-85 серпухоль по краснодар", - }, - FilterSpec { - generation: 8, - count: 1, - literal: "0б1 купить клин-себационные моя мультики на карта", - }, - FilterSpec { - generation: 9, - count: 17, - literal: "03-85 серпухоль по краснодар", - }, - FilterSpec { - generation: 10, - count: 3, - literal: "(http://kommer aspire", - }, - FilterSpec { - generation: 11, - count: 19, - literal: "(http://kommedium=cpc&utm_source=main происход", - }, - FilterSpec { - generation: 12, - count: 1, - literal: "(http://kommedium=cpc&utm_source=main произвестивозачать на автомобиле", - }, - FilterSpec { - generation: 13, - count: 9, - literal: "(http://kommed acce maximum 5*", - }, - FilterSpec { - generation: 14, - count: 2, - literal: "'kbnysq rbyjgjbcr ghjbpdjlcndf dbltj ujhs", - }, - FilterSpec { - generation: 15, - count: 25, - literal: "'kbnysq gbhj;rb gjkmpjdfz", - }, - FilterSpec { - generation: 16, - count: 155, - literal: "'kbnysq ctrcfylhtq d vfrcbvev", - }, - FilterSpec { - generation: 17, - count: 96, - literal: "'kbnysq ctrcdfqafq yjxm. usb-накомсомосква", - }, - FilterSpec { - generation: 18, - count: 131, - literal: "'kbnysq ctrcdfqafq yf cello (feat nival telligar", - }, - FilterSpec { - generation: 19, - count: 84, - literal: "'kbnysq ctrcdfqafq lj;lz", - }, - FilterSpec { - generation: 20, - count: 155, - literal: "'kbnysq cgtrnjh xperia mazda trash", - }, - FilterSpec { - generation: 21, - count: 294, - literal: "'kbnysq cgtrnjh xperia mazda temperia ua авторіа.ua автория", - }, - FilterSpec { - generation: 22, - count: 254, - literal: "'kbnysq cgtrnjh xperia mazda te editional", - }, - FilterSpec { - generation: 23, - count: 543, - literal: "'kbnyjuj ljvfiyb[ pfcnhjqrb bp vfccfl 634 картак", - }, - FilterSpec { - generation: 24, - count: 509, - literal: "'kbnyjuj ljvfiyb[ pfcnhjqcndj vfr 16523-28 днепродажа", - }, - FilterSpec { - generation: 25, - count: 2963, - literal: "'kbnyjuj ljvfiyb[ pfcjkbndf", - }, - FilterSpec { - generation: 26, - count: 284, - literal: "'exist.androit dogs/tags tuning dogg", - }, - FilterSpec { - generation: 27, - count: 250, - literal: "'exist.androit dogs купить памятников шарарок сатист онлайн", - }, - FilterSpec { - generation: 28, - count: 1, - literal: "'exist.androit dogs ever", - }, - FilterSpec { - generation: 29, - count: 226, - literal: "'cgjhnhtnm jnpsds madded", - }, - FilterSpec { - generation: 30, - count: 177, - literal: "'atrn dynami инстроги на добразовая", - }, - FilterSpec { - generation: 31, - count: 5193, - literal: "$dmini+7", - }, - FilterSpec { - generation: 32, - count: 144, - literal: "$_posten of greenjera mi 300 мегафонов (1944-105 отзывы", - }, - FilterSpec { - generation: 36, - count: 62, - literal: "$dmini+7", - }, - FilterSpec { - generation: 37, - count: 994, - literal: "$_posten of greenjera mi 300 мегафонов (1944-105 отзывы", - }, -]; - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(args.batch_size); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - let mut filters = build_filters(); - let scan_config = ScanConfig { - column: &args.column, - limit: args.limit, - }; - scan_column(&ctx, &scan_config, &mut filters).await; - - for filter in &filters { - let selectivity = filter.selected_rows as f64 / filter.distinct_rows as f64; - let ambiguous_rate = filter.ambiguous_rows as f64 / filter.distinct_rows as f64; - - println!( - r#""gen": {}, "filter": "{}", "count": {}, "selectivity": {}, "ambiguous_rate": {}"#, - filter.spec.generation, - filter.spec.literal, - filter.spec.count, - selectivity, - ambiguous_rate - ); - } -} - -fn build_filters() -> Vec { - FILTERS - .iter() - .map(|spec| FilterQuery { - spec: *spec, - op: Comparison::Lt, - needle: spec.literal.as_bytes().to_vec(), - selected_rows: 0, - ambiguous_rows: 0, - distinct_rows: 0, - }) - .collect() -} - -async fn scan_column( - ctx: &SessionContext, - config: &ScanConfig<'_>, - filters: &mut [FilterQuery], -) -> usize { - let sql = if let Some(limit) = config.limit { - format!("SELECT \"{}\" FROM \"hits\" LIMIT {}", config.column, limit) - } else { - format!("SELECT \"{}\" FROM \"hits\"", config.column) - }; - let df = ctx.sql(&sql).await.expect("create df"); - let mut stream = df.execute_stream().await.expect("execute stream"); - - let mut total_rows = 0usize; - while let Some(batch) = stream.next().await { - let batch = batch.expect("fetch batch"); - if batch.num_rows() == 0 { - continue; - } - let array = column_as_string_array(&batch, 0); - total_rows += array.len(); - - let (_compressor, byte_view) = LiquidByteViewArray::::train_from_arrow(&array); - for filter in filters.iter_mut() { - let (selected_rows, ambiguous_rows, distinct_rows) = - byte_view.prefix_compare_counts(&filter.needle, &filter.op); - filter.selected_rows += selected_rows; - filter.ambiguous_rows += ambiguous_rows; - filter.distinct_rows += distinct_rows; - } - } - - total_rows -} - -fn column_as_string_array(batch: &RecordBatch, index: usize) -> StringArray { - let array = batch.column(index).clone(); - let array = if array.data_type() == &DataType::Utf8 { - array - } else { - arrow::compute::cast(&array, &DataType::Utf8).expect("cast to Utf8") - }; - array.as_string::().clone() -} diff --git a/src/core/study/fsst_selectivity.rs b/src/core/study/fsst_selectivity.rs deleted file mode 100644 index 9e0ae16e4..000000000 --- a/src/core/study/fsst_selectivity.rs +++ /dev/null @@ -1,244 +0,0 @@ -use std::sync::Arc; -use std::time::Instant; - -use arrow::array::{Array, ArrayRef, StringArray, cast::AsArray}; -use arrow_schema::DataType; -use clap::Parser; -use datafusion::prelude::*; -use liquid_cache::liquid_array::raw::FsstArray; -use rand::SeedableRng; -use rand::rngs::StdRng; -use rand::seq::{SliceRandom, index::sample}; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Clone)] -#[command(name = "FSST Selected Decode Benchmark")] -#[command(about = "Benchmark FsstArray::to_uncompressed_selected at multiple selectivities")] -struct CliArgs { - /// Parquet file to read. - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Columns to process (comma-separated). - #[arg(long, value_delimiter = ',', default_value = "Title,URL")] - columns: Vec, - - /// Optional row limit for each column (useful for faster runs). - #[arg(long)] - limit: Option, - - /// Parquet batch size (rows per RecordBatch). - #[arg(long, default_value_t = 8192 * 2)] - batch_size: usize, - - /// Selectivities to benchmark (percent, comma-separated). - #[arg(long, value_delimiter = ',', default_value = "1,10,50,99")] - selectivities: Vec, - - /// Iterations per selectivity. - #[arg(long, default_value_t = 5)] - iterations: usize, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors. - #[arg(long, default_value = "false")] - bench: bool, -} - -struct Selection { - pct: u8, - indices: Vec, - approx_bytes: usize, -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(args.batch_size); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - let mut score = 0; - for (col_idx, column) in args.columns.iter().enumerate() { - let array = load_column_array(&ctx, column, args.limit).await; - let row_count = array.len(); - if row_count == 0 { - println!("Column {column}: no rows"); - continue; - } - - let compressor = Arc::new(FsstArray::train_compressor( - array.iter().flatten().map(|value| value.as_bytes()), - )); - let fsst = FsstArray::from_byte_array_with_compressor(&array, compressor); - let total_uncompressed = fsst.uncompressed_bytes(); - let avg_len = total_uncompressed as f64 / row_count as f64; - drop(array); - - println!( - "Column {column}\n rows: {row_count}\n uncompressed: {}\n avg_len: {:.2} bytes", - format_bytes(total_uncompressed), - avg_len - ); - - let selections = build_selections(row_count, avg_len, &args.selectivities, col_idx as u64); - - for selection in selections { - if selection.indices.is_empty() { - println!( - " selectivity {:>3}% -> selected: 0 (skipped)", - selection.pct - ); - continue; - } - - // Run once to reduce cold-start noise. - std::hint::black_box(fsst.to_uncompressed_selected(&selection.indices)); - - let mut total = 0.0; - let mut min = f64::MAX; - let mut max = 0.0_f64; - for _ in 0..args.iterations { - let start = Instant::now(); - let output = fsst.to_uncompressed_selected(&selection.indices); - std::hint::black_box(output); - let elapsed = start.elapsed().as_secs_f64(); - total += elapsed; - min = min.min(elapsed); - max = max.max(elapsed); - } - - let avg = total / args.iterations as f64; - let values_per_sec = selection.indices.len() as f64 / avg; - let mb_per_sec = selection.approx_bytes as f64 / avg / (1024.0 * 1024.0); - - println!( - " selectivity {:>3}% -> selected: {:>8} | avg: {:>8.6}s | min: {:>8.6}s | max: {:>8.6}s | {:>10.1} values/s | {:>8.1} MiB/s", - selection.pct, - selection.indices.len(), - avg, - min, - max, - values_per_sec, - mb_per_sec - ); - score += mb_per_sec as usize; - } - } - println!("Final score: {score}"); -} - -async fn load_column_array( - ctx: &SessionContext, - column: &str, - limit: Option, -) -> StringArray { - let sql = if let Some(limit) = limit { - format!("SELECT \"{column}\" FROM \"hits\" LIMIT {limit}") - } else { - format!("SELECT \"{column}\" FROM \"hits\"") - }; - let df = ctx.sql(&sql).await.expect("create df"); - let batches = df.collect().await.expect("collect"); - - if batches.is_empty() { - return StringArray::from(Vec::>::new()); - } - - let mut arrays = Vec::with_capacity(batches.len()); - for batch in batches { - let array = batch.column(0).clone(); - let array = if array.data_type() == &DataType::Utf8 { - array - } else { - arrow::compute::cast(&array, &DataType::Utf8).expect("cast to Utf8") - }; - arrays.push(array); - } - - concat_utf8_arrays(arrays) -} - -fn concat_utf8_arrays(arrays: Vec) -> StringArray { - if arrays.is_empty() { - return StringArray::from(Vec::>::new()); - } - - let refs: Vec<&dyn Array> = arrays.iter().map(|array| array.as_ref()).collect(); - let concatenated = arrow::compute::concat(&refs).expect("concat arrays"); - concatenated.as_string::().clone() -} - -fn build_selections(len: usize, avg_len: f64, selectivities: &[u8], seed: u64) -> Vec { - selectivities - .iter() - .copied() - .map(|pct| { - let mut rng = StdRng::seed_from_u64(seed ^ (pct as u64).wrapping_mul(0x9e37_79b9)); - let indices = build_selection(len, pct, &mut rng); - let approx_bytes = (avg_len * indices.len() as f64).round() as usize; - Selection { - pct, - indices, - approx_bytes, - } - }) - .collect() -} - -fn build_selection(len: usize, pct: u8, rng: &mut StdRng) -> Vec { - if len == 0 || pct == 0 { - return Vec::new(); - } - - let mut target = len.saturating_mul(pct as usize) / 100; - if target == 0 { - target = 1; - } - if target >= len { - return (0..len).collect(); - } - - let mut selected = if target > len / 2 { - let remove_count = len - target; - let mut remove = vec![false; len]; - for idx in sample(rng, len, remove_count) { - remove[idx] = true; - } - let mut indices = Vec::with_capacity(target); - for (idx, should_remove) in remove.iter().enumerate() { - if !should_remove { - indices.push(idx); - } - } - indices - } else { - sample(rng, len, target).into_vec() - }; - - selected.shuffle(rng); - selected -} - -fn format_bytes(bytes: usize) -> String { - const KB: f64 = 1024.0; - const MB: f64 = 1024.0 * KB; - const GB: f64 = 1024.0 * MB; - let b = bytes as f64; - if b >= GB { - format!("{:.2} GiB", b / GB) - } else if b >= MB { - format!("{:.2} MiB", b / MB) - } else if b >= KB { - format!("{:.2} KiB", b / KB) - } else { - format!("{bytes} B") - } -} diff --git a/src/core/study/fsst_view.rs b/src/core/study/fsst_view.rs deleted file mode 100644 index d1d142b25..000000000 --- a/src/core/study/fsst_view.rs +++ /dev/null @@ -1,765 +0,0 @@ -use std::collections::HashSet; -use std::fmt::Display; -use std::fs::File; -use std::io::{Cursor, Write}; -use std::sync::Arc; -use std::time::Instant; - -use clap::Parser; - -use arrow::array::{Array, AsArray, RecordBatch, StringViewArray}; -use arrow::compute::cast; -use arrow::datatypes::{DataType, Float64Type}; -use arrow_schema::{Field, Schema}; -use datafusion::prelude::*; -use fsst::Compressor; -use liquid_cache::liquid_array::byte_view_array::ByteViewArrayMemoryUsage; -use liquid_cache::liquid_array::byte_view_array::{ByteViewOperator, Comparison, Equality}; -use liquid_cache::liquid_array::raw::FsstArray; -use liquid_cache::liquid_array::{LiquidArray, LiquidByteViewArray}; -use rand::{rng, seq::SliceRandom}; -use serde::{Deserialize, Serialize}; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Debug, Clone, PartialEq)] -enum WorkloadType { - EncodeDecode, - FindNeedle, - CmpNeedle, -} - -impl WorkloadType { - fn from_str(s: &str) -> Result { - match s { - "encode_decode" => Ok(Self::EncodeDecode), - "find_needle" => Ok(Self::FindNeedle), - "cmp_needle" => Ok(Self::CmpNeedle), - _ => Err(format!("Unknown workload: {s}")), - } - } - - fn parse_workloads(s: &str) -> Result, String> { - if s == "all" { - Ok(vec![Self::EncodeDecode, Self::FindNeedle, Self::CmpNeedle]) - } else { - s.split(',') - .map(|w| Self::from_str(w.trim())) - .collect::, _>>() - } - } -} - -#[derive(Parser)] -#[command(name = "FSST View Benchmark")] -#[command(about = "A benchmark tool for comparing different array compression techniques")] -struct CliArgs { - /// Workload type to run - #[arg(long, default_value = "all")] - #[arg( - help = "Workload to run: encode_decode, find_needle, cmp_needle, sort, all, or comma-separated list (e.g., encode_decode,sort)" - )] - workload: String, - - /// Benchmark type to run - #[arg(long)] - #[arg(help = "Benchmark type to run: fsst_view, string_array, string_array_lz4, or all")] - benchmark: Option, - - /// Columns to process - #[arg(long)] - #[arg( - help = "Comma-separated list of columns to process. Available: Title,URL,SearchPhrase,Referer,OriginalURL" - )] - columns: Option, - - #[arg(long, default_value = "false")] - #[arg(help = "make cargo happy")] - bench: bool, -} - -#[derive(Clone)] -struct ColumnData { - data: Vec, - avg_str_length: f64, - distinct_count_ratio: f64, - non_empty_ratio: f64, -} - -async fn download_clickbench_column(column: &str) -> ColumnData { - let config = SessionConfig::default().with_batch_size(8192 * 2); - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet( - "hits", - "../../benchmark/clickbench/data/hits.parquet", - Default::default(), - ) - .await - .unwrap(); - - // Load the column data - let df = ctx - .sql(&format!("SELECT \"{column}\" from \"hits\" limit 10000000")) - .await - .unwrap(); - let batches = df.collect().await.unwrap(); - - let column_data = batches - .iter() - .map(|batch| { - batch - .column_by_name(column) - .unwrap() - .as_string_view() - .clone() - }) - .collect::>(); - - // Get average string length, distinct count ratio, and non-empty ratio in one query - let stats_batches = ctx - .sql(&format!( - "SELECT - AVG(LENGTH(\"{column}\")) AS avg_length, - COUNT(DISTINCT \"{column}\") * 1.0 / COUNT(\"{column}\") AS distinct_ratio, - COUNT(CASE WHEN \"{column}\" IS NOT NULL AND \"{column}\" != '' THEN 1 END) * 1.0 / COUNT(*) AS non_empty_ratio - FROM \"hits\" limit 10000000" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let avg_str_length = stats_batches[0] - .column_by_name("avg_length") - .unwrap() - .as_primitive::() - .clone() - .value(0); - - let distinct_count_ratio = stats_batches[0] - .column_by_name("distinct_ratio") - .unwrap() - .as_primitive::() - .clone() - .value(0); - - let non_empty_ratio = stats_batches[0] - .column_by_name("non_empty_ratio") - .unwrap() - .as_primitive::() - .clone() - .value(0); - - ColumnData { - data: column_data, - avg_str_length, - distinct_count_ratio, - non_empty_ratio, - } -} - -fn load_column_data(column: &str) -> ColumnData { - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(download_clickbench_column(column)) -} - -#[derive(Serialize, Deserialize, Clone)] -struct EncodeResult { - total_size: usize, - encode_time_sec: f64, - decode_time_sec: f64, - workload: String, -} - -#[derive(Serialize, Deserialize, Clone)] -struct FindNeedleResult { - needle_count: usize, - total_search_time_sec: f64, - avg_search_time_per_needle_sec: f64, - avg_search_time_per_needle_ms: f64, - workload: String, -} - -#[derive(Serialize, Deserialize, Clone)] -struct CmpNeedleResult { - needle_count: usize, - total_cmp_time_sec: f64, - avg_cmp_time_per_needle_sec: f64, - avg_cmp_time_per_needle_ms: f64, - workload: String, -} - -#[derive(Serialize, Deserialize, Clone)] -struct BenchmarkResults { - encode_results: Vec, - find_needle_results: Vec, - cmp_needle_results: Vec, -} - -/// Trait for running benchmarks on different array types -trait ArrayBenchmark { - type EncodedData; - - fn encode(&mut self, array: &StringViewArray) -> (Self::EncodedData, f64, usize); - fn run_decode(&self, encoded_data: &Self::EncodedData) -> f64; - fn run_find_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64; - fn run_cmp_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64; - fn workload_name(&self) -> String; -} - -struct BenchmarkRunner; - -impl BenchmarkRunner { - fn run_benchmark( - mut benchmark: T, - arrays: &[StringViewArray], - workloads: &[WorkloadType], - needles: &[String], - ) -> BenchmarkResults { - let mut encode_results = Vec::new(); - let mut find_needle_results = Vec::new(); - let mut cmp_needle_results = Vec::new(); - - // Repeat each workload 3 times - for iteration in 0..3 { - let mut total_encode_time = 0.0; - let mut total_decode_time = 0.0; - let mut total_size = 0; - let mut total_find_needle_time = 0.0; - let mut total_cmp_needle_time = 0.0; - - // First, encode all arrays (this is common for all workloads) - let mut encoded_arrays = Vec::new(); - for array in arrays { - let (encoded_data, encode_time, size) = benchmark.encode(array); - total_encode_time += encode_time; - total_size += size; - encoded_arrays.push(encoded_data); - } - - // Then run each specific workload - for workload in workloads { - match workload { - WorkloadType::EncodeDecode => { - for encoded_data in &encoded_arrays { - let decode_time = benchmark.run_decode(encoded_data); - total_decode_time += decode_time; - } - - let result = EncodeResult { - total_size, - encode_time_sec: total_encode_time, - decode_time_sec: total_decode_time, - workload: benchmark.workload_name(), - }; - println!( - "{} encode/decode (iteration {}): {}", - benchmark.workload_name(), - iteration + 1, - result - ); - encode_results.push(result); - } - WorkloadType::FindNeedle => { - for encoded_data in &encoded_arrays { - let search_time = benchmark.run_find_needle(encoded_data, needles); - total_find_needle_time += search_time; - } - - let needle_count = needles.len(); - let avg_search_time_per_needle_sec = - total_find_needle_time / needle_count as f64; - let result = FindNeedleResult { - needle_count, - total_search_time_sec: total_find_needle_time, - avg_search_time_per_needle_sec, - avg_search_time_per_needle_ms: avg_search_time_per_needle_sec * 1000.0, - workload: benchmark.workload_name(), - }; - println!( - "{} find needle (iteration {}): {}", - benchmark.workload_name(), - iteration + 1, - result - ); - find_needle_results.push(result); - } - WorkloadType::CmpNeedle => { - for encoded_data in &encoded_arrays { - let cmp_time = benchmark.run_cmp_needle(encoded_data, needles); - total_cmp_needle_time += cmp_time; - } - - let needle_count = needles.len(); - let avg_cmp_time_per_needle_sec = - total_cmp_needle_time / needle_count as f64; - let result = CmpNeedleResult { - needle_count, - total_cmp_time_sec: total_cmp_needle_time, - avg_cmp_time_per_needle_sec, - avg_cmp_time_per_needle_ms: avg_cmp_time_per_needle_sec * 1000.0, - workload: benchmark.workload_name(), - }; - println!( - "{} cmp needle (iteration {}): {}", - benchmark.workload_name(), - iteration + 1, - result - ); - cmp_needle_results.push(result); - } - } - } - } - - BenchmarkResults { - encode_results, - find_needle_results, - cmp_needle_results, - } - } - - fn run_workloads( - &self, - arrays: &[StringViewArray], - workloads: &[WorkloadType], - benchmark_filter: Option<&str>, - ) -> Vec { - let needles = select_random_needles(&arrays[0]); - let mut results = Vec::new(); - - match benchmark_filter { - Some("fsst_view") => { - results.push(Self::run_benchmark( - FsstViewBenchmark { compressor: None }, - arrays, - workloads, - &needles, - )); - } - Some("string_array") => { - results.push(Self::run_benchmark( - StringArrayBenchmark, - arrays, - workloads, - &needles, - )); - } - Some("string_array_lz4") => { - results.push(Self::run_benchmark( - StringArrayLz4Benchmark, - arrays, - workloads, - &needles, - )); - } - Some("all") | None => { - results.push(Self::run_benchmark( - FsstViewBenchmark { compressor: None }, - arrays, - workloads, - &needles, - )); - results.push(Self::run_benchmark( - StringArrayBenchmark, - arrays, - workloads, - &needles, - )); - results.push(Self::run_benchmark( - StringArrayLz4Benchmark, - arrays, - workloads, - &needles, - )); - } - Some(unknown) => { - eprintln!("Unknown benchmark type: {unknown}. Using all benchmarks."); - results.push(Self::run_benchmark( - FsstViewBenchmark { compressor: None }, - arrays, - workloads, - &needles, - )); - results.push(Self::run_benchmark( - StringArrayBenchmark, - arrays, - workloads, - &needles, - )); - results.push(Self::run_benchmark( - StringArrayLz4Benchmark, - arrays, - workloads, - &needles, - )); - } - } - - results - } -} - -impl Display for EncodeResult { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{} -- total_size: {} bytes, encode: {} s, decode: {} s", - self.workload, self.total_size, self.encode_time_sec, self.decode_time_sec - ) - } -} - -impl Display for FindNeedleResult { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{} -- needles: {}, total: {:.4} s, avg: {:.3} ms per needle", - self.workload, - self.needle_count, - self.total_search_time_sec, - self.avg_search_time_per_needle_ms - ) - } -} - -impl Display for CmpNeedleResult { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{} -- needles: {}, total: {:.4} s, avg: {:.3} ms per needle", - self.workload, - self.needle_count, - self.total_cmp_time_sec, - self.avg_cmp_time_per_needle_ms - ) - } -} - -/// Select 10 random different strings from the first batch of the array -fn select_random_needles(first_batch: &StringViewArray) -> Vec { - let mut unique_strings = HashSet::new(); - let mut all_strings = Vec::new(); - - // Collect all non-null unique strings from the first batch - for i in 0..first_batch.len() { - if !first_batch.is_null(i) { - let s = first_batch.value(i).to_string(); - if unique_strings.insert(s.clone()) { - all_strings.push(s); - } - } - } - - // Shuffle and take up to 10 strings - let mut rng = rng(); - all_strings.shuffle(&mut rng); - all_strings.into_iter().take(10).collect() -} - -struct FsstViewBenchmark { - compressor: Option>, -} - -impl ArrayBenchmark for FsstViewBenchmark { - type EncodedData = LiquidByteViewArray; - - fn workload_name(&self) -> String { - "FSSTView".to_string() - } - - fn encode(&mut self, array: &StringViewArray) -> (Self::EncodedData, f64, usize) { - // Train compressor only on the first call - let compressor = if let Some(cached_compressor) = &self.compressor { - cached_compressor.clone() - } else { - let (trained_compressor, _) = - LiquidByteViewArray::::train_from_string_view(array); - self.compressor = Some(trained_compressor.clone()); - trained_compressor - }; - - let start = Instant::now(); - let encoded_array = - LiquidByteViewArray::::from_string_view_array(array, compressor); - let encode_time = start.elapsed().as_secs_f64(); - let size = encoded_array.get_array_memory_size(); - (encoded_array, encode_time, size) - } - - fn run_decode(&self, encoded_data: &Self::EncodedData) -> f64 { - let start = Instant::now(); - let _array = encoded_data.to_arrow_array(); - start.elapsed().as_secs_f64() - } - - fn run_find_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - let needle_bytes = needle.as_bytes(); - let _result = - encoded_data.compare_with(needle_bytes, &ByteViewOperator::Equality(Equality::Eq)); - } - start.elapsed().as_secs_f64() - } - - fn run_cmp_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - let needle_bytes = needle.as_bytes(); - let _result = encoded_data - .compare_with(needle_bytes, &ByteViewOperator::Comparison(Comparison::Gt)); - } - start.elapsed().as_secs_f64() - } -} - -struct StringArrayBenchmark; - -impl ArrayBenchmark for StringArrayBenchmark { - type EncodedData = StringViewArray; - - fn workload_name(&self) -> String { - "StringArray".to_string() - } - - fn encode(&mut self, array: &StringViewArray) -> (Self::EncodedData, f64, usize) { - let start = Instant::now(); - let v = cast(array, &DataType::Utf8).unwrap(); - let v = v.as_string::().clone(); - let encode_time = start.elapsed().as_secs_f64(); - let size = v.get_array_memory_size(); - (array.clone(), encode_time, size) - } - - fn run_decode(&self, encoded_data: &Self::EncodedData) -> f64 { - let start = Instant::now(); - let _v = cast(encoded_data, &DataType::Utf8View).unwrap(); - start.elapsed().as_secs_f64() - } - - fn run_find_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - let needle_scalar = arrow::array::StringViewArray::new_scalar(needle.clone()); - let _result = arrow::compute::kernels::cmp::eq(encoded_data, &needle_scalar).unwrap(); - } - start.elapsed().as_secs_f64() - } - - fn run_cmp_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - let needle_scalar = arrow::array::StringViewArray::new_scalar(needle.clone()); - let _result = arrow::compute::kernels::cmp::gt(encoded_data, &needle_scalar).unwrap(); - } - start.elapsed().as_secs_f64() - } -} - -struct StringArrayLz4Benchmark; - -impl ArrayBenchmark for StringArrayLz4Benchmark { - type EncodedData = Vec; - - fn workload_name(&self) -> String { - "StringArrayLZ4".to_string() - } - - fn encode(&mut self, array: &StringViewArray) -> (Self::EncodedData, f64, usize) { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); - let compression = arrow::ipc::CompressionType::LZ4_FRAME; - let options = arrow::ipc::writer::IpcWriteOptions::default() - .try_with_compression(Some(compression)) - .unwrap(); - - let v = cast(array, &DataType::Utf8).unwrap(); - let mut file = vec![]; - let mut writer = arrow::ipc::writer::FileWriter::try_new_with_options( - &mut file, - &schema, - options.clone(), - ) - .unwrap(); - - let start = Instant::now(); - let batch = RecordBatch::try_new(schema.clone(), vec![v]).unwrap(); - writer.write(&batch).unwrap(); - writer.finish().unwrap(); - let encode_time = start.elapsed().as_secs_f64(); - let size = file.len(); - (file, encode_time, size) - } - - fn run_decode(&self, encoded_data: &Self::EncodedData) -> f64 { - let start = Instant::now(); - let mut file = Cursor::new(encoded_data); - let mut reader = arrow::ipc::reader::FileReader::try_new(&mut file, None).unwrap(); - let batch = reader.next().unwrap().unwrap(); - let _v = batch.column(0).as_string::().clone(); - start.elapsed().as_secs_f64() - } - - fn run_find_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - // Decompress and search - let mut file = Cursor::new(encoded_data); - let mut reader = arrow::ipc::reader::FileReader::try_new(&mut file, None).unwrap(); - let batch = reader.next().unwrap().unwrap(); - let string_array = batch.column(0).as_string::(); - - let needle_scalar = arrow::array::StringArray::new_scalar(needle.clone()); - let _result = arrow::compute::kernels::cmp::eq(&string_array, &needle_scalar).unwrap(); - } - start.elapsed().as_secs_f64() - } - - fn run_cmp_needle(&self, encoded_data: &Self::EncodedData, needles: &[String]) -> f64 { - let start = Instant::now(); - for needle in needles { - // Decompress and search - let mut file = Cursor::new(encoded_data); - let mut reader = arrow::ipc::reader::FileReader::try_new(&mut file, None).unwrap(); - let batch = reader.next().unwrap().unwrap(); - let string_array = batch.column(0).as_string::(); - - let needle_scalar = arrow::array::StringArray::new_scalar(needle.clone()); - let _result = arrow::compute::kernels::cmp::gt(&string_array, &needle_scalar).unwrap(); - } - start.elapsed().as_secs_f64() - } -} - -#[derive(Serialize, Deserialize, Clone)] -struct SerializableMemoryUsage { - dictionary_keys: usize, - prefix_keys: usize, - fsst_buffer: usize, - shared_prefix: usize, - string_fingerprints: usize, - struct_size: usize, - total: usize, -} - -impl From for SerializableMemoryUsage { - fn from(usage: ByteViewArrayMemoryUsage) -> Self { - Self { - dictionary_keys: usage.dictionary_key, - prefix_keys: usage.prefix_keys, - fsst_buffer: usage.fsst_buffer, - shared_prefix: usage.shared_prefix, - string_fingerprints: usage.string_fingerprints, - struct_size: usage.struct_size, - total: usage.total(), - } - } -} - -#[derive(Serialize, Deserialize)] -struct BenchmarkResult { - column_name: String, - avg_string_length: f64, - distinct_count_ratio: f64, - non_empty_ratio: f64, - benchmark_results: Vec, - fsst_view_memory_usage: SerializableMemoryUsage, -} - -#[derive(Serialize, Deserialize)] -struct CompleteResults { - benchmark_name: String, - timestamp: String, - columns: Vec, -} - -fn main() { - let args = CliArgs::parse(); - - // Parse and validate workloads - let workloads_to_run = WorkloadType::parse_workloads(&args.workload).unwrap(); - - let all_columns = ["Title", "URL", "SearchPhrase", "Referer", "OriginalURL"]; - let columns_to_process: Vec<&str> = if let Some(ref columns_str) = args.columns { - columns_str.split(',').map(|s| s.trim()).collect() - } else { - all_columns.to_vec() - }; - - let runner = BenchmarkRunner; - let mut all_column_results = Vec::new(); - - println!("Running workloads: {workloads_to_run:?}"); - if let Some(ref benchmark) = args.benchmark { - println!("Benchmark filter: {benchmark}"); - } - println!("Columns: {}", columns_to_process.join(", ")); - println!(); - - for c in columns_to_process { - println!("Loading column: {c}"); - let column_data = load_column_data(c); - println!( - "{c} average length: {:.2}, distinct ratio: {:.4}, non-empty ratio: {:.4}", - column_data.avg_str_length, - column_data.distinct_count_ratio, - column_data.non_empty_ratio - ); - - let benchmark_results = runner.run_workloads( - &column_data.data, - &workloads_to_run, - args.benchmark.as_deref(), - ); - - let (compressor, _) = - LiquidByteViewArray::::train_from_string_view(&column_data.data[0]); - let mut total_detailed_memory_usage = ByteViewArrayMemoryUsage { - dictionary_key: 0, - prefix_keys: 0, - fsst_buffer: 0, - string_fingerprints: 0, - shared_prefix: 0, - struct_size: 0, - }; - - for a in &column_data.data { - let array = - LiquidByteViewArray::::from_string_view_array(a, compressor.clone()); - total_detailed_memory_usage += array.get_detailed_memory_usage(); - } - - all_column_results.push(BenchmarkResult { - column_name: c.to_string(), - avg_string_length: column_data.avg_str_length, - distinct_count_ratio: column_data.distinct_count_ratio, - non_empty_ratio: column_data.non_empty_ratio, - benchmark_results, - fsst_view_memory_usage: total_detailed_memory_usage.into(), - }); - - println!("Finished processing {c}\n"); - } - - // Write all results to JSON file once at the end - let complete_results = CompleteResults { - benchmark_name: "FSST View Study".to_string(), - timestamp: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs() - .to_string(), - columns: all_column_results, - }; - - let json_output = serde_json::to_string_pretty(&complete_results).unwrap(); - let filename = "../../target/benchmark_results.json"; - let mut file = File::create(filename).unwrap(); - file.write_all(json_output.as_bytes()).unwrap(); - - println!("Benchmark results written to {filename}"); - - println!("Benchmark completed!"); -} diff --git a/src/core/study/linear_integer.rs b/src/core/study/linear_integer.rs deleted file mode 100644 index b4651d191..000000000 --- a/src/core/study/linear_integer.rs +++ /dev/null @@ -1,216 +0,0 @@ -use std::time::Instant; - -use arrow::array::{Array, ArrayRef, cast::AsArray}; -use arrow::datatypes::DataType; -use clap::Parser; -use datafusion::prelude::*; -use futures::StreamExt; -use liquid_cache::liquid_array::{ - LiquidArray, LiquidLinearArray, LiquidPrimitiveArray, LiquidPrimitiveType, -}; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Default, Clone)] -#[command(name = "Linear Integer Study")] -#[command(about = "Compare primitive bitpacking vs linear-model encoding for integer columns")] -struct CliArgs { - /// Parquet file to read - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Comma-separated list of columns. If not set, auto-detect integer-like columns. - #[arg(long)] - columns: Option, - - /// Optional row limit for each column (useful for faster runs) - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors - #[arg(long, default_value = "false")] - bench: bool, -} - -#[derive(Default, Debug, Clone)] -struct Stats { - rows: usize, - arrow_bytes: usize, - prim_bytes: usize, - linear_bytes: usize, - prim_encode_sec: f64, - prim_decode_sec: f64, - linear_encode_sec: f64, - linear_decode_sec: f64, -} - -impl Stats { - fn add(&mut self, other: &Stats) { - self.rows += other.rows; - self.arrow_bytes += other.arrow_bytes; - self.prim_bytes += other.prim_bytes; - self.linear_bytes += other.linear_bytes; - self.prim_encode_sec += other.prim_encode_sec; - self.prim_decode_sec += other.prim_decode_sec; - self.linear_encode_sec += other.linear_encode_sec; - self.linear_decode_sec += other.linear_decode_sec; - } -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(8192 * 2); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - // Identify columns - let columns = if let Some(cols) = args.columns.clone() { - cols.split(',') - .map(|s| s.trim().to_string()) - .collect::>() - } else { - autodetect_integer_columns(&ctx).await - }; - - println!("Linear Integer Study on {} column(s)", columns.len()); - - let mut grand = Stats::default(); - for col in columns { - let stats = run_for_column(&ctx, &col, args.limit).await; - println!( - "Column: {col}\n rows: {}\n sizes (bytes) -> arrow: {}, prim: {}, linear: {}\n encode (s) -> prim: {:.6}, linear: {:.6}\n decode (s) -> prim: {:.6}, linear: {:.6}", - stats.rows, - stats.arrow_bytes, - stats.prim_bytes, - stats.linear_bytes, - stats.prim_encode_sec, - stats.linear_encode_sec, - stats.prim_decode_sec, - stats.linear_decode_sec - ); - grand.add(&stats); - } - - println!( - "TOTAL\n rows: {}\n sizes (bytes) -> arrow: {}, prim: {}, linear: {}\n encode (s) -> prim: {:.6}, linear: {:.6}\n decode (s) -> prim: {:.6}, linear: {:.6}", - grand.rows, - grand.arrow_bytes, - grand.prim_bytes, - grand.linear_bytes, - grand.prim_encode_sec, - grand.linear_encode_sec, - grand.prim_decode_sec, - grand.linear_decode_sec - ); -} - -async fn autodetect_integer_columns(ctx: &SessionContext) -> Vec { - // Run a small query to fetch schema - let df = ctx.sql("SELECT * FROM \"hits\" LIMIT 1").await.unwrap(); - let batches = df.collect().await.unwrap(); - let schema = batches[0].schema(); - let mut cols = Vec::new(); - for f in schema.fields() { - if is_integer_like(f.data_type()) { - cols.push(f.name().to_string()); - } - } - cols -} - -fn is_integer_like(dt: &DataType) -> bool { - matches!( - dt, - DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::UInt8 - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::Date32 - | DataType::Date64 - ) -} - -async fn run_for_column(ctx: &SessionContext, column: &str, limit: Option) -> Stats { - let sql = if let Some(n) = limit { - format!("SELECT \"{column}\" FROM \"hits\" LIMIT {n}") - } else { - format!("SELECT \"{column}\" FROM \"hits\"") - }; - let df = ctx.sql(&sql).await.expect("create df"); - let mut stream = df.execute_stream().await.expect("execute stream"); - - let mut stats = Stats::default(); - while let Some(batch_res) = stream.next().await { - let batch = batch_res.expect("stream batch"); - let array: ArrayRef = batch.column(0).clone(); - stats.rows += array.len(); - stats.arrow_bytes += array.get_array_memory_size(); - - let dt = array.data_type().clone(); - match dt { - DataType::Int8 => accumulate::(&array, &mut stats), - DataType::Int16 => accumulate::(&array, &mut stats), - DataType::Int32 => accumulate::(&array, &mut stats), - DataType::Int64 => accumulate::(&array, &mut stats), - DataType::UInt8 => accumulate::(&array, &mut stats), - DataType::UInt16 => accumulate::(&array, &mut stats), - DataType::UInt32 => accumulate::(&array, &mut stats), - DataType::UInt64 => accumulate::(&array, &mut stats), - DataType::Date32 => accumulate::(&array, &mut stats), - DataType::Date64 => accumulate::(&array, &mut stats), - _ => {} - } - } - - stats -} - -fn accumulate(array: &ArrayRef, stats: &mut Stats) -where - ::Native: num_traits::cast::AsPrimitive - + num_traits::FromPrimitive - + num_traits::bounds::Bounded, -{ - let prim = array.as_primitive::().clone(); - - // Primitive bitpacking encode - let t0 = Instant::now(); - let lp = LiquidPrimitiveArray::::from_arrow_array(prim.clone()); - let enc_prim = t0.elapsed().as_secs_f64(); - let prim_bytes = lp.get_array_memory_size(); - - // Primitive decode - let t0 = Instant::now(); - let _ = lp.to_arrow_array(); - let dec_prim = t0.elapsed().as_secs_f64(); - - // Linear encode - let t0 = Instant::now(); - let ll = LiquidLinearArray::::from_arrow_array(prim); - let enc_linear = t0.elapsed().as_secs_f64(); - let linear_bytes = ll.get_array_memory_size(); - - // Linear decode - let t0 = Instant::now(); - let _ = ll.to_arrow_array(); - let dec_linear = t0.elapsed().as_secs_f64(); - - stats.prim_bytes += prim_bytes; - stats.linear_bytes += linear_bytes; - stats.prim_encode_sec += enc_prim; - stats.prim_decode_sec += dec_prim; - stats.linear_encode_sec += enc_linear; - stats.linear_decode_sec += dec_linear; -} diff --git a/src/core/study/prefix_differentiability.rs b/src/core/study/prefix_differentiability.rs deleted file mode 100644 index 6cad21064..000000000 --- a/src/core/study/prefix_differentiability.rs +++ /dev/null @@ -1,202 +0,0 @@ -use ahash::AHashSet; -use arrow::array::{Array, StringArray, cast::AsArray}; -use arrow::record_batch::RecordBatch; -use arrow_schema::DataType; -use clap::Parser; -use datafusion::prelude::*; -use futures::StreamExt; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Parser, Debug, Clone)] -#[command(name = "Prefix Differentiability Study")] -#[command(about = "Compute average prefix differentiability per batch for ClickBench columns")] -struct CliArgs { - /// Parquet file to read. - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Columns to process (comma-separated). - #[arg(long, value_delimiter = ',', default_value = "SearchPhrase,URL,Title")] - columns: Vec, - - /// Parquet batch size (rows per RecordBatch). - #[arg(long, default_value_t = 8192)] - batch_size: usize, - - /// Optional row limit (useful for faster runs). - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors. - #[arg(long, default_value = "false")] - bench: bool, -} - -struct ScanConfig<'a> { - column: &'a str, - limit: Option, - prefix_lengths: &'a [usize], -} - -struct PrefixStats { - prefix_lengths: Vec, - sum_ratios: Vec, - batches: usize, -} - -impl PrefixStats { - fn new(prefix_lengths: &[usize]) -> Self { - Self { - prefix_lengths: prefix_lengths.to_vec(), - sum_ratios: vec![0.0; prefix_lengths.len()], - batches: 0, - } - } - - fn add_batch(&mut self, ratios: &[f64]) { - for (idx, ratio) in ratios.iter().enumerate() { - self.sum_ratios[idx] += ratio; - } - self.batches += 1; - } -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(args.batch_size); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - let prefix_lengths: Vec = (1..=16).collect(); - for column in &args.columns { - let scan_config = ScanConfig { - column, - limit: args.limit, - prefix_lengths: &prefix_lengths, - }; - let stats = scan_column(&ctx, &scan_config).await; - - if stats.batches == 0 { - println!("Column {column}: no rows"); - continue; - } - - println!("Column {column} (batches: {})", stats.batches); - for (prefix_len, sum_ratio) in stats.prefix_lengths.iter().zip(stats.sum_ratios.iter()) { - let avg = sum_ratio / stats.batches as f64; - println!( - " prefix {:>2} -> avg differentiability {:.6}", - prefix_len, avg - ); - } - } -} - -async fn scan_column(ctx: &SessionContext, config: &ScanConfig<'_>) -> PrefixStats { - let sql = if let Some(limit) = config.limit { - format!("SELECT \"{}\" FROM \"hits\" LIMIT {}", config.column, limit) - } else { - format!("SELECT \"{}\" FROM \"hits\"", config.column) - }; - let df = ctx.sql(&sql).await.expect("create df"); - let mut stream = df.execute_stream().await.expect("execute stream"); - - let mut stats = PrefixStats::new(config.prefix_lengths); - while let Some(batch) = stream.next().await { - let batch = batch.expect("fetch batch"); - if batch.num_rows() == 0 { - continue; - } - let array = column_as_string_array(&batch, 0); - if let Some(ratios) = batch_prefix_differentiability(&array, config.prefix_lengths) { - stats.add_batch(&ratios); - } - } - - stats -} - -fn batch_prefix_differentiability( - array: &StringArray, - prefix_lengths: &[usize], -) -> Option> { - let mut values = Vec::with_capacity(array.len()); - for row in 0..array.len() { - if array.is_null(row) { - continue; - } - let value = array.value(row); - if value.is_empty() { - continue; - } - values.push(value.as_bytes().to_vec()); - } - - if values.is_empty() { - return None; - } - - let common_prefix_len = common_prefix_len(&values); - let mut sets: Vec>> = prefix_lengths - .iter() - .map(|_| AHashSet::with_capacity(values.len())) - .collect(); - let mut unique_values: AHashSet> = AHashSet::with_capacity(values.len()); - - for value in values { - let suffix = &value[common_prefix_len..]; - for (idx, prefix_len) in prefix_lengths.iter().enumerate() { - let end = suffix.len().min(*prefix_len); - sets[idx].insert(suffix[..end].to_vec()); - } - unique_values.insert(suffix.to_vec()); - } - - let total_rows = unique_values.len() as f64; - if total_rows == 0.0 { - return None; - } - Some( - sets.iter() - .map(|set| set.len() as f64 / total_rows) - .collect(), - ) -} - -fn common_prefix_len(values: &[Vec]) -> usize { - if values.is_empty() { - return 0; - } - let mut prefix_len = values[0].len(); - for value in values.iter().skip(1) { - let mut idx = 0; - let max_len = prefix_len.min(value.len()); - while idx < max_len && values[0][idx] == value[idx] { - idx += 1; - } - prefix_len = idx; - if prefix_len == 0 { - break; - } - } - prefix_len -} - -fn column_as_string_array(batch: &RecordBatch, index: usize) -> StringArray { - let array = batch.column(index).clone(); - let array = if array.data_type() == &DataType::Utf8 { - array - } else { - arrow::compute::cast(&array, &DataType::Utf8).expect("cast to Utf8") - }; - array.as_string::().clone() -} diff --git a/src/core/study/string-fingerprint.rs b/src/core/study/string-fingerprint.rs deleted file mode 100644 index 9a479f300..000000000 --- a/src/core/study/string-fingerprint.rs +++ /dev/null @@ -1,665 +0,0 @@ -use std::fs; -use std::path::PathBuf; - -use arrow::array::{ArrayRef, AsArray}; -use arrow::datatypes::DataType; -use arrow::record_batch::RecordBatch; -use clap::{Parser, ValueEnum}; -use datafusion::prelude::*; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct Fingerprint { - words: [u64; 4], -} - -impl Fingerprint { - #[inline] - fn from_bytes(bytes: &[u8], mapping: &ByteBucketMap) -> Self { - let mut words = [0u64; 4]; - for &b in bytes { - let bucket = mapping.table[b as usize] as usize; - debug_assert!(bucket < mapping.bucket_count as usize); - words[bucket >> 6] |= 1u64 << (bucket & 63); - } - Self { words } - } - - #[inline] - fn might_contain(&self, pattern: &Self) -> bool { - (self.words[0] & pattern.words[0]) == pattern.words[0] - && (self.words[1] & pattern.words[1]) == pattern.words[1] - && (self.words[2] & pattern.words[2]) == pattern.words[2] - && (self.words[3] & pattern.words[3]) == pattern.words[3] - } -} - -#[derive(Debug, Clone)] -struct ByteBucketMap { - bucket_count: u16, - table: [u8; 256], -} - -impl ByteBucketMap { - fn round_robin(bucket_count: u16) -> Self { - assert!((1..=256).contains(&bucket_count)); - let mut table = [0u8; 256]; - for (b, slot) in table.iter_mut().enumerate() { - *slot = (b as u16 % bucket_count) as u8; - } - Self { - bucket_count, - table, - } - } - - fn contiguous_range(bucket_count: u16) -> Self { - assert!((1..=256).contains(&bucket_count)); - let mut table = [0u8; 256]; - for (b, slot) in table.iter_mut().enumerate() { - *slot = (((b as u16) * bucket_count) >> 8) as u8; - } - Self { - bucket_count, - table, - } - } - - #[allow(dead_code)] - fn from_fn(bucket_count: u16, assign: impl Fn(u8) -> u8) -> Self { - assert!((1..=256).contains(&bucket_count)); - let mut table = [0u8; 256]; - for (b, slot) in table.iter_mut().enumerate() { - let bucket = assign(b as u8); - assert!( - (bucket as u16) < bucket_count, - "custom bucket assignment returned {bucket} (bucket_count={bucket_count})" - ); - *slot = bucket; - } - Self { - bucket_count, - table, - } - } - - fn from_table(bucket_count: u16, table: [u8; 256]) -> Self { - assert!((1..=256).contains(&bucket_count)); - for (b, &bucket) in table.iter().enumerate() { - assert!( - (bucket as u16) < bucket_count, - "custom mapping table[{b}]={bucket} out of range (bucket_count={bucket_count})" - ); - } - Self { - bucket_count, - table, - } - } -} - -#[derive(Clone, Copy, Debug, ValueEnum)] -enum MappingKind { - RoundRobin, - ContiguousRange, - CustomTable, - Optimized, -} - -#[derive(Parser, Debug, Clone)] -#[command(name = "String Fingerprint Study")] -#[command(about = "Compute string fingerprints and evaluate pattern filtering effectiveness")] -struct CliArgs { - /// Parquet file to read. - #[arg(long, default_value = "../../benchmark/clickbench/data/hits.parquet")] - parquet: String, - - /// Columns to process. - #[arg(long, value_delimiter = ',', default_value = "URL,Title,Referer")] - columns: Vec, - - /// Pattern to test (byte-based, default: google). - #[arg(long, default_value = "google")] - pattern: String, - - /// Number of buckets (n). More buckets => longer fingerprint (up to 256). - #[arg(long, value_delimiter = ',', default_value = "4,8,16,32,64")] - buckets: Vec, - - /// Sweep bucket counts from start to end (inclusive). - #[arg(long)] - bucket_start: Option, - - /// Sweep bucket counts from start to end (inclusive). - #[arg(long)] - bucket_end: Option, - - /// Step for --bucket-start/--bucket-end. - #[arg(long, default_value_t = 1)] - bucket_step: u16, - - /// Bucket mapping strategy. - #[arg(long, value_enum, default_value_t = MappingKind::RoundRobin)] - mapping: MappingKind, - - /// Path to a custom mapping table (256 integers in [0, n), separated by whitespace or commas). - #[arg(long)] - custom_map: Option, - - /// For `--mapping optimized`: number of (non-null) values sampled from the start of the column. - #[arg(long, default_value_t = 100)] - sample_values: usize, - - /// Optional row limit per column (useful for faster runs). - #[arg(long)] - limit: Option, - - /// Cargo passes --bench for harness=false binaries; accept it to avoid parse errors. - #[arg(long, default_value = "false")] - bench: bool, -} - -#[derive(Default, Debug, Clone)] -struct Stats { - rows: usize, - nulls: usize, - filtered_out: usize, - candidates: usize, - false_pos: usize, - actual_present: usize, -} - -impl Stats { - fn add(&mut self, other: &Stats) { - self.rows += other.rows; - self.nulls += other.nulls; - self.filtered_out += other.filtered_out; - self.candidates += other.candidates; - self.false_pos += other.false_pos; - self.actual_present += other.actual_present; - } -} - -#[tokio::main] -async fn main() { - let args = CliArgs::parse(); - - let mut config = SessionConfig::default().with_batch_size(8192 * 2); - let options = config.options_mut(); - options.execution.parquet.schema_force_view_types = false; - - let ctx = SessionContext::new_with_config(config); - ctx.register_parquet("hits", &args.parquet, Default::default()) - .await - .expect("register parquet"); - - println!( - "| Column | Pattern | Gram | Mapping | n | Rows | Nulls | Filtered Out | % | Candidates | % | False Pos | % | Actual Present | % |" - ); - println!( - "|--------|---------|------|---------|---|------|-------|--------------|---|------------|---|-----------|---|----------------|---|" - ); - - let bucket_counts = resolve_bucket_counts(&args); - - for column in &args.columns { - let sql = if let Some(limit) = args.limit { - format!("SELECT \"{column}\" FROM \"hits\" LIMIT {limit}") - } else { - format!("SELECT \"{column}\" FROM \"hits\"") - }; - let df = ctx.sql(&sql).await.expect("create df"); - let batches = df.collect().await.expect("collect"); - - let (histogram, pattern_unique) = if matches!(args.mapping, MappingKind::Optimized) { - ( - sample_histogram_from_batches(&batches, args.sample_values), - unique_bytes(args.pattern.as_bytes()), - ) - } else { - ([0u64; 256], [false; 256]) - }; - - for &bucket_count in &bucket_counts { - let mapping = build_mapping( - bucket_count, - args.mapping, - args.custom_map.as_ref(), - &histogram, - &pattern_unique, - ); - let pattern_fp = Fingerprint::from_bytes(args.pattern.as_bytes(), &mapping); - - let mut stats = Stats::default(); - for batch in &batches { - let array: ArrayRef = batch.column(0).clone(); - stats.add(&eval_array(&array, &args.pattern, &mapping, &pattern_fp)); - } - - let non_null = stats.rows.saturating_sub(stats.nulls); - let filtered_pct = pct(stats.filtered_out, non_null); - let candidates_pct = pct(stats.candidates, non_null); - let false_pos_pct = pct(stats.false_pos, stats.candidates); - let actual_present_pct = pct(stats.actual_present, non_null); - - println!( - "| {column} | {} | One | {} | {bucket_count} | {} | {} | {} | {:.2}% | {} | {:.2}% | {} | {:.2}% | {} | {:.2}% |", - args.pattern, - mapping_label(args.mapping), - stats.rows, - stats.nulls, - stats.filtered_out, - filtered_pct, - stats.candidates, - candidates_pct, - stats.false_pos, - false_pos_pct, - stats.actual_present, - actual_present_pct, - ); - } - } -} - -fn resolve_bucket_counts(args: &CliArgs) -> Vec { - match (args.bucket_start, args.bucket_end) { - (Some(start), Some(end)) => { - assert!(args.bucket_step > 0, "bucket_step must be > 0"); - assert!( - (1..=256).contains(&start) && (1..=256).contains(&end), - "bucket_start and bucket_end must be in 1..=256" - ); - assert!(start <= end, "bucket_start must be <= bucket_end"); - let mut buckets = Vec::new(); - let mut value = start; - while value <= end { - buckets.push(value); - let next = value.saturating_add(args.bucket_step); - if next == value { - break; - } - value = next; - } - buckets - } - (None, None) => args.buckets.clone(), - _ => panic!("bucket_start and bucket_end must be set together"), - } -} - -fn mapping_label(kind: MappingKind) -> &'static str { - match kind { - MappingKind::RoundRobin => "RoundRobin", - MappingKind::ContiguousRange => "ContiguousRange", - MappingKind::CustomTable => "CustomTable", - MappingKind::Optimized => "Optimized", - } -} - -fn pct(numer: usize, denom: usize) -> f64 { - if denom == 0 { - 0.0 - } else { - (numer as f64) * 100.0 / (denom as f64) - } -} - -fn build_mapping( - bucket_count: u16, - kind: MappingKind, - custom_map: Option<&PathBuf>, - histogram: &[u64; 256], - pattern_unique: &[bool; 256], -) -> ByteBucketMap { - match kind { - MappingKind::RoundRobin => ByteBucketMap::round_robin(bucket_count), - MappingKind::ContiguousRange => ByteBucketMap::contiguous_range(bucket_count), - MappingKind::CustomTable => { - let path = custom_map.expect("CustomTable mapping requires --custom-map "); - let table = read_custom_table(path); - ByteBucketMap::from_table(bucket_count, table) - } - MappingKind::Optimized => { - let table = optimized_mapping_table(bucket_count, histogram, pattern_unique); - ByteBucketMap::from_table(bucket_count, table) - } - } -} - -fn unique_bytes(bytes: &[u8]) -> [bool; 256] { - let mut present = [false; 256]; - for &b in bytes { - present[b as usize] = true; - } - present -} - -fn sample_histogram_from_batches(batches: &[RecordBatch], sample_values: usize) -> [u64; 256] { - let mut histogram = [0u64; 256]; - let mut remaining = sample_values; - if remaining == 0 { - return histogram; - } - - for batch in batches { - if remaining == 0 { - break; - } - let array = batch.column(0); - add_histogram_from_array(array, &mut histogram, &mut remaining); - } - histogram -} - -fn add_histogram_from_array( - array: &ArrayRef, - histogram: &mut [u64; 256], - remaining_values: &mut usize, -) { - if *remaining_values == 0 { - return; - } - match array.data_type() { - DataType::Utf8 => add_histogram_from_string_iter( - array.as_string::().iter(), - histogram, - remaining_values, - ), - DataType::LargeUtf8 => add_histogram_from_string_iter( - array.as_string::().iter(), - histogram, - remaining_values, - ), - DataType::Utf8View => add_histogram_from_string_iter( - array.as_string_view().iter(), - histogram, - remaining_values, - ), - DataType::Binary => add_histogram_from_binary_iter( - array.as_binary::().iter(), - histogram, - remaining_values, - ), - DataType::LargeBinary => add_histogram_from_binary_iter( - array.as_binary::().iter(), - histogram, - remaining_values, - ), - DataType::BinaryView => add_histogram_from_binary_iter( - array.as_binary_view().iter(), - histogram, - remaining_values, - ), - other => panic!("unsupported data type for histogram sampling: {other:?}"), - } -} - -fn add_histogram_from_string_iter<'a>( - iter: impl Iterator>, - histogram: &mut [u64; 256], - remaining_values: &mut usize, -) { - for value in iter { - if *remaining_values == 0 { - break; - } - let Some(s) = value else { - continue; - }; - for &b in s.as_bytes() { - histogram[b as usize] += 1; - } - *remaining_values -= 1; - } -} - -fn add_histogram_from_binary_iter<'a>( - iter: impl Iterator>, - histogram: &mut [u64; 256], - remaining_values: &mut usize, -) { - for value in iter { - if *remaining_values == 0 { - break; - } - let Some(bytes) = value else { - continue; - }; - for &b in bytes { - histogram[b as usize] += 1; - } - *remaining_values -= 1; - } -} - -fn optimized_mapping_table( - bucket_count: u16, - histogram: &[u64; 256], - pattern_unique: &[bool; 256], -) -> [u8; 256] { - assert!((1..=256).contains(&bucket_count)); - let bucket_count_usize = bucket_count as usize; - assert!(bucket_count_usize > 0); - - let mut table = [0u8; 256]; - let mut assigned = [false; 256]; - let mut bucket_load = vec![0u64; bucket_count_usize]; - let mut pattern_bucket_used = vec![false; bucket_count_usize]; - let mut cursor = 0usize; - - // Assign distinct pattern bytes to distinct, currently-lowest-mass buckets. - for b in 0u16..=255 { - let b = b as usize; - if !pattern_unique[b] { - continue; - } - let bucket = if pattern_bucket_used.iter().all(|&v| v) { - // More distinct pattern bytes than buckets: fall back to standard min-load assignment. - min_load_bucket_from(&bucket_load, cursor) - } else { - min_load_bucket_excluding(&bucket_load, &pattern_bucket_used, cursor) - }; - table[b] = bucket as u8; - assigned[b] = true; - pattern_bucket_used[bucket] = true; - bucket_load[bucket] += histogram[b]; - cursor = (bucket + 1) % bucket_count_usize; - } - - // Sort remaining bytes by descending frequency and greedily balance bucket loads. - let mut remaining: Vec = (0u16..=255) - .map(|b| b as u8) - .filter(|&b| !assigned[b as usize]) - .collect(); - remaining.sort_unstable_by(|&a, &b| { - let fa = histogram[a as usize]; - let fb = histogram[b as usize]; - fb.cmp(&fa).then_with(|| a.cmp(&b)) - }); - - for b in remaining { - let bucket = min_load_bucket_from(&bucket_load, cursor); - cursor = (bucket + 1) % bucket_count_usize; - let idx = b as usize; - table[idx] = bucket as u8; - bucket_load[bucket] += histogram[idx]; - } - - table -} - -#[inline] -fn min_load_bucket_from(bucket_load: &[u64], start: usize) -> usize { - let n = bucket_load.len(); - debug_assert!(n > 0); - let start = start % n; - let mut min_idx = start; - let mut min_val = bucket_load[start]; - for offset in 1..n { - let i = (start + offset) % n; - let v = bucket_load[i]; - if v < min_val { - min_val = v; - min_idx = i; - } - } - min_idx -} - -#[inline] -fn min_load_bucket_excluding(bucket_load: &[u64], excluded: &[bool], start: usize) -> usize { - debug_assert_eq!(bucket_load.len(), excluded.len()); - let n = bucket_load.len(); - debug_assert!(n > 0); - let start = start % n; - - let mut min_idx = None; - let mut min_val = 0u64; - for offset in 0..n { - let i = (start + offset) % n; - if excluded[i] { - continue; - } - let v = bucket_load[i]; - match min_idx { - None => { - min_idx = Some(i); - min_val = v; - } - Some(_) if v < min_val => { - min_idx = Some(i); - min_val = v; - } - _ => {} - } - } - min_idx.expect("excluded all buckets") -} - -fn read_custom_table(path: &PathBuf) -> [u8; 256] { - let raw = fs::read_to_string(path).expect("read custom map"); - let tokens = raw - .split(|c: char| c.is_whitespace() || c == ',') - .filter(|s| !s.is_empty()); - let mut table = [0u8; 256]; - let mut i = 0usize; - for tok in tokens { - assert!(i < 256, "custom map has more than 256 values: {path:?}"); - let v: u16 = tok - .parse() - .unwrap_or_else(|_| panic!("invalid integer token '{tok}' in {path:?}")); - assert!( - v <= 255, - "custom map value {v} out of range (0..=255): {path:?}" - ); - table[i] = v as u8; - i += 1; - } - assert!( - i == 256, - "custom map must contain exactly 256 values, got {i}: {path:?}" - ); - table -} - -fn eval_array( - array: &ArrayRef, - pattern: &str, - mapping: &ByteBucketMap, - pattern_fp: &Fingerprint, -) -> Stats { - match array.data_type() { - DataType::Utf8 => eval_string_iter( - array.as_string::().iter(), - pattern, - mapping, - pattern_fp, - ), - DataType::LargeUtf8 => eval_string_iter( - array.as_string::().iter(), - pattern, - mapping, - pattern_fp, - ), - DataType::Utf8View => { - eval_string_iter(array.as_string_view().iter(), pattern, mapping, pattern_fp) - } - DataType::Binary => eval_binary_iter( - array.as_binary::().iter(), - pattern, - mapping, - pattern_fp, - ), - DataType::LargeBinary => eval_binary_iter( - array.as_binary::().iter(), - pattern, - mapping, - pattern_fp, - ), - DataType::BinaryView => { - eval_binary_iter(array.as_binary_view().iter(), pattern, mapping, pattern_fp) - } - other => panic!("unsupported data type for string fingerprint study: {other:?}"), - } -} - -fn eval_string_iter<'a>( - iter: impl Iterator>, - pattern: &str, - mapping: &ByteBucketMap, - pattern_fp: &Fingerprint, -) -> Stats { - let mut stats = Stats::default(); - for value in iter { - stats.rows += 1; - let Some(s) = value else { - stats.nulls += 1; - continue; - }; - let fp = Fingerprint::from_bytes(s.as_bytes(), mapping); - if !fp.might_contain(pattern_fp) { - stats.filtered_out += 1; - continue; - } - stats.candidates += 1; - if s.contains(pattern) { - stats.actual_present += 1; - } else { - stats.false_pos += 1; - } - } - stats -} - -fn eval_binary_iter<'a>( - iter: impl Iterator>, - pattern: &str, - mapping: &ByteBucketMap, - pattern_fp: &Fingerprint, -) -> Stats { - let mut stats = Stats::default(); - for value in iter { - stats.rows += 1; - let Some(bytes) = value else { - stats.nulls += 1; - continue; - }; - let fp = Fingerprint::from_bytes(bytes, mapping); - if !fp.might_contain(pattern_fp) { - stats.filtered_out += 1; - continue; - } - stats.candidates += 1; - let present = std::str::from_utf8(bytes) - .map(|s| s.contains(pattern)) - .unwrap_or(false); - if present { - stats.actual_present += 1; - } else { - stats.false_pos += 1; - } - } - stats -} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 834e82e22..b5486a276 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -237,10 +237,12 @@ async fn test_url_prefix_filtering() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); test_runner(sql, &reference, cache_dir.path()).await; } @@ -260,10 +262,12 @@ async fn test_url_selection_and_ordering() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); test_runner(sql, &reference, cache_dir.path()).await; } @@ -283,10 +287,12 @@ async fn test_os_selection() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); test_runner(sql, &reference, cache_dir.path()).await; } @@ -307,10 +313,12 @@ async fn test_referer_filtering() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); test_runner(sql, &reference, cache_dir.path()).await; } @@ -331,10 +339,12 @@ async fn test_single_column_filter_projection() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); test_runner(sql, &reference, cache_dir.path()).await; } @@ -421,14 +431,9 @@ async fn test_provide_schema2() { } } - // FSST breaks equal-gain symbol ties using target-specific HashMap iteration order. - // Canonicalize the known AArch64 totals to x86_64 while leaving unexpected totals visible. - #[cfg(target_arch = "aarch64")] - let snapshot = snapshot - .replace("usage.memory_bytes: 999980", "usage.memory_bytes: 1000915") - .replace("usage.memory_bytes: 1035368", "usage.memory_bytes: 1036303"); - - insta::assert_snapshot!(snapshot); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(snapshot); + }); } #[tokio::test] @@ -447,10 +452,12 @@ async fn test_provide_schema_with_filter() { let reference = values.clone(); - insta::assert_snapshot!(format!( - "plan: \n{}\nvalues: \n{}\nstats:\n{}", - plan, values, stats - )); + insta::with_settings!({ filters => vec![(r"usage\.(memory|disk)_bytes: \d+", "usage.${1}_bytes: [bytes]")] }, { + insta::assert_snapshot!(format!( + "plan: \n{}\nvalues: \n{}\nstats:\n{}", + plan, values, stats + )); + }); let (ctx, _) = LiquidCacheLocalBuilder::new() .with_eviction_policy(Box::new(TranscodeEvict)) diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap index a2b47c338..ffd9e6af2 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap @@ -36,19 +36,19 @@ values: stats: entries.total: 8 entries.after_first_run: 8 -entries.memory.arrow: 0 -entries.memory.liquid: 7 -entries.disk.liquid: 1 +entries.memory.arrow: 3 +entries.memory.liquid: 5 +entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 329550 -usage.disk_bytes: 468740 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 3 get_with_selection: 3 eval_predicate: 4 try_read_liquid_calls: 0 - read_io_count: 5 - write_io_count: 3 + read_io_count: 3 + write_io_count: 0 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap index d6e659e63..1113f7c99 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema2.snap @@ -25,8 +25,8 @@ entries.memory.arrow: 5 entries.memory.liquid: 1 entries.disk.liquid: 2 entries.disk.arrow: 0 -usage.memory_bytes: 1000915 -usage.disk_bytes: 35000 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 4 get_with_selection: 4 @@ -174,8 +174,8 @@ entries.memory.arrow: 5 entries.memory.liquid: 1 entries.disk.liquid: 2 entries.disk.arrow: 0 -usage.memory_bytes: 1000915 -usage.disk_bytes: 35000 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 5 get_with_selection: 5 @@ -213,8 +213,8 @@ entries.memory.arrow: 5 entries.memory.liquid: 3 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 1036303 -usage.disk_bytes: 35000 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 4 get_with_selection: 4 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap index 6ff89dd3e..59dfbed6b 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__provide_schema_with_filter.snap @@ -48,8 +48,8 @@ entries.memory.arrow: 12 entries.memory.liquid: 0 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 574612 -usage.disk_bytes: 0 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 13 get_with_selection: 13 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap index 567317870..6309a155a 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap @@ -39,8 +39,8 @@ entries.memory.arrow: 2 entries.memory.liquid: 6 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 951300 -usage.disk_bytes: 877216 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 2 get_with_selection: 2 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap index 5973251b1..b9e654079 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__single_column_filter_projection.snap @@ -25,8 +25,8 @@ entries.memory.arrow: 4 entries.memory.liquid: 0 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 262528 -usage.disk_bytes: 0 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 1 get_with_selection: 1 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap index 7539741d9..6e881e549 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap @@ -55,12 +55,12 @@ entries.memory.arrow: 2 entries.memory.liquid: 2 entries.disk.liquid: 0 entries.disk.arrow: 0 -usage.memory_bytes: 1018212 -usage.disk_bytes: 139416 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: - get: 4 - get_with_selection: 4 - eval_predicate: 0 + get: 0 + get_with_selection: 0 + eval_predicate: 4 try_read_liquid_calls: 0 read_io_count: 1 write_io_count: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap index 3e480e87a..5dcdb137d 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap @@ -37,18 +37,18 @@ stats: entries.total: 4 entries.after_first_run: 4 entries.memory.arrow: 0 -entries.memory.liquid: 2 -entries.disk.liquid: 2 +entries.memory.liquid: 3 +entries.disk.liquid: 1 entries.disk.arrow: 0 -usage.memory_bytes: 147293 -usage.disk_bytes: 706252 +usage.memory_bytes: [bytes] +usage.disk_bytes: [bytes] RuntimeStatsSnapshot: get: 3 get_with_selection: 3 eval_predicate: 4 try_read_liquid_calls: 0 - read_io_count: 4 - write_io_count: 4 + read_io_count: 3 + write_io_count: 2 disk_evictions: 0 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion/README.md b/src/datafusion/README.md index f2f0a4e39..1b8e45922 100644 --- a/src/datafusion/README.md +++ b/src/datafusion/README.md @@ -1,11 +1,11 @@ # liquid-cache-datafusion -Parquet reader with liquid array caching and optimized data formats. +Parquet reader with a Vortex-backed LiquidCache tier. ## Lineage expression pushdown LiquidCache analyzes physical plans to record how each file column is consumed, -including date extraction, variant paths, predicates, and substring searches. +including date extraction, variant paths, and predicates. These expressions continue to flow through local and distributed execution even though the cache currently retains complete Arrow or Liquid arrays. Keeping the analysis separate preserves the information needed for future representation diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index c27c4ff8b..c08cf9af6 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -123,7 +123,6 @@ impl CachedColumn { let liquid_expr = LiquidExpr::try_new( Arc::clone(predicate.physical_expr()), self.field.data_type(), - self.expression.as_deref(), ); if let Some(liquid_expr) = liquid_expr @@ -170,7 +169,7 @@ impl CachedColumn { &self, expr: Arc, ) -> Option { - LiquidExpr::try_new(expr, self.field.data_type(), self.expression.as_deref()) + LiquidExpr::try_new(expr, self.field.data_type()) } pub(crate) fn liquid_expr_for_predicate( diff --git a/src/datafusion/src/io/mod.rs b/src/datafusion/src/io/mod.rs index f88a72d2d..32b0d60f9 100644 --- a/src/datafusion/src/io/mod.rs +++ b/src/datafusion/src/io/mod.rs @@ -4,13 +4,12 @@ use std::{ }; use ahash::AHashMap; -use liquid_cache::cache::{CacheExpression, EntryID, EntryMetadata, LiquidCompressorStates}; +use liquid_cache::cache::{CacheExpression, EntryID, EntryMetadata}; use crate::cache::{ColumnAccessPath, ParquetArrayID}; #[derive(Debug, Default)] pub(crate) struct ParquetCacheMetadata { - compressor_states: RwLock>>, expression_hints: RwLock>, } @@ -70,21 +69,12 @@ impl EntryMetadata for ParquetCacheMetadata { .get(&column_path) .and_then(ColumnExpressionTracker::majority) } - - fn get_compressor(&self, entry_id: &EntryID) -> Arc { - let column_path = ColumnAccessPath::from(ParquetArrayID::from(*entry_id)); - let mut states = self.compressor_states.write().unwrap(); - states - .entry(column_path) - .or_insert_with(|| Arc::new(LiquidCompressorStates::new())) - .clone() - } } #[cfg(test)] mod tests { use super::*; - use liquid_cache::liquid_array::Date32Field; + use liquid_cache::cache::Date32Field; fn entry(file: u64, rg: u64, col: u64) -> EntryID { let id = ParquetArrayID::new(file, rg, col, crate::cache::BatchID::from_raw(0)); diff --git a/src/datafusion/src/optimizers/lineage.rs b/src/datafusion/src/optimizers/lineage.rs index 2413c6f9c..97479a0c2 100644 --- a/src/datafusion/src/optimizers/lineage.rs +++ b/src/datafusion/src/optimizers/lineage.rs @@ -38,7 +38,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::PhysicalExpr; use datafusion::physical_plan::aggregates::AggregateExec; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::expressions::{CastExpr, Column, LikeExpr, Literal}; +use datafusion::physical_plan::expressions::{CastExpr, Column, Literal}; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::joins::HashJoinExec; use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; @@ -46,8 +46,7 @@ use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use liquid_cache::cache::CacheExpression; -use liquid_cache::liquid_array::Date32Field; +use liquid_cache::cache::{CacheExpression, Date32Field}; use crate::cache::ColumnLineages; @@ -69,7 +68,6 @@ enum Op { path: String, data_type: Option, }, - Substring, /// Any other consumption (arithmetic, comparison, cast, unknown op, …). Other, } @@ -460,20 +458,6 @@ fn lineage_for_expr(expr: &std::sync::Arc, input: &LineageMap) return propagate_other(expr, input); } - if let Some(like) = expr.downcast_ref::() { - if !like.case_insensitive() - && let Some(pattern) = literal_utf8(like.pattern()) - && is_substring_pattern(pattern.as_bytes()) - { - let mut usages = lineage_for_expr(like.expr(), input); - for usage in &mut usages { - usage.ops.push(Op::Substring); - } - return usages; - } - return propagate_other(expr, input); - } - if let Some(cast) = expr.downcast_ref::() { let mut usages = lineage_for_expr(cast.expr(), input); for usage in &mut usages { @@ -514,9 +498,6 @@ fn derive_hint(data_type: &DataType, usages: &[Vec]) -> Option]) -> Option { } } -fn derive_substring(data_type: &DataType, usages: &[Vec]) -> Option { - if !is_string_type(data_type) { - return None; - } - let mut saw_substring = false; - for chain in usages { - if chain.iter().any(|op| matches!(op, Op::Substring)) { - saw_substring = true; - continue; - } - if !chain.is_empty() { - return None; - } - } - saw_substring.then(CacheExpression::substring_search) -} - fn literal_utf8(expr: &std::sync::Arc) -> Option { let literal = expr.downcast_ref::()?; match literal.value() { @@ -629,28 +593,6 @@ fn literal_date_field(expr: &std::sync::Arc) -> Option bool { - if pattern.len() < 2 { - return false; - } - if pattern[0] != b'%' || pattern[pattern.len() - 1] != b'%' { - return false; - } - let inner = &pattern[1..pattern.len() - 1]; - if inner.is_empty() { - return false; - } - !inner.iter().any(|b| *b == b'%' || *b == b'_') -} - -fn is_string_type(data_type: &DataType) -> bool { - match data_type { - DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => true, - DataType::Dictionary(_, value_type) => is_string_type(value_type.as_ref()), - _ => false, - } -} - fn is_date_part_type(data_type: &DataType) -> bool { matches!(data_type, DataType::Date32 | DataType::Timestamp(_, _)) } @@ -809,24 +751,4 @@ mod tests { // needs the exact date. assert_eq!(hints.get("date"), None); } - - #[tokio::test] - async fn substring_search_in_filter() { - let hints = hints_for("SELECT date FROM t WHERE url LIKE '%example%'").await; - assert_eq!( - hints.get("url").map(|e| e.as_ref()), - Some(&CacheExpression::substring_search()) - ); - } - - #[tokio::test] - async fn anchored_like_is_not_substring() { - let hints = hints_for("SELECT date FROM t WHERE url LIKE 'https://%'").await; - // A prefix LIKE is not a substring search; no substring hint. - assert!( - hints - .get("url") - .is_none_or(|e| !matches!(e.as_ref(), CacheExpression::SubstringSearch)) - ); - } } From 7c4933fbedd872e2f38296d9cef0f4f31fc98e65 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Fri, 4 Sep 2026 13:51:14 -0400 Subject: [PATCH 10/24] update readme (#518) --- README.md | 54 ++++++++++++++++++++---------------------------------- 1 file changed, 20 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index be69b7874..1cf6210b2 100644 --- a/README.md +++ b/README.md @@ -16,19 +16,16 @@ [![TPC-DS](https://img.shields.io/badge/TPC--DS-passing-brightgreen)](https://github.com/XiangpengHao/liquid-cache/actions/workflows/ci.yml) -LiquidCache understands both your **data** and your **query**. -- It transcodes storage **data** into an optimized, cache-only format, so you can keep using your favorite formats without worrying about performance. -- It keeps the data that matters in memory and uses modern SSDs efficiently. For example, if your **query** groups by `year`, LiquidCache stores only the year in memory and keeps the full timestamp on disk. +LiquidCache understands both **data** and **queries**. -LiquidCache is a research project [funded](https://xiangpeng.systems/fund/) by [InfluxData](https://www.influxdata.com/), [SpiralDB](https://spiraldb.com/), and [Bauplan](https://www.bauplanlabs.com). - -You may want to consider [Foyer](https://github.com/foyer-rs/foyer) if you're looking for a black-box cache: easier to setup, but not as "smart" as LiquidCache. +- It caches data in an optimized, cache-only **data format**, so you can keep using existing storage formats without sacrificing performance. +- It keeps **query-relevant** data in memory and efficiently spills the rest to SSDs. For example, if a query groups by `year`, LiquidCache stores only the `year` in memory and keeps the full `timestamp` on disk. ## Quick start This quick start uses the core cache API from `src/core`. -Add these dependencies to your project: `liquid-cache`, `arrow`, and `datafusion`. -The example below shows insert, get, get with selection, and get with predicate pushdown. +Add `liquid-cache`, `arrow`, and `datafusion` to your project dependencies. +The example below demonstrates insertion, retrieval, selection pushdown, and predicate pushdown. ```rust use arrow::array::{BooleanArray, UInt64Array}; @@ -91,18 +88,18 @@ tokio_test::block_on(async { ## Performance troubleshooting -### LiquidCache uses DIRECT I/O +### LiquidCache uses direct I/O By default, LiquidCache bypasses the OS page cache using [O_DIRECT](https://man7.org/linux/man-pages/man2/open.2.html#:~:text=O_DIRECT) on Linux and `F_NOCACHE` on macOS. This avoids double-caching and bounds memory usage. -This also means LiquidCache can *appear slower* than other caches when most data fits in OS page cache, which is common in dev environments but unrealistic in production. - +This also means LiquidCache can *appear slower* than other caches when most of the data fits in the OS page cache, not a realistic scenario in production. -### Use LiquidCache with DataFusion +### Using LiquidCache with DataFusion LiquidCache requires a few non-default DataFusion configurations: -**ListingTable:** +**With `ListingTable`:** + ```rust let (ctx, _) = LiquidCacheLocalBuilder::new().build(config).await?; @@ -132,45 +129,34 @@ let (ctx, _) = LiquidCacheLocalBuilder::new() ### x86-64 optimization -LiquidCache is optimized for x86-64 with specific [instructions](https://github.com/XiangpengHao/liquid-cache/blob/f8d5b77829fa7996a56c031eb25503f7b0b0428d/src/liquid_parquet/src/utils.rs#L229-L327). On ARM (e.g., Apple Silicon), fallback implementations are used. Contributions are welcome. - +LiquidCache includes specific [x86-64 optimizations](https://github.com/XiangpengHao/liquid-cache/blob/f8d5b77829fa7996a56c031eb25503f7b0b0428d/src/liquid_parquet/src/utils.rs#L229-L327). On ARM platforms such as Apple Silicon, it uses fallback implementations. Contributions are welcome. ## Development -See [dev/README.md](./dev/README.md) +See [dev/README.md](./dev/README.md). ## Benchmark -See [benchmark/README.md](./benchmark/README.md) +See [benchmark/README.md](./benchmark/README.md). ## FAQ -#### Can I use LiquidCache in production today? +#### Is LiquidCache production-ready? -Not yet. Production readiness is our goal, but we are still implementing features and polishing the system. -LiquidCache began as a research project exploring new approaches to cost-effective caching. Like most research projects, it takes time to mature—we welcome your help. +Almost. LiquidCache began as a research project exploring new approaches to cost-effective caching. Like most research projects, it needs time to mature, and we welcome your help. -#### How does LiquidCache work? +#### How can I contribute? -See our [paper](/dev/doc/liquid-cache-vldb.pdf) for details. We are also working on a technical blog to introduce LiquidCache in a more accessible way. - -#### How can I get involved? - -We are always looking for contributors. Feedback and improvements are welcome—explore the issue list and contribute where you can. -If you want to get involved in the research side, [reach out](https://xiangpeng.systems/work-with-me/). +- We use LLMs to help write code. Please have an LLM review your code before submitting it. +- Be accountable for the code you propose, and help maintainers be accountable for the code they merge. +- A PR should take no more than 10 minutes to review, as measured by cognitive burden rather than lines of code. #### Who is behind LiquidCache? -LiquidCache is a research project funded by: -- [SpiralDB](https://spiraldb.com/) -- [InfluxData](https://www.influxdata.com/) -- [Bauplan](https://www.bauplanlabs.com) -- Taxpayers of the state of Wisconsin and the federal government. +LiquidCache grew out of [Xiangpeng Hao's PhD dissertation](https://typst.app/project/rWaXOvnkXMDr0nCjkCgbpp), which was generously supported by [SpiralDB](https://spiraldb.com/), [InfluxData](https://www.influxdata.com/), [Bauplan](https://www.bauplanlabs.com), and taxpayer funding from the state of Wisconsin and the federal government. LiquidCache is and will remain open source and free to use. -Your support for science is greatly appreciated! - ## License [Apache License 2.0](./LICENSE) From 328e17643bea0e44548e908c22f9601559dfc6f1 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Fri, 4 Sep 2026 15:29:19 -0400 Subject: [PATCH 11/24] clean metadata cache (#519) Latest data fusion already has metadata cache, so we don't need to do it again. --- Cargo.lock | 1 - src/datafusion/Cargo.toml | 1 - src/datafusion/src/cache/stats.rs | 10 +- src/datafusion/src/reader/plantime/mod.rs | 5 +- .../src/reader/plantime/morselizer.rs | 197 +++++++++++------- src/datafusion/src/reader/plantime/source.rs | 183 ++-------------- .../src/reader/runtime/liquid_cache_reader.rs | 46 ++-- src/datafusion/src/reader/runtime/morsel.rs | 6 +- 8 files changed, 165 insertions(+), 284 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8df28fc69..d3bb44a76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4481,7 +4481,6 @@ dependencies = [ "ahash", "arrow", "arrow-schema", - "bytes", "datafusion", "datafusion-datasource", "divan", diff --git a/src/datafusion/Cargo.toml b/src/datafusion/Cargo.toml index 4fc534077..53e89258f 100644 --- a/src/datafusion/Cargo.toml +++ b/src/datafusion/Cargo.toml @@ -16,7 +16,6 @@ datafusion-datasource = { workspace = true } futures = { workspace = true } tokio = { workspace = true } ahash = { workspace = true } -bytes = { workspace = true } log = { workspace = true } object_store = { workspace = true, features = ["http"] } liquid-cache-common = { workspace = true } diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index ac4b6e734..d1947fc64 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -160,7 +160,7 @@ impl LiquidCacheParquet { #[cfg(test)] mod tests { - use std::io::Read; + use std::fs::File; use crate::cache::{ParquetFileIdentity, id::BatchID}; @@ -169,7 +169,6 @@ mod tests { array::{Array, AsArray}, datatypes::UInt64Type, }; - use bytes::Bytes; use liquid_cache::{ cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, @@ -235,14 +234,11 @@ mod tests { } } - let mut tmp_file = NamedTempFile::new()?; + let tmp_file = NamedTempFile::new()?; cache.write_stats(tmp_file.path())?; // Read and verify stats - let mut bytes = Vec::new(); - tmp_file.read_to_end(&mut bytes)?; - let bytes = Bytes::from(bytes); - let reader = ParquetRecordBatchReader::try_new(bytes, 8192)?; + let reader = ParquetRecordBatchReader::try_new(File::open(tmp_file.path())?, 8192)?; let batch = reader.into_iter().next().unwrap()?; assert_eq!(batch.num_rows(), num_rows); diff --git a/src/datafusion/src/reader/plantime/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index ddc34eadc..e9d22f934 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -1,11 +1,8 @@ -#[cfg(test)] -pub(crate) use source::CachedMetaReaderFactory; pub use source::LiquidParquetSource; -pub(crate) use source::ParquetMetadataCacheReader; mod morselizer; mod row_filter; mod source; -pub(crate) use morselizer::{LiquidFileMetrics, LiquidMorselizer}; +pub(crate) use morselizer::{LiquidFileMetrics, LiquidFileReaderFactory, LiquidMorselizer}; pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 6334d39f2..a8ff2b80b 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -6,7 +6,7 @@ use datafusion::{ datasource::{ listing::{FileRange, PartitionedFile}, physical_plan::{ - ParquetFileMetrics, + ParquetFileMetrics, ParquetFileReaderFactory, parquet::{ BloomFilterStatistics, PagePruningAccessPlanFilter, ParquetAccessPlan, RowGroupAccessPlanFilter, @@ -15,6 +15,7 @@ use datafusion::{ table_schema::TableSchema, }, error::Result, + execution::object_store::ObjectStoreUrl, physical_expr::{ DynamicFilterTracking, PhysicalExpr, PhysicalExprSimplifier, projection::ProjectionExprs, utils::reassign_expr_columns, @@ -32,12 +33,13 @@ use parquet::{ arrow::{ ParquetRecordBatchStreamBuilder, ProjectionMask, arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection}, + async_reader::AsyncFileReader, parquet_column, }, + errors::ParquetError, file::metadata::PageIndexPolicy, }; -use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; use crate::{ cache::{ BatchID, ColumnLineages, InsertArrowArrayError, LiquidCacheParquetRef, ParquetFileIdentity, @@ -62,7 +64,8 @@ pub(crate) struct LiquidMorselizer { pub(crate) predicate: Option>, pub(crate) table_schema: TableSchema, pub(crate) metrics: ExecutionPlanMetricsSet, - pub(crate) parquet_file_reader_factory: Arc, + pub(crate) parquet_file_reader_factory: Arc, + pub(crate) object_store_url: ObjectStoreUrl, pub(crate) reorder_filters: bool, pub(crate) liquid_cache: LiquidCacheParquetRef, pub(crate) expr_adapter_factory: Arc, @@ -71,6 +74,32 @@ pub(crate) struct LiquidMorselizer { pub(crate) prefetch: bool, } +type ParquetInput = Box; + +#[derive(Clone)] +pub(crate) struct LiquidFileReaderFactory { + pub(crate) factory: Arc, + pub(crate) partition_index: usize, + pub(crate) partitioned_file: PartitionedFile, + pub(crate) metadata_size_hint: Option, + pub(crate) metrics: ExecutionPlanMetricsSet, +} + +impl LiquidFileReaderFactory { + pub(crate) fn create(&self) -> parquet::errors::Result { + let reader = self + .factory + .create_reader( + self.partition_index, + self.partitioned_file.clone(), + self.metadata_size_hint, + &self.metrics, + ) + .map_err(|error| ParquetError::External(Box::new(error)))?; + Ok(reader) + } +} + impl fmt::Debug for LiquidMorselizer { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("LiquidMorselizer") @@ -88,15 +117,16 @@ impl Morselizer for LiquidMorselizer { let metrics = LiquidFileMetrics::new(self.partition_index, &file_name, &self.metrics); let metadata_size_hint = partitioned_file.metadata_size_hint; let file_identity = ParquetFileIdentity::new( - self.parquet_file_reader_factory.object_store_url().clone(), + self.object_store_url.clone(), partitioned_file.object_meta.location.to_string(), ); - let reader = self.parquet_file_reader_factory.create_liquid_reader( - self.partition_index, - partitioned_file.clone(), + let reader_factory = Arc::new(LiquidFileReaderFactory { + factory: Arc::clone(&self.parquet_file_reader_factory), + partition_index: self.partition_index, + partitioned_file: partitioned_file.clone(), metadata_size_hint, - &self.metrics, - ); + metrics: self.metrics.clone(), + }); let logical_file_schema = Arc::clone(self.table_schema.file_schema()); let output_schema = Arc::new( @@ -151,7 +181,7 @@ impl Morselizer for LiquidMorselizer { file_name, metrics, file_pruner, - reader, + reader_factory, batch_size: self.batch_size, logical_file_schema, output_schema, @@ -201,7 +231,7 @@ struct PreparedLiquidOpen { file_name: String, metrics: LiquidFileMetrics, file_pruner: Option, - reader: ParquetMetadataCacheReader, + reader_factory: Arc, batch_size: usize, logical_file_schema: SchemaRef, output_schema: SchemaRef, @@ -232,7 +262,7 @@ struct RowGroupPlanningContext { reader_metadata: ArrowReaderMetadata, physical_file_schema: SchemaRef, cache_full_schema: SchemaRef, - builder: ParquetRecordBatchStreamBuilder, + builder: ParquetRecordBatchStreamBuilder, projection_mask: ProjectionMask, row_filter: Option, pruning_predicate: Option>, @@ -300,9 +330,9 @@ impl LiquidOpenState { let metadata_load_time = prepared.metrics.file_metrics.metadata_load_time.clone(); let mut timer = metadata_load_time.timer(); + let mut reader = prepared.reader_factory.create()?; let reader_metadata = - ArrowReaderMetadata::load_async(&mut prepared.reader, options.clone()) - .await?; + ArrowReaderMetadata::load_async(&mut reader, options.clone()).await?; timer.stop(); Ok(MetadataLoadedLiquidOpen { prepared, @@ -353,11 +383,6 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • 1, - "meta data must be cached already" - ); - let rewriter = loaded.prepared.expr_adapter_factory.create( Arc::clone(&loaded.prepared.logical_file_schema), Arc::clone(&physical_file_schema), @@ -381,7 +406,7 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result, + builder: &mut ParquetRecordBatchStreamBuilder, predicate: &PruningPredicate, file_metrics: &ParquetFileMetrics, row_groups: &RowGroupAccessPlanFilter, @@ -921,9 +946,12 @@ mod tests { common::ScalarValue, datasource::{ listing::PartitionedFile, - physical_plan::{FileScanConfigBuilder, FileSource, ParquetSource}, + physical_plan::{ + FileScanConfigBuilder, FileSource, ParquetSource, + parquet::{CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory}, + }, }, - execution::object_store::ObjectStoreUrl, + execution::{object_store::ObjectStoreUrl, runtime_env::RuntimeEnv}, logical_expr::Operator, physical_expr::{ PhysicalExpr, @@ -938,8 +966,8 @@ mod tests { cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, }; - use object_store::local::LocalFileSystem; - use parquet::arrow::{ArrowWriter, async_reader::AsyncFileReader}; + use object_store::{ObjectStore, local::LocalFileSystem, path::Path}; + use parquet::arrow::ArrowWriter; use crate::{ cache::{BatchID, CachedFileRef, CachedRowGroupRef, LiquidCacheParquet}, @@ -1094,6 +1122,7 @@ mod tests { ) .await; let metrics = ExecutionPlanMetricsSet::new(); + let object_store_url = ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(); let morselizer = LiquidMorselizer { partition_index: 0, projection: ProjectionExprs::from_indices(&options.projection_columns, schema.as_ref()), @@ -1101,10 +1130,10 @@ mod tests { predicate: options.predicate, table_schema: TableSchema::from(Arc::clone(&schema)), metrics: metrics.clone(), - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new( object_store, - ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(), )), + object_store_url: object_store_url.clone(), reorder_filters: false, liquid_cache: cache.clone(), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), @@ -1113,10 +1142,7 @@ mod tests { prefetch: true, }; let cached_file = cache.register_or_get_file( - ParquetFileIdentity::new( - ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(), - file_name, - ), + ParquetFileIdentity::new(object_store_url, file_name), schema, ); TestFilePlanner { @@ -1156,6 +1182,59 @@ mod tests { } } + async fn metadata_cache_hits(cache_limit: usize) -> (Option, Option) { + let schema = schema(); + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("data.parquet"); + write_two_row_group_file(&parquet_path, Arc::clone(&schema)); + let file = PartitionedFile::new( + "data.parquet", + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let object_store: Arc = + Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let metadata_cache = RuntimeEnv::default() + .cache_manager + .get_file_metadata_cache(); + metadata_cache.update_cache_limit(cache_limit); + let parquet_file_reader_factory = Arc::new(CachedParquetFileReaderFactory::new( + Arc::clone(&object_store), + Arc::clone(&metadata_cache), + )); + let parquet_source = ParquetSource::new(Arc::clone(&schema)) + .with_parquet_file_reader_factory(parquet_file_reader_factory); + let liquid_cache = create_test_cache(tmp_dir.path(), usize::MAX, usize::MAX).await; + let source = LiquidParquetSource::from_parquet_source(parquet_source, liquid_cache); + let base_config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()), + ) + .with_file(file.clone()) + .build(); + let morselizer = source + .create_morselizer(object_store, &base_config, 0) + .unwrap(); + let path = Path::from("data.parquet"); + + advance_to_row_group_chain(morselizer.plan_file(file.clone()).unwrap()).await; + let first_hits = metadata_cache + .list_entries() + .get(&path) + .map(|entry| entry.hits); + advance_to_row_group_chain(morselizer.plan_file(file).unwrap()).await; + let second_hits = metadata_cache + .list_entries() + .get(&path) + .map(|entry| entry.hits); + (first_hits, second_hits) + } + + #[tokio::test] + async fn datafusion_metadata_cache_ablation() { + assert_eq!(metadata_cache_hits(usize::MAX).await, (Some(0), Some(1))); + assert_eq!(metadata_cache_hits(0).await, (None, None)); + } + fn gt_expr(column_name: &str, column_index: usize, literal: i32) -> Arc { Arc::new(BinaryExpr::new( Arc::new(Column::new(column_name, column_index)), @@ -1172,44 +1251,6 @@ mod tests { )) } - #[tokio::test] - async fn metadata_cache_is_scoped_to_object_store() { - let schema = schema(); - let dir_a = tempfile::tempdir().unwrap(); - let dir_b = tempfile::tempdir().unwrap(); - let path_a = dir_a.path().join("data.parquet"); - let path_b = dir_b.path().join("data.parquet"); - write_single_row_group_file(&path_a, schema.clone(), vec![1]); - write_single_row_group_file(&path_b, schema, vec![1, 2]); - let metrics = ExecutionPlanMetricsSet::new(); - let mut reader_a = CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_a.path()).unwrap()), - ObjectStoreUrl::parse("store-a:///").unwrap(), - ) - .create_liquid_reader( - 0, - PartitionedFile::new("data.parquet", std::fs::metadata(path_a).unwrap().len()), - None, - &metrics, - ); - let mut reader_b = CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_b.path()).unwrap()), - ObjectStoreUrl::parse("store-b:///").unwrap(), - ) - .create_liquid_reader( - 0, - PartitionedFile::new("data.parquet", std::fs::metadata(path_b).unwrap().len()), - None, - &metrics, - ); - - let metadata_a = reader_a.get_metadata(None).await.unwrap(); - let metadata_b = reader_b.get_metadata(None).await.unwrap(); - - assert_eq!(metadata_a.file_metadata().num_rows(), 1); - assert_eq!(metadata_b.file_metadata().num_rows(), 2); - } - #[tokio::test] async fn data_cache_is_scoped_to_object_store() { let schema = schema(); @@ -1231,10 +1272,10 @@ mod tests { predicate: None, table_schema: TableSchema::from(Arc::clone(&schema)), metrics: metrics.clone(), - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_a.path()).unwrap()), - ObjectStoreUrl::parse("data-cache-a:///").unwrap(), - )), + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::new( + LocalFileSystem::new_with_prefix(dir_a.path()).unwrap(), + ))), + object_store_url: ObjectStoreUrl::parse("data-cache-a:///").unwrap(), reorder_filters: false, liquid_cache: Arc::clone(&cache), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), @@ -1249,10 +1290,10 @@ mod tests { predicate: None, table_schema: TableSchema::from(Arc::clone(&schema)), metrics, - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_b.path()).unwrap()), - ObjectStoreUrl::parse("data-cache-b:///").unwrap(), - )), + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::new( + LocalFileSystem::new_with_prefix(dir_b.path()).unwrap(), + ))), + object_store_url: ObjectStoreUrl::parse("data-cache-b:///").unwrap(), reorder_filters: false, liquid_cache: cache, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 5f0c764e1..69e0415b5 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -1,20 +1,16 @@ use super::LiquidMorselizer; use crate::cache::{ColumnLineages, LiquidCacheParquetRef}; -use ahash::{HashMap, HashMapExt}; -use bytes::Bytes; use datafusion::{ common::{internal_err, tree_node::TreeNodeRecursion}, config::{ConfigOptions, TableParquetOptions}, datasource::{ - listing::PartitionedFile, physical_plan::{ - FileScanConfig, FileSource, ParquetFileMetrics, ParquetFileReaderFactory, - ParquetSource, parquet::can_expr_be_pushed_down_with_schemas, + FileScanConfig, FileSource, ParquetFileReaderFactory, ParquetSource, + parquet::{DefaultParquetFileReaderFactory, can_expr_be_pushed_down_with_schemas}, }, table_schema::TableSchema, }, error::Result, - execution::object_store::ObjectStoreUrl, physical_expr::projection::ProjectionExprs, physical_expr::utils::conjunction, physical_expr_adapter::DefaultPhysicalExprAdapterFactory, @@ -25,168 +21,11 @@ use datafusion::{ }, }; use datafusion_datasource::morsel::Morselizer; -use futures::{FutureExt, future::BoxFuture}; -use object_store::{ObjectStore, ObjectStoreExt, path::Path}; -use parquet::{ - arrow::{arrow_reader::ArrowReaderOptions, async_reader::AsyncFileReader}, - errors::ParquetError, - file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}, -}; +use object_store::ObjectStore; use std::{ fmt::{self, Formatter}, - ops::Range, - sync::{Arc, LazyLock}, + sync::Arc, }; -use tokio::sync::RwLock; - -static META_CACHE: LazyLock = LazyLock::new(MetadataCache::new); - -#[derive(Debug)] -pub(crate) struct CachedMetaReaderFactory { - store: Arc, - store_url: ObjectStoreUrl, -} - -impl CachedMetaReaderFactory { - pub(crate) fn new(store: Arc, store_url: ObjectStoreUrl) -> Self { - Self { store, store_url } - } - - pub(crate) fn object_store_url(&self) -> &ObjectStoreUrl { - &self.store_url - } - - pub(crate) fn create_liquid_reader( - &self, - partition_index: usize, - partitioned_file: PartitionedFile, - metadata_size_hint: Option, - metrics: &ExecutionPlanMetricsSet, - ) -> ParquetMetadataCacheReader { - let path = partitioned_file.object_meta.location.clone(); - - ParquetMetadataCacheReader { - file_metrics: ParquetFileMetrics::new(partition_index, path.as_ref(), metrics), - store: Arc::clone(&self.store), - store_url: self.store_url.clone(), - file_size: partitioned_file.object_meta.size, - metadata_size_hint, - path, - } - } -} - -impl ParquetFileReaderFactory for CachedMetaReaderFactory { - fn create_reader( - &self, - partition_index: usize, - partitioned_file: PartitionedFile, - metadata_size_hint: Option, - metrics: &ExecutionPlanMetricsSet, - ) -> Result> { - let reader = self.create_liquid_reader( - partition_index, - partitioned_file, - metadata_size_hint, - metrics, - ); - Ok(Box::new(reader)) - } -} - -struct MetadataCache { - val: RwLock>>, -} - -impl MetadataCache { - fn new() -> Self { - Self { - val: RwLock::new(HashMap::new()), - } - } -} - -#[derive(Clone)] -pub struct ParquetMetadataCacheReader { - file_metrics: ParquetFileMetrics, - store: Arc, - store_url: ObjectStoreUrl, - file_size: u64, - metadata_size_hint: Option, - path: Path, -} - -fn to_parquet_err(error: object_store::Error) -> ParquetError { - ParquetError::External(Box::new(error)) -} - -impl AsyncFileReader for ParquetMetadataCacheReader { - fn get_byte_ranges( - &mut self, - ranges: Vec>, - ) -> BoxFuture<'_, parquet::errors::Result>> { - let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); - self.file_metrics.bytes_scanned.add(total as usize); - async move { - self.store - .get_ranges(&self.path, &ranges) - .await - .map_err(to_parquet_err) - } - .boxed() - } - - fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { - self.file_metrics - .bytes_scanned - .add((range.end - range.start) as usize); - async move { - self.store - .get_range(&self.path, range) - .await - .map_err(to_parquet_err) - } - .boxed() - } - - fn get_metadata( - &mut self, - options: Option<&ArrowReaderOptions>, - ) -> BoxFuture<'_, parquet::errors::Result>> { - let cache_key = (self.store_url.clone(), self.path.clone()); - let options = options.cloned(); - async move { - // First check with read lock - { - let cache = META_CACHE.val.read().await; - if let Some(meta) = cache.get(&cache_key) { - return Ok(meta.clone()); - } - } - - // Upgrade to write lock and double-check - let mut cache = META_CACHE.val.write().await; - match cache.entry(cache_key) { - std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()), - std::collections::hash_map::Entry::Vacant(entry) => { - let file_size = self.file_size; - let meta = ParquetMetaDataReader::new() - .with_arrow_reader_options(options.as_ref()) - .with_prefetch_hint(self.metadata_size_hint) - .load_and_finish(&mut *self, file_size) - .await?; - let mut reader = ParquetMetaDataReader::new_with_metadata(meta.clone()) - .with_page_index_policy(PageIndexPolicy::Optional); - reader.load_page_index(&mut *self).await?; - let meta = Arc::new(reader.finish()?); - entry.insert(meta.clone()); - Ok(meta) - } - } - } - .boxed() - } -} /// The data source for LiquidCache #[derive(Clone)] @@ -198,6 +37,7 @@ pub struct LiquidParquetSource { batch_size: Option, projection: ProjectionExprs, table_schema: TableSchema, + parquet_file_reader_factory: Option>, span: Option>, lineages: Arc, prefetch: bool, @@ -253,6 +93,7 @@ impl LiquidParquetSource { /// Create a new LiquidParquetSource from a ParquetSource pub fn from_parquet_source(source: ParquetSource, liquid_cache: LiquidCacheParquetRef) -> Self { let predicate = source.filter(); + let parquet_file_reader_factory = source.parquet_file_reader_factory().cloned(); let table_schema = source.table_schema().clone(); let projection = source.projection().cloned().unwrap_or_else(|| { @@ -269,6 +110,7 @@ impl LiquidParquetSource { liquid_cache, projection, metrics: source.metrics().clone(), + parquet_file_reader_factory, predicate: None, span: None, lineages: Arc::default(), @@ -311,10 +153,10 @@ impl FileSource for LiquidParquetSource { .clone() .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); - let reader_factory = Arc::new(CachedMetaReaderFactory::new( - object_store, - base_config.object_store_url.clone(), - )); + let parquet_file_reader_factory = self + .parquet_file_reader_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultParquetFileReaderFactory::new(object_store))); let execution_span = self .span @@ -330,7 +172,8 @@ impl FileSource for LiquidParquetSource { table_schema: self.table_schema.clone(), metrics: self.metrics.clone(), liquid_cache: self.liquid_cache.clone(), - parquet_file_reader_factory: reader_factory, + parquet_file_reader_factory, + object_store_url: base_config.object_store_url.clone(), reorder_filters: self.reorder_filters(), expr_adapter_factory, span: execution_span.map(Arc::new), diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 37babe4ad..84a03b55a 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -17,7 +17,7 @@ use parquet::errors::ParquetError; use parquet::file::metadata::ParquetMetaData; use crate::cache::{BatchID, CachedRowGroupRef, InsertArrowArrayError}; -use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; +use crate::reader::plantime::{LiquidFileReaderFactory, LiquidRowFilter}; use crate::reader::runtime::utils::take_next_batch; use crate::utils::{boolean_buffer_and_then, row_selector_to_boolean_buffer}; @@ -71,7 +71,7 @@ pub(crate) struct LiquidCacheReaderConfig { pub(crate) struct ParquetFallbackConfig { pub(crate) row_group_idx: usize, pub(crate) metadata: Arc, - pub(crate) input: ParquetMetadataCacheReader, + pub(crate) reader_factory: Arc, pub(crate) cache_projection: ProjectionMask, pub(crate) cache_column_ids: Vec, pub(crate) cache_batch_size: usize, @@ -81,7 +81,7 @@ pub(crate) struct ParquetFallbackConfig { pub(crate) struct ParquetFallback { row_group_idx: usize, metadata: Arc, - input: ParquetMetadataCacheReader, + reader_factory: Arc, cache_projection: ProjectionMask, cache_column_ids: Vec, cache_batch_size: usize, @@ -162,7 +162,7 @@ impl ParquetFallback { Self { row_group_idx: config.row_group_idx, metadata: config.metadata, - input: config.input, + reader_factory: config.reader_factory, cache_projection: config.cache_projection, cache_column_ids: config.cache_column_ids, cache_batch_size: config.cache_batch_size, @@ -199,14 +199,16 @@ impl ParquetFallback { let row_selection = build_row_selection_from(batch_id, self.cache_batch_size, self.row_count); - let stream = - ParquetRecordBatchStreamBuilder::new_with_metadata(self.input.clone(), reader_metadata) - .with_projection(self.cache_projection.clone()) - .with_row_groups(vec![self.row_group_idx]) - .with_batch_size(self.cache_batch_size) - .with_row_selection(row_selection) - .build()? - .boxed(); + let stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + self.reader_factory.create()?, + reader_metadata, + ) + .with_projection(self.cache_projection.clone()) + .with_row_groups(vec![self.row_group_idx]) + .with_batch_size(self.cache_batch_size) + .with_row_selection(row_selection) + .build()? + .boxed(); self.stream = Some(stream); self.next_batch_id = batch_id; @@ -496,13 +498,15 @@ mod tests { use super::*; use crate::{ cache::LiquidCacheParquet, - reader::plantime::CachedMetaReaderFactory, + reader::plantime::LiquidFileReaderFactory, reader::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}, }; use arrow::array::{ArrayRef, Int32Array}; use arrow::record_batch::RecordBatch; use arrow_schema::{DataType, Field, Schema, SchemaRef}; - use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::{ + listing::PartitionedFile, physical_plan::parquet::DefaultParquetFileReaderFactory, + }; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::{ logical_expr::Operator, @@ -573,11 +577,13 @@ mod tests { std::fs::metadata(&parquet_path).unwrap().len(), ); let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new( - object_store, - datafusion::execution::object_store::ObjectStoreUrl::parse("test-runtime:///").unwrap(), - ) - .create_liquid_reader(0, partitioned_file, None, &metrics); + let reader_factory = Arc::new(LiquidFileReaderFactory { + factory: Arc::new(DefaultParquetFileReaderFactory::new(object_store)), + partition_index: 0, + partitioned_file, + metadata_size_hint: None, + metrics, + }); let projection = ProjectionMask::roots( reader_metadata.metadata().file_metadata().schema_descr(), [0], @@ -622,7 +628,7 @@ mod tests { fallback: ParquetFallbackConfig { row_group_idx: 0, metadata: Arc::clone(reader_metadata.metadata()), - input, + reader_factory, cache_projection: projection, cache_column_ids: vec![0], cache_batch_size: batch_size, diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index 1c6931ffb..271f53fff 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -15,7 +15,7 @@ use parquet::{ use crate::{ cache::{CachedFileRef, CachedRowGroupRef, LiquidCacheParquetRef, RowGroupSnapshots}, - reader::plantime::{LiquidFileMetrics, LiquidRowFilter, ParquetMetadataCacheReader}, + reader::plantime::{LiquidFileMetrics, LiquidFileReaderFactory, LiquidRowFilter}, }; use super::{ @@ -27,7 +27,7 @@ use super::{ pub(crate) struct LiquidRowGroupPlanner { pub(crate) metadata: Arc, - pub(crate) input: ParquetMetadataCacheReader, + pub(crate) reader_factory: Arc, pub(crate) row_filter: Option, pub(crate) cached_file: CachedFileRef, pub(crate) projection: ProjectionMask, @@ -79,7 +79,7 @@ impl LiquidRowGroupPlanner { ParquetFallbackConfig { row_group_idx, metadata: Arc::clone(&self.metadata), - input: self.input.clone(), + reader_factory: Arc::clone(&self.reader_factory), cache_projection: details.cache_projection.clone(), cache_column_ids: details.cache_column_ids.clone(), cache_batch_size, From 0033b159b716a381dff59bb26358b992a21089cb Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Sat, 19 Sep 2026 14:46:52 -0700 Subject: [PATCH 12/24] fix client and server lineage pushdown (#520) it was only working in local mode --- AGENTS.md | 72 +++-- Cargo.lock | 2 + dev/design/05-compaction-cache.md | 129 ++++++++ src/core/src/cache/core.rs | 4 +- src/datafusion-client/src/lib.rs | 7 +- src/datafusion-local/src/lib.rs | 7 +- src/datafusion-server/Cargo.toml | 2 + src/datafusion-server/src/lib.rs | 7 +- src/datafusion-server/src/plan_codec.rs | 81 +++++ src/datafusion-server/src/tests/lineage.rs | 346 +++++++++++++++++++++ src/datafusion-server/src/tests/mod.rs | 1 + src/datafusion/src/lib.rs | 9 + src/datafusion/src/optimizers/lineage.rs | 7 +- 13 files changed, 644 insertions(+), 30 deletions(-) create mode 100644 dev/design/05-compaction-cache.md create mode 100644 src/datafusion-server/src/plan_codec.rs create mode 100644 src/datafusion-server/src/tests/lineage.rs diff --git a/AGENTS.md b/AGENTS.md index 2e8d8c225..41fb4877b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,25 +1,61 @@ -## Engineering principles +## Write code that requires minimal human attention -1. Think about minimal changes to complete the task. -2. Always cargo check after coding. -3. Consider refactoring when a function needs more than 3 args. -4. Whenever writing a if or edge case condition, (1) think why it is needed here, whether it is a real edge case, (2) think whether we can move the condition upper to the caller, i.e., is this the best place to handle edge case? -5. Never need to consider backward compatibility. +- Fact: every line of code requires human review, so write readable code not just working code. -## Code structure +- Build reusable, trustworthy components that has deep logic but thin interface, so that human only need to review the complexity once (reuse the attention). -- `src/core`, this is the LiquidCache core. -- `src/datafusion`, Parquet and DataFusion integration, this allows datafusion/parquet users to use LiquidCache with minimal effort. -- `src/datafusion-client` and `src/datafusion-server`, Client/Server library, this enables distributed LiquidCache. -- `src/datafusion-local`, this is a in-process LiquidCache, used for local DataFusion instances. +- Follow Rust best practices to save human attention from reviewing for common mistakes. -### Lineage-based cache expression +- Before implementing a new feature, identify similar existing features, and check if it's worth to refactor existing code so it's reusable. -1. The lineage_opt.rs analyze the input query's column usage, and passes it down to LiquidCache. -2. When creating a CachedColumn, it will register the expression to the cache; This info is used to determine the squeeze behavior only. -3. During cache read, a expression is passed in, the cache will decide whether the cached (squeezed) data can satisfy the expression, if so, it will return the cached data directly; otherwise, it will hydrate the data from disk. +## Engineering -## Testing +- Prefer explicit props and function arguments so ownership and data flow remain visible. Use context only when a dependency must cross at least two component layers and the intermediate parent does not use it. Keep state owned by the component that creates it. -1. All test cases are intended for human to read. They are part of the code and documents. -2. Test cases should be short, concise, yet have high coverage. Writing test cases almost requires highest level of thinking. + +- use descriptive, intention-revealing names for variables and functions, don't use generic names like `Built`, `Common`, `View`. + +- do not use free form functions, functions should have a clear owner, e.g., struct, enum; unless it is fully justified stateless function. + Example 1: `fn build_chart(model: &DataModel) -> Built`, instead use a member method like: `fn build_chart(&self) -> Built`. Even better, use a `From` trait to build the chart from the data model. + Example 2: `fn toggle_expand(code: crate::views::codemap::CodeState, key: (u32, u32))` instead use a member method like: `fn toggle_expand(&self, key: (u32, u32))`. + +- try to use minimal visibility for functions and variables, e.g., a pub function under a private mod is essentially private, but it is confusing, prefer to make it private. + +- try to avoid public fields, prefer to use getters and setters, unless the field is truly public. + +- organize the code by features, not by technical buckets, e.g., don't have `api.rs`, `views.rs`, `data.rs`, etc. Instead, do `settings.rs`, `data_panel.rs`, `code_map.rs`, etc. + +- No static variables. + +- Don't use large unconstrained struct constructions, e.g., following code is ugly and error prone. +Maybe there're invariants we want to enforce, but this code allows any fields to be set. +```rust +ghost_nodes.push(CrateInfo { + id: ghost_id.clone(), + name: ev.name.clone(), + version: ev.detail.clone().unwrap_or_default(), + is_member: false, + changed: false, + changed_files: 0, + manifest_changed: false, + affected_dist: None, + dependents: 0, + direct_deps: 0, + external_deps: 0, + ghost: true, + description: None, + license: None, + repository: None, + homepage: None, + documentation: None, + // A removed dependency's manifest is gone with + // it; the name is all we know. + crates_io: false, + rel_path: None, + }); +``` +Instead, we use use a much narrower constructor, with new(arg1, arg2, arg3), and check and enforce the invariants. + +- Typically a function should not have more than 3 parameters (including self if it is a member method). If it does, it is a sign of either too large function body, or a container struct should hold the parameters. + +- A struct should not have more than 7 fields, more than that adds cognitive burden. Use private structs to group related fields. diff --git a/Cargo.lock b/Cargo.lock index d3bb44a76..bf7f34ab5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4558,6 +4558,8 @@ dependencies = [ "liquid-cache", "liquid-cache-common", "liquid-cache-datafusion", + "liquid-cache-datafusion-client", + "liquid-cache-datafusion-local", "log", "object_store", "parquet", diff --git a/dev/design/05-compaction-cache.md b/dev/design/05-compaction-cache.md new file mode 100644 index 000000000..c6cc93d5f --- /dev/null +++ b/dev/design/05-compaction-cache.md @@ -0,0 +1,129 @@ +# File compaction in LiquidCache + +Status: direction agreed; implementation details remain TODO. DuckLake is the first integration. Step 2 focuses on compaction correctness; Step 3 adds cache reuse. + +The cited context describes existing systems and constraints, not additional implementation decisions. + +**Editing instructions:** Use simple, concise technical English. Organize agreed decisions, relevant context, cited research, and open questions under the corresponding step. Do not add undiscussed APIs, architecture, execution flows, algorithms, or implementation plans. Leave those as TODO; do not turn research findings or assumptions into agreed decisions. + +## Step 0: Revisit the LiquidCache client/server implementation + +Bring the LiquidCache client/server implementation to feature parity with local mode. In particular, query lineage must be pushed down to the server. This is a prerequisite for the DuckLake integration and compaction work. + +### Implementation + +The client analyzes lineage before splitting the physical plan into remote fragments. It sends each scan's expressions with the plan registration request. The server decodes these expressions and attaches them when rewriting Parquet scans to LiquidParquet scans. + +Local, client, and server sessions register the same variant functions. Client and server Parquet settings preserve Arrow metadata, and the client no longer forces binary columns to strings. The server restores variant return-field metadata omitted by DataFusion's plan serialization, so nested variant functions work remotely. + +Lineage analysis applies hash-join output projections when mapping columns back to their scans. The cache passes registered lineage to the eviction policy even when an incoming batch cannot fit in memory. + +Flight integration tests compare client/server queries with local mode and plain DataFusion. They cover date extraction, multiple variant paths, nested variant functions, raw binary data, mixed raw and derived usage, joins with overlapping column names, cache reuse, and disk hydration for a later raw query. They also check the expressions received by the eviction policy and cache reads. + +Validation: 200 library tests pass across the core, DataFusion integration, local, client, and server crates; one existing client test is ignored. `cargo check --workspace --all-targets` passes. + +### Current limits + +The current core retains lineage, but `TranscodeEvict` does not use it and there is no squeezed cache representation. The tests verify lineage propagation and ordinary disk hydration; squeeze-specific reuse and fallback cannot be tested until the core supports them. + +Distributed dynamic filter pushdown remains disabled because runtime updates from client-side operators do not reach remote scans. + +## Step 1: DuckLake–LiquidCache catalog provider + +Add a DataFusion catalog provider that integrates DuckLake with LiquidCache. Users register the catalog and query its tables normally, without manually wiring DuckLake and LiquidCache together. This step makes LiquidCache ergonomic for DuckLake users. + +Use [`datafusion-ducklake`](https://github.com/datafusion-contrib/datafusion-ducklake) for this integration. DuckLake is the first integration point; the core compaction feature should not depend on it. Other lakehouse formats could be supported through their own integrations. + +### Catalog context + +DuckLake defines its catalog through SQL tables and transactions. With PostgreSQL, clients connect directly to the database. [Catalog SQL](https://ducklake.select/docs/stable/specification/queries), [catalog backends](https://ducklake.select/docs/stable/duckdb/usage/choosing_a_catalog_database) + +### Integration test setup + +PostgreSQL is primarily a service for testing the catalog integration. In this test setup, it runs alongside LiquidCache and, once added, compaction. Compute runs remotely and uses the DuckLake–LiquidCache catalog provider. Object storage is separate. + +The focus is how LiquidCache compacts files. Hosting PostgreSQL is not a core requirement of the compaction feature. + +### TODO + +Provider API and implementation details. + +## Step 1.1: Direct LiquidParquet planning + +Investigate whether the table providers exposed by our catalog provider can produce LiquidParquet plans directly. This could remove the need for optimizer rules that rewrite Parquet plans into LiquidParquet plans and make the integration simpler. + +This is an open design question. Query lineage must still reach the server as required by Step 0. + +### TODO + +Determine whether direct planning can replace the scan-rewrite rules and how query lineage would be passed down. + +## Step 2: Compaction correctness + +LiquidCache will provide a compaction function, such as `merge_adjacent_files`. Users can call it from their existing maintenance jobs. Compaction runs on the LiquidCache server: the caller invokes the function, and the server performs the work. + +Focus on correct compacted output and accurate reporting of the input and output files. LiquidCache returns which files were compacted and which new files were produced. The client decides how to update its lakehouse catalog. LiquidCache does not update the catalog. + +PostgreSQL from Step 1 supports testing the catalog integration. The core compaction feature is independent of catalog hosting and scheduling. Cache reuse is outside this step's scope. + +### Compaction context + +Compaction exists across lakehouse formats: + +| Format | Operation | +| --- | --- | +| DuckLake | [`merge_adjacent_files`](https://ducklake.select/docs/stable/duckdb/maintenance/merge_adjacent_files) merges files and can sort the output. | +| Iceberg | [`rewrite_data_files`](https://iceberg.apache.org/docs/latest/spark-procedures/#rewrite_data_files) supports bin-packing and sorting. | +| Delta Lake | [`OPTIMIZE`](https://docs.delta.io/optimizations-oss/) combines small files. | + +In DuckLake's existing implementation, the client performs the Parquet rewrite and commits metadata changes; PostgreSQL stores the metadata. Compaction runs when a maintenance operation is invoked. The `auto_compact` option controls table eligibility, not automatic scheduling. [Catalog operations](https://ducklake.select/docs/stable/specification/queries), [compaction documentation](https://ducklake.select/docs/stable/duckdb/maintenance/merge_adjacent_files) + +DuckLake tracks deletes and snapshot visibility separately from file contents. It delays physical file cleanup to protect active reads and retained history. These constraints matter for compaction correctness and later cache reuse. [Delete files](https://ducklake.select/docs/stable/specification/tables/ducklake_delete_file), [transactions](https://ducklake.select/docs/stable/duckdb/advanced_features/transactions), [file cleanup](https://ducklake.select/docs/stable/duckdb/maintenance/cleanup_of_files) + +Related work: [AutoComp, SIGMOD 2025](https://arxiv.org/html/2504.04186v1) studies lakehouse compaction selection and scheduling. + +### TODO + +- Compaction flow. +- Function API and integration interface. +- Implementation details. + +## Step 3: Cache reuse across compaction + +Add cache reuse across compaction after Step 2 establishes correctness. Queries must use LiquidCache to benefit from it. + +### Motivation + +LiquidCache works well when Parquet files stay unchanged. In a lakehouse, files are added, rows are deleted, and files are compacted. Compaction is especially problematic: + +```text +a.parquet + b.parquet → c.parquet +``` + +The data may be unchanged, but its location and organization change. Cached data from `a.parquet` and `b.parquet` cannot currently be reused through the new file identity. The compaction operations in Step 2 can create this problem across lakehouse formats. + +LiquidCache stores column batches, typically 8,192 rows. Its current DataFusion cache key is `(file ID, column ID, row group ID, batch ID)`. [Cache keys](../../src/datafusion/src/cache/id.rs), [batch size](../../src/core/src/cache/builders.rs) + +The original idea is to track lineage so that data moved by compaction can be remapped to existing cached data. How to do this remains open. + +### Research context + +DuckLake row IDs survive both compaction and updates. Equal row IDs therefore do not prove equal values. Row IDs can be derived from a file's starting ID or stored explicitly inside Parquet. Catalog metadata alone does not describe every possible row reordering. [Row lineage](https://ducklake.select/docs/stable/duckdb/advanced_features/row_lineage), [file metadata](https://ducklake.select/docs/stable/specification/tables/ducklake_data_file) + +Merging files can split or combine existing cache batches. Sorting can also change row order. Unchanged logical data therefore does not imply that an output batch matches a cached source batch. [Sorted compaction](https://ducklake.select/docs/stable/duckdb/maintenance/merge_adjacent_files) + +We have not established whether `datafusion-ducklake`'s current APIs expose the information needed for cache reuse. + +### Related work + +| Work | Relevance | +| --- | --- | +| [Lance stable row IDs and remap separation](https://github.com/lance-format/lance/discussions/3694) | Discusses keeping indexes useful after compaction through stable IDs or row-address mappings. The linked discussion is a design proposal. | +| [Iceberg v3 row lineage](https://iceberg.apache.org/spec/#row-lineage) | Separates row identity from the sequence number of its last update. | +| [Delta Lake row tracking](https://docs.delta.io/delta-row-tracking/) | Tracks stable row IDs and row commit versions. | +| [dLSM, 2016](https://arxiv.org/abs/1606.02015) | Studies cache invalidation caused by compaction moving data in LSM trees. | +| [Databricks disk cache](https://docs.databricks.com/aws/en/optimizations/disk-cache) | Provides a baseline for invalidating cached Parquet data after file changes. The cited docs do not establish reuse across compaction. | + +### TODO + +Lineage tracking and the cache-reuse implementation. diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 7c36dd11a..d007d3164 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -291,7 +291,9 @@ impl LiquidCache { ) -> Result { match &batch { batch @ CacheEntry::MemoryArrow(_) => { - let outcome = self.eviction_policy.evict(batch, None); + let outcome = self + .eviction_policy + .evict(batch, self.metadata.lineage(&entry_id).as_deref()); let EvictionOutcome::Replace { entry: new_batch, bytes_to_write, diff --git a/src/datafusion-client/src/lib.rs b/src/datafusion-client/src/lib.rs index 02d8aef9b..4fa7f1561 100644 --- a/src/datafusion-client/src/lib.rs +++ b/src/datafusion-client/src/lib.rs @@ -91,7 +91,8 @@ impl LiquidCacheClientBuilder { .options_mut() .execution .parquet - .binary_as_string = true; + .skip_arrow_metadata = false; + session_config.options_mut().execution.parquet.skip_metadata = false; session_config.options_mut().execution.batch_size = ConfigNonZeroUsize::try_new(8192 * 2)?; // Dynamic filters (e.g. a hash join's runtime build-side filter) are pushed // into scan predicates by DataFusion. In distributed mode those scans are @@ -126,7 +127,9 @@ impl LiquidCacheClientBuilder { self.object_stores.clone(), ))) .build(); - Ok(SessionContext::new_with_state(session_state)) + let ctx = SessionContext::new_with_state(session_state); + liquid_cache_datafusion::register_variant_functions(&ctx); + Ok(ctx) } } diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index bfccbaa13..d8cbbadd5 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -7,7 +7,6 @@ mod tests; use std::path::PathBuf; use std::sync::Arc; -use datafusion::logical_expr::ScalarUDF; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion::{common::config::ConfigNonZeroUsize, error::Result}; use liquid_cache::cache::{AlwaysHydrate, HydrationPolicy, default_max_memory_bytes}; @@ -15,7 +14,7 @@ use liquid_cache::cache::{EvictionPolicy, TranscodeEvict}; use liquid_cache::cache_policies::{CachePolicy, LiquidPolicy}; use liquid_cache_datafusion::optimizers::LocalModeOptimizer; use liquid_cache_datafusion::{ - LiquidCacheParquet, LiquidCacheParquetRef, VariantGetUdf, VariantPretty, VariantToJsonUdf, + LiquidCacheParquet, LiquidCacheParquetRef, register_variant_functions, }; pub use liquid_cache as storage; @@ -207,9 +206,7 @@ impl LiquidCacheLocalBuilder { .build(); let ctx = SessionContext::new_with_state(state); - ctx.register_udf(ScalarUDF::new_from_impl(VariantGetUdf::default())); - ctx.register_udf(ScalarUDF::new_from_impl(VariantPretty::default())); - ctx.register_udf(ScalarUDF::new_from_impl(VariantToJsonUdf::default())); + register_variant_functions(&ctx); Ok((ctx, cache_ref)) } } diff --git a/src/datafusion-server/Cargo.toml b/src/datafusion-server/Cargo.toml index 67549800c..804b48dc3 100644 --- a/src/datafusion-server/Cargo.toml +++ b/src/datafusion-server/Cargo.toml @@ -44,3 +44,5 @@ hdrhistogram = "7.5.4" [dev-dependencies] insta = { workspace = true } parquet = { workspace = true } +liquid-cache-datafusion-client = { workspace = true } +liquid-cache-datafusion-local = { workspace = true } diff --git a/src/datafusion-server/src/lib.rs b/src/datafusion-server/src/lib.rs index fcaf0a772..d1cd5bd4d 100644 --- a/src/datafusion-server/src/lib.rs +++ b/src/datafusion-server/src/lib.rs @@ -33,7 +33,6 @@ use datafusion::{ execution::{SessionStateBuilder, object_store::ObjectStoreUrl}, prelude::{SessionConfig, SessionContext}, }; -use datafusion_proto::bytes::physical_plan_from_bytes; use fastrace::prelude::SpanContext; use futures::{Stream, TryStreamExt}; use liquid_cache::cache::CacheExpression; @@ -52,6 +51,7 @@ mod utils; use utils::FinalStream; mod admin_server; mod errors; +mod plan_codec; pub use admin_server::{models::*, run_admin_server}; pub use errors::{ LiquidCacheErrorExt, LiquidCacheResult, anyhow_to_status, df_error_to_status_with_trace, @@ -167,6 +167,8 @@ impl LiquidCacheService { let mut session_config = SessionConfig::from_env()?; let options_mut = session_config.options_mut(); options_mut.execution.parquet.pushdown_filters = true; + options_mut.execution.parquet.skip_arrow_metadata = false; + options_mut.execution.parquet.skip_metadata = false; options_mut.execution.batch_size = ConfigNonZeroUsize::try_new(8192 * 2)?; { @@ -186,6 +188,7 @@ impl LiquidCacheService { .build(); let ctx = SessionContext::new_with_state(state); + liquid_cache_datafusion::register_variant_functions(&ctx); Ok(ctx) } @@ -250,7 +253,7 @@ impl LiquidCacheService { Ok(Response::new(Box::pin(output))) } LiquidCacheActions::RegisterPlan(cmd) => { - let plan = physical_plan_from_bytes(&cmd.plan, &self.inner.get_ctx().task_ctx())?; + let plan = plan_codec::decode_plan(&cmd.plan, &self.inner.get_ctx().task_ctx())?; let handle = Uuid::from_bytes_ref(cmd.handle.as_ref().try_into()?); let mut lineages = ColumnLineages::default(); for hint in &cmd.lineages { diff --git a/src/datafusion-server/src/plan_codec.rs b/src/datafusion-server/src/plan_codec.rs new file mode 100644 index 000000000..66d95f376 --- /dev/null +++ b/src/datafusion-server/src/plan_codec.rs @@ -0,0 +1,81 @@ +//! Restore variant return fields that DataFusion's scalar-UDF protobuf omits. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion::{ + error::Result, + execution::TaskContext, + physical_expr::{PhysicalExpr, ScalarFunctionExpr}, + physical_plan::ExecutionPlan, +}; +use datafusion_proto::{ + bytes::physical_plan_from_bytes_with_proto_converter, + physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, + }, + protobuf, +}; +use liquid_cache_datafusion::VariantGetUdf; + +pub(crate) fn decode_plan(bytes: &[u8], ctx: &TaskContext) -> Result> { + physical_plan_from_bytes_with_proto_converter( + bytes, + ctx, + &DefaultPhysicalExtensionCodec {}, + &VariantFieldConverter, + ) +} + +struct VariantFieldConverter; + +impl PhysicalProtoConverterExtension for VariantFieldConverter { + fn proto_to_execution_plan( + &self, + proto: &protobuf::PhysicalPlanNode, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> Result> { + self.default_proto_to_execution_plan(proto, ctx) + } + + fn proto_to_physical_expr( + &self, + proto: &protobuf::PhysicalExprNode, + input_schema: &Schema, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> Result> { + let expr = self.default_proto_to_physical_expr(proto, input_schema, ctx)?; + let Some(function) = ScalarFunctionExpr::try_downcast_func::(expr.as_ref()) + else { + return Ok(expr); + }; + // The wire format carries the return data type, but not its extension metadata. + // Re-infer the field before a parent expression uses it, including nested UDFs. + Ok(Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(function.fun().clone()), + function.args().to_vec(), + input_schema, + ctx.task_ctx().session_config().options().clone(), + )? + .with_nullable(function.nullable()), + )) + } + + fn execution_plan_to_proto( + &self, + plan: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> Result { + DefaultPhysicalProtoConverter {}.execution_plan_to_proto(plan, codec) + } + + fn physical_expr_to_proto( + &self, + expr: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> Result { + DefaultPhysicalProtoConverter {}.physical_expr_to_proto(expr, codec) + } +} diff --git a/src/datafusion-server/src/tests/lineage.rs b/src/datafusion-server/src/tests/lineage.rs new file mode 100644 index 000000000..0d71cb883 --- /dev/null +++ b/src/datafusion-server/src/tests/lineage.rs @@ -0,0 +1,346 @@ +//! Exercise the real Flight boundary against local mode and plain DataFusion. + +use std::{ + fs::File, + path::Path, + sync::{Arc, Mutex}, + time::Duration, +}; + +use arrow::{ + array::{ArrayRef, BinaryArray, Date32Array, Int32Array, StringArray}, + compute::concat_batches, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, +}; +use arrow_flight::flight_service_server::FlightServiceServer; +use datafusion::{ + common::tree_node::{TreeNode, TreeNodeRecursion}, + prelude::{SessionConfig, SessionContext}, +}; +use liquid_cache::cache::{ + AlwaysHydrate, CacheEntry, CacheExpression, Date32Field, EvictionPolicy, LiquidPolicy, + TranscodeEvict, policies::EvictionOutcome, +}; +use liquid_cache_datafusion::{ + LiquidCacheParquetRef, optimizers::LineageHints, register_variant_functions, +}; +use liquid_cache_datafusion_client::{LiquidCacheClientBuilder, LiquidCacheClientExec}; +use liquid_cache_datafusion_local::LiquidCacheLocalBuilder; +use parquet::{arrow::ArrowWriter, variant::json_to_variant}; +use tokio::{net::TcpListener, task::JoinHandle}; + +use crate::LiquidCacheService; + +type ObservedLineages = Arc)>>>; + +#[derive(Debug)] +struct ObserveEviction(ObservedLineages); + +impl EvictionPolicy for ObserveEviction { + fn evict(&self, entry: &CacheEntry, lineage: Option<&CacheExpression>) -> EvictionOutcome { + if let CacheEntry::MemoryArrow(array) = entry { + self.0 + .lock() + .unwrap() + .push((array.clone(), lineage.cloned())); + } + TranscodeEvict.evict(entry, lineage) + } +} + +struct Fixture { + contexts: [SessionContext; 3], // local, Flight client, plain DataFusion + caches: [LiquidCacheParquetRef; 2], + observed: [ObservedLineages; 2], + server: JoinHandle<()>, + _dir: tempfile::TempDir, +} + +impl Drop for Fixture { + fn drop(&mut self) { + self.server.abort(); + } +} + +impl Fixture { + async fn new(memory_bytes: usize) -> Self { + let dir = tempfile::tempdir().unwrap(); + std::fs::create_dir(dir.path().join("local")).unwrap(); + let observed: [ObservedLineages; 2] = Default::default(); + let mut config = SessionConfig::new().with_target_partitions(1); + config + .options_mut() + .execution + .parquet + .schema_force_view_types = false; + config.options_mut().execution.parquet.skip_metadata = false; + let (local, local_cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(dir.path().join("local")) + .with_batch_size(8192 * 2) + .with_max_memory_bytes(memory_bytes) + .with_eviction_policy(Box::new(ObserveEviction(observed[0].clone()))) + .build(config.clone()) + .await + .unwrap(); + let service = LiquidCacheService::new( + LiquidCacheService::context().unwrap(), + Some(memory_bytes), + Some(dir.path().join("remote")), + Box::new(LiquidPolicy::new()), + Box::new(ObserveEviction(observed[1].clone())), + Box::new(AlwaysHydrate::new()), + ) + .await + .unwrap(); + let remote_cache = service.cache().clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::unfold(listener, |listener| async { + Some((listener.accept().await.map(|(socket, _)| socket), listener)) + }); + let server = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming(incoming) + .await + .unwrap(); + }); + let client = LiquidCacheClientBuilder::new(format!("http://{address}")) + .build(config.clone()) + .unwrap(); + let baseline = SessionContext::new_with_config(config); + register_variant_functions(&baseline); + let contexts = [local, client, baseline]; + for (table, offset) in [("a", 0), ("b", 400)] { + let path = dir.path().join(format!("{table}.parquet")); + write_fixture(&path, offset); + for ctx in &contexts { + ctx.register_parquet(table, path.to_str().unwrap(), Default::default()) + .await + .unwrap(); + } + } + Self { + contexts, + caches: [local_cache, remote_cache], + observed, + server, + _dir: dir, + } + } + + async fn assert_query(&self, sql: &str) { + let expected = query(&self.contexts[2], sql).await; + for (ctx, cache) in self.contexts[..2].iter().zip(&self.caches) { + assert_eq!(query(ctx, sql).await, expected, "cold: {sql}"); + let entries = cache.storage().stats().total_entries; + assert!(entries > 0, "query must exercise the cache: {sql}"); + assert_eq!(query(ctx, sql).await, expected, "warm: {sql}"); + let stats = cache.storage().stats(); + assert_eq!(stats.total_entries, entries); + assert!( + stats.runtime.get + stats.runtime.get_with_selection + stats.runtime.eval_predicate + > 0 + ); + } + } + + fn assert_lineages(&self, data_type: &DataType, expected: &[Option]) { + // Compare sets because a query can materialize and evict the same column repeatedly. + let canonical = |values: Vec>| { + let mut values: Vec<_> = values + .into_iter() + .map(|expr| expr.map(|expr| expr.to_metadata_value())) + .collect(); + values.sort(); + values.dedup(); + values + }; + for observed in &self.observed { + let values = observed + .lock() + .unwrap() + .iter() + .filter(|(array, _)| array.data_type() == data_type) + .map(|(_, expr)| expr.clone()) + .collect(); + assert_eq!(canonical(values), canonical(expected.to_vec())); + } + } +} + +async fn query(ctx: &SessionContext, sql: &str) -> RecordBatch { + let batches = tokio::time::timeout(Duration::from_secs(30), async { + ctx.sql(sql).await.unwrap().collect().await.unwrap() + }) + .await + .expect("query timed out"); + concat_batches(&batches[0].schema(), &batches).unwrap() +} + +fn write_fixture(path: &Path, offset: i32) { + let json: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"age":30,"name":"Alice"}"#), + Some(r#"{"age":25,"name":"Bob"}"#), + Some(r#"{"name":"Charlie"}"#), + None, + ])); + let variant = json_to_variant(&json).unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("date", DataType::Date32, true), + variant.field("data").as_ref().clone(), + Field::new("payload", DataType::Binary, true), + ])); + let dates = Date32Array::from(vec![ + Some(18_628 + offset), + Some(19_024 + offset), + Some(19_390 + offset), + None, + ]); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(dates), + ArrayRef::from(variant), + Arc::new(BinaryArray::from(vec![ + Some(&[0xff, 0x00][..]), + None, + Some(&[0x80][..]), + Some(&[][..]), + ])), + ], + ) + .unwrap(); + let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +#[tokio::test] +async fn client_projection_lineage_reaches_server_cache() { + let fixture = Fixture::new(0).await; + // The projection stays above the client-side join. Both components must travel. + let sql = "SELECT a.id, EXTRACT(YEAR FROM a.date) AS y, EXTRACT(MONTH FROM a.date) AS m FROM a JOIN b ON a.id = b.id ORDER BY a.id"; + let plan = fixture.contexts[1] + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let mut remote_scans = 0; + plan.apply(|node| { + if node.is::() { + // The server cannot recover this lineage from the fragment alone. + assert!( + LineageHints::analyze(node.children()[0]).is_empty(), + "{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); + remote_scans += 1; + return Ok(TreeNodeRecursion::Jump); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert_eq!(remote_scans, 2); + fixture.assert_query(sql).await; + fixture.assert_lineages( + &DataType::Date32, + &[CacheExpression::extract_date32_many([ + Date32Field::Year, + Date32Field::Month, + ])], + ); + for cache in &fixture.caches { + let trace = cache.consume_event_trace().to_string(); + assert!(trace.contains("expr=ExtractDate32[Year, Month]"), "{trace}"); + } +} + +#[tokio::test] +async fn mixed_raw_and_derived_usage_keeps_full_column() { + let fixture = Fixture::new(0).await; + fixture + .assert_query("SELECT id, date, payload, EXTRACT(YEAR FROM date) AS y FROM a ORDER BY id") + .await; + fixture.assert_lineages(&DataType::Date32, &[None]); +} + +#[tokio::test] +async fn join_lineage_is_scoped_to_each_scan() { + let fixture = Fixture::new(0).await; + fixture.assert_query("SELECT a.id, EXTRACT(YEAR FROM a.date) AS y, EXTRACT(MONTH FROM b.date) AS m FROM a JOIN b ON a.id = b.id ORDER BY a.id").await; + fixture.assert_lineages( + &DataType::Date32, + &[ + Some(CacheExpression::extract_date32(Date32Field::Year)), + Some(CacheExpression::extract_date32(Date32Field::Month)), + ], + ); + for observed in &fixture.observed { + for (array, expression) in observed.lock().unwrap().iter() { + if let Some(dates) = array.as_any().downcast_ref::() { + let field = match dates.value(0) { + 18_628 => Date32Field::Year, + 19_028 => Date32Field::Month, + value => panic!("unexpected date from join input: {value}"), + }; + assert_eq!(*expression, Some(CacheExpression::extract_date32(field))); + } + } + } +} + +#[tokio::test] +async fn variant_functions_and_lineage_cross_flight() { + let fixture = Fixture::new(0).await; + fixture.assert_query("SELECT id, variant_get(data, 'age', 'Int64') AS age, variant_get(data, 'name', 'Utf8') AS name FROM a ORDER BY id").await; + let variant_type = fixture.contexts[2] + .table("a") + .await + .unwrap() + .schema() + .field_with_unqualified_name("data") + .unwrap() + .data_type() + .clone(); + fixture.assert_lineages( + &variant_type, + &[Some(CacheExpression::variant_get_many([ + ("age", DataType::Int64), + ("name", DataType::Utf8), + ]))], + ); + // The predicate and aggregate execute on the server and require its UDF registry. + fixture + .assert_query("SELECT COUNT(*) FROM a WHERE variant_get(data, 'age', 'Int64') > 26") + .await; + fixture.assert_query("SELECT id, variant_pretty(data), variant_to_json(data), variant_to_json(variant_get(data, 'name')) FROM a ORDER BY id").await; + fixture + .assert_query("SELECT id, data FROM a ORDER BY id") + .await; +} + +#[tokio::test] +async fn later_raw_query_hydrates_cached_data() { + let fixture = Fixture::new(usize::MAX).await; + fixture.assert_query("SELECT id, EXTRACT(YEAR FROM date) AS y, variant_get(data, 'age', 'Int64') AS age FROM a ORDER BY id").await; + for cache in &fixture.caches { + cache.flush_data().await.unwrap(); + let stats = cache.storage().stats(); + assert!(stats.disk_arrow_entries > 0); + assert_eq!(stats.memory_arrow_entries, 0); + } + let sql = "SELECT id, date, variant_to_json(data) FROM a ORDER BY id"; + let expected = query(&fixture.contexts[2], sql).await; + for (ctx, cache) in fixture.contexts[..2].iter().zip(&fixture.caches) { + assert_eq!(query(ctx, sql).await, expected); + let stats = cache.storage().stats(); + assert!(stats.runtime.read_io_count > 0); + assert!(stats.memory_arrow_entries > 0); + } +} diff --git a/src/datafusion-server/src/tests/mod.rs b/src/datafusion-server/src/tests/mod.rs index fac8ea380..f83042d93 100644 --- a/src/datafusion-server/src/tests/mod.rs +++ b/src/datafusion-server/src/tests/mod.rs @@ -12,6 +12,7 @@ use liquid_cache::{ use uuid::Uuid; mod cases; +mod lineage; use crate::{LiquidCacheService, LiquidCacheServiceInner}; diff --git a/src/datafusion/src/lib.rs b/src/datafusion/src/lib.rs index 2c631f698..8d63d4279 100644 --- a/src/datafusion/src/lib.rs +++ b/src/datafusion/src/lib.rs @@ -14,3 +14,12 @@ pub use liquid_cache_common as common; pub use reader::variant_udf::{VariantGetUdf, VariantPretty, VariantToJsonUdf}; pub use reader::{FilterCandidateBuilder, LiquidParquetSource, LiquidPredicate, LiquidRowFilter}; pub use utils::{boolean_buffer_and_then, extract_execution_metrics}; + +/// Register the variant functions used by LiquidCache in a query session. +pub fn register_variant_functions(ctx: &datafusion::prelude::SessionContext) { + use datafusion::logical_expr::ScalarUDF; + + ctx.register_udf(ScalarUDF::new_from_impl(VariantGetUdf::default())); + ctx.register_udf(ScalarUDF::new_from_impl(VariantPretty::default())); + ctx.register_udf(ScalarUDF::new_from_impl(VariantToJsonUdf::default())); +} diff --git a/src/datafusion/src/optimizers/lineage.rs b/src/datafusion/src/optimizers/lineage.rs index 97479a0c2..c4b7b0fb0 100644 --- a/src/datafusion/src/optimizers/lineage.rs +++ b/src/datafusion/src/optimizers/lineage.rs @@ -368,8 +368,8 @@ impl HintAnalyzer { self.record(&ru); } - // Only equi-joins whose output is a straight concatenation of the two - // inputs, and that carry no residual filter, pass lineage through. Any + // Equi-joins without a residual filter pass lineage through, applying + // their output projection to the concatenated input columns. Any // other shape (semi/anti/mark joins, a residual filter we don't map) // is treated opaquely. let passthrough = join.filter().is_none() @@ -381,6 +381,9 @@ impl HintAnalyzer { if passthrough { let mut out = left; out.extend(right); + if let Some(projection) = &join.projection { + out = projection.iter().map(|&index| out[index].clone()).collect(); + } if out.len() == plan.schema().fields().len() { return out; } From 35636a3f9769114c8aef3284214da0699b97c020 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 13:38:09 +0530 Subject: [PATCH 13/24] feat(admission): footprint-based cache admission gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports hotdata-dev#1, #8 and #9 onto upstream's post-morsel optimizer. A parquet scan is routed through LiquidCache only when its estimated liquid footprint fits the cache; oversized scans stay on `ParquetSource` and read from the object store instead of evicting the working set to cache a scan that was never going to fit. The estimate is filter-aware and byte-accurate: per-file, the sum of the byte sizes of the columns the scan materializes (projection union predicate), over the files that survive the scan's own pruning predicate, deduped across the byte-range splits DataFusion makes of one file. A file with no stats, or a required column with an `Absent` size, falls back to the whole-file size, which over-counts in the safe direction. `Inexact` sizes are counted: they are real measurements, and rejecting them would make every scan against a catalog that labels them that way fall back to the whole file and bypass. The threshold is `memory × tolerance + disk`. Only the memory tier carries the overcommit, because LiquidCache compacts in RAM and keeps winning until roughly 5x over budget; the disk tier is counted at face value, since a scan that overflows RAM spills to it rather than thrashing and there is no evidence for extending the RAM crossover there. Computed in f64 so a budget near u64::MAX cannot wrap. The gate is a pure performance decision — admitting a scan or bypassing it returns identical rows — so estimation runs under a panic guard whose `strict` flag chooses between aborting the query (surface the bug) and caching normally (keep queries running). Every decision logs one line under `liquid_cache::admission` with the full breakdown; without it the gate is a black box and each tuning cycle costs a benchmark run. Off by default: `LocalModeOptimizer` and `LiquidCacheLocalBuilder` cache every scan unless `with_admission_gate` is called, so upstream behaviour is unchanged. --- src/datafusion-local/src/lib.rs | 30 +- src/datafusion/src/optimizers/mod.rs | 761 ++++++++++++++++++++++++++- 2 files changed, 785 insertions(+), 6 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index d8cbbadd5..898c6b0e0 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,6 +69,10 @@ pub struct LiquidCacheLocalBuilder { /// Hydration policy hydration_policy: Box, prefetch: bool, + /// Footprint-based admission gate `(expansion, safety, tolerance, strict)`. + /// When set, a scan is cached only if its estimated liquid footprint stays + /// within the admission threshold; `strict` toggles fail-loud panic handling. + admission: Option<(f64, f64, f64, bool)>, span: fastrace::Span, } @@ -85,6 +89,7 @@ impl Default for LiquidCacheLocalBuilder { eviction_policy: Box::new(TranscodeEvict), hydration_policy: Box::new(AlwaysHydrate::new()), prefetch: true, + admission: None, span: fastrace::Span::enter_with_local_parent("liquid_cache_datafusion_local_builder"), } } @@ -152,6 +157,26 @@ impl LiquidCacheLocalBuilder { self } + /// Enable the footprint-based admission gate. A scan is cached only when its + /// estimated liquid footprint (raw required bytes x `expansion` x `safety`) + /// stays within the admission threshold (`memory × tolerance + disk`); larger + /// scans are read directly from the parquet source, bypassing the cache. + /// `expansion`/`safety` are `>= 1.0` (inflate the estimate); `tolerance` is + /// `>= 1.0` (overcommit the memory tier, clamped to the measured ~5x + /// compaction crossover). `strict == true` lets a footprint-estimation panic + /// abort the query (fail loud); `false` catches it and caches the scan + /// normally. + pub fn with_admission_gate( + mut self, + expansion: f64, + safety: f64, + tolerance: f64, + strict: bool, + ) -> Self { + self.admission = Some((expansion, safety, tolerance, strict)); + self + } + /// Build a SessionContext with liquid cache configured /// Returns the SessionContext and the liquid cache reference pub async fn build( @@ -197,7 +222,10 @@ impl LiquidCacheLocalBuilder { .await; let cache_ref = Arc::new(cache); - let optimizer = LocalModeOptimizer::new(cache_ref.clone()).with_prefetch(self.prefetch); + let mut optimizer = LocalModeOptimizer::new(cache_ref.clone()).with_prefetch(self.prefetch); + if let Some((expansion, safety, tolerance, strict)) = self.admission { + optimizer = optimizer.with_admission_gate(expansion, safety, tolerance, strict); + } let state = datafusion::execution::SessionStateBuilder::new() .with_config(config) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 7e30acde2..dd20d02a4 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -2,14 +2,26 @@ mod lineage; +use std::collections::HashSet; use std::sync::Arc; use datafusion::{ + arrow::datatypes::SchemaRef, catalog::memory::DataSourceExec, - common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + common::{ + Statistics, + pruning::PrunableStatistics, + stats::Precision, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + }, config::ConfigOptions, - datasource::{physical_plan::ParquetSource, source::DataSource}, - physical_optimizer::PhysicalOptimizerRule, + datasource::{ + listing::PartitionedFile, + physical_plan::{FileScanConfig, FileSource, ParquetSource}, + source::DataSource, + }, + physical_expr::utils::collect_columns, + physical_optimizer::{PhysicalOptimizerRule, pruning::PruningPredicateBuilder}, physical_plan::ExecutionPlan, }; @@ -18,6 +30,44 @@ pub use lineage::LineageHints; use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnLineages}; +/// Parameters for the footprint-based admission gate. +/// +/// A parquet scan is routed through LiquidCache only if its estimated liquid +/// footprint stays within the admission threshold; larger scans are left as +/// vanilla parquet reads (which, if the object store is a cached mount, read +/// from it). +/// +/// The threshold spans both cache tiers, weighted differently: +/// `memory × tolerance + disk`. A scan that overflows RAM spills to the on-disk +/// liquid tier rather than thrashing, so disk capacity counts toward what fits — +/// at face value, since only the RAM tier has the measured compaction +/// overcommit. With the disk tier off it is just `memory × tolerance`. +/// +/// The estimate multiplies the raw required parquet bytes by `expansion` +/// (parquet -> liquid in-memory blow-up) and `safety` (extra margin); both are +/// `>= 1.0` so the estimate is conservative (over-counts). +#[derive(Debug, Clone, Copy)] +pub struct AdmissionGate { + /// Parquet-bytes -> liquid-in-memory-bytes multiplier (>= 1.0). Inflates the + /// estimate (conservative direction). + pub expansion: f64, + /// Extra safety margin on the estimate (>= 1.0). Conservative direction. + pub safety: f64, + /// Overcommit tolerance on the *memory* tier: how far the estimated liquid + /// footprint may exceed the RAM budget before it counts against the scan, in + /// multiples of the RAM budget. LiquidCache compacts in RAM up to ~5x over + /// budget before it starts thrashing, so caching still wins in that band. + /// This is the one *relaxing* knob, so it is clamped to `[1.0, 5.0]` (5.0 = + /// the measured crossover). It applies only to memory; the disk tier is + /// counted at face value (`memory × tolerance + disk`), never scaled by it. + pub tolerance: f64, + /// Fail-loud mode. When `true`, a panic during footprint estimation is *not* + /// caught — it aborts the query — so estimation bugs surface immediately. + /// When `false`, the panic is caught, logged at ERROR, and the scan is cached + /// normally so the query survives. The caller chooses the default. + pub strict: bool, +} + /// Physical optimizer rule for local mode liquid cache. /// /// Rewrites `DataSourceExec` parquet scans to use [`LiquidParquetSource`], and @@ -27,6 +77,10 @@ use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnLineages}; pub struct LocalModeOptimizer { cache: LiquidCacheParquetRef, prefetch: bool, + /// When set, a scan whose estimated footprint exceeds the cache budget is + /// left as a vanilla parquet read instead of being wrapped by LiquidCache. + /// `None` means cache every scan. + admission: Option, } impl LocalModeOptimizer { @@ -35,6 +89,7 @@ impl LocalModeOptimizer { Self { cache, prefetch: true, + admission: None, } } @@ -48,6 +103,41 @@ impl LocalModeOptimizer { self.prefetch = prefetch; self } + + /// Enable the footprint-based admission gate. A parquet scan is cached only + /// when its estimated liquid footprint (raw required bytes x `expansion` x + /// `safety`) stays within the admission threshold (`memory × tolerance + + /// disk`); otherwise it is read directly from the parquet source, bypassing + /// the cache. See [`AdmissionGate`]. + /// + /// Inputs are sanitized so a misconfigured value can never make the gate + /// unsound: `expansion`/`safety` are forced finite and `>= 1.0` (their only + /// effect is to inflate the estimate, the conservative direction), and + /// `tolerance` — the one *relaxing* knob — is forced finite and clamped to + /// `[1.0, 5.0]` (5.0 = the measured compaction crossover), defaulting to 3.0 + /// if non-finite. `strict` (see [`AdmissionGate::strict`]) turns off the + /// panic guard so estimation bugs fail loud. + pub fn with_admission_gate( + mut self, + expansion: f64, + safety: f64, + tolerance: f64, + strict: bool, + ) -> Self { + let estimate_factor = |v: f64| if v.is_finite() && v >= 1.0 { v } else { 1.0 }; + let tolerance = if tolerance.is_finite() { + tolerance.clamp(1.0, 5.0) + } else { + 3.0 + }; + self.admission = Some(AdmissionGate { + expansion: estimate_factor(expansion), + safety: estimate_factor(safety), + tolerance, + strict, + }); + self + } } impl PhysicalOptimizerRule for LocalModeOptimizer { @@ -59,7 +149,25 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { let analysis = HintAnalyzer::analyze(&plan); let cache = self.cache.clone(); let prefetch = self.prefetch; + let admission = self.admission; + // The gate sizes against both cache tiers, not just RAM: when a scan + // overflows memory its entries spill to the on-disk liquid tier (NVMe) + // instead of thrashing. They are weighted differently (see + // `admission_threshold`) — RAM carries the compaction overcommit, disk + // counts at face value — so they are passed through separately. With the + // disk tier off (`max_disk_bytes == 0`) the threshold is exactly the RAM + // budget, so gate behaviour is unchanged there. + let memory_budget = self.cache.max_memory_bytes() as u64; + let disk_budget = self.cache.max_disk_bytes() as u64; let mut convert = |node: &Arc, hints: ColumnLineages| { + // Leave scans whose estimated liquid footprint exceeds the budget as + // vanilla parquet reads, so oversized scans don't thrash the cache. + if let Some(gate) = admission + && let Some((cfg, src)) = parquet_scan_parts(node) + && should_bypass_guarded(cfg, src, gate, memory_budget, disk_budget) + { + return None; + } convert_parquet_scan(node, &cache, hints, prefetch) }; Ok(lineage::rewrite_with_hints(plan, &mut convert, &analysis)) @@ -107,6 +215,327 @@ pub fn rewrite_data_source_plan( rewrite_data_source_plan_with_hints(plan, cache, &ColumnLineages::default()) } +/// If `node` is a parquet `DataSourceExec`, return its `FileScanConfig` and +/// `ParquetSource`. +fn parquet_scan_parts(node: &Arc) -> Option<(&FileScanConfig, &ParquetSource)> { + let dse = node.downcast_ref::()?; + dse.downcast_to_file_source::() +} + +/// A scan's footprint estimate plus the breakdown behind it, for the admission +/// decision ([`should_bypass`]) and its diagnostic log line. +#[derive(Debug, Default, Clone, Copy)] +struct FootprintEstimate { + /// Raw required-column parquet bytes (the decision input). + raw_bytes: u64, + /// Number of file columns the scan materializes (projection ∪ predicate). + required_cols: usize, + /// Distinct physical files charged (survived pruning, deduped across the + /// byte-range splits of the same file). + charged_files: usize, + /// Raw `PartitionedFile` count before dedupe (one file may be split into + /// several byte-range partitions for parallelism). + partitioned_files: usize, + /// Surviving files charged the *whole file* size because they lacked + /// per-column byte sizes. A non-zero count means the catalog has no + /// per-column size stat for this table, so the estimate is coarse and + /// over-counts. + fallback_files: usize, +} + +/// Estimate the raw required-column parquet bytes a scan reads: the sum, over +/// the files that survive the scan's predicate, of the byte sizes of the columns +/// it materializes. This is the byte-accurate, filter-aware size the admission +/// decision is built on (see [`should_bypass`]); the expansion/safety/tolerance +/// factors are applied there, not here. +/// +/// "Required columns" is the output projection **unioned with the predicate +/// columns**, since LiquidCache materializes both. Byte sizing uses `Exact` or +/// `Inexact` per-column sizes; a column with an `Absent` size or a file with no +/// stats falls back to the whole file size. The sum saturates rather than +/// overflowing on pathological file lists. +fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> FootprintEstimate { + let num_file_cols = cfg.file_schema().fields().len(); + // Full table schema (file + partition columns). The pushed-down predicate + // and `PartitionedFile` stats are expressed against it, so file pruning must + // use it; byte accounting still charges only file columns. + let table_schema = cfg.file_source.table_schema().table_schema().clone(); + + // Columns the scan projects. `column_indices()` collects the source columns + // referenced by each projection expression, so it is correct for compound + // projections (e.g. `a * b` reads a and b) and does not panic on a + // non-column projection expression. Indices are table-schema-relative; keep + // only file columns (partition columns are literals, never materialized). + let mut required: Vec = match src.projection() { + Some(p) => p + .column_indices() + .into_iter() + .filter(|&i| i < num_file_cols) + .collect(), + None => (0..num_file_cols).collect(), + }; + if let Some(pred) = src.filter() { + for col in collect_columns(&pred) { + let idx = col.index(); + // Partition columns are appended after file columns in the table + // schema and are literals (never materialized in LiquidCache), so + // only file columns contribute to the footprint. + if idx < num_file_cols { + required.push(idx); + } + } + } + required.sort_unstable(); + required.dedup(); + + let files: Vec<&PartitionedFile> = cfg.file_groups.iter().flat_map(|g| g.files()).collect(); + let required_cols = required.len(); + if files.is_empty() { + return FootprintEstimate { + required_cols, + ..FootprintEstimate::default() + }; + } + + let surviving = surviving_files(src, &table_schema, &files); + + // Dedupe by physical file identity (path + size). DataFusion splits one file + // into several byte-range `PartitionedFile`s for parallelism, each a *clone* + // of the whole file's statistics (and whole-file object size), differing only + // in `range`. Counting every range would multiply a single file's footprint + // by the split count, so charge each distinct file once — pruning is + // file-granular, so a surviving file means caching its whole required columns + // regardless of how it was split. Size is part of the key so two genuinely + // distinct objects that share a path are never collapsed. + let mut seen: HashSet<(object_store::path::Path, u64)> = HashSet::new(); + let mut raw_bytes = 0u64; + let mut charged_files = 0usize; + let mut fallback_files = 0usize; + for (f, keep) in files.iter().zip(surviving.iter()) { + if !*keep { + continue; + } + if !seen.insert((f.object_meta.location.clone(), f.object_meta.size)) { + continue; + } + charged_files += 1; + let (bytes, fell_back) = + file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required); + if fell_back { + fallback_files += 1; + } + raw_bytes = raw_bytes.saturating_add(bytes); + } + + FootprintEstimate { + raw_bytes, + required_cols, + charged_files, + partitioned_files: files.len(), + fallback_files, + } +} + +/// The admission threshold in bytes: the largest estimated liquid footprint the +/// cache admits before a scan is bypassed. +/// +/// The two tiers are weighted differently. The memory tier carries the +/// compaction overcommit — `tolerance`, the measured RAM crossover where +/// LiquidCache still beats the fallback mount despite spilling. The on-disk +/// liquid tier is counted at **face value (1×)**: a scan that overflows RAM +/// spills to NVMe without thrashing, so its capacity needs no overcommit — and +/// extending the RAM crossover factor to disk has no evidence behind it. Hence +/// `memory × tolerance + disk` rather than `(memory + disk) × tolerance`. +/// +/// Computed in `f64` so the `memory × tolerance` product cannot overflow `u64`. +/// A zero threshold (both tiers unsized / cache disabled) bypasses any non-empty +/// scan and admits only zero-footprint ones. +fn admission_threshold(memory_budget: u64, disk_budget: u64, tolerance: f64) -> f64 { + (memory_budget as f64) * tolerance + (disk_budget as f64) +} + +/// Decide whether a parquet scan should bypass the cache (read as vanilla +/// parquet) rather than be transcoded into LiquidCache. +/// +/// The scan's estimated liquid footprint is `raw × expansion × safety`, where +/// `raw` is [`estimate_required_bytes`]. It bypasses when that footprint exceeds +/// the [`admission_threshold`] (`memory × tolerance + disk`). The comparison is +/// done in finite `f64` space to avoid integer overflow. +fn should_bypass( + cfg: &FileScanConfig, + src: &ParquetSource, + gate: AdmissionGate, + memory_budget: u64, + disk_budget: u64, +) -> bool { + let est = estimate_required_bytes(cfg, src); + // Multipliers are already sanitized to finite, >= 1.0 in `with_admission_gate`. + let footprint = (est.raw_bytes as f64) * gate.expansion * gate.safety; + let threshold = admission_threshold(memory_budget, disk_budget, gate.tolerance); + let bypass = footprint > threshold; + + // One line per admission decision. Without this the gate is a black box and + // every tuning cycle costs a benchmark run to infer decisions from side + // effects. `fallback_files > 0` is the flag that the catalog has no + // per-column sizes for this table (estimate is coarse / over-counts). + let path = cfg + .file_groups + .iter() + .flat_map(|g| g.files()) + .next() + .map(|f| f.object_meta.location.as_ref()) + .unwrap_or(""); + log::info!( + target: "liquid_cache::admission", + "admission {verdict}: file={path} projected_cols={cols} \ + charged_files={charged} partitioned_files={total} \ + fallback_files={fb} raw_bytes={raw} footprint_bytes={fp} \ + memory_bytes={mem} disk_bytes={disk} threshold_bytes={thr} \ + expansion={exp} safety={saf} tolerance={tol}", + verdict = if bypass { "BYPASS" } else { "ADMIT" }, + cols = est.required_cols, + charged = est.charged_files, + total = est.partitioned_files, + fb = est.fallback_files, + raw = est.raw_bytes, + fp = footprint as u64, + mem = memory_budget, + disk = disk_budget, + thr = threshold as u64, + exp = gate.expansion, + saf = gate.safety, + tol = gate.tolerance, + ); + + bypass +} + +/// [`should_bypass`], with an optional panic guard. +/// +/// The gate is a pure performance optimization: caching a scan or reading it as +/// vanilla parquet yields identical results. So if footprint estimation ever +/// panics — e.g. a DataFusion API that panics on an unusual plan shape — a +/// non-strict gate must not let it abort the query. +/// +/// Either way the panic is caught and logged at ERROR with its message (never +/// silently swallowed), and the log advises flipping the admission gate's +/// `strict` flag for the alternative behavior: +/// +/// - `gate.strict == true`: log, then re-raise so the panic aborts the query and +/// the bug surfaces immediately. The log advises turning the flag *off* to +/// keep queries running while it is fixed. +/// - `gate.strict == false`: log, then cache the scan normally so the query +/// survives. The log advises turning the flag *on* to fail loud instead. +fn should_bypass_guarded( + cfg: &FileScanConfig, + src: &ParquetSource, + gate: AdmissionGate, + memory_budget: u64, + disk_budget: u64, +) -> bool { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + should_bypass(cfg, src, gate, memory_budget, disk_budget) + })); + match result { + Ok(bypass) => bypass, + Err(payload) => { + let msg = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or(""); + if gate.strict { + log::error!( + "liquid-cache admission gate panicked during footprint estimation: \ + {msg}. Aborting query (admission gate is in strict mode). Configure the \ + admission gate with strict=false to fall back to caching and keep queries \ + running while this is fixed." + ); + std::panic::resume_unwind(payload); + } + log::error!( + "liquid-cache admission gate panicked during footprint estimation: {msg}; \ + caching scan normally. Configure the admission gate with strict=true to fail \ + loud instead." + ); + false + } + } +} + +/// Boolean per file: `true` if the file may match the predicate (keep it), +/// `false` if the predicate's stats prove it cannot match (prune it). +/// Conservative: no predicate, missing/mismatched stats, or any pruning error +/// keeps files. `table_schema` (file + partition columns) is used so predicates +/// on partition columns resolve, matching `PartitionedFile::statistics`. +fn surviving_files( + src: &ParquetSource, + table_schema: &SchemaRef, + files: &[&PartitionedFile], +) -> Vec { + let Some(pred) = src.filter() else { + return vec![true; files.len()]; + }; + let pruning = match PruningPredicateBuilder::new() + .with_file_schema(table_schema.clone()) + .try_build(pred) + { + Ok(p) => p, + Err(_) => return vec![true; files.len()], + }; + let expected = table_schema.fields().len(); + let stats: Vec> = files + .iter() + .map(|f| match &f.statistics { + // Only trust stats whose width matches the table schema; otherwise + // treat as unknown (kept, not pruned) to avoid a schema mismatch. + Some(s) if s.column_statistics.len() == expected => s.clone(), + _ => Arc::new(Statistics::new_unknown(table_schema)), + }) + .collect(); + let prunable = PrunableStatistics::new(stats, table_schema.clone()); + match pruning.prune(&prunable) { + Ok(mask) if mask.len() == files.len() => mask, + _ => vec![true; files.len()], + } +} + +/// Bytes the `required` columns of one file contribute to the footprint. +/// +/// Uses per-column byte sizes that are either `Exact` or `Inexact`. A catalog +/// may record the real compressed on-disk column size and still label it +/// `Inexact` (stats can go stale after deletes/compaction), so rejecting +/// `Inexact` would make the gate fall back to the whole-file size on every such +/// scan — charging all columns for a single-column read and bypassing +/// everything. `Inexact` is a real measurement, not a guess; the caller's +/// `expansion`/`safety` margin absorbs modest drift, and even a large stale-low +/// under-count only risks a too-eager admit (perf), never wrong results. +/// +/// If the file has no stats, or any required column's size is `Absent`, fall +/// back to the whole-file size — a deliberate over-estimate. +/// +/// Returns `(bytes, fell_back)`, where `fell_back` is `true` when the whole-file +/// over-estimate was used. +fn file_required_bytes( + stats: Option<&Statistics>, + object_size: u64, + required: &[usize], +) -> (u64, bool) { + let Some(stats) = stats else { + return (object_size, true); + }; + let mut sum: u64 = 0; + for &c in required { + match stats.column_statistics.get(c).map(|cs| &cs.byte_size) { + Some(Precision::Exact(n) | Precision::Inexact(n)) => { + sum = sum.saturating_add(*n as u64) + } + _ => return (object_size, true), + } + } + (sum, false) +} + /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent /// node backed by [`LiquidParquetSource`] carrying `hints`. fn convert_parquet_scan( @@ -157,12 +586,22 @@ mod tests { use super::*; async fn make_cache(path: &Path) -> LiquidCacheParquetRef { + make_cache_with_mem_disk(path, 1000000, usize::MAX).await + } + + /// The admission gate reads both tier budgets off the cache, so its tests + /// need them set independently rather than via [`make_cache`]'s defaults. + async fn make_cache_with_mem_disk( + path: &Path, + max_memory_bytes: usize, + max_disk_bytes: usize, + ) -> LiquidCacheParquetRef { let store = t4::mount(path.join("liquid_cache.t4")).await.unwrap(); Arc::new( LiquidCacheParquet::new( 8192, - 1000000, - usize::MAX, + max_memory_bytes, + max_disk_bytes, store, Box::new(LiquidPolicy::new()), Box::new(TranscodeEvict), @@ -172,6 +611,193 @@ mod tests { ) } + /// True if any parquet scan in `plan` was rewritten to `LiquidParquetSource`. + fn has_liquid_source(plan: &Arc) -> bool { + let mut found = false; + plan.apply(|node| { + if let Some(exec) = node.downcast_ref::() + && let Some(cfg) = exec.data_source().downcast_ref::() + && cfg + .file_source() + .downcast_ref::() + .is_some() + { + found = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + found + } + + async fn nano_hits_plan() -> Arc { + let ctx = SessionContext::new(); + ctx.register_parquet( + "nano_hits", + "../../examples/nano_hits.parquet", + Default::default(), + ) + .await + .unwrap(); + ctx.sql("SELECT * FROM nano_hits WHERE \"URL\" like 'https://%' limit 10") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap() + } + + /// The admission gate bypasses a scan whose estimated footprint exceeds the + /// budget (large expansion here forces that), and caches one that fits. + #[tokio::test] + async fn admission_gate_bypasses_an_oversized_scan() { + let plan = nano_hits_plan().await; + let config = ConfigOptions::default(); + let capped_dir = tempfile::tempdir().unwrap(); + let uncapped_dir = tempfile::tempdir().unwrap(); + + // A huge expansion inflates the estimate past the 1 MB budget → bypass. + let capped = LocalModeOptimizer::new( + make_cache_with_mem_disk(capped_dir.path(), 1_000_000, 0).await, + ) + .with_admission_gate(1e9, 1.0, 1.0, false) + .optimize(plan.clone(), &config) + .unwrap(); + assert!( + !has_liquid_source(&capped), + "oversized estimate should stay a plain ParquetSource" + ); + + // With a large budget the footprint fits (expansion 1.0) → cached. + let uncapped = LocalModeOptimizer::new( + make_cache_with_mem_disk(uncapped_dir.path(), usize::MAX, 0).await, + ) + .with_admission_gate(1.0, 1.0, 1.0, false) + .optimize(plan, &config) + .unwrap(); + assert!( + has_liquid_source(&uncapped), + "fitting scan should be wrapped in LiquidParquetSource" + ); + } + + /// The budget counts the on-disk liquid tier, not just memory: the same scan + /// that a memory-only budget bypasses is admitted once the disk tier has room + /// for it (evicted entries spill to disk instead of thrashing). + #[tokio::test] + async fn admission_gate_counts_the_disk_tier() { + let plan = nano_hits_plan().await; + let config = ConfigOptions::default(); + let mem_only_dir = tempfile::tempdir().unwrap(); + let with_disk_dir = tempfile::tempdir().unwrap(); + + // A huge expansion inflates the estimate past the 1 MB memory budget, and + // with no disk tier there is nowhere else for it to fit → bypass. + let mem_only = LocalModeOptimizer::new( + make_cache_with_mem_disk(mem_only_dir.path(), 1_000_000, 0).await, + ) + .with_admission_gate(1e9, 1.0, 1.0, false) + .optimize(plan.clone(), &config) + .unwrap(); + assert!( + !has_liquid_source(&mem_only), + "with a memory-only budget the oversized estimate should bypass" + ); + + // Same 1 MB memory and same estimate, but now a large disk tier — the + // threshold is memory × tolerance + disk, so the scan fits and is cached. + let with_disk = LocalModeOptimizer::new( + make_cache_with_mem_disk(with_disk_dir.path(), 1_000_000, usize::MAX).await, + ) + .with_admission_gate(1e9, 1.0, 1.0, false) + .optimize(plan, &config) + .unwrap(); + assert!( + has_liquid_source(&with_disk), + "the disk tier's capacity should count toward the budget and admit the scan" + ); + } + + /// A scan projection can hold a non-column expression (`SELECT s.a` pushes a + /// struct field access into the scan). The estimator must read it without + /// panicking — the class of bug that made the gate's panic guard necessary. + #[tokio::test] + async fn estimate_survives_non_column_scan_projection() { + use arrow::array::{ArrayRef, Int64Array, StructArray}; + use arrow_schema::Fields; + + let struct_fields = Fields::from(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Struct(struct_fields.clone()), false), + Field::new("p", DataType::Int64, false), + ])); + let struct_array = StructArray::new( + struct_fields, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(Int64Array::from(vec![4, 5, 6])) as ArrayRef, + ], + None, + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(struct_array) as ArrayRef, + Arc::new(Int64Array::from(vec![7, 8, 9])), + ], + ) + .unwrap(); + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("structs.parquet"); + let mut writer = ArrowWriter::try_new(File::create(&path).unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let ctx = SessionContext::new(); + ctx.register_parquet("t", path.to_str().unwrap(), Default::default()) + .await + .unwrap(); + let plan = ctx + .sql("SELECT s.a FROM t") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + // Locate the parquet scan and assert its projection really does carry a + // non-column expression — otherwise the test wouldn't exercise the bug. + let mut parts = None; + plan.apply(|node| { + if let Some((cfg, src)) = parquet_scan_parts(node) { + let has_expr = src + .projection() + .map(|p| p.iter().any(|e| e.expr.downcast_ref::().is_none())) + .unwrap_or(false); + assert!( + has_expr, + "expected a non-column scan projection to exercise the gate; \ + if this fails, DataFusion changed leaf-expression pushdown" + ); + parts = Some((cfg.clone(), src.clone())); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + let (cfg, src) = parts.expect("no parquet scan in plan"); + + // The call that used to panic. It must return a finite estimate; `s` + // (the struct column read by `s.a`) is the one required file column. + let est = estimate_required_bytes(&cfg, &src); + assert!(est.raw_bytes > 0, "estimate should be a real byte count"); + } + fn liquid_source(plan: &Arc) -> LiquidParquetSource { let mut source = None; plan.apply(|node| { @@ -342,3 +968,128 @@ mod tests { assert_eq!(pruning_metrics.pruned(), 1); } } + +/// Pure unit tests for the footprint byte-math (no cache / no t4 mount, so they +/// run everywhere). These cover the fallback direction: a column with an +/// `Absent` size (or a file with no stats) falls back to the whole-file size, +/// while `Exact`/`Inexact` per-column sizes are counted. +#[cfg(test)] +mod footprint_tests { + use super::file_required_bytes; + use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; + + fn col(byte_size: Precision) -> ColumnStatistics { + let mut c = ColumnStatistics::new_unknown(); + c.byte_size = byte_size; + c + } + + fn stats(cols: Vec) -> Statistics { + Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: cols, + } + } + + #[test] + fn no_stats_falls_back_to_full_file() { + assert_eq!(file_required_bytes(None, 5000, &[0, 1]), (5000, true)); + } + + #[test] + fn sums_only_required_exact_columns() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Exact(200)), + col(Precision::Exact(400)), + ]); + // required = cols 0 and 2 → 100 + 400, ignoring col 1. No fallback. + assert_eq!(file_required_bytes(Some(&s), 9999, &[0, 2]), (500, false)); + } + + #[test] + fn inexact_required_column_is_counted() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Inexact(200)), + ]); + // Inexact is a real (possibly-stale) size, so it counts (no fallback). + // A catalog that always labels byte_size Inexact would otherwise make + // every scan fall back to the whole file and bypass the cache. + assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), (300, false)); + assert_eq!(file_required_bytes(Some(&s), 7000, &[1]), (200, false)); + } + + #[test] + fn absent_among_sized_columns_falls_back_to_full_file() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Inexact(200)), + col(Precision::Absent), + ]); + // Any Absent required column → whole file (fallback), even mixed with sized ones. + assert_eq!( + file_required_bytes(Some(&s), 8000, &[0, 1, 2]), + (8000, true) + ); + // Dropping the absent column, Exact + Inexact are summed (no fallback). + assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1]), (300, false)); + } + + #[test] + fn missing_required_column_falls_back_to_full_file() { + let s = stats(vec![col(Precision::Exact(100))]); + // required col 5 doesn't exist → conservative whole file (fallback). + assert_eq!(file_required_bytes(Some(&s), 3000, &[0, 5]), (3000, true)); + } + + #[test] + fn absent_byte_size_falls_back_to_full_file() { + let s = stats(vec![col(Precision::Absent)]); + assert_eq!(file_required_bytes(Some(&s), 2000, &[0]), (2000, true)); + } +} + +/// Pure unit tests for the admission threshold arithmetic (no cache / no t4 +/// mount, so they run everywhere). +#[cfg(test)] +mod threshold_tests { + use super::admission_threshold; + + #[test] + fn memory_carries_tolerance_disk_is_face_value() { + // memory × tolerance + disk = 10×3 + 100 = 130. This pins the exact + // formula: it is NOT (memory+disk)×tolerance (=330), max(mem,disk) (=100), + // nor disk alone (=100). + assert_eq!(admission_threshold(10, 100, 3.0), 130.0); + } + + #[test] + fn disk_off_is_memory_times_tolerance() { + // With no disk tier the threshold is exactly the RAM budget × tolerance, + // so gate behaviour matches the pre-disk model. + assert_eq!(admission_threshold(1_000, 0, 3.0), 3_000.0); + } + + #[test] + fn disk_gets_no_overcommit() { + // A pure-disk budget contributes at 1×, never scaled by tolerance. + assert_eq!(admission_threshold(0, 500, 5.0), 500.0); + } + + #[test] + fn both_tiers_unsized_is_zero() { + // Zero threshold => any non-empty scan bypasses (cache disabled/unsized). + assert_eq!(admission_threshold(0, 0, 3.0), 0.0); + } + + #[test] + fn large_memory_times_tolerance_does_not_overflow() { + // memory × tolerance is computed in f64, so a budget near u64::MAX cannot + // wrap (the bug an integer `budget * tolerance` would have). + let t = admission_threshold(u64::MAX, 0, 5.0); + assert!(t > u64::MAX as f64); + assert!(t.is_finite()); + } +} From 749b6efd7e751956e09bae9f44282cc0f5f05324 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 13:45:55 +0530 Subject: [PATCH 14/24] test(reader): pin conjunct-pushdown correctness Ports the regression coverage from hotdata-dev#20 and #40 without their fix: upstream now evaluates nested-column predicates rather than dropping them, so the rows come out right by a better route than our decline-the-scan bypass, which cost the scan its cache. Three cases upstream's own nested_filter.rs does not reach: a nested conjunct alongside a pushable one, a nested column reached through OR (one candidate for the whole predicate, so refusing it applies no filter at all), and a conjunct on a column absent from the file schema. Each is checked cold and warm, since the cached path evaluates predicates separately from the source path. --- src/datafusion-local/src/tests/mod.rs | 1 + .../src/tests/unevaluable_conjunct.rs | 148 ++++++++++++++++++ 2 files changed, 149 insertions(+) create mode 100644 src/datafusion-local/src/tests/unevaluable_conjunct.rs diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index b5486a276..e57b302c0 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -22,6 +22,7 @@ use crate::LiquidCacheLocalBuilder; mod date_optimizer; mod filter_limit; mod nested_filter; +mod unevaluable_conjunct; mod variants; const TEST_FILE: &str = "../../examples/nano_hits.parquet"; diff --git a/src/datafusion-local/src/tests/unevaluable_conjunct.rs b/src/datafusion-local/src/tests/unevaluable_conjunct.rs new file mode 100644 index 000000000..8180b6bb4 --- /dev/null +++ b/src/datafusion-local/src/tests/unevaluable_conjunct.rs @@ -0,0 +1,148 @@ +//! A pushed-down conjunct the liquid row filter may not be able to evaluate. +//! +//! `build_row_filter` splits the pushed-down predicate into conjuncts and builds +//! one `FilterCandidate` per conjunct. A conjunct that is refused and then +//! dropped leaves the scan applying a *strictly weaker* filter than the query +//! asked for — and since DataFusion removes the `FilterExec` when it pushes a +//! predicate down, nothing re-applies what was dropped. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, Int64Array, RecordBatch, StructArray}; +use arrow_schema::{DataType, Field, Fields, Schema}; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::{ListingOptions, ListingTableUrl}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// Eight rows: `id` 0..8, `st` a `struct` whose `a` mirrors `id`. Exactly +/// one row has `st.a = 3`. +fn write_t(path: &Path) { + let struct_fields = Fields::from(vec![Field::new("a", DataType::Int32, false)]); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("st", DataType::Struct(struct_fields.clone()), false), + ])); + + let id: Int64Array = (0..8i64).collect::>().into(); + let a: ArrayRef = Arc::new((0..8i32).collect::()); + let st = StructArray::new(struct_fields, vec![a], None); + + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(id), Arc::new(st)]).unwrap(); + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +async fn liquid_ctx(cache_dir: &Path) -> SessionContext { + std::fs::create_dir_all(cache_dir).unwrap(); + let (ctx, _cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir.to_path_buf()) + .build(SessionConfig::new()) + .await + .unwrap(); + ctx +} + +async fn ids(ctx: &SessionContext, sql: &str) -> Vec { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let mut out = Vec::new(); + for batch in batches { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + out.extend(column.iter().map(|v| v.unwrap())); + } + out.sort_unstable(); + out +} + +/// One pushable conjunct (`id >= 0`, all eight rows) and one nested +/// (`st.a = 3`, one row). Dropping the second returns all eight. +#[tokio::test] +async fn nested_column_conjunct_is_still_applied() { + let dir = TempDir::new().unwrap(); + let parquet = dir.path().join("t.parquet"); + write_t(&parquet); + let ctx = liquid_ctx(&dir.path().join("cache")).await; + ctx.register_parquet( + "t", + parquet.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + let sql = "SELECT id FROM t WHERE id >= 0 AND st.a = 3"; + // Cold reads through the source and fills the cache; warm is served from it, + // a separate evaluation path. + for pass in ["cold", "warm"] { + assert_eq!(ids(&ctx, sql).await, vec![3], "{pass}"); + } +} + +/// A single conjunct that mixes a pushable and a nested column through `OR`, so +/// the whole predicate is one candidate. Refusing it and returning no filter at +/// all means the scan applies nothing, and the `FilterExec` is already gone. +#[tokio::test] +async fn sole_unevaluable_conjunct_is_still_applied() { + let dir = TempDir::new().unwrap(); + let parquet = dir.path().join("t.parquet"); + write_t(&parquet); + let ctx = liquid_ctx(&dir.path().join("cache")).await; + ctx.register_parquet( + "t", + parquet.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + let sql = "SELECT id FROM t WHERE id > 100 OR st.a = 3"; + for pass in ["cold", "warm"] { + assert_eq!(ids(&ctx, sql).await, vec![3], "{pass}"); + } +} + +/// A conjunct on a column that is not in the file schema. The table declares +/// `extra`, the file lacks it, so every row's `extra` is NULL and `extra = 3` is +/// never TRUE. A filter that dropped the conjunct would return all eight rows. +#[tokio::test] +async fn conjunct_on_column_outside_file_schema_is_still_applied() { + let dir = TempDir::new().unwrap(); + let table_dir = dir.path().join("t"); + std::fs::create_dir_all(&table_dir).unwrap(); + write_t(&table_dir.join("t.parquet")); + let ctx = liquid_ctx(&dir.path().join("cache")).await; + + // A declared schema wider than the file: `extra` exists in the table schema + // only. `st` is left out so this exercises the missing-column path alone. + let declared = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("extra", DataType::Int64, true), + ])); + let listing_options = + ListingOptions::new(Arc::new(ParquetFormat::default())).with_file_extension(".parquet"); + ctx.register_listing_table( + "t", + &ListingTableUrl::parse(table_dir.to_str().unwrap()).unwrap(), + listing_options, + Some(declared), + None, + ) + .await + .unwrap(); + + let sql = "SELECT id FROM t WHERE id >= 0 AND extra = 3"; + for pass in ["cold", "warm"] { + assert_eq!(ids(&ctx, sql).await, Vec::::new(), "{pass}"); + } +} From 007273ae033a66c15c2712c5f5ada8aca41d77f7 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 16:20:29 +0530 Subject: [PATCH 15/24] fix(reader): read the cache in cache-sized batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports hotdata-dev#15 (issue #13). `current_batch_id` indexes stored cache chunks, and both the cache read and the parquet fallback turn that id back into rows by multiplying it by the cache batch size. The reader walked the selection in windows of `datafusion.execution.batch_size` instead, so whenever a caller set that to anything other than the cache's size, batch id N named one range of rows to the reader and a different one to everything that resolved it. The scan then either ran off the end of the row group, or — with a LIMIT stopping it first — returned rows from the wrong offsets and succeeded. The fully-cached path misaligned the same way: a selection mask sized by the session batch size was applied to a chunk sized by the cache's, and `arrow::compute::filter` only rejects a mask longer than its target, so a shorter one silently took the wrong rows. Window by the cached row group's batch size, read once at the one place it is used. The caller is unaffected: `DataSourceExec` re-splits every source stream to the session batch size, which is what makes ignoring it inside the reader safe. The invariant this replaces was a `debug_assert_eq!`, so release builds carried the misalignment silently rather than failing. Two tests here are `#[ignore]`d against a separate upstream defect they uncovered: page-index pruning plus a pushed-down predicate panics in `boolean_buffer_and_then`. It reproduces with the session and cache batch sizes equal, so it is not a batch-size bug and is not this commit's to fix. --- .../src/tests/batch_size_alignment.rs | 309 ++++++++++++++++++ src/datafusion-local/src/tests/mod.rs | 1 + .../src/reader/runtime/liquid_cache_reader.rs | 16 +- 3 files changed, 320 insertions(+), 6 deletions(-) create mode 100644 src/datafusion-local/src/tests/batch_size_alignment.rs diff --git a/src/datafusion-local/src/tests/batch_size_alignment.rs b/src/datafusion-local/src/tests/batch_size_alignment.rs new file mode 100644 index 000000000..609f1baa4 --- /dev/null +++ b/src/datafusion-local/src/tests/batch_size_alignment.rs @@ -0,0 +1,309 @@ +//! Regression tests for issue #13: the reader indexes the cache by batch id, and +//! the parquet fallback turns that id back into rows using the *cache* batch size. +//! Reading at the session batch size (`datafusion.execution.batch_size`) instead +//! therefore addressed the wrong rows whenever the two differed — the scan either +//! ran off the end of the row group and panicked, or silently returned rows from +//! the wrong offsets. +//! +//! `LiquidCacheLocalBuilder::build` pins `execution.batch_size` to the cache batch +//! size, which is why the default configuration never hit this. A host that sets a +//! per-query batch size afterwards — as `SET datafusion.execution.batch_size` does +//! here — used to break the alignment. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array}; +use arrow::record_batch::RecordBatch; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::error::Result; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// One row group of `rows` rows, `id` ascending from 0. +fn write_single_row_group(path: &Path, rows: i64) { + write_parquet(path, rows, None, None); +} + +/// `rows` rows of ascending `id`, optionally split into row groups of +/// `row_group_rows` and data pages of `page_rows`. A small page size gives the +/// page index enough granularity to prune *within* a row group, which is what +/// produces a row selection that starts partway into it. +fn write_parquet(path: &Path, rows: i64, row_group_rows: Option, page_rows: Option) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from((0..rows).collect::>()))], + ) + .unwrap(); + let mut props = WriterProperties::builder(); + if let Some(n) = row_group_rows { + props = props.set_max_row_group_row_count(Some(n)); + } + if let Some(n) = page_rows { + props = props.set_data_page_row_count_limit(n); + } + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, Some(props.build())).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +/// Cache batch size stays at the 8192 default; the session batch size is then set +/// to `session_batch_size`, as a host sizing batches by row width would do. +async fn ctx_with_session_batch_size( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, +) -> Result { + ctx_with_session_batch_size_and_partitions(cache_dir, parquet_path, session_batch_size, 1).await +} + +async fn ctx_with_session_batch_size_and_partitions( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, + target_partitions: usize, +) -> Result { + std::fs::create_dir_all(cache_dir)?; + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = target_partitions; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir.to_path_buf()) + .build(config) + .await?; + assert_eq!(cache.batch_size(), 8192, "cache batch size should be 8192"); + + ctx.sql(&format!( + "SET datafusion.execution.batch_size = {session_batch_size}" + )) + .await? + .collect() + .await?; + + ctx.register_parquet( + "t", + parquet_path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await?; + Ok(ctx) +} + +async fn collect_batches(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql).await.unwrap().collect().await.unwrap() +} + +fn ids_of(batches: &[RecordBatch]) -> Vec { + let mut ids = Vec::new(); + for batch in batches { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..col.len()).map(|i| col.value(i))); + } + ids +} + +async fn collect_ids(ctx: &SessionContext, sql: &str) -> Vec { + ids_of(&collect_batches(ctx, sql).await) +} + +/// The reported failure: a 12000-row row group (fewer than `2 * 8192` rows) read +/// with a session batch size of 2145 used to drive the fallback past the end of +/// the row group on query batch 2, panicking with "parquet fallback ended before +/// batch 2". +#[tokio::test] +async fn narrowed_session_batch_size_reads_the_whole_row_group() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// Silent corruption, not an error: a LIMIT stopped the scan before the fallback +/// ran out of rows, so the query succeeded and returned rows 0-2144 followed by +/// rows 8192-10336 in place of rows 2145-4289. +#[tokio::test] +async fn narrowed_session_batch_size_with_limit_returns_aligned_rows() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let ids = collect_ids(&ctx, "SELECT id FROM t LIMIT 4290").await; + assert_eq!(ids, (0..4290).collect::>()); +} + +/// Issue #13 question 2: the fully-cached path misaligned as well. The cold run +/// populates cache batches 0 and 1; the warm run reads only from the cache and +/// used to emit rows 8192.. for query batch 1, because `read_arrow_array` applied +/// the 2145-bit selection to the 8192-row stored chunk and `arrow::compute::filter` +/// only rejects a mask *longer* than its target. +#[tokio::test] +async fn narrowed_session_batch_size_stays_aligned_on_the_fully_cached_path() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let sql = "SELECT id FROM t LIMIT 4290"; + let _cold = collect_ids(&ctx, sql).await; + let warm = collect_ids(&ctx, sql).await; + assert_eq!(warm, (0..4290).collect::>()); +} + +/// A session batch size *larger* than the cache batch size misaligned in the other +/// direction: the selection mask outran the stored chunk, and +/// `arrow::compute::filter` rejected it with "Filter predicate of length .. is +/// larger than target array of length ..". +#[tokio::test] +async fn widened_session_batch_size_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 20000) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// The scan now reads in cache-sized batches, but the caller must still receive +/// session-sized ones: `DataSourceExec::execute` wraps every source stream in a +/// `BatchSplitStream` sized by the session config. This is what makes ignoring the +/// session batch size inside the reader safe. +#[tokio::test] +async fn session_batch_size_still_bounds_the_batches_the_caller_receives() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let batches = collect_batches(&ctx, "SELECT id FROM t").await; + let oversized: Vec<_> = batches + .iter() + .map(|b| b.num_rows()) + .filter(|rows| *rows > 2145) + .collect(); + assert!( + oversized.is_empty(), + "batches exceeded the session batch size: {oversized:?}" + ); + assert_eq!(ids_of(&batches), (0..12000).collect::>()); + } +} + +/// `current_batch_id` restarts at 0 for each row group, so a file with several of +/// them exercises the mapping repeatedly rather than once. +#[tokio::test] +async fn multi_row_group_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000), None); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..30000).collect::>()); + } +} + +/// The subtle case: page-index pruning leaves a row selection that starts partway +/// *into* a row group, so the first window is mostly `RowSelector::skip`. Batch id +/// 0 must still mean physical rows `0..cache_batch_size` of that row group — +/// `take_next_batch` counts skipped rows toward the window, which is what keeps +/// the id aligned with the stored chunk. A pushed-down predicate also routes +/// through `evaluate_selection_with_predicate(current_batch_id, ..)`, so this +/// covers the filter path's use of the same id. +#[tokio::test] +#[ignore = "upstream bug, not a batch-size one: page-index pruning plus a pushed-down \ + predicate panics in boolean_buffer_and_then (left.count_set_bits()=904, \ + right.len()=5000). Reproduces with the session and cache batch sizes equal, \ + so it is independent of this file's subject."] +async fn pruned_filtered_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000), Some(500)); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t WHERE id >= 9500 AND id < 10500").await; + assert_eq!(ids, (9500..10500).collect::>()); + } + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t WHERE id >= 25000").await; + assert_eq!(ids, (25000..30000).collect::>()); + } +} + +/// Several partitions read the same file concurrently, each with its own reader +/// and its own `current_batch_id` sequence. +#[tokio::test] +async fn multi_partition_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000), None); + let ctx = ctx_with_session_batch_size_and_partitions( + &tmp.path().join("cache"), + &parquet_path, + 2145, + 4, + ) + .await + .unwrap(); + + for _ in 0..2 { + let mut ids = collect_ids(&ctx, "SELECT id FROM t").await; + ids.sort_unstable(); + assert_eq!(ids, (0..30000).collect::>()); + } +} + +/// Isolation probe: same pruned+filtered scan, but with the session batch size +/// left equal to the cache's, so batch windowing is not involved at all. +#[tokio::test] +#[ignore = "upstream bug, not a batch-size one: page-index pruning plus a pushed-down \ + predicate panics in boolean_buffer_and_then (left.count_set_bits()=904, \ + right.len()=5000). Reproduces with the session and cache batch sizes equal, \ + so it is independent of this file's subject."] +async fn pruned_filtered_scan_at_matching_batch_size() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000), Some(500)); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 8192) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t WHERE id >= 9500 AND id < 10500").await; + assert_eq!(ids, (9500..10500).collect::>()); + } +} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index e57b302c0..15d669466 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -19,6 +19,7 @@ use datafusion::{ }; use crate::LiquidCacheLocalBuilder; +mod batch_size_alignment; mod date_optimizer; mod filter_limit; mod nested_filter; diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 84a03b55a..9e6ded1f0 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -92,13 +92,17 @@ pub(crate) struct ParquetFallback { impl LiquidCacheReader { pub(crate) fn new(config: LiquidCacheReaderConfig) -> Self { - debug_assert_eq!( - config.batch_size, - config.cached_row_group.batch_size(), - "DataFusion and LiquidCache batch sizes must agree" - ); + // The selection is walked in cache-sized windows, not session-sized + // ones. `current_batch_id` indexes stored chunks, and both the cache + // read and the parquet fallback turn that id back into rows with the + // cache batch size; windowing at `datafusion.execution.batch_size` + // would address different rows than the id names whenever a caller + // sets it to anything other than the cache's own size. The caller is + // unaffected either way: `DataSourceExec` re-splits every source + // stream to the session batch size. + let batch_size = config.cached_row_group.batch_size(); let inner = LiquidCacheReaderInner::new( - config.batch_size, + batch_size, config.selection, config.cached_row_group, config.projection_columns, From 5fe914baf0a5113036f6aa266db0f65640fef198 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 16:22:06 +0530 Subject: [PATCH 16/24] fix(optimizer): keep scans that need virtual columns on ParquetSource MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports hotdata-dev#41. The liquid read path has no notion of DataFusion's virtual columns. It carries the `TableSchema` across faithfully but never produces one, so a scan whose projection includes a virtual column reads back a batch that simply lacks it, and resolving the column then fails with an Arrow schema error naming only the file's own fields. A predicate over one fails the same way: the reader rewrites it against the logical and physical file schemas, neither of which holds the column. Decline the swap for such scans, leaving them on `ParquetSource`, which derives virtual columns from the parquet reader. The guard keys on the virtual columns the scan actually reads — projection unioned with the pushed-down filter — not on what the table declares, so a provider that puts a row-position column on every table keeps the cache for the queries that never touch it. No projection at all is the one broad case, and it is not a guess: the scan then reads the whole table schema, virtual columns included. Positional reads are what reaches here in practice. Delete filtering and row lineage both project a reader-produced physical row position, an absolute index into the file, and a plausible but shifted position would associate a delete with the wrong row. Declining is the sound answer until the liquid reader can generate positions itself. --- src/datafusion/src/optimizers/mod.rs | 195 ++++++++++++++++++++++++++- 1 file changed, 194 insertions(+), 1 deletion(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index dd20d02a4..676b1bc5b 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -19,8 +19,9 @@ use datafusion::{ listing::PartitionedFile, physical_plan::{FileScanConfig, FileSource, ParquetSource}, source::DataSource, + table_schema::TableSchema, }, - physical_expr::utils::collect_columns, + physical_expr::{PhysicalExpr, projection::ProjectionExprs, utils::collect_columns}, physical_optimizer::{PhysicalOptimizerRule, pruning::PruningPredicateBuilder}, physical_plan::ExecutionPlan, }; @@ -536,6 +537,73 @@ fn file_required_bytes( (sum, false) } +/// The virtual columns this scan reads that the liquid path cannot produce. +/// +/// The liquid read path has no notion of DataFusion's virtual columns. It carries +/// the [`TableSchema`] across faithfully but never produces one, so a scan that +/// reads a virtual column gets back a batch which simply lacks it, and a predicate +/// over one cannot be rewritten against the file schemas either. Such a scan has +/// to stay on `ParquetSource`, which derives virtual columns from the parquet +/// reader. +/// +/// Declining a scan costs it the cache, so this stays narrow: a virtual column the +/// scan actually reads, not the mere presence of one on the table. A provider that +/// declares a row-position column on every table keeps the cache for the queries +/// that never project one. No projection at all is the one broad case, and it is +/// not a guess — the scan then reads the whole table schema, virtual columns +/// included. +/// +/// Positional reads are what reaches here: applying positional deletes, and row +/// lineage, both project a reader-produced physical row position, which is an +/// absolute index into the file. A plausible but shifted position would associate +/// a delete with the wrong row, so declining is the sound answer while the liquid +/// reader cannot generate positions itself. +fn unproducible_virtual_columns( + table_schema: &TableSchema, + projection: Option<&ProjectionExprs>, + filter: Option<&Arc>, +) -> Option { + let virtual_columns = table_schema.virtual_columns(); + if virtual_columns.is_empty() { + return None; + } + + let mut needed: Vec<&str> = Vec::new(); + match projection { + None => needed.extend(virtual_columns.iter().map(|field| field.name().as_str())), + Some(projection) => { + let mut read: HashSet = projection + .expr_iter() + .flat_map(|expr| collect_columns(&expr)) + .map(|column| column.name().to_string()) + .collect(); + // The pushed-down predicate is rewritten against the file schemas in + // the reader too, so a conjunct over a virtual column fails exactly as + // a projection over one does. The row filter's own check cannot catch + // it: that resolves against the table schema, which does hold the + // virtual columns. + if let Some(filter) = filter { + read.extend( + collect_columns(filter) + .into_iter() + .map(|column| column.name().to_string()), + ); + } + needed.extend( + virtual_columns + .iter() + .map(|field| field.name().as_str()) + .filter(|name| read.contains(*name)), + ); + } + } + + if needed.is_empty() { + return None; + } + Some(needed.join("`, `")) +} + /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent /// node backed by [`LiquidParquetSource`] carrying `hints`. fn convert_parquet_scan( @@ -548,6 +616,21 @@ fn convert_parquet_scan( let (file_scan_config, parquet_source) = data_source_exec.downcast_to_file_source::()?; + let pushed_filter = parquet_source.filter(); + if let Some(names) = unproducible_virtual_columns( + parquet_source.table_schema(), + parquet_source.projection(), + pushed_filter.as_ref(), + ) { + // At info, like the admission gate's BYPASS line: this silently turns the + // cache off for a scan, and the only symptom is that queries stop getting + // faster. + log::info!( + "liquid_cache scan BYPASS: the read path cannot produce virtual column(s) `{names}`" + ); + return None; + } + let new_source = LiquidParquetSource::from_parquet_source(parquet_source.clone(), cache.clone()) .with_lineages(Arc::new(hints)) @@ -648,6 +731,116 @@ mod tests { .unwrap() } + /// Declining a scan costs it the cache, so the guard must key on what the scan + /// reads, not on what the table declares. A row-position column present on the + /// table but absent from the projection keeps the cache; no projection at all + /// reads the whole table schema and does not. + #[test] + fn only_a_virtual_column_the_scan_reads_costs_the_cache() { + use arrow_schema::Fields; + use datafusion::physical_expr::expressions::{BinaryExpr, col, lit}; + use datafusion::physical_expr::projection::ProjectionExpr; + + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("val", DataType::Int64, true), + ])); + let row_pos = Field::new("__ducklake_row_pos", DataType::Int64, true); + + // No virtual columns at all: nothing to refuse, whatever the projection. + let plain = TableSchema::builder(Arc::clone(&file_schema)).build(); + assert_eq!(unproducible_virtual_columns(&plain, None, None), None); + + let positional = TableSchema::builder(Arc::clone(&file_schema)) + .with_virtual_columns(Fields::from(vec![row_pos.clone()])) + .build(); + let full = positional.table_schema(); + + // Declared but not read: still cached. + let only_id = + ProjectionExprs::new(vec![ProjectionExpr::new(col("id", full).unwrap(), "id")]); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&only_id), None), + None + ); + + // Read: refused, and named. + let reads_pos = ProjectionExprs::new(vec![ProjectionExpr::new( + col("__ducklake_row_pos", full).unwrap(), + "__ducklake_row_pos", + )]); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&reads_pos), None).as_deref(), + Some("__ducklake_row_pos") + ); + + // No projection reads the whole table schema, virtual columns included. + assert_eq!( + unproducible_virtual_columns(&positional, None, None).as_deref(), + Some("__ducklake_row_pos") + ); + + // Read only by the pushed-down predicate: refused too. The reader rewrites + // the predicate against the file schemas as well, and the row filter's own + // check passes it because that resolves against the table schema. + let pos_predicate: Arc = Arc::new(BinaryExpr::new( + col("__ducklake_row_pos", full).unwrap(), + Operator::Gt, + lit(0i64), + )); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&only_id), Some(&pos_predicate)) + .as_deref(), + Some("__ducklake_row_pos") + ); + } + + /// The guard at its call site, with an ordinary scan as a positive control, so + /// deleting it from `convert_parquet_scan` fails a test instead of silently + /// restoring a plan that cannot execute. + #[tokio::test] + async fn a_scan_reading_a_virtual_column_stays_on_parquet_source() { + use arrow_schema::Fields; + use datafusion::datasource::physical_plan::FileScanConfigBuilder; + use datafusion::execution::object_store::ObjectStoreUrl; + + let file_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + + let scan = |table_schema: TableSchema| -> Arc { + let source = Arc::new(ParquetSource::new(table_schema)) as Arc; + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(PartitionedFile::new("t.parquet", 16)) + .build(); + Arc::new(DataSourceExec::new(Arc::new(config))) + }; + + let tmp_dir = tempfile::tempdir().unwrap(); + let cache = make_cache(tmp_dir.path()).await; + + // Positive control: an ordinary scan is still handed to the cache. + let plain = scan(TableSchema::builder(Arc::clone(&file_schema)).build()); + assert!( + convert_parquet_scan(&plain, &cache, ColumnLineages::default(), true).is_some(), + "an ordinary scan must still convert to the liquid source" + ); + + // `ParquetSource::new` projects the whole table schema, so the positional + // scan reads the row-position column and cannot be served. + let positional = scan( + TableSchema::builder(file_schema) + .with_virtual_columns(Fields::from(vec![Field::new( + "__ducklake_row_pos", + DataType::Int64, + true, + )])) + .build(), + ); + assert!( + convert_parquet_scan(&positional, &cache, ColumnLineages::default(), true).is_none(), + "a scan reading a virtual column must stay on ParquetSource" + ); + } + /// The admission gate bypasses a scan whose estimated footprint exceeds the /// budget (large expansion here forces that), and caches one that fits. #[tokio::test] From 0ad86223e74c1ebde7ee9f97fbcfc9ff09cb7a84 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 16:54:03 +0530 Subject: [PATCH 17/24] fix(cache): serve entries only to their own file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports hotdata-dev#48, and with it the index-slot half of #45. A cache key packs the file id into 16 bits, so the 65,537th distinct file a process registers aliases the first. Entries recorded nothing about where they came from and the read API took no expected identity, so an aliased lookup returned the other file's data: a panic when the column types differed, silently wrong rows when they matched. Upstream guards the narrowing with a `debug_assert!` only, so release builds truncate and carry on. Record the unnarrowed file id alongside each entry and compare it on every read. A mismatch reads as a miss, so the caller re-reads from its source and gets correct data, and it is counted. Store objects are keyed by entry id *and* identity, so a write still in flight for one owner cannot overwrite bytes a later owner's entry names. Lease file ids instead of assigning them permanently. A lease is held by the file handle and by every row group and column derived from it, and returns to a pool when the last holder drops, so the id space tracks the files being read rather than every file ever read. Entries deliberately hold no lease: an id reused while its old entries are resident leaves them unreachable rather than readable, which is what lets the release stay out of index removal, where it would deadlock against `reset`. Two kinds of write, because they differ: a caller storing its own data takes a contested key over, while maintenance rewriting an entry it read earlier — evict, hydrate, spill, flush — lands only while the key still holds the identity it read, and is dropped otherwise. Adopting whatever holds the key by then would relabel one file's data as another's, and the new owner would read those rows as a hit. Two departures from the fork's version, both forced by upstream's newer code: The pool hands out a fresh id while the key's 16-bit file field has room and recycles only once it is exhausted. Recycling eagerly is correct but costs the previous holder every entry it cached, and upstream's join lineage test caught it: the two sides of a join took the same id and overwrote each other's lineage while 65,000 ids sat unused. The index slot also carries #45's fix — the payload is taken out of the slot when the index gives an entry up, so crossbeam-epoch's deferred destruction reclaims an empty shell rather than holding multi-megabyte arrays past the budget that counts them. `leased_file_ids`, `file_ids_over_key_width` and `identity_mismatches` expose the three counters that say whether any of this is being hit. --- examples/core.rs | 12 +- src/core/README.md | 16 +- src/core/src/cache/builders.rs | 33 +- src/core/src/cache/core.rs | 310 ++++++++---- src/core/src/cache/index.rs | 357 +++++++++++-- src/core/src/cache/io_context.rs | 17 +- src/core/src/cache/observer/stats.rs | 6 + src/core/src/cache/tests/policies.rs | 28 +- src/core/study/cache_storage.rs | 4 +- src/datafusion/src/cache/column.rs | 27 +- src/datafusion/src/cache/file_id.rs | 473 ++++++++++++++++++ src/datafusion/src/cache/mod.rs | 77 ++- src/datafusion/src/cache/stats.rs | 83 +-- .../src/reader/plantime/morselizer.rs | 5 +- .../src/reader/runtime/liquid_cache_reader.rs | 4 - src/datafusion/src/reader/runtime/morsel.rs | 2 - 16 files changed, 1212 insertions(+), 242 deletions(-) create mode 100644 src/datafusion/src/cache/file_id.rs diff --git a/examples/core.rs b/examples/core.rs index 1df164c7d..2d84688e3 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,13 +18,21 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); + // The identity names which source this data came from. A cache key packs + // its fields into fixed widths, so two sources can compute one key; the + // identity is what keeps a read from being served the other's data. One + // source here, so any constant will do. + let identity = 0; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); - storage.insert(entry_id, arrow_array.clone()).await.unwrap(); + storage + .insert(entry_id, identity, arrow_array.clone()) + .await + .unwrap(); // Move data to disk so the read demonstrates async I/O storage.flush_all_to_disk().await.unwrap(); - let retrieved = storage.get(&entry_id).await.unwrap(); + let retrieved = storage.get(&entry_id, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); Ok(()) diff --git a/src/core/README.md b/src/core/README.md index e0e752125..4352a42dc 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -25,9 +25,13 @@ let entry_id = EntryID::from(42); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); // Insert once; replacement/placement is handled by the cache policy -storage.insert(entry_id, arrow_array.clone()).await; +// `0` is the file identity: the unnarrowed name of the source this data +// belongs to. Cache keys pack their fields into fixed widths, so two sources +// can compute one key; the identity is what keeps a read from being served +// the other's data. One source here, so any constant will do. +storage.insert(entry_id, 0, arrow_array.clone()).await; -assert!(storage.contains(&entry_id)); +assert!(storage.is_cached(&entry_id, 0)); }); ``` @@ -43,13 +47,13 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(7); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, 0, arrow_array.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; // Read asynchronously -let retrieved = storage.get(&entry_id).await.unwrap(); +let retrieved = storage.get(&entry_id, 0).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); }); ``` @@ -74,7 +78,7 @@ let entry_id = EntryID::from(8); let data = Arc::new(StringArray::from(vec![ Some("apple"), Some("banana"), None, Some("apple"), Some("cherry"), ])); -storage.insert(entry_id, data.clone()).await; +storage.insert(entry_id, 0, data.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; @@ -93,7 +97,7 @@ let liquid_expr = liquid_cache::cache::LiquidExpr::try_new( // Read with predicate pushdown let mask = storage - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, 0, &liquid_expr) .with_selection(&selection) .await .unwrap(); diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 9bb4a8430..46d98e0e2 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -175,16 +175,25 @@ pub fn default_max_memory_bytes() -> usize { pub struct Insert<'a> { pub(super) storage: &'a Arc, pub(super) entry_id: EntryID, + /// The unnarrowed name of the source this data belongs to, recorded with + /// the entry so a key collision cannot serve it to anyone else. + pub(super) identity: u64, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, pub(super) lineage: Option>, } impl<'a> Insert<'a> { - pub(super) fn new(storage: &'a Arc, entry_id: EntryID, batch: ArrayRef) -> Self { + pub(super) fn new( + storage: &'a Arc, + entry_id: EntryID, + identity: u64, + batch: ArrayRef, + ) -> Self { Self { storage, entry_id, + identity, batch, skip_gc: false, lineage: None, @@ -213,7 +222,13 @@ impl<'a> Insert<'a> { self.storage.add_lineage(&self.entry_id, lineage); } let batch = CacheEntry::memory_arrow(batch); - self.storage.insert_inner(self.entry_id, batch).await + self.storage + .insert_inner( + self.entry_id, + crate::cache::index::WriteIdentity::Owned(self.identity), + batch, + ) + .await } } @@ -231,15 +246,17 @@ impl<'a> IntoFuture for Insert<'a> { pub struct Get<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) selection: Option<&'a BooleanBuffer>, pub(super) expression_hint: Option>, } impl<'a> Get<'a> { - pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID) -> Self { + pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID, identity: u64) -> Self { Self { storage, entry_id, + identity, selection: None, expression_hint: None, } @@ -271,6 +288,7 @@ impl<'a> Get<'a> { self.storage .read_arrow_array( self.entry_id, + self.identity, self.selection, self.expression_hint.as_deref(), ) @@ -326,6 +344,7 @@ fn maybe_gc_view_arrays(array: &ArrayRef) -> Option { pub struct EvaluatePredicate<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) predicate: &'a LiquidExpr, pub(super) selection: Option<&'a BooleanBuffer>, } @@ -334,11 +353,13 @@ impl<'a> EvaluatePredicate<'a> { pub(super) fn new( storage: &'a LiquidCache, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> Self { Self { storage, entry_id, + identity, predicate, selection: None, } @@ -353,7 +374,7 @@ impl<'a> EvaluatePredicate<'a> { /// Evaluate the predicate against the cached data. pub async fn read(self) -> Option { self.storage - .eval_predicate_internal(self.entry_id, self.selection, self.predicate) + .eval_predicate_internal(self.entry_id, self.identity, self.selection, self.predicate) .await } } @@ -451,9 +472,9 @@ mod tests { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(123usize); - cache.insert(entry_id, root.clone()).await.unwrap(); + cache.insert(entry_id, 0, root.clone()).await.unwrap(); - let stored = cache.get(&entry_id).await.expect("array present"); + let stored = cache.get(&entry_id, 0).await.expect("array present"); let post_size = stored.get_array_memory_size(); // GC should have compacted the view arrays, reducing memory footprint. diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index d007d3164..38a2dc705 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -17,7 +17,11 @@ use super::{ }; use crate::cache::policies::{EvictionOutcome, EvictionPolicy}; use crate::cache::utils::arrow_to_bytes; -use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; +use crate::cache::{ + CacheExpression, LiquidExpr, + index::{ArtIndex, WriteIdentity}, + utils::EntryID, +}; use crate::cache::{CacheFull, CacheStats, EventTrace}; use crate::sync::Arc; @@ -36,10 +40,13 @@ use crate::sync::Arc; /// /// let entry_id = EntryID::from(0); /// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); -/// storage.insert(entry_id, arrow_array.clone()).await; +/// // `0` is the file identity: cache keys pack their fields into fixed +/// // widths, so two sources can compute one key, and the identity is what +/// // keeps a read from being served the other's data. +/// storage.insert(entry_id, 0, arrow_array.clone()).await; /// /// // Get the arrow array back asynchronously -/// let retrieved = storage.get(&entry_id).await.unwrap(); +/// let retrieved = storage.get(&entry_id, 0).await.unwrap(); /// assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); /// }); /// ``` @@ -80,7 +87,7 @@ impl LiquidCache { let mut memory_arrow_bytes = 0usize; let mut memory_liquid_bytes = 0usize; - self.index.for_each(|_, batch| match batch { + self.index.for_each(|_, _, batch| match batch { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -99,6 +106,7 @@ impl LiquidCache { CacheStats { total_entries, + identity_mismatches: self.index.identity_mismatches(), memory_arrow_entries, memory_liquid_entries, disk_liquid_entries, @@ -117,28 +125,32 @@ impl LiquidCache { pub fn insert<'a>( self: &'a Arc, entry_id: EntryID, + identity: u64, batch_to_cache: ArrayRef, ) -> Insert<'a> { - Insert::new(self, entry_id, batch_to_cache) + Insert::new(self, entry_id, identity, batch_to_cache) } /// Create a [`Get`] builder for the provided entry. - pub fn get<'a>(&'a self, entry_id: &'a EntryID) -> Get<'a> { - Get::new(self, entry_id) + pub fn get<'a>(&'a self, entry_id: &'a EntryID, identity: u64) -> Get<'a> { + Get::new(self, entry_id, identity) } /// Create an [`EvaluatePredicate`] builder for evaluating predicates on cached data. pub fn eval_predicate<'a>( &'a self, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> EvaluatePredicate<'a> { - EvaluatePredicate::new(self, entry_id, predicate) + EvaluatePredicate::new(self, entry_id, identity, predicate) } /// Prefetch an entry into a memory-form snapshot without recording an access. - pub async fn prefetch(&self, entry_id: &EntryID) -> PrefetchResult { - let Some(entry) = self.index.get(entry_id) else { + pub async fn prefetch(&self, entry_id: &EntryID, identity: u64) -> PrefetchResult { + // Checked, not raw: a prefetch hands a snapshot to a caller, so an + // aliased key must read as absent rather than as someone else's rows. + let Some(entry) = self.index.get_checked(entry_id, identity) else { return PrefetchResult::Absent; }; match entry.as_ref() { @@ -146,19 +158,31 @@ impl LiquidCache { PrefetchResult::Snapshot(entry) } disk @ CacheEntry::DiskArrow { .. } => { - let Some(array) = self.read_disk_arrow_array(entry_id).await else { + let Some(array) = self.read_disk_arrow_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Arrow(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Arrow(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_arrow(array))) } disk @ CacheEntry::DiskLiquid { .. } => { - let Some(array) = self.read_disk_liquid_array(entry_id).await else { + let Some(array) = self.read_disk_liquid_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Liquid(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Liquid(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_liquid(array))) } } @@ -169,19 +193,26 @@ impl LiquidCache { pub async fn try_read_liquid( &self, entry_id: &EntryID, + identity: u64, ) -> Option { self.observer.on_try_read_liquid(); self.trace(InternalEvent::TryReadLiquid { entry: *entry_id }); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); match batch.as_ref() { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; Some(liquid) } CacheEntry::DiskArrow { .. } | CacheEntry::MemoryArrow(_) => None, @@ -191,7 +222,7 @@ impl LiquidCache { /// Iterate over all entries in the cache. /// No guarantees are made about the order of the entries. /// Isolation level: read-committed - pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { self.index.for_each(&mut f); } @@ -201,9 +232,13 @@ impl LiquidCache { self.budget.reset_usage(); } - /// Check whether the cache contains a batch. - pub fn contains(&self, entry_id: &EntryID) -> bool { - self.index.contains(entry_id) + /// Check whether the cache holds this entry *for this identity*. + /// + /// A key held by another identity reads as absent: it belongs to a source + /// that can no longer read it, and serving it here would hand one caller + /// another's rows. + pub fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.is_cached(entry_id, identity) } /// Get the config of the cache. @@ -234,18 +269,22 @@ impl LiquidCache { /// Flush all entries to disk. pub async fn flush_all_to_disk(&self) -> Result<(), CacheFull> { let mut entires = Vec::new(); - self.for_each_entry(|entry_id, batch| { - entires.push((*entry_id, batch.clone())); + self.for_each_entry(|entry_id, identity, batch| { + entires.push((*entry_id, identity, batch.clone())); }); - for (entry_id, batch) in entires { + for (entry_id, flush_identity, batch) in entires { match &batch { CacheEntry::MemoryArrow(array) => { let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); let disk_bytes = bytes.len(); - match self.write_batch_to_disk(entry_id, &batch, bytes).await { + match self + .write_batch_to_disk(entry_id, flush_identity, &batch, bytes) + .await + { Ok(()) => { self.try_insert( entry_id, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -257,12 +296,18 @@ impl LiquidCache { let liquid_bytes = liquid_array.to_bytes(); let disk_bytes = liquid_bytes.len(); match self - .write_batch_to_disk(entry_id, &batch, Bytes::from(liquid_bytes)) + .write_batch_to_disk( + entry_id, + flush_identity, + &batch, + Bytes::from(liquid_bytes), + ) .await { Ok(()) => { self.try_insert( entry_id, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), disk_bytes, @@ -287,6 +332,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: CacheEntry, ) -> Result { match &batch { @@ -302,7 +348,7 @@ impl LiquidCache { unreachable!("memory Arrow eviction cannot remove entry"); }; if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(entry_id, &new_batch, bytes_to_write) + self.write_batch_to_disk(entry_id, identity, &new_batch, bytes_to_write) .await?; } Ok(new_batch) @@ -310,7 +356,7 @@ impl LiquidCache { CacheEntry::MemoryLiquid(liquid_array) => { let liquid_bytes = Bytes::from(liquid_array.to_bytes()); let disk_bytes = liquid_bytes.len(); - self.write_batch_to_disk(entry_id, &batch, liquid_bytes) + self.write_batch_to_disk(entry_id, identity, &batch, liquid_bytes) .await?; Ok(CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), @@ -327,10 +373,11 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, + identity: WriteIdentity, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { loop { - let Err(not_inserted) = self.try_insert(entry_id, batch_to_cache) else { + let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { return Ok(()); }; self.trace(InternalEvent::InsertFailed { @@ -344,7 +391,7 @@ impl LiquidCache { // this can happen if the entry to be inserted is too large, in that case, // we write it to disk let on_disk_batch = self - .write_in_memory_batch_to_disk(entry_id, not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity.value(), not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -389,7 +436,12 @@ impl LiquidCache { } } - fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + fn try_insert( + &self, + entry_id: EntryID, + identity: WriteIdentity, + to_insert: CacheEntry, + ) -> Result<(), CacheEntry> { let new_memory_size = to_insert.memory_usage_bytes(); let cached_batch_type = if let Some(entry) = self.index.get(&entry_id) { let old_memory_size = entry.memory_usage_bytes(); @@ -401,14 +453,25 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + // A rewrite whose key changed hands since it was read is dropped by + // the index. Give the reservation back rather than counting memory + // for an entry that was never stored. + if !self.index.insert(&entry_id, identity, to_insert) { + self.budget + .try_update_memory_usage(new_memory_size, old_memory_size) + .ok(); + return Ok(()); + } batch_type } else { if self.budget.try_reserve_memory(new_memory_size).is_err() { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + if !self.index.insert(&entry_id, identity, to_insert) { + self.budget.try_update_memory_usage(new_memory_size, 0).ok(); + return Ok(()); + } batch_type }; @@ -439,7 +502,7 @@ impl LiquidCache { self.cache_policy.notify_remove(&entry_id); } - async fn remove_disk_entry(&self, entry_id: EntryID) { + async fn remove_disk_entry(&self, entry_id: EntryID, removed_identity: u64) { let Some(removed) = self.index.remove(&entry_id) else { return; }; @@ -449,7 +512,7 @@ impl LiquidCache { _ => panic!("remove_disk_entry called for non-disk entry"), }; self.store - .remove(&entry_id_to_key(&entry_id)) + .remove(&entry_id_to_key(&entry_id, removed_identity)) .await .expect("disk remove failed"); self.budget.release_disk(disk_bytes); @@ -496,7 +559,10 @@ impl LiquidCache { } async fn evict_victim_inner(&self, victim: EntryID) -> Result<(), CacheFull> { - let Some(mut victim_entry) = self.index.get(&victim) else { + // Read the identity alongside the entry: everything this loop writes + // back is a rewrite of what it just read, and must be dropped rather + // than relabelled if the key changes hands meanwhile. + let Some((identity, mut victim_entry)) = self.index.get_with_identity(&victim) else { return Ok(()); }; self.trace(InternalEvent::EvictionVictim { entry: victim }); @@ -512,10 +578,10 @@ impl LiquidCache { bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(victim, &new_batch, bytes_to_write) + self.write_batch_to_disk(victim, identity, &new_batch, bytes_to_write) .await?; } - match self.try_insert(victim, new_batch) { + match self.try_insert(victim, WriteIdentity::Rewrite(identity), new_batch) { Ok(()) => { break; } @@ -525,7 +591,7 @@ impl LiquidCache { } } EvictionOutcome::Remove => { - self.remove_disk_entry(victim).await; + self.remove_disk_entry(victim, identity).await; break; } } @@ -536,6 +602,7 @@ impl LiquidCache { async fn maybe_hydrate( &self, entry_id: &EntryID, + identity: u64, cached: &CacheEntry, materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, @@ -553,21 +620,24 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, new_entry).await; + let _ = self + .insert_inner(*entry_id, WriteIdentity::Rewrite(identity), new_entry) + .await; } } pub(crate) async fn read_arrow_array( &self, entry_id: &EntryID, + identity: u64, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.read_entry_inner(entry_id, batch.as_ref(), selection, expression) + self.read_entry_inner(entry_id, identity, batch.as_ref(), selection, expression) .await } @@ -575,18 +645,20 @@ impl LiquidCache { pub async fn read_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - self.read_entry_inner(entry_id, entry, selection, expression) + self.read_entry_inner(entry_id, identity, entry, selection, expression) .await } async fn read_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, @@ -613,7 +685,7 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(entry, entry_id, expression, selection) + self.read_disk_array(entry, entry_id, identity, expression, selection) .await } } @@ -623,6 +695,7 @@ impl LiquidCache { &self, entry: &CacheEntry, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -633,9 +706,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await?; + let full_array = self.read_disk_arrow_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Arrow(&full_array), expression, @@ -655,9 +729,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await?; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Liquid(&liquid), expression, @@ -675,6 +750,7 @@ impl LiquidCache { async fn write_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { @@ -688,7 +764,11 @@ impl LiquidCache { return Err(CacheFull); } for victim in victims { - self.remove_disk_entry(victim).await; + // Each victim's object is addressed by the identity that wrote + // it, so look that up rather than assuming this writer's. + if let Some((victim_identity, _)) = self.index.get_with_identity(&victim) { + self.remove_disk_entry(victim, victim_identity).await; + } } } self.trace(InternalEvent::IoWrite { @@ -697,14 +777,14 @@ impl LiquidCache { bytes: len, }); self.store - .put(entry_id_to_key(&entry_id), bytes.to_vec()) + .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) .await .expect("write failed"); Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> Option { + let bytes = match self.store.get(&entry_id_to_key(entry_id, identity)).await { Ok(bytes) => bytes, Err(t4::Error::NotFound) => return None, Err(error) => panic!("read failed: {error}"), @@ -725,8 +805,9 @@ impl LiquidCache { async fn read_disk_liquid_array( &self, entry_id: &EntryID, + identity: u64, ) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { + let bytes = match self.store.get(&entry_id_to_key(entry_id, identity)).await { Ok(bytes) => bytes, Err(t4::Error::NotFound) => return None, Err(error) => panic!("read failed: {error}"), @@ -743,33 +824,42 @@ impl LiquidCache { pub(crate) async fn eval_predicate_internal( &self, entry_id: &EntryID, + identity: u64, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.eval_predicate_on_entry_inner(entry_id, batch.as_ref(), selection_opt, predicate) - .await + self.eval_predicate_on_entry_inner( + entry_id, + identity, + batch.as_ref(), + selection_opt, + predicate, + ) + .await } /// Evaluate a predicate on an already-looked-up cache entry. pub async fn eval_predicate_on_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - self.eval_predicate_on_entry_inner(entry_id, entry, selection_opt, predicate) + self.eval_predicate_on_entry_inner(entry_id, identity, entry, selection_opt, predicate) .await } async fn eval_predicate_on_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, @@ -793,9 +883,15 @@ impl LiquidCache { Some(self.eval_predicate_on_array(filtered, predicate)) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) - .await; + let array = self.read_disk_arrow_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Arrow(&array), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(array.len())); @@ -815,9 +911,15 @@ impl LiquidCache { Some(array.try_eval_predicate(predicate, selection)) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(liquid.len())); @@ -897,7 +999,10 @@ mod tests { let entry_id1: EntryID = EntryID::from(1); let array1 = create_test_array(100); let size1 = array1.memory_usage_bytes(); - store.insert_inner(entry_id1, array1).await.unwrap(); + store + .insert_inner(entry_id1, WriteIdentity::Owned(0), array1) + .await + .unwrap(); // Verify budget usage and data correctness assert_eq!(store.budget.memory_usage_bytes(), size1); @@ -910,13 +1015,19 @@ mod tests { let entry_id2: EntryID = EntryID::from(2); let array2 = create_test_array(200); let size2 = array2.memory_usage_bytes(); - store.insert_inner(entry_id2, array2).await.unwrap(); + store + .insert_inner(entry_id2, WriteIdentity::Owned(0), array2) + .await + .unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size1 + size2); let array3 = create_test_array(150); let size3 = array3.memory_usage_bytes(); - store.insert_inner(entry_id1, array3).await.unwrap(); + store + .insert_inner(entry_id1, WriteIdentity::Owned(0), array3) + .await + .unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -936,7 +1047,7 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, create_test_array(800)) + .insert_inner(entry_id1, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -945,7 +1056,7 @@ mod tests { } store - .insert_inner(entry_id2, create_test_array(800)) + .insert_inner(entry_id2, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -995,7 +1106,7 @@ mod tests { let unique_id = thread_id * ops_per_thread + i; let entry_id: EntryID = EntryID::from(unique_id); let array = create_test_arrow_array(100); - store.insert(entry_id, array).await.unwrap(); + store.insert(entry_id, 0, array).await.unwrap(); } }); })); @@ -1029,8 +1140,14 @@ mod tests { // Insert two small batches let arr1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); let arr2: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128)); - storage.insert(EntryID::from(1usize), arr1).await.unwrap(); - storage.insert(EntryID::from(2usize), arr2).await.unwrap(); + storage + .insert(EntryID::from(1usize), 0, arr1) + .await + .unwrap(); + storage + .insert(EntryID::from(2usize), 0, arr2) + .await + .unwrap(); // Stats after insert: 2 entries, memory usage > 0, disk usage == 0 let s = storage.stats(); @@ -1054,14 +1171,14 @@ mod tests { let entry_id = EntryID::from(321usize); let array = create_test_arrow_array(8); - store.insert(entry_id, array.clone()).await.unwrap(); + store.insert(entry_id, 0, array.clone()).await.unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); assert!(matches!(entry.as_ref(), CacheEntry::DiskArrow { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); @@ -1079,11 +1196,14 @@ mod tests { .await; let id = EntryID::from(320usize); - cache.insert(id, create_test_arrow_array(8)).await.unwrap(); + cache + .insert(id, 0, create_test_arrow_array(8)) + .await + .unwrap(); cache.flush_all_to_disk().await.unwrap(); - store.remove(&entry_id_to_key(&id)).await.unwrap(); + store.remove(&entry_id_to_key(&id, 0)).await.unwrap(); - assert!(cache.get(&id).await.is_none()); + assert!(cache.get(&id, 0).await.is_none()); } #[tokio::test] @@ -1095,7 +1215,11 @@ mod tests { Arc::new(crate::liquid_array::LiquidArray::from_arrow_array(&arrow_array).unwrap()); store - .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) + .insert_inner( + entry_id, + WriteIdentity::Owned(0), + CacheEntry::memory_liquid(liquid.clone()), + ) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); @@ -1104,7 +1228,7 @@ mod tests { assert!(matches!(entry.as_ref(), CacheEntry::DiskLiquid { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), arrow_array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); @@ -1122,10 +1246,10 @@ mod tests { .await; let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let err = cache.insert(EntryID::from(900usize), array).await; + let err = cache.insert(EntryID::from(900usize), 0, array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.contains(&EntryID::from(900usize))); + assert!(!cache.is_cached(&EntryID::from(900usize), 0)); } #[tokio::test] @@ -1144,14 +1268,14 @@ mod tests { let first = EntryID::from(910usize); let second = EntryID::from(911usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.contains(&first)); + assert!(cache.is_cached(&first, 0)); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first)); + assert!(!cache.is_cached(&first, 0)); assert!(matches!( cache.index().get(&second).unwrap().as_ref(), CacheEntry::DiskArrow { .. } @@ -1172,13 +1296,13 @@ mod tests { .await; let first = EntryID::from(912usize); let second = EntryID::from(913usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first) || !cache.contains(&second)); + assert!(!cache.is_cached(&first, 0) || !cache.is_cached(&second, 0)); } #[tokio::test] @@ -1193,14 +1317,14 @@ mod tests { .build() .await; let entry = EntryID::from(914usize); - cache.insert(entry, array).await.unwrap(); + cache.insert(entry, 0, array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); let before = cache.stats().disk_usage_bytes; - cache.remove_disk_entry(entry).await; + cache.remove_disk_entry(entry, 0).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.contains(&entry)); + assert!(!cache.is_cached(&entry, 0)); } #[tokio::test] @@ -1213,11 +1337,11 @@ mod tests { .await; let entry_id = EntryID::from(901usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - cache.insert(entry_id, array).await.unwrap(); + cache.insert(entry_id, 0, array).await.unwrap(); let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.contains(&entry_id)); + assert!(!cache.is_cached(&entry_id, 0)); } } diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 62acca365..4ec252484 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -1,17 +1,84 @@ use congee::CongeeArc; use std::{ fmt::{Debug, Formatter}, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, + sync::atomic::{AtomicU64, AtomicUsize, Ordering}, }; use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; +use crate::sync::{Arc, RwLock}; + +/// The value stored in the ART. +/// +/// `CongeeArc` frees a removed or replaced value through crossbeam-epoch's +/// deferred destruction: it clones the `Arc` and drops the clone only when a +/// later pin collects that epoch's garbage — up to 64 objects per thread wait +/// in a thread-local bag, and the global queue drains 8 bags per 128 pins. +/// With multi-megabyte arrays as values, that kept every evicted entry alive +/// for an unbounded, budget-invisible stretch (liquid-cache#43: a tier +/// reporting at its limit while the process held several times that). +/// +/// So the tree stores a small slot and the payload is taken out of it the +/// moment the index gives the entry up. The deferred drop then reclaims only +/// an empty shell, and the array dies with the last caller-held reference. +/// +/// The slot also records the identity of what it holds. `EntryID` is a packed +/// integer whose fields are narrower than the values they encode, so two +/// distinct sources can compute the same key; the key alone therefore cannot +/// answer "is this the entry I asked for?". `identity` is the caller's +/// unnarrowed name for the data, compared on every read through +/// [`ArtIndex::get_checked`], which turns such aliasing into a miss rather +/// than a wrong answer. +struct Slot { + identity: u64, + entry: RwLock>>, +} + +impl Slot { + fn new(identity: u64, entry: CacheEntry) -> Arc { + Arc::new(Self { + identity, + entry: RwLock::new(Some(Arc::new(entry))), + }) + } + + fn load(&self) -> Option> { + self.entry.read().unwrap().clone() + } + + fn take(&self) -> Option> { + self.entry.write().unwrap().take() + } +} + +/// Whose data a write carries, and on what terms. +#[derive(Debug, Clone, Copy)] +pub(crate) enum WriteIdentity { + /// A caller storing its own data. Takes the key over if another identity + /// holds it: that identity belongs to a source that cannot read this key + /// any more, so leaving its entry there would cost the key to both. + Owned(u64), + /// Maintenance rewriting an entry it read earlier — transcode, squeeze, + /// hydrate, spill. It carries the identity the entry was read under and + /// lands only if the key still holds it. Adopting whatever is there + /// instead would relabel one source's data with another's whenever a + /// takeover lands between the read and the write, and the new owner would + /// then read those rows as its own. + Rewrite(u64), +} + +impl WriteIdentity { + /// The identity this write carries, whichever kind it is. + pub(crate) fn value(&self) -> u64 { + match self { + Self::Owned(id) | Self::Rewrite(id) => *id, + } + } +} pub(crate) struct ArtIndex { - art: CongeeArc, + art: CongeeArc, entry_count: AtomicUsize, + identity_mismatches: AtomicU64, } impl Debug for ArtIndex { @@ -22,62 +89,135 @@ impl Debug for ArtIndex { impl ArtIndex { pub(crate) fn new() -> Self { - let art: CongeeArc = CongeeArc::new(); Self { - art, + art: CongeeArc::new(), entry_count: AtomicUsize::new(0), + identity_mismatches: AtomicU64::new(0), } } + /// Look up an entry without checking whose it is. + /// + /// This is for maintenance that acts on whatever currently occupies a key — + /// eviction, squeezing, disk supersession, iteration for stats. A read + /// serving a caller must use [`Self::get_checked`] instead, so that a key + /// collision cannot return one caller another's data. pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); - let batch = self.art.get(*entry_id, &guard)?; - Some(batch) + // An empty slot means the entry was removed or replaced between the + // tree lookup and the load. A remove reading as a miss is exactly as + // if it had won the race outright, but after a replace the key is + // still present with a new slot, so look the key up once more rather + // than report a cached entry as absent. + let slot = self.art.get(*entry_id, &guard)?; + if let Some(entry) = slot.load() { + return Some(entry); + } + self.art.get(*entry_id, &guard)?.load() } - pub(crate) fn contains(&self, entry_id: &EntryID) -> bool { + /// Look up an entry, returning it only if it is the one `identity` names. + /// + /// A mismatch reads as a miss, so the caller re-reads from its source and + /// gets correct data. It is also counted: the tally is expected to stay at + /// zero, and a non-zero value means two sources are computing the same + /// `EntryID`. + pub(crate) fn get_checked(&self, entry_id: &EntryID, identity: u64) -> Option> { let guard = self.art.pin(); - self.art.get(*entry_id, &guard).is_some() + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + if let Some(entry) = slot.load() { + return Some(entry); + } + // Re-read as in `get`: a replace leaves the key present under a new + // slot, which carries its own identity and must be checked again. + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + slot.load() } - pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + /// Look up an entry together with the identity recorded against it, for + /// maintenance that must rewrite it under the identity it observed. + pub(crate) fn get_with_identity(&self, entry_id: &EntryID) -> Option<(u64, Arc)> { let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + if let Some(entry) = slot.load() { + return Some((identity, entry)); + } + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + slot.load().map(|entry| (identity, entry)) + } + + pub(crate) fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_checked(entry_id, identity).is_some() + } + + /// Store `batch` under `entry_id`, returning whether it was stored. + /// + /// See [`WriteIdentity`] for the two kinds of write and why they differ. + pub(crate) fn insert( + &self, + entry_id: &EntryID, + identity: WriteIdentity, + batch: CacheEntry, + ) -> bool { + let guard = self.art.pin(); + let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); + let identity = match (identity, existing_identity) { + (WriteIdentity::Owned(new), Some(old)) => { + if new != old { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + } + new + } + (WriteIdentity::Owned(new), None) => new, + // The key still holds what this rewrite was built from. + (WriteIdentity::Rewrite(expected), Some(old)) if expected == old => expected, + // It does not: the entry was taken over or removed while this + // rewrite was in flight, so the payload belongs to a source that + // no longer owns the key. Drop it. + (WriteIdentity::Rewrite(_), _) => return false, + }; let existing = self .art - .insert(*entry_id, Arc::new(batch), &guard) + .insert(*entry_id, Slot::new(identity, batch), &guard) .expect("Insertion failed"); - if existing.is_none() { - self.entry_count.fetch_add(1, Ordering::Relaxed); + match existing { + Some(replaced) => drop(replaced.take()), + None => { + self.entry_count.fetch_add(1, Ordering::Relaxed); + } } + true } pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); - let removed = self.art.remove(*entry_id, &guard); - if removed.is_some() { - self.entry_count.fetch_sub(1, Ordering::Relaxed); - } - removed + let removed = self.art.remove(*entry_id, &guard)?; + self.entry_count.fetch_sub(1, Ordering::Relaxed); + removed.take() } pub(crate) fn reset(&self) { - let guard = self.art.pin(); - self.art.keys().into_iter().for_each(|k| { - _ = self.art.remove(k, &guard).unwrap(); - }); + for k in self.art.keys() { + self.remove(&k); + } self.entry_count.store(0, Ordering::Relaxed); } - pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { - let guard = self.art.pin(); - for id in self.art.keys().into_iter() { - f( - &id, - &self - .art - .get(id, &guard) - .expect("Failed to get value from ART"), - ); + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { + for id in self.art.keys() { + if let Some((identity, entry)) = self.get_with_identity(&id) { + f(&id, identity, &entry); + } } } @@ -89,6 +229,15 @@ impl ArtIndex { pub(crate) fn entry_count(&self) -> usize { self.entry_count.load(Ordering::Relaxed) } + + /// How many lookups or inserts found a key held by a different identity. + /// + /// Expected to stay at zero. A non-zero value means two sources compute the + /// same `EntryID`, and every one of them was served correctly only because + /// the check turned it into a miss. + pub(crate) fn identity_mismatches(&self) -> u64 { + self.identity_mismatches.load(Ordering::Relaxed) + } } #[cfg(test)] @@ -99,24 +248,24 @@ mod tests { use super::*; #[test] - fn test_get_and_contains() { + fn test_get_and_is_cached() { let store = ArtIndex::new(); let entry_id1: EntryID = EntryID::from(1); let entry_id2: EntryID = EntryID::from(2); let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(!store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&entry_id2, 0)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached { - store.insert(&entry_id1, array1.clone()); + store.insert(&entry_id1, WriteIdentity::Owned(0), array1.clone()); } - assert!(store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&entry_id2, 0)); // Get should return the cached value match store.get(&entry_id1) { @@ -134,13 +283,135 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, array.clone()); + store.insert(&entry_id, WriteIdentity::Owned(0), array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.contains(&entry_id)); + assert!(store.is_cached(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.contains(&entry_id)); + assert!(!store.is_cached(&entry_id, 0)); + } + + /// The array behind a removed or replaced entry must die with the last + /// caller-held reference, not wait for epoch garbage collection. + #[test] + fn removed_and_replaced_entries_are_released_immediately() { + let store = ArtIndex::new(); + let id = EntryID::from(1); + + let first = create_test_array(100); + let CacheEntry::MemoryArrow(first_array) = &first else { + unreachable!() + }; + let weak_first = Arc::downgrade(first_array); + store.insert(&id, WriteIdentity::Owned(0), first); + store.insert(&id, WriteIdentity::Owned(0), create_test_array(200)); + assert!( + weak_first.upgrade().is_none(), + "replaced entry still alive: held by the index's deferred drop" + ); + + let second = store.get(&id).unwrap(); + let removed = store.remove(&id).unwrap(); + let CacheEntry::MemoryArrow(second_array) = removed.as_ref() else { + unreachable!() + }; + let weak_second = Arc::downgrade(second_array); + drop((second, removed)); + assert!( + weak_second.upgrade().is_none(), + "removed entry still alive: held by the index's deferred drop" + ); + assert_eq!(store.entry_count(), 0); + } + + /// Two files whose ids narrow to the same `EntryID` must not read each + /// other's data. Before the identity check this returned the incumbent's + /// array, which is a wrong answer whenever the two happen to share a type. + #[test] + fn an_entry_is_never_served_to_a_different_identity() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(7); + + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + + // The colliding file asks for the same key and is told nothing is there. + assert!(store.get_checked(&key, 2).is_none()); + assert!(!store.is_cached(&key, 2)); + assert_eq!(store.identity_mismatches(), 2); + + // The owner still reads its own entry. + assert!(store.get_checked(&key, 1).is_some()); + + // The colliding file takes the key over. It has to: it cannot read + // what is there, so leaving it would cost the key to both of them. + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + assert!( + store.get_checked(&key, 1).is_none(), + "the displaced file reads a miss, never the other file's rows" + ); + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!(array.len(), 200), + other => panic!("expected the new owner's array, found {other}"), + } + } + + /// A rewrite is built from an entry read earlier, and the key can be taken + /// over in between — a squeeze reads, awaits a disk write, then stores. If + /// the rewrite adopted whatever identity held the key by then, it would + /// relabel the old file's data as the new owner's, and the new owner would + /// read those rows as a hit. Carrying the identity it read under makes the + /// stale write drop instead. + #[test] + fn a_rewrite_does_not_land_on_a_key_taken_over_since_it_was_read() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(11); + + // File A caches, and something begins rewriting that entry. + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + let (observed, _read) = store.get_with_identity(&key).unwrap(); + assert_eq!(observed, 1); + + // File B takes the key over while that rewrite is in flight. + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + + // The rewrite lands too late and must be dropped, not relabelled. + assert!(!store.insert( + &key, + WriteIdentity::Rewrite(observed), + create_test_array(100) + )); + + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!( + array.len(), + 200, + "the new owner must still read its own rows, not the rewrite's" + ), + other => panic!("expected the new owner's array, found {other}"), + } + } + + /// Maintenance rewrites a key in place and must neither change whose the + /// entry is nor bring back one that has been removed — a stale reader that + /// misses goes on to insert what it read, and that write must not land + /// under a key nobody owns any more. + #[test] + fn maintenance_preserves_identity_and_cannot_resurrect_a_removed_key() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(9); + + assert!(store.insert(&key, WriteIdentity::Owned(5), create_test_array(10))); + assert!(store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(20))); + assert!( + store.get_checked(&key, 5).is_some(), + "rewriting in place kept the identity" + ); + + store.remove(&key); + assert!(!store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(30))); + assert!(store.get(&key).is_none()); + assert_eq!(store.entry_count(), 0); } } diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index 7a7f71d3d..12152a0a3 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -26,8 +26,21 @@ pub trait EntryMetadata: Debug + Send + Sync { } /// Convert an [`EntryID`] to a t4 key (8-byte little-endian representation). -pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { - usize::from(*entry_id).to_le_bytes().to_vec() +/// Convert an [`EntryID`] and the identity that owns it to a t4 key. +/// +/// Both halves, not just the entry id. `EntryID` is a packed integer whose +/// fields are narrower than the values they encode, so two sources can compute +/// one id — and a store object addressed by that id alone is *shared*. Scoping +/// only the index entry is not enough: a write in flight for one owner can land +/// after another has taken the key over and installed its own disk entry, +/// overwriting bytes the new owner's index entry agrees are its own. With the +/// identity in the key the two address different objects, so a late write +/// cannot reach the other's bytes at all. +pub(crate) fn entry_id_to_key(entry_id: &EntryID, identity: u64) -> Vec { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&usize::from(*entry_id).to_le_bytes()); + key.extend_from_slice(&identity.to_le_bytes()); + key } /// A default implementation of [`EntryMetadata`]. diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index 6c871583d..e8844923e 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -110,6 +110,12 @@ define_runtime_stats! { pub struct CacheStats { /// Total number of entries in the cache. pub total_entries: usize, + /// How many lookups or writes found a key held by a different file. + /// + /// Expected to stay at zero. A non-zero value means two sources computed + /// the same `EntryID`, and each was served correctly only because the + /// identity check turned the collision into a miss. + pub identity_mismatches: u64, /// Number of in-memory Arrow entries. pub memory_arrow_entries: usize, /// Number of in-memory Liquid entries. diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index ee3681a45..928ee86f5 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -26,12 +26,12 @@ async fn default_policies() { for i in 0..5 { let entry_id = EntryID::from(i); - cache.insert(entry_id, test_array.clone()).await.unwrap(); + cache.insert(entry_id, 0, test_array.clone()).await.unwrap(); } for i in 0..5 { let entry_id = EntryID::from(i); - let array = cache.get(&entry_id).read().await.unwrap(); + let array = cache.get(&entry_id, 0).read().await.unwrap(); assert_eq!(array.len(), test_array.len()); } @@ -54,17 +54,17 @@ async fn insert_wont_fit_cache() { .build() .await; cache - .insert(EntryID::from(0), test_array.clone()) + .insert(EntryID::from(0), 0, test_array.clone()) .await .unwrap(); let array_3x = arrow::compute::concat(&[&test_array, &test_array, &test_array]).unwrap(); let array_9x = arrow::compute::concat(&[&array_3x, &array_3x, &array_3x]).unwrap(); let array_27x = arrow::compute::concat(&[&array_9x, &array_9x, &array_9x]).unwrap(); cache - .insert(EntryID::from(1), array_27x.clone()) + .insert(EntryID::from(1), 0, array_27x.clone()) .await .unwrap(); - cache.get(&EntryID::from(1)).read().await.unwrap(); + cache.get(&EntryID::from(1), 0).read().await.unwrap(); let trace = cache.consume_event_trace(); let json_trace = serde_json::to_string(&trace).unwrap(); @@ -100,11 +100,11 @@ async fn liquid_eviction_reads_memory_and_disk() { .into_iter() .enumerate() { - cache.insert(EntryID::from(id), array).await.unwrap(); + cache.insert(EntryID::from(id), 0, array).await.unwrap(); } let mut states = Vec::new(); - cache.for_each_entry(|_, entry| states.push(CachedBatchType::from(entry))); + cache.for_each_entry(|_, _, entry| states.push(CachedBatchType::from(entry))); assert!( states .iter() @@ -113,13 +113,13 @@ async fn liquid_eviction_reads_memory_and_disk() { let id = EntryID::from(0); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); let selection = BooleanBuffer::from_iter((0..integers.len()).map(|index| index % 3 == 0)); let selected = cache - .get(&id) + .get(&id, 0) .with_selection(&selection) .read() .await @@ -137,7 +137,11 @@ async fn liquid_eviction_reads_memory_and_disk() { Arc::new(Literal::new(ScalarValue::Int64(Some(1_000)))), )); let predicate = LiquidExpr::try_new(physical, &DataType::Int64).unwrap(); - let actual = cache.eval_predicate(&id, &predicate).read().await.unwrap(); + let actual = cache + .eval_predicate(&id, 0, &predicate) + .read() + .await + .unwrap(); let expected = arrow::array::BooleanArray::from_iter( integers .as_any() @@ -150,14 +154,14 @@ async fn liquid_eviction_reads_memory_and_disk() { cache.flush_all_to_disk().await.unwrap(); let mut disk_states = Vec::new(); - cache.for_each_entry(|_, entry| disk_states.push(CachedBatchType::from(entry))); + cache.for_each_entry(|_, _, entry| disk_states.push(CachedBatchType::from(entry))); assert!(disk_states.iter().all(|state| matches!( state, CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow ))); assert!(disk_states.contains(&CachedBatchType::DiskLiquid)); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); } diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index 18e01d5c3..b25717534 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -108,7 +108,7 @@ fn main() { continue; }; if storage - .eval_predicate(id, &liquid_expr) + .eval_predicate(id, 0, &liquid_expr) .with_selection(&selection) .await .is_some() @@ -161,7 +161,7 @@ fn load_and_insert_referer( let id = EntryID::from(idx); ids.push(id); total_size += array.get_array_memory_size(); - storage.insert(id, array).await.unwrap(); + storage.insert(id, 0, array).await.unwrap(); idx += 1; } diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index c08cf9af6..61bffe002 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -20,6 +20,10 @@ use std::sync::Arc; #[derive(Debug)] pub struct CachedColumn { cache_store: Arc, + /// The lease its keys are built from. Held so the id cannot be recycled + /// while this column can still compute keys from it; `identity()` is what + /// every cache read and write is checked against. + file_id: Arc, field: Arc, column_path: ColumnAccessPath, expression: Option>, @@ -54,6 +58,7 @@ impl CachedColumn { pub(crate) fn new( field: Arc, cache_store: Arc, + file_id: Arc, column_access_path: ColumnAccessPath, expression: Option>, is_predicate_column: bool, @@ -78,6 +83,7 @@ impl CachedColumn { Self { field, cache_store, + file_id, column_path: column_access_path, expression, snapshots, @@ -85,12 +91,19 @@ impl CachedColumn { } /// row_id must be on a batch boundary. + /// The unnarrowed name of the file this column reads, recorded with every + /// entry it caches and checked on every read. + pub(crate) fn identity(&self) -> u64 { + self.file_id.identity() + } + pub(crate) fn entry_id(&self, batch_id: BatchID) -> ParquetArrayID { self.column_path.entry_id(batch_id) } pub(crate) fn contains(&self, batch_id: BatchID) -> bool { - self.cache_store.contains(&self.entry_id(batch_id).into()) + self.cache_store + .is_cached(&self.entry_id(batch_id).into(), self.identity()) } pub(crate) fn snapshot_entry(&self, batch_id: BatchID) -> Option> { @@ -131,6 +144,7 @@ impl CachedColumn { self.cache_store .eval_predicate_on_entry( &entry_id, + self.identity(), entry.as_ref(), Some(filter), &liquid_expr, @@ -139,7 +153,7 @@ impl CachedColumn { } None => { self.cache_store - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, self.identity(), &liquid_expr) .with_selection(filter) .await } @@ -191,6 +205,7 @@ impl CachedColumn { .cache_store .read_entry( &entry_id, + self.identity(), entry.as_ref(), Some(filter), self.expression.as_deref(), @@ -198,7 +213,7 @@ impl CachedColumn { .await; } self.cache_store - .get(&entry_id) + .get(&entry_id, self.identity()) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -208,7 +223,7 @@ impl CachedColumn { #[cfg(test)] pub(crate) async fn get_arrow_array_test_only(&self, batch_id: BatchID) -> Option { let entry_id = self.entry_id(batch_id).into(); - self.cache_store.get(&entry_id).await + self.cache_store.get(&entry_id, self.identity()).await } /// Insert an array into the cache. @@ -222,7 +237,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), array) + .insert(self.entry_id(batch_id).into(), self.identity(), array) .await?; Ok(()) } @@ -239,7 +254,7 @@ impl CachedColumn { if self.snapshots.get(&entry_id).is_some() { return PrefetchOutcome::AlreadySnapshotted; } - match self.cache_store.prefetch(&entry_id).await { + match self.cache_store.prefetch(&entry_id, self.identity()).await { PrefetchResult::Snapshot(entry) => { self.snapshots.insert(entry_id, entry); PrefetchOutcome::Snapshotted diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs new file mode 100644 index 000000000..94feea765 --- /dev/null +++ b/src/datafusion/src/cache/file_id.rs @@ -0,0 +1,473 @@ +//! Allocation of the file ids that name cached data. +//! +//! An id is the part of a cache key that says which file an entry came from. +//! It is narrowed into 16 bits by [`crate::cache::ColumnAccessPath`], so the +//! supply of *distinct* keys is finite while the number of files a process +//! opens is not. Handing ids out from a counter that only ever climbs means a +//! long-lived process eventually reuses a key while its previous owner's data +//! is still cached. +//! +//! So an id is a lease rather than a permanent assignment. It is held by +//! everything that can still compute a key from it — the file handle, the row +//! groups and columns derived from it — and returns to the pool when the last +//! of them is dropped. The live id count is then bounded by what is actually +//! being read, not by everything that has ever been read. +//! +//! Cache *entries* deliberately do not hold a lease. An id can be reused while +//! entries keyed from it are still resident, and each entry records the +//! identity of the file it came from, so the new owner's reads miss rather +//! than returning the previous owner's rows. +//! +//! Its writes are not refused, though. A key held by another identity belongs +//! to a file that has already let its id go, so nothing can read that entry +//! any more and the new owner takes the key over +//! (`liquid_cache::cache::ArtIndex::insert`). Refusing instead would leave the +//! key occupied by data nobody can use, and on a cache below its budget +//! nothing evicts it — the new owner would never cache that key again. +//! +//! The alternative, releasing ids from inside index removal, would take a +//! process-wide lock underneath a crossbeam-epoch pin and deadlock against +//! `reset`. Keeping id lifetime and entry lifetime separate is what avoids +//! that, and the identity check is what makes the overlap safe. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicU64, Ordering}; + +use ahash::AHashMap; +// Through `crate::sync`, not `std::sync`: under the shuttle test feature this +// resolves to shuttle's primitives, which is what lets the model checker +// explore interleavings across `acquire` and `release`. A `std::sync::Mutex` +// is opaque to it, so the pool would be excluded from the very job that is +// meant to cover it. +use crate::sync::{Arc, Mutex, Weak}; + +/// A leased file id. The id returns to its pool when this is dropped. +/// +/// Two numbers, because they answer different questions and only one of them +/// can be recycled: +/// +/// * `id` goes into the cache key, whose file field is 16 bits wide. It has to +/// be recycled or a long-lived process runs out. +/// * `identity` names *which file* an entry came from, and is never reused. +/// It cannot be the recycled id: a file that inherits id 0 from a file that +/// has finished would otherwise be indistinguishable from it, and would read +/// the entries it left behind — the exact aliasing the identity exists to +/// catch. +#[derive(Debug)] +pub(crate) struct FileId { + id: u64, + identity: u64, + path: String, + pool: Arc, +} + +impl FileId { + /// The narrow, recycled id the cache key is built from. + pub(crate) fn get(&self) -> u64 { + self.id + } + + /// The wide, never-reused name for this file, recorded alongside every + /// entry so a recycled key cannot serve one file another's data. + pub(crate) fn identity(&self) -> u64 { + self.identity + } +} + +impl Drop for FileId { + fn drop(&mut self) { + self.pool.release(&self.path, self.id, self.identity); + } +} + +/// Hands out file ids and takes them back. +#[derive(Debug, Default)] +pub(crate) struct FileIdPool { + inner: Mutex, + /// Ids ever allocated that did not fit the key's 16-bit file field. A + /// non-zero count means keys are aliasing and the cache is refusing to + /// serve entries across the alias, which is correct but costs hit rate. + over_key_width: AtomicU64, +} + +#[derive(Debug, Default)] +struct PoolInner { + /// Live leases by path, so concurrent readers of one file share an id. + /// Entries are weak: the map never keeps a file alive on its own, and a + /// path is removed when its lease is released. + leases: AHashMap>, + /// Released ids, reused oldest-first. FIFO rather than LIFO on purpose: a + /// just-released id is the one whose entries are most likely still + /// resident, and reusing it last gives them the longest window to be + /// evicted before anything keys over them. + /// + /// Each carries the path that released it and the identity it had. If the + /// same path comes back it keeps that identity, so its cached entries are + /// still its own and still readable — a file read twice is a cache hit, + /// not a collision. Any other path gets a fresh identity, so it cannot + /// read what the previous holder left behind. + free: VecDeque, + next: u64, + /// Only ever climbs. A `u64` of these is not a resource worth reclaiming: + /// at one a microsecond it outlasts the hardware. + next_identity: u64, +} + +#[derive(Debug)] +struct Released { + id: u64, + path: String, + identity: u64, +} + +impl FileIdPool { + pub(crate) fn new() -> Arc { + Arc::new(Self::default()) + } + + /// The lease for `path`, shared with any reader already holding one. + pub(crate) fn acquire(self: &Arc, path: &str) -> Arc { + let mut inner = self.inner.lock().unwrap(); + if let Some(existing) = inner.leases.get(path).and_then(Weak::upgrade) { + return existing; + } + // Prefer this path's own released record, wherever it sits in the + // queue. Matching only the front would restore an identity just when + // release order happens to match acquire order — release order is + // stream completion order and acquire order is partition open order, + // so for any scan over more than one file they diverge and every + // re-read would orphan the entries it cached last time. + let mine = inner.free.iter().position(|r| r.path == path); + let (id, reusable_identity) = match mine { + Some(at) => { + let released = inner.free.remove(at).expect("index came from the queue"); + (released.id, Some(released.identity)) + } + // A fresh id while the key's file field still has room, and only + // then the oldest released one. Recycling eagerly would be correct + // — the identity check turns an aliased read into a miss — but it + // costs the previous holder every entry it cached, since the new + // owner takes those keys over. Two files read one after another, + // or the two sides of a join, would evict each other for no reason + // while 65,000 ids sat unused. Reuse is what the narrow key field + // forces, not something to spend before it is needed. + None if inner.next <= u16::MAX as u64 => { + let id = inner.next; + inner.next += 1; + (id, None) + } + None => match inner.free.pop_front() { + Some(released) => (released.id, None), + None => { + let id = inner.next; + inner.next += 1; + (id, None) + } + }, + }; + if id > u16::MAX as u64 { + self.over_key_width.fetch_add(1, Ordering::Relaxed); + } + let identity = match reusable_identity { + Some(identity) => identity, + None => { + let identity = inner.next_identity; + inner.next_identity += 1; + identity + } + }; + let lease = Arc::new(FileId { + id, + identity, + path: path.to_string(), + pool: Arc::clone(self), + }); + inner + .leases + .insert(path.to_string(), Arc::downgrade(&lease)); + lease + } + + fn release(&self, path: &str, id: u64, self_identity: u64) { + let Ok(mut inner) = self.inner.lock() else { + // A poisoned pool means some other thread panicked holding it. + // Losing one id is better than panicking again inside a drop. + return; + }; + // Only drop the path if it still points at the lease being released. + // A new lease for the same path may already have replaced it, and + // removing that one would hand the same file two live ids. + if inner + .leases + .get(path) + .is_some_and(|weak| weak.strong_count() == 0) + { + inner.leases.remove(path); + } + inner.free.push_back(Released { + id, + path: path.to_string(), + identity: self_identity, + }); + } + + /// Ids currently leased. Bounded by what is being read, which is what + /// keeps the 16-bit key field from running out. + pub(crate) fn live_count(&self) -> usize { + self.inner.lock().map(|i| i.leases.len()).unwrap_or(0) + } + + /// Ids handed out that do not fit the key's file field. Expected to stay + /// at zero. + pub(crate) fn over_key_width(&self) -> u64 { + self.over_key_width.load(Ordering::Relaxed) + } + + /// Forget every lease and start ids from zero again. + /// + /// Only valid when nothing holds a lease; callers that still do would keep + /// computing keys from ids this pool is free to hand out again. + pub(crate) fn reset(&self) { + if let Ok(mut inner) = self.inner.lock() { + inner.leases.clear(); + inner.free.clear(); + inner.next = 0; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A file re-opened after other files have come and gone must still find + /// its own record. Release order is stream completion order and acquire + /// order is partition open order, so the two rarely line up; matching only + /// the front of the queue would hand a re-read a fresh identity and orphan + /// everything it cached before. + #[test] + fn a_reopened_path_finds_its_record_anywhere_in_the_queue() { + let pool = FileIdPool::new(); + + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + let (a_id, a_identity) = (a.get(), a.identity()); + let (b_id, b_identity) = (b.get(), b.identity()); + + // Released a-then-b, so b's record sits behind a's. + drop(a); + drop(b); + + // Re-open b first: the queue front is a's record, not b's. + let b_again = pool.acquire("b.parquet"); + assert_eq!(b_again.get(), b_id, "b should get its own id back"); + assert_eq!( + b_again.identity(), + b_identity, + "b should keep its identity, or its cached entries are orphaned" + ); + + let a_again = pool.acquire("a.parquet"); + assert_eq!(a_again.get(), a_id); + assert_eq!(a_again.identity(), a_identity); + } + + /// Two leases alive at the same time must never share an id, and never + /// share an identity. That is the property the whole scheme rests on: + /// a shared id means two files computing one key, and a shared identity + /// means the check that catches it cannot tell them apart. + /// + /// Run under the model checker because `acquire` and `release` race by + /// construction — a lease is released from `Drop`, on whatever thread + /// happened to hold it last. + fn concurrent_leases_stay_distinct() { + let pool = FileIdPool::new(); + let mut threads = Vec::new(); + + for t in 0..3 { + let pool = Arc::clone(&pool); + threads.push(crate::sync::thread::spawn(move || { + for i in 0..3 { + let mine = pool.acquire(&format!("f{t}-{i}.parquet")); + + // Held at the same time, so they cannot be the same file. + let probe = pool.acquire("probe.parquet"); + assert_ne!(mine.get(), probe.get(), "two live leases shared an id"); + assert_ne!( + mine.identity(), + probe.identity(), + "two live leases shared an identity" + ); + drop(probe); + + // The same path always resolves to the same lease. + let again = pool.acquire(&format!("f{t}-{i}.parquet")); + assert_eq!(mine.get(), again.get()); + assert_eq!(mine.identity(), again.identity()); + } + })); + } + + for thread in threads { + thread.join().unwrap(); + } + } + + #[test] + fn concurrent_leases_stay_distinct_single_threaded() { + concurrent_leases_stay_distinct(); + } + + #[cfg(feature = "shuttle")] + #[test] + fn shuttle_concurrent_leases_stay_distinct() { + let mut runner = shuttle::PortfolioRunner::new(true, Default::default()); + let cores = std::thread::available_parallelism().unwrap().get().min(4); + for _ in 0..cores { + runner.add(shuttle::scheduler::PctScheduler::new(10, 1_000)); + } + runner.run(concurrent_leases_stay_distinct); + } + + #[test] + fn concurrent_readers_of_one_path_share_a_lease() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + assert_eq!(first.get(), second.get()); + assert_eq!(pool.live_count(), 1); + } + + #[test] + fn an_id_returns_to_the_pool_when_its_last_holder_drops() { + let pool = FileIdPool::new(); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + assert_eq!((a.get(), b.get()), (0, 1)); + assert_eq!(pool.live_count(), 2); + + drop(a); + assert_eq!(pool.live_count(), 1, "the released path is forgotten"); + + // The released id is available, but a fresh one is preferred while the + // key field has room: reuse costs the previous holder its entries, so + // it is spent only once there is nothing else to hand out. + let c = pool.acquire("c.parquet"); + assert_eq!(c.get(), 2); + assert_eq!(pool.live_count(), 2); + } + + #[test] + fn a_second_holder_keeps_the_id_alive() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + drop(first); + assert_eq!(pool.live_count(), 1); + // Bound, not a temporary: an unheld lease is released the moment the + // expression ends, which would put its id back before the next call. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), 1, "id 0 is still leased"); + drop(second); + assert_eq!(pool.live_count(), 1, "the last holder released id 0"); + } + + /// The two numbers have to move independently. Reusing an id is how the + /// key space stays bounded; reusing an *identity* for a different file is + /// how one file reads another's entries. Re-opening the same path must + /// keep its identity, or every lease boundary silently empties the cache. + #[test] + fn identity_follows_the_path_while_the_id_is_recycled() { + let pool = FileIdPool::new(); + // Past the key width, so every acquire below takes the recycling path + // rather than a fresh id. Held, so the queue holds only what this test + // releases. + let _fillers = exhaust_key_width(&pool); + + let first = pool.acquire("a.parquet"); + let (a_id, a_identity) = (first.get(), first.identity()); + drop(first); + + // Same file again: same id and the same name, so its cached entries + // are still its own. + let reopened = pool.acquire("a.parquet"); + assert_eq!(reopened.get(), a_id); + assert_eq!( + reopened.identity(), + a_identity, + "re-opening a file must keep its identity, or its cache is dead" + ); + drop(reopened); + + // A different file inherits the id but must not inherit the name. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), a_id, "the id is recycled"); + assert_ne!( + other.identity(), + a_identity, + "a different file must not be able to read what the last one left" + ); + } + + /// Once the key field is exhausted there is nothing to hand out but + /// released ids, and the oldest goes first — giving the most recently + /// released file's entries the longer window before they are taken over. + #[test] + fn released_ids_are_reused_oldest_first() { + let pool = FileIdPool::new(); + let _fillers = exhaust_key_width(&pool); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + let (a_id, b_id) = (a.get(), b.get()); + drop(a); + drop(b); + assert_eq!(pool.acquire("x.parquet").get(), a_id); + assert_eq!(pool.acquire("y.parquet").get(), b_id); + } + + /// Drive `next` past the 16-bit key field so the pool has no fresh ids to + /// hand out and must recycle. + /// + /// The fillers are returned rather than dropped, and the caller has to keep + /// them: dropping them here would leave 65,536 released records queued + /// ahead of whatever the test then releases, and FIFO reuse would hand back + /// a filler's id instead of the one under test. + #[must_use] + fn exhaust_key_width(pool: &Arc) -> Vec> { + (0..=(u16::MAX as u64)) + .map(|i| pool.acquire(&format!("filler_{i}.parquet"))) + .collect() + } + + #[test] + fn ids_beyond_the_key_width_are_counted() { + let pool = FileIdPool::new(); + { + let mut inner = pool.inner.lock().unwrap(); + inner.next = u16::MAX as u64; + } + let _fits = pool.acquire("fits.parquet"); + assert_eq!(pool.over_key_width(), 0); + let _over = pool.acquire("over.parquet"); + assert_eq!(pool.over_key_width(), 1); + } + + /// A path re-registered while its old lease is being dropped must not lose + /// the new lease's entry in the map — that would give one file two live + /// ids and split its cache. + #[test] + fn releasing_a_stale_lease_leaves_a_newer_one_alone() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let first_id = first.get(); + drop(first); + let second = pool.acquire("a.parquet"); + assert_eq!(second.get(), first_id, "the id came back round"); + assert_eq!(pool.live_count(), 1); + assert_eq!( + pool.acquire("a.parquet").get(), + second.get(), + "the live lease is still the one the map points at" + ); + } +} diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index fed837f6f..162774eb2 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -3,7 +3,9 @@ use crate::io::ParquetCacheMetadata; use crate::reader::{LiquidPredicate, extract_multi_column_or}; -use crate::sync::{Mutex, RwLock}; +use crate::sync::RwLock; +mod file_id; + use ahash::AHashMap; use arrow::array::{BooleanArray, RecordBatch}; use arrow::buffer::BooleanBuffer; @@ -21,7 +23,6 @@ use parquet::arrow::arrow_reader::ArrowPredicate; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; mod column; mod id; @@ -50,6 +51,15 @@ pub struct ParquetFileIdentity { } impl ParquetFileIdentity { + /// The key a file id is leased against. + /// + /// Both components, for the same reason the identity carries both: a path + /// is unique only within its object store, so keying on it alone would let + /// two stores share one lease and therefore one set of cache keys. + pub(crate) fn lease_key(&self) -> String { + format!("{}{}", self.object_store_url.as_str(), self.path) + } + /// Create an identity from an object store URL and an object path. pub fn new(object_store_url: ObjectStoreUrl, path: String) -> Self { Self { @@ -108,16 +118,18 @@ impl CachedRowGroup { fn new( cache_store: Arc, row_group_idx: u64, - file_idx: u64, + file_id: Arc, columns: &[CachedColumnSpec], snapshots: Arc, ) -> Self { let mut column_maps = ColumnMaps::default(); for (column_id, field, expression, is_predicate_column) in columns { - let column_access_path = ColumnAccessPath::new(file_idx, row_group_idx, *column_id); + let column_access_path = + ColumnAccessPath::new(file_id.get(), row_group_idx, *column_id); let column = Arc::new(CachedColumn::new( Arc::clone(field), Arc::clone(&cache_store), + Arc::clone(&file_id), column_access_path, expression.clone(), *is_predicate_column, @@ -218,7 +230,11 @@ impl CachedRowGroup { let entry_id = column.entry_id(batch_id).into(); let liquid_array = match snapshot_liquid { Some(array) => Some(array), - None => self.cache_store.try_read_liquid(&entry_id).await, + None => { + self.cache_store + .try_read_liquid(&entry_id, column.identity()) + .await + } }; let liquid_array = match liquid_array { None => { @@ -266,7 +282,9 @@ pub(crate) type CachedRowGroupRef = Arc; #[derive(Debug)] pub struct CachedFile { cache_store: Arc, - file_id: u64, + /// Held, not copied: the id returns to the pool when the last holder — this + /// file and everything derived from it — is dropped. + file_id: Arc, file_schema: SchemaRef, lineages: Arc, } @@ -274,7 +292,7 @@ pub struct CachedFile { impl CachedFile { fn new( cache_store: Arc, - file_id: u64, + file_id: Arc, file_schema: SchemaRef, lineages: Arc, ) -> Self { @@ -321,7 +339,7 @@ impl CachedFile { Arc::new(CachedRowGroup::new( self.cache_store.clone(), row_group_id, - self.file_id, + Arc::clone(&self.file_id), &columns, snapshots, )) @@ -345,11 +363,12 @@ pub(crate) type CachedFileRef = Arc; #[derive(Debug)] pub struct LiquidCacheParquet { /// Map object-store-qualified file identity to file id. - files: Mutex>, + /// Leases the file ids that cache keys are built from, so the id space + /// tracks the files being read rather than every file ever read — see + /// [`file_id`]. + file_ids: Arc, cache_store: Arc, - - current_file_id: AtomicU64, } /// A reference to the main cache structure. @@ -408,9 +427,8 @@ impl LiquidCacheParquet { .await; LiquidCacheParquet { - files: Mutex::new(AHashMap::new()), + file_ids: file_id::FileIdPool::new(), cache_store: cache_storage, - current_file_id: AtomicU64::new(0), } } @@ -431,20 +449,36 @@ impl LiquidCacheParquet { full_file_schema: SchemaRef, lineages: Arc, ) -> CachedFileRef { - let mut files = self.files.lock().unwrap(); - let file_id = *files - .entry(file_identity) - .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); - drop(files); - Arc::new(CachedFile::new( self.cache_store.clone(), - file_id, + self.file_ids.acquire(&file_identity.lease_key()), full_file_schema, lineages, )) } + /// How many file ids are currently leased. + /// + /// Bounded by the files being read, not by everything ever read. Rising + /// without bound means leases are being held longer than the reads that + /// need them. + pub fn leased_file_ids(&self) -> usize { + self.file_ids.live_count() + } + + /// How many ids have been handed out that do not fit the cache key's + /// 16-bit file field. Non-zero means more files are being read at once + /// than the key can name, and ids are being recycled under entries that + /// are still resident. + pub fn file_ids_over_key_width(&self) -> u64 { + self.file_ids.over_key_width() + } + + /// How many cache lookups or writes found a key held by another file. + pub fn identity_mismatches(&self) -> u64 { + self.cache_store.stats().identity_mismatches + } + /// Get the batch size of the cache. pub fn batch_size(&self) -> usize { self.cache_store.config().batch_size() @@ -491,8 +525,7 @@ impl LiquidCacheParquet { /// This is unsafe because resetting the cache while other threads are using the cache may cause undefined behavior. /// You should only call this when no one else is using the cache. pub unsafe fn reset(&self) { - let mut files = self.files.lock().unwrap(); - files.clear(); + self.file_ids.reset(); self.cache_store.reset(); } diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index d1947fc64..fadba50b9 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -123,35 +123,36 @@ impl LiquidCacheParquet { /// Write the stats of the cache to a parquet file. pub fn write_stats(&self, parquet_file_path: impl AsRef) -> Result<(), ParquetError> { let mut writer = StatsWriter::new(parquet_file_path)?; - self.cache_store.for_each_entry(|entry_id, cached_batch| { - let memory_size = cached_batch.memory_usage_bytes(); - let row_count = match cached_batch { - CacheEntry::MemoryArrow(array) => Some(array.len() as u64), - CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), - CacheEntry::DiskLiquid { .. } => None, - CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count - }; - let cache_type = match cached_batch { - CacheEntry::MemoryArrow(_) => "InMemory", - CacheEntry::MemoryLiquid(_) => "LiquidMemory", - CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", - CacheEntry::DiskArrow { .. } => "OnDiskArrow", - }; - let reference_count = cached_batch.reference_count(); - let entry_id = ParquetArrayID::from(*entry_id); - writer - .append_entry( - &entry_id.display_path(), - entry_id.row_group_id_inner(), - entry_id.column_id_inner(), - entry_id.batch_id_inner() * self.batch_size() as u64, - row_count, - memory_size as u64, - cache_type, - reference_count as u64, - ) - .unwrap(); - }); + self.cache_store + .for_each_entry(|entry_id, _, cached_batch| { + let memory_size = cached_batch.memory_usage_bytes(); + let row_count = match cached_batch { + CacheEntry::MemoryArrow(array) => Some(array.len() as u64), + CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), + CacheEntry::DiskLiquid { .. } => None, + CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count + }; + let cache_type = match cached_batch { + CacheEntry::MemoryArrow(_) => "InMemory", + CacheEntry::MemoryLiquid(_) => "LiquidMemory", + CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", + CacheEntry::DiskArrow { .. } => "OnDiskArrow", + }; + let reference_count = cached_batch.reference_count(); + let entry_id = ParquetArrayID::from(*entry_id); + writer + .append_entry( + &entry_id.display_path(), + entry_id.row_group_id_inner(), + entry_id.column_id_inner(), + entry_id.batch_id_inner() * self.batch_size() as u64, + row_count, + memory_size as u64, + cache_type, + reference_count as u64, + ) + .unwrap(); + }); writer.finish()?; Ok(()) @@ -204,15 +205,21 @@ mod tests { let mut row_start_id_sum = 0; let mut row_count_sum = 0; let mut memory_size_sum = 0; - for file_no in 0..8 { - let file_name = format!("test_{file_no}.parquet"); - let file = cache.register_or_get_file( - ParquetFileIdentity::new( - datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), - file_name, - ), - schema.clone(), - ); + // Held for the whole loop, not per iteration: a file id is leased and + // comes back when its handle drops, so releasing each file before + // opening the next would hand them all the same id. + let files: Vec<_> = (0..8) + .map(|file_no| { + cache.register_or_get_file( + ParquetFileIdentity::new( + datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + format!("test_{file_no}.parquet"), + ), + schema.clone(), + ) + }) + .collect(); + for file in &files { for rg in 0..8 { let row_group = file.create_row_group(rg, vec![]); for col in 0..8 { diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index a8ff2b80b..8e9f1a3b7 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -182,7 +182,6 @@ impl Morselizer for LiquidMorselizer { metrics, file_pruner, reader_factory, - batch_size: self.batch_size, logical_file_schema, output_schema, projection, @@ -232,7 +231,6 @@ struct PreparedLiquidOpen { metrics: LiquidFileMetrics, file_pruner: Option, reader_factory: Arc, - batch_size: usize, logical_file_schema: SchemaRef, output_schema: SchemaRef, projection: ProjectionExprs, @@ -599,7 +597,6 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result Option { let mut kind = None; - cache.for_each_entry(|entry_id, entry| { + cache.for_each_entry(|entry_id, _, entry| { if entry_id == id { kind = Some(CachedBatchType::from(entry)); } diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 9e6ded1f0..2b6f55f18 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -58,7 +58,6 @@ struct LiquidCacheReaderInner { } pub(crate) struct LiquidCacheReaderConfig { - pub(crate) batch_size: usize, pub(crate) selection: RowSelection, pub(crate) row_filter: Option, pub(crate) cached_row_group: CachedRowGroupRef, @@ -530,7 +529,6 @@ mod tests { use std::sync::Arc; struct TestRowGroup { - batch_size: usize, row_group: CachedRowGroupRef, schema: SchemaRef, fallback: ParquetFallbackConfig, @@ -547,7 +545,6 @@ mod tests { impl TestRowGroup { fn reader(&self, request: ReaderRequest) -> LiquidCacheReader { LiquidCacheReader::new(LiquidCacheReaderConfig { - batch_size: self.batch_size, selection: request.selection, row_filter: request.row_filter, cached_row_group: Arc::clone(&self.row_group), @@ -626,7 +623,6 @@ mod tests { } TestRowGroup { - batch_size, row_group, schema, fallback: ParquetFallbackConfig { diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index 271f53fff..035a0a4e5 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -31,7 +31,6 @@ pub(crate) struct LiquidRowGroupPlanner { pub(crate) row_filter: Option, pub(crate) cached_file: CachedFileRef, pub(crate) projection: ProjectionMask, - pub(crate) batch_size: usize, pub(crate) stream_schema: SchemaRef, pub(crate) output_schema: SchemaRef, pub(crate) projector: Arc, @@ -151,7 +150,6 @@ impl LiquidRowGroupPlanner { Some(LiquidRowGroupMorsel { config: LiquidCacheReaderConfig { - batch_size: self.batch_size, selection, row_filter: self.row_filter.clone(), cached_row_group, From e048b66f65a80ab9c530bc2d14bd064549c6d8bb Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 16:55:51 +0530 Subject: [PATCH 18/24] chore: add default CODEOWNERS Ports hotdata-dev#14. Fork-local: this rides with our stack and is never offered upstream. --- .github/CODEOWNERS | 1 + 1 file changed, 1 insertion(+) create mode 100644 .github/CODEOWNERS diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 000000000..43693f296 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1 @@ +* @hotdata-dev/engineers From c70c6fdf9b8fb3840b3d6a9b0cfdee9185764f85 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 18:13:06 +0530 Subject: [PATCH 19/24] fix(cache): reclaim disk a write strands when a key changes hands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found by review of the file-id commit; three defects, one cause. Scoping store keys by identity means an object stops being reachable through the index the moment its key changes hands — `release_disk` is driven by an index entry, and by then no index entry names it. Under the shared key this port replaced, the next write simply overwrote the same object, so none of this had to be handled and there was nothing to copy. Three paths stranded bytes and objects for the life of the process: A rewrite dropped as stale had already written to the store. `try_insert` refunded its memory and returned Ok, and the bytes stayed charged with no entry naming them. It now reports what it left behind and each caller — evict, flush, insert — deletes the object and releases the reservation. An `Owned` takeover replaced a disk-resident entry and dropped it. The index now hands the displaced entry back with the identity that held it, for the same treatment. Only across identities: a write under the identity that already held the key addresses the same object and its put overwrote it, so reclaiming there would delete the bytes just written — which the policy snapshots caught. `remove_disk_entry` took an identity to address the object but removed the index record unchecked, so a caller holding a stale identity deleted the current owner's record while deleting the old identity's object. That is the window fork #49 closed on the write path, left open on the removal path. Removal is identity-checked now, and re-checked after the tree removal rather than only before. Two tests cover the reclamation and the refused removal. Disabling the reclamation strands 968 bytes in the first of them. --- src/core/src/cache/core.rs | 246 +++++++++++++++++++++++++++++++----- src/core/src/cache/index.rs | 142 ++++++++++++++++++--- 2 files changed, 338 insertions(+), 50 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 38a2dc705..4477f5204 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -73,6 +73,51 @@ pub enum PrefetchResult { } /// Builder returned by [`LiquidCache::insert`] for configuring cache writes. +/// Disk an insert left for its caller to reclaim. +/// +/// A store object is addressed by entry id *and* the identity that wrote it, so +/// an object stops being reachable through the index the moment the key changes +/// hands or the write that produced it is dropped. Nothing else will ever delete +/// it or release its bytes: `release_disk` is driven by an index entry, and by +/// then no index entry names it. +#[derive(Default)] +struct DiskResidue { + /// `(identity, disk_bytes)` of a disk-resident entry this write displaced. + displaced: Option<(u64, usize)>, + /// The write itself did not land, so whatever the caller already wrote to + /// the store under its own identity is unreachable too. + dropped: bool, +} + +impl DiskResidue { + fn dropped() -> Self { + Self { + displaced: None, + dropped: true, + } + } + + /// Only a displacement by a *different* identity strands anything. + /// + /// The store key is `(entry id, identity)`, so a write under the identity + /// that already held the key addresses the very same object: the put + /// overwrote it, and reclaiming here would delete the bytes just written. + /// Across identities the two keys differ, and the old object becomes + /// unreachable the moment the index stops naming it. + fn displacing(displaced: Option<&(u64, Arc)>, writer: u64) -> Self { + Self { + displaced: displaced + .filter(|(identity, _)| *identity != writer) + .and_then(|(identity, entry)| match entry.as_ref() { + CacheEntry::DiskLiquid { disk_bytes, .. } + | CacheEntry::DiskArrow { disk_bytes, .. } => Some((*identity, *disk_bytes)), + CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => None, + }), + dropped: false, + } + } +} + impl LiquidCache { /// Return current cache statistics: counts and resource usage. pub fn stats(&self) -> CacheStats { @@ -282,12 +327,15 @@ impl LiquidCache { .await { Ok(()) => { - self.try_insert( - entry_id, - WriteIdentity::Rewrite(flush_identity), - CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), - ) - .expect("failed to insert disk arrow entry"); + let residue = self + .try_insert( + entry_id, + WriteIdentity::Rewrite(flush_identity), + CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), + ) + .expect("failed to insert disk arrow entry"); + self.settle(entry_id, residue, Some((flush_identity, disk_bytes))) + .await; } Err(CacheFull) => self.drop_memory_entry(entry_id, &batch), } @@ -305,15 +353,18 @@ impl LiquidCache { .await { Ok(()) => { - self.try_insert( - entry_id, - WriteIdentity::Rewrite(flush_identity), - CacheEntry::disk_liquid( - liquid_array.original_arrow_data_type(), - disk_bytes, - ), - ) - .expect("failed to insert disk liquid entry"); + let residue = self + .try_insert( + entry_id, + WriteIdentity::Rewrite(flush_identity), + CacheEntry::disk_liquid( + liquid_array.original_arrow_data_type(), + disk_bytes, + ), + ) + .expect("failed to insert disk liquid entry"); + self.settle(entry_id, residue, Some((flush_identity, disk_bytes))) + .await; } Err(CacheFull) => self.drop_memory_entry(entry_id, &batch), } @@ -376,9 +427,16 @@ impl LiquidCache { identity: WriteIdentity, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { + // Set once this loop spills the entry to disk itself: those bytes are the + // caller's own write, so a rewrite dropped as stale has to reclaim them. + let mut wrote = None; loop { - let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { - return Ok(()); + let not_inserted = match self.try_insert(entry_id, identity, batch_to_cache) { + Ok(residue) => { + self.settle(entry_id, residue, wrote).await; + return Ok(()); + } + Err(not_inserted) => not_inserted, }; self.trace(InternalEvent::InsertFailed { entry: entry_id, @@ -393,6 +451,11 @@ impl LiquidCache { let on_disk_batch = self .write_in_memory_batch_to_disk(entry_id, identity.value(), not_inserted) .await?; + if let CacheEntry::DiskLiquid { disk_bytes, .. } + | CacheEntry::DiskArrow { disk_bytes, .. } = &on_disk_batch + { + wrote = Some((identity.value(), *disk_bytes)); + } batch_to_cache = on_disk_batch; continue; } @@ -441,9 +504,9 @@ impl LiquidCache { entry_id: EntryID, identity: WriteIdentity, to_insert: CacheEntry, - ) -> Result<(), CacheEntry> { + ) -> Result { let new_memory_size = to_insert.memory_usage_bytes(); - let cached_batch_type = if let Some(entry) = self.index.get(&entry_id) { + let (cached_batch_type, outcome) = if let Some(entry) = self.index.get(&entry_id) { let old_memory_size = entry.memory_usage_bytes(); if self .budget @@ -453,26 +516,29 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - // A rewrite whose key changed hands since it was read is dropped by - // the index. Give the reservation back rather than counting memory - // for an entry that was never stored. - if !self.index.insert(&entry_id, identity, to_insert) { + let outcome = self.index.insert(&entry_id, identity, to_insert); + if !outcome.stored { + // A rewrite whose key changed hands since it was read. Give the + // reservation back rather than counting memory for an entry that + // was never stored, and tell the caller its disk write is now + // unreachable. self.budget .try_update_memory_usage(new_memory_size, old_memory_size) .ok(); - return Ok(()); + return Ok(DiskResidue::dropped()); } - batch_type + (batch_type, outcome) } else { if self.budget.try_reserve_memory(new_memory_size).is_err() { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - if !self.index.insert(&entry_id, identity, to_insert) { + let outcome = self.index.insert(&entry_id, identity, to_insert); + if !outcome.stored { self.budget.try_update_memory_usage(new_memory_size, 0).ok(); - return Ok(()); + return Ok(DiskResidue::dropped()); } - batch_type + (batch_type, outcome) }; self.trace(InternalEvent::InsertSuccess { @@ -482,7 +548,51 @@ impl LiquidCache { self.cache_policy .notify_insert(&entry_id, cached_batch_type); - Ok(()) + Ok(DiskResidue::displacing( + outcome.displaced.as_ref(), + identity.value(), + )) + } + + /// Delete a store object nothing can reach any more and give its bytes back. + /// + /// Reached on the two paths where an object outlives the index entry that + /// named it: a write dropped as stale after its bytes were already written, + /// and an entry displaced by a write under a different identity. Both are + /// consequences of the store key carrying the identity — under a shared key + /// the next write simply overwrote the same object. + async fn reclaim_orphaned_disk(&self, entry_id: EntryID, identity: u64, disk_bytes: usize) { + match self + .store + .remove(&entry_id_to_key(&entry_id, identity)) + .await + { + // `false` means the object was already gone, which is fine: the + // bytes still have to be given back either way. + Ok(_) | Err(t4::Error::NotFound) => {} + Err(error) => panic!("orphan remove failed: {error}"), + } + self.budget.release_disk(disk_bytes); + self.trace(InternalEvent::DiskEvict { + entry: entry_id, + bytes: disk_bytes, + }); + } + + /// Reclaim whatever an insert left unreachable, including the caller's own + /// write when it was dropped as stale. + /// + /// `wrote` is the (identity, bytes) the caller put in the store before the + /// insert, if any. + async fn settle(&self, entry_id: EntryID, residue: DiskResidue, wrote: Option<(u64, usize)>) { + if let Some((identity, bytes)) = residue.displaced { + self.reclaim_orphaned_disk(entry_id, identity, bytes).await; + } + if residue.dropped + && let Some((identity, bytes)) = wrote + { + self.reclaim_orphaned_disk(entry_id, identity, bytes).await; + } } fn drop_memory_entry(&self, entry_id: EntryID, _expected: &CacheEntry) { @@ -503,7 +613,11 @@ impl LiquidCache { } async fn remove_disk_entry(&self, entry_id: EntryID, removed_identity: u64) { - let Some(removed) = self.index.remove(&entry_id) else { + // Checked: the caller read this entry earlier, and between then and now + // the key can change hands. Removing the new owner's record would strand + // its store object while releasing a byte count taken from the record + // just destroyed. + let Some(removed) = self.index.remove_checked(&entry_id, removed_identity) else { return; }; let disk_bytes = match removed.as_ref() { @@ -577,12 +691,19 @@ impl LiquidCache { entry: new_batch, bytes_to_write, } => { + // Remember what went to the store: if the rewrite is then + // dropped as stale, these bytes are unreachable and have to + // be reclaimed here. + let mut wrote = None; if let Some(bytes_to_write) = bytes_to_write { + let len = bytes_to_write.len(); self.write_batch_to_disk(victim, identity, &new_batch, bytes_to_write) .await?; + wrote = Some((identity, len)); } match self.try_insert(victim, WriteIdentity::Rewrite(identity), new_batch) { - Ok(()) => { + Ok(residue) => { + self.settle(victim, residue, wrote).await; break; } Err(batch) => { @@ -986,6 +1107,69 @@ mod tests { } } + /// A rewrite that loses its key must not leave its disk write behind. + /// + /// The bytes were already in the store when the index refused the write, and + /// the store key carries the identity that wrote them, so nothing reachable + /// through the index names them afterwards: neither the object nor its share + /// of `used_disk_bytes` would ever come back. + #[tokio::test] + async fn a_dropped_rewrite_reclaims_the_disk_it_already_wrote() { + let store = create_cache_store(10 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(1usize); + + // The owner caches an entry and flushes it, so a disk object exists. + store + .insert(entry_id, 7, create_test_arrow_array(64)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let after_flush = store.budget.disk_usage_bytes(); + assert!(after_flush > 0, "flush must have written bytes"); + + // Another identity takes the key over, so the earlier owner's rewrite is + // now stale. Its disk object is unreachable and must be reclaimed. + store + .insert(entry_id, 9, create_test_arrow_array(64)) + .await + .unwrap(); + assert_eq!( + store.budget.disk_usage_bytes(), + 0, + "the displaced owner's disk bytes must be released, not stranded" + ); + assert!( + store.get(&entry_id, 9).await.is_some(), + "the new owner must still read its own entry" + ); + assert!( + store.get(&entry_id, 7).await.is_none(), + "the displaced owner must not read the new owner's rows" + ); + } + + /// The removal path is identity-checked: a caller that read an entry earlier + /// must not destroy the record of whoever holds the key now. + #[tokio::test] + async fn removing_a_disk_entry_under_a_stale_identity_is_refused() { + let store = create_cache_store(10 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(2usize); + + store + .insert(entry_id, 3, create_test_arrow_array(64)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + + // A stale identity tries to evict it. Nothing of the current owner's may + // be touched. + store.remove_disk_entry(entry_id, 999).await; + assert!( + store.get(&entry_id, 3).await.is_some(), + "a stale remove must leave the current owner's entry readable" + ); + } + #[tokio::test] async fn test_basic_cache_operations() { // Test basic insert, get, and size tracking in one test diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 4ec252484..54c642330 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -75,6 +75,30 @@ impl WriteIdentity { } } +/// What an [`ArtIndex::insert`] did. +/// +/// `displaced` is the entry this write replaced, with the identity that was +/// recorded against it. It is handed back rather than dropped here because a +/// store object is addressed by entry id *and* identity +/// (`entry_id_to_key`): once another identity holds the key, nothing reachable +/// through the index names the old object any more, so the caller has to delete +/// it and give its disk bytes back or both leak for the life of the process. +pub(crate) struct InsertOutcome { + /// Whether the batch was stored. False only for a stale rewrite. + pub(crate) stored: bool, + /// The entry this write replaced, and the identity that held it. + pub(crate) displaced: Option<(u64, Arc)>, +} + +impl InsertOutcome { + fn dropped() -> Self { + Self { + stored: false, + displaced: None, + } + } +} + pub(crate) struct ArtIndex { art: CongeeArc, entry_count: AtomicUsize, @@ -160,15 +184,17 @@ impl ArtIndex { self.get_checked(entry_id, identity).is_some() } - /// Store `batch` under `entry_id`, returning whether it was stored. + /// Store `batch` under `entry_id`, reporting what the write did. /// - /// See [`WriteIdentity`] for the two kinds of write and why they differ. + /// See [`WriteIdentity`] for the two kinds of write and why they differ, and + /// [`InsertOutcome`] for why a displaced entry has to be handed back rather + /// than dropped here. pub(crate) fn insert( &self, entry_id: &EntryID, identity: WriteIdentity, batch: CacheEntry, - ) -> bool { + ) -> InsertOutcome { let guard = self.art.pin(); let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); let identity = match (identity, existing_identity) { @@ -184,19 +210,65 @@ impl ArtIndex { // It does not: the entry was taken over or removed while this // rewrite was in flight, so the payload belongs to a source that // no longer owns the key. Drop it. - (WriteIdentity::Rewrite(_), _) => return false, + (WriteIdentity::Rewrite(_), _) => return InsertOutcome::dropped(), }; let existing = self .art .insert(*entry_id, Slot::new(identity, batch), &guard) .expect("Insertion failed"); - match existing { - Some(replaced) => drop(replaced.take()), + let displaced = match existing { + Some(replaced) => { + let was = replaced.identity; + replaced.take().map(|entry| (was, entry)) + } None => { self.entry_count.fetch_add(1, Ordering::Relaxed); + None + } + }; + InsertOutcome { + stored: true, + displaced, + } + } + + /// Remove an entry only if `identity` is the one that holds it. + /// + /// The unchecked [`Self::remove`] is for maintenance acting on whatever + /// occupies a key. A caller that read an entry earlier and then removes it + /// must come through here: between its read and its removal the key can + /// change hands, and removing the new owner's record would strand that + /// owner's store object while releasing a byte count taken from the record + /// it just destroyed. + pub(crate) fn remove_checked( + &self, + entry_id: &EntryID, + identity: u64, + ) -> Option> { + let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + let removed = self.art.remove(*entry_id, &guard)?; + if removed.identity != identity { + // Lost the race between the check and the removal: put it back + // rather than destroying a record this caller has no claim on. + if let Some(entry) = removed.take() { + self.art + .insert( + *entry_id, + Slot::new(removed.identity, (*entry).clone()), + &guard, + ) + .expect("Insertion failed"); } + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; } - true + self.entry_count.fetch_sub(1, Ordering::Relaxed); + removed.take() } pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { @@ -334,7 +406,11 @@ mod tests { let store = ArtIndex::new(); let key: EntryID = EntryID::from(7); - assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + assert!( + store + .insert(&key, WriteIdentity::Owned(1), create_test_array(100)) + .stored + ); // The colliding file asks for the same key and is told nothing is there. assert!(store.get_checked(&key, 2).is_none()); @@ -346,7 +422,11 @@ mod tests { // The colliding file takes the key over. It has to: it cannot read // what is there, so leaving it would cost the key to both of them. - assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + assert!( + store + .insert(&key, WriteIdentity::Owned(2), create_test_array(200)) + .stored + ); assert!( store.get_checked(&key, 1).is_none(), "the displaced file reads a miss, never the other file's rows" @@ -369,19 +449,31 @@ mod tests { let key: EntryID = EntryID::from(11); // File A caches, and something begins rewriting that entry. - assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + assert!( + store + .insert(&key, WriteIdentity::Owned(1), create_test_array(100)) + .stored + ); let (observed, _read) = store.get_with_identity(&key).unwrap(); assert_eq!(observed, 1); // File B takes the key over while that rewrite is in flight. - assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + assert!( + store + .insert(&key, WriteIdentity::Owned(2), create_test_array(200)) + .stored + ); // The rewrite lands too late and must be dropped, not relabelled. - assert!(!store.insert( - &key, - WriteIdentity::Rewrite(observed), - create_test_array(100) - )); + assert!( + !store + .insert( + &key, + WriteIdentity::Rewrite(observed), + create_test_array(100) + ) + .stored + ); match store.get_checked(&key, 2).unwrap().as_ref() { CacheEntry::MemoryArrow(array) => assert_eq!( @@ -402,15 +494,27 @@ mod tests { let store = ArtIndex::new(); let key: EntryID = EntryID::from(9); - assert!(store.insert(&key, WriteIdentity::Owned(5), create_test_array(10))); - assert!(store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(20))); + assert!( + store + .insert(&key, WriteIdentity::Owned(5), create_test_array(10)) + .stored + ); + assert!( + store + .insert(&key, WriteIdentity::Rewrite(5), create_test_array(20)) + .stored + ); assert!( store.get_checked(&key, 5).is_some(), "rewriting in place kept the identity" ); store.remove(&key); - assert!(!store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(30))); + assert!( + !store + .insert(&key, WriteIdentity::Rewrite(5), create_test_array(30)) + .stored + ); assert!(store.get(&key).is_none()); assert_eq!(store.entry_count(), 0); } From 5769ea172aabd54d46971e73ae846c001fd1ff93 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 23 Sep 2026 22:54:25 +0530 Subject: [PATCH 20/24] fix(reader): keep column-free conjuncts in the row filter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the code half of hotdata-dev#20, which the sync dropped. `749b6ef` brought #20 and #40 across as tests only, on the finding that upstream had superseded both. That was right for #40 and wrong for #20: upstream refuses nested columns and columns outside the file schema through `try_pushdown_filters`, which is what those tests cover, but it still drops a conjunct that references no column at all. Nothing tested that case, so the gap looked closed. A literal `Boolean(NULL)` conjunct is exactly what expression simplification leaves behind: `NOT (s = s)` becomes `s IS NULL AND NULL`. `pushdown_columns` returns an empty column set for it, the `is_empty()` bail dropped it, and by then DataFusion has removed the `FilterExec` on the strength of the predicate being fully pushed down — so the scan is the only place it is applied and it applies a strictly weaker filter. `SELECT s FROM t WHERE NOT (s = s)` returned every row with a NULL `s` instead of none. Keep the conjunct. A column-free candidate now builds with an empty projection mask, and the cached path evaluates it against a batch that carries only the row count the selection implies — `RecordBatch` needs that explicitly, since no array is there to imply it. Found by runtimedb's query fuzzer as a ternary-partition violation: the union of `P`, `NOT P` and `P IS NULL` returned 14664 rows where the unfiltered scan returned 12000. The new tests cover both that identity and the direct `NOT (s = s)` case, cold and warm; both fail without this change. The defect is upstream's and predates the sync — our fork carried this fix, upstream still does not. --- .../src/tests/constant_conjunct.rs | 131 ++++++++++++++++++ src/datafusion-local/src/tests/mod.rs | 1 + src/datafusion/src/cache/mod.rs | 13 +- .../src/reader/plantime/row_filter.rs | 9 +- 4 files changed, 149 insertions(+), 5 deletions(-) create mode 100644 src/datafusion-local/src/tests/constant_conjunct.rs diff --git a/src/datafusion-local/src/tests/constant_conjunct.rs b/src/datafusion-local/src/tests/constant_conjunct.rs new file mode 100644 index 000000000..6641372b9 --- /dev/null +++ b/src/datafusion-local/src/tests/constant_conjunct.rs @@ -0,0 +1,131 @@ +//! A pushed-down conjunct that references no column. +//! +//! Expression simplification turns `NOT (s = s)` into `s IS NULL AND NULL`: a +//! column conjunct and a literal `Boolean(NULL)` one. The literal reads no +//! column, and `build_row_filter` used to drop such a conjunct. Since DataFusion +//! removes the `FilterExec` when it pushes a predicate down, the scan is the only +//! place the predicate is applied, so dropping a conjunct *widens* the filter and +//! rows that cannot match come back. +//! +//! The visible symptom is a three-way partition that does not reconstruct the +//! scan: for a predicate `P`, `WHERE P`, `WHERE NOT P` and `WHERE P IS NULL` must +//! together return each row exactly once. Rows whose `P` is NULL were returned by +//! both `WHERE NOT P` and `WHERE P IS NULL`. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Array, Float64Array, Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// The predicate under test. `s = s` is NULL wherever `s` is NULL, so a row with +/// a NULL `s`, an `f` above the threshold and an `id` outside the range makes the +/// whole predicate NULL. +const P: &str = "((t1.f <= 333.0 OR t1.s = t1.s) OR t1.id BETWEEN 1 AND 7)"; + +/// 12000 rows, more than the 8192-row default cache batch, so the scan spans more +/// than one cached batch. Every third row has a NULL `s`; `f` cycles 1..1000, so +/// two thirds of those NULL rows sit above the 333.0 threshold. +fn write_t1(path: &Path) { + let rows = 12000i64; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("f", DataType::Float64, true), + Field::new("s", DataType::Utf8, true), + ])); + let id: Int64Array = (1..=rows).collect::>().into(); + let f: Float64Array = (1..=rows) + .map(|i| Some((i % 1000) as f64)) + .collect::>() + .into(); + let s: StringArray = (1..=rows) + .map(|i| (i % 3 != 0).then(|| format!("str{}", i % 17))) + .collect::>() + .into(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(id), Arc::new(f), Arc::new(s)]).unwrap(); + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +async fn liquid_ctx(dir: &Path) -> SessionContext { + let parquet = dir.join("t1.parquet"); + write_t1(&parquet); + let cache_dir = dir.join("cache"); + std::fs::create_dir_all(&cache_dir).unwrap(); + let (ctx, _cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir) + .build(SessionConfig::new()) + .await + .unwrap(); + ctx.register_parquet( + "t1", + parquet.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + ctx +} + +/// The `s` column as a sorted multiset, NULL rendered as ``. `s` is +/// projected as a string view, hence the cast. +async fn s_values(ctx: &SessionContext, sql: &str) -> Vec { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let mut out = Vec::new(); + for batch in batches { + let column = arrow::compute::cast(batch.column(0), &DataType::Utf8).unwrap(); + let column = column.as_any().downcast_ref::().unwrap(); + for row in 0..column.len() { + out.push(match column.is_null(row) { + true => "".to_string(), + false => column.value(row).to_string(), + }); + } + } + out.sort(); + out +} + +/// `NOT (s = s)` is NULL where `s` is NULL and FALSE everywhere else, so it +/// matches no row. Simplification leaves it as `s IS NULL AND NULL`, and dropping +/// the literal conjunct returns every row with a NULL `s`. +#[tokio::test] +async fn constant_null_conjunct_is_still_applied() { + let dir = TempDir::new().unwrap(); + let ctx = liquid_ctx(dir.path()).await; + + // Cold reads through the source and fills the cache; warm is served from it, + // a separate evaluation path. + for pass in ["cold", "warm"] { + let rows = s_values(&ctx, "SELECT t1.s FROM t1 WHERE NOT (t1.s = t1.s)").await; + assert_eq!(rows, Vec::::new(), "{pass}"); + } +} + +/// `WHERE P`, `WHERE NOT P` and `WHERE P IS NULL` partition the table: together +/// they must return exactly the rows of the unfiltered scan, each once. +#[tokio::test] +async fn three_way_partition_reconstructs_the_scan() { + let dir = TempDir::new().unwrap(); + let ctx = liquid_ctx(dir.path()).await; + + for pass in ["cold", "warm"] { + let unfiltered = s_values(&ctx, "SELECT t1.s FROM t1").await; + + let mut partitioned = s_values(&ctx, &format!("SELECT t1.s FROM t1 WHERE {P}")).await; + partitioned.extend(s_values(&ctx, &format!("SELECT t1.s FROM t1 WHERE NOT {P}")).await); + partitioned.extend(s_values(&ctx, &format!("SELECT t1.s FROM t1 WHERE {P} IS NULL")).await); + partitioned.sort(); + + assert_eq!(partitioned.len(), unfiltered.len(), "{pass}"); + assert_eq!(partitioned, unfiltered, "{pass}"); + } +} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 15d669466..ef8197d3e 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -20,6 +20,7 @@ use datafusion::{ use crate::LiquidCacheLocalBuilder; mod batch_size_alignment; +mod constant_conjunct; mod date_optimizer; mod filter_limit; mod nested_filter; diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 162774eb2..4501158f2 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -7,7 +7,7 @@ use crate::sync::RwLock; mod file_id; use ahash::AHashMap; -use arrow::array::{BooleanArray, RecordBatch}; +use arrow::array::{BooleanArray, RecordBatch, RecordBatchOptions}; use arrow::buffer::BooleanBuffer; use arrow_schema::{ArrowError, Field, Schema, SchemaRef}; use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -258,6 +258,17 @@ impl CachedRowGroup { } } } + // A conjunct that reads no column still has to be evaluated, over a batch + // that carries only the row count the selection implies. + if column_ids.is_empty() { + let options = + RecordBatchOptions::new().with_row_count(Some(selection.count_set_bits())); + let record_batch = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), Vec::new(), &options) + .ok()?; + return Some(predicate.evaluate(record_batch)); + } + // Otherwise, we need to first convert the data into arrow arrays. let mut arrays = Vec::new(); let mut fields = Vec::new(); diff --git a/src/datafusion/src/reader/plantime/row_filter.rs b/src/datafusion/src/reader/plantime/row_filter.rs index 54c778cec..5d17d410b 100644 --- a/src/datafusion/src/reader/plantime/row_filter.rs +++ b/src/datafusion/src/reader/plantime/row_filter.rs @@ -275,10 +275,11 @@ impl FilterCandidateBuilder { return Ok(None); }; - if required_indices_into_file_schema.is_empty() { - return Ok(None); - } - + // A conjunct that references no column - a literal `NULL` or `false` left + // behind by expression simplification, for instance - is still a conjunct. + // Dropping it here widens the filter, because by the time this runs + // DataFusion has removed the `FilterExec` on the assumption the predicate + // was fully pushed down, so the scan is the only place it is applied. let projected_file_schema = Arc::new( self.file_schema .project(&required_indices_into_file_schema)?, From fb7bf08c3138f71adf56bcfce231db9e00b66468 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 24 Sep 2026 01:01:16 +0530 Subject: [PATCH 21/24] fix(cache): reclaim the disk copy a memory entry displaces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A read that materializes a disk entry replaces it with a memory one, and a caller overwriting an entry does the same. Neither puts anything in the store, so the object under `(entry id, identity)` and its share of `used_disk_bytes` survive a replacement that nothing else will ever release: `release_disk` is only reached from an index entry, and by then no index entry names it. `DiskResidue::displacing` treated every same-identity displacement as an overwrite of the same object. That holds only when the write is itself disk-resident — those are the paths that put bytes under that key first. A memory entry displacing a disk one strands the copy instead, so the next spill reserves the same byte count again for an object the put overwrites. Over a fixed working set the disk tally climbs with every read/spill round, and once it reaches the limit the tier evicts entries that are genuinely there. Deleting the copy rather than keeping it for the next spill to reuse: the spill path serializes and writes unconditionally, so reuse needs a record of which form of the entry is on disk, kept correct across overwrite, transcode and takeover. Releasing the reservation without deleting is worse than either — the object then outlives any entry that could ever name it, and the store grows with nothing counting it. The datafusion-local snapshots move because `disk_evictions` now counts these reclaims; no result, plan or IO count changes. --- src/core/src/cache/core.rs | 331 +++++++++++++++++- ...he__tests__policies__default_policies.snap | 3 + ...datafusion_local__tests__os_selection.snap | 2 +- ...usion_local__tests__referer_filtering.snap | 2 +- ...on_local__tests__url_prefix_filtering.snap | 2 +- ...al__tests__url_selection_and_ordering.snap | 2 +- 6 files changed, 325 insertions(+), 17 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 4477f5204..db26efe2a 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -97,17 +97,28 @@ impl DiskResidue { } } - /// Only a displacement by a *different* identity strands anything. + /// A displaced disk entry's object survives this insert unless the insert + /// put its own bytes over it. /// - /// The store key is `(entry id, identity)`, so a write under the identity - /// that already held the key addresses the very same object: the put - /// overwrote it, and reclaiming here would delete the bytes just written. - /// Across identities the two keys differ, and the old object becomes - /// unreachable the moment the index stops naming it. - fn displacing(displaced: Option<&(u64, Arc)>, writer: u64) -> Self { + /// The store key is `(entry id, identity)`, so only a disk-resident entry + /// written under the identity that already held the key addresses the very + /// same object: there the put overwrote it, and reclaiming would delete the + /// bytes just written. Otherwise the object is left behind — a different + /// identity addresses a different key, and an entry that lives in memory + /// wrote nothing at all — and it becomes unreachable the moment the index + /// stops naming it. + fn displacing( + displaced: Option<&(u64, Arc)>, + writer: u64, + written: CachedBatchType, + ) -> Self { + let overwrites_in_place = matches!( + written, + CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow + ); Self { displaced: displaced - .filter(|(identity, _)| *identity != writer) + .filter(|(identity, _)| !(overwrites_in_place && *identity == writer)) .and_then(|(identity, entry)| match entry.as_ref() { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => Some((*identity, *disk_bytes)), @@ -551,16 +562,17 @@ impl LiquidCache { Ok(DiskResidue::displacing( outcome.displaced.as_ref(), identity.value(), + cached_batch_type, )) } /// Delete a store object nothing can reach any more and give its bytes back. /// - /// Reached on the two paths where an object outlives the index entry that - /// named it: a write dropped as stale after its bytes were already written, - /// and an entry displaced by a write under a different identity. Both are - /// consequences of the store key carrying the identity — under a shared key - /// the next write simply overwrote the same object. + /// Reached on the paths where an object outlives the index entry that named + /// it: a write dropped as stale after its bytes were already written, an + /// entry displaced by a write under a different identity, and a disk entry + /// replaced by a memory one — hydration, or a caller overwriting the value + /// — which puts nothing in the store and so leaves the old object whole. async fn reclaim_orphaned_disk(&self, entry_id: EntryID, identity: u64, disk_bytes: usize) { match self .store @@ -1107,6 +1119,299 @@ mod tests { } } + /// Hydrating a disk entry must not charge its bytes to the disk budget + /// twice. + /// + /// The read that materializes a `DiskArrow`/`DiskLiquid` entry replaces it + /// with a memory entry. The store object under `(entry id, identity)` and + /// its share of `used_disk_bytes` outlive that replacement, so the next + /// spill reserves the same byte count again for an object the put simply + /// overwrites. Over a fixed working set, repeated read/spill rounds make + /// `disk_usage_bytes` climb without a byte more being written. + #[tokio::test] + async fn hydrating_a_disk_entry_does_not_recharge_its_disk_bytes() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(500usize); + let array = create_test_arrow_array(1024); + + store.insert(entry_id, 0, array.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + let charged = store.budget.disk_usage_bytes(); + assert!(charged > 0, "flush must have written bytes"); + + let mut usage = vec![charged]; + for _ in 0..3 { + // Reading a disk entry hydrates it back into memory ... + let read = store.get(&entry_id, 0).await.expect("present"); + assert_eq!(read.as_ref(), array.as_ref()); + assert!(matches!( + store.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::MemoryArrow(_) + )); + // ... and the next flush spills the very same bytes again. + store.flush_all_to_disk().await.unwrap(); + usage.push(store.budget.disk_usage_bytes()); + } + + assert_eq!( + usage, + vec![charged; 4], + "one entry of a fixed size occupies the same disk across read/spill rounds" + ); + } + + /// Every byte counted against the disk budget must belong to an index + /// entry that is on disk. Anything else can never be released: + /// `release_disk` is only ever reached from an index entry. + fn charged_disk_bytes_match_the_index(cache: &LiquidCache) -> (usize, usize) { + let mut named = 0usize; + cache.for_each_entry(|_, _, entry| match entry { + CacheEntry::DiskLiquid { disk_bytes, .. } + | CacheEntry::DiskArrow { disk_bytes, .. } => named += *disk_bytes, + CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => {} + }); + (named, cache.budget.disk_usage_bytes()) + } + + /// Overwriting a disk-resident entry under the identity that already holds + /// the key must not leave the superseded copy charged. + /// + /// The store key is `(entry id, identity)`, so the object the old entry + /// named is still there — but this insert wrote nothing to the store, so + /// it did not overwrite it, and the index no longer names it. + #[tokio::test] + async fn overwriting_a_disk_entry_releases_the_superseded_copy() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(501usize); + + store + .insert(entry_id, 5, create_test_arrow_array(1024)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!( + store.budget.disk_usage_bytes() > 0, + "flush must have written" + ); + + // Same identity, new value: the index entry becomes a memory one. + store + .insert(entry_id, 5, create_test_arrow_array(2048)) + .await + .unwrap(); + + let (named, charged) = charged_disk_bytes_match_the_index(&store); + assert_eq!( + charged, named, + "the superseded copy is still charged but no index entry names it" + ); + } + + /// An overwrite must never be readable as the value it replaced, and the + /// form recorded in the index must be the form the store actually holds. + #[tokio::test] + async fn an_overwritten_entry_never_reads_back_the_superseded_bytes() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(502usize); + let first = create_test_arrow_array(1024); + let second: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( + (0..2048).map(|v| v + 1_000_000), + )); + + store.insert(entry_id, 5, first.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!(matches!( + store.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::DiskArrow { .. } + )); + + store.insert(entry_id, 5, second.clone()).await.unwrap(); + assert_eq!( + store.get(&entry_id, 5).await.expect("present").as_ref(), + second.as_ref(), + "the overwrite must be what a read returns" + ); + + // Spill again: whatever form the index records has to match the bytes + // the store now holds, or the next read decodes one as the other. + store.flush_all_to_disk().await.unwrap(); + assert!( + matches!( + store.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::DiskArrow { .. } + ), + "an Arrow flush must be recorded as an Arrow copy" + ); + assert_eq!( + store.get(&entry_id, 5).await.expect("present").as_ref(), + second.as_ref(), + "reading the spilled copy must not return the superseded value" + ); + } + + /// A hydrated entry that is then transcoded and spilled must be recorded + /// as the form it was written in, not the form it was hydrated from. + #[tokio::test] + async fn a_hydrated_then_transcoded_entry_is_recorded_as_what_was_written() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(503usize); + let array = create_test_arrow_array(1024); + + store.insert(entry_id, 5, array.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + + // Hydrate the Arrow copy back into memory ... + store.get(&entry_id, 5).await.expect("present"); + // ... transcode it to liquid, then spill it as liquid. + store + .evict_victim_inner(entry_id) + .await + .expect("transcode must fit"); + assert!(matches!( + store.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::MemoryLiquid(_) + )); + store.flush_all_to_disk().await.unwrap(); + assert!( + matches!( + store.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::DiskLiquid { .. } + ), + "a liquid flush must be recorded as a liquid copy" + ); + + assert_eq!( + store.get(&entry_id, 5).await.expect("present").as_ref(), + array.as_ref(), + "a liquid stub must not be decoded over Arrow IPC bytes" + ); + let (named, charged) = charged_disk_bytes_match_the_index(&store); + assert_eq!( + charged, named, + "the Arrow copy it was hydrated from is still charged" + ); + } + + /// A flush that cannot place an entry on disk drops it. Nothing of that + /// entry may stay charged against the disk budget afterwards. + #[tokio::test] + async fn a_flush_that_drops_an_entry_leaves_no_disk_charged_to_it() { + let array = create_test_arrow_array(1024); + let one_copy = arrow_to_bytes(&array).unwrap().len(); + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .with_max_disk_bytes(one_copy) + .with_eviction_policy(Box::new(TranscodeEvict)) + .with_hydration_policy(Box::new(crate::cache::AlwaysHydrate::new())) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build() + .await; + let first = EntryID::from(504usize); + let second = EntryID::from(505usize); + + cache.insert(first, 0, array.clone()).await.unwrap(); + cache.flush_all_to_disk().await.unwrap(); + // Reading it brings it back into memory; the disk tier should now be + // empty, so the next flush has room for both entries. + cache.get(&first, 0).await.expect("present"); + + cache.insert(second, 0, array.clone()).await.unwrap(); + cache.flush_all_to_disk().await.unwrap(); + + let (named, charged) = charged_disk_bytes_match_the_index(&cache); + assert_eq!( + charged, named, + "disk charged to entries the flush dropped can never be released" + ); + } + + /// A takeover landing *during* a rewrite's `store.put` must leave the new + /// owner whole: its bytes, its record, and its share of the budget. + /// + /// The steps below are what `evict_victim_inner` does — read the entry with + /// the identity it holds, write it to the store, then swap the record — with + /// the takeover injected between the read and the write, which is the one + /// interleaving that ordering cannot be produced by calling it once. + /// + /// Three things hold it together, and the first is the decisive one: + /// `entry_id_to_key` puts the writer's identity in the store key, so the + /// stale put addresses its own object and can never reach the new owner's; + /// `WriteIdentity::Rewrite` makes the index refuse the record swap; and + /// `settle` reclaims the object the refused write had already put there. + #[tokio::test] + async fn a_takeover_during_a_rewrites_disk_write_leaves_the_new_owner_whole() { + let cache = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(600usize); + let theirs = create_test_arrow_array(1024); + let ours: ArrayRef = Arc::new(arrow::array::Int64Array::from_iter_values( + (0..512).map(|v| v + 7_000_000), + )); + + // Identity 7 caches an entry, and eviction picks it up. + cache.insert(entry_id, 7, theirs.clone()).await.unwrap(); + let (observed, _read) = cache.index().get_with_identity(&entry_id).unwrap(); + assert_eq!(observed, 7); + let stale_bytes = arrow_to_bytes(&theirs).unwrap(); + let stale_len = stale_bytes.len(); + let stale_rewrite = CacheEntry::disk_arrow(theirs.data_type().clone(), stale_len); + + // Identity 9 takes the key over and puts its own copy on disk. + cache.insert(entry_id, 9, ours.clone()).await.unwrap(); + cache.flush_all_to_disk().await.unwrap(); + let owner_bytes = match cache.index().get(&entry_id).unwrap().as_ref() { + CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, + other => panic!("expected the new owner on disk, found {other}"), + }; + + // Only now does the stale rewrite's write complete ... + cache + .write_batch_to_disk(entry_id, observed, &stale_rewrite, stale_bytes) + .await + .unwrap(); + // ... and reach the record swap. + let residue = cache + .try_insert(entry_id, WriteIdentity::Rewrite(observed), stale_rewrite) + .expect("a refused rewrite is not a failure to insert"); + cache + .settle(entry_id, residue, Some((observed, stale_len))) + .await; + + // The stale writer's object is gone; the new owner's is not. + assert!( + matches!( + cache.store.get(&entry_id_to_key(&entry_id, observed)).await, + Err(t4::Error::NotFound) + ), + "the refused rewrite must take its own write back" + ); + assert!( + cache + .store + .get(&entry_id_to_key(&entry_id, 9)) + .await + .is_ok(), + "the new owner's object must survive a stale writer" + ); + let (named, charged) = charged_disk_bytes_match_the_index(&cache); + assert_eq!(charged, named); + assert_eq!(charged, owner_bytes, "only the new owner's copy is charged"); + + // And the new owner's record still names bytes that decode to its rows. + assert!(matches!( + cache.index().get(&entry_id).unwrap().as_ref(), + CacheEntry::DiskArrow { .. } + )); + assert_eq!( + cache.get(&entry_id, 9).await.expect("present").as_ref(), + ours.as_ref(), + "the new owner must read its own rows, never the stale writer's" + ); + assert!( + cache.get(&entry_id, 7).await.is_none(), + "the displaced identity reads a miss" + ); + } + /// A rewrite that loses its key must not leave its disk write behind. /// /// The bytes were already in the store when the index refused the write, and diff --git a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap index fc702d26f..5775f0191 100644 --- a/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap +++ b/src/core/src/cache/tests/snapshots/liquid_cache__cache__tests__policies__default_policies.snap @@ -35,14 +35,17 @@ event=read entry=0 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=0 bytes=[bytes] event=hydrate entry=0 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=0 kind=MemoryLiquid +event=disk_evict entry=0 bytes=[bytes] event=read entry=1 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=1 bytes=[bytes] event=hydrate entry=1 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=1 kind=MemoryLiquid +event=disk_evict entry=1 bytes=[bytes] event=read entry=2 selection=false expr=None cached=DiskLiquid event=io_read_liquid entry=2 bytes=[bytes] event=hydrate entry=2 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=2 kind=MemoryLiquid +event=disk_evict entry=2 bytes=[bytes] event=read entry=3 selection=false expr=None cached=MemoryLiquid event=read entry=4 selection=false expr=None cached=MemoryArrow ] diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap index ffd9e6af2..212649c19 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__os_selection.snap @@ -49,6 +49,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 3 write_io_count: 0 - disk_evictions: 0 + disk_evictions: 3 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap index 6309a155a..f9bee9ede 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__referer_filtering.snap @@ -48,6 +48,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 4 write_io_count: 0 - disk_evictions: 0 + disk_evictions: 4 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap index 6e881e549..4a85ae85f 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_prefix_filtering.snap @@ -64,6 +64,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 1 write_io_count: 0 - disk_evictions: 0 + disk_evictions: 1 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap index 5dcdb137d..1923cf5a8 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__url_selection_and_ordering.snap @@ -49,6 +49,6 @@ RuntimeStatsSnapshot: try_read_liquid_calls: 0 read_io_count: 3 write_io_count: 2 - disk_evictions: 0 + disk_evictions: 3 disk_reservation_failures: 0 eval_predicate_on_liquid_failed: 0 From 5450d044ffffd6b6213fc6c9125252cd81db03df Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 24 Sep 2026 12:13:35 +0530 Subject: [PATCH 22/24] fix(cache): release a disk copy this write overwrote in place MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review found a third shape of the same defect. `DiskResidue::displacing` treated a same-identity disk write as needing no reclamation at all, on the grounds that its put landed on the very object the displaced entry named — true of the object, but not of the reservation. The superseded entry's bytes stayed counted, so one object could be charged twice. The three cases are now explicit rather than one filter: a different identity or a memory write leaves an object nothing can reach, so it is deleted and its bytes released; a same-identity disk write keeps the object it overwrote and gives back only the superseded reservation; a displaced memory entry has nothing to reclaim. Reachability: I could not construct it. Four attempts — shrinking the memory budget to 64 bytes and inserting over a disk-resident key — all ended with the entry transcoded to liquid and resident in memory, so `insert_inner` never took the spill path that would displace a disk entry with another disk entry. Every displacement I could actually produce on that path displaces a memory entry, where the old predicate was already correct. The fix is here because the case is real in the code and the correct behaviour is cheap, not because a test forced it. --- src/core/src/cache/core.rs | 88 +++++++++++++++++++++++++++++++++----- 1 file changed, 78 insertions(+), 10 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index db26efe2a..3751d1f86 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -82,8 +82,13 @@ pub enum PrefetchResult { /// then no index entry names it. #[derive(Default)] struct DiskResidue { - /// `(identity, disk_bytes)` of a disk-resident entry this write displaced. + /// `(identity, disk_bytes)` of a disk-resident entry this write displaced, + /// whose object is now unreachable and has to be deleted. displaced: Option<(u64, usize)>, + /// Bytes of a displaced disk entry whose object this write *overwrote in + /// place*. The object is current and must be kept; only the superseded + /// entry's reservation is given back, or one object is charged twice. + superseded: Option, /// The write itself did not land, so whatever the caller already wrote to /// the store under its own identity is unreachable too. dropped: bool, @@ -93,6 +98,7 @@ impl DiskResidue { fn dropped() -> Self { Self { displaced: None, + superseded: None, dropped: true, } } @@ -116,15 +122,25 @@ impl DiskResidue { written, CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow ); - Self { - displaced: displaced - .filter(|(identity, _)| !(overwrites_in_place && *identity == writer)) - .and_then(|(identity, entry)| match entry.as_ref() { - CacheEntry::DiskLiquid { disk_bytes, .. } - | CacheEntry::DiskArrow { disk_bytes, .. } => Some((*identity, *disk_bytes)), - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => None, - }), - dropped: false, + let disk_bytes = displaced.and_then(|(identity, entry)| match entry.as_ref() { + CacheEntry::DiskLiquid { disk_bytes, .. } + | CacheEntry::DiskArrow { disk_bytes, .. } => Some((*identity, *disk_bytes)), + CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => None, + }); + // Same identity and a disk-resident write means this put landed on the + // very object the displaced entry named: keep the object, give back only + // its reservation. Anything else leaves an object nothing can reach. + match disk_bytes { + Some((identity, bytes)) if overwrites_in_place && identity == writer => Self { + displaced: None, + superseded: Some(bytes), + dropped: false, + }, + other => Self { + displaced: other, + superseded: None, + dropped: false, + }, } } } @@ -600,6 +616,11 @@ impl LiquidCache { if let Some((identity, bytes)) = residue.displaced { self.reclaim_orphaned_disk(entry_id, identity, bytes).await; } + if let Some(bytes) = residue.superseded { + // The object stays — this write overwrote it — so only the byte + // count the superseded entry held is returned. + self.budget.release_disk(bytes); + } if residue.dropped && let Some((identity, bytes)) = wrote { @@ -1160,6 +1181,53 @@ mod tests { ); } + /// A spill that overwrites an entry's own disk object in place must release + /// the copy it superseded. + /// + /// Reported by review. The store key is `(entry id, identity)`, so this + /// insert's put landed on the very object the old entry named — deleting it + /// would destroy the bytes just written. But the old entry's reservation is + /// still counted, so one object ends up charged twice. + #[tokio::test] + async fn an_in_place_disk_overwrite_releases_the_copy_it_supersedes() { + // Tiny, so the second insert cannot stay in memory and — with only a + // disk entry present — has no memory victim to evict. `insert_inner` + // then spills the batch itself and re-inserts it over its own object. + let store = create_cache_store(64, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(1usize); + + // Get the key onto disk first, so the next insert displaces a disk entry. + store + .insert(entry_id, 7, create_test_arrow_array(512)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let (named_once, charged_once) = charged_disk_bytes_match_the_index(&store); + assert!( + named_once > 0, + "the entry must be on disk for this to test anything" + ); + assert_eq!(named_once, charged_once, "baseline must be consistent"); + + // Same key, same identity, written to disk again over its own object. + store + .insert(entry_id, 7, create_test_arrow_array(4096)) + .await + .unwrap(); + + let (named, charged) = charged_disk_bytes_match_the_index(&store); + let mut kinds = Vec::new(); + store.for_each_entry(|_, _, e| kinds.push(CachedBatchType::from(e))); + println!( + "PROBE first: named={named_once} charged={charged_once}; \ + second: named={named} charged={charged}; entries={kinds:?}" + ); + assert_eq!( + charged, named, + "the superseded copy must be released: one object, one reservation" + ); + } + /// Every byte counted against the disk budget must belong to an index /// entry that is on disk. Anything else can never be released: /// `release_disk` is only ever reached from an index entry. From 315f474f1d31c5cbab143aae6dfc565bb99f9214 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 24 Sep 2026 12:36:03 +0530 Subject: [PATCH 23/24] fix(dev-tools): teach the trace parser about disk_evict MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reclamation work made `disk_evict` appear in the cache trace snapshots for the first time. dev-tools parses those snapshots and panics on any event it does not recognise, so `parse_all_snapshot_traces` failed — the event type has existed in the core's trace enum all along, it had simply never been emitted into a recorded trace before. Adds the variant, its parse arm, and the two match arms that were exhaustive. The simulator deliberately moves no I/O counter for it: the event frees disk space rather than reading or writing any. Found by CI, not locally: dev-tools does not build here without its Tailwind asset, so every local run of mine excluded it. Dropping a placeholder into dev/dev-tools/assets/ is enough to type-check and test the crate locally, which is how this was verified before pushing. --- dev/dev-tools/src/trace/parser.rs | 21 +++++++++++++++++++++ dev/dev-tools/src/trace/simulator.rs | 6 ++++++ 2 files changed, 27 insertions(+) diff --git a/dev/dev-tools/src/trace/parser.rs b/dev/dev-tools/src/trace/parser.rs index 99f918f51..966821f69 100644 --- a/dev/dev-tools/src/trace/parser.rs +++ b/dev/dev-tools/src/trace/parser.rs @@ -56,6 +56,13 @@ pub enum TraceEvent { kind: CacheKind, bytes: u64, }, + /// A disk object was deleted and its bytes returned to the budget — either + /// an entry evicted from the disk tier, or an object left unreachable when + /// its key changed hands. + DiskEvict { + entry: u64, + bytes: u64, + }, IoReadSqueezedBacking { entry: u64, bytes: u64, @@ -101,6 +108,9 @@ impl TraceEvent { TraceEvent::InsertSuccess { entry, kind } => { format!("Insert entry {} as {}", entry, kind.display_name()) } + TraceEvent::DiskEvict { entry, bytes } => { + format!("Free {} bytes of entry {} from disk", bytes, entry) + } TraceEvent::InsertFailed { entry, kind } => { format!( "Failed to insert entry {} as {}", @@ -198,6 +208,7 @@ impl TraceEvent { TraceEvent::EvictionBegin { .. } => "eviction_begin", TraceEvent::EvictionVictim { .. } => "eviction_victim", TraceEvent::IoWrite { .. } => "io_write", + TraceEvent::DiskEvict { .. } => "disk_evict", TraceEvent::IoReadSqueezedBacking { .. } => "io_r_squeezed", TraceEvent::IoReadArrow { .. } => "io_read_arrow", TraceEvent::IoReadLiquid { .. } => "io_read_liquid", @@ -308,6 +319,16 @@ fn parse_event_line(line: &str) -> TraceEvent { .and_then(|s| s.parse().ok()) .unwrap_or(0), }, + Some("disk_evict") => TraceEvent::DiskEvict { + entry: fields + .get("entry") + .and_then(|s| s.parse().ok()) + .unwrap_or(0), + bytes: fields + .get("bytes") + .and_then(|s| s.parse().ok()) + .unwrap_or(0), + }, Some("io_read_squeezed_backing") => TraceEvent::IoReadSqueezedBacking { entry: fields .get("entry") diff --git a/dev/dev-tools/src/trace/simulator.rs b/dev/dev-tools/src/trace/simulator.rs index ea5f48cb8..799cd9f12 100644 --- a/dev/dev-tools/src/trace/simulator.rs +++ b/dev/dev-tools/src/trace/simulator.rs @@ -223,6 +223,12 @@ impl CacheSimulator { .insert(*victim, VictimStatus::Selected); } } + TraceEvent::DiskEvict { .. } => { + // The entry'''s disk copy was deleted and its bytes returned to + // the budget. No I/O counter moves — this frees space rather + // than reading or writing it — and the entry keeps whatever + // state it has in memory, so there is nothing to mark here. + } TraceEvent::EvictionVictim { entry } => { // Remove from squeeze victims list self.state.eviction_victims.retain(|v| v != entry); From be1797ed6b12a3e8d8f4e3db7a359b3d55242b86 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 24 Sep 2026 13:35:39 +0530 Subject: [PATCH 24/24] chore(cache): drop a stray doc line and test debug output Both from review. The doc line was left over from the Insert builder and had become the first line of DiskResidue's rustdoc, so the summary named the wrong type. The println and its kinds collection were debugging I left in while chasing whether the in-place overwrite path was reachable; nothing reads them. --- src/core/src/cache/core.rs | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 3751d1f86..b95f306bc 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -72,7 +72,6 @@ pub enum PrefetchResult { Absent, } -/// Builder returned by [`LiquidCache::insert`] for configuring cache writes. /// Disk an insert left for its caller to reclaim. /// /// A store object is addressed by entry id *and* the identity that wrote it, so @@ -1216,12 +1215,6 @@ mod tests { .unwrap(); let (named, charged) = charged_disk_bytes_match_the_index(&store); - let mut kinds = Vec::new(); - store.for_each_entry(|_, _, e| kinds.push(CachedBatchType::from(e))); - println!( - "PROBE first: named={named_once} charged={charged_once}; \ - second: named={named} charged={charged}; entries={kinds:?}" - ); assert_eq!( charged, named, "the superseded copy must be released: one object, one reservation"