From 98e1f51926dd8beb67e1a2bfb3d841d9d5e3e48b Mon Sep 17 00:00:00 2001 From: Thomas Peiselt Date: Fri, 20 Feb 2026 11:24:17 +0100 Subject: [PATCH 1/6] Working build --- .github/workflows/arrow.yml | 169 ------------------------ .github/workflows/arrow_flight.yml | 28 ---- .github/workflows/dev.yml | 61 --------- .github/workflows/docs.yml | 91 ------------- .github/workflows/integration.yml | 182 -------------------------- .github/workflows/miri.sh | 20 --- .github/workflows/miri.yaml | 62 --------- .github/workflows/parquet-variant.yml | 18 --- .github/workflows/parquet.yml | 71 ---------- .github/workflows/parquet_derive.yml | 14 -- .github/workflows/rust.yml | 74 ----------- arrow-array/Cargo.toml | 2 +- arrow-avro/Cargo.toml | 2 +- arrow-cast/Cargo.toml | 6 +- arrow-data/Cargo.toml | 2 +- arrow-flight/Cargo.toml | 4 +- arrow-ipc/Cargo.toml | 2 +- arrow-json/Cargo.toml | 4 +- arrow-ord/Cargo.toml | 2 +- arrow-row/Cargo.toml | 2 +- arrow/Cargo.toml | 4 +- parquet/Cargo.toml | 6 +- 22 files changed, 18 insertions(+), 808 deletions(-) delete mode 100644 .github/workflows/dev.yml delete mode 100644 .github/workflows/docs.yml delete mode 100644 .github/workflows/integration.yml delete mode 100755 .github/workflows/miri.sh delete mode 100644 .github/workflows/miri.yaml diff --git a/.github/workflows/arrow.yml b/.github/workflows/arrow.yml index 9b8147326186..09b1a3fd04be 100644 --- a/.github/workflows/arrow.yml +++ b/.github/workflows/arrow.yml @@ -135,172 +135,3 @@ jobs: - name: Check compilation --no-default-features --all-targets --features chrono-tz run: cargo check -p arrow --no-default-features --all-targets --features chrono-tz - - # test the arrow crate builds against wasm32 in nightly rust - wasm32-build: - name: Build wasm32 - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Install wasm32 targets - run: | - rustup target add wasm32-unknown-unknown - rustup target add wasm32-wasip1 - - name: Build wasm32-unknown-unknown - run: cargo build -p arrow --no-default-features --features=json,csv,ipc,ffi --target wasm32-unknown-unknown - - name: Build wasm32-wasip1 - run: cargo build -p arrow --no-default-features --features=json,csv,ipc,ffi --target wasm32-wasip1 - - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - - name: Clippy arrow-buffer - run: | - mod=arrow-buffer - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-data - run: | - mod=arrow-data - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-schema - run: | - mod=arrow-schema - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-array - run: | - mod=arrow-array - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-select - run: | - mod=arrow-select - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-cast - run: | - mod=arrow-cast - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-ipc - run: | - mod=arrow-ipc - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-csv - run: | - mod=arrow-csv - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-json - run: | - mod=arrow-json - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-avro - run: | - mod=arrow-avro - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-string - run: | - mod=arrow-string - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-ord - run: | - mod=arrow-ord - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-arith - run: | - mod=arrow-arith - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-row - run: | - mod=arrow-row - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow - run: | - mod=arrow - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-integration-test - run: | - mod=arrow-integration-test - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies - - name: Clippy arrow-integration-testing - run: | - mod=arrow-integration-testing - cargo clippy -p "$mod" --all-targets --all-features -- -D warnings - # Dependency checks excluding tests & benches. - cargo clippy -p "$mod" -- -D unused_crate_dependencies - cargo clippy -p "$mod" --all-features -- -D unused_crate_dependencies - cargo clippy -p "$mod" --no-default-features -- -D unused_crate_dependencies diff --git a/.github/workflows/arrow_flight.yml b/.github/workflows/arrow_flight.yml index e6aba901aa22..58474ec2dc3e 100644 --- a/.github/workflows/arrow_flight.yml +++ b/.github/workflows/arrow_flight.yml @@ -61,31 +61,3 @@ jobs: - name: Test --examples run: | cargo test -p arrow-flight --features=flight-sql,tls-ring --examples - - vendor: - name: Verify Vendored Code - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Run gen - run: ./arrow-flight/regen.sh - - name: Verify workspace clean (if this fails, run ./arrow-flight/regen.sh and check in results) - run: git diff --exit-code - - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - - name: Run clippy - run: cargo clippy -p arrow-flight --all-targets --all-features -- -D warnings diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml deleted file mode 100644 index c3ca6f4cd682..000000000000 --- a/.github/workflows/dev.yml +++ /dev/null @@ -1,61 +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. - -name: dev - -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs and changes to main -on: - push: - branches: - - main - pull_request: - -env: - ARCHERY_DOCKER_USER: ${{ secrets.DOCKERHUB_USER }} - ARCHERY_DOCKER_PASSWORD: ${{ secrets.DOCKERHUB_TOKEN }} - -jobs: - - rat: - name: Release Audit Tool (RAT) - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: 3.8 - - name: Audit licenses - run: ./dev/release/run-rat.sh . - - prettier: - name: Markdown format - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - - uses: actions/setup-node@v6 - with: - node-version: "14" - - name: Prettier check - run: | - # if you encounter error, run the command below and commit the changes - npx prettier@2.3.2 --write {arrow,arrow-flight,dev,arrow-integration-testing,parquet}/**/*.md README.md CODE_OF_CONDUCT.md CONTRIBUTING.md - git diff --exit-code diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml deleted file mode 100644 index 4eaf62d95de2..000000000000 --- a/.github/workflows/docs.yml +++ /dev/null @@ -1,91 +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. - -name: docs - -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs and changes to main -on: - push: - branches: - - main - pull_request: - -jobs: - - # test doc links still work - docs: - name: Rustdocs are clean - runs-on: ubuntu-latest - container: - image: amd64/rust - env: - RUSTDOCFLAGS: "-Dwarnings --enable-index-page -Zunstable-options" - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Install Nightly Rust - run: rustup install nightly - - name: Run cargo doc - run: cargo +nightly doc --document-private-items --no-deps --workspace --all-features - - name: Fix file permissions - shell: sh - run: | - chmod -c -R +rX "target/doc" | - while read line; do - echo "::warning title=Invalid file permissions automatically fixed::$line" - done - - name: Upload artifacts - uses: actions/upload-pages-artifact@v4 - with: - name: crate-docs - path: target/doc - - deploy: - # Only deploy if a push to main - if: github.ref_name == 'main' && github.event_name == 'push' - needs: docs - permissions: - contents: write - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - - name: Download crate docs - uses: actions/download-artifact@v5 - with: - name: crate-docs - path: website/build - - name: Prepare website - run: | - tar -xf website/build/artifact.tar -C website/build - rm website/build/artifact.tar - cp .asf.yaml ./website/build/.asf.yaml - - name: Deploy to gh-pages - uses: peaceiris/actions-gh-pages@v4.0.0 - if: github.event_name == 'push' && github.ref_name == 'main' - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: website/build - publish_branch: asf-site - # Avoid accumulating history of in progress API jobs: https://github.com/apache/arrow-rs/issues/5908 - force_orphan: true diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml deleted file mode 100644 index b3f5b166abaa..000000000000 --- a/.github/workflows/integration.yml +++ /dev/null @@ -1,182 +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. - -name: integration - -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main -on: - push: - branches: - - main - pull_request: - paths: - - .github/** - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-csv/** - - arrow-data/** - - arrow-integration-test/** - - arrow-integration-testing/** - - arrow-ipc/** - - arrow-json/** - - arrow-avro/** - - arrow-ord/** - - arrow-pyarrow-integration-testing/** - - arrow-pyarrow/** - - arrow-schema/** - - arrow-select/** - - arrow-sort/** - - arrow-string/** - - arrow/** - -jobs: - integration: - name: Archery test With other arrows - runs-on: ubuntu-latest - container: - image: apache/arrow-dev:amd64-conda-integration - env: - ARROW_USE_CCACHE: OFF - ARROW_CPP_EXE_PATH: /build/cpp/debug - ARROW_NANOARROW_PATH: /build/nanoarrow - ARROW_RUST_EXE_PATH: /build/rust/debug - BUILD_DOCS_CPP: OFF - ARROW_INTEGRATION_CPP: ON - ARROW_INTEGRATION_CSHARP: ON - ARCHERY_INTEGRATION_TARGET_IMPLEMENTATIONS: "rust" - ARCHERY_INTEGRATION_WITH_DOTNET: "1" - ARCHERY_INTEGRATION_WITH_GO: "1" - ARCHERY_INTEGRATION_WITH_JAVA: "1" - ARCHERY_INTEGRATION_WITH_JS: "1" - ARCHERY_INTEGRATION_WITH_NANOARROW: "1" - # https://github.com/apache/arrow/pull/38403/files#r1371281630 - ARCHERY_INTEGRATION_WITH_RUST: "1" - # These are necessary because the github runner overrides $HOME - # https://github.com/actions/runner/issues/863 - RUSTUP_HOME: /root/.rustup - CARGO_HOME: /root/.cargo - defaults: - run: - shell: bash - steps: - # This is necessary so that actions/checkout can find git - - name: Export conda path - run: echo "/opt/conda/envs/arrow/bin" >> $GITHUB_PATH - # This is necessary so that Rust can find cargo - - name: Export cargo path - run: echo "/root/.cargo/bin" >> $GITHUB_PATH - - name: Check rustup - run: which rustup - - name: Check cmake - run: which cmake - - name: Checkout Arrow - uses: actions/checkout@v5 - with: - repository: apache/arrow - submodules: true - fetch-depth: 0 - - name: Checkout Arrow Rust - uses: actions/checkout@v5 - with: - path: rust - submodules: true - fetch-depth: 0 - - name: Checkout Arrow .NET - uses: actions/checkout@v5 - with: - repository: apache/arrow-dotnet - path: dotnet - - name: Checkout Arrow Go - uses: actions/checkout@v5 - with: - repository: apache/arrow-go - path: go - - name: Checkout Arrow Java - uses: actions/checkout@v5 - with: - repository: apache/arrow-java - path: java - - name: Checkout Arrow JavaScript - uses: actions/checkout@v5 - with: - repository: apache/arrow-js - path: js - - name: Checkout Arrow nanoarrow - uses: actions/checkout@v5 - with: - repository: apache/arrow-nanoarrow - path: nanoarrow - - name: Build - run: conda run --no-capture-output ci/scripts/integration_arrow_build.sh $PWD /build - - name: Run - run: conda run --no-capture-output ci/scripts/integration_arrow.sh $PWD /build - - # test FFI against the C-Data interface exposed by pyarrow - pyarrow-integration-test: - name: Pyarrow C Data Interface - runs-on: ubuntu-latest - strategy: - matrix: - rust: [stable] - # PyArrow 15 was the first version to introduce StringView/BinaryView support - pyarrow: ["15", "16", "17"] - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Setup Rust toolchain - run: | - rustup toolchain install ${{ matrix.rust }} - rustup default ${{ matrix.rust }} - rustup component add rustfmt clippy - - name: Cache Cargo - uses: actions/cache@v4 - with: - path: /home/runner/.cargo - key: cargo-maturin-cache- - - name: Cache Rust dependencies - uses: actions/cache@v4 - with: - path: /home/runner/target - # this key is not equal because maturin uses different compilation flags. - key: ${{ runner.os }}-${{ matrix.arch }}-target-maturin-cache-${{ matrix.rust }}- - - uses: actions/setup-python@v6 - with: - python-version: '3.8' - - name: Upgrade pip and setuptools - run: pip install --upgrade pip setuptools wheel virtualenv - - name: Create virtualenv and install dependencies - run: | - virtualenv venv - source venv/bin/activate - pip install maturin toml pytest pytz pyarrow==${{ matrix.pyarrow }} - - name: Run Rust tests - run: | - source venv/bin/activate - cd arrow-pyarrow-testing - cargo test - - name: Run Python tests - run: | - source venv/bin/activate - cd arrow-pyarrow-integration-testing - maturin develop - pytest -v . diff --git a/.github/workflows/miri.sh b/.github/workflows/miri.sh deleted file mode 100755 index 86be2100ee67..000000000000 --- a/.github/workflows/miri.sh +++ /dev/null @@ -1,20 +0,0 @@ -#!/bin/bash -# -# Script -# -# Must be run with nightly rust for example -# rustup default nightly - -set -e - -export MIRIFLAGS="-Zmiri-disable-isolation" -cargo miri setup -cargo clean - -echo "Starting Arrow MIRI run..." -cargo miri test -p arrow-buffer -cargo miri test -p arrow-data --features ffi -cargo miri test -p arrow-schema --features ffi -cargo miri test -p arrow-ord -cargo miri test -p arrow-array -cargo miri test -p arrow-arith \ No newline at end of file diff --git a/.github/workflows/miri.yaml b/.github/workflows/miri.yaml deleted file mode 100644 index 92c432dc893b..000000000000 --- a/.github/workflows/miri.yaml +++ /dev/null @@ -1,62 +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. - -name: miri - -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true - -# trigger for all PRs that touch certain files and changes to main -on: - push: - branches: - - main - pull_request: - paths: - - .github/** - - arrow-array/** - - arrow-buffer/** - - arrow-cast/** - - arrow-csv/** - - arrow-data/** - - arrow-ipc/** - - arrow-json/** - - arrow-avro/** - - arrow-schema/** - - arrow-select/** - - arrow-string/** - - arrow/** - -jobs: - miri-checks: - name: MIRI - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Setup Rust toolchain - run: | - rustup toolchain install nightly --component miri - rustup override set nightly - cargo miri setup - - name: Run Miri Checks - env: - RUST_BACKTRACE: full - RUST_LOG: "trace" - run: bash .github/workflows/miri.sh diff --git a/.github/workflows/parquet-variant.yml b/.github/workflows/parquet-variant.yml index 26cd73ea24e5..559e235dceca 100644 --- a/.github/workflows/parquet-variant.yml +++ b/.github/workflows/parquet-variant.yml @@ -73,21 +73,3 @@ jobs: run: cargo check -p parquet-variant-json - name: Check compilation (parquet-variant-compute) run: cargo check -p parquet-variant-compute - - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - - name: Run clippy (parquet-variant) - run: cargo clippy -p parquet-variant --all-targets --all-features -- -D warnings - - name: Run clippy (parquet-variant-json) - run: cargo clippy -p parquet-variant-json --all-targets --all-features -- -D warnings - - name: Run clippy (parquet-variant-compute) - run: cargo clippy -p parquet-variant-compute --all-targets --all-features -- -D warnings diff --git a/.github/workflows/parquet.yml b/.github/workflows/parquet.yml index 09fc18e351d9..0421bfce65dd 100644 --- a/.github/workflows/parquet.yml +++ b/.github/workflows/parquet.yml @@ -125,74 +125,3 @@ jobs: - name: Check compilation --no-default-features --features variant_experimental run: cargo check -p parquet --no-default-features --features variant_experimental - - # test the parquet crate builds against wasm32 in stable rust - wasm32-build: - name: Build wasm32 - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Install wasm32 targets - run: | - rustup target add wasm32-unknown-unknown - rustup target add wasm32-wasip1 - - name: Install clang # Needed for zlib compilation - run: apt-get update && apt-get install -y clang gcc-multilib - - name: Build wasm32-unknown-unknown - run: cargo build -p parquet --target wasm32-unknown-unknown - - name: Build wasm32-wasip1 - run: cargo build -p parquet --target wasm32-wasip1 - - pyspark-integration-test: - name: PySpark Integration Test - runs-on: ubuntu-latest - strategy: - matrix: - rust: [ stable ] - steps: - - uses: actions/checkout@v5 - - name: Setup Python - uses: actions/setup-python@v6 - with: - python-version: "3.10" - cache: "pip" - - name: Install Python dependencies - run: | - cd parquet/pytest - pip install -r requirements.txt - - name: Black check the test files - run: | - cd parquet/pytest - black --check *.py --verbose - - name: Setup Rust toolchain - run: | - rustup toolchain install ${{ matrix.rust }} - rustup default ${{ matrix.rust }} - - name: Install binary for checking - run: | - cargo install --path parquet --bin parquet-show-bloom-filter --features=cli - cargo install --path parquet --bin parquet-fromcsv --features=arrow,cli - - name: Run pytest - run: | - cd parquet/pytest - pytest -v - - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - - name: Run clippy - run: cargo clippy -p parquet --all-targets --all-features -- -D warnings diff --git a/.github/workflows/parquet_derive.yml b/.github/workflows/parquet_derive.yml index 98c3168cc1be..189922e03516 100644 --- a/.github/workflows/parquet_derive.yml +++ b/.github/workflows/parquet_derive.yml @@ -50,17 +50,3 @@ jobs: uses: ./.github/actions/setup-builder - name: Test run: cargo test -p parquet_derive - - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - - name: Run clippy - run: cargo clippy -p parquet_derive --all-features -- -D warnings diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 9cd33b296da1..96bdfecac211 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -30,59 +30,6 @@ on: pull_request: jobs: - # Check workspace wide compile and test with default features for - # mac - macos: - name: Test on Mac - runs-on: macos-latest - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Install protoc with brew - run: brew install protobuf - - name: Setup Rust toolchain - run: | - rustup toolchain install stable --no-self-update - rustup default stable - - name: Run tests - shell: bash - run: | - # do not produce debug symbols to keep memory usage down - export RUSTFLAGS="-C debuginfo=0" - # PyArrow tests happen in integration.yml. - cargo test --workspace - - # Check workspace wide compile and test with default features for - # windows - windows: - name: Test on Windows - runs-on: windows-latest - steps: - - uses: actions/checkout@v5 - with: - submodules: true - - name: Install protobuf compiler in /d/protoc - shell: bash - run: | - mkdir /d/protoc - cd /d/protoc - curl -LO https://github.com/protocolbuffers/protobuf/releases/download/v21.4/protoc-21.4-win64.zip - unzip protoc-21.4-win64.zip - export PATH=$PATH:/d/protoc/bin - protoc --version - - - name: Setup Rust toolchain - run: | - rustup toolchain install stable --no-self-update - rustup default stable - - name: Run tests - shell: bash - run: | - # do not produce debug symbols to keep memory usage down - export RUSTFLAGS="-C debuginfo=0" - export PATH=$PATH:/d/protoc/bin - cargo test --workspace # Run cargo fmt for all crates lint: @@ -106,24 +53,3 @@ jobs: # if this fails, run this from the parquet directory: # cargo fmt -p parquet -- --config skip_children=true `find . -name "*.rs" \! -name format.rs` cargo fmt -p parquet -- --check --config skip_children=true `find . -name "*.rs" \! -name format.rs` - - msrv: - name: Verify MSRV (Minimum Supported Rust Version) - runs-on: ubuntu-latest - container: - image: amd64/rust - steps: - - uses: actions/checkout@v5 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Install cargo-msrv (if needed) - # cargo-msrv binary may be cached by the cargo cache step in setup-builder, and cargo install will error if it is already installed - run: if which cargo-msrv ; then echo "using existing cargo-msrv binary" ; else cargo install cargo-msrv ; fi - - name: Check all packages - run: | - # run `cargo msrv verify --manifest-path "path/to/Cargo.toml"` to see problematic dependencies - find . -mindepth 2 -name Cargo.toml | while read -r dir - do - echo "Checking package '$dir'" - cargo msrv verify --manifest-path "$dir" --output-format=json || exit 1 - done diff --git a/arrow-array/Cargo.toml b/arrow-array/Cargo.toml index 876872ce4c03..975063617f76 100644 --- a/arrow-array/Cargo.toml +++ b/arrow-array/Cargo.toml @@ -47,7 +47,7 @@ chrono-tz = { version = "0.10", optional = true } num-complex = { version = "0.4.6", default-features = false, features = ["std"] } num-integer = { version = "0.1.46", default-features = false, features = ["std"] } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } -half = { version = "2.1", default-features = false, features = ["num-traits"] } +half = { version = ">=2.1.0, <=2.7.1", default-features = false, features = ["num-traits"] } hashbrown = { version = "0.16.0", default-features = false } [package.metadata.docs.rs] diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index dc59d337a0f4..add99577a90f 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -59,7 +59,7 @@ snap = { version = "1.0", default-features = false, optional = true } zstd = { version = "0.13", default-features = false, optional = true } bzip2 = { version = "0.6.0", optional = true } xz = { version = "0.1", default-features = false, optional = true } -crc = { version = "3.0", optional = true } +crc = { version = "=3.0.1", optional = true } strum_macros = "0.27" uuid = "1.17" indexmap = "2.10" diff --git a/arrow-cast/Cargo.toml b/arrow-cast/Cargo.toml index 12da1af79fe0..c70bf4a7baa0 100644 --- a/arrow-cast/Cargo.toml +++ b/arrow-cast/Cargo.toml @@ -46,17 +46,17 @@ arrow-data = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } chrono = { workspace = true } -half = { version = "2.1", default-features = false } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } lexical-core = { version = "1.0", default-features = false, features = ["write-integers", "write-floats", "parse-integers", "parse-floats"] } atoi = "2.0.0" -comfy-table = { version = "7", optional = true, default-features = false } +comfy-table = { version = "=7.1.0", optional = true, default-features = false } base64 = "0.22" ryu = "1.0.16" [dev-dependencies] criterion = { version = "0.5", default-features = false } -half = { version = "2.1", default-features = false } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } rand = "0.9" [[bench]] diff --git a/arrow-data/Cargo.toml b/arrow-data/Cargo.toml index 9c7a5206b2f4..3c68c7446fa4 100644 --- a/arrow-data/Cargo.toml +++ b/arrow-data/Cargo.toml @@ -50,7 +50,7 @@ arrow-schema = { workspace = true } num-integer = { version = "0.1.46", default-features = false, features = ["std"] } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } -half = { version = "2.1", default-features = false } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } [dev-dependencies] diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml index 8f95e1995a67..a700118aad26 100644 --- a/arrow-flight/Cargo.toml +++ b/arrow-flight/Cargo.toml @@ -75,10 +75,10 @@ cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-root [dev-dependencies] arrow-cast = { workspace = true, features = ["prettyprint"] } -assert_cmd = "2.0.8" +assert_cmd = ">=2.0.8, <2.1.0" http = "1.1.0" http-body = "1.0.0" -hyper-util = "0.1" +hyper-util = "=0.1.4" pin-project-lite = "0.2" tempfile = "3.3" tracing-log = { version = "0.2" } diff --git a/arrow-ipc/Cargo.toml b/arrow-ipc/Cargo.toml index eb42a1ea9589..63ab6301b424 100644 --- a/arrow-ipc/Cargo.toml +++ b/arrow-ipc/Cargo.toml @@ -42,7 +42,7 @@ arrow-data = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true} flatbuffers = { version = "25.2.10", default-features = false } -lz4_flex = { version = "0.11", default-features = false, features = ["std", "frame"], optional = true } +lz4_flex = { version = ">=0.11.0, <=0.11.6", default-features = false, features = ["std", "frame"], optional = true } zstd = { version = "0.13.0", default-features = false, optional = true } [features] diff --git a/arrow-json/Cargo.toml b/arrow-json/Cargo.toml index 2f9e58406025..7707c037d230 100644 --- a/arrow-json/Cargo.toml +++ b/arrow-json/Cargo.toml @@ -41,8 +41,8 @@ arrow-buffer = { workspace = true } arrow-cast = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } -half = { version = "2.1", default-features = false } -indexmap = { version = "2.0", default-features = false, features = ["std"] } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } +indexmap = { version = ">=2.0.2, <=2.12.1", default-features = false, features = ["std"] } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } serde_core = { version = "1.0", default-features = false } serde_json = { version = "1.0", default-features = false, features = ["std"] } diff --git a/arrow-ord/Cargo.toml b/arrow-ord/Cargo.toml index ae76841bda39..c246636a9844 100644 --- a/arrow-ord/Cargo.toml +++ b/arrow-ord/Cargo.toml @@ -43,5 +43,5 @@ arrow-schema = { workspace = true } arrow-select = { workspace = true } [dev-dependencies] -half = { version = "2.1", default-features = false, features = ["num-traits"] } +half = { version = ">=2.1.0, <=2.7.1", default-features = false, features = ["num-traits"] } rand = { version = "0.9", default-features = false, features = ["std", "std_rng"] } diff --git a/arrow-row/Cargo.toml b/arrow-row/Cargo.toml index cd854aa3d48f..04b44eef0864 100644 --- a/arrow-row/Cargo.toml +++ b/arrow-row/Cargo.toml @@ -41,7 +41,7 @@ arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } -half = { version = "2.1", default-features = false } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } [dev-dependencies] arrow-cast = { workspace = true } diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index 743628c8c7d1..893a9cace0de 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -55,7 +55,7 @@ arrow-select = { workspace = true } arrow-string = { workspace = true } rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"], optional = true } -half = { version = "2.1", default-features = false, optional = true } +half = { version = ">=2.1.0, <=2.7.1", default-features = false, optional = true } [package.metadata.docs.rs] all-features = true @@ -85,7 +85,7 @@ canonical_extension_types = ["arrow-schema/canonical_extension_types"] [dev-dependencies] chrono = { workspace = true } criterion = { version = "0.6", default-features = false } -half = { version = "2.1", default-features = false } +half = { version = ">=2.1.0, <=2.7.1", default-features = false } rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } serde = { version = "1.0", default-features = false, features = ["derive"] } # used in examples diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index aa0071ca38e5..3576fff251b0 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -58,7 +58,7 @@ snap = { version = "1.0", default-features = false, optional = true } brotli = { version = "8.0", default-features = false, features = ["std"], optional = true } # To use `flate2` you must enable either the `flate2-zlib-rs` or `flate2-rust_backened` backends flate2 = { version = "1.1", default-features = false, optional = true } -lz4_flex = { version = "0.11", default-features = false, features = ["std", "frame"], optional = true } +lz4_flex = { version = ">=0.11.0, <=0.11.6", default-features = false, features = ["std", "frame"], optional = true } zstd = { version = "0.13", optional = true, default-features = false } chrono = { workspace = true } num-bigint = { version = "0.4", default-features = false } @@ -74,7 +74,7 @@ tokio = { version = "1.0", optional = true, default-features = false, features = hashbrown = { version = "0.16", default-features = false } twox-hash = { version = "2.0", default-features = false, features = ["xxhash64"] } paste = { version = "1.0" } -half = { version = "2.1", default-features = false, features = ["num-traits"] } +half = { version = ">=2.1.0, <=2.7.1", default-features = false, features = ["num-traits"] } crc32fast = { version = "1.4.2", optional = true, default-features = false } simdutf8 = { workspace = true , optional = true } ring = { version = "0.17", default-features = false, features = ["std"], optional = true } @@ -87,7 +87,7 @@ tempfile = { version = "3.0", default-features = false } insta = "1.43.1" brotli = { version = "8.0", default-features = false, features = ["std"] } flate2 = { version = "1.0", default-features = false, features = ["rust_backend"] } -lz4_flex = { version = "0.11", default-features = false, features = ["std", "frame"] } +lz4_flex = { version = ">=0.11.0, <=0.11.6", default-features = false, features = ["std", "frame"] } zstd = { version = "0.13", default-features = false } serde_json = { version = "1.0", features = ["std"], default-features = false } arrow = { workspace = true, features = ["ipc", "test_utils", "prettyprint", "json"] } From 2ce02e91b2fb667595f5afd1435221ca30899d2f Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 7 Apr 2026 11:45:57 -0600 Subject: [PATCH 2/6] Add coop yield in parquet reader (fork-only) Adds with_predicate_coop() to ReadPlanBuilder, an async version of with_predicate() that cooperatively yields after decoding 2MB of data, preventing long predicate evaluations from blocking the async runtime. Co-Authored-By: Dan Harris --- parquet/src/arrow/arrow_reader/read_plan.rs | 49 +++++++++++++++++++++ parquet/src/arrow/async_reader/mod.rs | 4 +- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 754fcd339c5b..0d8ed9f4f988 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -27,6 +27,11 @@ use arrow_array::Array; use arrow_select::filter::prep_null_mask_filter; use std::collections::VecDeque; +/// Maximum number of bytes that can be decoded during predicate evaluation +/// before yielding back to the async scheduler (fork-only) +#[cfg(feature = "async")] +const DECODE_BUDGET: usize = 2 * 1024 * 1024; + /// A builder for [`ReadPlan`] #[derive(Clone)] pub struct ReadPlanBuilder { @@ -122,6 +127,50 @@ impl ReadPlanBuilder { Ok(self) } + /// Async version of [`Self::with_predicate`] that cooperatively yields + /// back to the async scheduler after decoding [`DECODE_BUDGET`] bytes, + /// preventing long predicate evaluations from blocking the runtime. + #[cfg(feature = "async")] + pub(crate) async fn with_predicate_coop( + mut self, + array_reader: Box, + predicate: &mut dyn ArrowPredicate, + ) -> Result { + let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); + let mut filters = vec![]; + let mut budget = DECODE_BUDGET; + for maybe_batch in reader { + let maybe_batch = maybe_batch?; + budget = budget.saturating_sub(maybe_batch.get_array_memory_size()); + + let input_rows = maybe_batch.num_rows(); + let filter = predicate.evaluate(maybe_batch)?; + // Since user supplied predicate, check error here to catch bugs quickly + if filter.len() != input_rows { + return Err(arrow_err!( + "ArrowPredicate predicate returned {} rows, expected {input_rows}", + filter.len() + )); + } + match filter.null_count() { + 0 => filters.push(filter), + _ => filters.push(prep_null_mask_filter(&filter)), + }; + + if budget == 0 { + budget = DECODE_BUDGET; + tokio::task::yield_now().await; + } + } + + let raw = RowSelection::from_filters(&filters); + self.selection = match self.selection.take() { + Some(selection) => Some(selection.and_then(&raw)), + None => Some(raw), + }; + Ok(self) + } + /// Create a final `ReadPlan` the read plan for the scan pub fn build(mut self) -> ReadPlan { // If selection is empty, truncate diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 27f90e3d7bc6..6807da9eecbb 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -651,7 +651,9 @@ where .with_cache_options(Some(&cache_options)) .build_array_reader(self.fields.as_deref(), predicate.projection())?; - plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?; + plan_builder = plan_builder + .with_predicate_coop(array_reader, predicate.as_mut()) + .await?; } } From 9ef74ed84dd6f90890cbf150981e4e5e86e976d8 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 7 Apr 2026 12:20:08 -0600 Subject: [PATCH 3/6] Add row_id column to parquet reader for index-based row lookups (#8715 v57) Adds RowId support to ParquetRecordBatchReader and the async stream builder. When enabled via with_row_id(), a UInt64 column is prepended to each RecordBatch containing the file-level row offset for each row. Co-Authored-By: Thomas Peiselt --- parquet/src/arrow/array_reader/mod.rs | 5 +- parquet/src/arrow/arrow_reader/mod.rs | 120 ++++- parquet/src/arrow/arrow_reader/read_plan.rs | 4 +- parquet/src/arrow/async_reader/mod.rs | 462 +++++++++++++++++++- 4 files changed, 571 insertions(+), 20 deletions(-) diff --git a/parquet/src/arrow/array_reader/mod.rs b/parquet/src/arrow/array_reader/mod.rs index b3595e58d695..aee4a62e5b29 100644 --- a/parquet/src/arrow/array_reader/mod.rs +++ b/parquet/src/arrow/array_reader/mod.rs @@ -48,7 +48,9 @@ mod struct_array; mod test_util; // Note that this crate is public under the `experimental` feature flag. -pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder}; +pub use builder::ArrayReaderBuilder; +#[allow(unused_imports)] +pub use builder::{CacheOptions, CacheOptionsBuilder}; pub use byte_array::make_byte_array_reader; pub use byte_array_dictionary::make_byte_array_dictionary_reader; #[allow(unused_imports)] // Only used for benchmarks @@ -60,6 +62,7 @@ pub use list_array::ListArrayReader; pub use map_array::MapArrayReader; pub use null_array::NullArrayReader; pub use primitive_array::PrimitiveArrayReader; +#[allow(unused_imports)] pub use row_group_cache::RowGroupCache; pub use struct_array::StructArrayReader; diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 4e91685519f3..f4dd0d43d338 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -18,9 +18,11 @@ //! Contains reader which reads parquet data into arrow [`RecordBatch`] use arrow_array::Array; +use arrow_array::ArrayRef; +use arrow_array::builder::UInt64Builder; use arrow_array::cast::AsArray; use arrow_array::{RecordBatch, RecordBatchReader}; -use arrow_schema::{ArrowError, DataType as ArrowType, Schema, SchemaRef}; +use arrow_schema::{ArrowError, DataType as ArrowType, Field, FieldRef, Schema, SchemaRef}; pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter}; pub use selection::{RowSelection, RowSelector}; use std::fmt::{Debug, Formatter}; @@ -121,7 +123,11 @@ pub struct ArrowReaderBuilder { pub(crate) metrics: ArrowReaderMetrics, + #[allow(dead_code)] pub(crate) max_predicate_cache_size: usize, + + #[allow(unused)] + pub(crate) row_id: Option, } impl Debug for ArrowReaderBuilder { @@ -159,6 +165,7 @@ impl ArrowReaderBuilder { offset: None, metrics: ArrowReaderMetrics::Disabled, max_predicate_cache_size: 100 * 1024 * 1024, // 100MB default cache size + row_id: None, } } @@ -185,6 +192,15 @@ impl ArrowReaderBuilder { Self { batch_size, ..self } } + /// Project a column into the result with name `field_name` that will contain the row ID + /// for each row. The row ID will be the row offset of the row in the underlying file + pub fn with_row_id(self, field_name: impl Into) -> Self { + Self { + row_id: Some(RowId::field_ref(field_name)), + ..self + } + } + /// Only read data from the provided row group indexes /// /// This is also called row group filtering @@ -869,6 +885,7 @@ impl ParquetRecordBatchReaderBuilder { metrics, // Not used for the sync reader, see https://github.com/apache/arrow-rs/issues/8000 max_predicate_cache_size: _, + row_id: _row_id, } = self; // Try to avoid allocate large buffer @@ -914,7 +931,8 @@ impl ParquetRecordBatchReaderBuilder { .build_limited() .build(); - Ok(ParquetRecordBatchReader::new(array_reader, read_plan)) + // TODO what do we do here? + Ok(ParquetRecordBatchReader::new(array_reader, read_plan, None)) } } @@ -990,12 +1008,55 @@ impl Iterator for ReaderPageIterator { impl PageIterator for ReaderPageIterator {} +pub(crate) struct RowId { + offset: u64, + field: FieldRef, + buffer: UInt64Builder, +} + +impl RowId { + #[allow(unused)] + pub fn new(offset: u64, field: FieldRef, batch_size: usize) -> Self { + Self { + offset, + field, + buffer: UInt64Builder::with_capacity(batch_size), + } + } + + pub fn field_ref(name: impl Into) -> FieldRef { + Arc::new(Field::new(name, ArrowType::UInt64, false)) + } + + pub fn skip(&mut self, n: usize) { + self.offset += n as u64; + } + + pub fn field(&self) -> FieldRef { + self.field.clone() + } + + fn read(&mut self, n: usize) { + // SAFETY: We are appending a `Range` which has a trusted length + unsafe { + self.buffer + .append_trusted_len_iter(self.offset..self.offset + n as u64) + } + self.offset += n as u64; + } + + fn consume(&mut self) -> ArrayRef { + Arc::new(self.buffer.finish()) + } +} + /// An `Iterator>` that yields [`RecordBatch`] /// read from a parquet data source pub struct ParquetRecordBatchReader { array_reader: Box, schema: SchemaRef, read_plan: ReadPlan, + row_id: Option, } impl Iterator for ParquetRecordBatchReader { @@ -1024,6 +1085,10 @@ impl ParquetRecordBatchReader { if front.skip { let skipped = self.array_reader.skip_records(front.row_count)?; + if let Some(row_id) = self.row_id.as_mut() { + row_id.skip(skipped); + } + if skipped != front.row_count { return Err(general_err!( "failed to skip rows, expected {}, got {}", @@ -1053,13 +1118,23 @@ impl ParquetRecordBatchReader { }; match self.array_reader.read_records(to_read)? { 0 => break, - rec => read_records += rec, + rec => { + if let Some(rowid) = self.row_id.as_mut() { + rowid.read(rec); + } + read_records += rec + } }; } } - None => { - self.array_reader.read_records(batch_size)?; - } + None => match self.array_reader.read_records(batch_size) { + Ok(n) => { + if let Some(rowid) = self.row_id.as_mut() { + rowid.read(n); + } + } + Err(error) => return Err(error), + }, }; let array = self.array_reader.consume_batch()?; @@ -1068,7 +1143,16 @@ impl ParquetRecordBatchReader { })?; Ok(if struct_array.len() > 0 { - Some(RecordBatch::from(struct_array)) + Some(match self.row_id.as_mut() { + Some(rowid) => { + let columns = std::iter::once(rowid.consume()) + .chain(struct_array.columns().iter().cloned()) + .collect(); + + RecordBatch::try_new(self.schema.clone(), columns).expect("invalid schema") + } + None => RecordBatch::from(struct_array), + }) } else { None }) @@ -1118,22 +1202,38 @@ impl ParquetRecordBatchReader { array_reader, schema: Arc::new(Schema::new(levels.fields.clone())), read_plan, + row_id: None, }) } /// Create a new [`ParquetRecordBatchReader`] that will read at most `batch_size` rows at /// a time from [`ArrayReader`] based on the configured `selection`. If `selection` is `None` /// all rows will be returned - pub(crate) fn new(array_reader: Box, read_plan: ReadPlan) -> Self { - let schema = match array_reader.get_data_type() { - ArrowType::Struct(fields) => Schema::new(fields.clone()), + pub(crate) fn new( + array_reader: Box, + read_plan: ReadPlan, + rowid: Option, + ) -> Self { + let struct_fields = match array_reader.get_data_type() { + ArrowType::Struct(fields) => fields.clone(), _ => unreachable!("Struct array reader's data type is not struct!"), }; + let schema = match rowid.as_ref() { + Some(rowid) => { + let fields: Vec<_> = std::iter::once(rowid.field()) + .chain(struct_fields.iter().cloned()) + .collect(); + Schema::new(fields) + } + None => Schema::new(struct_fields), + }; + Self { array_reader, schema: Arc::new(schema), read_plan, + row_id: rowid, } } diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 0d8ed9f4f988..d46ab335a538 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -100,7 +100,7 @@ impl ReadPlanBuilder { array_reader: Box, predicate: &mut dyn ArrowPredicate, ) -> Result { - let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); + let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build(), None); let mut filters = vec![]; for maybe_batch in reader { let maybe_batch = maybe_batch?; @@ -136,7 +136,7 @@ impl ReadPlanBuilder { array_reader: Box, predicate: &mut dyn ArrowPredicate, ) -> Result { - let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); + let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build(), None); let mut filters = vec![]; let mut budget = DECODE_BUDGET; for maybe_batch in reader { diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 6807da9eecbb..d732e13380d4 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -36,7 +36,7 @@ use futures::stream::Stream; use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; use arrow_array::RecordBatch; -use arrow_schema::{DataType, Fields, Schema, SchemaRef}; +use arrow_schema::{DataType, FieldRef, Fields, Schema, SchemaRef}; use crate::arrow::ProjectionMask; use crate::arrow::array_reader::{ @@ -44,7 +44,7 @@ use crate::arrow::array_reader::{ }; use crate::arrow::arrow_reader::{ ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, - RowFilter, RowSelection, + RowFilter, RowId, RowSelection, }; use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash}; @@ -516,17 +516,27 @@ impl ParquetRecordBatchStreamBuilder { offset: self.offset, metrics: self.metrics, max_predicate_cache_size: self.max_predicate_cache_size, + rowid: self.row_id.clone(), }; // Ensure schema of ParquetRecordBatchStream respects projection, and does // not store metadata (same as for ParquetRecordBatchReader and emitted RecordBatches) - let projected_fields = match reader_factory.fields.as_deref().map(|pf| &pf.arrow_type) { + let mut projected_fields = match reader_factory.fields.as_deref().map(|pf| &pf.arrow_type) { Some(DataType::Struct(fields)) => { fields.filter_leaves(|idx, _| self.projection.leaf_included(idx)) } None => Fields::empty(), _ => unreachable!("Must be Struct for root type"), }; + + if let Some(field) = &self.row_id { + projected_fields = Fields::from( + std::iter::once(field.clone()) + .chain(projected_fields.iter().cloned()) + .collect::>(), + ); + } + let schema = Arc::new(Schema::new(projected_fields)); Ok(ParquetRecordBatchStream { @@ -561,7 +571,7 @@ struct ReaderFactory { /// Optional filter filter: Option, - /// Limit to apply to remaining row groups. + /// Limit to apply to remaining row groups. limit: Option, /// Offset to apply to the next @@ -572,6 +582,9 @@ struct ReaderFactory { /// Maximum size of the predicate cache max_predicate_cache_size: usize, + + /// Optional row ID field + rowid: Option, } impl ReaderFactory @@ -703,6 +716,14 @@ where ) .await?; + let rowid = self.rowid.clone().map(|field| { + let offset = self.metadata.row_groups()[..row_group_idx] + .iter() + .map(|rg| rg.num_rows() as u64) + .sum::(); + RowId::new(offset, field, batch_size) + }); + let plan = plan_builder.build(); let cache_options = cache_options_builder.consumer(); @@ -710,7 +731,7 @@ where .with_cache_options(Some(&cache_options)) .build_array_reader(self.fields.as_deref(), &projection)?; - let reader = ParquetRecordBatchReader::new(array_reader, plan); + let reader = ParquetRecordBatchReader::new(array_reader, plan, rowid); Ok((self, Some(reader))) } @@ -1221,14 +1242,15 @@ mod tests { use crate::file::properties::WriterProperties; use arrow::compute::kernels::cmp::eq; use arrow::error::Result as ArrowResult; - use arrow_array::builder::{ListBuilder, StringBuilder}; + use arrow_array::builder::{ListBuilder, StringBuilder, UInt64Builder}; use arrow_array::cast::AsArray; - use arrow_array::types::Int32Type; + use arrow_array::types::{Int32Type, UInt64Type}; use arrow_array::{ Array, ArrayRef, Int8Array, Int32Array, RecordBatchReader, Scalar, StringArray, StructArray, UInt64Array, }; use arrow_schema::{DataType, Field, Schema}; + use arrow_select::concat::concat; use futures::{StreamExt, TryStreamExt}; use rand::{Rng, rng}; use std::collections::HashMap; @@ -1992,6 +2014,7 @@ mod tests { offset: None, metrics: ArrowReaderMetrics::disabled(), max_predicate_cache_size: 0, + rowid: None, }; let mut skip = true; @@ -2457,6 +2480,7 @@ mod tests { offset: None, metrics: ArrowReaderMetrics::disabled(), max_predicate_cache_size: 0, + rowid: None, }; // Provide an output projection that also selects the same nested leaf @@ -2570,6 +2594,430 @@ mod tests { assert_eq!(result.len(), 1); } + #[tokio::test] + async fn test_async_reader_with_rowid() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let metadata = Arc::new(metadata); + + assert_eq!(metadata.num_row_groups(), 1); + + let async_reader = TestReader { + data: data.clone(), + metadata: Some(metadata.clone()), + requests: Default::default(), + }; + + let requests = async_reader.requests.clone(); + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_id("_rowid") + .build() + .unwrap(); + + assert_eq!( + stream + .schema() + .fields() + .first() + .expect("no fields in schema") + .name(), + "_rowid" + ); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + assert!(async_batches.iter().all(|batch| { + batch + .schema() + .fields() + .first() + .expect("no fields in schema") + .name() + == "_rowid" + })); + + let rowid_arrays = async_batches + .iter() + .map(|batch| batch.column(0).as_ref()) + .collect::>(); + let rowids = concat(&rowid_arrays).expect("concat rowids"); + + let expected_rowids = UInt64Array::from_iter_values(0..rowids.len() as u64); + + assert_eq!(rowids.as_primitive::(), &expected_rowids); + + let requests = requests.lock().unwrap(); + let (offset_1, length_1) = metadata.row_group(0).column(1).byte_range(); + let (offset_2, length_2) = metadata.row_group(0).column(2).byte_range(); + + assert_eq!( + &requests[..], + &[ + offset_1 as usize..(offset_1 + length_1) as usize, + offset_2 as usize..(offset_2 + length_2) as usize + ] + ); + } + + #[tokio::test] + async fn test_async_reader_with_rowid_offset() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let metadata = Arc::new(metadata); + + assert_eq!(metadata.num_row_groups(), 1); + + let async_reader = TestReader { + data: data.clone(), + metadata: Some(metadata.clone()), + requests: Default::default(), + }; + + let builder = ParquetRecordBatchStreamBuilder::new(async_reader) + .await + .unwrap(); + + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![1, 2]); + let stream = builder + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_offset(3) + .with_row_id("_rowid") + .build() + .unwrap(); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + assert!(async_batches.iter().all(|batch| { + batch + .schema() + .fields() + .first() + .expect("no fields in schema") + .name() + == "_rowid" + })); + + let rowid_arrays = async_batches + .iter() + .map(|batch| batch.column(0).as_ref()) + .collect::>(); + let rowids = concat(&rowid_arrays).expect("concat rowids"); + + let expected_rowids = UInt64Array::from_iter_values(3..rowids.len() as u64 + 3); + assert_eq!(rowids.as_primitive::(), &expected_rowids); + } + + #[tokio::test] + async fn test_fuzz_async_reader_with_rowid_and_selection() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let metadata = Arc::new(metadata); + + assert_eq!(metadata.num_row_groups(), 1); + + let mut rand = rng(); + + for _ in 0..100 { + let mut expected_rowids_builder = UInt64Builder::new(); + let mut offset = 0; + + let mut expected_rows = 0; + let mut total_rows = 0; + let mut skip = false; + let mut selectors = vec![]; + + while total_rows < 7300 { + let row_count: usize = rand.random_range(1..100); + + let row_count = row_count.min(7300 - total_rows); + + selectors.push(RowSelector { row_count, skip }); + + total_rows += row_count; + if !skip { + expected_rowids_builder.append_slice( + (offset..offset + row_count as u64) + .collect::>() + .as_slice(), + ); + expected_rows += row_count; + } + + offset += row_count as u64; + + skip = !skip; + } + + let selection = RowSelection::from(selectors); + + let async_reader = TestReader { + data: data.clone(), + metadata: Some(metadata.clone()), + requests: Default::default(), + }; + + let options = ArrowReaderOptions::new().with_page_index(true); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + let col_idx: usize = rand.random_range(0..13); + let mask = ProjectionMask::leaves(builder.parquet_schema(), vec![col_idx]); + + let stream = builder + .with_projection(mask.clone()) + .with_row_selection(selection.clone()) + .with_row_id("_rowid") + .build() + .expect("building stream"); + + let async_batches: Vec<_> = stream.try_collect().await.unwrap(); + + let expected_rowids = expected_rowids_builder.finish(); + + assert!(async_batches.iter().all(|batch| { + batch + .schema() + .fields() + .first() + .expect("no fields in schema") + .name() + == "_rowid" + })); + + let rowid_arrays = async_batches + .iter() + .map(|batch| batch.column(0).as_ref()) + .collect::>(); + let rowids = concat(&rowid_arrays).expect("concat rowids"); + + assert_eq!(rowids.as_primitive::(), &expected_rowids); + + let actual_rows: usize = async_batches.into_iter().map(|b| b.num_rows()).sum(); + + assert_eq!(actual_rows, expected_rows); + } + } + + #[tokio::test] + async fn test_async_reader_with_row_id_and_row_filter() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]); + let c = Int32Array::from_iter(0..6); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("c", Arc::new(c) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), None).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + let parquet_schema = metadata.file_metadata().schema_descr_ptr(); + + let test = TestReader { + data, + metadata: Some(Arc::new(metadata)), + requests: Default::default(), + }; + let requests = test.requests.clone(); + + let a_scalar = StringArray::from_iter_values(["b"]); + let a_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![0]), + move |batch| eq(batch.column(0), &Scalar::new(&a_scalar)), + ); + + let b_scalar = StringArray::from_iter_values(["4"]); + let b_filter = ArrowPredicateFn::new( + ProjectionMask::leaves(&parquet_schema, vec![1]), + move |batch| eq(batch.column(0), &Scalar::new(&b_scalar)), + ); + + let filter = RowFilter::new(vec![Box::new(a_filter), Box::new(b_filter)]); + + let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 2]); + let stream = ParquetRecordBatchStreamBuilder::new(test) + .await + .unwrap() + .with_projection(mask.clone()) + .with_batch_size(1024) + .with_row_filter(filter) + .with_row_id("_rowid") + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 3); + + let col = batch.column(0); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, 3); + + let col = batch.column(1); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, "b"); + + let col = batch.column(2); + let val = col.as_any().downcast_ref::().unwrap().value(0); + assert_eq!(val, 3); + + // Should only have made 3 requests + assert_eq!(requests.lock().unwrap().len(), 3); + } + + #[tokio::test] + async fn test_async_reader_with_rowid_limit_multiple_row_groups() { + let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]); + let b = StringArray::from_iter_values(["1", "2", "3", "4", "5", "6"]); + let c = Int32Array::from_iter(0..6); + let data = RecordBatch::try_from_iter([ + ("a", Arc::new(a) as ArrayRef), + ("b", Arc::new(b) as ArrayRef), + ("c", Arc::new(c) as ArrayRef), + ]) + .unwrap(); + + let mut buf = Vec::with_capacity(1024); + let props = WriterProperties::builder() + .set_max_row_group_size(3) + .build(); + let mut writer = ArrowWriter::try_new(&mut buf, data.schema(), Some(props)).unwrap(); + writer.write(&data).unwrap(); + writer.close().unwrap(); + + let data: Bytes = buf.into(); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&data) + .unwrap(); + + assert_eq!(metadata.num_row_groups(), 2); + + let test = TestReader { + data, + metadata: Some(Arc::new(metadata)), + requests: Default::default(), + }; + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_batch_size(1024) + .with_limit(4) + .with_row_id("_rowid") + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Expect one batch for each row group + assert_eq!(batches.len(), 2); + + let batch = &batches[0]; + // First batch should contain all rows + assert_eq!(batch.num_rows(), 3); + assert_eq!(batch.num_columns(), 4); + let rowids = batch.column(0).as_primitive::(); + assert_eq!(rowids.values(), &[0, 1, 2]); + let col3 = batch.column(3).as_primitive::(); + assert_eq!(col3.values(), &[0, 1, 2]); + + let batch = &batches[1]; + // Second batch should trigger the limit and only have one row + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 4); + let rowids = batch.column(0).as_primitive::(); + assert_eq!(rowids.values(), &[3]); + let col3 = batch.column(3).as_primitive::(); + assert_eq!(col3.values(), &[3]); + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_offset(2) + .with_limit(3) + .with_row_id("_rowid") + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Expect one batch for each row group + assert_eq!(batches.len(), 2); + + let batch = &batches[0]; + // First batch should contain one row + assert_eq!(batch.num_rows(), 1); + assert_eq!(batch.num_columns(), 4); + let rowids = batch.column(0).as_primitive::(); + assert_eq!(rowids.values(), &[2]); + let col3 = batch.column(3).as_primitive::(); + assert_eq!(col3.values(), &[2]); + + let batch = &batches[1]; + // Second batch should contain two rows + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 4); + let rowids = batch.column(0).as_primitive::(); + assert_eq!(rowids.values(), &[3, 4]); + let col3 = batch.column(3).as_primitive::(); + assert_eq!(col3.values(), &[3, 4]); + + let stream = ParquetRecordBatchStreamBuilder::new(test.clone()) + .await + .unwrap() + .with_offset(4) + .with_limit(20) + .with_row_id("_rowid") + .build() + .unwrap(); + + let batches: Vec<_> = stream.try_collect().await.unwrap(); + // Should skip first row group + assert_eq!(batches.len(), 1); + + let batch = &batches[0]; + // First batch should contain two rows + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 4); + let rowids = batch.column(0).as_primitive::(); + assert_eq!(rowids.values(), &[4, 5]); + let col3 = batch.column(3).as_primitive::(); + assert_eq!(col3.values(), &[4, 5]); + } + #[tokio::test] async fn test_cached_array_reader_sparse_offset_error() { use futures::TryStreamExt; From cd1fb46dfc675bd1aee03ad8142e3099089aba60 Mon Sep 17 00:00:00 2001 From: Georgi Krastev Date: Thu, 12 Mar 2026 17:17:01 +0200 Subject: [PATCH 4/6] Bulk RLE encode (fork-only) --- parquet/src/arrow/arrow_writer/levels.rs | 32 +++++ parquet/src/column/writer/mod.rs | 144 ++++++++++++++++++++++- parquet/src/encodings/levels.rs | 22 ++++ parquet/src/encodings/rle.rs | 90 ++++++++++++++ 4 files changed, 285 insertions(+), 3 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/levels.rs b/parquet/src/arrow/arrow_writer/levels.rs index 3c283bcbe3d2..478fc78dab40 100644 --- a/parquet/src/arrow/arrow_writer/levels.rs +++ b/parquet/src/arrow/arrow_writer/levels.rs @@ -619,6 +619,9 @@ pub(crate) struct ArrayLevels { /// cached logical nulls of the array. logical_nulls: Option, + + #[allow(dead_code)] + def_levels_runs: Option>, } impl PartialEq for ArrayLevels { @@ -652,6 +655,7 @@ impl ArrayLevels { max_rep_level, array, logical_nulls, + def_levels_runs: (max_def_level != 0).then(Vec::new), } } @@ -726,6 +730,7 @@ mod tests { max_rep_level: 2, array: Arc::new(primitives), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected); } @@ -747,6 +752,7 @@ mod tests { max_rep_level: 0, array, logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); } @@ -775,6 +781,7 @@ mod tests { max_rep_level: 0, array, logical_nulls, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); } @@ -810,6 +817,7 @@ mod tests { max_rep_level: 1, array: Arc::new(leaf_array), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); @@ -844,6 +852,7 @@ mod tests { max_rep_level: 1, array: Arc::new(leaf_array), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); } @@ -894,6 +903,7 @@ mod tests { max_rep_level: 1, array: Arc::new(leaf), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); @@ -945,6 +955,7 @@ mod tests { max_rep_level: 2, array: Arc::new(leaf), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); @@ -983,6 +994,7 @@ mod tests { max_rep_level: 1, array: Arc::new(leaf), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); @@ -1016,6 +1028,7 @@ mod tests { max_rep_level: 1, array: Arc::new(leaf), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); @@ -1065,6 +1078,7 @@ mod tests { max_rep_level: 2, array: Arc::new(leaf), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); } @@ -1106,6 +1120,7 @@ mod tests { max_rep_level: 0, array: leaf, logical_nulls, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_levels); } @@ -1146,6 +1161,7 @@ mod tests { max_rep_level: 1, array: Arc::new(a_values), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); } @@ -1239,6 +1255,7 @@ mod tests { max_rep_level: 0, array: Arc::new(a), logical_nulls: None, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); @@ -1254,6 +1271,7 @@ mod tests { max_rep_level: 0, array: Arc::new(b), logical_nulls: b_logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); @@ -1269,6 +1287,7 @@ mod tests { max_rep_level: 0, array: Arc::new(d), logical_nulls: d_logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); @@ -1284,6 +1303,7 @@ mod tests { max_rep_level: 0, array: Arc::new(f), logical_nulls: f_logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); } @@ -1392,6 +1412,7 @@ mod tests { max_rep_level: 1, array: map.keys().clone(), logical_nulls: map_keys_logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); @@ -1407,6 +1428,7 @@ mod tests { max_rep_level: 1, array: map.values().clone(), logical_nulls: map_values_logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); } @@ -1494,6 +1516,7 @@ mod tests { max_rep_level: 1, array: values, logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); @@ -1536,6 +1559,7 @@ mod tests { max_rep_level: 1, array: values, logical_nulls, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_level); @@ -1623,6 +1647,7 @@ mod tests { max_rep_level: 2, array: a1_values, logical_nulls: a1_logical_nulls, + def_levels_runs: None, }; assert_eq!(&levels[0], &expected_level); @@ -1636,6 +1661,7 @@ mod tests { max_rep_level: 1, array: a2_values, logical_nulls: a2_logical_nulls, + def_levels_runs: None, }; assert_eq!(&levels[1], &expected_level); @@ -1676,6 +1702,7 @@ mod tests { max_rep_level: 1, array: values, logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); } @@ -1828,6 +1855,7 @@ mod tests { max_rep_level: 1, array: values_a, logical_nulls: values_a_logical_nulls, + def_levels_runs: None, }; // [[{b: 2}, null], null, [null, null], [{b: 3}, {b: 4}]] let values_b_logical_nulls = values_b.logical_nulls(); @@ -1839,6 +1867,7 @@ mod tests { max_rep_level: 1, array: values_b, logical_nulls: values_b_logical_nulls, + def_levels_runs: None, }; assert_eq!(a_levels, &expected_a); @@ -1872,6 +1901,7 @@ mod tests { max_rep_level: 1, array: values, logical_nulls, + def_levels_runs: None, }; assert_eq!(list_level, &expected_level); } @@ -1909,6 +1939,7 @@ mod tests { max_rep_level: 2, array: values, logical_nulls, + def_levels_runs: None, }; assert_eq!(levels[0], expected_level); @@ -1942,6 +1973,7 @@ mod tests { max_rep_level: 0, array: Arc::new(dict), logical_nulls, + def_levels_runs: None, }; assert_eq!(levels[0], expected_level); } diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index fdb94962b6fd..a334954cccac 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -356,6 +356,9 @@ pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> { data_page_boundary_descending: bool, /// (min, max) last_non_null_data_page_min_max: Option<(E::T, E::T)>, + + def_levels_runs_sink: Vec<(i16, usize)>, + num_levels: usize, } impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { @@ -421,6 +424,8 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { data_page_boundary_ascending: true, data_page_boundary_descending: true, last_non_null_data_page_min_max: None, + def_levels_runs_sink: vec![], + num_levels: 0, } } @@ -435,6 +440,9 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { max: Option<&E::T>, distinct_count: Option, ) -> Result { + if !self.def_levels_runs_sink.is_empty() { + self.add_data_page()?; + } // Check if number of definition levels is the same as number of repetition levels. if let (Some(def), Some(rep)) = (def_levels, rep_levels) { if def.len() != rep.len() { @@ -551,6 +559,112 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { ) } + /// Write a batch of values with a uniform definition level (bulk RLE optimization). + pub fn write_def_level_range_batch( + &mut self, + values: &E::Values, + def_level: i16, + num_levels: usize, + min: Option<&E::T>, + max: Option<&E::T>, + distinct_count: Option, + ) -> Result { + if self.descr.max_rep_level() > 0 { + return Err(general_err!( + "Cannot write def level range when rep level > 0", + )); + } + if self.statistics_enabled == EnabledStatistics::Chunk { + match (min, max) { + (Some(min), Some(max)) => { + update_min(&self.descr, min, &mut self.column_metrics.min_column_value); + update_max(&self.descr, max, &mut self.column_metrics.max_column_value); + } + (None, Some(_)) | (Some(_), None) => { + panic!("min/max should be both set or both None") + } + (None, None) => {} + }; + } + + // We can only set the distinct count if there are no other writes + if self.encoder.num_values() == 0 { + self.column_metrics.column_distinct_count = distinct_count; + } else { + self.column_metrics.column_distinct_count = None; + } + + let mut values_offset = 0; + let mut levels_offset = 0; + let base_batch_size = self.props.write_batch_size(); + while levels_offset < num_levels { + let end_offset = num_levels.min(levels_offset + base_batch_size); + + values_offset += self.write_mini_batch_with_def_level( + values, + values_offset, + None, + end_offset - levels_offset, + def_level, + )?; + levels_offset = end_offset; + } + + // Return total number of values processed. + Ok(values_offset) + } + + fn write_mini_batch_with_def_level( + &mut self, + values: &E::Values, + values_offset: usize, + value_indices: Option<&[usize]>, + num_levels: usize, + def_level: i16, + ) -> Result { + if !self.def_levels_sink.is_empty() { + self.add_data_page()?; + } + if let Some((last_def_level, last_count)) = self.def_levels_runs_sink.last_mut() { + if *last_def_level == def_level { + *last_count += num_levels; + } else { + self.def_levels_runs_sink.push((def_level, num_levels)); + } + } else { + self.def_levels_runs_sink.push((def_level, num_levels)); + } + self.num_levels += num_levels; + let values_to_write = if def_level == self.descr.max_def_level() { + num_levels + } else { + self.page_metrics.num_page_nulls += num_levels as u64; + 0 + }; + + self.page_metrics.num_buffered_rows += num_levels as u32; + + match value_indices { + Some(indices) => { + let indices = &indices[values_offset..values_offset + values_to_write]; + self.encoder.write_gather(values, indices)?; + } + None => self.encoder.write(values, values_offset, values_to_write)?, + } + + self.page_metrics.num_buffered_values += num_levels as u32; + + if self.should_add_data_page() { + self.add_data_page()?; + } + + if self.should_dict_fallback() { + self.dict_fallback()?; + } + + Ok(values_to_write) + } + /// Returns the estimated total memory usage. /// /// Unlike [`Self::get_estimated_total_bytes`] this is an estimate @@ -1060,13 +1174,21 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } if max_def_level > 0 { - buffer.extend_from_slice( + let encoded_levels = if self.def_levels_runs_sink.is_empty() { &self.encode_levels_v1( Encoding::RLE, &self.def_levels_sink[..], max_def_level, - )[..], - ); + )[..] + } else { + &self.encode_levels_v1_bulk( + Encoding::RLE, + &self.def_levels_runs_sink[..], + self.num_levels, + max_def_level, + )[..] + }; + buffer.extend_from_slice(encoded_levels); } buffer.extend_from_slice(&values_data.buf); @@ -1159,6 +1281,8 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { self.rep_levels_sink.clear(); self.def_levels_sink.clear(); self.page_metrics.new_page(); + self.def_levels_runs_sink.clear(); + self.num_levels = 0; Ok(()) } @@ -1242,6 +1366,20 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { encoder.consume() } + /// Encodes definition or repetition levels for Data Page v1. + #[inline] + fn encode_levels_v1_bulk( + &self, + encoding: Encoding, + level_runs: &[(i16, usize)], + num_levels: usize, + max_level: i16, + ) -> Vec { + let mut encoder = LevelEncoder::v1(encoding, max_level, num_levels); + encoder.put_bulk(level_runs); + encoder.consume() + } + /// Encodes definition or repetition levels for Data Page v2. /// Encoding is always RLE. #[inline] diff --git a/parquet/src/encodings/levels.rs b/parquet/src/encodings/levels.rs index 8c95e5ca51aa..2d019da8a238 100644 --- a/parquet/src/encodings/levels.rs +++ b/parquet/src/encodings/levels.rs @@ -108,6 +108,28 @@ impl LevelEncoder { num_encoded } + /// Put/encode levels vector into this level encoder. + /// Returns number of encoded values that are less than or equal to length of the + /// input buffer. + /// Write multiple level runs as `(level, count)` pairs. + #[inline] + pub fn put_bulk(&mut self, runs: &[(i16, usize)]) -> usize { + let mut num_encoded = 0; + match *self { + LevelEncoder::Rle(ref mut encoder) | LevelEncoder::RleV2(ref mut encoder) => { + for &(value, count) in runs { + encoder.put_bulk(value as u64, count); + num_encoded += count; + } + encoder.flush(); + } + LevelEncoder::BitPacked(_, _) => { + unimplemented!() + } + } + num_encoded + } + /// Finalizes level encoder, flush all intermediate buffers and return resulting /// encoded buffer. Returned buffer is already truncated to encoded bytes only. #[inline] diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index db8227fcac3a..43c696cc44f4 100644 --- a/parquet/src/encodings/rle.rs +++ b/parquet/src/encodings/rle.rs @@ -152,6 +152,96 @@ impl RleEncoder { } } + /// Encode a single value repeated `count` times. + #[inline] + pub fn put_bulk(&mut self, value: u64, count: usize) { + assert!(count > 0, "Count must be positive"); + let remaining = 8 - self.num_buffered_values; + if self.current_value == value { + if self.repeat_count >= 8 { + self.repeat_count += count; + return; + } + let n = remaining.min(count); + for _ in 0..n { + self.buffered_values[self.num_buffered_values] = value; + self.num_buffered_values += 1; + self.repeat_count += 1; + } + if self.num_buffered_values == 8 { + assert_eq!(self.bit_packed_count % 8, 0); + self.flush_buffered_values(); + if count > n { + let mut remaining_run = count - n; + // Fill buffer + for _ in 0..remaining_run.min(8) { + self.repeat_count += 1; + remaining_run -= 1; + if self.repeat_count > 8 { + break; + } + self.buffered_values[self.num_buffered_values] = value; + self.num_buffered_values += 1; + } + if self.num_buffered_values == 8 && self.repeat_count <= 8 { + // Buffered values are full. Flush them. + assert_eq!(self.bit_packed_count % 8, 0); + self.flush_buffered_values(); + } + + self.repeat_count += remaining_run; + } + } + } else { + if self.repeat_count >= 8 { + // The current RLE run has ended and we've gathered enough. Flush first. + assert_eq!( + self.bit_packed_count, 0, + "rc = {}, value = {}, num_buffered = {}, c = {count}, {:?}", + self.repeat_count, value, self.num_buffered_values, self.buffered_values + ); + self.flush_rle_run(); + } + self.current_value = value; + let n = remaining.min(count); + self.repeat_count = 0; + for _ in 0..n { + self.buffered_values[self.num_buffered_values] = value; + self.num_buffered_values += 1; + self.repeat_count += 1; + } + if self.num_buffered_values == 8 { + let mut new_count = count; + if self.repeat_count < 8 { + new_count = count - self.repeat_count; + } + // Buffered values are full. Flush them. + assert_eq!(self.bit_packed_count % 8, 0); + self.flush_buffered_values(); + if count > n { + if self.repeat_count <= 8 { + for _ in 0..(count - n).min(8) { + self.repeat_count += 1; + if self.repeat_count > 8 { + break; + } + self.buffered_values[self.num_buffered_values] = value; + self.num_buffered_values += 1; + } + if self.num_buffered_values == 8 && self.repeat_count <= 8 { + // Buffered values are full. Flush them. + assert_eq!(self.bit_packed_count % 8, 0); + self.flush_buffered_values(); + } + } + self.repeat_count = new_count; + } + } else if count > 8 { + self.repeat_count = count; + } + } + } + #[inline] #[allow(unused)] pub fn buffer(&self) -> &[u8] { From f10dd19595699cf1061c78aa45c6ca3b97574a5e Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 14 Apr 2026 11:14:14 -0600 Subject: [PATCH 5/6] Expose option to set line terminator (#9617) --- arrow-csv/src/writer.rs | 108 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 104 insertions(+), 4 deletions(-) diff --git a/arrow-csv/src/writer.rs b/arrow-csv/src/writer.rs index 3088c12c20f4..be71f6ff2b86 100644 --- a/arrow-csv/src/writer.rs +++ b/arrow-csv/src/writer.rs @@ -197,6 +197,8 @@ pub struct WriterBuilder { quote: u8, /// Optional escape character. Defaults to `b'\\'` escape: u8, + /// Optional line terminator. Defaults to `LF` (`\n`) + terminator: Terminator, /// Enable double quote escapes. Defaults to `true` double_quote: bool, /// Optional date format for date arrays @@ -213,6 +215,15 @@ pub struct WriterBuilder { null_value: Option, } +/// The line terminator to use when writing CSV files. +#[derive(Clone, Debug)] +pub enum Terminator { + /// Use CRLF (`\r\n`) as the line terminator + CRLF, + /// Use the specified byte character as the line terminator + Any(u8), +} + impl Default for WriterBuilder { fn default() -> Self { WriterBuilder { @@ -220,6 +231,7 @@ impl Default for WriterBuilder { has_header: true, quote: b'"', escape: b'\\', + terminator: Terminator::Any(b'\n'), double_quote: true, date_format: None, datetime_format: None, @@ -389,14 +401,32 @@ impl WriterBuilder { self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE) } + /// Set the CSV file's line terminator + pub fn with_line_terminator(mut self, terminator: Terminator) -> Self { + self.terminator = terminator; + self + } + + /// Get the CSV file's line terminator, defaults to `LF` (`\n`) + pub fn line_terminator(&self) -> &Terminator { + &self.terminator + } + /// Create a new `Writer` pub fn build(self, writer: W) -> Writer { let mut builder = csv::WriterBuilder::new(); + + let terminator = match self.terminator { + Terminator::CRLF => csv::Terminator::CRLF, + Terminator::Any(byte) => csv::Terminator::Any(byte), + }; + let writer = builder .delimiter(self.delimiter) .quote(self.quote) .double_quote(self.double_quote) .escape(self.escape) + .terminator(terminator) .from_writer(writer); Writer { writer, @@ -804,10 +834,80 @@ sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo let mut buffer: Vec = vec![]; file.read_to_end(&mut buffer).unwrap(); - assert_eq!( - "c1,c2\n00:02,46:17\n00:02,\n", - String::from_utf8(buffer).unwrap() - ); + let output = String::from_utf8(buffer).unwrap(); + assert_eq!(output, "c1,c2\n00:02,46:17\n00:02,\n"); + } + + #[test] + fn test_write_csv_with_lf_terminator() { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Utf8, false), + Field::new("c2", DataType::UInt32, false), + ]); + + let c1 = StringArray::from(vec!["hello", "world"]); + let c2 = PrimitiveArray::::from(vec![1, 2]); + + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap(); + + let mut buf = Vec::new(); + let mut writer = WriterBuilder::new() + .with_line_terminator(Terminator::Any(b'\n')) + .build(&mut buf); + writer.write(&batch).unwrap(); + drop(writer); + + let output = String::from_utf8(buf).unwrap(); + assert_eq!(output, "c1,c2\nhello,1\nworld,2\n"); + } + + #[test] + fn test_write_csv_with_crlf_terminator() { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Utf8, false), + Field::new("c2", DataType::UInt32, false), + ]); + + let c1 = StringArray::from(vec!["hello", "world"]); + let c2 = PrimitiveArray::::from(vec![1, 2]); + + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap(); + + let mut buf = Vec::new(); + let mut writer = WriterBuilder::new() + .with_line_terminator(Terminator::CRLF) + .build(&mut buf); + writer.write(&batch).unwrap(); + drop(writer); + + let output = String::from_utf8(buf).unwrap(); + assert_eq!(output, "c1,c2\r\nhello,1\r\nworld,2\r\n"); + } + + #[test] + fn test_write_csv_with_any_terminator() { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Utf8, false), + Field::new("c2", DataType::UInt32, false), + ]); + + let c1 = StringArray::from(vec!["hello", "world"]); + let c2 = PrimitiveArray::::from(vec![1, 2]); + + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap(); + + let mut buf = Vec::new(); + let mut writer = WriterBuilder::new() + .with_line_terminator(Terminator::Any(b'|')) + .build(&mut buf); + writer.write(&batch).unwrap(); + drop(writer); + + let output = String::from_utf8(buf).unwrap(); + assert_eq!(output, "c1,c2|hello,1|world,2|"); } #[test] From 7d5c1c9732311b5eacfc449c117185c5f6829f01 Mon Sep 17 00:00:00 2001 From: Georgi Krastev Date: Mon, 27 Apr 2026 16:58:05 +0300 Subject: [PATCH 6/6] Fix `extend_nulls` panic for UnionArray (#9607) (#82) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Fix `MutableArrayData::extend_nulls` which previously panicked unconditionally for both sparse and dense Union arrays - For sparse unions: append the first type_id and extend nulls in all children - For dense unions: append the first type_id, compute offsets into the first child, and extend nulls in that child only ## Background This bug was discovered via DataFusion. `CaseExpr` uses `MutableArrayData` via `scatter()` to build result arrays. When a `CASE` expression returns a Union type (e.g., from `json_get` which returns a JSON union) and there are rows where no `WHEN` branch matches (implicit `ELSE NULL`), `scatter` calls `extend_nulls` which panics with "cannot call extend_nulls on UnionArray as cannot infer type". Any query like: ```sql SELECT CASE WHEN condition THEN returns_union(col, 'key') END FROM table ``` would panic if `condition` is false for any row. ## Root Cause The `extend_nulls` implementation for Union arrays unconditionally panicked because it claimed it "cannot infer type". However, the Union's field definitions (child types and type IDs) are available in the `MutableArrayData`'s data type — there's enough information to produce valid null entries by picking the first declared type_id. ## Test plan - [x] Added test for sparse union `extend_nulls` - [x] Added test for dense union `extend_nulls` - [x] Existing `test_union_dense` continues to pass - [x] All `array_transform` tests pass 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Jeffrey Vo --- arrow-data/src/transform/mod.rs | 4 +- arrow-data/src/transform/union.rs | 41 +++++++++++++-- arrow/tests/array_transform.rs | 86 +++++++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 6 deletions(-) diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index 5b994046e6ca..30fcca8d5c83 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -804,8 +804,8 @@ impl<'a> MutableArrayData<'a> { }; let nulls = match data.data_type { - // RunEndEncoded and Null arrays cannot have top-level null bitmasks - DataType::RunEndEncoded(_, _) | DataType::Null => None, + // RunEndEncoded, Null, and Union arrays cannot have top-level null bitmasks + DataType::RunEndEncoded(_, _) | DataType::Null | DataType::Union(_, _) => None, _ => data .null_buffer .map(|nulls| { diff --git a/arrow-data/src/transform/union.rs b/arrow-data/src/transform/union.rs index f6f291e3f05d..d1301249d326 100644 --- a/arrow-data/src/transform/union.rs +++ b/arrow-data/src/transform/union.rs @@ -17,6 +17,7 @@ use super::{_MutableArrayData, Extend}; use crate::ArrayData; +use arrow_schema::DataType; pub(super) fn build_extend_sparse(array: &ArrayData) -> Extend<'_> { let type_ids = array.buffer::(0); @@ -68,10 +69,42 @@ pub(super) fn build_extend_dense(array: &ArrayData) -> Extend<'_> { ) } -pub(super) fn extend_nulls_dense(_mutable: &mut _MutableArrayData, _len: usize) { - panic!("cannot call extend_nulls on UnionArray as cannot infer type"); +pub(super) fn extend_nulls_dense(mutable: &mut _MutableArrayData, len: usize) { + let DataType::Union(fields, _) = &mutable.data_type else { + unreachable!() + }; + let first_type_id = fields + .iter() + .next() + .expect("union must have at least one field") + .0; + + // Extend type_ids buffer + mutable.buffer1.extend_from_slice(&vec![first_type_id; len]); + + // Dense: extend offsets pointing into the first child, then extend nulls in that child + let child_offset = mutable.child_data[0].len(); + let (start, end) = (child_offset as i32, (child_offset + len) as i32); + mutable.buffer2.extend(start..end); + mutable.child_data[0].extend_nulls(len); } -pub(super) fn extend_nulls_sparse(_mutable: &mut _MutableArrayData, _len: usize) { - panic!("cannot call extend_nulls on UnionArray as cannot infer type"); +pub(super) fn extend_nulls_sparse(mutable: &mut _MutableArrayData, len: usize) { + let DataType::Union(fields, _) = &mutable.data_type else { + unreachable!() + }; + let first_type_id = fields + .iter() + .next() + .expect("union must have at least one field") + .0; + + // Extend type_ids buffer + mutable.buffer1.extend_from_slice(&vec![first_type_id; len]); + + // Sparse: extend nulls in ALL children + mutable + .child_data + .iter_mut() + .for_each(|child| child.extend_nulls(len)); } diff --git a/arrow/tests/array_transform.rs b/arrow/tests/array_transform.rs index 511dc1e8bfcd..9a43152c2816 100644 --- a/arrow/tests/array_transform.rs +++ b/arrow/tests/array_transform.rs @@ -1151,3 +1151,89 @@ fn test_fixed_size_list_append() { .unwrap(); assert_eq!(finished, expected_fixed_size_list_data); } + +#[test] +fn test_extend_nulls_sparse_union() { + let fields = UnionFields::new( + vec![0, 1], + vec![ + Field::new("null", DataType::Null, true), + Field::new("str", DataType::Utf8, true), + ], + ); + + let type_ids = ScalarBuffer::from(vec![1i8]); + let child_null = Arc::new(NullArray::new(1)) as ArrayRef; + let child_str = Arc::new(StringArray::from(vec![Some("hello")])) as ArrayRef; + let union_array = UnionArray::try_new( + fields.clone(), + type_ids, + None, // sparse + vec![child_null, child_str], + ) + .unwrap(); + + let data = union_array.to_data(); + let mut mutable = MutableArrayData::new(vec![&data], true, 4); + mutable.extend(0, 0, 1); // copy the first element + mutable.extend_nulls(2); // add two nulls + let result = mutable.freeze(); + + // Union arrays must not have a null bitmap per Arrow spec + assert!(result.nulls().is_none()); + + let result_array = UnionArray::from(result); + assert_eq!(result_array.len(), 3); + // First element should be type_id 1 (str) + assert_eq!(result_array.type_id(0), 1); + // Null elements use the first type_id (0) + assert_eq!(result_array.type_id(1), 0); + assert_eq!(result_array.type_id(2), 0); + // All children should have length 3 (sparse invariant) + assert_eq!(result_array.child(0).len(), 3); + assert_eq!(result_array.child(1).len(), 3); +} + +#[test] +fn test_extend_nulls_dense_union() { + let fields = UnionFields::new( + vec![0, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ); + + let type_ids = ScalarBuffer::from(vec![1i8]); + let offsets = ScalarBuffer::from(vec![0i32]); + let child_int = Arc::new(Int32Array::new_null(0)) as ArrayRef; + let child_str = Arc::new(StringArray::from(vec![Some("hello")])) as ArrayRef; + let union_array = UnionArray::try_new( + fields.clone(), + type_ids, + Some(offsets), + vec![child_int, child_str], + ) + .unwrap(); + + let data = union_array.to_data(); + let mut mutable = MutableArrayData::new(vec![&data], true, 4); + mutable.extend(0, 0, 1); // copy the first element + mutable.extend_nulls(2); // add two nulls + let result = mutable.freeze(); + + // Union arrays must not have a null bitmap per Arrow spec + assert!(result.nulls().is_none()); + + let result_array = UnionArray::from(result); + assert_eq!(result_array.len(), 3); + // First element is type_id 1 (str) + assert_eq!(result_array.type_id(0), 1); + // Null elements use the first type_id (0) + assert_eq!(result_array.type_id(1), 0); + assert_eq!(result_array.type_id(2), 0); + // First child (int) should have 2 null entries from extend_nulls + assert_eq!(result_array.child(0).len(), 2); + // Second child (str) should have 1 entry from extend + assert_eq!(result_array.child(1).len(), 1); +}