diff --git a/.asf.yaml b/.asf.yaml index 366c719597aa..dd4975435cf0 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -16,12 +16,13 @@ # under the License. # Documentation can be found here: -# https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=127405038 +# https://github.com/apache/infrastructure-asfyaml/blob/main/README.md notifications: commits: commits@arrow.apache.org issues: github@arrow.apache.org pullrequests: github@arrow.apache.org + discussions: github@arrow.apache.org jira_options: link label worklog github: description: "Official Rust implementation of Apache Arrow" @@ -29,14 +30,15 @@ github: labels: - arrow - parquet - - object-store - rust enabled_merge_buttons: squash: true + squash_commit_message: PR_TITLE_AND_DESC merge: false rebase: false features: issues: true + discussions: true protected_branches: main: required_status_checks: diff --git a/.github/dependabot.yml b/.github/dependabot.yml index b22c01f8a1b9..2da398d7d861 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,22 +1,26 @@ version: 2 updates: - package-ecosystem: cargo - directory: "/" - schedule: - interval: daily - open-pull-requests-limit: 10 - target-branch: main - labels: [ auto-dependencies, arrow ] - - package-ecosystem: cargo - directory: "/object_store" + directories: + - "/" + - "/arrow-pyarrow-integration-testing" schedule: interval: daily open-pull-requests-limit: 10 target-branch: main - labels: [ auto-dependencies, object_store ] + labels: [auto-dependencies, arrow] + groups: + prost: + applies-to: version-updates + patterns: + - "prost*" + tonic: + applies-to: version-updates + patterns: + - "tonic*" - package-ecosystem: "github-actions" directory: "/" schedule: interval: "daily" open-pull-requests-limit: 10 - labels: [ auto-dependencies ] + labels: [auto-dependencies] diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index b275cf64af1a..49b34c6137f7 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,31 +1,28 @@ # Which issue does this PR close? - +We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. -Closes #. +- Closes #NNN. # Rationale for this change - - # What changes are included in this PR? - -# Are there any user-facing changes? +# Are these changes tested? + +We typically require tests for all PRs in order to: +1. Prevent the code from being accidentally broken by subsequent changes +2. Serve as another way to document the expected behavior of the code +If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? + +# Are there any user-facing changes? - - diff --git a/.github/workflows/arrow.yml b/.github/workflows/arrow.yml index daf38f2523fc..0b90a78577e5 100644 --- a/.github/workflows/arrow.yml +++ b/.github/workflows/arrow.yml @@ -146,11 +146,11 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: - target: wasm32-unknown-unknown,wasm32-wasi + target: wasm32-unknown-unknown,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-wasi - run: cargo build -p arrow --no-default-features --features=json,csv,ipc,ffi --target wasm32-wasi + - name: Build wasm32-wasip1 + run: cargo build -p arrow --no-default-features --features=json,csv,ipc,ffi --target wasm32-wasip1 clippy: name: Clippy diff --git a/.github/workflows/arrow_flight.yml b/.github/workflows/arrow_flight.yml index 79627448ca40..a76d721b4948 100644 --- a/.github/workflows/arrow_flight.yml +++ b/.github/workflows/arrow_flight.yml @@ -60,7 +60,7 @@ jobs: cargo test -p arrow-flight --all-features - name: Test --examples run: | - cargo test -p arrow-flight --features=flight-sql-experimental,tls --examples + cargo test -p arrow-flight --features=flight-sql,tls-ring --examples vendor: name: Verify Vendored Code diff --git a/.github/workflows/dev_pr/labeler.yml b/.github/workflows/dev_pr/labeler.yml index cae015018eac..64299bd507d3 100644 --- a/.github/workflows/dev_pr/labeler.yml +++ b/.github/workflows/dev_pr/labeler.yml @@ -44,12 +44,10 @@ arrow-flight: parquet: - changed-files: - - any-glob-to-any-file: [ 'parquet/**/*' ] + - any-glob-to-any-file: + - 'parquet/**/*' + - 'parquet-variant/**/*' parquet-derive: - changed-files: - any-glob-to-any-file: [ 'parquet_derive/**/*' ] - -object-store: - - changed-files: - - any-glob-to-any-file: [ 'object_store/**/*' ] diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 9b23b1b5ad2e..09711719296c 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -41,6 +41,7 @@ on: - arrow-avro/** - arrow-ord/** - arrow-pyarrow-integration-testing/** + - arrow-pyarrow/** - arrow-schema/** - arrow-select/** - arrow-sort/** @@ -61,12 +62,11 @@ jobs: BUILD_DOCS_CPP: OFF ARROW_INTEGRATION_CPP: ON ARROW_INTEGRATION_CSHARP: ON - ARROW_INTEGRATION_GO: ON - ARROW_INTEGRATION_JAVA: ON - ARROW_INTEGRATION_JS: ON ARCHERY_INTEGRATION_TARGET_IMPLEMENTATIONS: "rust" - # Disable nanoarrow integration, due to https://github.com/apache/arrow-rs/issues/5052 - ARCHERY_INTEGRATION_WITH_NANOARROW: "0" + 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 @@ -98,16 +98,26 @@ jobs: with: path: rust fetch-depth: 0 + - name: Checkout Arrow Go + uses: actions/checkout@v4 + with: + repository: apache/arrow-go + path: go + - name: Checkout Arrow Java + uses: actions/checkout@v4 + with: + repository: apache/arrow-java + path: java + - name: Checkout Arrow JavaScript + uses: actions/checkout@v4 + with: + repository: apache/arrow-js + path: js - name: Checkout Arrow nanoarrow uses: actions/checkout@v4 with: repository: apache/arrow-nanoarrow path: nanoarrow - fetch-depth: 0 - # Workaround https://github.com/rust-lang/rust/issues/125067 - - name: Downgrade rust - working-directory: rust - run: rustup override set 1.77 - name: Build run: conda run --no-capture-output ci/scripts/integration_arrow_build.sh $PWD /build - name: Run @@ -155,8 +165,9 @@ jobs: - name: Run Rust tests run: | source venv/bin/activate - cargo test -p arrow --test pyarrow --features pyarrow - - name: Run tests + cd arrow-pyarrow-testing + cargo test + - name: Run Python tests run: | source venv/bin/activate cd arrow-pyarrow-integration-testing diff --git a/.github/workflows/object_store.yml b/.github/workflows/object_store.yml deleted file mode 100644 index 93f809aaabd4..000000000000 --- a/.github/workflows/object_store.yml +++ /dev/null @@ -1,220 +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. - ---- -# tests for `object_store` crate -name: object_store - -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: - - object_store/** - - .github/** - -jobs: - clippy: - name: Clippy - runs-on: ubuntu-latest - container: - image: amd64/rust - defaults: - run: - working-directory: object_store - steps: - - uses: actions/checkout@v4 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - - name: Setup Clippy - run: rustup component add clippy - # Run different tests for the library on its own as well as - # all targets to ensure that it still works in the absence of - # features that might be enabled by dev-dependencies of other - # targets. - - name: Run clippy with default features - run: cargo clippy -- -D warnings - - name: Run clippy with aws feature - run: cargo clippy --features aws -- -D warnings - - name: Run clippy with gcp feature - run: cargo clippy --features gcp -- -D warnings - - name: Run clippy with azure feature - run: cargo clippy --features azure -- -D warnings - - name: Run clippy with http feature - run: cargo clippy --features http -- -D warnings - - name: Run clippy with all features - run: cargo clippy --all-features -- -D warnings - - name: Run clippy with all features and all targets - run: cargo clippy --all-features --all-targets -- -D warnings - - # test doc links still work - # - # Note that since object_store is not part of the main workspace, - # this needs a separate docs job as it is not covered by - # `cargo doc --workspace` - docs: - name: Rustdocs - runs-on: ubuntu-latest - defaults: - run: - working-directory: object_store - env: - RUSTDOCFLAGS: "-Dwarnings" - steps: - - uses: actions/checkout@v4 - - name: Run cargo doc - run: cargo doc --document-private-items --no-deps --workspace --all-features - - # test the crate - # This runs outside a container to workaround lack of support for passing arguments - # to service containers - https://github.com/orgs/community/discussions/26688 - linux-test: - name: Emulator Tests - runs-on: ubuntu-latest - defaults: - run: - working-directory: object_store - env: - # Disable full debug symbol generation to speed up CI build and keep memory down - # "1" means line tables only, which is useful for panic tracebacks. - RUSTFLAGS: "-C debuginfo=1" - RUST_BACKTRACE: "1" - # Run integration tests - TEST_INTEGRATION: 1 - EC2_METADATA_ENDPOINT: http://localhost:1338 - AZURE_CONTAINER_NAME: test-bucket - AZURE_STORAGE_USE_EMULATOR: "1" - AZURITE_BLOB_STORAGE_URL: "http://localhost:10000" - AZURITE_QUEUE_STORAGE_URL: "http://localhost:10001" - AWS_BUCKET: test-bucket - AWS_DEFAULT_REGION: "us-east-1" - AWS_ACCESS_KEY_ID: test - AWS_SECRET_ACCESS_KEY: test - AWS_ENDPOINT: http://localhost:4566 - AWS_ALLOW_HTTP: true - AWS_COPY_IF_NOT_EXISTS: dynamo:test-table:2000 - AWS_CONDITIONAL_PUT: dynamo:test-table:2000 - AWS_SERVER_SIDE_ENCRYPTION: aws:kms - HTTP_URL: "http://localhost:8080" - GOOGLE_BUCKET: test-bucket - GOOGLE_SERVICE_ACCOUNT: "/tmp/gcs.json" - - steps: - - uses: actions/checkout@v4 - - # We are forced to use docker commands instead of service containers as we need to override the entrypoints - # which is currently not supported - https://github.com/actions/runner/discussions/1872 - - name: Configure Fake GCS Server (GCP emulation) - # Custom image - see fsouza/fake-gcs-server#1164 - run: | - echo "GCS_CONTAINER=$(docker run -d -p 4443:4443 tustvold/fake-gcs-server -scheme http -backend memory -public-host localhost:4443)" >> $GITHUB_ENV - # Give the container a moment to start up prior to configuring it - sleep 1 - curl -v -X POST --data-binary '{"name":"test-bucket"}' -H "Content-Type: application/json" "http://localhost:4443/storage/v1/b" - echo '{"gcs_base_url": "http://localhost:4443", "disable_oauth": true, "client_email": "", "private_key": "", "private_key_id": ""}' > "$GOOGLE_SERVICE_ACCOUNT" - - - name: Setup WebDav - run: docker run -d -p 8080:80 rclone/rclone serve webdav /data --addr :80 - - - name: Setup LocalStack (AWS emulation) - run: | - echo "LOCALSTACK_CONTAINER=$(docker run -d -p 4566:4566 localstack/localstack:4.0.3)" >> $GITHUB_ENV - echo "EC2_METADATA_CONTAINER=$(docker run -d -p 1338:1338 amazon/amazon-ec2-metadata-mock:v1.9.2 --imdsv2)" >> $GITHUB_ENV - aws --endpoint-url=http://localhost:4566 s3 mb s3://test-bucket - aws --endpoint-url=http://localhost:4566 s3api create-bucket --bucket test-object-lock --object-lock-enabled-for-bucket - aws --endpoint-url=http://localhost:4566 dynamodb create-table --table-name test-table --key-schema AttributeName=path,KeyType=HASH AttributeName=etag,KeyType=RANGE --attribute-definitions AttributeName=path,AttributeType=S AttributeName=etag,AttributeType=S --provisioned-throughput ReadCapacityUnits=5,WriteCapacityUnits=5 - - KMS_KEY=$(aws --endpoint-url=http://localhost:4566 kms create-key --description "test key") - echo "AWS_SSE_KMS_KEY_ID=$(echo $KMS_KEY | jq -r .KeyMetadata.KeyId)" >> $GITHUB_ENV - - - name: Configure Azurite (Azure emulation) - # the magical connection string is from - # https://docs.microsoft.com/en-us/azure/storage/common/storage-use-azurite?tabs=visual-studio#http-connection-strings - run: | - echo "AZURITE_CONTAINER=$(docker run -d -p 10000:10000 -p 10001:10001 -p 10002:10002 mcr.microsoft.com/azure-storage/azurite)" >> $GITHUB_ENV - az storage container create -n test-bucket --connection-string 'DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://localhost:10000/devstoreaccount1;QueueEndpoint=http://localhost:10001/devstoreaccount1;' - - - name: Setup Rust toolchain - run: | - rustup toolchain install stable - rustup default stable - - - name: Run object_store tests - run: cargo test --features=aws,azure,gcp,http - - - name: Run object_store tests (AWS native conditional put) - run: cargo test --features=aws - env: - AWS_CONDITIONAL_PUT: etag - AWS_COPY_IF_NOT_EXISTS: multipart - - - name: GCS Output - if: ${{ !cancelled() }} - run: docker logs $GCS_CONTAINER - - - name: LocalStack Output - if: ${{ !cancelled() }} - run: docker logs $LOCALSTACK_CONTAINER - - - name: EC2 Metadata Output - if: ${{ !cancelled() }} - run: docker logs $EC2_METADATA_CONTAINER - - - name: Azurite Output - if: ${{ !cancelled() }} - run: docker logs $AZURITE_CONTAINER - - # test the object_store crate builds against wasm32 in stable rust - wasm32-build: - name: Build wasm32 - runs-on: ubuntu-latest - container: - image: amd64/rust - defaults: - run: - working-directory: object_store - steps: - - uses: actions/checkout@v4 - with: - submodules: true - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - with: - target: wasm32-unknown-unknown,wasm32-wasi - - name: Build wasm32-unknown-unknown - run: cargo build --target wasm32-unknown-unknown - - name: Build wasm32-wasi - run: cargo build --target wasm32-wasi - - windows: - name: cargo test LocalFileSystem (win64) - runs-on: windows-latest - defaults: - run: - working-directory: object_store - steps: - - uses: actions/checkout@v4 - with: - submodules: true - - name: Run LocalFileSystem tests - run: cargo test local::tests diff --git a/.github/workflows/parquet-variant.yml b/.github/workflows/parquet-variant.yml new file mode 100644 index 000000000000..9e4003f3645f --- /dev/null +++ b/.github/workflows/parquet-variant.yml @@ -0,0 +1,93 @@ +# 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. + +--- +# tests for parquet-variant crate +name: "parquet-variant" + +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: + - parquet-variant/** + - parquet-variant-json/** + - parquet-variant-compute/** + - .github/** + +jobs: + # test the crate + linux-test: + name: Test + runs-on: ubuntu-latest + container: + image: amd64/rust + steps: + - uses: actions/checkout@v4 + with: + submodules: true + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + - name: Test parquet-variant + run: cargo test -p parquet-variant + - name: Test parquet-variant-json + run: cargo test -p parquet-variant-json + - name: Test parquet-variant-compute + run: cargo test -p parquet-variant-compute + + # test compilation + linux-features: + name: Check Compilation + runs-on: ubuntu-latest + container: + image: amd64/rust + steps: + - uses: actions/checkout@v4 + with: + submodules: true + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + - name: Check compilation (parquet-variant) + run: cargo check -p parquet-variant + - name: Check compilation (parquet-variant-json) + 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@v4 + - 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 2269950fd235..946aef75db19 100644 --- a/.github/workflows/parquet.yml +++ b/.github/workflows/parquet.yml @@ -97,6 +97,8 @@ jobs: run: cargo check -p parquet --no-default-features - name: Check compilation --no-default-features --features arrow run: cargo check -p parquet --no-default-features --features arrow + - name: Check compilation --no-default-features --features simdutf8 + run: cargo check -p parquet --no-default-features --features simdutf8 - name: Check compilation --no-default-features --all-features run: cargo check -p parquet --all-features - name: Check compilation --all-targets @@ -109,6 +111,15 @@ jobs: run: cargo check -p parquet --all-targets --all-features - name: Check compilation --all-targets --no-default-features --features json run: cargo check -p parquet --all-targets --no-default-features --features json + - name: Check compilation --no-default-features --features encryption --features async + run: cargo check -p parquet --no-default-features --features encryption --features async + - name: Check compilation --no-default-features --features flate2, this is expected to fail + run: if `cargo check -p parquet --no-default-features --features flate2 2>/dev/null`; then false; else true; fi + - name: Check compilation --no-default-features --features flate2 --features flate2-rust_backened + run: cargo check -p parquet --no-default-features --features flate2 --features flate2-rust_backened + - name: Check compilation --no-default-features --features flate2 --features flate2-zlib-rs + run: cargo check -p parquet --no-default-features --features flate2 --features flate2-zlib-rs + # test the parquet crate builds against wasm32 in stable rust wasm32-build: @@ -123,13 +134,13 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: - target: wasm32-unknown-unknown,wasm32-wasi + target: wasm32-unknown-unknown,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-wasi - run: cargo build -p parquet --target wasm32-wasi + - name: Build wasm32-wasip1 + run: cargo build -p parquet --target wasm32-wasip1 pyspark-integration-test: name: PySpark Integration Test diff --git a/object_store/deny.toml b/.github/workflows/release.yml similarity index 52% rename from object_store/deny.toml rename to .github/workflows/release.yml index bfd060a0b94d..8f87c50649d3 100644 --- a/object_store/deny.toml +++ b/.github/workflows/release.yml @@ -15,31 +15,31 @@ # specific language governing permissions and limitations # under the License. -# Configuration documentation: -#  https://embarkstudios.github.io/cargo-deny/index.html - -[advisories] -vulnerability = "deny" -yanked = "deny" -unmaintained = "warn" -notice = "warn" -ignore = [ -] -git-fetch-with-cli = true - -[licenses] -default = "allow" -unlicensed = "allow" -copyleft = "allow" - -[bans] -multiple-versions = "warn" -deny = [ - # We are using rustls as the TLS implementation, so we shouldn't be linking - # in OpenSSL too. - # - # If you're hitting this, you might want to take a look at what new - # dependencies you have introduced and check if there's a way to depend on - # rustls instead of OpenSSL (tip: check the crate's feature flags). - { name = "openssl-sys" } -] +# Creates a github release on https://github.com/apache/arrow-rs/releases +# when a tag is pushed to the repository +name: Release +on: + push: + tags: + - '*' + - '!*-rc*' +permissions: + contents: write +env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} +jobs: + publish: + name: Publish + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + - name: Create GitHub Releases + run: | + version=${GITHUB_REF_NAME} + title="arrow ${version}" + notes_file=CHANGELOG.md + gh release create ${GITHUB_REF_NAME} \ + --title "${title}" \ + --notes-file ${notes_file} \ + --verify-tag diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 044250b70435..38cccdec3c70 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -30,7 +30,6 @@ on: pull_request: jobs: - # Check workspace wide compile and test with default features for # mac macos: @@ -51,8 +50,8 @@ jobs: run: | # do not produce debug symbols to keep memory usage down export RUSTFLAGS="-C debuginfo=0" - cargo test - + # PyArrow tests happen in integration.yml. + cargo test --workspace # Check workspace wide compile and test with default features for # windows @@ -83,8 +82,7 @@ jobs: # do not produce debug symbols to keep memory usage down export RUSTFLAGS="-C debuginfo=0" export PATH=$PATH:/d/protoc/bin - cargo test - + cargo test --workspace # Run cargo fmt for all crates lint: @@ -108,9 +106,6 @@ 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` - - name: Format object_store - working-directory: object_store - run: cargo fmt --all -- --check msrv: name: Verify MSRV (Minimum Supported Rust Version) @@ -123,32 +118,11 @@ jobs: uses: ./.github/actions/setup-builder - name: Install cargo-msrv run: cargo install cargo-msrv - - name: Downgrade arrow dependencies - run: cargo update -p ahash --precise 0.8.7 - - name: Check arrow - working-directory: arrow - run: | - # run `cd arrow; cargo msrv verify` to see problematic dependencies - cargo msrv verify --output-format=json - - name: Check parquet - working-directory: parquet - run: | - # run `cd parquet; cargo msrv verify` to see problematic dependencies - cargo msrv verify --output-format=json - - name: Check arrow-flight - working-directory: arrow-flight - run: | - # run `cd arrow-flight; cargo msrv verify` to see problematic dependencies - cargo msrv verify --output-format=json - - name: Downgrade object_store dependencies - working-directory: object_store - # Necessary because tokio 1.30.0 updates MSRV to 1.63 - # and url 2.5.1, updates to 1.67 - run: | - cargo update -p tokio --precise 1.29.1 - cargo update -p url --precise 2.5.0 - - name: Check object_store - working-directory: object_store + - name: Check all packages run: | - # run `cd object_store; cargo msrv verify` to see problematic dependencies - cargo msrv verify --output-format=json + # 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/.gitignore b/.gitignore index 0788daea0166..05091a4e975d 100644 --- a/.gitignore +++ b/.gitignore @@ -97,3 +97,6 @@ $RECYCLE.BIN/ # Python virtual env in parquet crate parquet/pytest/venv/ __pycache__/ + +# Parquet file from arrow_reader_clickbench +hits_1.parquet \ No newline at end of file diff --git a/CHANGELOG-old.md b/CHANGELOG-old.md index 3fb17b390ac1..941c9f26382c 100644 --- a/CHANGELOG-old.md +++ b/CHANGELOG-old.md @@ -19,6 +19,498 @@ # Historical Changelog +## [55.1.0](https://github.com/apache/arrow-rs/tree/55.1.0) (2025-05-09) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/55.0.0...55.1.0) + +**Breaking changes:** + +- refactor!: do not default the struct array length to 0 in Struct::try\_new [\#7247](https://github.com/apache/arrow-rs/pull/7247) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([westonpace](https://github.com/westonpace)) + +**Implemented enhancements:** + +- Add a way to get max `usize` from `OffsetSizeTrait` [\#7474](https://github.com/apache/arrow-rs/issues/7474) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Deterministic metadata encoding [\#7448](https://github.com/apache/arrow-rs/issues/7448) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support Arrow type Dictionary with value FixedSizeBinary in Parquet [\#7445](https://github.com/apache/arrow-rs/issues/7445) +- Parquet: Add ability to project rowid in parquet reader [\#7444](https://github.com/apache/arrow-rs/issues/7444) +- Move parquet::file::metadata::reader::FooterTail to parquet::file::metadata so that it is public [\#7438](https://github.com/apache/arrow-rs/issues/7438) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Speedup take\_bytes by precalculating capacity [\#7432](https://github.com/apache/arrow-rs/issues/7432) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Improve performance of interleave\_primitive and interleave\_bytes [\#7421](https://github.com/apache/arrow-rs/issues/7421) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Implement `Eq` and `Default` for `ScalarBuffer` [\#7411](https://github.com/apache/arrow-rs/issues/7411) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add decryption support for column index and offset index [\#7390](https://github.com/apache/arrow-rs/issues/7390) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Support writing encrypted Parquet files with plaintext footers [\#7320](https://github.com/apache/arrow-rs/issues/7320) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Support Parquet key management tools [\#7256](https://github.com/apache/arrow-rs/issues/7256) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Verify footer tags when reading encrypted Parquet files with plaintext footers [\#7255](https://github.com/apache/arrow-rs/issues/7255) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- StructArray::try\_new behavior can be unexpected when there are no child arrays [\#7246](https://github.com/apache/arrow-rs/issues/7246) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Parquet performance: improve performance of reading int8/int16 [\#7097](https://github.com/apache/arrow-rs/issues/7097) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] + +**Fixed bugs:** + +- StructArray::try\_new validation incorrectly returns an error when `logical_nulls()` returns Some\(\) && null\_count == 0 [\#7435](https://github.com/apache/arrow-rs/issues/7435) +- Reading empty DataPageV2 fails with `snappy: corrupt input (empty)` [\#7388](https://github.com/apache/arrow-rs/issues/7388) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] + +**Documentation updates:** + +- Improve documentation and add examples for ArrowPredicateFn [\#7480](https://github.com/apache/arrow-rs/pull/7480) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Document Arrow \<--\> Parquet schema conversion better [\#7479](https://github.com/apache/arrow-rs/pull/7479) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Fix a typo in arrow/examples/README.md [\#7473](https://github.com/apache/arrow-rs/pull/7473) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Mottl](https://github.com/Mottl)) + +**Closed issues:** + +- Refactor Parquet DecryptionPropertiesBuilder to fix use of unreachable [\#7476](https://github.com/apache/arrow-rs/issues/7476) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Implement `Eq` and `Default` for `OffsetBuffer` [\#7417](https://github.com/apache/arrow-rs/issues/7417) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] + +**Merged pull requests:** + +- Add Parquet `arrow_reader` benchmarks for {u}int{8,16} columns [\#7484](https://github.com/apache/arrow-rs/pull/7484) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- fix: `rustdoc::unportable_markdown` was removed [\#7483](https://github.com/apache/arrow-rs/pull/7483) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([crepererum](https://github.com/crepererum)) +- Support round trip reading / writing Arrow `Duration` type to parquet [\#7482](https://github.com/apache/arrow-rs/pull/7482) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([Liyixin95](https://github.com/Liyixin95)) +- Add const MAX\_OFFSET to OffsetSizeTrait [\#7478](https://github.com/apache/arrow-rs/pull/7478) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([thinkharderdev](https://github.com/thinkharderdev)) +- Refactor Parquet DecryptionPropertiesBuilder [\#7477](https://github.com/apache/arrow-rs/pull/7477) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Support parsing and display pretty for StructType [\#7469](https://github.com/apache/arrow-rs/pull/7469) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([goldmedal](https://github.com/goldmedal)) +- chore\(deps\): update sysinfo requirement from 0.34.0 to 0.35.0 [\#7462](https://github.com/apache/arrow-rs/pull/7462) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Verify footer tags when reading encrypted Parquet files with plaintext footers [\#7459](https://github.com/apache/arrow-rs/pull/7459) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([rok](https://github.com/rok)) +- Improve comments for avro [\#7449](https://github.com/apache/arrow-rs/pull/7449) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kumarlokesh](https://github.com/kumarlokesh)) +- feat: Support round trip reading/writing Arrow type `Dictionary(_, FixedSizeBinary(_))` to Parquet [\#7446](https://github.com/apache/arrow-rs/pull/7446) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([albertlockett](https://github.com/albertlockett)) +- Fix out of bounds crash in RleValueDecoder [\#7441](https://github.com/apache/arrow-rs/pull/7441) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([apilloud](https://github.com/apilloud)) +- Make `FooterTail` public [\#7440](https://github.com/apache/arrow-rs/pull/7440) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([masonh22](https://github.com/masonh22)) +- Support writing encrypted Parquet files with plaintext footers [\#7439](https://github.com/apache/arrow-rs/pull/7439) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([rok](https://github.com/rok)) +- feat: deterministic metadata encoding [\#7437](https://github.com/apache/arrow-rs/pull/7437) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([timsaucer](https://github.com/timsaucer)) +- Fix validation logic in `StructArray::try_new` to account for array.logical\_nulls\(\) returning Some\(\) and null\_count == 0 [\#7436](https://github.com/apache/arrow-rs/pull/7436) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([phillipleblanc](https://github.com/phillipleblanc)) +- Minor: Fix typo in async\_reader comment [\#7433](https://github.com/apache/arrow-rs/pull/7433) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([amoeba](https://github.com/amoeba)) +- feat: coerce fixed size binary to binary view [\#7431](https://github.com/apache/arrow-rs/pull/7431) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([chenkovsky](https://github.com/chenkovsky)) +- chore\(deps\): update brotli requirement from 7.0 to 8.0 [\#7430](https://github.com/apache/arrow-rs/pull/7430) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Speedup take\_bytes \(-35% -69%\) by precalculating capacity [\#7422](https://github.com/apache/arrow-rs/pull/7422) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Improve performance of interleave\_primitive \(-15% - 45%\) / interleave\_bytes \(-10-25%\) [\#7420](https://github.com/apache/arrow-rs/pull/7420) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Implement `Eq` and `Default` for `OffsetBuffer` [\#7418](https://github.com/apache/arrow-rs/pull/7418) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kylebarron](https://github.com/kylebarron)) +- Implement `Default` for `Buffer` & `ScalarBuffer` [\#7413](https://github.com/apache/arrow-rs/pull/7413) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([emilk](https://github.com/emilk)) +- Implement `Eq` for `ScalarBuffer` when `T: Eq` [\#7412](https://github.com/apache/arrow-rs/pull/7412) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([emilk](https://github.com/emilk)) +- Skip page should also support skip dict page [\#7409](https://github.com/apache/arrow-rs/pull/7409) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Replace `RecordBatch::with_schema_unchecked` with `RecordBatch::new_unchecked` [\#7405](https://github.com/apache/arrow-rs/pull/7405) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold)) +- feat: Adding `with_schema_unchecked` method for `RecordBatch` [\#7402](https://github.com/apache/arrow-rs/pull/7402) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([comphead](https://github.com/comphead)) +- Add benchmark for parquet reader with row\_filter and project settings [\#7401](https://github.com/apache/arrow-rs/pull/7401) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Parquet: Expose accessors from `ArrowReaderOptions` [\#7400](https://github.com/apache/arrow-rs/pull/7400) ([kylebarron](https://github.com/kylebarron)) +- Support decryption of Parquet column and offset indexes [\#7399](https://github.com/apache/arrow-rs/pull/7399) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Handle compressed empty DataPage v2 [\#7389](https://github.com/apache/arrow-rs/pull/7389) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([EnricoMi](https://github.com/EnricoMi)) +- Improve performance of reading int8/int16 Parquet data [\#7055](https://github.com/apache/arrow-rs/pull/7055) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) + +## [55.0.0](https://github.com/apache/arrow-rs/tree/55.0.0) (2025-04-08) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/54.3.1...55.0.0) + +**Breaking changes:** + +- Change Parquet API interaction to use `u64` \(support files larger than 4GB in WASM\) [\#7371](https://github.com/apache/arrow-rs/pull/7371) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([kylebarron](https://github.com/kylebarron)) +- Remove `AsyncFileReader::get_metadata_with_options`, add `options` to `AsyncFileReader::get_metadata` [\#7342](https://github.com/apache/arrow-rs/pull/7342) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([corwinjoy](https://github.com/corwinjoy)) +- Parquet: Support reading Parquet metadata via suffix range requests [\#7334](https://github.com/apache/arrow-rs/pull/7334) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([kylebarron](https://github.com/kylebarron)) +- Upgrade to `object_store` to `0.12.0` [\#7328](https://github.com/apache/arrow-rs/pull/7328) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbrobbel](https://github.com/mbrobbel)) +- Upgrade `pyo3` to `0.24` [\#7324](https://github.com/apache/arrow-rs/pull/7324) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([mbrobbel](https://github.com/mbrobbel)) +- Reapply Box `FlightErrror::tonic` to reduce size \(fixes nightly clippy\) [\#7277](https://github.com/apache/arrow-rs/pull/7277) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([alamb](https://github.com/alamb)) +- Improve parquet gzip compression performance using zlib-rs [\#7200](https://github.com/apache/arrow-rs/pull/7200) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([psvri](https://github.com/psvri)) +- Fix: `date_part` to extract only the requested part \(not the overall interval\) [\#7189](https://github.com/apache/arrow-rs/pull/7189) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([delamarch3](https://github.com/delamarch3)) +- chore: upgrade flatbuffer version to `25.2.10` [\#7134](https://github.com/apache/arrow-rs/pull/7134) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tisonkun](https://github.com/tisonkun)) +- Add hooks to json encoder to override default encoding or add support for unsupported types [\#7015](https://github.com/apache/arrow-rs/pull/7015) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([adriangb](https://github.com/adriangb)) + +**Implemented enhancements:** + +- Improve the performance of `concat` [\#7357](https://github.com/apache/arrow-rs/issues/7357) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Pushdown predictions to Parquet in-memory row group fetches [\#7348](https://github.com/apache/arrow-rs/issues/7348) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Improve CSV parsing errors: Print the row that makes csv parsing fails [\#7344](https://github.com/apache/arrow-rs/issues/7344) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support ColumnMetaData `encoding_stats` in Parquet Writing [\#7341](https://github.com/apache/arrow-rs/issues/7341) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Support writing Parquet with modular encryption [\#7327](https://github.com/apache/arrow-rs/issues/7327) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Parquet Use U64 Instead of Usize \(wasm support for files greater than 4GB\) [\#7238](https://github.com/apache/arrow-rs/issues/7238) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Support different TimeUnits and timezones when reading Timestamps from INT96 [\#7220](https://github.com/apache/arrow-rs/issues/7220) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] + +**Fixed bugs:** + +- New clippy failures in code base with release of rustc 1.86 [\#7381](https://github.com/apache/arrow-rs/issues/7381) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Fix bug in `ParquetMetaDataReader` and add test of suffix metadata reads with encryption [\#7372](https://github.com/apache/arrow-rs/pull/7372) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) + +**Documentation updates:** + +- Improve documentation on `ArrayData::offset` [\#7385](https://github.com/apache/arrow-rs/pull/7385) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Improve documentation for `AsyncFileReader::get_metadata` [\#7380](https://github.com/apache/arrow-rs/pull/7380) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Improve documentation on implementing Parquet predicate pushdown [\#7370](https://github.com/apache/arrow-rs/pull/7370) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Add documentation and examples for pretty printing, make `pretty_format_columns_with_options` pub [\#7346](https://github.com/apache/arrow-rs/pull/7346) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Improve documentation on writing parquet, including multiple threads [\#7321](https://github.com/apache/arrow-rs/pull/7321) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) + +**Merged pull requests:** + +- chore: apply clippy suggestions newly introduced in rust 1.86 [\#7382](https://github.com/apache/arrow-rs/pull/7382) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([westonpace](https://github.com/westonpace)) +- bench: add more {boolean, string, int} benchmarks for concat kernel [\#7376](https://github.com/apache/arrow-rs/pull/7376) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Add more examples of using Parquet encryption [\#7374](https://github.com/apache/arrow-rs/pull/7374) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Clean up `ArrowReaderMetadata::load_async` [\#7369](https://github.com/apache/arrow-rs/pull/7369) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- bump pyo3 for RUSTSEC-2025-0020 [\#7368](https://github.com/apache/arrow-rs/pull/7368) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([onursatici](https://github.com/onursatici)) +- Test int96 Parquet file from Spark [\#7367](https://github.com/apache/arrow-rs/pull/7367) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbutrovich](https://github.com/mbutrovich)) +- fix: respect offset/length when converting ArrayData to StructArray [\#7366](https://github.com/apache/arrow-rs/pull/7366) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([westonpace](https://github.com/westonpace)) +- Print row, data present, expected type, and row number in error messages for arrow-csv [\#7361](https://github.com/apache/arrow-rs/pull/7361) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([psiayn](https://github.com/psiayn)) +- Use rust builtins for round\_upto\_multiple\_of\_64 and ceil [\#7358](https://github.com/apache/arrow-rs/pull/7358) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([psvri](https://github.com/psvri)) +- Write parquet PageEncodingStats [\#7354](https://github.com/apache/arrow-rs/pull/7354) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jhorstmann](https://github.com/jhorstmann)) +- Move `sysinfo` to `dev-dependencies` [\#7353](https://github.com/apache/arrow-rs/pull/7353) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbrobbel](https://github.com/mbrobbel)) +- chore\(deps\): update sysinfo requirement from 0.33.0 to 0.34.0 [\#7352](https://github.com/apache/arrow-rs/pull/7352) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Add additional benchmarks for utf8view comparison kernels [\#7351](https://github.com/apache/arrow-rs/pull/7351) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Upgrade to twox-hash 2.0 [\#7347](https://github.com/apache/arrow-rs/pull/7347) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- refactor: apply borrowed chunk reader to Sbbf::read\_from\_column\_chunk [\#7345](https://github.com/apache/arrow-rs/pull/7345) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([ethe](https://github.com/ethe)) +- Merge changelog and version from 54.3.1 into main [\#7340](https://github.com/apache/arrow-rs/pull/7340) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([timsaucer](https://github.com/timsaucer)) +- Remove `object-store` label from `.asf.yaml` [\#7339](https://github.com/apache/arrow-rs/pull/7339) ([mbrobbel](https://github.com/mbrobbel)) +- Encapsulate encryption code more in readers [\#7337](https://github.com/apache/arrow-rs/pull/7337) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Bump MSRV to 1.81 [\#7336](https://github.com/apache/arrow-rs/pull/7336) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([mbrobbel](https://github.com/mbrobbel)) +- Add an option to show column type [\#7335](https://github.com/apache/arrow-rs/pull/7335) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([blaginin](https://github.com/blaginin)) +- Add missing type annotation [\#7326](https://github.com/apache/arrow-rs/pull/7326) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbrobbel](https://github.com/mbrobbel)) +- Minor: Improve parallel parquet encoding example [\#7323](https://github.com/apache/arrow-rs/pull/7323) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- feat: allow if expressions for fallbacks in downcast macro [\#7322](https://github.com/apache/arrow-rs/pull/7322) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Minor: rename `ParquetRecordBatchStream::reader` to `ParquetRecordBatchStream::reader_factory` [\#7319](https://github.com/apache/arrow-rs/pull/7319) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- bugfix: correct offsets when serializing a list of fixed sized list and non-zero start offset [\#7318](https://github.com/apache/arrow-rs/pull/7318) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([timsaucer](https://github.com/timsaucer)) +- Remove object\_store references in Readme.md [\#7317](https://github.com/apache/arrow-rs/pull/7317) ([alamb](https://github.com/alamb)) +- Adopt MSRV policy [\#7314](https://github.com/apache/arrow-rs/pull/7314) ([psvri](https://github.com/psvri)) +- fix: correct array length validation error message [\#7313](https://github.com/apache/arrow-rs/pull/7313) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([wkalt](https://github.com/wkalt)) +- chore: remove trailing space in debug print [\#7311](https://github.com/apache/arrow-rs/pull/7311) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([xxchan](https://github.com/xxchan)) +- Improve `concat` performance, and add `append_array` for some array builder implementations [\#7309](https://github.com/apache/arrow-rs/pull/7309) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- feat: add `append_buffer` for `NullBufferBuilder` [\#7308](https://github.com/apache/arrow-rs/pull/7308) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- MINOR: fix incorrect method name in deprecate node [\#7306](https://github.com/apache/arrow-rs/pull/7306) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([waynexia](https://github.com/waynexia)) +- Allow retrieving Parquet decryption keys using the key metadata [\#7286](https://github.com/apache/arrow-rs/pull/7286) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Support different TimeUnits and timezones when reading Timestamps from INT96 [\#7285](https://github.com/apache/arrow-rs/pull/7285) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbutrovich](https://github.com/mbutrovich)) +- Add Parquet Modular encryption support \(write\) [\#7111](https://github.com/apache/arrow-rs/pull/7111) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([rok](https://github.com/rok)) + +## [54.3.1](https://github.com/apache/arrow-rs/tree/54.3.1) (2025-03-26) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/54.3.0...54.3.1) + +**Fixed bugs:** + +- Round trip encoding of list of fixed list fails when offset is not zero [\#7315](https://github.com/apache/arrow-rs/issues/7315) + +**Merged pull requests:** + +- Add missing type annotation [\#7326](https://github.com/apache/arrow-rs/pull/7326) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([mbrobbel](https://github.com/mbrobbel)) +- bugfix: correct offsets when serializing a list of fixed sized list and non-zero start offset [\#7318](https://github.com/apache/arrow-rs/pull/7318) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([timsaucer](https://github.com/timsaucer)) +## [54.3.0](https://github.com/apache/arrow-rs/tree/54.3.0) (2025-03-17) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/53.4.1...54.3.0) + +**Implemented enhancements:** + +- Using column chunk offset index in `InMemoryRowGroup::fetch` [\#7300](https://github.com/apache/arrow-rs/issues/7300) +- Support reading parquet with modular encryption [\#7296](https://github.com/apache/arrow-rs/issues/7296) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Add example for how to read/write encrypted parquet files [\#7281](https://github.com/apache/arrow-rs/issues/7281) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Have writer return parsed `ParquetMetadata` [\#7254](https://github.com/apache/arrow-rs/issues/7254) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- feat: Support Utf8View in JSON reader [\#7244](https://github.com/apache/arrow-rs/issues/7244) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- StructBuilder should provide a way to get a &dyn ArrayBuilder of a field builder [\#7193](https://github.com/apache/arrow-rs/issues/7193) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support div\_wrapping/rem\_wrapping for numeric arithmetic kernels [\#7158](https://github.com/apache/arrow-rs/issues/7158) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Improve RleDecoder performance [\#7195](https://github.com/apache/arrow-rs/pull/7195) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([Dandandan](https://github.com/Dandandan)) +- Improve arrow-json deserialization performance by 30% [\#7157](https://github.com/apache/arrow-rs/pull/7157) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([mwylde](https://github.com/mwylde)) +- Add `with_skip_validation` flag to IPC `StreamReader`, `FileReader` and `FileDecoder` [\#7120](https://github.com/apache/arrow-rs/pull/7120) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) + +**Fixed bugs:** + +- Archery integration CI test is failing on main: error: package `half v2.5.0` cannot be built because it requires rustc 1.81 or newer, while the currently active rustc version is 1.77.2 [\#7291](https://github.com/apache/arrow-rs/issues/7291) +- MSRV CI check is failing on main [\#7289](https://github.com/apache/arrow-rs/issues/7289) +- Incorrect IPC schema encoding for multiple dictionaries [\#7058](https://github.com/apache/arrow-rs/issues/7058) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] + +**Documentation updates:** + +- Add example for how to read encrypted parquet files [\#7283](https://github.com/apache/arrow-rs/pull/7283) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([rok](https://github.com/rok)) +- Update the relative path of the test data in docs [\#7221](https://github.com/apache/arrow-rs/pull/7221) ([Ziy1-Tan](https://github.com/Ziy1-Tan)) +- Minor: fix doc and remove unused code [\#7194](https://github.com/apache/arrow-rs/pull/7194) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([lewiszlw](https://github.com/lewiszlw)) +- doc: modify wrong comment [\#7190](https://github.com/apache/arrow-rs/pull/7190) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([YichiZhang0613](https://github.com/YichiZhang0613)) +- doc: fix IPC file reader/writer docs [\#7178](https://github.com/apache/arrow-rs/pull/7178) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Jefffrey](https://github.com/Jefffrey)) + +**Merged pull requests:** + +- chore: require ffi feature in arrow-schema benchmark [\#7298](https://github.com/apache/arrow-rs/pull/7298) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ethe](https://github.com/ethe)) +- Fix archery integration test [\#7292](https://github.com/apache/arrow-rs/pull/7292) ([alamb](https://github.com/alamb)) +- Minor: run `test_decimal_list` again [\#7282](https://github.com/apache/arrow-rs/pull/7282) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Move Parquet encryption tests into the arrow\_reader integration tests [\#7279](https://github.com/apache/arrow-rs/pull/7279) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Include license and notice files in published crates, part 2 [\#7275](https://github.com/apache/arrow-rs/pull/7275) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ankane](https://github.com/ankane)) +- feat: Support Utf8View in JSON reader [\#7263](https://github.com/apache/arrow-rs/pull/7263) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- feat: use `force_validate` feature flag when creating an arrays [\#7241](https://github.com/apache/arrow-rs/pull/7241) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- fix: take on empty struct array returns empty array [\#7224](https://github.com/apache/arrow-rs/pull/7224) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([westonpace](https://github.com/westonpace)) +- fix: correct `bloom_filter_position` description [\#7223](https://github.com/apache/arrow-rs/pull/7223) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([romanz](https://github.com/romanz)) +- Minor: Move `make_builder` into mod.rs [\#7218](https://github.com/apache/arrow-rs/pull/7218) ([lewiszlw](https://github.com/lewiszlw)) +- Expose `field_builders` in `StructBuilder` [\#7217](https://github.com/apache/arrow-rs/pull/7217) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([lewiszlw](https://github.com/lewiszlw)) +- Minor: Fix json StructMode docs links [\#7215](https://github.com/apache/arrow-rs/pull/7215) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([gstvg](https://github.com/gstvg)) +- \[main\] Bump arrow version to 54.2.1 \(\#7207\) [\#7212](https://github.com/apache/arrow-rs/pull/7212) ([alamb](https://github.com/alamb)) +- feat: add `downcast_integer_array` macro helper [\#7211](https://github.com/apache/arrow-rs/pull/7211) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Remove zstd pin [\#7199](https://github.com/apache/arrow-rs/pull/7199) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([tustvold](https://github.com/tustvold)) +- fix: Use chrono's quarter\(\) to avoid conflict [\#7198](https://github.com/apache/arrow-rs/pull/7198) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([yutannihilation](https://github.com/yutannihilation)) +- Fix some Clippy 1.85 warnings [\#7167](https://github.com/apache/arrow-rs/pull/7167) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([mbrobbel](https://github.com/mbrobbel)) +- feat: add to concat different data types error message the data types [\#7166](https://github.com/apache/arrow-rs/pull/7166) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Add Week ISO, Year ISO computation [\#7163](https://github.com/apache/arrow-rs/pull/7163) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kosiew](https://github.com/kosiew)) +- fix: create\_random\_batch fails with timestamp types having a timezone [\#7162](https://github.com/apache/arrow-rs/pull/7162) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([niebayes](https://github.com/niebayes)) +- Avoid overflow of remainder [\#7159](https://github.com/apache/arrow-rs/pull/7159) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([wForget](https://github.com/wForget)) +- fix: Data type inference for NaN, inf and -inf in csv files [\#7150](https://github.com/apache/arrow-rs/pull/7150) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Mottl](https://github.com/Mottl)) +- Preserve null dictionary values in `interleave` and `concat` kernels [\#7144](https://github.com/apache/arrow-rs/pull/7144) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kawadakk](https://github.com/kawadakk)) +- Support casting `Date` to a time zone-specific timestamp [\#7141](https://github.com/apache/arrow-rs/pull/7141) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([friendlymatthew](https://github.com/friendlymatthew)) +- Minor: Add doctest to ArrayDataBuilder::build\_unchecked [\#7139](https://github.com/apache/arrow-rs/pull/7139) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([gstvg](https://github.com/gstvg)) +- arrow-ord: add support for nested types to `partition` [\#7131](https://github.com/apache/arrow-rs/pull/7131) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([asubiotto](https://github.com/asubiotto)) +- Update prost-build requirement from =0.13.4 to =0.13.5 [\#7127](https://github.com/apache/arrow-rs/pull/7127) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Avoid use of `flatbuffers::size_prefixed_root`, fix validation error in arrow-flight [\#7109](https://github.com/apache/arrow-rs/pull/7109) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([bkietz](https://github.com/bkietz)) +- Optimise decimal casting for infallible conversions [\#7021](https://github.com/apache/arrow-rs/pull/7021) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([aweltsch](https://github.com/aweltsch)) + +## [53.4.1](https://github.com/apache/arrow-rs/tree/53.4.1) (2025-03-04) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/54.2.1...53.4.1) + +**Fixed bugs:** + +- Take empty struct array would get array with length 0 [\#7225](https://github.com/apache/arrow-rs/issues/7225) + +**Closed issues:** + +- Release arrow-rs / parquet patch version 54.2.1 \(Feb 2025\) \(HOTFIX\) [\#7209](https://github.com/apache/arrow-rs/issues/7209) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +## [54.2.1](https://github.com/apache/arrow-rs/tree/54.2.1) (2025-02-27) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/54.2.0...54.2.1) + +**Fixed bugs:** + +- Use chrono >= 0.4.34, < 0.4.40 to avoid breaking [\#7210](https://github.com/apache/arrow-rs/pull/7210) + +**Fixed bugs:** +## [54.2.0](https://github.com/apache/arrow-rs/tree/54.2.0) (2025-02-12) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/54.1.0...54.2.0) + +**Implemented enhancements:** + +- Casting from Utf8View to Dict\(k, Utf8View\) [\#7114](https://github.com/apache/arrow-rs/issues/7114) +- Support creating map arrays with key metadata [\#7100](https://github.com/apache/arrow-rs/issues/7100) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- \[parquet\] Print Parquet BasicTypeInfo id when present [\#7081](https://github.com/apache/arrow-rs/issues/7081) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Add arrow-ipc benchmarks for the IPC reader and writer [\#6968](https://github.com/apache/arrow-rs/issues/6968) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] + +**Fixed bugs:** + +- NullBufferBuilder::allocated_size Returns Size in Bits [\#7121](https://github.com/apache/arrow-rs/issues/7121) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- \[Regression in 54.0.0\]. Decimal cast to smaller precision gives invalid \(off-by-one\) result in some cases [\#7069](https://github.com/apache/arrow-rs/issues/7069) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Minor: Fix deprecated note to point to the correct const [\#7067](https://github.com/apache/arrow-rs/issues/7067) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- incorrect error message for reading definition levels [\#7056](https://github.com/apache/arrow-rs/issues/7056) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- First None in ListArray panics in `cast_with_options` [\#7043](https://github.com/apache/arrow-rs/issues/7043) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] + +**Documentation updates:** + +- Minor: Clarify documentation on `NullBufferBuilder::allocated_size` [\#7089](https://github.com/apache/arrow-rs/pull/7089) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: Update release schedule [\#7086](https://github.com/apache/arrow-rs/pull/7086) ([alamb](https://github.com/alamb)) +- Improve `ListArray` documentation for slices [\#7039](https://github.com/apache/arrow-rs/pull/7039) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) + +**Merged pull requests:** + +- fix: NullBufferBuilder::allocated_size should return Size in Bytes [\#7122](https://github.com/apache/arrow-rs/pull/7122) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([shuozel](https://github.com/shuozel)) +- minor: fix deprecated_note [\#7105](https://github.com/apache/arrow-rs/pull/7105) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Chen-Yuan-Lai](https://github.com/Chen-Yuan-Lai)) +- Minor: Fix ArrayDataBuilder::build_unchecked docs [\#7103](https://github.com/apache/arrow-rs/pull/7103) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([gstvg](https://github.com/gstvg)) +- Support setting key field in MapBuilder [\#7101](https://github.com/apache/arrow-rs/pull/7101) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rshkv](https://github.com/rshkv)) +- Add tests that arrow IPC data is validated [\#7096](https://github.com/apache/arrow-rs/pull/7096) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Print Parquet BasicTypeInfo id when present [\#7094](https://github.com/apache/arrow-rs/pull/7094) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([devinrsmith](https://github.com/devinrsmith)) +- Expose record boundary information in JSON decoder [\#7092](https://github.com/apache/arrow-rs/pull/7092) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([scovich](https://github.com/scovich)) +- Benchmarks for Arrow IPC reader [\#7091](https://github.com/apache/arrow-rs/pull/7091) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Benchmarks for Arrow IPC writer [\#7090](https://github.com/apache/arrow-rs/pull/7090) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add another decimal cast edge test case [\#7078](https://github.com/apache/arrow-rs/pull/7078) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- minor: re-export `OffsetBufferBuilder` in `arrow` crate [\#7077](https://github.com/apache/arrow-rs/pull/7077) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Support converting large dates \(i.e. +10999-12-31\) from string to Date32 [\#7074](https://github.com/apache/arrow-rs/pull/7074) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([phillipleblanc](https://github.com/phillipleblanc)) +- fix: issue introduced in \#6833 - less than equal check for scale in decimal conversion [\#7070](https://github.com/apache/arrow-rs/pull/7070) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([himadripal](https://github.com/himadripal)) +- perf: inline `from_iter` for `ScalarBuffer` [\#7066](https://github.com/apache/arrow-rs/pull/7066) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([0ax1](https://github.com/0ax1)) +- fix: first none/empty list in `ListArray` panics in `cast_with_options` [\#7065](https://github.com/apache/arrow-rs/pull/7065) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([irenjj](https://github.com/irenjj)) +- Minor: add ticket reference for todo [\#7064](https://github.com/apache/arrow-rs/pull/7064) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Refactor some decimal-related code and tests [\#7062](https://github.com/apache/arrow-rs/pull/7062) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([CurtHagenlocher](https://github.com/CurtHagenlocher)) +- fix error message for reading definition levels [\#7057](https://github.com/apache/arrow-rs/pull/7057) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jp0317](https://github.com/jp0317)) +- Update release schedule README.md [\#7053](https://github.com/apache/arrow-rs/pull/7053) ([alamb](https://github.com/alamb)) +- Support both 0x01 and 0x02 as type for list of booleans in thrift metadata [\#7052](https://github.com/apache/arrow-rs/pull/7052) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jhorstmann](https://github.com/jhorstmann)) +- Refactor arrow-ipc: Move `create_*_array` methods into `RecordBatchDecoder` [\#7029](https://github.com/apache/arrow-rs/pull/7029) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +## [54.1.0](https://github.com/apache/arrow-rs/tree/54.1.0) (2025-01-29) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/53.4.0...54.1.0) + +**Implemented enhancements:** + +- Create GitHub releases automatically on tagging [\#7041](https://github.com/apache/arrow-rs/issues/7041) +- Add required methods to access inner builder for `NullBufferBuilder` [\#7002](https://github.com/apache/arrow-rs/issues/7002) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Re-export `NullBufferBuilder` in the arrow crate [\#6975](https://github.com/apache/arrow-rs/issues/6975) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- `arrow-string` function should support binary input as well [\#6923](https://github.com/apache/arrow-rs/issues/6923) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- MMap support for IPC files [\#6709](https://github.com/apache/arrow-rs/issues/6709) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- fix: mark \(Large\)ListView as nested and support in equal data type [\#6995](https://github.com/apache/arrow-rs/pull/6995) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Expose min/max values for Decimal128/256 and improve docs [\#6992](https://github.com/apache/arrow-rs/pull/6992) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- \[Parquet\] Improve speed of dictionary encoding NaN float values [\#6953](https://github.com/apache/arrow-rs/pull/6953) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- Optimize `BooleanBufferBuilder` for non nullable columns [\#6973](https://github.com/apache/arrow-rs/issues/6973) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- `arrow::compute::concat` should merge dictionary type when concatenating list of dictionaries [\#6888](https://github.com/apache/arrow-rs/issues/6888) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Improve error message for unsupported cast between struct and other types [\#6724](https://github.com/apache/arrow-rs/issues/6724) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- implement regexp\_match, regexp\_scalar\_match and regexp\_array\_match for StringViewArray [\#6717](https://github.com/apache/arrow-rs/issues/6717) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Speed up Parquet utf8 validation [\#6667](https://github.com/apache/arrow-rs/issues/6667) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] + +**Fixed bugs:** + +- Regression: Concatenating sliced `ListArray`s is broken [\#7034](https://github.com/apache/arrow-rs/issues/7034) +- `PrimitiveDictionaryBuilder` with specific value data type and capacity [\#7011](https://github.com/apache/arrow-rs/issues/7011) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Arrow IPC Writer Panics for sliced nested arrays [\#6997](https://github.com/apache/arrow-rs/issues/6997) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- RecordBatch with no columns cannot be roundtripped through Parquet [\#6988](https://github.com/apache/arrow-rs/issues/6988) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- StringView: Using the Interleave kernel \(and potentially others\) results in many repeated buffers in variadic\_buffers [\#6780](https://github.com/apache/arrow-rs/issues/6780) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- fix prefetch of page index [\#6999](https://github.com/apache/arrow-rs/pull/6999) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adriangb](https://github.com/adriangb)) +- fix: Parquet column writer `Dictionary(_, Decimal128)` and `Dictionary(_, Decimal256)` [\#6987](https://github.com/apache/arrow-rs/pull/6987) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([korowa](https://github.com/korowa)) +- Writing floating point values containing NaN to Parquet is slow when using dictionary encoding [\#6952](https://github.com/apache/arrow-rs/issues/6952) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Public API using private types: `Buffer::from_bytes` takes unexported `Bytes` [\#6754](https://github.com/apache/arrow-rs/issues/6754) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Some MSRVs are inaccurate [\#6741](https://github.com/apache/arrow-rs/issues/6741) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] + +**Documentation updates:** + +- docs: add to bit slice iterator docs that the start value is inclusive and end value is exclusive [\#7022](https://github.com/apache/arrow-rs/pull/7022) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Fix duplicate link references in README [\#7020](https://github.com/apache/arrow-rs/pull/7020) ([Jefffrey](https://github.com/Jefffrey)) +- Enhance ListViewArray related docs [\#7007](https://github.com/apache/arrow-rs/pull/7007) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Jefffrey](https://github.com/Jefffrey)) +- Document data type support and examples to predicates `*like`, `starts_with`, `ends_with`, `contains` [\#7003](https://github.com/apache/arrow-rs/pull/7003) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: improve documentation on timezone representations [\#7000](https://github.com/apache/arrow-rs/pull/7000) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add additional documentation for UTC representation of timestamps [\#6994](https://github.com/apache/arrow-rs/pull/6994) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Abdullahsab3](https://github.com/Abdullahsab3)) +- Improve `ParquetRecordBatchStreamBuilder` docs / examples [\#6948](https://github.com/apache/arrow-rs/pull/6948) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Document the `ParquetRecordBatchStream` buffering [\#6947](https://github.com/apache/arrow-rs/pull/6947) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Minor: improve `zip` kernel docs, add examples [\#6928](https://github.com/apache/arrow-rs/pull/6928) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add doctest example for `Buffer::from_bytes` [\#6920](https://github.com/apache/arrow-rs/pull/6920) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kylebarron](https://github.com/kylebarron)) +- \[object store\] Add planned object\_store release schedule to crate readme [\#6904](https://github.com/apache/arrow-rs/pull/6904) ([alamb](https://github.com/alamb)) +- Avoid panics? [\#6737](https://github.com/apache/arrow-rs/issues/6737) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] + +**Merged pull requests:** + +- Create GitHub releases automatically on tagging [\#7042](https://github.com/apache/arrow-rs/pull/7042) ([kou](https://github.com/kou)) +- Fix `concat` for sliced `ListArrays` [\#7037](https://github.com/apache/arrow-rs/pull/7037) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: Clarify NullBufferBuilder::new capacity parameter [\#7016](https://github.com/apache/arrow-rs/pull/7016) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add `is_valid` and `truncate` methods to `NullBufferBuilder` [\#7013](https://github.com/apache/arrow-rs/pull/7013) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Chen-Yuan-Lai](https://github.com/Chen-Yuan-Lai)) +- fix: use the values builder capacity for the hash map in `PrimitiveDictionaryBuilder::new_from_builders` [\#7012](https://github.com/apache/arrow-rs/pull/7012) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Refactor ipc reading code into methods on `ArrayReader` [\#7006](https://github.com/apache/arrow-rs/pull/7006) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: make it clear Predicate is crate private [\#7001](https://github.com/apache/arrow-rs/pull/7001) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- fix: Panic on reencoding offsets in arrow-ipc with sliced nested arrays [\#6998](https://github.com/apache/arrow-rs/pull/6998) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([HawaiianSpork](https://github.com/HawaiianSpork)) +- Add check for empty schema in `parquet::schema::types::from_thrift_helper` [\#6990](https://github.com/apache/arrow-rs/pull/6990) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add example reading data from an `mmap`ed IPC file [\#6986](https://github.com/apache/arrow-rs/pull/6986) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Improve `arrow-ipc` documentation [\#6983](https://github.com/apache/arrow-rs/pull/6983) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add `simdutf8` feature to make `simdutf8` optional, consolidate `check_valid_utf8` [\#6979](https://github.com/apache/arrow-rs/pull/6979) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Export NullBufferBuilder along with BooleanBufferBuilder in `arrow` crate [\#6976](https://github.com/apache/arrow-rs/pull/6976) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: improve the documentation of NullBuffer and BooleanBuffer [\#6974](https://github.com/apache/arrow-rs/pull/6974) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Simplify Validation/Alignment APIs of `ArrayDataBuilder`: validate and align [\#6966](https://github.com/apache/arrow-rs/pull/6966) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Fix WASM CI for Rust 1.84 release [\#6963](https://github.com/apache/arrow-rs/pull/6963) ([alamb](https://github.com/alamb)) +- \[Parquet\] Add benchmark and test for writing NaNs to Parquet [\#6955](https://github.com/apache/arrow-rs/pull/6955) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([adamreeve](https://github.com/adamreeve)) +- Add `peek_next_page_offset` to `SerializedPageReader` [\#6945](https://github.com/apache/arrow-rs/pull/6945) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([XiangpengHao](https://github.com/XiangpengHao)) +- Improve `Buffer` documentation, deprecate `Buffer::from_bytes` add `From` and `From` impls [\#6939](https://github.com/apache/arrow-rs/pull/6939) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([alamb](https://github.com/alamb)) +- minor: fix test and remove println in tests [\#6935](https://github.com/apache/arrow-rs/pull/6935) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([himadripal](https://github.com/himadripal)) +- Document how to use Extend for generic methods on ArrayBuilders [\#6932](https://github.com/apache/arrow-rs/pull/6932) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([wiedld](https://github.com/wiedld)) +- \[Parquet\] Add projection utility functions [\#6931](https://github.com/apache/arrow-rs/pull/6931) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([XiangpengHao](https://github.com/XiangpengHao)) +- \[Parquet\] Reuse buffer in `ByteViewArrayDecoderPlain` [\#6930](https://github.com/apache/arrow-rs/pull/6930) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([XiangpengHao](https://github.com/XiangpengHao)) +- Support `Binary` arrays in `starts_with`, `ends_with` and `contains` [\#6926](https://github.com/apache/arrow-rs/pull/6926) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Improve the error message for casting between struct and non-struct types [\#6919](https://github.com/apache/arrow-rs/pull/6919) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([takaebato](https://github.com/takaebato)) +- Fix error message typos with Parquet compression [\#6918](https://github.com/apache/arrow-rs/pull/6918) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([orf](https://github.com/orf)) +- Expose arrow-schema methods, for use when writing parquet outside of ArrowWriter [\#6916](https://github.com/apache/arrow-rs/pull/6916) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([wiedld](https://github.com/wiedld)) +- feat\(arrow-ord\): support boolean in `rank` and add tests for sorting lists of booleans [\#6912](https://github.com/apache/arrow-rs/pull/6912) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- chore\(arrow-ord\): move `can_rank` to the `rank` file [\#6910](https://github.com/apache/arrow-rs/pull/6910) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- feat\(parquet\): Add next\_row\_group API for ParquetRecordBatchStream [\#6907](https://github.com/apache/arrow-rs/pull/6907) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([Xuanwo](https://github.com/Xuanwo)) +- feat\(arrow-select\): `concat` kernel will merge dictionary values for list of dictionaries [\#6893](https://github.com/apache/arrow-rs/pull/6893) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- add `extend_dictionary` in dictionary builder for improved performance [\#6875](https://github.com/apache/arrow-rs/pull/6875) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- \[arrow-string\] Implement string view support for `regexp_match` [\#6849](https://github.com/apache/arrow-rs/pull/6849) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tlm365](https://github.com/tlm365)) +- Add support `StringView` / `BinaryView` in `interleave` kernel [\#6779](https://github.com/apache/arrow-rs/pull/6779) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([onursatici](https://github.com/onursatici)) +- `RecordBatch` normalization \(flattening\) [\#6758](https://github.com/apache/arrow-rs/pull/6758) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ngli-me](https://github.com/ngli-me)) +## [54.0.0](https://github.com/apache/arrow-rs/tree/54.0.0) (2024-12-18) + +[Full Changelog](https://github.com/apache/arrow-rs/compare/53.3.0...54.0.0) + +**Breaking changes:** + +- avoid redundant parsing of repeated value in RleDecoder [\#6834](https://github.com/apache/arrow-rs/pull/6834) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jp0317](https://github.com/jp0317)) +- Handling nullable DictionaryArray in CSV parser [\#6830](https://github.com/apache/arrow-rs/pull/6830) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([edmondop](https://github.com/edmondop)) +- fix\(flightsql\): remove Any encoding of DoPutUpdateResult [\#6825](https://github.com/apache/arrow-rs/pull/6825) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([davisp](https://github.com/davisp)) +- arrow-ipc: Default to not preserving dict IDs [\#6788](https://github.com/apache/arrow-rs/pull/6788) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) +- Remove some very old deprecated functions [\#6774](https://github.com/apache/arrow-rs/pull/6774) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- update to pyo3 0.23.0 [\#6745](https://github.com/apache/arrow-rs/pull/6745) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([psvri](https://github.com/psvri)) +- Remove APIs deprecated since v 4.4.0 [\#6722](https://github.com/apache/arrow-rs/pull/6722) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) +- Return `None` when Parquet page indexes are not present in file [\#6639](https://github.com/apache/arrow-rs/pull/6639) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add `ParquetError::NeedMoreData` mark `ParquetError` as `non_exhaustive` [\#6630](https://github.com/apache/arrow-rs/pull/6630) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Remove APIs deprecated since v 2.0.0 [\#6609](https://github.com/apache/arrow-rs/pull/6609) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) + +**Implemented enhancements:** + +- Parquet schema hint doesn't support integer types upcasting [\#6891](https://github.com/apache/arrow-rs/issues/6891) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Parquet UTF-8 max statistics are overly pessimistic [\#6867](https://github.com/apache/arrow-rs/issues/6867) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Add builder support for Int8 keys [\#6844](https://github.com/apache/arrow-rs/issues/6844) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Formalize the name of the nested `Field` in a list [\#6784](https://github.com/apache/arrow-rs/issues/6784) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Allow disabling the writing of Parquet Offset Index [\#6778](https://github.com/apache/arrow-rs/issues/6778) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- `parquet::record::make_row` is not exposed to users, leaving no option to users to manually create `Row` objects [\#6761](https://github.com/apache/arrow-rs/issues/6761) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Avoid `from_num_days_from_ce_opt` calls in `timestamp_s_to_datetime` if we don't need [\#6746](https://github.com/apache/arrow-rs/issues/6746) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support Temporal -\> Utf8View casting [\#6734](https://github.com/apache/arrow-rs/issues/6734) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add Option To Coerce List Type on Parquet Write [\#6733](https://github.com/apache/arrow-rs/issues/6733) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support Numeric -\> Utf8View casting [\#6714](https://github.com/apache/arrow-rs/issues/6714) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support Utf8View \<=\> boolean casting [\#6713](https://github.com/apache/arrow-rs/issues/6713) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] + +**Fixed bugs:** + +- `Buffer::bit_slice` loses length with byte-aligned offsets [\#6895](https://github.com/apache/arrow-rs/issues/6895) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- parquet arrow writer doesn't track memory size correctly for fixed sized lists [\#6839](https://github.com/apache/arrow-rs/issues/6839) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Casting Decimal128 to Decimal128 with smaller precision produces incorrect results in some cases [\#6833](https://github.com/apache/arrow-rs/issues/6833) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Should empty nullable dictionary be parsed as null from arrow-csv? [\#6821](https://github.com/apache/arrow-rs/issues/6821) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Array take doesn't make fields nullable [\#6809](https://github.com/apache/arrow-rs/issues/6809) +- Arrow Flight Encodes a Slice's List Offsets If the slice offset is starts with zero [\#6803](https://github.com/apache/arrow-rs/issues/6803) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Parquet readers incorrectly interpret legacy nested lists [\#6756](https://github.com/apache/arrow-rs/issues/6756) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- filter\_bits under-allocates resulting boolean buffer [\#6750](https://github.com/apache/arrow-rs/issues/6750) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Multi-language support issues with Arrow FlightSQL client's execute\_update and execute\_ingest methods [\#6545](https://github.com/apache/arrow-rs/issues/6545) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] + +**Documentation updates:** + +- Should we document at what rate deprecated APIs are removed? [\#6851](https://github.com/apache/arrow-rs/issues/6851) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Fix docstring for `Format::with_header` in `arrow-csv` [\#6856](https://github.com/apache/arrow-rs/pull/6856) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kylebarron](https://github.com/kylebarron)) +- Add deprecation / API removal policy [\#6852](https://github.com/apache/arrow-rs/pull/6852) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Minor: add example for creating `SchemaDescriptor` [\#6841](https://github.com/apache/arrow-rs/pull/6841) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- chore: enrich panic context when BooleanBuffer fails to create [\#6810](https://github.com/apache/arrow-rs/pull/6810) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tisonkun](https://github.com/tisonkun)) + +**Closed issues:** + +- \[FlightSQL\] GetCatalogsBuilder does not sort the catalog names [\#6807](https://github.com/apache/arrow-rs/issues/6807) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Add a lint to automatically check for unused dependencies [\#6796](https://github.com/apache/arrow-rs/issues/6796) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] + +**Merged pull requests:** + +- doc: add comment for timezone string [\#6899](https://github.com/apache/arrow-rs/pull/6899) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([xxchan](https://github.com/xxchan)) +- docs: fix typo [\#6890](https://github.com/apache/arrow-rs/pull/6890) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) +- Minor: Fix deprecation notice for `arrow_to_parquet_schema` [\#6889](https://github.com/apache/arrow-rs/pull/6889) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add Field::with\_dict\_is\_ordered [\#6885](https://github.com/apache/arrow-rs/pull/6885) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Deprecate "max statistics size" property in `WriterProperties` [\#6884](https://github.com/apache/arrow-rs/pull/6884) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add deprecation warnings for everything related to `dict_id` [\#6873](https://github.com/apache/arrow-rs/pull/6873) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([brancz](https://github.com/brancz)) +- Enable matching temporal as from\_type to Utf8View [\#6872](https://github.com/apache/arrow-rs/pull/6872) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Kev1n8](https://github.com/Kev1n8)) +- Enable string-based column projections from Parquet files [\#6871](https://github.com/apache/arrow-rs/pull/6871) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Improvements to UTF-8 statistics truncation [\#6870](https://github.com/apache/arrow-rs/pull/6870) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- fix: make GetCatalogsBuilder sort catalog names [\#6864](https://github.com/apache/arrow-rs/pull/6864) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([niebayes](https://github.com/niebayes)) +- add buffered data\_pages to parquet column writer total bytes estimation [\#6862](https://github.com/apache/arrow-rs/pull/6862) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([onursatici](https://github.com/onursatici)) +- Update prost-build requirement from =0.13.3 to =0.13.4 [\#6860](https://github.com/apache/arrow-rs/pull/6860) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Minor: add comments explaining bad MSRV, output in json [\#6857](https://github.com/apache/arrow-rs/pull/6857) ([alamb](https://github.com/alamb)) +- perf: Use Cow in get\_format\_string in FFI\_ArrowSchema [\#6853](https://github.com/apache/arrow-rs/pull/6853) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([andygrove](https://github.com/andygrove)) +- chore: add cast\_decimal benchmark [\#6850](https://github.com/apache/arrow-rs/pull/6850) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([andygrove](https://github.com/andygrove)) +- arrow-array::builder: support Int8, Int16 and Int64 keys [\#6845](https://github.com/apache/arrow-rs/pull/6845) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ajwerner](https://github.com/ajwerner)) +- Add `ArrowToParquetSchemaConverter`, deprecate `arrow_to_parquet_schema` [\#6840](https://github.com/apache/arrow-rs/pull/6840) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Remove APIs deprecated in 50.0.0 [\#6838](https://github.com/apache/arrow-rs/pull/6838) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- fix: decimal conversion looses value on lower precision [\#6836](https://github.com/apache/arrow-rs/pull/6836) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([himadripal](https://github.com/himadripal)) +- Update sysinfo requirement from 0.32.0 to 0.33.0 [\#6835](https://github.com/apache/arrow-rs/pull/6835) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Optionally coerce names of maps and lists to match Parquet specification [\#6828](https://github.com/apache/arrow-rs/pull/6828) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Remove deprecated unary\_dyn and try\_unary\_dyn [\#6824](https://github.com/apache/arrow-rs/pull/6824) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Remove deprecated flight\_data\_from\_arrow\_batch [\#6823](https://github.com/apache/arrow-rs/pull/6823) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) +- \[arrow-cast\] Support cast boolean from/to string view [\#6822](https://github.com/apache/arrow-rs/pull/6822) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tlm365](https://github.com/tlm365)) +- Hook up Avro Decoder [\#6820](https://github.com/apache/arrow-rs/pull/6820) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold)) +- Fix arrow-avro compilation without default features [\#6819](https://github.com/apache/arrow-rs/pull/6819) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Support shrink to empty [\#6817](https://github.com/apache/arrow-rs/pull/6817) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold)) +- \[arrow-cast\] Support cast numeric to string view \(alternate\) [\#6816](https://github.com/apache/arrow-rs/pull/6816) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Hide implicit optional dependency features in arrow-flight [\#6806](https://github.com/apache/arrow-rs/pull/6806) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) +- fix: Encoding of List offsets was incorrect when slice offsets begin with zero [\#6805](https://github.com/apache/arrow-rs/pull/6805) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([HawaiianSpork](https://github.com/HawaiianSpork)) +- Enable unused\_crate\_dependencies Rust lint, remove unused dependencies [\#6804](https://github.com/apache/arrow-rs/pull/6804) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) +- Minor: Fix docstrings for `ColumnProperties::statistics_enabled` property [\#6798](https://github.com/apache/arrow-rs/pull/6798) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add option to disable writing of Parquet offset index [\#6797](https://github.com/apache/arrow-rs/pull/6797) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Remove unused dependencies [\#6792](https://github.com/apache/arrow-rs/pull/6792) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) +- Add `Array::shrink_to_fit(&mut self)` [\#6790](https://github.com/apache/arrow-rs/pull/6790) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([emilk](https://github.com/emilk)) +- Formalize the default nested list field name to `item` [\#6785](https://github.com/apache/arrow-rs/pull/6785) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([gruuya](https://github.com/gruuya)) +- Improve UnionArray logical\_nulls tests [\#6781](https://github.com/apache/arrow-rs/pull/6781) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([gstvg](https://github.com/gstvg)) +- Improve list builder usage example in docs [\#6775](https://github.com/apache/arrow-rs/pull/6775) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Update proc-macro2 requirement from =1.0.89 to =1.0.92 [\#6772](https://github.com/apache/arrow-rs/pull/6772) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([dependabot[bot]](https://github.com/apps/dependabot)) +- Allow NullBuffer construction directly from array [\#6769](https://github.com/apache/arrow-rs/pull/6769) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Include license and notice files in published crates [\#6767](https://github.com/apache/arrow-rs/pull/6767) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([ankane](https://github.com/ankane)) +- fix: remove redundant `bit_util::ceil` [\#6766](https://github.com/apache/arrow-rs/pull/6766) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([miroim](https://github.com/miroim)) +- Remove 'make\_row', expose a 'Row::new' method instead. [\#6763](https://github.com/apache/arrow-rs/pull/6763) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jonded94](https://github.com/jonded94)) +- Read nested Parquet 2-level lists correctly [\#6757](https://github.com/apache/arrow-rs/pull/6757) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Split `timestamp_s_to_datetime` to `date` and `time` to avoid unnecessary computation [\#6755](https://github.com/apache/arrow-rs/pull/6755) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([jayzhan211](https://github.com/jayzhan211)) +- More trivial implementation of `Box` and `Box` [\#6748](https://github.com/apache/arrow-rs/pull/6748) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([ethe](https://github.com/ethe)) +- Update cache action to v4 [\#6744](https://github.com/apache/arrow-rs/pull/6744) ([findepi](https://github.com/findepi)) +- Remove redundant implementation of `StringArrayType` [\#6743](https://github.com/apache/arrow-rs/pull/6743) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tlm365](https://github.com/tlm365)) +- Fix Dictionary logical nulls for RunArray/UnionArray Values [\#6740](https://github.com/apache/arrow-rs/pull/6740) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Allow reading Parquet maps that lack a `values` field [\#6730](https://github.com/apache/arrow-rs/pull/6730) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Improve default implementation of Array::is\_nullable [\#6721](https://github.com/apache/arrow-rs/pull/6721) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +- Fix Buffer::bit\_slice losing length with byte-aligned offsets [\#6707](https://github.com/apache/arrow-rs/pull/6707) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([itsjunetime](https://github.com/itsjunetime)) + ## [53.3.0](https://github.com/apache/arrow-rs/tree/53.3.0) (2024-11-17) [Full Changelog](https://github.com/apache/arrow-rs/compare/53.2.0...53.3.0) diff --git a/CHANGELOG.md b/CHANGELOG.md index a7f2a4ff34d1..03c5f6436fd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,116 +19,177 @@ # Changelog -## [54.0.0](https://github.com/apache/arrow-rs/tree/54.0.0) (2024-12-18) +## [55.2.0](https://github.com/apache/arrow-rs/tree/55.2.0) (2025-06-22) -[Full Changelog](https://github.com/apache/arrow-rs/compare/53.3.0...54.0.0) - -**Breaking changes:** - -- avoid redundant parsing of repeated value in RleDecoder [\#6834](https://github.com/apache/arrow-rs/pull/6834) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jp0317](https://github.com/jp0317)) -- Handling nullable DictionaryArray in CSV parser [\#6830](https://github.com/apache/arrow-rs/pull/6830) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([edmondop](https://github.com/edmondop)) -- fix\(flightsql\): remove Any encoding of DoPutUpdateResult [\#6825](https://github.com/apache/arrow-rs/pull/6825) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([davisp](https://github.com/davisp)) -- arrow-ipc: Default to not preserving dict IDs [\#6788](https://github.com/apache/arrow-rs/pull/6788) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) -- Remove some very old deprecated functions [\#6774](https://github.com/apache/arrow-rs/pull/6774) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) -- update to pyo3 0.23.0 [\#6745](https://github.com/apache/arrow-rs/pull/6745) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([psvri](https://github.com/psvri)) -- Remove APIs deprecated since v 4.4.0 [\#6722](https://github.com/apache/arrow-rs/pull/6722) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) -- Return `None` when Parquet page indexes are not present in file [\#6639](https://github.com/apache/arrow-rs/pull/6639) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Add `ParquetError::NeedMoreData` mark `ParquetError` as `non_exhaustive` [\#6630](https://github.com/apache/arrow-rs/pull/6630) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Remove APIs deprecated since v 2.0.0 [\#6609](https://github.com/apache/arrow-rs/pull/6609) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) +[Full Changelog](https://github.com/apache/arrow-rs/compare/55.1.0...55.2.0) **Implemented enhancements:** -- Parquet schema hint doesn't support integer types upcasting [\#6891](https://github.com/apache/arrow-rs/issues/6891) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- Parquet UTF-8 max statistics are overly pessimistic [\#6867](https://github.com/apache/arrow-rs/issues/6867) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- Add builder support for Int8 keys [\#6844](https://github.com/apache/arrow-rs/issues/6844) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Formalize the name of the nested `Field` in a list [\#6784](https://github.com/apache/arrow-rs/issues/6784) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] -- Allow disabling the writing of Parquet Offset Index [\#6778](https://github.com/apache/arrow-rs/issues/6778) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- `parquet::record::make_row` is not exposed to users, leaving no option to users to manually create `Row` objects [\#6761](https://github.com/apache/arrow-rs/issues/6761) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- Avoid `from_num_days_from_ce_opt` calls in `timestamp_s_to_datetime` if we don't need [\#6746](https://github.com/apache/arrow-rs/issues/6746) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Support Temporal -\> Utf8View casting [\#6734](https://github.com/apache/arrow-rs/issues/6734) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Add Option To Coerce List Type on Parquet Write [\#6733](https://github.com/apache/arrow-rs/issues/6733) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Support Numeric -\> Utf8View casting [\#6714](https://github.com/apache/arrow-rs/issues/6714) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Support Utf8View \<=\> boolean casting [\#6713](https://github.com/apache/arrow-rs/issues/6713) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Do not populate nulls for `NullArray` for `MutableArrayData` [\#7725](https://github.com/apache/arrow-rs/issues/7725) +- Implement `PartialEq` for RunArray [\#7691](https://github.com/apache/arrow-rs/issues/7691) +- `interleave_views` is really slow [\#7688](https://github.com/apache/arrow-rs/issues/7688) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add min max aggregates for FixedSizeBinary [\#7674](https://github.com/apache/arrow-rs/issues/7674) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Deliver pyarrow as a standalone crate [\#7668](https://github.com/apache/arrow-rs/issues/7668) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- \[Variant\] Implement `VariantObject::field` and `VariantObject::fields` [\#7665](https://github.com/apache/arrow-rs/issues/7665) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- \[Variant\] Implement read support for remaining primitive types [\#7630](https://github.com/apache/arrow-rs/issues/7630) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Fast and ergonomic method to add metadata to a `RecordBatch` [\#7628](https://github.com/apache/arrow-rs/issues/7628) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add efficient way to change the keys of string dictionary builder [\#7610](https://github.com/apache/arrow-rs/issues/7610) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Support `add_nulls` on additional builder types [\#7605](https://github.com/apache/arrow-rs/issues/7605) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add `into_inner` for `AsyncArrowWriter` [\#7603](https://github.com/apache/arrow-rs/issues/7603) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Optimize `PrimitiveBuilder::append_trusted_len_iter` [\#7591](https://github.com/apache/arrow-rs/issues/7591) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Benchmark for filter+concat and take+concat into even sized record batches [\#7589](https://github.com/apache/arrow-rs/issues/7589) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- `max_statistics_truncate_length` is ignored when writing statistics to data page headers [\#7579](https://github.com/apache/arrow-rs/issues/7579) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Feature Request: Encoding in `parquet-rewrite` [\#7575](https://github.com/apache/arrow-rs/issues/7575) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Add a `strong_count` method to `Buffer` [\#7568](https://github.com/apache/arrow-rs/issues/7568) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Create version of LexicographicalComparator that compares fixed number of columns [\#7531](https://github.com/apache/arrow-rs/issues/7531) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- parquet-show-bloom-filter should work with integer typed columns [\#7528](https://github.com/apache/arrow-rs/issues/7528) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Allow merging primitive dictionary values in concat and interleave kernels [\#7518](https://github.com/apache/arrow-rs/issues/7518) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Add efficient concatenation of StructArrays [\#7516](https://github.com/apache/arrow-rs/issues/7516) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Rename `flight-sql-experimental` to `flight-sql` [\#7498](https://github.com/apache/arrow-rs/issues/7498) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Consider moving from ryu to lexical-core for string formatting / casting floats to string. [\#7496](https://github.com/apache/arrow-rs/issues/7496) +- Arithmetic kernels can be safer and faster [\#7494](https://github.com/apache/arrow-rs/issues/7494) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Speedup `filter_bytes` by precalculating capacity [\#7465](https://github.com/apache/arrow-rs/issues/7465) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- \[Variant\]: Rust API to Create Variant Values [\#7424](https://github.com/apache/arrow-rs/issues/7424) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- \[Variant\] Rust API to Read Variant Values [\#7423](https://github.com/apache/arrow-rs/issues/7423) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Release arrow-rs / parquet Minor version `55.1.0` \(May 2025\) [\#7393](https://github.com/apache/arrow-rs/issues/7393) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Support create\_random\_array for Decimal data types [\#7343](https://github.com/apache/arrow-rs/issues/7343) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Truncate Parquet page data page statistics [\#7555](https://github.com/apache/arrow-rs/pull/7555) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) **Fixed bugs:** -- `Buffer::bit_slice` loses length with byte-aligned offsets [\#6895](https://github.com/apache/arrow-rs/issues/6895) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- parquet arrow writer doesn't track memory size correctly for fixed sized lists [\#6839](https://github.com/apache/arrow-rs/issues/6839) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- Casting Decimal128 to Decimal128 with smaller precision produces incorrect results in some cases [\#6833](https://github.com/apache/arrow-rs/issues/6833) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Should empty nullable dictionary be parsed as null from arrow-csv? [\#6821](https://github.com/apache/arrow-rs/issues/6821) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Array take doesn't make fields nullable [\#6809](https://github.com/apache/arrow-rs/issues/6809) -- Arrow Flight Encodes a Slice's List Offsets If the slice offset is starts with zero [\#6803](https://github.com/apache/arrow-rs/issues/6803) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Parquet readers incorrectly interpret legacy nested lists [\#6756](https://github.com/apache/arrow-rs/issues/6756) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] -- filter\_bits under-allocates resulting boolean buffer [\#6750](https://github.com/apache/arrow-rs/issues/6750) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Multi-language support issues with Arrow FlightSQL client's execute\_update and execute\_ingest methods [\#6545](https://github.com/apache/arrow-rs/issues/6545) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- In arrow\_json, Decoder::decode can panic if it encounters two high surrogates in a row. [\#7712](https://github.com/apache/arrow-rs/issues/7712) +- FlightSQL "GetDbSchemas" and "GetTables" schemas do not fully match the protocol [\#7637](https://github.com/apache/arrow-rs/issues/7637) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Cannot read encrypted Parquet file if page index reading is enabled [\#7629](https://github.com/apache/arrow-rs/issues/7629) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- `encoding_stats` not present in Parquet generated by `parquet-rewrite` [\#7616](https://github.com/apache/arrow-rs/issues/7616) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- When writing parquet plaintext footer files `footer_signing_key_metadata` is not included, encryption alghoritm is always written in footer [\#7599](https://github.com/apache/arrow-rs/issues/7599) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- `new_null_array` panics when constructing a struct of a dictionary [\#7571](https://github.com/apache/arrow-rs/issues/7571) +- Parquet derive fails to build when Result is aliased [\#7547](https://github.com/apache/arrow-rs/issues/7547) +- Unable to read `Dictionary(u8, FixedSizeBinary(_))` using datafusion. [\#7545](https://github.com/apache/arrow-rs/issues/7545) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- filter\_record\_batch panics with empty struct array. [\#7538](https://github.com/apache/arrow-rs/issues/7538) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Panic in `pretty_format` function when displaying DurationSecondsArray with `i64::MIN` / `i64::MAX` [\#7533](https://github.com/apache/arrow-rs/issues/7533) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Record API unable to parse TIME\_MILLIS when encoded as INT32 [\#7510](https://github.com/apache/arrow-rs/issues/7510) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- The `read_record_batch` func of the `RecordBatchDecoder` does not respect the `skip_validation` property [\#7508](https://github.com/apache/arrow-rs/issues/7508) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- `arrow-55.1.0` breaks `filter_record_batch` [\#7500](https://github.com/apache/arrow-rs/issues/7500) +- Files containing binary data with \>=8\_388\_855 bytes per row written with `arrow-rs` can't be read with `pyarrow` [\#7489](https://github.com/apache/arrow-rs/issues/7489) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- \[Bug\] Ingestion with Arrow Flight Sql panic when the input stream is empty or fallible [\#7329](https://github.com/apache/arrow-rs/issues/7329) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- Ensure page encoding statistics are written to Parquet file [\#7643](https://github.com/apache/arrow-rs/pull/7643) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) **Documentation updates:** -- Should we document at what rate deprecated APIs are removed? [\#6851](https://github.com/apache/arrow-rs/issues/6851) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] -- Fix docstring for `Format::with_header` in `arrow-csv` [\#6856](https://github.com/apache/arrow-rs/pull/6856) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kylebarron](https://github.com/kylebarron)) -- Add deprecation / API removal policy [\#6852](https://github.com/apache/arrow-rs/pull/6852) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) -- Minor: add example for creating `SchemaDescriptor` [\#6841](https://github.com/apache/arrow-rs/pull/6841) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) -- chore: enrich panic context when BooleanBuffer fails to create [\#6810](https://github.com/apache/arrow-rs/pull/6810) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tisonkun](https://github.com/tisonkun)) +- arrow\_reader\_row\_filter benchmark doesn't capture page cache improvements [\#7460](https://github.com/apache/arrow-rs/issues/7460) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- chore: fix a typo in `ExtensionType::supports_data_type` docs [\#7682](https://github.com/apache/arrow-rs/pull/7682) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([mbrobbel](https://github.com/mbrobbel)) +- \[Variant\] Add variant docs and examples [\#7661](https://github.com/apache/arrow-rs/pull/7661) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Minor: Add version to deprecation notice for `ParquetMetaDataReader::decode_footer` [\#7639](https://github.com/apache/arrow-rs/pull/7639) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Add references for defaults in `WriterPropertiesBuilder` [\#7558](https://github.com/apache/arrow-rs/pull/7558) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Clarify Docs: NullBuffer::len is in bits [\#7556](https://github.com/apache/arrow-rs/pull/7556) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- docs: fix typo for `Decimal128Array` [\#7525](https://github.com/apache/arrow-rs/pull/7525) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([burmecia](https://github.com/burmecia)) +- Minor: Add examples to ProjectionMask documentation [\#7523](https://github.com/apache/arrow-rs/pull/7523) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Improve documentation for Parquet `WriterProperties` [\#7491](https://github.com/apache/arrow-rs/pull/7491) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) **Closed issues:** -- \[FlightSQL\] GetCatalogsBuilder does not sort the catalog names [\#6807](https://github.com/apache/arrow-rs/issues/6807) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] -- Add a lint to automatically check for unused dependencies [\#6796](https://github.com/apache/arrow-rs/issues/6796) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] +- \[Variant\] More efficient determination of String vs ShortString [\#7700](https://github.com/apache/arrow-rs/issues/7700) +- \[Variant\] Improve API for iterating over values of a VariantList [\#7685](https://github.com/apache/arrow-rs/issues/7685) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- \[Variant\] Consider validating variants on creation \(rather than read\) [\#7684](https://github.com/apache/arrow-rs/issues/7684) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] +- Miri test\_native\_type\_pow test failing [\#7641](https://github.com/apache/arrow-rs/issues/7641) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Improve performance of `coalesce` and `concat` for views [\#7615](https://github.com/apache/arrow-rs/issues/7615) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] +- Bad min value in row group statistics in some special cases [\#7593](https://github.com/apache/arrow-rs/issues/7593) +- Feature Request: BloomFilter Position Flexibility in `parquet-rewrite` [\#7552](https://github.com/apache/arrow-rs/issues/7552) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] **Merged pull requests:** -- doc: add comment for timezone string [\#6899](https://github.com/apache/arrow-rs/pull/6899) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([xxchan](https://github.com/xxchan)) -- docs: fix typo [\#6890](https://github.com/apache/arrow-rs/pull/6890) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([rluvaton](https://github.com/rluvaton)) -- Minor: Fix deprecation notice for `arrow_to_parquet_schema` [\#6889](https://github.com/apache/arrow-rs/pull/6889) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Add Field::with\_dict\_is\_ordered [\#6885](https://github.com/apache/arrow-rs/pull/6885) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) -- Deprecate "max statistics size" property in `WriterProperties` [\#6884](https://github.com/apache/arrow-rs/pull/6884) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Add deprecation warnings for everything related to `dict_id` [\#6873](https://github.com/apache/arrow-rs/pull/6873) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([brancz](https://github.com/brancz)) -- Enable matching temporal as from\_type to Utf8View [\#6872](https://github.com/apache/arrow-rs/pull/6872) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Kev1n8](https://github.com/Kev1n8)) -- Enable string-based column projections from Parquet files [\#6871](https://github.com/apache/arrow-rs/pull/6871) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Improvements to UTF-8 statistics truncation [\#6870](https://github.com/apache/arrow-rs/pull/6870) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- fix: make GetCatalogsBuilder sort catalog names [\#6864](https://github.com/apache/arrow-rs/pull/6864) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([niebayes](https://github.com/niebayes)) -- add buffered data\_pages to parquet column writer total bytes estimation [\#6862](https://github.com/apache/arrow-rs/pull/6862) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([onursatici](https://github.com/onursatici)) -- Update prost-build requirement from =0.13.3 to =0.13.4 [\#6860](https://github.com/apache/arrow-rs/pull/6860) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([dependabot[bot]](https://github.com/apps/dependabot)) -- Minor: add comments explaining bad MSRV, output in json [\#6857](https://github.com/apache/arrow-rs/pull/6857) ([alamb](https://github.com/alamb)) -- perf: Use Cow in get\_format\_string in FFI\_ArrowSchema [\#6853](https://github.com/apache/arrow-rs/pull/6853) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([andygrove](https://github.com/andygrove)) -- chore: add cast\_decimal benchmark [\#6850](https://github.com/apache/arrow-rs/pull/6850) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([andygrove](https://github.com/andygrove)) -- arrow-array::builder: support Int8, Int16 and Int64 keys [\#6845](https://github.com/apache/arrow-rs/pull/6845) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ajwerner](https://github.com/ajwerner)) -- Add `ArrowToParquetSchemaConverter`, deprecate `arrow_to_parquet_schema` [\#6840](https://github.com/apache/arrow-rs/pull/6840) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) -- Remove APIs deprecated in 50.0.0 [\#6838](https://github.com/apache/arrow-rs/pull/6838) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- fix: decimal conversion looses value on lower precision [\#6836](https://github.com/apache/arrow-rs/pull/6836) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([himadripal](https://github.com/himadripal)) -- Update sysinfo requirement from 0.32.0 to 0.33.0 [\#6835](https://github.com/apache/arrow-rs/pull/6835) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([dependabot[bot]](https://github.com/apps/dependabot)) -- Optionally coerce names of maps and lists to match Parquet specification [\#6828](https://github.com/apache/arrow-rs/pull/6828) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Remove deprecated unary\_dyn and try\_unary\_dyn [\#6824](https://github.com/apache/arrow-rs/pull/6824) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Remove deprecated flight\_data\_from\_arrow\_batch [\#6823](https://github.com/apache/arrow-rs/pull/6823) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) -- \[arrow-cast\] Support cast boolean from/to string view [\#6822](https://github.com/apache/arrow-rs/pull/6822) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tlm365](https://github.com/tlm365)) -- Hook up Avro Decoder [\#6820](https://github.com/apache/arrow-rs/pull/6820) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold)) -- Fix arrow-avro compilation without default features [\#6819](https://github.com/apache/arrow-rs/pull/6819) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Support shrink to empty [\#6817](https://github.com/apache/arrow-rs/pull/6817) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tustvold](https://github.com/tustvold)) -- \[arrow-cast\] Support cast numeric to string view \(alternate\) [\#6816](https://github.com/apache/arrow-rs/pull/6816) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) -- Hide implicit optional dependency features in arrow-flight [\#6806](https://github.com/apache/arrow-rs/pull/6806) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) -- fix: Encoding of List offsets was incorrect when slice offsets begin with zero [\#6805](https://github.com/apache/arrow-rs/pull/6805) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([HawaiianSpork](https://github.com/HawaiianSpork)) -- Enable unused\_crate\_dependencies Rust lint, remove unused dependencies [\#6804](https://github.com/apache/arrow-rs/pull/6804) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) -- Minor: Fix docstrings for `ColumnProperties::statistics_enabled` property [\#6798](https://github.com/apache/arrow-rs/pull/6798) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Add option to disable writing of Parquet offset index [\#6797](https://github.com/apache/arrow-rs/pull/6797) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Remove unused dependencies [\#6792](https://github.com/apache/arrow-rs/pull/6792) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([findepi](https://github.com/findepi)) -- Add `Array::shrink_to_fit(&mut self)` [\#6790](https://github.com/apache/arrow-rs/pull/6790) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([emilk](https://github.com/emilk)) -- Formalize the default nested list field name to `item` [\#6785](https://github.com/apache/arrow-rs/pull/6785) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([gruuya](https://github.com/gruuya)) -- Improve UnionArray logical\_nulls tests [\#6781](https://github.com/apache/arrow-rs/pull/6781) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([gstvg](https://github.com/gstvg)) -- Improve list builder usage example in docs [\#6775](https://github.com/apache/arrow-rs/pull/6775) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Update proc-macro2 requirement from =1.0.89 to =1.0.92 [\#6772](https://github.com/apache/arrow-rs/pull/6772) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([dependabot[bot]](https://github.com/apps/dependabot)) -- Allow NullBuffer construction directly from array [\#6769](https://github.com/apache/arrow-rs/pull/6769) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Include license and notice files in published crates [\#6767](https://github.com/apache/arrow-rs/pull/6767) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([ankane](https://github.com/ankane)) -- fix: remove redundant `bit_util::ceil` [\#6766](https://github.com/apache/arrow-rs/pull/6766) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([miroim](https://github.com/miroim)) -- Remove 'make\_row', expose a 'Row::new' method instead. [\#6763](https://github.com/apache/arrow-rs/pull/6763) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jonded94](https://github.com/jonded94)) -- Read nested Parquet 2-level lists correctly [\#6757](https://github.com/apache/arrow-rs/pull/6757) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Split `timestamp_s_to_datetime` to `date` and `time` to avoid unnecessary computation [\#6755](https://github.com/apache/arrow-rs/pull/6755) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([jayzhan211](https://github.com/jayzhan211)) -- More trivial implementation of `Box` and `Box` [\#6748](https://github.com/apache/arrow-rs/pull/6748) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([ethe](https://github.com/ethe)) -- Update cache action to v4 [\#6744](https://github.com/apache/arrow-rs/pull/6744) ([findepi](https://github.com/findepi)) -- Remove redundant implementation of `StringArrayType` [\#6743](https://github.com/apache/arrow-rs/pull/6743) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([tlm365](https://github.com/tlm365)) -- Fix Dictionary logical nulls for RunArray/UnionArray Values [\#6740](https://github.com/apache/arrow-rs/pull/6740) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Allow reading Parquet maps that lack a `values` field [\#6730](https://github.com/apache/arrow-rs/pull/6730) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) -- Improve default implementation of Array::is\_nullable [\#6721](https://github.com/apache/arrow-rs/pull/6721) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([findepi](https://github.com/findepi)) -- Fix Buffer::bit\_slice losing length with byte-aligned offsets [\#6707](https://github.com/apache/arrow-rs/pull/6707) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([itsjunetime](https://github.com/itsjunetime)) +- arrow-array: Implement PartialEq for RunArray [\#7727](https://github.com/apache/arrow-rs/pull/7727) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) +- fix: Do not add null buffer for `NullArray` in MutableArrayData [\#7726](https://github.com/apache/arrow-rs/pull/7726) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([comphead](https://github.com/comphead)) +- fix JSON decoder error checking for UTF16 / surrogate parsing panic [\#7721](https://github.com/apache/arrow-rs/pull/7721) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([nicklan](https://github.com/nicklan)) +- \[Variant\] Introduce new type over &str for ShortString [\#7718](https://github.com/apache/arrow-rs/pull/7718) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([friendlymatthew](https://github.com/friendlymatthew)) +- Split out variant code into several new sub-modules [\#7717](https://github.com/apache/arrow-rs/pull/7717) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([scovich](https://github.com/scovich)) +- Support write to buffer api for SerializedFileWriter [\#7714](https://github.com/apache/arrow-rs/pull/7714) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Make variant iterators safely infallible [\#7704](https://github.com/apache/arrow-rs/pull/7704) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([scovich](https://github.com/scovich)) +- Speedup `interleave_views` \(4-7x faster\) [\#7695](https://github.com/apache/arrow-rs/pull/7695) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Define a "arrow-pyrarrow" crate to implement the "pyarrow" feature. [\#7694](https://github.com/apache/arrow-rs/pull/7694) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brunal](https://github.com/brunal)) +- Document REE row format and add some more tests [\#7680](https://github.com/apache/arrow-rs/pull/7680) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- feat: add min max aggregate support for FixedSizeBinary [\#7675](https://github.com/apache/arrow-rs/pull/7675) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alexwilcoxson-rel](https://github.com/alexwilcoxson-rel)) +- arrow-data: Add REE support for `build_extend` and `build_extend_nulls` [\#7671](https://github.com/apache/arrow-rs/pull/7671) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) +- Remove `lazy_static` dependency [\#7669](https://github.com/apache/arrow-rs/pull/7669) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Expyron](https://github.com/Expyron)) +- Finish implementing Variant::Object and Variant::List [\#7666](https://github.com/apache/arrow-rs/pull/7666) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([scovich](https://github.com/scovich)) +- Add `RecordBatch::schema_metadata_mut` and `Field::metadata_mut` [\#7664](https://github.com/apache/arrow-rs/pull/7664) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([emilk](https://github.com/emilk)) +- \[Variant\] Simplify creation of Variants from metadata and value [\#7663](https://github.com/apache/arrow-rs/pull/7663) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- chore: group prost dependabot updates [\#7659](https://github.com/apache/arrow-rs/pull/7659) ([mbrobbel](https://github.com/mbrobbel)) +- Initial Builder API for Creating Variant Values [\#7653](https://github.com/apache/arrow-rs/pull/7653) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([PinkCrow007](https://github.com/PinkCrow007)) +- Add `BatchCoalescer::push_filtered_batch` and docs [\#7652](https://github.com/apache/arrow-rs/pull/7652) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Optimize coalesce kernel for StringView \(10-50% faster\) [\#7650](https://github.com/apache/arrow-rs/pull/7650) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- arrow-row: Add support for REE [\#7649](https://github.com/apache/arrow-rs/pull/7649) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) +- Use approximate comparisons for pow tests [\#7646](https://github.com/apache/arrow-rs/pull/7646) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([adamreeve](https://github.com/adamreeve)) +- \[Variant\] Implement read support for remaining primitive types [\#7644](https://github.com/apache/arrow-rs/pull/7644) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([superserious-dev](https://github.com/superserious-dev)) +- Add `pretty_format_batches_with_schema` function [\#7642](https://github.com/apache/arrow-rs/pull/7642) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([lewiszlw](https://github.com/lewiszlw)) +- Deprecate old Parquet page index parsing functions [\#7640](https://github.com/apache/arrow-rs/pull/7640) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Update FlightSQL `GetDbSchemas` and `GetTables` schemas to fully match the protocol [\#7638](https://github.com/apache/arrow-rs/pull/7638) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([sgrebnov](https://github.com/sgrebnov)) +- Minor: Remove outdated FIXME from `ParquetMetaDataReader` [\#7635](https://github.com/apache/arrow-rs/pull/7635) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Fix the error info of `StructArray::try_new` [\#7634](https://github.com/apache/arrow-rs/pull/7634) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([xudong963](https://github.com/xudong963)) +- Fix reading encrypted Parquet pages when using the page index [\#7633](https://github.com/apache/arrow-rs/pull/7633) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- \[Variant\] Add commented out primitive test casees [\#7631](https://github.com/apache/arrow-rs/pull/7631) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Improve `coalesce` kernel tests [\#7626](https://github.com/apache/arrow-rs/pull/7626) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Revert "Revert "Improve `coalesce` and `concat` performance for views… [\#7625](https://github.com/apache/arrow-rs/pull/7625) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Revert "Improve `coalesce` and `concat` performance for views \(\#7614\)" [\#7623](https://github.com/apache/arrow-rs/pull/7623) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Improve coalesce\_kernel benchmark to capture inline vs non inline views [\#7619](https://github.com/apache/arrow-rs/pull/7619) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Improve `coalesce` and `concat` performance for views [\#7614](https://github.com/apache/arrow-rs/pull/7614) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- feat: add constructor to help efficiently upgrade key for GenericBytesDictionaryBuilder [\#7611](https://github.com/apache/arrow-rs/pull/7611) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([albertlockett](https://github.com/albertlockett)) +- feat: support append\_nulls on additional builders [\#7606](https://github.com/apache/arrow-rs/pull/7606) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([albertlockett](https://github.com/albertlockett)) +- feat: add AsyncArrowWriter::into\_inner [\#7604](https://github.com/apache/arrow-rs/pull/7604) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([jpopesculian](https://github.com/jpopesculian)) +- Move variant interop test to Rust integration test [\#7602](https://github.com/apache/arrow-rs/pull/7602) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Include footer key metadata when writing encrypted Parquet with a plaintext footer [\#7600](https://github.com/apache/arrow-rs/pull/7600) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([rok](https://github.com/rok)) +- Add `coalesce` kernel and`BatchCoalescer` for statefully combining selected b…atches: [\#7597](https://github.com/apache/arrow-rs/pull/7597) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add FixedSizeBinary to `take_kernel` benchmark [\#7592](https://github.com/apache/arrow-rs/pull/7592) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Fix GenericBinaryArray docstring. [\#7588](https://github.com/apache/arrow-rs/pull/7588) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brunal](https://github.com/brunal)) +- fix: error reading multiple batches of `Dict(_, FixedSizeBinary(_))` [\#7585](https://github.com/apache/arrow-rs/pull/7585) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([albertlockett](https://github.com/albertlockett)) +- Revert "Minor: remove filter code deprecated in 2023 \(\#7554\)" [\#7583](https://github.com/apache/arrow-rs/pull/7583) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Fixed a warning build build: function never used. [\#7577](https://github.com/apache/arrow-rs/pull/7577) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([JigaoLuo](https://github.com/JigaoLuo)) +- Adding Encoding argument in `parquet-rewrite` [\#7576](https://github.com/apache/arrow-rs/pull/7576) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([JigaoLuo](https://github.com/JigaoLuo)) +- feat: add `row_group_is_[max/min]_value_exact` to StatisticsConverter [\#7574](https://github.com/apache/arrow-rs/pull/7574) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([CookiePieWw](https://github.com/CookiePieWw)) +- \[array\] Remove unwrap checks from GenericByteArray::value\_unchecked [\#7573](https://github.com/apache/arrow-rs/pull/7573) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ctsk](https://github.com/ctsk)) +- \[benches/row\_format\] fix typo in array lengths [\#7572](https://github.com/apache/arrow-rs/pull/7572) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ctsk](https://github.com/ctsk)) +- Add a strong\_count method to Buffer [\#7569](https://github.com/apache/arrow-rs/pull/7569) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([westonpace](https://github.com/westonpace)) +- Minor: Enable byte view for clickbench benchmark [\#7565](https://github.com/apache/arrow-rs/pull/7565) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Optimize length calculation in row encoding for fixed-length columns [\#7564](https://github.com/apache/arrow-rs/pull/7564) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([ctsk](https://github.com/ctsk)) +- Use PR title and description for commit message [\#7563](https://github.com/apache/arrow-rs/pull/7563) ([kou](https://github.com/kou)) +- Use apache/arrow-{go,java,js} in integration test [\#7561](https://github.com/apache/arrow-rs/pull/7561) ([kou](https://github.com/kou)) +- Implement Array Decoding in arrow-avro [\#7559](https://github.com/apache/arrow-rs/pull/7559) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([jecsand838](https://github.com/jecsand838)) +- Minor: remove filter code deprecated in 2023 [\#7554](https://github.com/apache/arrow-rs/pull/7554) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- fix: Correct docs for `WriterPropertiesBuilder::set_column_index_truncate_length` [\#7553](https://github.com/apache/arrow-rs/pull/7553) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([etseidl](https://github.com/etseidl)) +- Adding Bloom Filter Position argument in parquet-rewrite [\#7550](https://github.com/apache/arrow-rs/pull/7550) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([JigaoLuo](https://github.com/JigaoLuo)) +- Fix `Result` name collision in parquet\_derive [\#7548](https://github.com/apache/arrow-rs/pull/7548) ([jspaezp](https://github.com/jspaezp)) +- Fix: Converted feature flight-sql-experimental to flight-sql [\#7546](https://github.com/apache/arrow-rs/pull/7546) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([kunalsinghdadhwal](https://github.com/kunalsinghdadhwal)) +- Fix CI on main due to logical conflict [\#7542](https://github.com/apache/arrow-rs/pull/7542) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Fix `filter_record_batch` panics with empty struct array [\#7539](https://github.com/apache/arrow-rs/pull/7539) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([thorfour](https://github.com/thorfour)) +- \[Variant\] Initial API for reading Variant data and metadata [\#7535](https://github.com/apache/arrow-rs/pull/7535) ([mkarbo](https://github.com/mkarbo)) +- fix: Panic in pretty\_format function when displaying DurationSecondsA… [\#7534](https://github.com/apache/arrow-rs/pull/7534) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([zhuqi-lucas](https://github.com/zhuqi-lucas)) +- Create version of LexicographicalComparator that compares fixed number of columns \(~ -15%\) [\#7530](https://github.com/apache/arrow-rs/pull/7530) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Make parquet-show-bloom-filter work with integer typed columns [\#7529](https://github.com/apache/arrow-rs/pull/7529) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([adamreeve](https://github.com/adamreeve)) +- chore\(deps\): update criterion requirement from 0.5 to 0.6 [\#7527](https://github.com/apache/arrow-rs/pull/7527) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([mbrobbel](https://github.com/mbrobbel)) +- Minor: Add a parquet row\_filter test, reduce some test boiler plate [\#7522](https://github.com/apache/arrow-rs/pull/7522) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Refactor `build_array_reader` into a struct [\#7521](https://github.com/apache/arrow-rs/pull/7521) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- arrow: add concat structs benchmark [\#7520](https://github.com/apache/arrow-rs/pull/7520) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([asubiotto](https://github.com/asubiotto)) +- arrow-select: add support for merging primitive dictionary values [\#7519](https://github.com/apache/arrow-rs/pull/7519) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([asubiotto](https://github.com/asubiotto)) +- arrow-select: add support for optimized concatenation of struct arrays [\#7517](https://github.com/apache/arrow-rs/pull/7517) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([asubiotto](https://github.com/asubiotto)) +- Fix Clippy in CI for Rust 1.87 release [\#7514](https://github.com/apache/arrow-rs/pull/7514) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([alamb](https://github.com/alamb)) +- Simplify `ParquetRecordBatchReader::next` control logic [\#7512](https://github.com/apache/arrow-rs/pull/7512) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Fix record API support for reading INT32 encoded TIME\_MILLIS [\#7511](https://github.com/apache/arrow-rs/pull/7511) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([njaremko](https://github.com/njaremko)) +- RecordBatchDecoder: skip RecordBatch validation when `skip_validation` property is enabled [\#7509](https://github.com/apache/arrow-rs/pull/7509) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([nilskch](https://github.com/nilskch)) +- Introduce `ReadPlan` to encapsulate the calculation of what parquet rows to decode [\#7502](https://github.com/apache/arrow-rs/pull/7502) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Update documentation for ParquetReader [\#7501](https://github.com/apache/arrow-rs/pull/7501) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Improve `Field` docs, add missing `Field::set_*` methods [\#7497](https://github.com/apache/arrow-rs/pull/7497) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Speed up arithmetic kernels, reduce `unsafe` usage [\#7493](https://github.com/apache/arrow-rs/pull/7493) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Prevent FlightSQL server panics for `do_put` when stream is empty or 1st stream element is an Err [\#7492](https://github.com/apache/arrow-rs/pull/7492) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] [[arrow-flight](https://github.com/apache/arrow-rs/labels/arrow-flight)] ([superserious-dev](https://github.com/superserious-dev)) +- arrow-ipc: add `StreamDecoder::schema` [\#7488](https://github.com/apache/arrow-rs/pull/7488) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([lidavidm](https://github.com/lidavidm)) +- arrow-select: Implement concat for `RunArray`s [\#7487](https://github.com/apache/arrow-rs/pull/7487) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([brancz](https://github.com/brancz)) +- \[Variant\] Add \(empty\) `parquet-variant` crate, update `parquet-testing` pin [\#7485](https://github.com/apache/arrow-rs/pull/7485) ([alamb](https://github.com/alamb)) +- Improve error messages if schema hint mismatches with parquet schema [\#7481](https://github.com/apache/arrow-rs/pull/7481) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([alamb](https://github.com/alamb)) +- Add `arrow_reader_clickbench` benchmark [\#7470](https://github.com/apache/arrow-rs/pull/7470) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Speedup `filter_bytes` ~-20-40%, `filter_native` low selectivity \(~-37%\) [\#7463](https://github.com/apache/arrow-rs/pull/7463) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Dandandan](https://github.com/Dandandan)) +- Update arrow\_reader\_row\_filter benchmark to reflect ClickBench distribution [\#7461](https://github.com/apache/arrow-rs/pull/7461) [[parquet](https://github.com/apache/arrow-rs/labels/parquet)] ([alamb](https://github.com/alamb)) +- Add Map support to arrow-avro [\#7451](https://github.com/apache/arrow-rs/pull/7451) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([jecsand838](https://github.com/jecsand838)) +- Support Utf8View for Avro [\#7434](https://github.com/apache/arrow-rs/pull/7434) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([kumarlokesh](https://github.com/kumarlokesh)) +- Add support for creating random Decimal128 and Decimal256 arrays [\#7427](https://github.com/apache/arrow-rs/pull/7427) [[arrow](https://github.com/apache/arrow-rs/labels/arrow)] ([Weijun-H](https://github.com/Weijun-H)) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 38236ee39125..07ed5e010c40 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -112,16 +112,16 @@ git submodule update --init This populates data in two git submodules: -- `../parquet-testing/data` (sourced from https://github.com/apache/parquet-testing.git) -- `../testing` (sourced from https://github.com/apache/arrow-testing) +- `./parquet-testing/data` (sourced from https://github.com/apache/parquet-testing.git) +- `./testing` (sourced from https://github.com/apache/arrow-testing) By default, `cargo test` will look for these directories at their standard location. The following environment variables can be used to override the location: ```bash # Optionally specify a different location for test data -export PARQUET_TEST_DATA=$(cd ../parquet-testing/data; pwd) -export ARROW_TEST_DATA=$(cd ../testing/data; pwd) +export PARQUET_TEST_DATA=$(cd ./parquet-testing/data; pwd) +export ARROW_TEST_DATA=$(cd ./testing/data; pwd) ``` From here on, this is a pure Rust project and `cargo` can be used to run tests, benchmarks, docs and examples as usual. diff --git a/Cargo.toml b/Cargo.toml index 75ba410f12a6..73c0f7058b44 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,11 +33,15 @@ members = [ "arrow-ipc", "arrow-json", "arrow-ord", + "arrow-pyarrow", "arrow-row", "arrow-schema", "arrow-select", "arrow-string", "parquet", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "parquet_derive", "parquet_derive_test", ] @@ -53,16 +57,17 @@ members = [ resolver = "2" exclude = [ + # arrow-pyarrow-testing is excluded because it requires a Python interpreter with the pyarrow package installed, + # which makes running `cargo test --all` fail if the appropriate Python environment is not set up. + "arrow-pyarrow-testing", # arrow-pyarrow-integration-testing is excluded because it requires different compilation flags, thereby # significantly changing how it is compiled within the workspace, causing the whole workspace to be compiled from # scratch this way, this is a stand-alone package that compiles independently of the others. "arrow-pyarrow-integration-testing", - # object_store is excluded because it follows a separate release cycle from the other arrow crates - "object_store" ] [workspace.package] -version = "54.0.0" +version = "55.2.0" homepage = "https://github.com/apache/arrow-rs" repository = "https://github.com/apache/arrow-rs" authors = ["Apache Arrow "] @@ -72,25 +77,42 @@ include = [ "benches/*.rs", "src/**/*.rs", "Cargo.toml", + "LICENSE.txt", + "NOTICE.txt", ] edition = "2021" -rust-version = "1.62" +rust-version = "1.84" [workspace.dependencies] -arrow = { version = "54.0.0", path = "./arrow", default-features = false } -arrow-arith = { version = "54.0.0", path = "./arrow-arith" } -arrow-array = { version = "54.0.0", path = "./arrow-array" } -arrow-buffer = { version = "54.0.0", path = "./arrow-buffer" } -arrow-cast = { version = "54.0.0", path = "./arrow-cast" } -arrow-csv = { version = "54.0.0", path = "./arrow-csv" } -arrow-data = { version = "54.0.0", path = "./arrow-data" } -arrow-ipc = { version = "54.0.0", path = "./arrow-ipc" } -arrow-json = { version = "54.0.0", path = "./arrow-json" } -arrow-ord = { version = "54.0.0", path = "./arrow-ord" } -arrow-row = { version = "54.0.0", path = "./arrow-row" } -arrow-schema = { version = "54.0.0", path = "./arrow-schema" } -arrow-select = { version = "54.0.0", path = "./arrow-select" } -arrow-string = { version = "54.0.0", path = "./arrow-string" } -parquet = { version = "54.0.0", path = "./parquet", default-features = false } +arrow = { version = "55.2.0", path = "./arrow", default-features = false } +arrow-arith = { version = "55.2.0", path = "./arrow-arith" } +arrow-array = { version = "55.2.0", path = "./arrow-array" } +arrow-buffer = { version = "55.2.0", path = "./arrow-buffer" } +arrow-cast = { version = "55.2.0", path = "./arrow-cast" } +arrow-csv = { version = "55.2.0", path = "./arrow-csv" } +arrow-data = { version = "55.2.0", path = "./arrow-data" } +arrow-ipc = { version = "55.2.0", path = "./arrow-ipc" } +arrow-json = { version = "55.2.0", path = "./arrow-json" } +arrow-ord = { version = "55.2.0", path = "./arrow-ord" } +arrow-pyarrow = { version = "55.2.0", path = "./arrow-pyarrow" } +arrow-row = { version = "55.2.0", path = "./arrow-row" } +arrow-schema = { version = "55.2.0", path = "./arrow-schema" } +arrow-select = { version = "55.2.0", path = "./arrow-select" } +arrow-string = { version = "55.2.0", path = "./arrow-string" } +parquet = { version = "55.2.0", path = "./parquet", default-features = false } -chrono = { version = "0.4.34", default-features = false, features = ["clock"] } +# These crates have not yet been released and thus do not use the workspace version +parquet-variant = { version = "0.1.0", path = "./parquet-variant" } +parquet-variant-json = { version = "0.1.0", path = "./parquet-variant-json" } +parquet-variant-compute = { version = "0.1.0", path = "./parquet-variant-json" } + +chrono = { version = "0.4.40", default-features = false, features = ["clock"] } + +simdutf8 = { version = "0.1.5", default-features = false } + +# release inherited profile keeping debug information and symbols +# for mem/cpu profiling +[profile.profiling] +inherits = "release" +debug = true +strip = false diff --git a/README.md b/README.md index ed42f630514b..7e7b3b6cf0d8 100644 --- a/README.md +++ b/README.md @@ -21,29 +21,29 @@ Welcome to the [Rust][rust] implementation of [Apache Arrow], the popular in-memory columnar format. -This repo contains the following main components: +This repository contains the following crates: | Crate | Description | Latest API Docs | README | | ------------------ | ---------------------------------------------------------------------------- | ------------------------------------------------ | --------------------------------- | | [`arrow`] | Core functionality (memory layout, arrays, low level computations) | [docs.rs](https://docs.rs/arrow/latest) | [(README)][arrow-readme] | | [`arrow-flight`] | Support for Arrow-Flight IPC protocol | [docs.rs](https://docs.rs/arrow-flight/latest) | [(README)][flight-readme] | -| [`object-store`] | Support for object store interactions (aws, azure, gcp, local, in-memory) | [docs.rs](https://docs.rs/object_store/latest) | [(README)][objectstore-readme] | | [`parquet`] | Support for Parquet columnar file format | [docs.rs](https://docs.rs/parquet/latest) | [(README)][parquet-readme] | | [`parquet_derive`] | A crate for deriving RecordWriter/RecordReader for arbitrary, simple structs | [docs.rs](https://docs.rs/parquet-derive/latest) | [(README)][parquet-derive-readme] | The current development version the API documentation in this repo can be found [here](https://arrow.apache.org/rust). +Note: previously the [`object_store`] crate was also part of this repository, +but it has been moved to the [arrow-rs-object-store repository] + [apache arrow]: https://arrow.apache.org/ [`arrow`]: https://crates.io/crates/arrow [`parquet`]: https://crates.io/crates/parquet [`parquet_derive`]: https://crates.io/crates/parquet-derive [`arrow-flight`]: https://crates.io/crates/arrow-flight -[`object-store`]: https://crates.io/crates/object-store +[arrow-rs-object-store repository]: https://github.com/apache/arrow-rs-object-store ## Release Versioning and Schedule -### `arrow` and `parquet` crates - The Arrow Rust project releases approximately monthly and follows [Semantic Versioning]. @@ -53,7 +53,7 @@ as the [`parquet`] and [`parquet-derive`] crates. This crate releases every month. We release new major versions (with potentially breaking API changes) at most once a quarter, and release incremental minor -versions in the intervening months. See [this ticket] for more details. +versions in the intervening months. See [ticket #5368] for more details. To keep our maintenance burden down, we do regularly scheduled releases (major and minor) from the `main` branch. How we handle PRs with breaking API changes @@ -63,32 +63,35 @@ is described in the [contributing] guide. Planned Release Schedule -| Approximate Date | Version | Notes | -| ---------------- | -------- | ------------------------------------------ | -| Nov 2024 | `53.3.0` | Minor, NO breaking API changes | -| Dec 2024 | `54.0.0` | Major, potentially breaking API changes | -| Jan 2025 | `53.4.0` | Minor, NO breaking API changes (`53` line) | -| Jan 2025 | `54.1.0` | Minor, NO breaking API changes | -| Feb 2025 | `54.2.0` | Minor, NO breaking API changes | -| Mar 2025 | `55.0.0` | Major, potentially breaking API changes | - -[this ticket]: https://github.com/apache/arrow-rs/issues/5368 +| Approximate Date | Version | Notes | +| ---------------- | ---------- | --------------------------------------- | +| July 2025 | [`56.0.0`] | Major, potentially breaking API changes | +| August 2025 | [`56.1.0`] | Minor, NO breaking API changes | +| September 2025 | [`56.2.0`] | Minor, NO breaking API changes | +| October 2025 | [`57.0.0`] | Major, potentially breaking API changes | + +[`56.0.0`]: https://github.com/apache/arrow-rs/issues/7395 +[`56.1.0`]: https://github.com/apache/arrow-rs/issues/7837 +[`56.2.0`]: https://github.com/apache/arrow-rs/issues/7836 +[`57.0.0`]: https://github.com/apache/arrow-rs/issues/7835 +[ticket #5368]: https://github.com/apache/arrow-rs/issues/5368 [semantic versioning]: https://semver.org/ -### `object_store` crate +### Rust Version Compatibility Policy -The [`object_store`] crate is released independently of the `arrow` and -`parquet` crates and follows [Semantic Versioning]. We aim to release new -versions approximately every 2 months. +arrow-rs and parquet are built and tested with stable Rust, and will keep a rolling MSRV (minimum supported Rust version) that can only be updated in major releases on a need by basis (e.g. project dependencies bump their MSRV or a particular Rust feature is useful for us etc.). The new MSRV if selected will be at least 6 months old. The minor releases are guaranteed to have the same MSRV. -[`object_store`]: https://crates.io/crates/object_store +Note: If a Rust hotfix is released for the current MSRV, the MSRV will be updated to the specific minor version that includes all applicable hotfixes preceding other policies. -Planned Release Schedule +### Guidelines for `panic` vs `Result` + +In general, use panics for bad states that are unreachable, unrecoverable or harmful. +For those caused by invalid user input, however, we prefer to report that invalidity +gracefully as an error result instead of panicking. In general, invalid input should result +in an `Error` as soon as possible. It _is_ ok for code paths after validation to assume +validation has already occurred and panic if not. See [ticket #6737] for more nuances. -| Approximate Date | Version | Notes | -| ---------------- | -------- | --------------------------------------- | -| Dec 2024 | `0.11.2` | Minor, NO breaking API changes | -| Feb 2025 | `0.12.0` | Major, potentially breaking API changes | +[ticket #6737]: https://github.com/apache/arrow-rs/issues/6737 ### Deprecation Guidelines @@ -121,20 +124,18 @@ maintainers. There are several related crates in different repositories -| Crate | Description | Documentation | -| ------------------------ | ------------------------------------------- | --------------------------------------- | -| [`datafusion`] | In-memory query engine with SQL support | [(README)][datafusion-readme] | -| [`ballista`] | Distributed query execution | [(README)][ballista-readme] | -| [`object_store_opendal`] | Use [`opendal`] as [`object_store`] backend | [(README)][object_store_opendal-readme] | -| [`parquet_opendal`] | Use [`opendal`] for [`parquet`] Arrow IO | [(README)][parquet_opendal-readme] | +| Crate | Description | Documentation | +| ------------------- | ------------------------------------------------------------ | ---------------------------------- | +| [`object_store`] | Object Storage (aws, azure, gcp, local, in-memory) interface | [(README)](object_store-readme) | +| [`datafusion`] | In-memory query engine with SQL support | [(README)][datafusion-readme] | +| [`ballista`] | Distributed query execution | [(README)][ballista-readme] | +| [`parquet_opendal`] | Use [`opendal`] for [`parquet`] Arrow IO | [(README)][parquet_opendal-readme] | [`datafusion`]: https://crates.io/crates/datafusion [`ballista`]: https://crates.io/crates/ballista -[`object_store_opendal`]: https://crates.io/crates/object_store_opendal -[`opendal`]: https://crates.io/crates/opendal -[object_store_opendal-readme]: https://github.com/apache/opendal/blob/main/integrations/object_store/README.md [`parquet_opendal`]: https://crates.io/crates/parquet_opendal [parquet_opendal-readme]: https://github.com/apache/opendal/blob/main/integrations/parquet/README.md +[object_store-readme]: https://github.com/apache/arrow-rs-object-store/blob/main/README.md Collectively, these crates support a wider array of functionality for analytic computations in Rust. @@ -163,18 +164,18 @@ a great place to meet other contributors and get guidance on where to contribute The Rust implementation uses [GitHub issues][issues] as the system of record for new features and bug fixes and this plays a critical role in the release process. -For design discussions we generally collaborate on Google documents and file a GitHub issue linking to the document. +For design discussions we generally use GitHub issues. There is more information in the [contributing] guide. [rust]: https://www.rust-lang.org/ +[`object_store`]: https://crates.io/crates/object-store [arrow-readme]: arrow/README.md [contributing]: CONTRIBUTING.md [parquet-readme]: parquet/README.md [flight-readme]: arrow-flight/README.md [datafusion-readme]: https://github.com/apache/datafusion/blob/main/README.md [ballista-readme]: https://github.com/apache/datafusion-ballista/blob/main/README.md -[objectstore-readme]: object_store/README.md [parquet-derive-readme]: parquet_derive/README.md [issues]: https://github.com/apache/arrow-rs/issues [discussions]: https://github.com/apache/arrow-rs/discussions diff --git a/arrow-arith/Cargo.toml b/arrow-arith/Cargo.toml index 66696df8aa04..a3fdafa823a2 100644 --- a/arrow-arith/Cargo.toml +++ b/arrow-arith/Cargo.toml @@ -30,9 +30,11 @@ rust-version = { workspace = true } [lib] name = "arrow_arith" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow-array = { workspace = true } arrow-buffer = { workspace = true } @@ -40,5 +42,3 @@ arrow-data = { workspace = true } arrow-schema = { workspace = true } chrono = { workspace = true } num = { version = "0.4", default-features = false, features = ["std"] } - -[dev-dependencies] diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs index ef0fddeb0b8e..9a19b5d8a1f1 100644 --- a/arrow-arith/src/aggregate.rs +++ b/arrow-arith/src/aggregate.rs @@ -513,6 +513,11 @@ pub fn max_binary_view(array: &BinaryViewArray) -> Option<&[u8]> { min_max_view_helper(array, Ordering::Greater) } +/// Returns the maximum value in the fixed size binary array, according to the natural order. +pub fn max_fixed_size_binary(array: &FixedSizeBinaryArray) -> Option<&[u8]> { + min_max_helper::<&[u8], _, _>(array, |a, b| *a < *b) +} + /// Returns the minimum value in the binary array, according to the natural order. pub fn min_binary(array: &GenericBinaryArray) -> Option<&[u8]> { min_max_helper::<&[u8], _, _>(array, |a, b| *a > *b) @@ -523,6 +528,11 @@ pub fn min_binary_view(array: &BinaryViewArray) -> Option<&[u8]> { min_max_view_helper(array, Ordering::Less) } +/// Returns the minimum value in the fixed size binary array, according to the natural order. +pub fn min_fixed_size_binary(array: &FixedSizeBinaryArray) -> Option<&[u8]> { + min_max_helper::<&[u8], _, _>(array, |a, b| *a > *b) +} + /// Returns the maximum value in the string array, according to the natural order. pub fn max_string(array: &GenericStringArray) -> Option<&str> { min_max_helper::<&str, _, _>(array, |a, b| *a < *b) @@ -1240,6 +1250,41 @@ mod tests { assert!(max(&a).unwrap().is_nan()); } + fn pad_inputs_and_test_fixed_size_binary( + input: Vec>, + expected_min: Option<&[u8]>, + expected_max: Option<&[u8]>, + ) { + fn pad_slice(slice: &[u8], len: usize) -> Vec { + let mut padded = vec![0; len]; + padded[..slice.len()].copy_from_slice(slice); + padded + } + + let max_len = input + .iter() + .filter_map(|x| x.as_ref().map(|b| b.len())) + .max() + .unwrap_or(0); + let padded_input = input + .iter() + .map(|x| x.as_ref().map(|b| pad_slice(b, max_len))); + let input_arr = + FixedSizeBinaryArray::try_from_sparse_iter_with_size(padded_input, max_len as i32) + .unwrap(); + let padded_expected_min = expected_min.map(|b| pad_slice(b, max_len)); + let padded_expected_max = expected_max.map(|b| pad_slice(b, max_len)); + + assert_eq!( + padded_expected_min.as_deref(), + min_fixed_size_binary(&input_arr) + ); + assert_eq!( + padded_expected_max.as_deref(), + max_fixed_size_binary(&input_arr) + ); + } + macro_rules! test_binary { ($NAME:ident, $ARRAY:expr, $EXPECTED_MIN:expr, $EXPECTED_MAX: expr) => { #[test] @@ -1255,6 +1300,8 @@ mod tests { let binary_view = BinaryViewArray::from($ARRAY); assert_eq!($EXPECTED_MIN, min_binary_view(&binary_view)); assert_eq!($EXPECTED_MAX, max_binary_view(&binary_view)); + + pad_inputs_and_test_fixed_size_binary($ARRAY, $EXPECTED_MIN, $EXPECTED_MAX); } }; } diff --git a/arrow-arith/src/arithmetic.rs b/arrow-arith/src/arithmetic.rs index febf5ceabdd9..768fd798c04c 100644 --- a/arrow-arith/src/arithmetic.rs +++ b/arrow-arith/src/arithmetic.rs @@ -43,8 +43,7 @@ fn get_fixed_point_info( if required_scale > product_scale { return Err(ArrowError::ComputeError(format!( - "Required scale {} is greater than product scale {}", - required_scale, product_scale + "Required scale {required_scale} is greater than product scale {product_scale}", ))); } @@ -122,7 +121,7 @@ pub fn multiply_fixed_point_checked( let mut mul = a.wrapping_mul(b); mul = divide_and_round::(mul, divisor); mul.to_i128().ok_or_else(|| { - ArrowError::ArithmeticOverflow(format!("Overflow happened on: {:?} * {:?}", a, b)) + ArrowError::ArithmeticOverflow(format!("Overflow happened on: {a:?} * {b:?}")) }) }) .and_then(|a| a.with_precision_and_scale(precision, required_scale)) diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs index 9b3272abb617..d1bf1abcb269 100644 --- a/arrow-arith/src/arity.rs +++ b/arrow-arith/src/arity.rs @@ -21,7 +21,7 @@ use arrow_array::builder::BufferBuilder; use arrow_array::*; use arrow_buffer::buffer::NullBuffer; use arrow_buffer::ArrowNativeType; -use arrow_buffer::{Buffer, MutableBuffer}; +use arrow_buffer::MutableBuffer; use arrow_data::ArrayData; use arrow_schema::ArrowError; @@ -124,13 +124,13 @@ where let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()); - let values = a.values().iter().zip(b.values()).map(|(l, r)| op(*l, *r)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size from a PrimitiveArray - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; + let values = a + .values() + .into_iter() + .zip(b.values()) + .map(|(l, r)| op(*l, *r)); + + let buffer: Vec<_> = values.collect(); Ok(PrimitiveArray::new(buffer.into(), nulls)) } diff --git a/arrow-arith/src/lib.rs b/arrow-arith/src/lib.rs index c8b6412e5efc..63640c51c3ce 100644 --- a/arrow-arith/src/lib.rs +++ b/arrow-arith/src/lib.rs @@ -17,6 +17,11 @@ //! Arrow arithmetic and aggregation kernels +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] pub mod aggregate; #[doc(hidden)] // Kernels to be removed in a future release diff --git a/arrow-arith/src/numeric.rs b/arrow-arith/src/numeric.rs index b6af40f7d7c2..198447b4db7b 100644 --- a/arrow-arith/src/numeric.rs +++ b/arrow-arith/src/numeric.rs @@ -70,8 +70,10 @@ pub fn div(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { /// Perform `lhs % rhs` /// -/// Overflow or division by zero will result in an error, with exception to +/// Division by zero will result in an error, with exception to /// floating point numbers, which instead follow the IEEE 754 rules +/// +/// `signed_integer::MIN % -1` will not result in an error but return 0 pub fn rem(lhs: &dyn Datum, rhs: &dyn Datum) -> Result { arithmetic_op(Op::Rem, lhs, rhs) } @@ -109,6 +111,20 @@ pub fn neg(array: &dyn Array) -> Result { Float16 => neg_wrapping!(Float16Type, array), Float32 => neg_wrapping!(Float32Type, array), Float64 => neg_wrapping!(Float64Type, array), + Decimal32(p, s) => { + let a = array + .as_primitive::() + .try_unary::<_, Decimal32Type, _>(|x| x.neg_checked())?; + + Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) + } + Decimal64(p, s) => { + let a = array + .as_primitive::() + .try_unary::<_, Decimal64Type, _>(|x| x.neg_checked())?; + + Ok(Arc::new(a.with_precision_and_scale(*p, *s)?)) + } Decimal128(p, s) => { let a = array .as_primitive::() @@ -232,6 +248,8 @@ fn arithmetic_op(op: Op, lhs: &dyn Datum, rhs: &dyn Datum) -> Result interval_op::(op, l, l_scalar, r, r_scalar), (Date32, _) => date_op::(op, l, l_scalar, r, r_scalar), (Date64, _) => date_op::(op, l, l_scalar, r, r_scalar), + (Decimal32(_, _), Decimal32(_, _)) => decimal_op::(op, l, l_scalar, r, r_scalar), + (Decimal64(_, _), Decimal64(_, _)) => decimal_op::(op, l, l_scalar, r, r_scalar), (Decimal128(_, _), Decimal128(_, _)) => decimal_op::(op, l, l_scalar, r, r_scalar), (Decimal256(_, _), Decimal256(_, _)) => decimal_op::(op, l, l_scalar, r, r_scalar), (l_t, r_t) => match (l_t, r_t) { @@ -313,7 +331,13 @@ fn integer_op( Op::MulWrapping => op!(l, l_s, r, r_s, l.mul_wrapping(r)), Op::Mul => try_op!(l, l_s, r, r_s, l.mul_checked(r)), Op::Div => try_op!(l, l_s, r, r_s, l.div_checked(r)), - Op::Rem => try_op!(l, l_s, r, r_s, l.mod_checked(r)), + Op::Rem => try_op!(l, l_s, r, r_s, { + if r.is_zero() { + Err(ArrowError::DivideByZero) + } else { + Ok(l.mod_wrapping(r)) + } + }), }; Ok(Arc::new(array)) } @@ -502,49 +526,116 @@ fn timestamp_op( } /// Arithmetic trait for date arrays -/// -/// Note: these should be fallible (#4456) trait DateOp: ArrowTemporalType { - fn add_year_month(timestamp: Self::Native, delta: i32) -> Self::Native; - fn add_day_time(timestamp: Self::Native, delta: IntervalDayTime) -> Self::Native; - fn add_month_day_nano(timestamp: Self::Native, delta: IntervalMonthDayNano) -> Self::Native; + fn add_year_month(timestamp: Self::Native, delta: i32) -> Result; + fn add_day_time( + timestamp: Self::Native, + delta: IntervalDayTime, + ) -> Result; + fn add_month_day_nano( + timestamp: Self::Native, + delta: IntervalMonthDayNano, + ) -> Result; + + fn sub_year_month(timestamp: Self::Native, delta: i32) -> Result; + fn sub_day_time( + timestamp: Self::Native, + delta: IntervalDayTime, + ) -> Result; + fn sub_month_day_nano( + timestamp: Self::Native, + delta: IntervalMonthDayNano, + ) -> Result; +} + +impl DateOp for Date32Type { + fn add_year_month(left: Self::Native, right: i32) -> Result { + // Date32Type functions don't have _opt variants and should be safe + Ok(Self::add_year_months(left, right)) + } + + fn add_day_time( + left: Self::Native, + right: IntervalDayTime, + ) -> Result { + Ok(Self::add_day_time(left, right)) + } + + fn add_month_day_nano( + left: Self::Native, + right: IntervalMonthDayNano, + ) -> Result { + Ok(Self::add_month_day_nano(left, right)) + } + + fn sub_year_month(left: Self::Native, right: i32) -> Result { + Ok(Self::subtract_year_months(left, right)) + } + + fn sub_day_time( + left: Self::Native, + right: IntervalDayTime, + ) -> Result { + Ok(Self::subtract_day_time(left, right)) + } - fn sub_year_month(timestamp: Self::Native, delta: i32) -> Self::Native; - fn sub_day_time(timestamp: Self::Native, delta: IntervalDayTime) -> Self::Native; - fn sub_month_day_nano(timestamp: Self::Native, delta: IntervalMonthDayNano) -> Self::Native; + fn sub_month_day_nano( + left: Self::Native, + right: IntervalMonthDayNano, + ) -> Result { + Ok(Self::subtract_month_day_nano(left, right)) + } } -macro_rules! date { - ($t:ty) => { - impl DateOp for $t { - fn add_year_month(left: Self::Native, right: i32) -> Self::Native { - Self::add_year_months(left, right) - } +impl DateOp for Date64Type { + fn add_year_month(left: Self::Native, right: i32) -> Result { + Self::add_year_months_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right} months",)) + }) + } - fn add_day_time(left: Self::Native, right: IntervalDayTime) -> Self::Native { - Self::add_day_time(left, right) - } + fn add_day_time( + left: Self::Native, + right: IntervalDayTime, + ) -> Result { + Self::add_day_time_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}")) + }) + } - fn add_month_day_nano(left: Self::Native, right: IntervalMonthDayNano) -> Self::Native { - Self::add_month_day_nano(left, right) - } + fn add_month_day_nano( + left: Self::Native, + right: IntervalMonthDayNano, + ) -> Result { + Self::add_month_day_nano_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} + {right:?}")) + }) + } - fn sub_year_month(left: Self::Native, right: i32) -> Self::Native { - Self::subtract_year_months(left, right) - } + fn sub_year_month(left: Self::Native, right: i32) -> Result { + Self::subtract_year_months_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right} months",)) + }) + } - fn sub_day_time(left: Self::Native, right: IntervalDayTime) -> Self::Native { - Self::subtract_day_time(left, right) - } + fn sub_day_time( + left: Self::Native, + right: IntervalDayTime, + ) -> Result { + Self::subtract_day_time_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}")) + }) + } - fn sub_month_day_nano(left: Self::Native, right: IntervalMonthDayNano) -> Self::Native { - Self::subtract_month_day_nano(left, right) - } - } - }; + fn sub_month_day_nano( + left: Self::Native, + right: IntervalMonthDayNano, + ) -> Result { + Self::subtract_month_day_nano_opt(left, right).ok_or_else(|| { + ArrowError::ComputeError(format!("Date arithmetic overflow: {left} - {right:?}")) + }) + } } -date!(Date32Type); -date!(Date64Type); /// Arithmetic trait for interval arrays trait IntervalOp: ArrowPrimitiveType { @@ -681,29 +772,29 @@ fn date_op( match (op, r_t) { (Op::Add | Op::AddWrapping, Interval(YearMonth)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::add_year_month(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_year_month(l, r))) } (Op::Sub | Op::SubWrapping, Interval(YearMonth)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::sub_year_month(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_year_month(l, r))) } (Op::Add | Op::AddWrapping, Interval(DayTime)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::add_day_time(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_day_time(l, r))) } (Op::Sub | Op::SubWrapping, Interval(DayTime)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::sub_day_time(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_day_time(l, r))) } (Op::Add | Op::AddWrapping, Interval(MonthDayNano)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::add_month_day_nano(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::add_month_day_nano(l, r))) } (Op::Sub | Op::SubWrapping, Interval(MonthDayNano)) => { let r = r.as_primitive::(); - Ok(op_ref!(T, l, l_s, r, r_s, T::sub_month_day_nano(l, r))) + Ok(try_op_ref!(T, l, l_s, r, r_s, T::sub_month_day_nano(l, r))) } _ => Err(ArrowError::InvalidArgumentError(format!( @@ -726,6 +817,8 @@ fn decimal_op( let r = r.as_primitive::(); let (p1, s1, p2, s2) = match (l.data_type(), r.data_type()) { + (DataType::Decimal32(p1, s1), DataType::Decimal32(p2, s2)) => (p1, s1, p2, s2), + (DataType::Decimal64(p1, s1), DataType::Decimal64(p2, s2)) => (p1, s1, p2, s2), (DataType::Decimal128(p1, s1), DataType::Decimal128(p2, s2)) => (p1, s1, p2, s2), (DataType::Decimal256(p1, s1), DataType::Decimal256(p2, s2)) => (p1, s1, p2, s2), _ => unreachable!(), @@ -914,6 +1007,28 @@ mod tests { "Arithmetic overflow: Overflow happened on: - -9223372036854775808" ); + let a = Decimal32Array::from(vec![1, 3, -44, 2, 4]) + .with_precision_and_scale(9, 6) + .unwrap(); + + let r = neg(&a).unwrap(); + assert_eq!(r.data_type(), a.data_type()); + assert_eq!( + r.as_primitive::().values(), + &[-1, -3, 44, -2, -4] + ); + + let a = Decimal64Array::from(vec![1, 3, -44, 2, 4]) + .with_precision_and_scale(9, 6) + .unwrap(); + + let r = neg(&a).unwrap(); + assert_eq!(r.data_type(), a.data_type()); + assert_eq!( + r.as_primitive::().values(), + &[-1, -3, 44, -2, -4] + ); + let a = Decimal128Array::from(vec![1, 3, -44, 2, 4]) .with_precision_and_scale(9, 6) .unwrap(); @@ -1042,6 +1157,11 @@ mod tests { "Arithmetic overflow: Overflow happened on: -32768 / -1" ); + let a = Int16Array::from(vec![i16::MIN]); + let b = Int16Array::from(vec![-1]); + let result = rem(&a, &b).unwrap(); + assert_eq!(result.as_ref(), &Int16Array::from(vec![0])); + let a = Int16Array::from(vec![21]); let b = Int16Array::from(vec![0]); let err = div(&a, &b).unwrap_err().to_string(); @@ -1520,4 +1640,536 @@ mod tests { "Arithmetic overflow: Overflow happened on: 9223372036854775807 - -1" ); } + + #[test] + fn test_date64_to_naive_date_opt_boundary_values() { + use arrow_array::types::Date64Type; + + // Date64Type::to_naive_date_opt has boundaries determined by NaiveDate's supported range. + // The valid date range is from January 1, -262143 to December 31, 262142 (Gregorian calendar). + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + let ms_per_day = 24 * 60 * 60 * 1000i64; + + // Define the boundary dates using NaiveDate::from_ymd_opt + let max_valid_date = NaiveDate::from_ymd_opt(262142, 12, 31).unwrap(); + let min_valid_date = NaiveDate::from_ymd_opt(-262143, 1, 1).unwrap(); + + // Calculate their millisecond values from epoch + let max_valid_millis = (max_valid_date - epoch).num_milliseconds(); + let min_valid_millis = (min_valid_date - epoch).num_milliseconds(); + + // Verify these match the expected boundaries in milliseconds + assert_eq!( + max_valid_millis, 8210266790400000i64, + "December 31, 262142 should be 8210266790400000 ms from epoch" + ); + assert_eq!( + min_valid_millis, -8334601228800000i64, + "January 1, -262143 should be -8334601228800000 ms from epoch" + ); + + // Test that the boundary dates work + assert!( + Date64Type::to_naive_date_opt(max_valid_millis).is_some(), + "December 31, 262142 should return Some" + ); + assert!( + Date64Type::to_naive_date_opt(min_valid_millis).is_some(), + "January 1, -262143 should return Some" + ); + + // Test that one day beyond the boundaries fails + assert!( + Date64Type::to_naive_date_opt(max_valid_millis + ms_per_day).is_none(), + "January 1, 262143 should return None" + ); + assert!( + Date64Type::to_naive_date_opt(min_valid_millis - ms_per_day).is_none(), + "December 31, -262144 should return None" + ); + + // Test some values well within the valid range + assert!( + Date64Type::to_naive_date_opt(0).is_some(), + "Epoch (1970-01-01) should return Some" + ); + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + assert!( + Date64Type::to_naive_date_opt(year_2000_millis).is_some(), + "Year 2000 should return Some" + ); + + // Test extreme values that definitely fail due to Duration constraints + assert!( + Date64Type::to_naive_date_opt(i64::MAX).is_none(), + "i64::MAX should return None" + ); + assert!( + Date64Type::to_naive_date_opt(i64::MIN).is_none(), + "i64::MIN should return None" + ); + } + + #[test] + fn test_date64_add_year_months_opt_boundary_values() { + use arrow_array::types::Date64Type; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + assert!( + Date64Type::add_year_months_opt(year_2000_millis, 120).is_some(), + "Adding 10 years to year 2000 should succeed" + ); + + // Test with moderate years that are within chrono's safe range + let large_year = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap(); + let large_year_millis = (large_year - epoch).num_milliseconds(); + assert!( + Date64Type::add_year_months_opt(large_year_millis, 12).is_some(), + "Adding 12 months to year 5000 should succeed" + ); + + let neg_year = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap(); + let neg_year_millis = (neg_year - epoch).num_milliseconds(); + assert!( + Date64Type::add_year_months_opt(neg_year_millis, -12).is_some(), + "Subtracting 12 months from year -5000 should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::add_year_months_opt(i64::MAX, 1).is_none(), + "Adding months to i64::MAX should fail" + ); + assert!( + Date64Type::add_year_months_opt(i64::MIN, -1).is_none(), + "Subtracting months from i64::MIN should fail" + ); + + // Test edge case: adding zero should always work for valid dates + assert!( + Date64Type::add_year_months_opt(year_2000_millis, 0).is_some(), + "Adding zero months should always succeed for valid dates" + ); + } + + #[test] + fn test_date64_add_day_time_opt_boundary_values() { + use arrow_array::types::Date64Type; + use arrow_buffer::IntervalDayTime; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test with a date far from the boundary but still testing the function + let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap(); + let near_max_millis = (near_max_date - epoch).num_milliseconds(); + + // Adding 30 days should succeed + let interval_30_days = IntervalDayTime::new(30, 0); + assert!( + Date64Type::add_day_time_opt(near_max_millis, interval_30_days).is_some(), + "Adding 30 days to large year should succeed" + ); + + // Adding a very large number of days should fail + let interval_large_days = IntervalDayTime::new(100000000, 0); + assert!( + Date64Type::add_day_time_opt(near_max_millis, interval_large_days).is_none(), + "Adding 100M days to large year should fail" + ); + + // Test with a date far from the boundary in the negative direction + let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap(); + let near_min_millis = (near_min_date - epoch).num_milliseconds(); + + // Subtracting 30 days should succeed + let interval_minus_30_days = IntervalDayTime::new(-30, 0); + assert!( + Date64Type::add_day_time_opt(near_min_millis, interval_minus_30_days).is_some(), + "Subtracting 30 days from large negative year should succeed" + ); + + // Subtracting a very large number of days should fail + let interval_minus_large_days = IntervalDayTime::new(-100000000, 0); + assert!( + Date64Type::add_day_time_opt(near_min_millis, interval_minus_large_days).is_none(), + "Subtracting 100M days from large negative year should fail" + ); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + let interval_1000_days = IntervalDayTime::new(1000, 12345); + assert!( + Date64Type::add_day_time_opt(year_2000_millis, interval_1000_days).is_some(), + "Adding 1000 days and time to year 2000 should succeed" + ); + + // Test with extreme input values that would cause overflow + let interval_one_day = IntervalDayTime::new(1, 0); + assert!( + Date64Type::add_day_time_opt(i64::MAX, interval_one_day).is_none(), + "Adding interval to i64::MAX should fail" + ); + assert!( + Date64Type::add_day_time_opt(i64::MIN, IntervalDayTime::new(-1, 0)).is_none(), + "Subtracting interval from i64::MIN should fail" + ); + + // Test with extreme interval values + let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX); + assert!( + Date64Type::add_day_time_opt(0, max_interval).is_none(), + "Adding extreme interval should fail" + ); + + let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN); + assert!( + Date64Type::add_day_time_opt(0, min_interval).is_none(), + "Adding extreme negative interval should fail" + ); + + // Test millisecond overflow within a day + let large_ms_interval = IntervalDayTime::new(0, i32::MAX); + assert!( + Date64Type::add_day_time_opt(year_2000_millis, large_ms_interval).is_some(), + "Adding large milliseconds within valid range should succeed" + ); + } + + #[test] + fn test_date64_add_month_day_nano_opt_boundary_values() { + use arrow_array::types::Date64Type; + use arrow_buffer::IntervalMonthDayNano; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test with a large year that is still within chrono's safe range + let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap(); + let near_max_millis = (near_max_date - epoch).num_milliseconds(); + + // Adding 1 month and 30 days should succeed + let interval_safe = IntervalMonthDayNano::new(1, 30, 0); + assert!( + Date64Type::add_month_day_nano_opt(near_max_millis, interval_safe).is_some(), + "Adding 1 month 30 days to large year should succeed" + ); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + + // Test edge case: adding zero should always work for valid dates + let zero_interval = IntervalMonthDayNano::new(0, 0, 0); + assert!( + Date64Type::add_month_day_nano_opt(year_2000_millis, zero_interval).is_some(), + "Adding zero interval should always succeed for valid dates" + ); + + // Test with a negative year that is still within chrono's safe range + let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap(); + let near_min_millis = (near_min_date - epoch).num_milliseconds(); + + // Subtracting 1 month and 30 days should succeed + let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0); + assert!( + Date64Type::add_month_day_nano_opt(near_min_millis, interval_safe_neg).is_some(), + "Subtracting 1 month 30 days from large negative year should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0)) + .is_none(), + "Adding interval to i64::MAX should fail" + ); + + let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000); + assert!( + Date64Type::add_month_day_nano_opt(year_2000_millis, interval_normal).is_some(), + "Adding 2 months, 10 days, and nanos to year 2000 should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::add_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(1, 0, 0)) + .is_none(), + "Adding interval to i64::MAX should fail" + ); + assert!( + Date64Type::add_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(-1, 0, 0)) + .is_none(), + "Subtracting interval from i64::MIN should fail" + ); + + // Test with invalid timestamp input (the _opt function should handle these gracefully) + + // Test nanosecond precision (should not affect boundary since it's < 1ms) + let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999); + assert!( + Date64Type::add_month_day_nano_opt(year_2000_millis, nano_interval).is_some(), + "Adding nanoseconds within valid range should succeed" + ); + + // Test large nanosecond values that convert to milliseconds + let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos + assert!( + Date64Type::add_month_day_nano_opt(year_2000_millis, large_nano_interval).is_some(), + "Adding 1 day worth of nanoseconds should succeed" + ); + } + + #[test] + fn test_date64_subtract_year_months_opt_boundary_values() { + use arrow_array::types::Date64Type; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test with a negative year that is still within chrono's safe range + let near_min_date = NaiveDate::from_ymd_opt(-5000, 12, 31).unwrap(); + let near_min_millis = (near_min_date - epoch).num_milliseconds(); + + // Subtracting 12 months should succeed + assert!( + Date64Type::subtract_year_months_opt(near_min_millis, 12).is_some(), + "Subtracting 12 months from year -5000 should succeed" + ); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + + // Test edge case: subtracting zero should always work for valid dates + assert!( + Date64Type::subtract_year_months_opt(year_2000_millis, 0).is_some(), + "Subtracting zero months should always succeed for valid dates" + ); + + // Test with a large year that is still within chrono's safe range + let near_max_date = NaiveDate::from_ymd_opt(5000, 1, 1).unwrap(); + let near_max_millis = (near_max_date - epoch).num_milliseconds(); + + // Adding 12 months (subtracting negative) should succeed + assert!( + Date64Type::subtract_year_months_opt(near_max_millis, -12).is_some(), + "Adding 12 months to year 5000 should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(), + "Adding months to i64::MAX should fail" + ); + + assert!( + Date64Type::subtract_year_months_opt(year_2000_millis, 12).is_some(), + "Subtracting 1 year from year 2000 should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::subtract_year_months_opt(i64::MAX, -1).is_none(), + "Adding months to i64::MAX should fail" + ); + assert!( + Date64Type::subtract_year_months_opt(i64::MIN, 1).is_none(), + "Subtracting months from i64::MIN should fail" + ); + + // Test edge case: subtracting zero should always work for valid dates + let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); + let valid_millis = (valid_date - epoch).num_milliseconds(); + assert!( + Date64Type::subtract_year_months_opt(valid_millis, 0).is_some(), + "Subtracting zero months should always succeed for valid dates" + ); + } + + #[test] + fn test_date64_subtract_day_time_opt_boundary_values() { + use arrow_array::types::Date64Type; + use arrow_buffer::IntervalDayTime; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test with a date far from the boundary in the negative direction + let near_min_date = NaiveDate::from_ymd_opt(-200000, 2, 1).unwrap(); + let near_min_millis = (near_min_date - epoch).num_milliseconds(); + + // Subtracting 30 days should succeed + let interval_30_days = IntervalDayTime::new(30, 0); + assert!( + Date64Type::subtract_day_time_opt(near_min_millis, interval_30_days).is_some(), + "Subtracting 30 days from large negative year should succeed" + ); + + // Subtracting a very large number of days should fail + let interval_large_days = IntervalDayTime::new(100000000, 0); + assert!( + Date64Type::subtract_day_time_opt(near_min_millis, interval_large_days).is_none(), + "Subtracting 100M days from large negative year should fail" + ); + + // Test with a date far from the boundary but still testing the function + let near_max_date = NaiveDate::from_ymd_opt(200000, 12, 1).unwrap(); + let near_max_millis = (near_max_date - epoch).num_milliseconds(); + + // Adding 30 days (subtracting negative) should succeed + let interval_minus_30_days = IntervalDayTime::new(-30, 0); + assert!( + Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_30_days).is_some(), + "Adding 30 days to large year should succeed" + ); + + // Adding a very large number of days should fail + let interval_minus_large_days = IntervalDayTime::new(-100000000, 0); + assert!( + Date64Type::subtract_day_time_opt(near_max_millis, interval_minus_large_days).is_none(), + "Adding 100M days to large year should fail" + ); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + let interval_1000_days = IntervalDayTime::new(1000, 12345); + assert!( + Date64Type::subtract_day_time_opt(year_2000_millis, interval_1000_days).is_some(), + "Subtracting 1000 days and time from year 2000 should succeed" + ); + + // Test with extreme input values that would cause overflow + let interval_one_day = IntervalDayTime::new(1, 0); + assert!( + Date64Type::subtract_day_time_opt(i64::MIN, interval_one_day).is_none(), + "Subtracting interval from i64::MIN should fail" + ); + assert!( + Date64Type::subtract_day_time_opt(i64::MAX, IntervalDayTime::new(-1, 0)).is_none(), + "Adding interval to i64::MAX should fail" + ); + + // Test with extreme interval values + let max_interval = IntervalDayTime::new(i32::MAX, i32::MAX); + assert!( + Date64Type::subtract_day_time_opt(0, max_interval).is_none(), + "Subtracting extreme interval should fail" + ); + + let min_interval = IntervalDayTime::new(i32::MIN, i32::MIN); + assert!( + Date64Type::subtract_day_time_opt(0, min_interval).is_none(), + "Subtracting extreme negative interval should fail" + ); + + // Test millisecond precision + let large_ms_interval = IntervalDayTime::new(0, i32::MAX); + assert!( + Date64Type::subtract_day_time_opt(year_2000_millis, large_ms_interval).is_some(), + "Subtracting large milliseconds within valid range should succeed" + ); + + // Test edge case: subtracting zero should always work for valid dates + let zero_interval = IntervalDayTime::new(0, 0); + let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); + let valid_millis = (valid_date - epoch).num_milliseconds(); + assert!( + Date64Type::subtract_day_time_opt(valid_millis, zero_interval).is_some(), + "Subtracting zero interval should always succeed for valid dates" + ); + } + + #[test] + fn test_date64_subtract_month_day_nano_opt_boundary_values() { + use arrow_array::types::Date64Type; + use arrow_buffer::IntervalMonthDayNano; + + let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + + // Test with a negative year that is still within chrono's safe range + let near_min_date = NaiveDate::from_ymd_opt(-5000, 2, 28).unwrap(); + let near_min_millis = (near_min_date - epoch).num_milliseconds(); + + // Subtracting 1 month and 30 days should succeed + let interval_safe = IntervalMonthDayNano::new(1, 30, 0); + assert!( + Date64Type::subtract_month_day_nano_opt(near_min_millis, interval_safe).is_some(), + "Subtracting 1 month 30 days from large negative year should succeed" + ); + + // Test normal case within valid range + let year_2000 = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap(); + let year_2000_millis = (year_2000 - epoch).num_milliseconds(); + + // Test edge case: subtracting zero should always work for valid dates + let zero_interval = IntervalMonthDayNano::new(0, 0, 0); + assert!( + Date64Type::subtract_month_day_nano_opt(year_2000_millis, zero_interval).is_some(), + "Subtracting zero interval should always succeed for valid dates" + ); + + // Test with a large year that is still within chrono's safe range + let near_max_date = NaiveDate::from_ymd_opt(5000, 11, 1).unwrap(); + let near_max_millis = (near_max_date - epoch).num_milliseconds(); + + // Adding 1 month and 30 days (subtracting negative) should succeed + let interval_safe_neg = IntervalMonthDayNano::new(-1, -30, 0); + assert!( + Date64Type::subtract_month_day_nano_opt(near_max_millis, interval_safe_neg).is_some(), + "Adding 1 month 30 days to large year should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0)) + .is_none(), + "Subtracting interval from i64::MIN should fail" + ); + + let interval_normal = IntervalMonthDayNano::new(2, 10, 123_456_789_000); + assert!( + Date64Type::subtract_month_day_nano_opt(year_2000_millis, interval_normal).is_some(), + "Subtracting 2 months, 10 days, and nanos from year 2000 should succeed" + ); + + // Test with extreme input values that would cause overflow + assert!( + Date64Type::subtract_month_day_nano_opt(i64::MIN, IntervalMonthDayNano::new(1, 0, 0)) + .is_none(), + "Subtracting interval from i64::MIN should fail" + ); + assert!( + Date64Type::subtract_month_day_nano_opt(i64::MAX, IntervalMonthDayNano::new(-1, 0, 0)) + .is_none(), + "Adding interval to i64::MAX should fail" + ); + + // Test nanosecond precision (should not affect boundary since it's < 1ms) + let nano_interval = IntervalMonthDayNano::new(0, 0, 999_999_999); + assert!( + Date64Type::subtract_month_day_nano_opt(year_2000_millis, nano_interval).is_some(), + "Subtracting nanoseconds within valid range should succeed" + ); + + // Test large nanosecond values that convert to milliseconds + let large_nano_interval = IntervalMonthDayNano::new(0, 0, 86_400_000_000_000); // 1 day in nanos + assert!( + Date64Type::subtract_month_day_nano_opt(year_2000_millis, large_nano_interval) + .is_some(), + "Subtracting 1 day worth of nanoseconds should succeed" + ); + + // Test edge case: subtracting zero should always work for valid dates + let zero_interval = IntervalMonthDayNano::new(0, 0, 0); + let valid_date = NaiveDate::from_ymd_opt(2020, 6, 15).unwrap(); + let valid_millis = (valid_date - epoch).num_milliseconds(); + assert!( + Date64Type::subtract_month_day_nano_opt(valid_millis, zero_interval).is_some(), + "Subtracting zero interval should always succeed for valid dates" + ); + } } diff --git a/arrow-arith/src/temporal.rs b/arrow-arith/src/temporal.rs index 3458669a6fd1..a9682742bbf0 100644 --- a/arrow-arith/src/temporal.rs +++ b/arrow-arith/src/temporal.rs @@ -31,7 +31,6 @@ use arrow_array::temporal_conversions::{ use arrow_array::timezone::Tz; use arrow_array::types::*; use arrow_array::*; -use arrow_buffer::ArrowNativeType; use arrow_schema::{ArrowError, DataType, IntervalUnit, TimeUnit}; /// Valid parts to extract from date/time/timestamp arrays. @@ -47,10 +46,14 @@ pub enum DatePart { Quarter, /// Calendar year Year, + /// ISO year, computed as per ISO 8601 + YearISO, /// Month in the year, in range `1..=12` Month, - /// ISO week of the year, in range `1..=53` + /// week of the year, in range `1..=53`, computed as per ISO 8601 Week, + /// ISO week of the year, in range `1..=53` + WeekISO, /// Day of the month, in range `1..=31` Day, /// Day of the week, in range `0..=6`, where Sunday is `0` @@ -75,7 +78,7 @@ pub enum DatePart { impl std::fmt::Display for DatePart { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{:?}", self) + write!(f, "{self:?}") } } @@ -91,8 +94,9 @@ where match part { DatePart::Quarter => |d| d.quarter() as i32, DatePart::Year => |d| d.year(), + DatePart::YearISO => |d| d.iso_week().year(), DatePart::Month => |d| d.month() as i32, - DatePart::Week => |d| d.iso_week().week() as i32, + DatePart::Week | DatePart::WeekISO => |d| d.iso_week().week() as i32, DatePart::Day => |d| d.day() as i32, DatePart::DayOfWeekSunday0 => |d| d.num_days_from_sunday(), DatePart::DayOfWeekMonday0 => |d| d.num_days_from_monday(), @@ -102,7 +106,7 @@ where DatePart::Second => |d| d.second() as i32, DatePart::Millisecond => |d| (d.nanosecond() / 1_000_000) as i32, DatePart::Microsecond => |d| (d.nanosecond() / 1_000) as i32, - DatePart::Nanosecond => |d| (d.nanosecond()) as i32, + DatePart::Nanosecond => |d| d.nanosecond() as i32, } } @@ -130,9 +134,14 @@ where /// let input: TimestampMicrosecondArray = /// vec![Some(1612025847000000), None, Some(1722015847000000)].into(); /// -/// let actual = date_part(&input, DatePart::Week).unwrap(); +/// let week = date_part(&input, DatePart::Week).unwrap(); +/// let week_iso = date_part(&input, DatePart::WeekISO).unwrap(); /// let expected: Int32Array = vec![Some(4), None, Some(30)].into(); -/// assert_eq!(actual.as_ref(), &expected); +/// assert_eq!(week.as_ref(), &expected); +/// assert_eq!(week_iso.as_ref(), &expected); +/// let year_iso = date_part(&input, DatePart::YearISO).unwrap(); +/// let expected: Int32Array = vec![Some(2021), None, Some(2024)].into(); +/// assert_eq!(year_iso.as_ref(), &expected); /// ``` pub fn date_part(array: &dyn Array, part: DatePart) -> Result { downcast_temporal_array!( @@ -187,16 +196,6 @@ pub fn date_part(array: &dyn Array, part: DatePart) -> Result( - array: &PrimitiveArray, - part: DatePart, -) -> Result { - let array = date_part(array, part)?; - Ok(array.as_primitive::().to_owned()) -} - /// Extract optional [`Tz`] from timestamp data types, returning error /// if called with a non-timestamp type. fn get_tz(dt: &DataType) -> Result, ArrowError> { @@ -430,6 +429,8 @@ impl ExtractDatePartExt for PrimitiveArray { DatePart::Quarter | DatePart::Week + | DatePart::WeekISO + | DatePart::YearISO | DatePart::Day | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 @@ -452,14 +453,20 @@ impl ExtractDatePartExt for PrimitiveArray { DatePart::Week => Ok(self.unary_opt(|d| Some(d.days / 7))), DatePart::Day => Ok(self.unary_opt(|d| Some(d.days))), DatePart::Hour => Ok(self.unary_opt(|d| Some(d.milliseconds / (60 * 60 * 1_000)))), - DatePart::Minute => Ok(self.unary_opt(|d| Some(d.milliseconds / (60 * 1_000)))), - DatePart::Second => Ok(self.unary_opt(|d| Some(d.milliseconds / 1_000))), - DatePart::Millisecond => Ok(self.unary_opt(|d| Some(d.milliseconds))), - DatePart::Microsecond => Ok(self.unary_opt(|d| d.milliseconds.checked_mul(1_000))), - DatePart::Nanosecond => Ok(self.unary_opt(|d| d.milliseconds.checked_mul(1_000_000))), + DatePart::Minute => Ok(self.unary_opt(|d| Some(d.milliseconds / (60 * 1_000) % 60))), + DatePart::Second => Ok(self.unary_opt(|d| Some(d.milliseconds / 1_000 % 60))), + DatePart::Millisecond => Ok(self.unary_opt(|d| Some(d.milliseconds % (60 * 1_000)))), + DatePart::Microsecond => { + Ok(self.unary_opt(|d| (d.milliseconds % (60 * 1_000)).checked_mul(1_000))) + } + DatePart::Nanosecond => { + Ok(self.unary_opt(|d| (d.milliseconds % (60 * 1_000)).checked_mul(1_000_000))) + } DatePart::Quarter | DatePart::Year + | DatePart::YearISO + | DatePart::WeekISO | DatePart::Month | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 @@ -474,27 +481,35 @@ impl ExtractDatePartExt for PrimitiveArray { fn date_part(&self, part: DatePart) -> Result { match part { DatePart::Year => Ok(self.unary_opt(|d: IntervalMonthDayNano| Some(d.months / 12))), - DatePart::Month => Ok(self.unary_opt(|d: IntervalMonthDayNano| Some(d.months))), + DatePart::Month => Ok(self.unary_opt(|d: IntervalMonthDayNano| Some(d.months % 12))), DatePart::Week => Ok(self.unary_opt(|d: IntervalMonthDayNano| Some(d.days / 7))), DatePart::Day => Ok(self.unary_opt(|d: IntervalMonthDayNano| Some(d.days))), DatePart::Hour => { Ok(self.unary_opt(|d| (d.nanoseconds / (60 * 60 * 1_000_000_000)).try_into().ok())) } DatePart::Minute => { - Ok(self.unary_opt(|d| (d.nanoseconds / (60 * 1_000_000_000)).try_into().ok())) + Ok(self.unary_opt(|d| (d.nanoseconds / (60 * 1_000_000_000) % 60).try_into().ok())) } DatePart::Second => { - Ok(self.unary_opt(|d| (d.nanoseconds / 1_000_000_000).try_into().ok())) - } - DatePart::Millisecond => { - Ok(self.unary_opt(|d| (d.nanoseconds / 1_000_000).try_into().ok())) + Ok(self.unary_opt(|d| ((d.nanoseconds / 1_000_000_000) % 60).try_into().ok())) } - DatePart::Microsecond => { - Ok(self.unary_opt(|d| (d.nanoseconds / 1_000).try_into().ok())) + DatePart::Millisecond => Ok(self.unary_opt(|d| { + (d.nanoseconds % (60 * 1_000_000_000) / 1_000_000) + .try_into() + .ok() + })), + DatePart::Microsecond => Ok(self.unary_opt(|d| { + (d.nanoseconds % (60 * 1_000_000_000) / 1_000) + .try_into() + .ok() + })), + DatePart::Nanosecond => { + Ok(self.unary_opt(|d| (d.nanoseconds % (60 * 1_000_000_000)).try_into().ok())) } - DatePart::Nanosecond => Ok(self.unary_opt(|d| d.nanoseconds.try_into().ok())), DatePart::Quarter + | DatePart::WeekISO + | DatePart::YearISO | DatePart::DayOfWeekSunday0 | DatePart::DayOfWeekMonday0 | DatePart::DayOfYear => { @@ -523,6 +538,8 @@ impl ExtractDatePartExt for PrimitiveArray { ), DatePart::Year + | DatePart::YearISO + | DatePart::WeekISO | DatePart::Quarter | DatePart::Month | DatePart::DayOfWeekSunday0 @@ -553,6 +570,8 @@ impl ExtractDatePartExt for PrimitiveArray { } DatePart::Year + | DatePart::YearISO + | DatePart::WeekISO | DatePart::Quarter | DatePart::Month | DatePart::DayOfWeekSunday0 @@ -583,6 +602,8 @@ impl ExtractDatePartExt for PrimitiveArray { } DatePart::Year + | DatePart::YearISO + | DatePart::WeekISO | DatePart::Quarter | DatePart::Month | DatePart::DayOfWeekSunday0 @@ -613,6 +634,8 @@ impl ExtractDatePartExt for PrimitiveArray { DatePart::Nanosecond => Ok(self.unary_opt(|d| d.try_into().ok())), DatePart::Year + | DatePart::YearISO + | DatePart::WeekISO | DatePart::Quarter | DatePart::Month | DatePart::DayOfWeekSunday0 @@ -634,12 +657,6 @@ pub(crate) use return_compute_error_with; // Internal trait, which is used for mapping values from DateLike structures trait ChronoDateExt { - /// Returns a value in range `1..=4` indicating the quarter this date falls into - fn quarter(&self) -> u32; - - /// Returns a value in range `0..=3` indicating the quarter (zero-based) this date falls into - fn quarter0(&self) -> u32; - /// Returns the day of week; Monday is encoded as `0`, Tuesday as `1`, etc. fn num_days_from_monday(&self) -> i32; @@ -648,14 +665,6 @@ trait ChronoDateExt { } impl ChronoDateExt for T { - fn quarter(&self) -> u32 { - self.quarter0() + 1 - } - - fn quarter0(&self) -> u32 { - self.month0() / 3 - } - fn num_days_from_monday(&self) -> i32 { self.weekday().num_days_from_monday() as i32 } @@ -665,300 +674,26 @@ impl ChronoDateExt for T { } } -/// Extracts the hours of a given array as an array of integers within -/// the range of [0, 23]. If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn hour_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Hour) -} - -/// Extracts the hours of a given temporal primitive array as an array of integers within -/// the range of [0, 23]. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn hour(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Hour) -} - -/// Extracts the years of a given temporal array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn year_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Year) -} - -/// Extracts the years of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn year(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Year) -} - -/// Extracts the quarter of a given temporal array as an array of integersa within -/// the range of [1, 4]. If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn quarter_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Quarter) -} - -/// Extracts the quarter of a given temporal primitive array as an array of integers within -/// the range of [1, 4]. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn quarter(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Quarter) -} - -/// Extracts the month of a given temporal array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn month_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Month) -} - -/// Extracts the month of a given temporal primitive array as an array of integers within -/// the range of [1, 12]. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn month(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Month) -} - -/// Extracts the day of week of a given temporal array as an array of -/// integers. -/// -/// Monday is encoded as `0`, Tuesday as `1`, etc. -/// -/// See also [`num_days_from_sunday`] which starts at Sunday. -/// -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn num_days_from_monday_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::DayOfWeekMonday0) -} - -/// Extracts the day of week of a given temporal primitive array as an array of -/// integers. -/// -/// Monday is encoded as `0`, Tuesday as `1`, etc. -/// -/// See also [`num_days_from_sunday`] which starts at Sunday. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn num_days_from_monday(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::DayOfWeekMonday0) -} - -/// Extracts the day of week of a given temporal array as an array of -/// integers, starting at Sunday. -/// -/// Sunday is encoded as `0`, Monday as `1`, etc. -/// -/// See also [`num_days_from_monday`] which starts at Monday. -/// -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn num_days_from_sunday_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::DayOfWeekSunday0) -} - -/// Extracts the day of week of a given temporal primitive array as an array of -/// integers, starting at Sunday. -/// -/// Sunday is encoded as `0`, Monday as `1`, etc. -/// -/// See also [`num_days_from_monday`] which starts at Monday. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn num_days_from_sunday(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::DayOfWeekSunday0) -} - -/// Extracts the day of a given temporal array as an array of integers. -/// -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn day_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Day) -} - -/// Extracts the day of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn day(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Day) -} - -/// Extracts the day of year of a given temporal array as an array of integers. -/// -/// The day of year that ranges from 1 to 366. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn doy_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::DayOfYear) -} - -/// Extracts the day of year of a given temporal primitive array as an array of integers. -/// -/// The day of year that ranges from 1 to 366 -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn doy(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - T::Native: ArrowNativeType, - i64: From, -{ - date_part_primitive(array, DatePart::DayOfYear) -} - -/// Extracts the minutes of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn minute(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Minute) -} - -/// Extracts the week of a given temporal array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn week_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Week) -} - -/// Extracts the week of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn week(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Week) -} - -/// Extracts the seconds of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn second(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Second) -} - -/// Extracts the nanoseconds of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn nanosecond(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Nanosecond) -} - -/// Extracts the nanoseconds of a given temporal primitive array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn nanosecond_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Nanosecond) -} - -/// Extracts the microseconds of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn microsecond(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Microsecond) -} - -/// Extracts the microseconds of a given temporal primitive array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn microsecond_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Microsecond) -} - -/// Extracts the milliseconds of a given temporal primitive array as an array of integers -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn millisecond(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - date_part_primitive(array, DatePart::Millisecond) -} - -/// Extracts the milliseconds of a given temporal primitive array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn millisecond_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Millisecond) -} - -/// Extracts the minutes of a given temporal array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn minute_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Minute) -} - -/// Extracts the seconds of a given temporal array as an array of integers. -/// If the given array isn't temporal primitive or dictionary array, -/// an `Err` will be returned. -#[deprecated(since = "51.0.0", note = "Use `date_part` instead")] -pub fn second_dyn(array: &dyn Array) -> Result { - date_part(array, DatePart::Second) -} - #[cfg(test)] -#[allow(deprecated)] mod tests { use super::*; + /// Used to integrate new [`date_part()`] method with deprecated shims such as + /// [`hour()`] and [`week()`]. + fn date_part_primitive( + array: &PrimitiveArray, + part: DatePart, + ) -> Result { + let array = date_part(array, part)?; + Ok(array.as_primitive::().to_owned()) + } + #[test] fn test_temporal_array_date64_hour() { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(4, b.value(2)); @@ -968,7 +703,7 @@ mod tests { fn test_temporal_array_date32_hour() { let a: PrimitiveArray = vec![Some(15147), None, Some(15148)].into(); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(0, b.value(2)); @@ -978,7 +713,7 @@ mod tests { fn test_temporal_array_time32_second_hour() { let a: PrimitiveArray = vec![37800, 86339].into(); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(10, b.value(0)); assert_eq!(23, b.value(1)); } @@ -987,7 +722,7 @@ mod tests { fn test_temporal_array_time64_micro_hour() { let a: PrimitiveArray = vec![37800000000, 86339000000].into(); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(10, b.value(0)); assert_eq!(23, b.value(1)); } @@ -996,7 +731,7 @@ mod tests { fn test_temporal_array_timestamp_micro_hour() { let a: TimestampMicrosecondArray = vec![37800000000, 86339000000].into(); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(10, b.value(0)); assert_eq!(23, b.value(1)); } @@ -1006,7 +741,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = year(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Year).unwrap(); assert_eq!(2018, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2019, b.value(2)); @@ -1016,7 +751,7 @@ mod tests { fn test_temporal_array_date32_year() { let a: PrimitiveArray = vec![Some(15147), None, Some(15448)].into(); - let b = year(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Year).unwrap(); assert_eq!(2011, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2012, b.value(2)); @@ -1029,7 +764,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1566275025000)].into(); - let b = quarter(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Quarter).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(3, b.value(2)); @@ -1039,7 +774,7 @@ mod tests { fn test_temporal_array_date32_quarter() { let a: PrimitiveArray = vec![Some(1), None, Some(300)].into(); - let b = quarter(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Quarter).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(4, b.value(2)); @@ -1049,10 +784,10 @@ mod tests { fn test_temporal_array_timestamp_quarter_with_timezone() { // 24 * 60 * 60 = 86400 let a = TimestampSecondArray::from(vec![86400 * 90]).with_timezone("+00:00".to_string()); - let b = quarter(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Quarter).unwrap(); assert_eq!(2, b.value(0)); let a = TimestampSecondArray::from(vec![86400 * 90]).with_timezone("-10:00".to_string()); - let b = quarter(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Quarter).unwrap(); assert_eq!(1, b.value(0)); } @@ -1063,7 +798,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = month(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Month).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2, b.value(2)); @@ -1073,7 +808,7 @@ mod tests { fn test_temporal_array_date32_month() { let a: PrimitiveArray = vec![Some(1), None, Some(31)].into(); - let b = month(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Month).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2, b.value(2)); @@ -1083,10 +818,10 @@ mod tests { fn test_temporal_array_timestamp_month_with_timezone() { // 24 * 60 * 60 = 86400 let a = TimestampSecondArray::from(vec![86400 * 31]).with_timezone("+00:00".to_string()); - let b = month(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Month).unwrap(); assert_eq!(2, b.value(0)); let a = TimestampSecondArray::from(vec![86400 * 31]).with_timezone("-10:00".to_string()); - let b = month(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Month).unwrap(); assert_eq!(1, b.value(0)); } @@ -1094,10 +829,10 @@ mod tests { fn test_temporal_array_timestamp_day_with_timezone() { // 24 * 60 * 60 = 86400 let a = TimestampSecondArray::from(vec![86400]).with_timezone("+00:00".to_string()); - let b = day(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Day).unwrap(); assert_eq!(2, b.value(0)); let a = TimestampSecondArray::from(vec![86400]).with_timezone("-10:00".to_string()); - let b = day(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Day).unwrap(); assert_eq!(1, b.value(0)); } @@ -1108,7 +843,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = num_days_from_monday(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::DayOfWeekMonday0).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2, b.value(2)); @@ -1127,7 +862,7 @@ mod tests { ] .into(); - let b = num_days_from_sunday(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::DayOfWeekSunday0).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(1, b.value(2)); @@ -1141,7 +876,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = day(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Day).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(20, b.value(2)); @@ -1151,7 +886,7 @@ mod tests { fn test_temporal_array_date32_day() { let a: PrimitiveArray = vec![Some(0), None, Some(31)].into(); - let b = day(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Day).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(1, b.value(2)); @@ -1170,7 +905,7 @@ mod tests { ] .into(); - let b = doy(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::DayOfYear).unwrap(); assert_eq!(1, b.value(0)); assert_eq!(1, b.value(1)); assert!(!b.is_valid(2)); @@ -1182,7 +917,7 @@ mod tests { let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); - let b = year(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Year).unwrap(); assert_eq!(2021, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2024, b.value(2)); @@ -1193,7 +928,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = minute(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Minute).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(23, b.value(2)); @@ -1204,7 +939,7 @@ mod tests { let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); - let b = minute(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Minute).unwrap(); assert_eq!(57, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(44, b.value(2)); @@ -1214,7 +949,7 @@ mod tests { fn test_temporal_array_date32_week() { let a: PrimitiveArray = vec![Some(0), None, Some(7)].into(); - let b = week(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Week).unwrap(); assert_eq!(1, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(2, b.value(2)); @@ -1232,7 +967,7 @@ mod tests { ] .into(); - let b = week(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Week).unwrap(); assert_eq!(9, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(1, b.value(2)); @@ -1246,7 +981,7 @@ mod tests { let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); - let b = week(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Week).unwrap(); assert_eq!(4, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(30, b.value(2)); @@ -1257,7 +992,7 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - let b = second(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Second).unwrap(); assert_eq!(0, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(45, b.value(2)); @@ -1268,7 +1003,7 @@ mod tests { let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); - let b = second(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Second).unwrap(); assert_eq!(27, b.value(0)); assert!(!b.is_valid(1)); assert_eq!(7, b.value(2)); @@ -1277,7 +1012,7 @@ mod tests { #[test] fn test_temporal_array_timestamp_second_with_timezone() { let a = TimestampSecondArray::from(vec![10, 20]).with_timezone("+00:00".to_string()); - let b = second(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Second).unwrap(); assert_eq!(10, b.value(0)); assert_eq!(20, b.value(1)); } @@ -1285,7 +1020,7 @@ mod tests { #[test] fn test_temporal_array_timestamp_minute_with_timezone() { let a = TimestampSecondArray::from(vec![0, 60]).with_timezone("+00:50".to_string()); - let b = minute(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Minute).unwrap(); assert_eq!(50, b.value(0)); assert_eq!(51, b.value(1)); } @@ -1293,42 +1028,46 @@ mod tests { #[test] fn test_temporal_array_timestamp_minute_with_negative_timezone() { let a = TimestampSecondArray::from(vec![60 * 55]).with_timezone("-00:50".to_string()); - let b = minute(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Minute).unwrap(); assert_eq!(5, b.value(0)); } #[test] fn test_temporal_array_timestamp_hour_with_timezone() { let a = TimestampSecondArray::from(vec![60 * 60 * 10]).with_timezone("+01:00".to_string()); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(11, b.value(0)); } #[test] fn test_temporal_array_timestamp_hour_with_timezone_without_colon() { let a = TimestampSecondArray::from(vec![60 * 60 * 10]).with_timezone("+0100".to_string()); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(11, b.value(0)); } #[test] fn test_temporal_array_timestamp_hour_with_timezone_without_minutes() { let a = TimestampSecondArray::from(vec![60 * 60 * 10]).with_timezone("+01".to_string()); - let b = hour(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Hour).unwrap(); assert_eq!(11, b.value(0)); } #[test] fn test_temporal_array_timestamp_hour_with_timezone_without_initial_sign() { let a = TimestampSecondArray::from(vec![60 * 60 * 10]).with_timezone("0100".to_string()); - let err = hour(&a).unwrap_err().to_string(); + let err = date_part_primitive(&a, DatePart::Hour) + .unwrap_err() + .to_string(); assert!(err.contains("Invalid timezone"), "{}", err); } #[test] fn test_temporal_array_timestamp_hour_with_timezone_with_only_colon() { let a = TimestampSecondArray::from(vec![60 * 60 * 10]).with_timezone("01:00".to_string()); - let err = hour(&a).unwrap_err().to_string(); + let err = date_part_primitive(&a, DatePart::Hour) + .unwrap_err() + .to_string(); assert!(err.contains("Invalid timezone"), "{}", err); } @@ -1338,7 +1077,7 @@ mod tests { // 1970-01-01T00:00:00 + 4 days -> 1970-01-05T00:00:00 Monday (week 2) // 1970-01-01T00:00:00 + 4 days - 1 second -> 1970-01-04T23:59:59 Sunday (week 1) let a = TimestampSecondArray::from(vec![0, 86400 * 4, 86400 * 4 - 1]); - let b = week(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Week).unwrap(); assert_eq!(1, b.value(0)); assert_eq!(2, b.value(1)); assert_eq!(1, b.value(2)); @@ -1351,7 +1090,7 @@ mod tests { // 1970-01-01T01:00:00+01:00 + 4 days - 1 second -> 1970-01-05T00:59:59+01:00 Monday (week 2) let a = TimestampSecondArray::from(vec![0, 86400 * 4, 86400 * 4 - 1]) .with_timezone("+01:00".to_string()); - let b = week(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Week).unwrap(); assert_eq!(1, b.value(0)); assert_eq!(2, b.value(1)); assert_eq!(2, b.value(2)); @@ -1369,7 +1108,7 @@ mod tests { let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2, 1]); let dict = DictionaryArray::try_new(keys.clone(), Arc::new(a)).unwrap(); - let b = hour_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Hour).unwrap(); let expected_dict = DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![11, 21, 7]))); @@ -1378,7 +1117,7 @@ mod tests { let b = date_part(&dict, DatePart::Minute).unwrap(); - let b_old = minute_dyn(&dict).unwrap(); + let b_old = date_part(&dict, DatePart::Minute).unwrap(); let expected_dict = DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![1, 2, 3]))); @@ -1388,7 +1127,7 @@ mod tests { let b = date_part(&dict, DatePart::Second).unwrap(); - let b_old = second_dyn(&dict).unwrap(); + let b_old = date_part(&dict, DatePart::Second).unwrap(); let expected_dict = DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![1, 2, 3]))); @@ -1411,7 +1150,7 @@ mod tests { let keys = Int8Array::from_iter_values([0_i8, 1, 1, 0]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = year_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Year).unwrap(); let expected_dict = DictionaryArray::new( keys, @@ -1430,13 +1169,13 @@ mod tests { let keys = Int8Array::from_iter_values([0_i8, 1, 1, 0]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = quarter_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Quarter).unwrap(); let expected = DictionaryArray::new(keys.clone(), Arc::new(Int32Array::from(vec![1, 3, 3, 1]))); assert_eq!(b.as_ref(), &expected); - let b = month_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Month).unwrap(); let expected = DictionaryArray::new(keys, Arc::new(Int32Array::from(vec![1, 8, 8, 1]))); assert_eq!(b.as_ref(), &expected); @@ -1451,31 +1190,31 @@ mod tests { let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1), Some(0), None]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = num_days_from_monday_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::DayOfWeekMonday0).unwrap(); let a = Int32Array::from(vec![Some(0), Some(2), Some(2), Some(0), None]); let expected = DictionaryArray::new(keys.clone(), Arc::new(a)); assert_eq!(b.as_ref(), &expected); - let b = num_days_from_sunday_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::DayOfWeekSunday0).unwrap(); let a = Int32Array::from(vec![Some(1), Some(3), Some(3), Some(1), None]); let expected = DictionaryArray::new(keys.clone(), Arc::new(a)); assert_eq!(b.as_ref(), &expected); - let b = day_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Day).unwrap(); let a = Int32Array::from(vec![Some(1), Some(20), Some(20), Some(1), None]); let expected = DictionaryArray::new(keys.clone(), Arc::new(a)); assert_eq!(b.as_ref(), &expected); - let b = doy_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::DayOfYear).unwrap(); let a = Int32Array::from(vec![Some(1), Some(51), Some(51), Some(1), None]); let expected = DictionaryArray::new(keys.clone(), Arc::new(a)); assert_eq!(b.as_ref(), &expected); - let b = week_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Week).unwrap(); let a = Int32Array::from(vec![Some(1), Some(8), Some(8), Some(1), None]); let expected = DictionaryArray::new(keys, Arc::new(a)); @@ -1492,13 +1231,13 @@ mod tests { let a: PrimitiveArray = vec![None, Some(1667328721453)].into(); - let b = nanosecond(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Nanosecond).unwrap(); assert!(!b.is_valid(0)); assert_eq!(453_000_000, b.value(1)); let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = nanosecond_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Nanosecond).unwrap(); let a = Int32Array::from(vec![None, Some(453_000_000)]); let expected_dict = DictionaryArray::new(keys, Arc::new(a)); @@ -1510,13 +1249,13 @@ mod tests { fn test_temporal_array_date64_microsecond() { let a: PrimitiveArray = vec![None, Some(1667328721453)].into(); - let b = microsecond(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Microsecond).unwrap(); assert!(!b.is_valid(0)); assert_eq!(453_000, b.value(1)); let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = microsecond_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Microsecond).unwrap(); let a = Int32Array::from(vec![None, Some(453_000)]); let expected_dict = DictionaryArray::new(keys, Arc::new(a)); @@ -1528,13 +1267,13 @@ mod tests { fn test_temporal_array_date64_millisecond() { let a: PrimitiveArray = vec![None, Some(1667328721453)].into(); - let b = millisecond(&a).unwrap(); + let b = date_part_primitive(&a, DatePart::Millisecond).unwrap(); assert!(!b.is_valid(0)); assert_eq!(453, b.value(1)); let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]); let dict = DictionaryArray::new(keys.clone(), Arc::new(a)); - let b = millisecond_dyn(&dict).unwrap(); + let b = date_part(&dict, DatePart::Millisecond).unwrap(); let a = Int32Array::from(vec![None, Some(453)]); let expected_dict = DictionaryArray::new(keys, Arc::new(a)); @@ -1754,9 +1493,14 @@ mod tests { fn test_interval_day_time_array() { let input: IntervalDayTimeArray = vec![ IntervalDayTime::ZERO, - IntervalDayTime::new(10, 42), - IntervalDayTime::new(10, 1042), - IntervalDayTime::new(10, MILLISECONDS_IN_DAY as i32 + 1), + IntervalDayTime::new(10, 42), // 10d, 42ms + IntervalDayTime::new(10, 1042), // 10d, 1s, 42ms + IntervalDayTime::new(10, MILLISECONDS_IN_DAY as i32 + 1), // 10d, 24h, 1ms + IntervalDayTime::new( + 6, + (MILLISECONDS * 60 * 60 * 4 + MILLISECONDS * 60 * 22 + MILLISECONDS * 11 + 3) + as i32, + ), // 6d, 4h, 22m, 11s, 3ms ] .into(); @@ -1767,6 +1511,7 @@ mod tests { assert_eq!(10, actual.value(1)); assert_eq!(10, actual.value(2)); assert_eq!(10, actual.value(3)); + assert_eq!(6, actual.value(4)); let actual = date_part(&input, DatePart::Week).unwrap(); let actual = actual.as_primitive::(); @@ -1774,51 +1519,57 @@ mod tests { assert_eq!(1, actual.value(1)); assert_eq!(1, actual.value(2)); assert_eq!(1, actual.value(3)); + assert_eq!(0, actual.value(4)); // Days doesn't affect time. - let actual = date_part(&input, DatePart::Nanosecond).unwrap(); + let actual = date_part(&input, DatePart::Hour).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); - assert_eq!(42_000_000, actual.value(1)); - assert_eq!(1_042_000_000, actual.value(2)); - // Overflow returns zero. + assert_eq!(0, actual.value(1)); + assert_eq!(0, actual.value(2)); + assert_eq!(24, actual.value(3)); + assert_eq!(4, actual.value(4)); + + let actual = date_part(&input, DatePart::Minute).unwrap(); + let actual = actual.as_primitive::(); + assert_eq!(0, actual.value(0)); + assert_eq!(0, actual.value(1)); + assert_eq!(0, actual.value(2)); assert_eq!(0, actual.value(3)); + assert_eq!(22, actual.value(4)); - let actual = date_part(&input, DatePart::Microsecond).unwrap(); + let actual = date_part(&input, DatePart::Second).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); - assert_eq!(42_000, actual.value(1)); - assert_eq!(1_042_000, actual.value(2)); - // Overflow returns zero. + assert_eq!(0, actual.value(1)); + assert_eq!(1, actual.value(2)); assert_eq!(0, actual.value(3)); + assert_eq!(11, actual.value(4)); let actual = date_part(&input, DatePart::Millisecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(42, actual.value(1)); assert_eq!(1042, actual.value(2)); - assert_eq!(MILLISECONDS_IN_DAY as i32 + 1, actual.value(3)); - - let actual = date_part(&input, DatePart::Second).unwrap(); - let actual = actual.as_primitive::(); - assert_eq!(0, actual.value(0)); - assert_eq!(0, actual.value(1)); - assert_eq!(1, actual.value(2)); - assert_eq!(24 * 60 * 60, actual.value(3)); + assert_eq!(1, actual.value(3)); + assert_eq!(11003, actual.value(4)); - let actual = date_part(&input, DatePart::Minute).unwrap(); + let actual = date_part(&input, DatePart::Microsecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); - assert_eq!(0, actual.value(1)); - assert_eq!(0, actual.value(2)); - assert_eq!(24 * 60, actual.value(3)); + assert_eq!(42_000, actual.value(1)); + assert_eq!(1_042_000, actual.value(2)); + assert_eq!(1_000, actual.value(3)); + assert_eq!(11_003_000, actual.value(4)); - let actual = date_part(&input, DatePart::Hour).unwrap(); + let actual = date_part(&input, DatePart::Nanosecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); - assert_eq!(0, actual.value(1)); - assert_eq!(0, actual.value(2)); - assert_eq!(24, actual.value(3)); + assert_eq!(42_000_000, actual.value(1)); + assert_eq!(1_042_000_000, actual.value(2)); + assert_eq!(1_000_000, actual.value(3)); + // Overflow returns zero. + assert_eq!(0, actual.value(4)); // Month and year are not valid (since days in month varies). assert!(date_part(&input, DatePart::Month).is_err()); @@ -1831,8 +1582,18 @@ mod tests { fn test_interval_month_day_nano_array() { let input: IntervalMonthDayNanoArray = vec![ IntervalMonthDayNano::ZERO, - IntervalMonthDayNano::new(5, 10, 42), - IntervalMonthDayNano::new(16, 35, MILLISECONDS_IN_DAY * 1_000_000 + 1), + IntervalMonthDayNano::new(5, 10, 42), // 5m, 1w, 3d, 42ns + IntervalMonthDayNano::new(16, 35, NANOSECONDS_IN_DAY + 1), // 1y, 4m, 5w, 24h, 1ns + IntervalMonthDayNano::new( + 0, + 0, + NANOSECONDS * 60 * 60 * 4 + + NANOSECONDS * 60 * 22 + + NANOSECONDS * 11 + + 1_000_000 * 33 + + 1_000 * 44 + + 5, + ), // 4hr, 22m, 11s, 33ms, 44us, 5ns ] .into(); @@ -1842,12 +1603,14 @@ mod tests { assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); assert_eq!(1, actual.value(2)); + assert_eq!(0, actual.value(3)); let actual = date_part(&input, DatePart::Month).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(5, actual.value(1)); - assert_eq!(16, actual.value(2)); + assert_eq!(4, actual.value(2)); + assert_eq!(0, actual.value(3)); // Week and day follow from day, but are not affected by months or nanos. let actual = date_part(&input, DatePart::Week).unwrap(); @@ -1855,12 +1618,14 @@ mod tests { assert_eq!(0, actual.value(0)); assert_eq!(1, actual.value(1)); assert_eq!(5, actual.value(2)); + assert_eq!(0, actual.value(3)); let actual = date_part(&input, DatePart::Day).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(10, actual.value(1)); assert_eq!(35, actual.value(2)); + assert_eq!(0, actual.value(3)); // Times follow from nanos, but are not affected by months or days. let actual = date_part(&input, DatePart::Hour).unwrap(); @@ -1868,38 +1633,43 @@ mod tests { assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); assert_eq!(24, actual.value(2)); + assert_eq!(4, actual.value(3)); let actual = date_part(&input, DatePart::Minute).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); - assert_eq!(24 * 60, actual.value(2)); + assert_eq!(0, actual.value(2)); + assert_eq!(22, actual.value(3)); let actual = date_part(&input, DatePart::Second).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); - assert_eq!(24 * 60 * 60, actual.value(2)); + assert_eq!(0, actual.value(2)); + assert_eq!(11, actual.value(3)); let actual = date_part(&input, DatePart::Millisecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); - assert_eq!(24 * 60 * 60 * 1_000, actual.value(2)); + assert_eq!(0, actual.value(2)); + assert_eq!(11_033, actual.value(3)); let actual = date_part(&input, DatePart::Microsecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(0, actual.value(1)); - // Overflow gives zero. assert_eq!(0, actual.value(2)); + assert_eq!(11_033_044, actual.value(3)); let actual = date_part(&input, DatePart::Nanosecond).unwrap(); let actual = actual.as_primitive::(); assert_eq!(0, actual.value(0)); assert_eq!(42, actual.value(1)); - // Overflow gives zero. - assert_eq!(0, actual.value(2)); + assert_eq!(1, actual.value(2)); + // Overflow returns zero. + assert_eq!(0, actual.value(3)); } #[test] @@ -2072,4 +1842,130 @@ mod tests { ensure_returns_error(&DurationMicrosecondArray::from(vec![0])); ensure_returns_error(&DurationNanosecondArray::from(vec![0])); } + + const TIMESTAMP_SECOND_1970_01_01: i64 = 0; + const TIMESTAMP_SECOND_2018_01_01: i64 = 1_514_764_800; + const TIMESTAMP_SECOND_2019_02_20: i64 = 1_550_636_625; + const SECONDS_IN_DAY: i64 = 24 * 60 * 60; + // In 2018 the ISO year and calendar year start on the same date— 2018-01-01 or 2018-W01-1 + #[test] + fn test_temporal_array_date64_week_iso() { + let a: PrimitiveArray = vec![ + Some(TIMESTAMP_SECOND_2018_01_01 * 1000), + Some(TIMESTAMP_SECOND_2019_02_20 * 1000), + ] + .into(); + + let b = date_part(&a, DatePart::WeekISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(1, actual.value(0)); + assert_eq!(8, actual.value(1)); + } + + #[test] + fn test_temporal_array_date64_year_iso() { + let a: PrimitiveArray = vec![ + Some(TIMESTAMP_SECOND_2018_01_01 * 1000), + Some(TIMESTAMP_SECOND_2019_02_20 * 1000), + ] + .into(); + + let b = date_part(&a, DatePart::YearISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(2018, actual.value(0)); + assert_eq!(2019, actual.value(1)); + } + + #[test] + fn test_temporal_array_timestamp_week_iso() { + let a = TimestampSecondArray::from(vec![ + TIMESTAMP_SECOND_1970_01_01, // 0 and is Thursday + SECONDS_IN_DAY * 4, // Monday of week 2 + SECONDS_IN_DAY * 4 - 1, // Sunday of week 1 + ]); + let b = date_part(&a, DatePart::WeekISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(1, actual.value(0)); + assert_eq!(2, actual.value(1)); + assert_eq!(1, actual.value(2)); + } + + #[test] + fn test_temporal_array_timestamp_year_iso() { + let a = TimestampSecondArray::from(vec![ + TIMESTAMP_SECOND_1970_01_01, + SECONDS_IN_DAY * 4, + SECONDS_IN_DAY * 4 - 1, + ]); + let b = date_part(&a, DatePart::YearISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(1970, actual.value(0)); + assert_eq!(1970, actual.value(1)); + assert_eq!(1970, actual.value(2)); + } + + const TIMESTAMP_SECOND_2015_12_28: i64 = 1_451_260_800; + const TIMESTAMP_SECOND_2016_01_03: i64 = 1_451_779_200; + // January 1st 2016 is a Friday, so 2015 week 53 runs from + // 2015-12-28 to 2016-01-03 inclusive, and + // 2016 week 1 runs from 2016-01-04 to 2016-01-10 inclusive. + #[test] + fn test_temporal_array_date64_week_iso_edge_cases() { + let a: PrimitiveArray = vec![ + Some(TIMESTAMP_SECOND_2015_12_28 * 1000), + Some(TIMESTAMP_SECOND_2016_01_03 * 1000), + Some((TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY) * 1000), + ] + .into(); + + let b = date_part(&a, DatePart::WeekISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(53, actual.value(0)); + assert_eq!(53, actual.value(1)); + assert_eq!(1, actual.value(2)); + } + + #[test] + fn test_temporal_array_date64_year_iso_edge_cases() { + let a: PrimitiveArray = vec![ + Some(TIMESTAMP_SECOND_2015_12_28 * 1000), + Some(TIMESTAMP_SECOND_2016_01_03 * 1000), + Some((TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY) * 1000), + ] + .into(); + + let b = date_part(&a, DatePart::YearISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(2015, actual.value(0)); + assert_eq!(2015, actual.value(1)); + assert_eq!(2016, actual.value(2)); + } + + #[test] + fn test_temporal_array_timestamp_week_iso_edge_cases() { + let a = TimestampSecondArray::from(vec![ + TIMESTAMP_SECOND_2015_12_28, + TIMESTAMP_SECOND_2016_01_03, + TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY, + ]); + let b = date_part(&a, DatePart::WeekISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(53, actual.value(0)); + assert_eq!(53, actual.value(1)); + assert_eq!(1, actual.value(2)); + } + + #[test] + fn test_temporal_array_timestamp_year_iso_edge_cases() { + let a = TimestampSecondArray::from(vec![ + TIMESTAMP_SECOND_2015_12_28, + TIMESTAMP_SECOND_2016_01_03, + TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY, + ]); + let b = date_part(&a, DatePart::YearISO).unwrap(); + let actual = b.as_primitive::(); + assert_eq!(2015, actual.value(0)); + assert_eq!(2015, actual.value(1)); + assert_eq!(2016, actual.value(2)); + } } diff --git a/arrow-array/Cargo.toml b/arrow-array/Cargo.toml index 6eae8e24677d..8ebe21c70772 100644 --- a/arrow-array/Cargo.toml +++ b/arrow-array/Cargo.toml @@ -30,10 +30,8 @@ rust-version = { workspace = true } [lib] name = "arrow_array" -path = "src/lib.rs" bench = false - [target.'cfg(target_arch = "wasm32")'.dependencies] ahash = { version = "0.8", default-features = false, features = ["compile-time-rng"] } @@ -50,22 +48,23 @@ num = { version = "0.4.1", default-features = false, features = ["std"] } half = { version = "2.1", default-features = false, features = ["num-traits"] } hashbrown = { version = "0.15.1", default-features = false } +[package.metadata.docs.rs] +all-features = true + [features] ffi = ["arrow-schema/ffi", "arrow-data/ffi"] force_validate = [] [dev-dependencies] -rand = { version = "0.8", default-features = false, features = ["std", "std_rng"] } +rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } criterion = { version = "0.5", default-features = false } -[build-dependencies] - [[bench]] name = "occupancy" harness = false [[bench]] -name = "gc_view_types" +name = "view_types" harness = false [[bench]] diff --git a/arrow-array/benches/decimal_overflow.rs b/arrow-array/benches/decimal_overflow.rs index 8f22b4b47c31..f9438b76de07 100644 --- a/arrow-array/benches/decimal_overflow.rs +++ b/arrow-array/benches/decimal_overflow.rs @@ -18,6 +18,7 @@ use arrow_array::builder::{Decimal128Builder, Decimal256Builder}; use arrow_buffer::i256; use criterion::*; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let len = 8192; @@ -36,16 +37,16 @@ fn criterion_benchmark(c: &mut Criterion) { let array_256 = builder_256.finish(); c.bench_function("validate_decimal_precision_128", |b| { - b.iter(|| black_box(array_128.validate_decimal_precision(8))); + b.iter(|| hint::black_box(array_128.validate_decimal_precision(8))); }); c.bench_function("null_if_overflow_precision_128", |b| { - b.iter(|| black_box(array_128.null_if_overflow_precision(8))); + b.iter(|| hint::black_box(array_128.null_if_overflow_precision(8))); }); c.bench_function("validate_decimal_precision_256", |b| { - b.iter(|| black_box(array_256.validate_decimal_precision(8))); + b.iter(|| hint::black_box(array_256.validate_decimal_precision(8))); }); c.bench_function("null_if_overflow_precision_256", |b| { - b.iter(|| black_box(array_256.null_if_overflow_precision(8))); + b.iter(|| hint::black_box(array_256.null_if_overflow_precision(8))); }); } diff --git a/arrow-array/benches/fixed_size_list_array.rs b/arrow-array/benches/fixed_size_list_array.rs index 5270a4a5def3..2bdb0c252b8a 100644 --- a/arrow-array/benches/fixed_size_list_array.rs +++ b/arrow-array/benches/fixed_size_list_array.rs @@ -18,13 +18,13 @@ use arrow_array::{Array, FixedSizeListArray, Int32Array}; use arrow_schema::Field; use criterion::*; -use rand::{thread_rng, Rng}; -use std::sync::Arc; +use rand::{rng, Rng}; +use std::{hint, sync::Arc}; fn gen_fsl(len: usize, value_len: usize) -> FixedSizeListArray { - let mut rng = thread_rng(); + let mut rng = rng(); let values = Arc::new(Int32Array::from( - (0..len).map(|_| rng.gen::()).collect::>(), + (0..len).map(|_| rng.random::()).collect::>(), )); let field = Arc::new(Field::new_list_field(values.data_type().clone(), true)); FixedSizeListArray::new(field, value_len as i32, values, None) @@ -39,7 +39,7 @@ fn criterion_benchmark(c: &mut Criterion) { |b| { b.iter(|| { for i in 0..len / value_len { - black_box(fsl.value(i)); + hint::black_box(fsl.value(i)); } }); }, diff --git a/arrow-array/benches/occupancy.rs b/arrow-array/benches/occupancy.rs index ed4b94351c28..283020364199 100644 --- a/arrow-array/benches/occupancy.rs +++ b/arrow-array/benches/occupancy.rs @@ -19,8 +19,8 @@ use arrow_array::types::Int32Type; use arrow_array::{DictionaryArray, Int32Array}; use arrow_buffer::NullBuffer; use criterion::*; -use rand::{thread_rng, Rng}; -use std::sync::Arc; +use rand::{rng, Rng}; +use std::{hint, sync::Arc}; fn gen_dict( len: usize, @@ -28,11 +28,11 @@ fn gen_dict( occupancy: f64, null_percent: f64, ) -> DictionaryArray { - let mut rng = thread_rng(); + let mut rng = rng(); let values = Int32Array::from(vec![0; values_len]); let max_key = (values_len as f64 * occupancy) as i32; - let keys = (0..len).map(|_| rng.gen_range(0..max_key)).collect(); - let nulls = (0..len).map(|_| !rng.gen_bool(null_percent)).collect(); + let keys = (0..len).map(|_| rng.random_range(0..max_key)).collect(); + let nulls = (0..len).map(|_| !rng.random_bool(null_percent)).collect(); let keys = Int32Array::new(keys, Some(NullBuffer::new(nulls))); DictionaryArray::new(keys, Arc::new(values)) @@ -45,7 +45,7 @@ fn criterion_benchmark(c: &mut Criterion) { let dict = gen_dict(1024, values, occupancy, null_percent); c.bench_function(&format!("occupancy(values: {values}, occupancy: {occupancy}, null_percent: {null_percent})"), |b| { b.iter(|| { - black_box(&dict).occupancy() + hint::black_box(&dict).occupancy() }); }); } diff --git a/arrow-array/benches/union_array.rs b/arrow-array/benches/union_array.rs index c5b2ec0f7752..d63eb9e43419 100644 --- a/arrow-array/benches/union_array.rs +++ b/arrow-array/benches/union_array.rs @@ -15,26 +15,23 @@ // specific language governing permissions and limitations // under the License. -use std::{ - iter::{repeat, repeat_with}, - sync::Arc, -}; +use std::{hint, iter::repeat_with, sync::Arc}; use arrow_array::{Array, ArrayRef, Int32Array, UnionArray}; use arrow_buffer::{NullBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, UnionFields}; use criterion::*; -use rand::{thread_rng, Rng}; +use rand::{rng, Rng}; fn array_with_nulls() -> ArrayRef { - let mut rng = thread_rng(); + let mut rng = rng(); - let values = ScalarBuffer::from_iter(repeat_with(|| rng.gen()).take(4096)); + let values = ScalarBuffer::from_iter(repeat_with(|| rng.random()).take(4096)); // nulls with at least one null and one valid let nulls: NullBuffer = [true, false] .into_iter() - .chain(repeat_with(|| rng.gen())) + .chain(repeat_with(|| rng.random())) .take(4096) .collect(); @@ -42,9 +39,9 @@ fn array_with_nulls() -> ArrayRef { } fn array_without_nulls() -> ArrayRef { - let mut rng = thread_rng(); + let mut rng = rng(); - let values = ScalarBuffer::from_iter(repeat_with(|| rng.gen()).take(4096)); + let values = ScalarBuffer::from_iter(repeat_with(|| rng.random()).take(4096)); Arc::new(Int32Array::new(values.clone(), None)) } @@ -66,14 +63,13 @@ fn criterion_benchmark(c: &mut Criterion) { fields, type_ids.cycle().take(4096).collect(), None, - repeat(array_with_nulls()) - .take(with_nulls as usize) - .chain(repeat(array_without_nulls()).take(without_nulls as usize)) + std::iter::repeat_n(array_with_nulls(), with_nulls as usize) + .chain(std::iter::repeat_n(array_without_nulls(), without_nulls as usize)) .collect(), ) .unwrap(); - b.iter(|| black_box(array.logical_nulls())) + b.iter(|| hint::black_box(array.logical_nulls())) }, ); } diff --git a/arrow-array/benches/view_types.rs b/arrow-array/benches/view_types.rs new file mode 100644 index 000000000000..986d4c65c1b1 --- /dev/null +++ b/arrow-array/benches/view_types.rs @@ -0,0 +1,111 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::StringViewArray; +use criterion::*; + +fn gen_view_array(size: usize) -> StringViewArray { + StringViewArray::from_iter((0..size).map(|v| match v % 3 { + 0 => Some("small"), + 1 => Some("larger than 12 bytes array"), + 2 => None, + _ => unreachable!("unreachable"), + })) +} + +fn gen_view_array_without_nulls(size: usize) -> StringViewArray { + StringViewArray::from_iter((0..size).map(|v| { + let s = match v % 3 { + 0 => "small".to_string(), // < 12 bytes + 1 => "larger than 12 bytes array".to_string(), // >12 bytes + 2 => "x".repeat(300), // 300 bytes (>256) + _ => unreachable!(), + }; + Some(s) + })) +} + +fn criterion_benchmark(c: &mut Criterion) { + let array = gen_view_array(100_000); + + c.bench_function("view types slice", |b| { + b.iter(|| { + black_box(array.slice(0, 100_000 / 2)); + }); + }); + + c.bench_function("gc view types all[100000]", |b| { + b.iter(|| { + black_box(array.gc()); + }); + }); + + let sliced = array.slice(0, 100_000 / 2); + c.bench_function("gc view types slice half[100000]", |b| { + b.iter(|| { + black_box(sliced.gc()); + }); + }); + + let array = gen_view_array_without_nulls(100_000); + + c.bench_function("gc view types all without nulls[100000]", |b| { + b.iter(|| { + black_box(array.gc()); + }); + }); + + let sliced = array.slice(0, 100_000 / 2); + c.bench_function("gc view types slice half without nulls[100000]", |b| { + b.iter(|| { + black_box(sliced.gc()); + }); + }); + + let array = gen_view_array(8000); + + c.bench_function("gc view types all[8000]", |b| { + b.iter(|| { + black_box(array.gc()); + }); + }); + + let sliced = array.slice(0, 8000 / 2); + c.bench_function("gc view types slice half[8000]", |b| { + b.iter(|| { + black_box(sliced.gc()); + }); + }); + + let array = gen_view_array_without_nulls(8000); + + c.bench_function("gc view types all without nulls[8000]", |b| { + b.iter(|| { + black_box(array.gc()); + }); + }); + + let sliced = array.slice(0, 8000 / 2); + c.bench_function("gc view types slice half without nulls[8000]", |b| { + b.iter(|| { + black_box(sliced.gc()); + }); + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/arrow-array/src/arithmetic.rs b/arrow-array/src/arithmetic.rs index fb9c868fb6c0..031864cb0809 100644 --- a/arrow-array/src/arithmetic.rs +++ b/arrow-array/src/arithmetic.rs @@ -418,21 +418,59 @@ native_type_float_op!( f32, 0., 1., - unsafe { std::mem::transmute(-1_i32) }, - unsafe { std::mem::transmute(i32::MAX) } + unsafe { + // Need to allow in clippy because + // current MSRV (Minimum Supported Rust Version) is `1.84.0` but this item is stable since `1.87.0` + #[allow(unnecessary_transmutes)] + std::mem::transmute(-1_i32) + }, + unsafe { + // Need to allow in clippy because + // current MSRV (Minimum Supported Rust Version) is `1.84.0` but this item is stable since `1.87.0` + #[allow(unnecessary_transmutes)] + std::mem::transmute(i32::MAX) + } ); native_type_float_op!( f64, 0., 1., - unsafe { std::mem::transmute(-1_i64) }, - unsafe { std::mem::transmute(i64::MAX) } + unsafe { + // Need to allow in clippy because + // current MSRV (Minimum Supported Rust Version) is `1.84.0` but this item is stable since `1.87.0` + #[allow(unnecessary_transmutes)] + std::mem::transmute(-1_i64) + }, + unsafe { + // Need to allow in clippy because + // current MSRV (Minimum Supported Rust Version) is `1.84.0` but this item is stable since `1.87.0` + #[allow(unnecessary_transmutes)] + std::mem::transmute(i64::MAX) + } ); #[cfg(test)] mod tests { use super::*; + macro_rules! assert_approx_eq { + ( $x: expr, $y: expr ) => {{ + assert_approx_eq!($x, $y, 1.0e-4) + }}; + ( $x: expr, $y: expr, $tol: expr ) => {{ + let x_val = $x; + let y_val = $y; + let diff = f64::from((x_val - y_val).abs()); + assert!( + diff <= $tol, + "{} != {} (with tolerance = {})", + x_val, + y_val, + $tol + ); + }}; + } + #[test] fn test_native_type_is_zero() { assert!(0_i8.is_zero()); @@ -803,9 +841,9 @@ mod tests { assert_eq!(8_u16.pow_wrapping(2_u32), 64_u16); assert_eq!(8_u32.pow_wrapping(2_u32), 64_u32); assert_eq!(8_u64.pow_wrapping(2_u32), 64_u64); - assert_eq!(f16::from_f32(8.0).pow_wrapping(2_u32), f16::from_f32(64.0)); - assert_eq!(8.0_f32.pow_wrapping(2_u32), 64_f32); - assert_eq!(8.0_f64.pow_wrapping(2_u32), 64_f64); + assert_approx_eq!(f16::from_f32(8.0).pow_wrapping(2_u32), f16::from_f32(64.0)); + assert_approx_eq!(8.0_f32.pow_wrapping(2_u32), 64_f32); + assert_approx_eq!(8.0_f64.pow_wrapping(2_u32), 64_f64); // pow_checked assert_eq!(8_i8.pow_checked(2_u32).unwrap(), 64_i8); @@ -821,12 +859,12 @@ mod tests { assert_eq!(8_u16.pow_checked(2_u32).unwrap(), 64_u16); assert_eq!(8_u32.pow_checked(2_u32).unwrap(), 64_u32); assert_eq!(8_u64.pow_checked(2_u32).unwrap(), 64_u64); - assert_eq!( + assert_approx_eq!( f16::from_f32(8.0).pow_checked(2_u32).unwrap(), f16::from_f32(64.0) ); - assert_eq!(8.0_f32.pow_checked(2_u32).unwrap(), 64_f32); - assert_eq!(8.0_f64.pow_checked(2_u32).unwrap(), 64_f64); + assert_approx_eq!(8.0_f32.pow_checked(2_u32).unwrap(), 64_f32); + assert_approx_eq!(8.0_f64.pow_checked(2_u32).unwrap(), 64_f64); } #[test] diff --git a/arrow-array/src/array/binary_array.rs b/arrow-array/src/array/binary_array.rs index 0e8a7a7cb618..8e2158416f49 100644 --- a/arrow-array/src/array/binary_array.rs +++ b/arrow-array/src/array/binary_array.rs @@ -20,7 +20,7 @@ use crate::{Array, GenericByteArray, GenericListArray, GenericStringArray, Offse use arrow_data::ArrayData; use arrow_schema::DataType; -/// A [`GenericBinaryArray`] for storing `[u8]` +/// A [`GenericByteArray`] for storing `[u8]` pub type GenericBinaryArray = GenericByteArray>; impl GenericBinaryArray { diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs index 9c2d4af8c454..fcebf5a0f718 100644 --- a/arrow-array/src/array/boolean_array.rs +++ b/arrow-array/src/array/boolean_array.rs @@ -479,7 +479,7 @@ impl From for BooleanArray { mod tests { use super::*; use arrow_buffer::Buffer; - use rand::{thread_rng, Rng}; + use rand::{rng, Rng}; #[test] fn test_boolean_fmt_debug() { @@ -667,11 +667,11 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] // Takes too long fn test_true_false_count() { - let mut rng = thread_rng(); + let mut rng = rng(); for _ in 0..10 { // No nulls - let d: Vec<_> = (0..2000).map(|_| rng.gen_bool(0.5)).collect(); + let d: Vec<_> = (0..2000).map(|_| rng.random_bool(0.5)).collect(); let b = BooleanArray::from(d.clone()); let expected_true = d.iter().filter(|x| **x).count(); @@ -680,7 +680,7 @@ mod tests { // With nulls let d: Vec<_> = (0..2000) - .map(|_| rng.gen_bool(0.5).then(|| rng.gen_bool(0.5))) + .map(|_| rng.random_bool(0.5).then(|| rng.random_bool(0.5))) .collect(); let b = BooleanArray::from(d.clone()); diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs index f2b22507081d..192c9654b055 100644 --- a/arrow-array/src/array/byte_array.rs +++ b/arrow-array/src/array/byte_array.rs @@ -164,6 +164,9 @@ impl GenericByteArray { values: Buffer, nulls: Option, ) -> Self { + if cfg!(feature = "force_validate") { + return Self::new(offsets, values, nulls); + } Self { data_type: T::DATA_TYPE, value_offsets: offsets, @@ -289,8 +292,10 @@ impl GenericByteArray { // both of which should cleanly cast to isize on an architecture that supports // 32/64-bit offsets let b = std::slice::from_raw_parts( - self.value_data.as_ptr().offset(start.to_isize().unwrap()), - (end - start).to_usize().unwrap(), + self.value_data + .as_ptr() + .offset(start.to_isize().unwrap_unchecked()), + (end - start).to_usize().unwrap_unchecked(), ); // SAFETY: diff --git a/arrow-array/src/array/byte_view_array.rs b/arrow-array/src/array/byte_view_array.rs index 9d2d396a5266..43ff3f76369f 100644 --- a/arrow-array/src/array/byte_view_array.rs +++ b/arrow-array/src/array/byte_view_array.rs @@ -22,11 +22,12 @@ use crate::types::bytes::ByteArrayNativeType; use crate::types::{BinaryViewType, ByteViewType, StringViewType}; use crate::{Array, ArrayAccessor, ArrayRef, GenericByteArray, OffsetSizeTrait, Scalar}; use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, ScalarBuffer}; -use arrow_data::{ArrayData, ArrayDataBuilder, ByteView}; +use arrow_data::{ArrayData, ArrayDataBuilder, ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::{ArrowError, DataType}; use core::str; use num::ToPrimitive; use std::any::Any; +use std::cmp::Ordering; use std::fmt::Debug; use std::marker::PhantomData; use std::sync::Arc; @@ -77,8 +78,9 @@ use super::ByteArrayType; /// 0 31 63 95 127 /// ``` /// -/// * Strings with length <= 12 are stored directly in the view. See -/// [`Self::inline_value`] to access the inlined prefix from a short view. +/// * Strings with length <= 12 ([`MAX_INLINE_VIEW_LEN`]) are stored directly in +/// the view. See [`Self::inline_value`] to access the inlined prefix from a +/// short view. /// /// * Strings with length > 12: The first four bytes are stored inline in the /// view and the entire string is stored in one of the buffers. See [`ByteView`] @@ -128,6 +130,7 @@ use super::ByteArrayType; /// assert_eq!(value, "this string is also longer than 12 bytes"); /// ``` /// +/// [`MAX_INLINE_VIEW_LEN`]: arrow_data::MAX_INLINE_VIEW_LEN /// [`arrow_compute`]: https://docs.rs/arrow/latest/arrow/compute/index.html /// /// Unlike [`GenericByteArray`], there are no constraints on the offsets other @@ -150,7 +153,7 @@ use super::ByteArrayType; /// "CrumpleFacedFish" │ 16 │ Crum │ 0 │ 103 │─ ─│─ ─ ─ ┘ │eFa│ /// └──────┴──────┴──────┴──────┘ │ced│ /// ┌──────┬────────────────────┐ └ ─ ─ ─ ─ ─ ─ ─ ─ ▶│Fis│ -/// "LavaMonster" │ 11 │ LavaMonster\0 │ │hWa│ +/// "LavaMonster" │ 11 │ LavaMonster │ │hWa│ /// └──────┴────────────────────┘ offset │sIn│ /// 115 │Tow│ /// │nTo│ @@ -232,6 +235,10 @@ impl GenericByteViewArray { buffers: Vec, nulls: Option, ) -> Self { + if cfg!(feature = "force_validate") { + return Self::new(views, buffers, nulls); + } + Self { data_type: T::DATA_TYPE, phantom: Default::default(), @@ -312,7 +319,7 @@ impl GenericByteViewArray { pub unsafe fn value_unchecked(&self, idx: usize) -> &T::Native { let v = self.views.get_unchecked(idx); let len = *v as u32; - let b = if len <= 12 { + let b = if len <= MAX_INLINE_VIEW_LEN { Self::inline_value(v, len as usize) } else { let view = ByteView::from(*v); @@ -327,10 +334,10 @@ impl GenericByteViewArray { /// /// # Safety /// - The `view` must be a valid element from `Self::views()` that adheres to the view layout. - /// - The `len` must be the length of the inlined value. It should never be larger than 12. + /// - The `len` must be the length of the inlined value. It should never be larger than [`MAX_INLINE_VIEW_LEN`]. #[inline(always)] pub unsafe fn inline_value(view: &u128, len: usize) -> &[u8] { - debug_assert!(len <= 12); + debug_assert!(len <= MAX_INLINE_VIEW_LEN as usize); std::slice::from_raw_parts((view as *const u128 as *const u8).wrapping_add(4), len) } @@ -343,7 +350,7 @@ impl GenericByteViewArray { pub fn bytes_iter(&self) -> impl Iterator { self.views.iter().map(move |v| { let len = *v as u32; - if len <= 12 { + if len <= MAX_INLINE_VIEW_LEN { unsafe { Self::inline_value(v, len as usize) } } else { let view = ByteView::from(*v); @@ -367,7 +374,7 @@ impl GenericByteViewArray { return &[] as &[u8]; } - if prefix_len <= 4 || len <= 12 { + if prefix_len <= 4 || len as u32 <= MAX_INLINE_VIEW_LEN { unsafe { StringViewArray::inline_value(v, prefix_len) } } else { let view = ByteView::from(*v); @@ -397,7 +404,7 @@ impl GenericByteViewArray { return &[] as &[u8]; } - if len <= 12 { + if len as u32 <= MAX_INLINE_VIEW_LEN { unsafe { &StringViewArray::inline_value(v, len)[len - suffix_len..] } } else { let view = ByteView::from(*v); @@ -466,13 +473,115 @@ impl GenericByteViewArray { /// Note: this function does not attempt to canonicalize / deduplicate values. For this /// feature see [`GenericByteViewBuilder::with_deduplicate_strings`]. pub fn gc(&self) -> Self { - let mut builder = GenericByteViewBuilder::::with_capacity(self.len()); + // 1) Read basic properties once + let len = self.len(); // number of elements + let nulls = self.nulls().cloned(); // reuse & clone existing null bitmap + + // 1.5) Fast path: if there are no buffers, just reuse original views and no data blocks + if self.data_buffers().is_empty() { + return unsafe { + GenericByteViewArray::new_unchecked( + self.views().clone(), + vec![], // empty data blocks + nulls, + ) + }; + } - for v in self.iter() { - builder.append_option(v); + // 2) Calculate total size of all non-inline data and detect if any exists + let total_large = self.total_buffer_bytes_used(); + + // 2.5) Fast path: if there is no non-inline data, avoid buffer allocation & processing + if total_large == 0 { + // Views are inline-only or all null; just reuse original views and no data blocks + return unsafe { + GenericByteViewArray::new_unchecked( + self.views().clone(), + vec![], // empty data blocks + nulls, + ) + }; } - builder.finish() + // 3) Allocate exactly capacity for all non-inline data + let mut data_buf = Vec::with_capacity(total_large); + + // 4) Iterate over views and process each inline/non-inline view + let views_buf: Vec = (0..len) + .map(|i| unsafe { self.copy_view_to_buffer(i, &mut data_buf) }) + .collect(); + + // 5) Wrap up buffers + let data_block = Buffer::from_vec(data_buf); + let views_scalar = ScalarBuffer::from(views_buf); + let data_blocks = vec![data_block]; + + // SAFETY: views_scalar, data_blocks, and nulls are correctly aligned and sized + unsafe { GenericByteViewArray::new_unchecked(views_scalar, data_blocks, nulls) } + } + + /// Copy the i‑th view into `data_buf` if it refers to an out‑of‑line buffer. + /// + /// # Safety + /// + /// - `i < self.len()`. + /// - Every element in `self.views()` must currently refer to a valid slice + /// inside one of `self.buffers`. + /// - `data_buf` must be ready to have additional bytes appended. + /// - After this call, the returned view will have its + /// `buffer_index` reset to `0` and its `offset` updated so that it points + /// into the bytes just appended at the end of `data_buf`. + #[inline(always)] + unsafe fn copy_view_to_buffer(&self, i: usize, data_buf: &mut Vec) -> u128 { + // SAFETY: `i < self.len()` ensures this is in‑bounds. + let raw_view = *self.views().get_unchecked(i); + let mut bv = ByteView::from(raw_view); + + // Inline‑small views stay as‑is. + if bv.length <= MAX_INLINE_VIEW_LEN { + raw_view + } else { + // SAFETY: `bv.buffer_index` and `bv.offset..bv.offset+bv.length` + // must both lie within valid ranges for `self.buffers`. + let buffer = self.buffers.get_unchecked(bv.buffer_index as usize); + let start = bv.offset as usize; + let end = start + bv.length as usize; + let slice = buffer.get_unchecked(start..end); + + // Copy out‑of‑line data into our single “0” buffer. + let new_offset = data_buf.len() as u32; + data_buf.extend_from_slice(slice); + + bv.buffer_index = 0; + bv.offset = new_offset; + bv.into() + } + } + + /// Returns the total number of bytes used by all non inlined views in all + /// buffers. + /// + /// Note this does not account for views that point at the same underlying + /// data in buffers + /// + /// For example, if the array has three strings views: + /// * View with length = 9 (inlined) + /// * View with length = 32 (non inlined) + /// * View with length = 16 (non inlined) + /// + /// Then this method would report 48 + pub fn total_buffer_bytes_used(&self) -> usize { + self.views() + .iter() + .map(|v| { + let len = *v as u32; + if len > MAX_INLINE_VIEW_LEN { + len as usize + } else { + 0 + } + }) + .sum() } /// Compare two [`GenericByteViewArray`] at index `left_idx` and `right_idx` @@ -481,11 +590,11 @@ impl GenericByteViewArray { /// It takes a bit of patience to understand why we don't just compare two &[u8] directly. /// /// ByteView types give us the following two advantages, and we need to be careful not to lose them: - /// (1) For string/byte smaller than 12 bytes, the entire data is inlined in the view. + /// (1) For string/byte smaller than [`MAX_INLINE_VIEW_LEN`] bytes, the entire data is inlined in the view. /// Meaning that reading one array element requires only one memory access /// (two memory access required for StringArray, one for offset buffer, the other for value buffer). /// - /// (2) For string/byte larger than 12 bytes, we can still be faster than (for certain operations) StringArray/ByteArray, + /// (2) For string/byte larger than [`MAX_INLINE_VIEW_LEN`] bytes, we can still be faster than (for certain operations) StringArray/ByteArray, /// thanks to the inlined 4 bytes. /// Consider equality check: /// If the first four bytes of the two strings are different, we can return false immediately (with just one memory access). @@ -495,8 +604,8 @@ impl GenericByteViewArray { /// e.g., if the inlined 4 bytes are different, we can directly return unequal without looking at the full string. /// /// # Order check flow - /// (1) if both string are smaller than 12 bytes, we can directly compare the data inlined to the view. - /// (2) if any of the string is larger than 12 bytes, we need to compare the full string. + /// (1) if both string are smaller than [`MAX_INLINE_VIEW_LEN`] bytes, we can directly compare the data inlined to the view. + /// (2) if any of the string is larger than [`MAX_INLINE_VIEW_LEN`] bytes, we need to compare the full string. /// (2.1) if the inlined 4 bytes are different, we can return the result immediately. /// (2.2) o.w., we need to compare the full string. /// @@ -507,25 +616,30 @@ impl GenericByteViewArray { left_idx: usize, right: &GenericByteViewArray, right_idx: usize, - ) -> std::cmp::Ordering { + ) -> Ordering { let l_view = left.views().get_unchecked(left_idx); - let l_len = *l_view as u32; + let l_byte_view = ByteView::from(*l_view); let r_view = right.views().get_unchecked(right_idx); - let r_len = *r_view as u32; + let r_byte_view = ByteView::from(*r_view); + + let l_len = l_byte_view.length; + let r_len = r_byte_view.length; if l_len <= 12 && r_len <= 12 { - let l_data = unsafe { GenericByteViewArray::::inline_value(l_view, l_len as usize) }; - let r_data = unsafe { GenericByteViewArray::::inline_value(r_view, r_len as usize) }; - return l_data.cmp(r_data); + return Self::inline_key_fast(*l_view).cmp(&Self::inline_key_fast(*r_view)); } // one of the string is larger than 12 bytes, // we then try to compare the inlined data first - let l_inlined_data = unsafe { GenericByteViewArray::::inline_value(l_view, 4) }; - let r_inlined_data = unsafe { GenericByteViewArray::::inline_value(r_view, 4) }; - if r_inlined_data != l_inlined_data { - return l_inlined_data.cmp(r_inlined_data); + + // Note: In theory, ByteView is only used for string which is larger than 12 bytes, + // but we can still use it to get the inlined prefix for shorter strings. + // The prefix is always the first 4 bytes of the view, for both short and long strings. + let l_inlined_be = l_byte_view.prefix.swap_bytes(); + let r_inlined_be = r_byte_view.prefix.swap_bytes(); + if l_inlined_be != r_inlined_be { + return l_inlined_be.cmp(&r_inlined_be); } // unfortunately, we need to compare the full data @@ -534,6 +648,119 @@ impl GenericByteViewArray { l_full_data.cmp(r_full_data) } + + /// Builds a 128-bit composite key for an inline value: + /// + /// - High 96 bits: the inline data in big-endian byte order (for correct lexicographical sorting). + /// - Low 32 bits: the length in big-endian byte order, acting as a tiebreaker so shorter strings + /// (or those with fewer meaningful bytes) always numerically sort before longer ones. + /// + /// This function extracts the length and the 12-byte inline string data from the raw + /// little-endian `u128` representation, converts them to big-endian ordering, and packs them + /// into a single `u128` value suitable for fast, branchless comparisons. + /// + /// # Why include length? + /// + /// A pure 96-bit content comparison can’t distinguish between two values whose inline bytes + /// compare equal—either because one is a true prefix of the other or because zero-padding + /// hides extra bytes. By tucking the 32-bit length into the lower bits, a single `u128` compare + /// handles both content and length in one go. + /// + /// Example: comparing "bar" (3 bytes) vs "bar\0" (4 bytes) + /// + /// | String | Bytes 0–4 (length LE) | Bytes 4–16 (data + padding) | + /// |------------|-----------------------|---------------------------------| + /// | `"bar"` | `03 00 00 00` | `62 61 72` + 9 × `00` | + /// | `"bar\0"`| `04 00 00 00` | `62 61 72 00` + 8 × `00` | + /// + /// Both inline parts become `62 61 72 00…00`, so they tie on content. The length field + /// then differentiates: + /// + /// ```text + /// key("bar") = 0x0000000000000000000062617200000003 + /// key("bar\0") = 0x0000000000000000000062617200000004 + /// ⇒ key("bar") < key("bar\0") + /// ``` + /// # Inlining and Endianness + /// + /// - We start by calling `.to_le_bytes()` on the `raw` `u128`, because Rust’s native in‑memory + /// representation is little‑endian on x86/ARM. + /// - We extract the low 32 bits numerically (`raw as u32`)—this step is endianness‑free. + /// - We copy the 12 bytes of inline data (original order) into `buf[0..12]`. + /// - We serialize `length` as big‑endian into `buf[12..16]`. + /// - Finally, `u128::from_be_bytes(buf)` treats `buf[0]` as the most significant byte + /// and `buf[15]` as the least significant, producing a `u128` whose integer value + /// directly encodes “inline data then length” in big‑endian form. + /// + /// This ensures that a simple `u128` comparison is equivalent to the desired + /// lexicographical comparison of the inline bytes followed by length. + #[inline(always)] + pub fn inline_key_fast(raw: u128) -> u128 { + // 1. Decompose `raw` into little‑endian bytes: + // - raw_bytes[0..4] = length in LE + // - raw_bytes[4..16] = inline string data + let raw_bytes = raw.to_le_bytes(); + + // 2. Numerically truncate to get the low 32‑bit length (endianness‑free). + let length = raw as u32; + + // 3. Build a 16‑byte buffer in big‑endian order: + // - buf[0..12] = inline string bytes (in original order) + // - buf[12..16] = length.to_be_bytes() (BE) + let mut buf = [0u8; 16]; + buf[0..12].copy_from_slice(&raw_bytes[4..16]); // inline data + + // Why convert length to big-endian for comparison? + // + // Rust (on most platforms) stores integers in little-endian format, + // meaning the least significant byte is at the lowest memory address. + // For example, an u32 value like 0x22345677 is stored in memory as: + // + // [0x77, 0x56, 0x34, 0x22] // little-endian layout + // ^ ^ ^ ^ + // LSB ↑↑↑ MSB + // + // This layout is efficient for arithmetic but *not* suitable for + // lexicographic (dictionary-style) comparison of byte arrays. + // + // To compare values by byte order—e.g., for sorted keys or binary trees— + // we must convert them to **big-endian**, where: + // + // - The most significant byte (MSB) comes first (index 0) + // - The least significant byte (LSB) comes last (index N-1) + // + // In big-endian, the same u32 = 0x22345677 would be represented as: + // + // [0x22, 0x34, 0x56, 0x77] + // + // This ordering aligns with natural string/byte sorting, so calling + // `.to_be_bytes()` allows us to construct + // keys where standard numeric comparison (e.g., `<`, `>`) behaves + // like lexicographic byte comparison. + buf[12..16].copy_from_slice(&length.to_be_bytes()); // length in BE + + // 4. Deserialize the buffer as a big‑endian u128: + // buf[0] is MSB, buf[15] is LSB. + // Details: + // Note on endianness and layout: + // + // Although `buf[0]` is stored at the lowest memory address, + // calling `u128::from_be_bytes(buf)` interprets it as the **most significant byte (MSB)**, + // and `buf[15]` as the **least significant byte (LSB)**. + // + // This is the core principle of **big-endian decoding**: + // - Byte at index 0 maps to bits 127..120 (highest) + // - Byte at index 1 maps to bits 119..112 + // - ... + // - Byte at index 15 maps to bits 7..0 (lowest) + // + // So even though memory layout goes from low to high (left to right), + // big-endian treats the **first byte** as highest in value. + // + // This guarantees that comparing two `u128` keys is equivalent to lexicographically + // comparing the original inline bytes, followed by length. + u128::from_be_bytes(buf) + } } impl Debug for GenericByteViewArray { @@ -842,9 +1069,16 @@ impl From>> for StringViewArray { #[cfg(test)] mod tests { use crate::builder::{BinaryViewBuilder, StringViewBuilder}; - use crate::{Array, BinaryViewArray, StringViewArray}; + use crate::types::BinaryViewType; + use crate::{ + Array, BinaryViewArray, GenericBinaryArray, GenericByteViewArray, StringViewArray, + }; use arrow_buffer::{Buffer, ScalarBuffer}; - use arrow_data::ByteView; + use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; + use rand::prelude::StdRng; + use rand::{Rng, SeedableRng}; + + const BLOCK_SIZE: u32 = 8; #[test] fn try_new_string() { @@ -1034,6 +1268,130 @@ mod tests { check_gc(&array.slice(3, 1)); } + /// 1) Empty array: no elements, expect gc to return empty with no data buffers + #[test] + fn test_gc_empty_array() { + let array = StringViewBuilder::new() + .with_fixed_block_size(BLOCK_SIZE) + .finish(); + let gced = array.gc(); + // length and null count remain zero + assert_eq!(gced.len(), 0); + assert_eq!(gced.null_count(), 0); + // no underlying data buffers should be allocated + assert!( + gced.data_buffers().is_empty(), + "Expected no data buffers for empty array" + ); + } + + /// 2) All inline values (<= INLINE_LEN): capacity-only data buffer, same values + #[test] + fn test_gc_all_inline() { + let mut builder = StringViewBuilder::new().with_fixed_block_size(BLOCK_SIZE); + // append many short strings, each exactly INLINE_LEN long + for _ in 0..100 { + let s = "A".repeat(MAX_INLINE_VIEW_LEN as usize); + builder.append_option(Some(&s)); + } + let array = builder.finish(); + let gced = array.gc(); + // Since all views fit inline, data buffer is empty + assert_eq!( + gced.data_buffers().len(), + 0, + "Should have no data buffers for inline values" + ); + assert_eq!(gced.len(), 100); + // verify element-wise equality + array.iter().zip(gced.iter()).for_each(|(orig, got)| { + assert_eq!(orig, got, "Inline value mismatch after gc"); + }); + } + + /// 3) All large values (> INLINE_LEN): each must be copied into the new data buffer + #[test] + fn test_gc_all_large() { + let mut builder = StringViewBuilder::new().with_fixed_block_size(BLOCK_SIZE); + let large_str = "X".repeat(MAX_INLINE_VIEW_LEN as usize + 5); + // append multiple large strings + for _ in 0..50 { + builder.append_option(Some(&large_str)); + } + let array = builder.finish(); + let gced = array.gc(); + // New data buffers should be populated (one or more blocks) + assert!( + !gced.data_buffers().is_empty(), + "Expected data buffers for large values" + ); + assert_eq!(gced.len(), 50); + // verify that every large string emerges unchanged + array.iter().zip(gced.iter()).for_each(|(orig, got)| { + assert_eq!(orig, got, "Large view mismatch after gc"); + }); + } + + /// 4) All null elements: ensure null bitmap handling path is correct + #[test] + fn test_gc_all_nulls() { + let mut builder = StringViewBuilder::new().with_fixed_block_size(BLOCK_SIZE); + for _ in 0..20 { + builder.append_null(); + } + let array = builder.finish(); + let gced = array.gc(); + // length and null count match + assert_eq!(gced.len(), 20); + assert_eq!(gced.null_count(), 20); + // data buffers remain empty for null-only array + assert!( + gced.data_buffers().is_empty(), + "No data should be stored for nulls" + ); + } + + /// 5) Random mix of inline, large, and null values with slicing tests + #[test] + fn test_gc_random_mixed_and_slices() { + let mut rng = StdRng::seed_from_u64(42); + let mut builder = StringViewBuilder::new().with_fixed_block_size(BLOCK_SIZE); + // Keep a Vec of original Option for later comparison + let mut original: Vec> = Vec::new(); + + for _ in 0..200 { + if rng.random_bool(0.1) { + // 10% nulls + builder.append_null(); + original.push(None); + } else { + // random length between 0 and twice the inline limit + let len = rng.random_range(0..(MAX_INLINE_VIEW_LEN * 2)); + let s: String = "A".repeat(len as usize); + builder.append_option(Some(&s)); + original.push(Some(s)); + } + } + + let array = builder.finish(); + // Test multiple slice ranges to ensure offset logic is correct + for (offset, slice_len) in &[(0, 50), (10, 100), (150, 30)] { + let sliced = array.slice(*offset, *slice_len); + let gced = sliced.gc(); + // Build expected slice of Option<&str> + let expected: Vec> = original[*offset..(*offset + *slice_len)] + .iter() + .map(|opt| opt.as_deref()) + .collect(); + + assert_eq!(gced.len(), *slice_len, "Slice length mismatch"); + // Compare element-wise + gced.iter().zip(expected.iter()).for_each(|(got, expect)| { + assert_eq!(got, *expect, "Value mismatch in mixed slice after gc"); + }); + } + } + #[test] fn test_eq() { let test_data = [ @@ -1058,4 +1416,105 @@ mod tests { assert_eq!(array2, array2.clone()); assert_eq!(array1, array2); } + + /// Integration tests for `inline_key_fast` covering: + /// + /// 1. Monotonic ordering across increasing lengths and lexical variations. + /// 2. Cross-check against `GenericBinaryArray` comparison to ensure semantic equivalence. + /// + /// This also includes a specific test for the “bar” vs. “bar\0” case, demonstrating why + /// the length field is required even when all inline bytes fit in 12 bytes. + /// + /// The test includes strings that verify correct byte order (prevent reversal bugs), + /// and length-based tie-breaking in the composite key. + /// + /// The test confirms that `inline_key_fast` produces keys which sort consistently + /// with the expected lexicographical order of the raw byte arrays. + #[test] + fn test_inline_key_fast_various_lengths_and_lexical() { + /// Helper to create a raw u128 value representing an inline ByteView: + /// - `length`: number of meaningful bytes (must be ≤ 12) + /// - `data`: the actual inline data bytes + /// + /// The first 4 bytes encode length in little-endian, + /// the following 12 bytes contain the inline string data (unpadded). + fn make_raw_inline(length: u32, data: &[u8]) -> u128 { + assert!(length as usize <= 12, "Inline length must be ≤ 12"); + assert!( + data.len() == length as usize, + "Data length must match `length`" + ); + + let mut raw_bytes = [0u8; 16]; + raw_bytes[0..4].copy_from_slice(&length.to_le_bytes()); // length stored little-endian + raw_bytes[4..(4 + data.len())].copy_from_slice(data); // inline data + u128::from_le_bytes(raw_bytes) + } + + // Test inputs: various lengths and lexical orders, + // plus special cases for byte order and length tie-breaking + let test_inputs: Vec<&[u8]> = vec![ + b"a", + b"aa", + b"aaa", + b"aab", + b"abcd", + b"abcde", + b"abcdef", + b"abcdefg", + b"abcdefgh", + b"abcdefghi", + b"abcdefghij", + b"abcdefghijk", + b"abcdefghijkl", + // Tests for byte-order reversal bug: + // Without the fix, "backend one" would compare as "eno dnekcab", + // causing incorrect sort order relative to "backend two". + b"backend one", + b"backend two", + // Tests length-tiebreaker logic: + // "bar" (3 bytes) and "bar\0" (4 bytes) have identical inline data, + // so only the length differentiates their ordering. + b"bar", + b"bar\0", + // Additional lexical and length tie-breaking cases with same prefix, in correct lex order: + b"than12Byt", + b"than12Bytes", + b"than12Bytes\0", + b"than12Bytesx", + b"than12Bytex", + b"than12Bytez", + // Additional lexical tests + b"xyy", + b"xyz", + b"xza", + ]; + + // Create a GenericBinaryArray for cross-comparison of lex order + let array: GenericBinaryArray = + GenericBinaryArray::from(test_inputs.iter().map(|s| Some(*s)).collect::>()); + + for i in 0..array.len() - 1 { + let v1 = array.value(i); + let v2 = array.value(i + 1); + + // Assert the array's natural lexical ordering is correct + assert!(v1 < v2, "Array compare failed: {v1:?} !< {v2:?}"); + + // Assert the keys produced by inline_key_fast reflect the same ordering + let key1 = GenericByteViewArray::::inline_key_fast(make_raw_inline( + v1.len() as u32, + v1, + )); + let key2 = GenericByteViewArray::::inline_key_fast(make_raw_inline( + v2.len() as u32, + v2, + )); + + assert!( + key1 < key2, + "Key compare failed: key({v1:?})=0x{key1:032x} !< key({v2:?})=0x{key2:032x}", + ); + } + } } diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs index f852b57fb65e..acbdcb8b60fa 100644 --- a/arrow-array/src/array/dictionary_array.rs +++ b/arrow-array/src/array/dictionary_array.rs @@ -327,6 +327,10 @@ impl DictionaryArray { /// /// Safe provided [`Self::try_new`] would not return an error pub unsafe fn new_unchecked(keys: PrimitiveArray, values: ArrayRef) -> Self { + if cfg!(feature = "force_validate") { + return Self::new(keys, values); + } + let data_type = DataType::Dictionary( Box::new(keys.data_type().clone()), Box::new(values.data_type().clone()), @@ -481,6 +485,7 @@ impl DictionaryArray { /// Returns `PrimitiveDictionaryBuilder` of this dictionary array for mutating /// its keys and values if the underlying data buffer is not shared by others. + #[allow(clippy::result_large_err)] pub fn into_primitive_dict_builder(self) -> Result, Self> where V: ArrowPrimitiveType, @@ -537,6 +542,7 @@ impl DictionaryArray { /// assert_eq!(typed.value(1), 11); /// assert_eq!(typed.value(2), 21); /// ``` + #[allow(clippy::result_large_err)] pub fn unary_mut(self, op: F) -> Result, DictionaryArray> where V: ArrowPrimitiveType, diff --git a/arrow-array/src/array/fixed_size_binary_array.rs b/arrow-array/src/array/fixed_size_binary_array.rs index 576b8012491b..55973a58f2cb 100644 --- a/arrow-array/src/array/fixed_size_binary_array.rs +++ b/arrow-array/src/array/fixed_size_binary_array.rs @@ -87,7 +87,7 @@ impl FixedSizeBinaryArray { ) -> Result { let data_type = DataType::FixedSizeBinary(size); let s = size.to_usize().ok_or_else(|| { - ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {}", size)) + ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}")) })?; let len = values.len() / s; diff --git a/arrow-array/src/array/fixed_size_list_array.rs b/arrow-array/src/array/fixed_size_list_array.rs index 44be442c9f85..f807cc88fbca 100644 --- a/arrow-array/src/array/fixed_size_list_array.rs +++ b/arrow-array/src/array/fixed_size_list_array.rs @@ -149,7 +149,7 @@ impl FixedSizeListArray { nulls: Option, ) -> Result { let s = size.to_usize().ok_or_else(|| { - ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {}", size)) + ArrowError::InvalidArgumentError(format!("Size cannot be negative, got {size}")) })?; let len = match s { @@ -343,8 +343,8 @@ impl From for FixedSizeListArray { fn from(data: ArrayData) -> Self { let value_length = match data.data_type() { DataType::FixedSizeList(_, len) => *len, - _ => { - panic!("FixedSizeListArray data should contain a FixedSizeList data type") + data_type => { + panic!("FixedSizeListArray data should contain a FixedSizeList data type, got {data_type:?}") } }; diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs index bed0bdf889b2..832a1c0a9ad8 100644 --- a/arrow-array/src/array/list_array.rs +++ b/arrow-array/src/array/list_array.rs @@ -42,20 +42,26 @@ pub trait OffsetSizeTrait: ArrowNativeType + std::ops::AddAssign + Integer { const IS_LARGE: bool; /// Prefix for the offset size const PREFIX: &'static str; + /// The max `usize` offset + const MAX_OFFSET: usize; } impl OffsetSizeTrait for i32 { const IS_LARGE: bool = false; const PREFIX: &'static str = ""; + const MAX_OFFSET: usize = i32::MAX as usize; } impl OffsetSizeTrait for i64 { const IS_LARGE: bool = true; const PREFIX: &'static str = "Large"; + const MAX_OFFSET: usize = i64::MAX as usize; } /// An array of [variable length lists], similar to JSON arrays -/// (e.g. `["A", "B", "C"]`). +/// (e.g. `["A", "B", "C"]`). This struct specifically represents +/// the [list layout]. Refer to [`GenericListViewArray`] for the +/// [list-view layout]. /// /// Lists are represented using `offsets` into a `values` child /// array. Offsets are stored in two adjacent entries of an @@ -118,12 +124,48 @@ impl OffsetSizeTrait for i64 { /// (offsets[i], │ ListArray (Array) /// offsets[i+1]) └ ─ ─ ─ ─ ─ ─ ┘ │ /// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ``` +/// +/// # Slicing +/// +/// Slicing a `ListArray` creates a new `ListArray` without copying any data, +/// but this means the [`Self::values`] and [`Self::offsets`] may have "unused" data /// +/// For example, calling `slice(1, 3)` on the `ListArray` in the above example +/// would result in the following. Note /// +/// 1. `Values` array is unchanged +/// 2. `Offsets` do not start at `0`, nor cover all values in the Values array. +/// +/// ```text +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ┌ ─ ─ ─ ─ ─ ─ ┐ │ ╔═══╗ +/// │ ╔═══╗ ╔═══╗ ║ ║ Not used +/// │ ║ 1 ║ ║ A ║ │ 0 │ ╚═══╝ +/// ┌─────────────┐ ┌───────┐ │ ┌───┐ ┌───┐ ╠═══╣ ╠═══╣ +/// │ [] (empty) │ │ (3,3) │ │ 1 │ │ 3 │ │ ║ 1 ║ ║ B ║ │ 1 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ╠═══╣ ╠═══╣ +/// │ NULL │ │ (3,4) │ │ 0 │ │ 3 │ │ ║ 1 ║ ║ C ║ │ 2 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ╠───╣ ╠───╣ +/// │ [D] │ │ (4,5) │ │ 1 │ │ 4 │ │ │ 0 │ │ ? │ │ 3 │ +/// └─────────────┘ └───────┘ │ └───┘ ├───┤ ├───┤ ├───┤ +/// │ 5 │ │ │ 1 │ │ D │ │ 4 │ +/// │ └───┘ ├───┤ ├───┤ +/// │ │ 0 │ │ ? │ │ 5 │ +/// │ Validity ╠═══╣ ╠═══╣ +/// Logical Logical (nulls) Offsets │ ║ 1 ║ ║ F ║ │ 6 │ +/// Values Offsets │ ╚═══╝ ╚═══╝ +/// │ Values │ │ +/// (offsets[i], │ ListArray (Array) +/// offsets[i+1]) └ ─ ─ ─ ─ ─ ─ ┘ │ +/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ /// ``` /// /// [`StringArray`]: crate::array::StringArray +/// [`GenericListViewArray`]: crate::array::GenericListViewArray /// [variable length lists]: https://arrow.apache.org/docs/format/Columnar.html#variable-size-list-layout +/// [list layout]: https://arrow.apache.org/docs/format/Columnar.html#list-layout +/// [list-view layout]: https://arrow.apache.org/docs/format/Columnar.html#listview-layout pub struct GenericListArray { data_type: DataType, nulls: Option, @@ -258,13 +300,22 @@ impl GenericListArray { /// Returns a reference to the offsets of this list /// /// Unlike [`Self::value_offsets`] this returns the [`OffsetBuffer`] - /// allowing for zero-copy cloning + /// allowing for zero-copy cloning. + /// + /// Notes: The `offsets` may not start at 0 and may not cover all values in + /// [`Self::values`]. This can happen when the list array was sliced via + /// [`Self::slice`]. See documentation for [`Self`] for more details. #[inline] pub fn offsets(&self) -> &OffsetBuffer { &self.value_offsets } /// Returns a reference to the values of this list + /// + /// Note: The list array may not refer to all values in the `values` array. + /// For example if the list array was sliced via [`Self::slice`] values will + /// still contain values both before and after the slice. See documentation + /// for [`Self`] for more details. #[inline] pub fn values(&self) -> &ArrayRef { &self.values @@ -291,7 +342,9 @@ impl GenericListArray { self.values.slice(start, end - start) } - /// Returns the offset values in the offsets buffer + /// Returns the offset values in the offsets buffer. + /// + /// See [`Self::offsets`] for more details. #[inline] pub fn value_offsets(&self) -> &[OffsetSize] { &self.value_offsets @@ -320,6 +373,10 @@ impl GenericListArray { } /// Returns a zero-copy slice of this array with the indicated offset and length. + /// + /// Notes: this method does *NOT* slice the underlying values array or modify + /// the values in the offsets buffer. See [`Self::values`] and + /// [`Self::offsets`] for more information. pub fn slice(&self, offset: usize, length: usize) -> Self { Self { data_type: self.data_type.clone(), @@ -397,7 +454,7 @@ impl From for GenericListArray< _ => unreachable!(), }; - let offsets = OffsetBuffer::from_lengths(std::iter::repeat(size).take(value.len())); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(size, value.len())); Self { data_type: Self::DATA_TYPE_CONSTRUCTOR(field.clone()), @@ -551,12 +608,12 @@ impl std::fmt::Debug for GenericListArray; /// A [`GenericListArray`] of variable size lists, storing offsets as `i64`. /// -// See [`LargeListBuilder`](crate::builder::LargeListBuilder) for how to construct a [`LargeListArray`] +/// See [`LargeListBuilder`](crate::builder::LargeListBuilder) for how to construct a [`LargeListArray`] pub type LargeListArray = GenericListArray; #[cfg(test)] diff --git a/arrow-array/src/array/list_view_array.rs b/arrow-array/src/array/list_view_array.rs index 7e52a6f3e457..a239ea1e5e73 100644 --- a/arrow-array/src/array/list_view_array.rs +++ b/arrow-array/src/array/list_view_array.rs @@ -32,16 +32,81 @@ pub type ListViewArray = GenericListViewArray; /// A [`GenericListViewArray`] of variable size lists, storing offsets as `i64`. pub type LargeListViewArray = GenericListViewArray; +/// An array of [variable length lists], specifically in the [list-view layout]. /// -/// Different from [`crate::GenericListArray`] as it stores both an offset and length -/// meaning that take / filter operations can be implemented without copying the underlying data. +/// Differs from [`GenericListArray`] (which represents the [list layout]) in that +/// the sizes of the child arrays are explicitly encoded in a separate buffer, instead +/// of being derived from the difference between subsequent offsets in the offset buffer. /// -/// [Variable-size List Layout: ListView Layout]: https://arrow.apache.org/docs/format/Columnar.html#listview-layout +/// This allows the offsets (and subsequently child data) to be out of order. It also +/// allows take / filter operations to be implemented without copying the underlying data. +/// +/// # Representation +/// +/// Given the same example array from [`GenericListArray`], it would be represented +/// as such via a list-view layout array: +/// +/// ```text +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ┌ ─ ─ ─ ─ ─ ─ ┐ │ +/// ┌─────────────┐ ┌───────┐ │ ┌───┐ ┌───┐ ┌───┐ ┌───┐ ┌───┐ +/// │ [A,B,C] │ │ (0,3) │ │ 1 │ │ 0 │ │ 3 │ │ │ 1 │ │ A │ │ 0 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [] │ │ (3,0) │ │ 1 │ │ 3 │ │ 0 │ │ │ 1 │ │ B │ │ 1 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ NULL │ │ (?,?) │ │ 0 │ │ ? │ │ ? │ │ │ 1 │ │ C │ │ 2 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [D] │ │ (4,1) │ │ 1 │ │ 4 │ │ 1 │ │ │ ? │ │ ? │ │ 3 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [NULL, F] │ │ (5,2) │ │ 1 │ │ 5 │ │ 2 │ │ │ 1 │ │ D │ │ 4 │ +/// └─────────────┘ └───────┘ │ └───┘ └───┘ └───┘ ├───┤ ├───┤ +/// │ │ 0 │ │ ? │ │ 5 │ +/// Logical Logical │ Validity Offsets Sizes ├───┤ ├───┤ +/// Values Offset (nulls) │ │ 1 │ │ F │ │ 6 │ +/// & Size │ └───┘ └───┘ +/// │ Values │ │ +/// (offsets[i], │ ListViewArray (Array) +/// sizes[i]) └ ─ ─ ─ ─ ─ ─ ┘ │ +/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ``` +/// +/// Another way of representing the same array but taking advantage of the offsets being out of order: +/// +/// ```text +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ┌ ─ ─ ─ ─ ─ ─ ┐ │ +/// ┌─────────────┐ ┌───────┐ │ ┌───┐ ┌───┐ ┌───┐ ┌───┐ ┌───┐ +/// │ [A,B,C] │ │ (2,3) │ │ 1 │ │ 2 │ │ 3 │ │ │ 0 │ │ ? │ │ 0 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [] │ │ (0,0) │ │ 1 │ │ 0 │ │ 0 │ │ │ 1 │ │ F │ │ 1 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ NULL │ │ (?,?) │ │ 0 │ │ ? │ │ ? │ │ │ 1 │ │ A │ │ 2 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [D] │ │ (5,1) │ │ 1 │ │ 5 │ │ 1 │ │ │ 1 │ │ B │ │ 3 │ +/// ├─────────────┤ ├───────┤ │ ├───┤ ├───┤ ├───┤ ├───┤ ├───┤ +/// │ [NULL, F] │ │ (0,2) │ │ 1 │ │ 0 │ │ 2 │ │ │ 1 │ │ C │ │ 4 │ +/// └─────────────┘ └───────┘ │ └───┘ └───┘ └───┘ ├───┤ ├───┤ +/// │ │ 1 │ │ D │ │ 5 │ +/// Logical Logical │ Validity Offsets Sizes └───┘ └───┘ +/// Values Offset (nulls) │ Values │ │ +/// & Size │ (Array) +/// └ ─ ─ ─ ─ ─ ─ ┘ │ +/// (offsets[i], │ ListViewArray +/// sizes[i]) │ +/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// ``` +/// +/// [`GenericListArray`]: crate::array::GenericListArray +/// [variable length lists]: https://arrow.apache.org/docs/format/Columnar.html#variable-size-list-layout +/// [list layout]: https://arrow.apache.org/docs/format/Columnar.html#list-layout +/// [list-view layout]: https://arrow.apache.org/docs/format/Columnar.html#listview-layout #[derive(Clone)] pub struct GenericListViewArray { data_type: DataType, nulls: Option, values: ArrayRef, + // Unlike GenericListArray, we do not use OffsetBuffer here as offsets are not + // guaranteed to be monotonically increasing. value_offsets: ScalarBuffer, value_sizes: ScalarBuffer, } @@ -410,7 +475,7 @@ impl From for GenericListViewAr _ => unreachable!(), }; let mut acc = 0_usize; - let iter = std::iter::repeat(size).take(value.len()); + let iter = std::iter::repeat_n(size, value.len()); let mut sizes = Vec::with_capacity(iter.size_hint().0); let mut offsets = Vec::with_capacity(iter.size_hint().0); @@ -830,8 +895,8 @@ mod tests { .build() .unwrap(), ); - assert_eq!(string.value_offsets(), &[]); - assert_eq!(string.value_sizes(), &[]); + assert_eq!(string.value_offsets(), &[] as &[i32; 0]); + assert_eq!(string.value_sizes(), &[] as &[i32; 0]); let string = LargeListViewArray::from( ArrayData::builder(DataType::LargeListView(f)) @@ -841,8 +906,8 @@ mod tests { .unwrap(), ); assert_eq!(string.len(), 0); - assert_eq!(string.value_offsets(), &[]); - assert_eq!(string.value_sizes(), &[]); + assert_eq!(string.value_offsets(), &[] as &[i64; 0]); + assert_eq!(string.value_sizes(), &[] as &[i64; 0]); } #[test] diff --git a/arrow-array/src/array/mod.rs b/arrow-array/src/array/mod.rs index 23b3cb628aaf..5fdfb9fb2244 100644 --- a/arrow-array/src/array/mod.rs +++ b/arrow-array/src/array/mod.rs @@ -620,6 +620,29 @@ impl<'a> StringArrayType<'a> for &'a StringViewArray { } } +/// A trait for Arrow String Arrays, currently three types are supported: +/// - `BinaryArray` +/// - `LargeBinaryArray` +/// - `BinaryViewArray` +/// +/// This trait helps to abstract over the different types of binary arrays +/// so that we don't need to duplicate the implementation for each type. +pub trait BinaryArrayType<'a>: ArrayAccessor + Sized { + /// Constructs a new iterator + fn iter(&self) -> ArrayIter; +} + +impl<'a, O: OffsetSizeTrait> BinaryArrayType<'a> for &'a GenericBinaryArray { + fn iter(&self) -> ArrayIter { + GenericBinaryArray::::iter(self) + } +} +impl<'a> BinaryArrayType<'a> for &'a BinaryViewArray { + fn iter(&self) -> ArrayIter { + BinaryViewArray::iter(self) + } +} + impl PartialEq for dyn Array + '_ { fn eq(&self, other: &Self) -> bool { self.to_data().eq(&other.to_data()) @@ -710,6 +733,12 @@ impl PartialEq for GenericByteViewArray { } } +impl PartialEq for RunArray { + fn eq(&self, other: &Self) -> bool { + self.to_data().eq(&other.to_data()) + } +} + /// Constructs an array using the input `data`. /// Returns a reference-counted `Array` instance. pub fn make_array(data: ArrayData) -> ArrayRef { @@ -804,6 +833,8 @@ pub fn make_array(data: ArrayData) -> ArrayRef { dt => panic!("Unexpected data type for run_ends array {dt:?}"), }, DataType::Null => Arc::new(NullArray::from(data)) as ArrayRef, + DataType::Decimal32(_, _) => Arc::new(Decimal32Array::from(data)) as ArrayRef, + DataType::Decimal64(_, _) => Arc::new(Decimal64Array::from(data)) as ArrayRef, DataType::Decimal128(_, _) => Arc::new(Decimal128Array::from(data)) as ArrayRef, DataType::Decimal256(_, _) => Arc::new(Decimal256Array::from(data)) as ArrayRef, dt => panic!("Unexpected data type {dt:?}"), diff --git a/arrow-array/src/array/null_array.rs b/arrow-array/src/array/null_array.rs index 9a7a5ebe17fe..2dd9570a0e94 100644 --- a/arrow-array/src/array/null_array.rs +++ b/arrow-array/src/array/null_array.rs @@ -170,6 +170,9 @@ impl std::fmt::Debug for NullArray { #[cfg(test)] mod tests { use super::*; + use crate::{make_array, Int64Array, StructArray}; + use arrow_data::transform::MutableArrayData; + use arrow_schema::Field; #[test] fn test_null_array() { @@ -201,4 +204,32 @@ mod tests { let array = NullArray::new(1024 * 1024); assert_eq!(format!("{array:?}"), "NullArray(1048576)"); } + + #[test] + fn test_null_array_with_parent_null_buffer() { + let null_array = NullArray::new(1); + let int_array = Int64Array::from(vec![42]); + + let fields = vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Null, true), + ]; + + let struct_array_data = ArrayData::builder(DataType::Struct(fields.into())) + .len(1) + .add_child_data(int_array.to_data()) + .add_child_data(null_array.to_data()) + .build() + .unwrap(); + + let mut mutable = MutableArrayData::new(vec![&struct_array_data], true, 1); + + // Simulate a NULL value in the parent array, for instance, if array being queried by + // invalid index + mutable.extend_nulls(1); + let data = mutable.freeze(); + + let struct_array = Arc::new(StructArray::from(data.clone())); + assert!(make_array(data) == struct_array); + } } diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 57aa23bf9040..9327668824f8 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -410,6 +410,44 @@ pub type DurationMicrosecondArray = PrimitiveArray; /// A [`PrimitiveArray`] of elapsed durations in nanoseconds pub type DurationNanosecondArray = PrimitiveArray; +/// A [`PrimitiveArray`] of 32-bit fixed point decimals +/// +/// # Examples +/// +/// Construction +/// +/// ``` +/// # use arrow_array::Decimal32Array; +/// // Create from Vec> +/// let arr = Decimal32Array::from(vec![Some(1), None, Some(2)]); +/// // Create from Vec +/// let arr = Decimal32Array::from(vec![1, 2, 3]); +/// // Create iter/collect +/// let arr: Decimal32Array = std::iter::repeat(42).take(10).collect(); +/// ``` +/// +/// See [`PrimitiveArray`] for more information and examples +pub type Decimal32Array = PrimitiveArray; + +/// A [`PrimitiveArray`] of 64-bit fixed point decimals +/// +/// # Examples +/// +/// Construction +/// +/// ``` +/// # use arrow_array::Decimal64Array; +/// // Create from Vec> +/// let arr = Decimal64Array::from(vec![Some(1), None, Some(2)]); +/// // Create from Vec +/// let arr = Decimal64Array::from(vec![1, 2, 3]); +/// // Create iter/collect +/// let arr: Decimal64Array = std::iter::repeat(42).take(10).collect(); +/// ``` +/// +/// See [`PrimitiveArray`] for more information and examples +pub type Decimal64Array = PrimitiveArray; + /// A [`PrimitiveArray`] of 128-bit fixed point decimals /// /// # Examples @@ -418,7 +456,7 @@ pub type DurationNanosecondArray = PrimitiveArray; /// /// ``` /// # use arrow_array::Decimal128Array; -/// // Create from Vec> +/// // Create from Vec> /// let arr = Decimal128Array::from(vec![Some(1), None, Some(2)]); /// // Create from Vec /// let arr = Decimal128Array::from(vec![1, 2, 3]); @@ -672,6 +710,8 @@ impl PrimitiveArray { DataType::Timestamp(t1, _) => { matches!(data_type, DataType::Timestamp(t2, _) if &t1 == t2) } + DataType::Decimal32(_, _) => matches!(data_type, DataType::Decimal32(_, _)), + DataType::Decimal64(_, _) => matches!(data_type, DataType::Decimal64(_, _)), DataType::Decimal128(_, _) => matches!(data_type, DataType::Decimal128(_, _)), DataType::Decimal256(_, _) => matches!(data_type, DataType::Decimal256(_, _)), _ => T::DATA_TYPE.eq(data_type), @@ -729,10 +769,8 @@ impl PrimitiveArray { /// Creates a PrimitiveArray based on a constant value with `count` elements pub fn from_value(value: T::Native, count: usize) -> Self { - unsafe { - let val_buf = Buffer::from_trusted_len_iter((0..count).map(|_| value)); - Self::new(val_buf.into(), None) - } + let val_buf: Vec<_> = vec![value; count]; + Self::new(val_buf.into(), None) } /// Returns an iterator that returns the values of `array.value(i)` for an iterator with each element `i` @@ -827,13 +865,8 @@ impl PrimitiveArray { F: Fn(T::Native) -> O::Native, { let nulls = self.nulls().cloned(); - let values = self.values().iter().map(|v| op(*v)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size because arrays are sized. - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; + let values = self.values().into_iter().map(|v| op(*v)); + let buffer: Vec<_> = values.collect(); PrimitiveArray::new(buffer.into(), nulls) } @@ -1035,13 +1068,10 @@ impl PrimitiveArray { F: FnMut(U::Item) -> T::Native, { let nulls = left.logical_nulls(); - let buffer = unsafe { + let buffer: Vec<_> = (0..left.len()) // SAFETY: i in range 0..left.len() - let iter = (0..left.len()).map(|i| op(left.value_unchecked(i))); - // SAFETY: upper bound is trusted because `iter` is over a range - Buffer::from_trusted_len_iter(iter) - }; - + .map(|i| op(unsafe { left.value_unchecked(i) })) + .collect(); PrimitiveArray::new(buffer.into(), nulls) } @@ -1353,6 +1383,8 @@ def_from_for_primitive!(UInt64Type, u64); def_from_for_primitive!(Float16Type, f16); def_from_for_primitive!(Float32Type, f32); def_from_for_primitive!(Float64Type, f64); +def_from_for_primitive!(Decimal32Type, i32); +def_from_for_primitive!(Decimal64Type, i64); def_from_for_primitive!(Decimal128Type, i128); def_from_for_primitive!(Decimal256Type, i256); @@ -1465,6 +1497,8 @@ def_numeric_from_vec!(UInt64Type); def_numeric_from_vec!(Float16Type); def_numeric_from_vec!(Float32Type); def_numeric_from_vec!(Float64Type); +def_numeric_from_vec!(Decimal32Type); +def_numeric_from_vec!(Decimal64Type); def_numeric_from_vec!(Decimal128Type); def_numeric_from_vec!(Decimal256Type); @@ -1573,6 +1607,26 @@ impl PrimitiveArray { /// Returns the decimal precision of this array pub fn precision(&self) -> u8 { match T::BYTE_LENGTH { + 4 => { + if let DataType::Decimal32(p, _) = self.data_type() { + *p + } else { + unreachable!( + "Decimal32Array datatype is not DataType::Decimal32 but {}", + self.data_type() + ) + } + } + 8 => { + if let DataType::Decimal64(p, _) = self.data_type() { + *p + } else { + unreachable!( + "Decimal64Array datatype is not DataType::Decimal64 but {}", + self.data_type() + ) + } + } 16 => { if let DataType::Decimal128(p, _) = self.data_type() { *p @@ -1600,6 +1654,26 @@ impl PrimitiveArray { /// Returns the decimal scale of this array pub fn scale(&self) -> i8 { match T::BYTE_LENGTH { + 4 => { + if let DataType::Decimal32(_, s) = self.data_type() { + *s + } else { + unreachable!( + "Decimal32Array datatype is not DataType::Decimal32 but {}", + self.data_type() + ) + } + } + 8 => { + if let DataType::Decimal64(_, s) = self.data_type() { + *s + } else { + unreachable!( + "Decimal64Array datatype is not DataType::Decimal64 but {}", + self.data_type() + ) + } + } 16 => { if let DataType::Decimal128(_, s) = self.data_type() { *s @@ -1628,7 +1702,9 @@ impl PrimitiveArray { #[cfg(test)] mod tests { use super::*; - use crate::builder::{Decimal128Builder, Decimal256Builder}; + use crate::builder::{ + Decimal128Builder, Decimal256Builder, Decimal32Builder, Decimal64Builder, + }; use crate::cast::downcast_array; use crate::BooleanArray; use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano}; @@ -2024,7 +2100,7 @@ mod tests { .with_timezone("Asia/Taipei".to_string()); assert_eq!( "PrimitiveArray\n[\n 2018-12-31T08:00:00+08:00,\n 2018-12-31T08:00:00+08:00,\n 1921-01-02T08:00:00+08:00,\n]", - format!("{:?}", arr) + format!("{arr:?}") ); } @@ -2077,7 +2153,7 @@ mod tests { .with_timezone("America/Denver".to_string()); assert_eq!( "PrimitiveArray\n[\n 2022-03-13T01:59:59-07:00,\n 2022-03-13T03:00:00-06:00,\n 2022-11-06T00:59:59-06:00,\n 2022-11-06T01:00:00-06:00,\n]", - format!("{:?}", arr) + format!("{arr:?}") ); } @@ -2238,6 +2314,42 @@ mod tests { let _ = PrimitiveArray::::from(foo.into_data()); } + #[test] + fn test_decimal32() { + let values: Vec<_> = vec![0, 1, -1, i32::MIN, i32::MAX]; + let array: PrimitiveArray = + PrimitiveArray::from_iter(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array: PrimitiveArray = + PrimitiveArray::from_iter_values(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(values.clone()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(array.to_data()); + assert_eq!(array.values(), &values); + } + + #[test] + fn test_decimal64() { + let values: Vec<_> = vec![0, 1, -1, i64::MIN, i64::MAX]; + let array: PrimitiveArray = + PrimitiveArray::from_iter(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array: PrimitiveArray = + PrimitiveArray::from_iter_values(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(values.clone()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(array.to_data()); + assert_eq!(array.values(), &values); + } + #[test] fn test_decimal128() { let values: Vec<_> = vec![0, 1, -1, i128::MIN, i128::MAX]; @@ -2509,6 +2621,74 @@ mod tests { assert!(!array.is_null(2)); } + #[test] + fn test_decimal64_iter() { + let mut builder = Decimal64Builder::with_capacity(30); + let decimal1 = 12345; + builder.append_value(decimal1); + + builder.append_null(); + + let decimal2 = 56789; + builder.append_value(decimal2); + + let array: Decimal64Array = builder.finish().with_precision_and_scale(18, 4).unwrap(); + + let collected: Vec<_> = array.iter().collect(); + assert_eq!(vec![Some(decimal1), None, Some(decimal2)], collected); + } + + #[test] + fn test_from_iter_decimal64array() { + let value1 = 12345; + let value2 = 56789; + + let mut array: Decimal64Array = + vec![Some(value1), None, Some(value2)].into_iter().collect(); + array = array.with_precision_and_scale(18, 4).unwrap(); + assert_eq!(array.len(), 3); + assert_eq!(array.data_type(), &DataType::Decimal64(18, 4)); + assert_eq!(value1, array.value(0)); + assert!(!array.is_null(0)); + assert!(array.is_null(1)); + assert_eq!(value2, array.value(2)); + assert!(!array.is_null(2)); + } + + #[test] + fn test_decimal32_iter() { + let mut builder = Decimal32Builder::with_capacity(30); + let decimal1 = 12345; + builder.append_value(decimal1); + + builder.append_null(); + + let decimal2 = 56789; + builder.append_value(decimal2); + + let array: Decimal32Array = builder.finish().with_precision_and_scale(9, 2).unwrap(); + + let collected: Vec<_> = array.iter().collect(); + assert_eq!(vec![Some(decimal1), None, Some(decimal2)], collected); + } + + #[test] + fn test_from_iter_decimal32array() { + let value1 = 12345; + let value2 = 56789; + + let mut array: Decimal32Array = + vec![Some(value1), None, Some(value2)].into_iter().collect(); + array = array.with_precision_and_scale(9, 2).unwrap(); + assert_eq!(array.len(), 3); + assert_eq!(array.data_type(), &DataType::Decimal32(9, 2)); + assert_eq!(value1, array.value(0)); + assert!(!array.is_null(0)); + assert!(array.is_null(1)); + assert_eq!(value2, array.value(2)); + assert!(!array.is_null(2)); + } + #[test] fn test_unary_opt() { let array = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7]); @@ -2651,7 +2831,7 @@ mod tests { None, ] .into(); - let debug_str = format!("{:?}", array); + let debug_str = format!("{array:?}"); assert_eq!("PrimitiveArray\n[\n Cast error: Failed to convert -1 to temporal for Time32(Second),\n 00:00:00,\n 23:59:59,\n Cast error: Failed to convert 86400 to temporal for Time32(Second),\n Cast error: Failed to convert 86401 to temporal for Time32(Second),\n null,\n]", debug_str ); @@ -2668,7 +2848,7 @@ mod tests { None, ] .into(); - let debug_str = format!("{:?}", array); + let debug_str = format!("{array:?}"); assert_eq!("PrimitiveArray\n[\n Cast error: Failed to convert -1 to temporal for Time32(Millisecond),\n 00:00:00,\n 23:59:59,\n Cast error: Failed to convert 86400000 to temporal for Time32(Millisecond),\n Cast error: Failed to convert 86401000 to temporal for Time32(Millisecond),\n null,\n]", debug_str ); @@ -2685,7 +2865,7 @@ mod tests { None, ] .into(); - let debug_str = format!("{:?}", array); + let debug_str = format!("{array:?}"); assert_eq!( "PrimitiveArray\n[\n Cast error: Failed to convert -1 to temporal for Time64(Nanosecond),\n 00:00:00,\n 23:59:59,\n Cast error: Failed to convert 86400000000000 to temporal for Time64(Nanosecond),\n Cast error: Failed to convert 86401000000000 to temporal for Time64(Nanosecond),\n null,\n]", debug_str @@ -2703,7 +2883,7 @@ mod tests { None, ] .into(); - let debug_str = format!("{:?}", array); + let debug_str = format!("{array:?}"); assert_eq!("PrimitiveArray\n[\n Cast error: Failed to convert -1 to temporal for Time64(Microsecond),\n 00:00:00,\n 23:59:59,\n Cast error: Failed to convert 86400000000 to temporal for Time64(Microsecond),\n Cast error: Failed to convert 86401000000 to temporal for Time64(Microsecond),\n null,\n]", debug_str); } diff --git a/arrow-array/src/array/run_array.rs b/arrow-array/src/array/run_array.rs index b340bf9a9065..05cfa2d17135 100644 --- a/arrow-array/src/array/run_array.rs +++ b/arrow-array/src/array/run_array.rs @@ -662,15 +662,15 @@ where #[cfg(test)] mod tests { + use rand::rng; use rand::seq::SliceRandom; - use rand::thread_rng; use rand::Rng; use super::*; use crate::builder::PrimitiveRunBuilder; use crate::cast::AsArray; use crate::types::{Int8Type, UInt32Type}; - use crate::{Int32Array, StringArray}; + use crate::{Int16Array, Int32Array, StringArray}; fn build_input_array(size: usize) -> Vec> { // The input array is created by shuffling and repeating @@ -691,7 +691,7 @@ mod tests { ]; let mut result: Vec> = Vec::with_capacity(size); let mut ix = 0; - let mut rng = thread_rng(); + let mut rng = rng(); // run length can go up to 8. Cap the max run length for smaller arrays to size / 2. let max_run_length = 8_usize.min(1_usize.max(size / 2)); while result.len() < size { @@ -700,7 +700,7 @@ mod tests { seed.shuffle(&mut rng); } // repeat the items between 1 and 8 times. Cap the length for smaller sized arrays - let num = max_run_length.min(rand::thread_rng().gen_range(1..=max_run_length)); + let num = max_run_length.min(rng.random_range(1..=max_run_length)); for _ in 0..num { result.push(seed[ix]); } @@ -1000,7 +1000,7 @@ mod tests { let mut logical_indices: Vec = (0_u32..(logical_len as u32)).collect(); // add same indices once more logical_indices.append(&mut logical_indices.clone()); - let mut rng = thread_rng(); + let mut rng = rng(); logical_indices.shuffle(&mut rng); let physical_indices = run_array.get_physical_indices(&logical_indices).unwrap(); @@ -1036,7 +1036,7 @@ mod tests { let mut logical_indices: Vec = (0_u32..(slice_len as u32)).collect(); // add same indices once more logical_indices.append(&mut logical_indices.clone()); - let mut rng = thread_rng(); + let mut rng = rng(); logical_indices.shuffle(&mut rng); // test for offset = 0 and slice length = slice_len @@ -1104,4 +1104,69 @@ mod tests { assert_eq!(&n, &expected[offset..offset + length], "{offset} {length}"); } } + + #[test] + fn test_run_array_eq_identical() { + let run_ends1 = Int32Array::from(vec![2, 4, 6]); + let values1 = StringArray::from(vec!["a", "b", "c"]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int32Array::from(vec![2, 4, 6]); + let values2 = StringArray::from(vec!["a", "b", "c"]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + assert_eq!(array1, array2); + } + + #[test] + fn test_run_array_ne_different_run_ends() { + let run_ends1 = Int32Array::from(vec![2, 4, 6]); + let values1 = StringArray::from(vec!["a", "b", "c"]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int32Array::from(vec![1, 4, 6]); + let values2 = StringArray::from(vec!["a", "b", "c"]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + assert_ne!(array1, array2); + } + + #[test] + fn test_run_array_ne_different_values() { + let run_ends1 = Int32Array::from(vec![2, 4, 6]); + let values1 = StringArray::from(vec!["a", "b", "c"]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int32Array::from(vec![2, 4, 6]); + let values2 = StringArray::from(vec!["a", "b", "d"]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + assert_ne!(array1, array2); + } + + #[test] + fn test_run_array_eq_with_nulls() { + let run_ends1 = Int32Array::from(vec![2, 4, 6]); + let values1 = StringArray::from(vec![Some("a"), None, Some("c")]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int32Array::from(vec![2, 4, 6]); + let values2 = StringArray::from(vec![Some("a"), None, Some("c")]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + assert_eq!(array1, array2); + } + + #[test] + fn test_run_array_eq_different_run_end_types() { + let run_ends_i16_1 = Int16Array::from(vec![2_i16, 4, 6]); + let values_i16_1 = StringArray::from(vec!["a", "b", "c"]); + let array_i16_1 = RunArray::::try_new(&run_ends_i16_1, &values_i16_1).unwrap(); + + let run_ends_i16_2 = Int16Array::from(vec![2_i16, 4, 6]); + let values_i16_2 = StringArray::from(vec!["a", "b", "c"]); + let array_i16_2 = RunArray::::try_new(&run_ends_i16_2, &values_i16_2).unwrap(); + + assert_eq!(array_i16_1, array_i16_2); + } } diff --git a/arrow-array/src/array/struct_array.rs b/arrow-array/src/array/struct_array.rs index de6d9c699d22..fbc34ef0c85b 100644 --- a/arrow-array/src/array/struct_array.rs +++ b/arrow-array/src/array/struct_array.rs @@ -91,6 +91,28 @@ impl StructArray { Self::try_new(fields, arrays, nulls).unwrap() } + /// Create a new [`StructArray`] from the provided parts, returning an error on failure + /// + /// The length will be inferred from the length of the child arrays. Returns an error if + /// there are no child arrays. Consider using [`Self::try_new_with_length`] if the length + /// is known to avoid this. + /// + /// # Errors + /// + /// Errors if + /// + /// * `fields.len() == 0` + /// * Any reason that [`Self::try_new_with_length`] would error + pub fn try_new( + fields: Fields, + arrays: Vec, + nulls: Option, + ) -> Result { + let len = arrays.first().map(|x| x.len()).ok_or_else(||ArrowError::InvalidArgumentError("use StructArray::try_new_with_length or StructArray::new_empty_fields to create a struct array with no fields so that the length can be set correctly".to_string()))?; + + Self::try_new_with_length(fields, arrays, nulls, len) + } + /// Create a new [`StructArray`] from the provided parts, returning an error on failure /// /// # Errors @@ -102,10 +124,11 @@ impl StructArray { /// * `arrays[i].len() != arrays[j].len()` /// * `arrays[i].len() != nulls.len()` /// * `!fields[i].is_nullable() && !nulls.contains(arrays[i].nulls())` - pub fn try_new( + pub fn try_new_with_length( fields: Fields, arrays: Vec, nulls: Option, + len: usize, ) -> Result { if fields.len() != arrays.len() { return Err(ArrowError::InvalidArgumentError(format!( @@ -114,7 +137,6 @@ impl StructArray { arrays.len() ))); } - let len = arrays.first().map(|x| x.len()).unwrap_or_default(); if let Some(n) = nulls.as_ref() { if n.len() != len { @@ -146,7 +168,9 @@ impl StructArray { if !f.is_nullable() { if let Some(a) = a.logical_nulls() { - if !nulls.as_ref().map(|n| n.contains(&a)).unwrap_or_default() { + if !nulls.as_ref().map(|n| n.contains(&a)).unwrap_or_default() + && a.null_count() > 0 + { return Err(ArrowError::InvalidArgumentError(format!( "Found unmasked nulls for non-nullable StructArray field {:?}", f.name() @@ -181,6 +205,10 @@ impl StructArray { /// Create a new [`StructArray`] from the provided parts without validation /// + /// The length will be inferred from the length of the child arrays. Panics if there are no + /// child arrays. Consider using [`Self::new_unchecked_with_length`] if the length is known + /// to avoid this. + /// /// # Safety /// /// Safe if [`Self::new`] would not panic with the given arguments @@ -189,7 +217,36 @@ impl StructArray { arrays: Vec, nulls: Option, ) -> Self { - let len = arrays.first().map(|x| x.len()).unwrap_or_default(); + if cfg!(feature = "force_validate") { + return Self::new(fields, arrays, nulls); + } + + let len = arrays.first().map(|x| x.len()).expect( + "cannot use StructArray::new_unchecked if there are no fields, length is unknown", + ); + Self { + len, + data_type: DataType::Struct(fields), + nulls, + fields: arrays, + } + } + + /// Create a new [`StructArray`] from the provided parts without validation + /// + /// # Safety + /// + /// Safe if [`Self::new`] would not panic with the given arguments + pub unsafe fn new_unchecked_with_length( + fields: Fields, + arrays: Vec, + nulls: Option, + len: usize, + ) -> Self { + if cfg!(feature = "force_validate") { + return Self::try_new_with_length(fields, arrays, nulls, len).unwrap(); + } + Self { len, data_type: DataType::Struct(fields), @@ -290,10 +347,19 @@ impl StructArray { impl From for StructArray { fn from(data: ArrayData) -> Self { + let parent_offset = data.offset(); + let parent_len = data.len(); + let fields = data .child_data() .iter() - .map(|cd| make_array(cd.clone())) + .map(|cd| { + if parent_offset != 0 || parent_len != cd.len() { + make_array(cd.slice(parent_offset, parent_len)) + } else { + make_array(cd.clone()) + } + }) .collect(); Self { @@ -412,7 +478,7 @@ impl From> for StructArray { impl std::fmt::Debug for StructArray { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { writeln!(f, "StructArray")?; - writeln!(f, "-- validity: ")?; + writeln!(f, "-- validity:")?; writeln!(f, "[")?; print_long_array(self, f, |_array, _index, f| write!(f, "valid"))?; writeln!(f, "]\n[")?; @@ -519,6 +585,81 @@ mod tests { assert_eq!(0, struct_array.offset()); } + #[test] + fn test_struct_array_from_data_with_offset_and_length() { + // Various ways to make the struct array: + // + // [{x: 2}, {x: 3}, None] + // + // from slicing larger buffers/arrays with offsets and lengths + let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]); + let int_field = Field::new("x", DataType::Int32, false); + let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false])); + let int_data = int_arr.to_data(); + // Case 1: Offset + length, nulls are not sliced + let case1 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()]))) + .len(3) + .offset(1) + .nulls(Some(struct_nulls)) + .add_child_data(int_data.clone()) + .build() + .unwrap(); + + // Case 2: Offset + length, nulls are sliced + let struct_nulls = + NullBuffer::new(BooleanBuffer::from(vec![true, true, true, false, true]).slice(1, 3)); + let case2 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()]))) + .len(3) + .offset(1) + .nulls(Some(struct_nulls.clone())) + .add_child_data(int_data.clone()) + .build() + .unwrap(); + + // Case 3: struct length is smaller than child length but no offset + let offset_int_data = int_data.slice(1, 4); + let case3 = ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()]))) + .len(3) + .nulls(Some(struct_nulls)) + .add_child_data(offset_int_data) + .build() + .unwrap(); + + let expected = StructArray::new( + Fields::from(vec![int_field.clone()]), + vec![Arc::new(int_arr)], + Some(NullBuffer::new(BooleanBuffer::from(vec![ + true, true, true, false, true, + ]))), + ) + .slice(1, 3); + + for case in [case1, case2, case3] { + let struct_arr_from_data = StructArray::from(case); + assert_eq!(struct_arr_from_data, expected); + assert_eq!(struct_arr_from_data.column(0), expected.column(0)); + } + } + + #[test] + #[should_panic(expected = "assertion failed: (offset + length) <= self.len()")] + fn test_struct_array_from_data_with_offset_and_length_error() { + let int_arr = Int32Array::from(vec![1, 2, 3, 4, 5]); + let int_field = Field::new("x", DataType::Int32, false); + let struct_nulls = NullBuffer::new(BooleanBuffer::from(vec![true, true, false])); + let int_data = int_arr.to_data(); + // If parent offset is 3 and len is 3 then child must have 6 items + let struct_data = + ArrayData::builder(DataType::Struct(Fields::from(vec![int_field.clone()]))) + .len(3) + .offset(3) + .nulls(Some(struct_nulls)) + .add_child_data(int_data) + .build() + .unwrap(); + let _ = StructArray::from(struct_data); + } + /// validates that struct can be accessed using `column_name` as index i.e. `struct_array["column_name"]`. #[test] fn test_struct_array_index_access() { @@ -729,9 +870,38 @@ mod tests { } #[test] + #[should_panic(expected = "use StructArray::try_new_with_length")] fn test_struct_array_from_empty() { - let sa = StructArray::from(vec![]); - assert!(sa.is_empty()) + // This can't work because we don't know how many rows the array should have. Previously we inferred 0 but + // that often led to bugs. + let _ = StructArray::from(vec![]); + } + + #[test] + fn test_empty_struct_array() { + assert!(StructArray::try_new(Fields::empty(), vec![], None).is_err()); + + let arr = StructArray::new_empty_fields(10, None); + assert_eq!(arr.len(), 10); + assert_eq!(arr.null_count(), 0); + assert_eq!(arr.num_columns(), 0); + + let arr2 = StructArray::try_new_with_length(Fields::empty(), vec![], None, 10).unwrap(); + assert_eq!(arr2.len(), 10); + + let arr = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10))); + assert_eq!(arr.len(), 10); + assert_eq!(arr.null_count(), 10); + assert_eq!(arr.num_columns(), 0); + + let arr2 = StructArray::try_new_with_length( + Fields::empty(), + vec![], + Some(NullBuffer::new_null(10)), + 10, + ) + .unwrap(); + assert_eq!(arr2.len(), 10); } #[test] @@ -752,6 +922,25 @@ mod tests { (0..30).map(|i| i % 2 == 0).collect::>(), ))), ); - assert_eq!(format!("{arr:?}"), "StructArray\n-- validity: \n[\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n ...10 elements...,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n]\n[\n-- child 0: \"c\" (Int32)\nPrimitiveArray\n[\n 0,\n 1,\n 2,\n 3,\n 4,\n 5,\n 6,\n 7,\n 8,\n 9,\n ...10 elements...,\n 20,\n 21,\n 22,\n 23,\n 24,\n 25,\n 26,\n 27,\n 28,\n 29,\n]\n]") + assert_eq!(format!("{arr:?}"), "StructArray\n-- validity:\n[\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n ...10 elements...,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n valid,\n null,\n]\n[\n-- child 0: \"c\" (Int32)\nPrimitiveArray\n[\n 0,\n 1,\n 2,\n 3,\n 4,\n 5,\n 6,\n 7,\n 8,\n 9,\n ...10 elements...,\n 20,\n 21,\n 22,\n 23,\n 24,\n 25,\n 26,\n 27,\n 28,\n 29,\n]\n]") + } + + #[test] + fn test_struct_array_logical_nulls() { + // Field is non-nullable + let field = Field::new("a", DataType::Int32, false); + let values = vec![1, 2, 3]; + // Create a NullBuffer with all bits set to valid (true) + let nulls = NullBuffer::from(vec![true, true, true]); + let array = Int32Array::new(values.into(), Some(nulls)); + let child = Arc::new(array) as ArrayRef; + assert!(child.logical_nulls().is_some()); + assert_eq!(child.logical_nulls().unwrap().null_count(), 0); + + let fields = Fields::from(vec![field]); + let arrays = vec![child]; + let nulls = None; + + StructArray::try_new(fields, arrays, nulls).expect("should not error"); } } diff --git a/arrow-array/src/array/union_array.rs b/arrow-array/src/array/union_array.rs index b442395b4978..1350cae3a38b 100644 --- a/arrow-array/src/array/union_array.rs +++ b/arrow-array/src/array/union_array.rs @@ -781,13 +781,18 @@ impl Array for UnionArray { }; if fields.len() <= 1 { - return self - .fields - .iter() - .flatten() - .map(Array::logical_nulls) - .next() - .flatten(); + return self.fields.iter().find_map(|field_opt| { + field_opt + .as_ref() + .and_then(|field| field.logical_nulls()) + .map(|logical_nulls| { + if self.is_dense() { + self.gather_nulls(vec![(0, logical_nulls)]).into() + } else { + logical_nulls + } + }) + }); } let logical_nulls = self.fields_logical_nulls(); @@ -940,7 +945,7 @@ impl std::fmt::Debug for UnionArray { if let Some(offsets) = &self.offsets { writeln!(f, "-- offsets buffer:")?; - writeln!(f, "{:?}", offsets)?; + writeln!(f, "{offsets:?}")?; } let fields = match self.data_type() { @@ -994,7 +999,7 @@ fn selection_mask(type_ids_chunk: &[i8], type_id: i8) -> u64 { .copied() .enumerate() .fold(0, |packed, (bit_idx, v)| { - packed | ((v == type_id) as u64) << bit_idx + packed | (((v == type_id) as u64) << bit_idx) }) } @@ -1074,6 +1079,30 @@ mod tests { } } + #[test] + fn slice_union_array_single_field() { + // Dense Union + // [1, null, 3, null, 4] + let union_array = { + let mut builder = UnionBuilder::new_dense(); + builder.append::("a", 1).unwrap(); + builder.append_null::("a").unwrap(); + builder.append::("a", 3).unwrap(); + builder.append_null::("a").unwrap(); + builder.append::("a", 4).unwrap(); + builder.build().unwrap() + }; + + // [null, 3, null] + let union_slice = union_array.slice(1, 3); + let logical_nulls = union_slice.logical_nulls().unwrap(); + + assert_eq!(logical_nulls.len(), 3); + assert!(logical_nulls.is_null(0)); + assert!(logical_nulls.is_valid(1)); + assert!(logical_nulls.is_null(2)); + } + #[test] #[cfg_attr(miri, ignore)] fn test_dense_i32_large() { diff --git a/arrow-array/src/builder/boolean_builder.rs b/arrow-array/src/builder/boolean_builder.rs index 60ed86ce80b4..a0bd5745d21d 100644 --- a/arrow-array/src/builder/boolean_builder.rs +++ b/arrow-array/src/builder/boolean_builder.rs @@ -16,7 +16,7 @@ // under the License. use crate::builder::{ArrayBuilder, BooleanBufferBuilder}; -use crate::{ArrayRef, BooleanArray}; +use crate::{Array, ArrayRef, BooleanArray}; use arrow_buffer::Buffer; use arrow_buffer::NullBufferBuilder; use arrow_data::ArrayData; @@ -146,6 +146,18 @@ impl BooleanBuilder { } } + /// Appends array values and null to this builder as is + /// (this means that underlying null values are copied as is). + #[inline] + pub fn append_array(&mut self, array: &BooleanArray) { + self.values_builder.append_buffer(array.values()); + if let Some(null_buffer) = array.nulls() { + self.null_buffer_builder.append_buffer(null_buffer); + } else { + self.null_buffer_builder.append_n_non_nulls(array.len()); + } + } + /// Builds the [BooleanArray] and reset this builder. pub fn finish(&mut self) -> BooleanArray { let len = self.len(); @@ -232,6 +244,7 @@ impl Extend> for BooleanBuilder { mod tests { use super::*; use crate::Array; + use arrow_buffer::{BooleanBuffer, NullBuffer}; #[test] fn test_boolean_array_builder() { @@ -346,4 +359,50 @@ mod tests { let values = array.iter().map(|x| x.unwrap()).collect::>(); assert_eq!(&values, &[true, true, true, false, false]) } + + #[test] + fn test_append_array() { + let input = vec![ + Some(true), + None, + Some(true), + None, + Some(false), + None, + None, + None, + Some(false), + Some(false), + Some(false), + Some(true), + Some(false), + ]; + let arr1 = BooleanArray::from(input[..5].to_vec()); + let arr2 = BooleanArray::from(input[5..8].to_vec()); + let arr3 = BooleanArray::from(input[8..].to_vec()); + + let mut builder = BooleanBuilder::new(); + builder.append_array(&arr1); + builder.append_array(&arr2); + builder.append_array(&arr3); + let actual = builder.finish(); + let expected = BooleanArray::from(input); + + assert_eq!(actual, expected); + } + + #[test] + fn test_append_array_add_underlying_null_values() { + let array = BooleanArray::new( + BooleanBuffer::from(vec![true, false, true, false]), + Some(NullBuffer::from(&[true, true, false, false])), + ); + + let mut builder = BooleanBuilder::new(); + builder.append_array(&array); + let actual = builder.finish(); + + assert_eq!(actual, array); + assert_eq!(actual.values(), array.values()) + } } diff --git a/arrow-array/src/builder/buffer_builder.rs b/arrow-array/src/builder/buffer_builder.rs index ab67669febb8..5975654667ce 100644 --- a/arrow-array/src/builder/buffer_builder.rs +++ b/arrow-array/src/builder/buffer_builder.rs @@ -16,6 +16,8 @@ // under the License. pub use arrow_buffer::BufferBuilder; +pub use arrow_buffer::OffsetBufferBuilder; + use half::f16; use crate::types::*; @@ -43,6 +45,10 @@ pub type Float32BufferBuilder = BufferBuilder; /// Buffer builder for 64-bit floating point type. pub type Float64BufferBuilder = BufferBuilder; +/// Buffer builder for 32-bit decimal type. +pub type Decimal32BufferBuilder = BufferBuilder<::Native>; +/// Buffer builder for 64-bit decimal type. +pub type Decimal64BufferBuilder = BufferBuilder<::Native>; /// Buffer builder for 128-bit decimal type. pub type Decimal128BufferBuilder = BufferBuilder<::Native>; /// Buffer builder for 256-bit decimal type. diff --git a/arrow-array/src/builder/fixed_size_binary_builder.rs b/arrow-array/src/builder/fixed_size_binary_builder.rs index 65072a09f603..b5f268917c92 100644 --- a/arrow-array/src/builder/fixed_size_binary_builder.rs +++ b/arrow-array/src/builder/fixed_size_binary_builder.rs @@ -93,6 +93,19 @@ impl FixedSizeBinaryBuilder { self.null_buffer_builder.append_null(); } + /// Appends `n` `null`s into the builder. + #[inline] + pub fn append_nulls(&mut self, n: usize) { + self.values_builder + .append_slice(&vec![0u8; self.value_length as usize * n][..]); + self.null_buffer_builder.append_n_nulls(n); + } + + /// Returns the current values buffer as a slice + pub fn values_slice(&self) -> &[u8] { + self.values_builder.as_slice() + } + /// Builds the [`FixedSizeBinaryArray`] and reset this builder. pub fn finish(&mut self) -> FixedSizeBinaryArray { let array_length = self.len(); @@ -164,17 +177,22 @@ mod tests { fn test_fixed_size_binary_builder() { let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 5); - // [b"hello", null, "arrow"] + // [b"hello", null, "arrow", null, null, "world"] builder.append_value(b"hello").unwrap(); builder.append_null(); builder.append_value(b"arrow").unwrap(); + builder.append_nulls(2); + builder.append_value(b"world").unwrap(); let array: FixedSizeBinaryArray = builder.finish(); assert_eq!(&DataType::FixedSizeBinary(5), array.data_type()); - assert_eq!(3, array.len()); - assert_eq!(1, array.null_count()); + assert_eq!(6, array.len()); + assert_eq!(3, array.null_count()); assert_eq!(10, array.value_offset(2)); + assert_eq!(15, array.value_offset(3)); assert_eq!(5, array.value_length()); + assert!(array.is_null(3)); + assert!(array.is_null(4)); } #[test] diff --git a/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs b/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs new file mode 100644 index 000000000000..21e842723b4a --- /dev/null +++ b/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs @@ -0,0 +1,511 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::builder::{ArrayBuilder, FixedSizeBinaryBuilder, PrimitiveBuilder}; +use crate::types::ArrowDictionaryKeyType; +use crate::{Array, ArrayRef, DictionaryArray, PrimitiveArray}; +use arrow_buffer::ArrowNativeType; +use arrow_schema::DataType::FixedSizeBinary; +use arrow_schema::{ArrowError, DataType}; +use hashbrown::HashTable; +use num::NumCast; +use std::any::Any; +use std::sync::Arc; + +/// Builder for [`DictionaryArray`] of [`FixedSizeBinaryArray`] +/// +/// The output array has a dictionary of unique, fixed-size binary values. The +/// builder handles deduplication. +/// +/// # Example +/// ``` +/// # use arrow_array::builder::{FixedSizeBinaryDictionaryBuilder}; +/// # use arrow_array::array::{Array, FixedSizeBinaryArray}; +/// # use arrow_array::DictionaryArray; +/// # use arrow_array::types::Int8Type; +/// // Build 3 byte FixedBinaryArrays +/// let byte_width = 3; +/// let mut builder = FixedSizeBinaryDictionaryBuilder::::new(3); +/// builder.append("abc").unwrap(); +/// builder.append_null(); +/// builder.append(b"def").unwrap(); +/// builder.append(b"def").unwrap(); // duplicate value +/// // Result is a Dictionary Array +/// let array = builder.finish(); +/// let dict_array = array.as_any().downcast_ref::>().unwrap(); +/// // The array represents "abc", null, "def", "def" +/// assert_eq!(array.keys().len(), 4); +/// // but there are only 2 unique values +/// assert_eq!(array.values().len(), 2); +/// let values = dict_array.values().as_any().downcast_ref::().unwrap(); +/// assert_eq!(values.value(0), "abc".as_bytes()); +/// assert_eq!(values.value(1), "def".as_bytes()); +/// ``` +/// +/// [`FixedSizeBinaryArray`]: crate::FixedSizeBinaryArray +#[derive(Debug)] +pub struct FixedSizeBinaryDictionaryBuilder +where + K: ArrowDictionaryKeyType, +{ + state: ahash::RandomState, + dedup: HashTable, + + keys_builder: PrimitiveBuilder, + values_builder: FixedSizeBinaryBuilder, + byte_width: i32, +} + +impl FixedSizeBinaryDictionaryBuilder +where + K: ArrowDictionaryKeyType, +{ + /// Creates a new `FixedSizeBinaryDictionaryBuilder` + pub fn new(byte_width: i32) -> Self { + let keys_builder = PrimitiveBuilder::new(); + let values_builder = FixedSizeBinaryBuilder::new(byte_width); + Self { + state: Default::default(), + dedup: HashTable::with_capacity(keys_builder.capacity()), + keys_builder, + values_builder, + byte_width, + } + } + + /// Creates a new `FixedSizeBinaryDictionaryBuilder` with the provided capacities + /// + /// `keys_capacity`: the number of keys, i.e. length of array to build + /// `value_capacity`: the number of distinct dictionary values, i.e. size of dictionary + /// `byte_width`: the byte width for individual values in the values array + pub fn with_capacity(keys_capacity: usize, value_capacity: usize, byte_width: i32) -> Self { + Self { + state: Default::default(), + dedup: Default::default(), + keys_builder: PrimitiveBuilder::with_capacity(keys_capacity), + values_builder: FixedSizeBinaryBuilder::with_capacity(value_capacity, byte_width), + byte_width, + } + } + + /// Creates a new `FixedSizeBinaryDictionaryBuilder` from the existing builder with the same + /// keys and values, but with a new data type for the keys. + /// + /// # Example + /// ``` + /// # use arrow_array::builder::FixedSizeBinaryDictionaryBuilder; + /// # use arrow_array::types::{UInt8Type, UInt16Type, UInt64Type}; + /// # use arrow_array::UInt16Array; + /// # use arrow_schema::ArrowError; + /// + /// let mut u8_keyed_builder = FixedSizeBinaryDictionaryBuilder::::new(2); + /// // appending too many values causes the dictionary to overflow + /// for i in 0..=255 { + /// u8_keyed_builder.append_value(vec![0, i]); + /// } + /// let result = u8_keyed_builder.append(vec![1, 0]); + /// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{}))); + /// + /// // we need to upgrade to a larger key type + /// let mut u16_keyed_builder = FixedSizeBinaryDictionaryBuilder::::try_new_from_builder(u8_keyed_builder).unwrap(); + /// let dictionary_array = u16_keyed_builder.finish(); + /// let keys = dictionary_array.keys(); + /// + /// assert_eq!(keys, &UInt16Array::from_iter(0..256)); + /// ``` + pub fn try_new_from_builder( + mut source: FixedSizeBinaryDictionaryBuilder, + ) -> Result + where + K::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast, + { + let state = source.state; + let dedup = source.dedup; + let values_builder = source.values_builder; + let byte_width = source.byte_width; + + let source_keys = source.keys_builder.finish(); + let new_keys: PrimitiveArray = source_keys.try_unary(|value| { + num::cast::cast::(value).ok_or_else(|| { + ArrowError::CastError(format!( + "Can't cast dictionary keys from source type {:?} to type {:?}", + K2::DATA_TYPE, + K::DATA_TYPE + )) + }) + })?; + + // drop source key here because currently source_keys and new_keys are holding reference to + // the same underlying null_buffer. Below we want to call new_keys.into_builder() it must + // be the only reference holder. + drop(source_keys); + + Ok(Self { + state, + dedup, + keys_builder: new_keys + .into_builder() + .expect("underlying buffer has no references"), + values_builder, + byte_width, + }) + } +} + +impl ArrayBuilder for FixedSizeBinaryDictionaryBuilder +where + K: ArrowDictionaryKeyType, +{ + /// Returns the builder as an non-mutable `Any` reference. + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns the builder as an mutable `Any` reference. + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + /// Returns the boxed builder as a box of `Any`. + fn into_box_any(self: Box) -> Box { + self + } + + /// Returns the number of array slots in the builder + fn len(&self) -> usize { + self.keys_builder.len() + } + + /// Builds the array and reset this builder. + fn finish(&mut self) -> ArrayRef { + Arc::new(self.finish()) + } + + /// Builds the array without resetting the builder. + fn finish_cloned(&self) -> ArrayRef { + Arc::new(self.finish_cloned()) + } +} + +impl FixedSizeBinaryDictionaryBuilder +where + K: ArrowDictionaryKeyType, +{ + fn get_or_insert_key(&mut self, value: impl AsRef<[u8]>) -> Result { + let value_bytes: &[u8] = value.as_ref(); + + let state = &self.state; + let storage = &mut self.values_builder; + let hash = state.hash_one(value_bytes); + + let idx = *self + .dedup + .entry( + hash, + |idx| value_bytes == get_bytes(storage, self.byte_width, *idx), + |idx| state.hash_one(get_bytes(storage, self.byte_width, *idx)), + ) + .or_insert_with(|| { + let idx = storage.len(); + let _ = storage.append_value(value); + idx + }) + .get(); + + let key = K::Native::from_usize(idx).ok_or(ArrowError::DictionaryKeyOverflowError)?; + + Ok(key) + } + + /// Append a value to the array. Return an existing index + /// if already present in the values array or a new index if the + /// value is appended to the values array. + /// + /// Returns an error if the new index would overflow the key type. + pub fn append(&mut self, value: impl AsRef<[u8]>) -> Result { + if self.byte_width != value.as_ref().len() as i32 { + Err(ArrowError::InvalidArgumentError(format!( + "Invalid input length passed to FixedSizeBinaryBuilder. Expected {} got {}", + self.byte_width, + value.as_ref().len() + ))) + } else { + let key = self.get_or_insert_key(value)?; + self.keys_builder.append_value(key); + Ok(key) + } + } + + /// Appends a null slot into the builder + #[inline] + pub fn append_null(&mut self) { + self.keys_builder.append_null() + } + + /// Appends `n` `null`s into the builder. + #[inline] + pub fn append_nulls(&mut self, n: usize) { + self.keys_builder.append_nulls(n); + } + + /// Infallibly append a value to this builder + /// + /// # Panics + /// + /// Panics if the resulting length of the dictionary values array would exceed `T::Native::MAX` + pub fn append_value(&mut self, value: impl AsRef<[u8]>) { + self.append(value).expect("dictionary key overflow"); + } + + /// Builds the `DictionaryArray` and reset this builder. + pub fn finish(&mut self) -> DictionaryArray { + self.dedup.clear(); + let values = self.values_builder.finish(); + let keys = self.keys_builder.finish(); + + let data_type = DataType::Dictionary( + Box::new(K::DATA_TYPE), + Box::new(FixedSizeBinary(self.byte_width)), + ); + + let builder = keys + .into_data() + .into_builder() + .data_type(data_type) + .child_data(vec![values.into_data()]); + + DictionaryArray::from(unsafe { builder.build_unchecked() }) + } + + /// Builds the `DictionaryArray` without resetting the builder. + pub fn finish_cloned(&self) -> DictionaryArray { + let values = self.values_builder.finish_cloned(); + let keys = self.keys_builder.finish_cloned(); + + let data_type = DataType::Dictionary( + Box::new(K::DATA_TYPE), + Box::new(FixedSizeBinary(self.byte_width)), + ); + + let builder = keys + .into_data() + .into_builder() + .data_type(data_type) + .child_data(vec![values.into_data()]); + + DictionaryArray::from(unsafe { builder.build_unchecked() }) + } +} + +fn get_bytes(values: &FixedSizeBinaryBuilder, byte_width: i32, idx: usize) -> &[u8] { + let values = values.values_slice(); + let start = idx * byte_width.as_usize(); + let end = idx * byte_width.as_usize() + byte_width.as_usize(); + &values[start..end] +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::types::{Int16Type, Int32Type, Int8Type, UInt16Type, UInt8Type}; + use crate::{ArrowPrimitiveType, FixedSizeBinaryArray, Int8Array}; + + #[test] + fn test_fixed_size_dictionary_builder() { + let values = ["abc", "def"]; + + let mut b = FixedSizeBinaryDictionaryBuilder::::new(3); + assert_eq!(b.append(values[0]).unwrap(), 0); + b.append_null(); + assert_eq!(b.append(values[1]).unwrap(), 1); + assert_eq!(b.append(values[1]).unwrap(), 1); + assert_eq!(b.append(values[0]).unwrap(), 0); + b.append_nulls(2); + assert_eq!(b.append(values[0]).unwrap(), 0); + let array = b.finish(); + + assert_eq!( + array.keys(), + &Int8Array::from(vec![ + Some(0), + None, + Some(1), + Some(1), + Some(0), + None, + None, + Some(0) + ]), + ); + + // Values are polymorphic and so require a downcast. + let ava = array + .values() + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(ava.value(0), values[0].as_bytes()); + assert_eq!(ava.value(1), values[1].as_bytes()); + } + + #[test] + fn test_fixed_size_dictionary_builder_wrong_size() { + let mut b = FixedSizeBinaryDictionaryBuilder::::new(3); + let err = b.append(b"too long").unwrap_err().to_string(); + assert_eq!(err, "Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 8"); + let err = b.append("").unwrap_err().to_string(); + assert_eq!(err, "Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 0"); + } + + #[test] + fn test_fixed_size_dictionary_builder_finish_cloned() { + let values = ["abc", "def", "ghi"]; + + let mut builder = FixedSizeBinaryDictionaryBuilder::::new(3); + + builder.append(values[0]).unwrap(); + builder.append_null(); + builder.append(values[1]).unwrap(); + builder.append(values[1]).unwrap(); + builder.append(values[0]).unwrap(); + let mut array = builder.finish_cloned(); + + assert_eq!( + array.keys(), + &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)]) + ); + + // Values are polymorphic and so require a downcast. + let ava = array + .values() + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(ava.value(0), values[0].as_bytes()); + assert_eq!(ava.value(1), values[1].as_bytes()); + + builder.append(values[0]).unwrap(); + builder.append(values[2]).unwrap(); + builder.append(values[1]).unwrap(); + + array = builder.finish(); + + assert_eq!( + array.keys(), + &Int8Array::from(vec![ + Some(0), + None, + Some(1), + Some(1), + Some(0), + Some(0), + Some(2), + Some(1) + ]) + ); + + // Values are polymorphic and so require a downcast. + let ava2 = array + .values() + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(ava2.value(0), values[0].as_bytes()); + assert_eq!(ava2.value(1), values[1].as_bytes()); + assert_eq!(ava2.value(2), values[2].as_bytes()); + } + + fn _test_try_new_from_builder_generic_for_key_types(values: Vec<[u8; 3]>) + where + K1: ArrowDictionaryKeyType, + K1::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast + From, + { + let mut source = FixedSizeBinaryDictionaryBuilder::::new(3); + source.append_value(values[0]); + source.append_null(); + source.append_value(values[1]); + source.append_value(values[2]); + + let mut result = + FixedSizeBinaryDictionaryBuilder::::try_new_from_builder(source).unwrap(); + let array = result.finish(); + + let mut expected_keys_builder = PrimitiveBuilder::::new(); + expected_keys_builder + .append_value(<::Native as From>::from(0u8)); + expected_keys_builder.append_null(); + expected_keys_builder + .append_value(<::Native as From>::from(1u8)); + expected_keys_builder + .append_value(<::Native as From>::from(2u8)); + let expected_keys = expected_keys_builder.finish(); + assert_eq!(array.keys(), &expected_keys); + + let av = array.values(); + let ava = av.as_any().downcast_ref::().unwrap(); + assert_eq!(ava.value(0), values[0]); + assert_eq!(ava.value(1), values[1]); + assert_eq!(ava.value(2), values[2]); + } + + #[test] + fn test_try_new_from_builder() { + let values = vec![[1, 2, 3], [5, 6, 7], [6, 7, 8]]; + // test cast to bigger size unsigned + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test cast going to smaller size unsigned + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test cast going to bigger size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test cast going to smaller size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test going from signed to signed for different size changes + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + } + + #[test] + fn test_try_new_from_builder_cast_fails() { + let mut source_builder = FixedSizeBinaryDictionaryBuilder::::new(2); + for i in 0u16..257u16 { + source_builder.append_value(vec![(i >> 8) as u8, i as u8]); + } + + // there should be too many values that we can't downcast to the underlying type + // we have keys that wouldn't fit into UInt8Type + let result = + FixedSizeBinaryDictionaryBuilder::::try_new_from_builder(source_builder); + assert!(result.is_err()); + if let Err(e) = result { + assert!(matches!(e, ArrowError::CastError(_))); + assert_eq!( + e.to_string(), + "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8" + ); + } + } +} diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs index e2be96615b61..91ac2a483ef4 100644 --- a/arrow-array/src/builder/generic_bytes_builder.rs +++ b/arrow-array/src/builder/generic_bytes_builder.rs @@ -17,7 +17,7 @@ use crate::builder::{ArrayBuilder, BufferBuilder, UInt8BufferBuilder}; use crate::types::{ByteArrayType, GenericBinaryType, GenericStringType}; -use crate::{ArrayRef, GenericByteArray, OffsetSizeTrait}; +use crate::{Array, ArrayRef, GenericByteArray, OffsetSizeTrait}; use arrow_buffer::NullBufferBuilder; use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer}; use arrow_data::ArrayDataBuilder; @@ -129,6 +129,56 @@ impl GenericByteBuilder { self.offsets_builder.append(self.next_offset()); } + /// Appends `n` `null`s into the builder. + #[inline] + pub fn append_nulls(&mut self, n: usize) { + self.null_buffer_builder.append_n_nulls(n); + let next_offset = self.next_offset(); + self.offsets_builder.append_n(n, next_offset); + } + + /// Appends array values and null to this builder as is + /// (this means that underlying null values are copied as is). + #[inline] + pub fn append_array(&mut self, array: &GenericByteArray) { + if array.len() == 0 { + return; + } + + let offsets = array.offsets(); + + // If the offsets are contiguous, we can append them directly avoiding the need to align + // for example, when the first appended array is not sliced (starts at offset 0) + if self.next_offset() == offsets[0] { + self.offsets_builder.append_slice(&offsets[1..]); + } else { + // Shifting all the offsets + let shift: T::Offset = self.next_offset() - offsets[0]; + + // Creating intermediate offsets instead of pushing each offset is faster + // (even if we make MutableBuffer to avoid updating length on each push + // and reserve the necessary capacity, it's still slower) + let mut intermediate = Vec::with_capacity(offsets.len() - 1); + + for &offset in &offsets[1..] { + intermediate.push(offset + shift) + } + + self.offsets_builder.append_slice(&intermediate); + } + + // Append underlying values, starting from the first offset and ending at the last offset + self.value_builder.append_slice( + &array.values().as_slice()[offsets[0].as_usize()..offsets[array.len()].as_usize()], + ); + + if let Some(null_buffer) = array.nulls() { + self.null_buffer_builder.append_buffer(null_buffer); + } else { + self.null_buffer_builder.append_n_non_nulls(array.len()); + } + } + /// Builds the [`GenericByteArray`] and reset this builder. pub fn finish(&mut self) -> GenericByteArray { let array_type = T::DATA_TYPE; @@ -358,6 +408,7 @@ mod tests { use super::*; use crate::array::Array; use crate::GenericStringArray; + use arrow_buffer::NullBuffer; use std::fmt::Write as _; use std::io::Write as _; @@ -396,15 +447,18 @@ mod tests { builder.append_null(); builder.append_null(); builder.append_null(); - assert_eq!(3, builder.len()); + builder.append_nulls(2); + assert_eq!(5, builder.len()); assert!(!builder.is_empty()); let array = builder.finish(); - assert_eq!(3, array.null_count()); - assert_eq!(3, array.len()); + assert_eq!(5, array.null_count()); + assert_eq!(5, array.len()); assert!(array.is_null(0)); assert!(array.is_null(1)); assert!(array.is_null(2)); + assert!(array.is_null(3)); + assert!(array.is_null(4)); } #[test] @@ -432,16 +486,23 @@ mod tests { builder.append_null(); builder.append_value(b"arrow"); builder.append_value(b""); + builder.append_nulls(2); + builder.append_value(b"hi"); let array = builder.finish(); - assert_eq!(4, array.len()); - assert_eq!(1, array.null_count()); + assert_eq!(7, array.len()); + assert_eq!(3, array.null_count()); assert_eq!(b"parquet", array.value(0)); assert!(array.is_null(1)); + assert!(array.is_null(4)); + assert!(array.is_null(5)); assert_eq!(b"arrow", array.value(2)); assert_eq!(b"", array.value(1)); + assert_eq!(b"hi", array.value(6)); + assert_eq!(O::zero(), array.value_offsets()[0]); assert_eq!(O::from_usize(7).unwrap(), array.value_offsets()[2]); + assert_eq!(O::from_usize(14).unwrap(), array.value_offsets()[7]); assert_eq!(O::from_usize(5).unwrap(), array.value_length(2)); } @@ -466,7 +527,9 @@ mod tests { builder.append_option(Some("rust")); builder.append_option(None::<&str>); builder.append_option(None::); - assert_eq!(7, builder.len()); + builder.append_nulls(2); + builder.append_value("parquet"); + assert_eq!(10, builder.len()); assert_eq!( GenericStringArray::::from(vec![ @@ -476,7 +539,10 @@ mod tests { None, Some("rust"), None, - None + None, + None, + None, + Some("parquet") ]), builder.finish() ); @@ -593,4 +659,178 @@ mod tests { &["foo".as_bytes(), "bar\n".as_bytes(), "fizbuz".as_bytes()] ) } + + #[test] + fn test_append_array_without_nulls() { + let input = vec![ + "hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing", "well", + "thank", "you", "for", "asking", + ]; + let arr1 = GenericStringArray::::from(input[..3].to_vec()); + let arr2 = GenericStringArray::::from(input[3..7].to_vec()); + let arr3 = GenericStringArray::::from(input[7..].to_vec()); + + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&arr1); + builder.append_array(&arr2); + builder.append_array(&arr3); + + let actual = builder.finish(); + let expected = GenericStringArray::::from(input); + + assert_eq!(actual, expected); + } + + #[test] + fn test_append_array_with_nulls() { + let input = vec![ + Some("hello"), + None, + Some("how"), + None, + None, + None, + None, + Some("I"), + Some("am"), + Some("doing"), + Some("well"), + ]; + let arr1 = GenericStringArray::::from(input[..3].to_vec()); + let arr2 = GenericStringArray::::from(input[3..7].to_vec()); + let arr3 = GenericStringArray::::from(input[7..].to_vec()); + + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&arr1); + builder.append_array(&arr2); + builder.append_array(&arr3); + + let actual = builder.finish(); + let expected = GenericStringArray::::from(input); + + assert_eq!(actual, expected); + } + + #[test] + fn test_append_empty_array() { + let arr = GenericStringArray::::from(Vec::<&str>::new()); + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&arr); + let result = builder.finish(); + assert_eq!(result.len(), 0); + } + + #[test] + fn test_append_array_with_offset_not_starting_at_0() { + let input = vec![ + Some("hello"), + None, + Some("how"), + None, + None, + None, + None, + Some("I"), + Some("am"), + Some("doing"), + Some("well"), + ]; + let full_array = GenericStringArray::::from(input); + let sliced = full_array.slice(1, 4); + + assert_ne!(sliced.offsets()[0].as_usize(), 0); + assert_ne!(sliced.offsets().last(), full_array.offsets().last()); + + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&sliced); + let actual = builder.finish(); + + let expected = GenericStringArray::::from(vec![None, Some("how"), None, None]); + + assert_eq!(actual, expected); + } + + #[test] + fn test_append_underlying_null_values_added_as_is() { + let input_1_array_with_nulls = { + let input = vec![ + "hello", "world", "how", "are", "you", "doing", "today", "I", "am", + ]; + let (offsets, buffer, _) = GenericStringArray::::from(input).into_parts(); + + GenericStringArray::::new( + offsets, + buffer, + Some(NullBuffer::from(&[ + true, false, true, false, false, true, true, true, false, + ])), + ) + }; + let input_2_array_with_nulls = { + let input = vec!["doing", "well", "thank", "you", "for", "asking"]; + let (offsets, buffer, _) = GenericStringArray::::from(input).into_parts(); + + GenericStringArray::::new( + offsets, + buffer, + Some(NullBuffer::from(&[false, false, true, false, true, true])), + ) + }; + + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&input_1_array_with_nulls); + builder.append_array(&input_2_array_with_nulls); + + let actual = builder.finish(); + let expected = GenericStringArray::::from(vec![ + Some("hello"), + None, // world + Some("how"), + None, // are + None, // you + Some("doing"), + Some("today"), + Some("I"), + None, // am + None, // doing + None, // well + Some("thank"), + None, // "you", + Some("for"), + Some("asking"), + ]); + + assert_eq!(actual, expected); + + let expected_underlying_buffer = Buffer::from( + [ + "hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing", + "well", "thank", "you", "for", "asking", + ] + .join("") + .as_bytes(), + ); + assert_eq!(actual.values(), &expected_underlying_buffer); + } + + #[test] + fn append_array_with_continues_indices() { + let input = vec![ + "hello", "world", "how", "are", "you", "doing", "today", "I", "am", "doing", "well", + "thank", "you", "for", "asking", + ]; + let full_array = GenericStringArray::::from(input); + let slice1 = full_array.slice(0, 3); + let slice2 = full_array.slice(3, 4); + let slice3 = full_array.slice(7, full_array.len() - 7); + + let mut builder = GenericStringBuilder::::new(); + builder.append_array(&slice1); + builder.append_array(&slice2); + builder.append_array(&slice3); + + let actual = builder.finish(); + + assert_eq!(actual, full_array); + } } diff --git a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs index ead151d5ceea..a2ed91ac905d 100644 --- a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs +++ b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs @@ -17,10 +17,13 @@ use crate::builder::{ArrayBuilder, GenericByteBuilder, PrimitiveBuilder}; use crate::types::{ArrowDictionaryKeyType, ByteArrayType, GenericBinaryType, GenericStringType}; -use crate::{Array, ArrayRef, DictionaryArray, GenericByteArray, TypedDictionaryArray}; +use crate::{ + Array, ArrayRef, DictionaryArray, GenericByteArray, PrimitiveArray, TypedDictionaryArray, +}; use arrow_buffer::ArrowNativeType; use arrow_schema::{ArrowError, DataType}; use hashbrown::HashTable; +use num::NumCast; use std::any::Any; use std::sync::Arc; @@ -152,6 +155,71 @@ where values_builder, }) } + + /// Creates a new `GenericByteDictionaryBuilder` from the existing builder with the same + /// keys and values, but with a new data type for the keys. + /// + /// # Example + /// ``` + /// # + /// # use arrow_array::builder::StringDictionaryBuilder; + /// # use arrow_array::types::{UInt8Type, UInt16Type}; + /// # use arrow_array::UInt16Array; + /// # use arrow_schema::ArrowError; + /// + /// let mut u8_keyed_builder = StringDictionaryBuilder::::new(); + /// + /// // appending too many values causes the dictionary to overflow + /// for i in 0..256 { + /// u8_keyed_builder.append_value(format!("{}", i)); + /// } + /// let result = u8_keyed_builder.append("256"); + /// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{}))); + /// + /// // we need to upgrade to a larger key type + /// let mut u16_keyed_builder = StringDictionaryBuilder::::try_new_from_builder(u8_keyed_builder).unwrap(); + /// let dictionary_array = u16_keyed_builder.finish(); + /// let keys = dictionary_array.keys(); + /// + /// assert_eq!(keys, &UInt16Array::from_iter(0..256)); + /// ``` + pub fn try_new_from_builder( + mut source: GenericByteDictionaryBuilder, + ) -> Result + where + K::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast, + { + let state = source.state; + let dedup = source.dedup; + let values_builder = source.values_builder; + + let source_keys = source.keys_builder.finish(); + let new_keys: PrimitiveArray = source_keys.try_unary(|value| { + num::cast::cast::(value).ok_or_else(|| { + ArrowError::CastError(format!( + "Can't cast dictionary keys from source type {:?} to type {:?}", + K2::DATA_TYPE, + K::DATA_TYPE + )) + }) + })?; + + // drop source key here because currently source_keys and new_keys are holding reference to + // the same underlying null_buffer. Below we want to call new_keys.into_builder() it must + // be the only reference holder. + drop(source_keys); + + Ok(Self { + state, + dedup, + keys_builder: new_keys + .into_builder() + .expect("underlying buffer has no references"), + values_builder, + }) + } } impl ArrayBuilder for GenericByteDictionaryBuilder @@ -503,7 +571,7 @@ mod tests { use crate::array::Int8Array; use crate::cast::AsArray; - use crate::types::{Int16Type, Int32Type, Int8Type, Utf8Type}; + use crate::types::{Int16Type, Int32Type, Int8Type, UInt16Type, UInt8Type, Utf8Type}; use crate::{ArrowPrimitiveType, BinaryArray, StringArray}; fn test_bytes_dictionary_builder(values: Vec<&T::Native>) @@ -614,6 +682,97 @@ mod tests { ]); } + fn _test_try_new_from_builder_generic_for_key_types(values: Vec<&T::Native>) + where + K1: ArrowDictionaryKeyType, + K1::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast + From, + T: ByteArrayType, + ::Native: PartialEq + AsRef<::Native>, + { + let mut source = GenericByteDictionaryBuilder::::new(); + source.append(values[0]).unwrap(); + source.append(values[1]).unwrap(); + source.append_null(); + source.append(values[2]).unwrap(); + + let mut result = + GenericByteDictionaryBuilder::::try_new_from_builder(source).unwrap(); + let array = result.finish(); + + let mut expected_keys_builder = PrimitiveBuilder::::new(); + expected_keys_builder + .append_value(<::Native as From>::from(0u8)); + expected_keys_builder + .append_value(<::Native as From>::from(1u8)); + expected_keys_builder.append_null(); + expected_keys_builder + .append_value(<::Native as From>::from(2u8)); + let expected_keys = expected_keys_builder.finish(); + assert_eq!(array.keys(), &expected_keys); + + let av = array.values(); + let ava: &GenericByteArray = av.as_any().downcast_ref::>().unwrap(); + assert_eq!(ava.value(0), values[0]); + assert_eq!(ava.value(1), values[1]); + assert_eq!(ava.value(2), values[2]); + } + + fn test_try_new_from_builder(values: Vec<&T::Native>) + where + T: ByteArrayType, + ::Native: PartialEq + AsRef<::Native>, + { + // test cast to bigger size unsigned + _test_try_new_from_builder_generic_for_key_types::( + values.clone(), + ); + // test cast going to smaller size unsigned + _test_try_new_from_builder_generic_for_key_types::( + values.clone(), + ); + // test cast going to bigger size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test cast going to smaller size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test going from signed to signed for different size changes + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + } + + #[test] + fn test_string_dictionary_builder_try_new_from_builder() { + test_try_new_from_builder::>(vec!["abc", "def", "ghi"]); + } + + #[test] + fn test_binary_dictionary_builder_try_new_from_builder() { + test_try_new_from_builder::>(vec![b"abc", b"def", b"ghi"]); + } + + #[test] + fn test_try_new_from_builder_cast_fails() { + let mut source_builder = StringDictionaryBuilder::::new(); + for i in 0..257 { + source_builder.append_value(format!("val{i}")); + } + + // there should be too many values that we can't downcast to the underlying type + // we have keys that wouldn't fit into UInt8Type + let result = StringDictionaryBuilder::::try_new_from_builder(source_builder); + assert!(result.is_err()); + if let Err(e) = result { + assert!(matches!(e, ArrowError::CastError(_))); + assert_eq!( + e.to_string(), + "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8" + ); + } + } + fn test_bytes_dictionary_builder_with_existing_dictionary( dictionary: GenericByteArray, values: Vec<&T::Native>, diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs index 7268e751b149..cba2bb428e53 100644 --- a/arrow-array/src/builder/generic_bytes_view_builder.rs +++ b/arrow-array/src/builder/generic_bytes_view_builder.rs @@ -19,8 +19,8 @@ use std::any::Any; use std::marker::PhantomData; use std::sync::Arc; -use arrow_buffer::{Buffer, BufferBuilder, NullBufferBuilder, ScalarBuffer}; -use arrow_data::ByteView; +use arrow_buffer::{Buffer, NullBufferBuilder, ScalarBuffer}; +use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN}; use arrow_schema::ArrowError; use hashbrown::hash_table::Entry; use hashbrown::HashTable; @@ -28,7 +28,7 @@ use hashbrown::HashTable; use crate::builder::ArrayBuilder; use crate::types::bytes::ByteArrayNativeType; use crate::types::{BinaryViewType, ByteViewType, StringViewType}; -use crate::{ArrayRef, GenericByteViewArray}; +use crate::{Array, ArrayRef, GenericByteViewArray}; const STARTING_BLOCK_SIZE: u32 = 8 * 1024; // 8KiB const MAX_BLOCK_SIZE: u32 = 2 * 1024 * 1024; // 2MiB @@ -68,8 +68,8 @@ impl BlockSizeGrowthStrategy { /// /// To avoid bump allocating, this builder allocates data in fixed size blocks, configurable /// using [`GenericByteViewBuilder::with_fixed_block_size`]. [`GenericByteViewBuilder::append_value`] -/// writes values larger than 12 bytes to the current in-progress block, with values smaller -/// than 12 bytes inlined into the views. If a value is appended that will not fit in the +/// writes values larger than [`MAX_INLINE_VIEW_LEN`] bytes to the current in-progress block, with values smaller +/// than [`MAX_INLINE_VIEW_LEN`] bytes inlined into the views. If a value is appended that will not fit in the /// in-progress block, it will be closed, and a new block of sufficient size allocated /// /// # Append Views @@ -79,7 +79,7 @@ impl BlockSizeGrowthStrategy { /// using [`GenericByteViewBuilder::append_block`] and then views into this block appended /// using [`GenericByteViewBuilder::try_append_view`] pub struct GenericByteViewBuilder { - views_builder: BufferBuilder, + views_buffer: Vec, null_buffer_builder: NullBufferBuilder, completed: Vec, in_progress: Vec, @@ -99,7 +99,7 @@ impl GenericByteViewBuilder { /// Creates a new [`GenericByteViewBuilder`] with space for `capacity` string values. pub fn with_capacity(capacity: usize) -> Self { Self { - views_builder: BufferBuilder::new(capacity), + views_buffer: Vec::with_capacity(capacity), null_buffer_builder: NullBufferBuilder::new(capacity), completed: vec![], in_progress: vec![], @@ -114,7 +114,7 @@ impl GenericByteViewBuilder { /// Set a fixed buffer size for variable length strings /// /// The block size is the size of the buffer used to store values greater - /// than 12 bytes. The builder allocates new buffers when the current + /// than [`MAX_INLINE_VIEW_LEN`] bytes. The builder allocates new buffers when the current /// buffer is full. /// /// By default the builder balances buffer size and buffer count by @@ -134,13 +134,6 @@ impl GenericByteViewBuilder { } } - /// Override the size of buffers to allocate for holding string data - /// Use `with_fixed_block_size` instead. - #[deprecated(since = "53.0.0", note = "Use `with_fixed_block_size` instead")] - pub fn with_block_size(self, block_size: u32) -> Self { - self.with_fixed_block_size(block_size) - } - /// Deduplicate strings while building the array /// /// This will potentially decrease the memory usage if the array have repeated strings @@ -148,7 +141,7 @@ impl GenericByteViewBuilder { pub fn with_deduplicate_strings(self) -> Self { Self { string_tracker: Some(( - HashTable::with_capacity(self.views_builder.capacity()), + HashTable::with_capacity(self.views_buffer.capacity()), Default::default(), )), ..self @@ -201,10 +194,42 @@ impl GenericByteViewBuilder { let b = b.get_unchecked(start..end); let view = make_view(b, block, offset); - self.views_builder.append(view); + self.views_buffer.push(view); self.null_buffer_builder.append_non_null(); } + /// Appends an array to the builder. + /// This will flush any in-progress block and append the data buffers + /// and add the (adapted) views. + pub fn append_array(&mut self, array: &GenericByteViewArray) { + self.flush_in_progress(); + // keep original views if this array is the first to be added or if there are no data buffers (all inline views) + let keep_views = self.completed.is_empty() || array.data_buffers().is_empty(); + let starting_buffer = self.completed.len() as u32; + + self.completed.extend(array.data_buffers().iter().cloned()); + + if keep_views { + self.views_buffer.extend_from_slice(array.views()); + } else { + self.views_buffer.extend(array.views().iter().map(|v| { + let mut byte_view = ByteView::from(*v); + if byte_view.length > MAX_INLINE_VIEW_LEN { + // Small views (<=12 bytes) are inlined, so only need to update large views + byte_view.buffer_index += starting_buffer; + }; + + byte_view.as_u128() + })); + } + + if let Some(null_buffer) = array.nulls() { + self.null_buffer_builder.append_buffer(null_buffer); + } else { + self.null_buffer_builder.append_n_non_nulls(array.len()); + } + } + /// Try to append a view of the given `block`, `offset` and `length` /// /// See [`Self::append_block`] @@ -255,9 +280,9 @@ impl GenericByteViewBuilder { /// Useful if we want to know what value has been inserted to the builder /// The index has to be smaller than `self.len()`, otherwise it will panic pub fn get_value(&self, index: usize) -> &[u8] { - let view = self.views_builder.as_slice().get(index).unwrap(); + let view = self.views_buffer.as_slice().get(index).unwrap(); let len = *view as u32; - if len <= 12 { + if len <= MAX_INLINE_VIEW_LEN { // # Safety // The view is valid from the builder unsafe { GenericByteViewArray::::inline_value(view, len as usize) } @@ -283,11 +308,11 @@ impl GenericByteViewBuilder { pub fn append_value(&mut self, value: impl AsRef) { let v: &[u8] = value.as_ref().as_ref(); let length: u32 = v.len().try_into().unwrap(); - if length <= 12 { + if length <= MAX_INLINE_VIEW_LEN { let mut view_buffer = [0; 16]; view_buffer[0..4].copy_from_slice(&length.to_le_bytes()); view_buffer[4..4 + v.len()].copy_from_slice(v); - self.views_builder.append(u128::from_le_bytes(view_buffer)); + self.views_buffer.push(u128::from_le_bytes(view_buffer)); self.null_buffer_builder.append_non_null(); return; } @@ -311,8 +336,7 @@ impl GenericByteViewBuilder { Entry::Occupied(occupied) => { // If the string already exists, we will directly use the view let idx = occupied.get(); - self.views_builder - .append(self.views_builder.as_slice()[*idx]); + self.views_buffer.push(self.views_buffer[*idx]); self.null_buffer_builder.append_non_null(); self.string_tracker = Some((ht, hasher)); return; @@ -320,7 +344,7 @@ impl GenericByteViewBuilder { Entry::Vacant(vacant) => { // o.w. we insert the (string hash -> view index) // the idx is current length of views_builder, as we are inserting a new view - vacant.insert(self.views_builder.len()); + vacant.insert(self.views_buffer.len()); } } self.string_tracker = Some((ht, hasher)); @@ -341,7 +365,7 @@ impl GenericByteViewBuilder { buffer_index: self.completed.len() as u32, offset, }; - self.views_builder.append(view.into()); + self.views_buffer.push(view.into()); self.null_buffer_builder.append_non_null(); } @@ -358,21 +382,20 @@ impl GenericByteViewBuilder { #[inline] pub fn append_null(&mut self) { self.null_buffer_builder.append_null(); - self.views_builder.append(0); + self.views_buffer.push(0); } /// Builds the [`GenericByteViewArray`] and reset this builder pub fn finish(&mut self) -> GenericByteViewArray { self.flush_in_progress(); let completed = std::mem::take(&mut self.completed); - let len = self.views_builder.len(); - let views = ScalarBuffer::new(self.views_builder.finish(), 0, len); let nulls = self.null_buffer_builder.finish(); if let Some((ref mut ht, _)) = self.string_tracker.as_mut() { ht.clear(); } + let views = std::mem::take(&mut self.views_buffer); // SAFETY: valid by construction - unsafe { GenericByteViewArray::new_unchecked(views, completed, nulls) } + unsafe { GenericByteViewArray::new_unchecked(views.into(), completed, nulls) } } /// Builds the [`GenericByteViewArray`] without resetting the builder @@ -381,8 +404,8 @@ impl GenericByteViewBuilder { if !self.in_progress.is_empty() { completed.push(Buffer::from_slice_ref(&self.in_progress)); } - let len = self.views_builder.len(); - let views = Buffer::from_slice_ref(self.views_builder.as_slice()); + let len = self.views_buffer.len(); + let views = Buffer::from_slice_ref(self.views_buffer.as_slice()); let views = ScalarBuffer::new(views, 0, len); let nulls = self.null_buffer_builder.finish_cloned(); // SAFETY: valid by construction @@ -396,7 +419,7 @@ impl GenericByteViewBuilder { /// Return the allocated size of this builder in bytes, useful for memory accounting. pub fn allocated_size(&self) -> usize { - let views = self.views_builder.capacity() * std::mem::size_of::(); + let views = self.views_buffer.capacity() * std::mem::size_of::(); let null = self.null_buffer_builder.allocated_size(); let buffer_size = self.completed.iter().map(|b| b.capacity()).sum::(); let in_progress = self.in_progress.capacity(); @@ -418,7 +441,7 @@ impl std::fmt::Debug for GenericByteViewBuilder { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}ViewBuilder", T::PREFIX)?; f.debug_struct("") - .field("views_builder", &self.views_builder) + .field("views_buffer", &self.views_buffer) .field("in_progress", &self.in_progress) .field("completed", &self.completed) .field("null_buffer_builder", &self.null_buffer_builder) diff --git a/arrow-array/src/builder/generic_list_builder.rs b/arrow-array/src/builder/generic_list_builder.rs index a9c88ec6c586..463b498c55ba 100644 --- a/arrow-array/src/builder/generic_list_builder.rs +++ b/arrow-array/src/builder/generic_list_builder.rs @@ -270,6 +270,14 @@ where self.null_buffer_builder.append_null(); } + /// Appends `n` `null`s into the builder. + #[inline] + pub fn append_nulls(&mut self, n: usize) { + let next_offset = self.next_offset(); + self.offsets_builder.append_n(n, next_offset); + self.null_buffer_builder.append_n_nulls(n); + } + /// Appends an optional value into this [`GenericListBuilder`] /// /// If `Some` calls [`Self::append_value`] otherwise calls [`Self::append_null`] @@ -406,7 +414,7 @@ mod tests { let values_builder = Int32Builder::with_capacity(10); let mut builder = GenericListBuilder::::new(values_builder); - // [[0, 1, 2], null, [3, null, 5], [6, 7]] + // [[0, 1, 2], null, [3, null, 5], [6, 7], null, null, [8]] builder.values().append_value(0); builder.values().append_value(1); builder.values().append_value(2); @@ -419,14 +427,20 @@ mod tests { builder.values().append_value(6); builder.values().append_value(7); builder.append(true); + builder.append_nulls(2); + builder.values().append_value(8); + builder.append(true); let list_array = builder.finish(); assert_eq!(DataType::Int32, list_array.value_type()); - assert_eq!(4, list_array.len()); - assert_eq!(1, list_array.null_count()); + assert_eq!(7, list_array.len()); + assert_eq!(3, list_array.null_count()); assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]); + assert_eq!(O::from_usize(9).unwrap(), list_array.value_offsets()[7]); assert_eq!(O::from_usize(3).unwrap(), list_array.value_length(2)); + assert!(list_array.is_null(4)); + assert!(list_array.is_null(5)); } #[test] diff --git a/arrow-array/src/builder/generic_list_view_builder.rs b/arrow-array/src/builder/generic_list_view_builder.rs new file mode 100644 index 000000000000..5aaf9efefe24 --- /dev/null +++ b/arrow-array/src/builder/generic_list_view_builder.rs @@ -0,0 +1,707 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::builder::ArrayBuilder; +use crate::{ArrayRef, GenericListViewArray, OffsetSizeTrait}; +use arrow_buffer::{Buffer, BufferBuilder, NullBufferBuilder, ScalarBuffer}; +use arrow_schema::{Field, FieldRef}; +use std::any::Any; +use std::sync::Arc; + +/// Builder for [`GenericListViewArray`] +#[derive(Debug)] +pub struct GenericListViewBuilder { + offsets_builder: BufferBuilder, + sizes_builder: BufferBuilder, + null_buffer_builder: NullBufferBuilder, + values_builder: T, + field: Option, + current_offset: OffsetSize, +} + +impl Default for GenericListViewBuilder { + fn default() -> Self { + Self::new(T::default()) + } +} + +impl ArrayBuilder + for GenericListViewBuilder +{ + /// Returns the builder as a non-mutable `Any` reference. + fn as_any(&self) -> &dyn Any { + self + } + + /// Returns the builder as a mutable `Any` reference. + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + /// Returns the boxed builder as a box of `Any`. + fn into_box_any(self: Box) -> Box { + self + } + + /// Returns the number of array slots in the builder + fn len(&self) -> usize { + self.null_buffer_builder.len() + } + + /// Builds the array and reset this builder. + fn finish(&mut self) -> ArrayRef { + Arc::new(self.finish()) + } + + /// Builds the array without resetting the builder. + fn finish_cloned(&self) -> ArrayRef { + Arc::new(self.finish_cloned()) + } +} + +impl GenericListViewBuilder { + /// Creates a new [`GenericListViewBuilder`] from a given values array builder + pub fn new(values_builder: T) -> Self { + let capacity = values_builder.len(); + Self::with_capacity(values_builder, capacity) + } + + /// Creates a new [`GenericListViewBuilder`] from a given values array builder + /// `capacity` is the number of items to pre-allocate space for in this builder + pub fn with_capacity(values_builder: T, capacity: usize) -> Self { + let offsets_builder = BufferBuilder::::new(capacity); + let sizes_builder = BufferBuilder::::new(capacity); + Self { + offsets_builder, + null_buffer_builder: NullBufferBuilder::new(capacity), + values_builder, + sizes_builder, + field: None, + current_offset: OffsetSize::zero(), + } + } + + /// + /// By default a nullable field is created with the name `item` + /// + /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the + /// field's data type does not match that of `T` + pub fn with_field(self, field: impl Into) -> Self { + Self { + field: Some(field.into()), + ..self + } + } +} + +impl GenericListViewBuilder +where + T: 'static, +{ + /// Returns the child array builder as a mutable reference. + /// + /// This mutable reference can be used to append values into the child array builder, + /// but you must call [`append`](#method.append) to delimit each distinct list value. + pub fn values(&mut self) -> &mut T { + &mut self.values_builder + } + + /// Returns the child array builder as an immutable reference + pub fn values_ref(&self) -> &T { + &self.values_builder + } + + /// Finish the current variable-length list array slot + /// + /// # Panics + /// + /// Panics if the length of [`Self::values`] exceeds `OffsetSize::MAX` + #[inline] + pub fn append(&mut self, is_valid: bool) { + self.offsets_builder.append(self.current_offset); + self.sizes_builder.append( + OffsetSize::from_usize( + self.values_builder.len() - self.current_offset.to_usize().unwrap(), + ) + .unwrap(), + ); + self.null_buffer_builder.append(is_valid); + self.current_offset = OffsetSize::from_usize(self.values_builder.len()).unwrap(); + } + + /// Append value into this [`GenericListViewBuilder`] + #[inline] + pub fn append_value(&mut self, i: I) + where + T: Extend>, + I: IntoIterator>, + { + self.extend(std::iter::once(Some(i))) + } + + /// Append a null to this [`GenericListViewBuilder`] + /// + /// See [`Self::append_value`] for an example use. + #[inline] + pub fn append_null(&mut self) { + self.offsets_builder.append(self.current_offset); + self.sizes_builder + .append(OffsetSize::from_usize(0).unwrap()); + self.null_buffer_builder.append_null(); + } + + /// Appends an optional value into this [`GenericListViewBuilder`] + /// + /// If `Some` calls [`Self::append_value`] otherwise calls [`Self::append_null`] + #[inline] + pub fn append_option(&mut self, i: Option) + where + T: Extend>, + I: IntoIterator>, + { + match i { + Some(i) => self.append_value(i), + None => self.append_null(), + } + } + + /// Builds the [`GenericListViewArray`] and reset this builder. + pub fn finish(&mut self) -> GenericListViewArray { + let values = self.values_builder.finish(); + let nulls = self.null_buffer_builder.finish(); + let offsets = self.offsets_builder.finish(); + self.current_offset = OffsetSize::zero(); + + // Safety: Safe by construction + let offsets = ScalarBuffer::from(offsets); + let sizes = self.sizes_builder.finish(); + let sizes = ScalarBuffer::from(sizes); + let field = match &self.field { + Some(f) => f.clone(), + None => Arc::new(Field::new("item", values.data_type().clone(), true)), + }; + GenericListViewArray::new(field, offsets, sizes, values, nulls) + } + + /// Builds the [`GenericListViewArray`] without resetting the builder. + pub fn finish_cloned(&self) -> GenericListViewArray { + let values = self.values_builder.finish_cloned(); + let nulls = self.null_buffer_builder.finish_cloned(); + + let offsets = Buffer::from_slice_ref(self.offsets_builder.as_slice()); + // Safety: safe by construction + let offsets = ScalarBuffer::from(offsets); + + let sizes = Buffer::from_slice_ref(self.sizes_builder.as_slice()); + let sizes = ScalarBuffer::from(sizes); + + let field = match &self.field { + Some(f) => f.clone(), + None => Arc::new(Field::new("item", values.data_type().clone(), true)), + }; + + GenericListViewArray::new(field, offsets, sizes, values, nulls) + } + + /// Returns the current offsets buffer as a slice + pub fn offsets_slice(&self) -> &[OffsetSize] { + self.offsets_builder.as_slice() + } +} + +impl Extend> for GenericListViewBuilder +where + O: OffsetSizeTrait, + B: ArrayBuilder + Extend, + V: IntoIterator, +{ + #[inline] + fn extend>>(&mut self, iter: T) { + for v in iter { + match v { + Some(elements) => { + self.values_builder.extend(elements); + self.append(true); + } + None => self.append(false), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::builder::{make_builder, Int32Builder, ListViewBuilder}; + use crate::cast::AsArray; + use crate::types::Int32Type; + use crate::{Array, Int32Array}; + use arrow_schema::DataType; + + fn test_generic_list_view_array_builder_impl() { + let values_builder = Int32Builder::with_capacity(10); + let mut builder = GenericListViewBuilder::::new(values_builder); + + // [[0, 1, 2], [3, 4, 5], [6, 7]] + builder.values().append_value(0); + builder.values().append_value(1); + builder.values().append_value(2); + builder.append(true); + builder.values().append_value(3); + builder.values().append_value(4); + builder.values().append_value(5); + builder.append(true); + builder.values().append_value(6); + builder.values().append_value(7); + builder.append(true); + let list_array = builder.finish(); + + let list_values = list_array.values().as_primitive::(); + assert_eq!(list_values.values(), &[0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(list_array.value_offsets(), [0, 3, 6].map(O::usize_as)); + assert_eq!(list_array.value_sizes(), [3, 3, 2].map(O::usize_as)); + assert_eq!(DataType::Int32, list_array.value_type()); + assert_eq!(3, list_array.len()); + assert_eq!(0, list_array.null_count()); + assert_eq!(O::from_usize(6).unwrap(), list_array.value_offsets()[2]); + assert_eq!(O::from_usize(2).unwrap(), list_array.value_sizes()[2]); + for i in 0..2 { + assert!(list_array.is_valid(i)); + assert!(!list_array.is_null(i)); + } + } + + #[test] + fn test_list_view_array_builder() { + test_generic_list_view_array_builder_impl::() + } + + #[test] + fn test_large_list_view_array_builder() { + test_generic_list_view_array_builder_impl::() + } + + fn test_generic_list_view_array_builder_nulls_impl() { + let values_builder = Int32Builder::with_capacity(10); + let mut builder = GenericListViewBuilder::::new(values_builder); + + // [[0, 1, 2], null, [3, null, 5], [6, 7]] + builder.values().append_value(0); + builder.values().append_value(1); + builder.values().append_value(2); + builder.append(true); + builder.append(false); + builder.values().append_value(3); + builder.values().append_null(); + builder.values().append_value(5); + builder.append(true); + builder.values().append_value(6); + builder.values().append_value(7); + builder.append(true); + + let list_array = builder.finish(); + + assert_eq!(DataType::Int32, list_array.value_type()); + assert_eq!(4, list_array.len()); + assert_eq!(1, list_array.null_count()); + assert_eq!(O::from_usize(3).unwrap(), list_array.value_offsets()[2]); + assert_eq!(O::from_usize(3).unwrap(), list_array.value_sizes()[2]); + } + + #[test] + fn test_list_view_array_builder_nulls() { + test_generic_list_view_array_builder_nulls_impl::() + } + + #[test] + fn test_large_list_view_array_builder_nulls() { + test_generic_list_view_array_builder_nulls_impl::() + } + + #[test] + fn test_list_view_array_builder_finish() { + let values_builder = Int32Array::builder(5); + let mut builder = ListViewBuilder::new(values_builder); + + builder.values().append_slice(&[1, 2, 3]); + builder.append(true); + builder.values().append_slice(&[4, 5, 6]); + builder.append(true); + + let mut arr = builder.finish(); + assert_eq!(2, arr.len()); + assert!(builder.is_empty()); + + builder.values().append_slice(&[7, 8, 9]); + builder.append(true); + arr = builder.finish(); + assert_eq!(1, arr.len()); + assert!(builder.is_empty()); + } + + #[test] + fn test_list_view_array_builder_finish_cloned() { + let values_builder = Int32Array::builder(5); + let mut builder = ListViewBuilder::new(values_builder); + + builder.values().append_slice(&[1, 2, 3]); + builder.append(true); + builder.values().append_slice(&[4, 5, 6]); + builder.append(true); + + let mut arr = builder.finish_cloned(); + assert_eq!(2, arr.len()); + assert!(!builder.is_empty()); + + builder.values().append_slice(&[7, 8, 9]); + builder.append(true); + arr = builder.finish(); + assert_eq!(3, arr.len()); + assert!(builder.is_empty()); + } + + #[test] + fn test_list_view_list_view_array_builder() { + let primitive_builder = Int32Builder::with_capacity(10); + let values_builder = ListViewBuilder::new(primitive_builder); + let mut builder = ListViewBuilder::new(values_builder); + + // [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]] + builder.values().values().append_value(1); + builder.values().values().append_value(2); + builder.values().append(true); + builder.values().values().append_value(3); + builder.values().values().append_value(4); + builder.values().append(true); + builder.append(true); + + builder.values().values().append_value(5); + builder.values().values().append_value(6); + builder.values().values().append_value(7); + builder.values().append(true); + builder.values().append(false); + builder.values().values().append_value(8); + builder.values().append(true); + builder.append(true); + + builder.append(false); + + builder.values().values().append_value(9); + builder.values().values().append_value(10); + builder.values().append(true); + builder.append(true); + + let l1 = builder.finish(); + + assert_eq!(4, l1.len()); + assert_eq!(1, l1.null_count()); + + assert_eq!(l1.value_offsets(), &[0, 2, 5, 5]); + assert_eq!(l1.value_sizes(), &[2, 3, 0, 1]); + + let l2 = l1.values().as_list_view::(); + + assert_eq!(6, l2.len()); + assert_eq!(1, l2.null_count()); + assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8]); + assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2]); + + let i1 = l2.values().as_primitive::(); + assert_eq!(10, i1.len()); + assert_eq!(0, i1.null_count()); + assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + } + + #[test] + fn test_extend() { + let mut builder = ListViewBuilder::new(Int32Builder::new()); + builder.extend([ + Some(vec![Some(1), Some(2), Some(7), None]), + Some(vec![]), + Some(vec![Some(4), Some(5)]), + None, + ]); + + let array = builder.finish(); + assert_eq!(array.value_offsets(), [0, 4, 4, 6]); + assert_eq!(array.value_sizes(), [4, 0, 2, 0]); + assert_eq!(array.null_count(), 1); + assert!(array.is_null(3)); + let elements = array.values().as_primitive::(); + assert_eq!(elements.values(), &[1, 2, 7, 0, 4, 5]); + assert_eq!(elements.null_count(), 1); + assert!(elements.is_null(3)); + } + + #[test] + fn test_boxed_primitive_array_builder() { + let values_builder = make_builder(&DataType::Int32, 5); + let mut builder = ListViewBuilder::new(values_builder); + + builder + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_slice(&[1, 2, 3]); + builder.append(true); + + builder + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_slice(&[4, 5, 6]); + builder.append(true); + + let arr = builder.finish(); + assert_eq!(2, arr.len()); + + let elements = arr.values().as_primitive::(); + assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]); + } + + #[test] + fn test_boxed_list_view_list_view_array_builder() { + // This test is same as `test_list_list_array_builder` but uses boxed builders. + let values_builder = make_builder( + &DataType::ListView(Arc::new(Field::new("item", DataType::Int32, true))), + 10, + ); + test_boxed_generic_list_view_generic_list_view_array_builder::(values_builder); + } + + #[test] + fn test_boxed_large_list_view_large_list_view_array_builder() { + // This test is same as `test_list_list_array_builder` but uses boxed builders. + let values_builder = make_builder( + &DataType::LargeListView(Arc::new(Field::new("item", DataType::Int32, true))), + 10, + ); + test_boxed_generic_list_view_generic_list_view_array_builder::(values_builder); + } + + fn test_boxed_generic_list_view_generic_list_view_array_builder( + values_builder: Box, + ) where + O: OffsetSizeTrait + PartialEq, + { + let mut builder: GenericListViewBuilder> = + GenericListViewBuilder::>::new(values_builder); + + // [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]] + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(1); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(2); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(true); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(3); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(4); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(true); + builder.append(true); + + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(5); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(6); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an (Large)ListViewBuilder") + .append_value(7); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(true); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(false); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(8); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(true); + builder.append(true); + + builder.append(false); + + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(9); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .values() + .as_any_mut() + .downcast_mut::() + .expect("should be an Int32Builder") + .append_value(10); + builder + .values() + .as_any_mut() + .downcast_mut::>>() + .expect("should be an (Large)ListViewBuilder") + .append(true); + builder.append(true); + + let l1 = builder.finish(); + assert_eq!(4, l1.len()); + assert_eq!(1, l1.null_count()); + assert_eq!(l1.value_offsets(), &[0, 2, 5, 5].map(O::usize_as)); + assert_eq!(l1.value_sizes(), &[2, 3, 0, 1].map(O::usize_as)); + + let l2 = l1.values().as_list_view::(); + assert_eq!(6, l2.len()); + assert_eq!(1, l2.null_count()); + assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8].map(O::usize_as)); + assert_eq!(l2.value_sizes(), &[2, 2, 3, 0, 1, 2].map(O::usize_as)); + + let i1 = l2.values().as_primitive::(); + assert_eq!(10, i1.len()); + assert_eq!(0, i1.null_count()); + assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + } + + #[test] + fn test_with_field() { + let field = Arc::new(Field::new("bar", DataType::Int32, false)); + let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone()); + builder.append_value([Some(1), Some(2), Some(3)]); + builder.append_null(); // This is fine as nullability refers to nullability of values + builder.append_value([Some(4)]); + let array = builder.finish(); + assert_eq!(array.len(), 3); + assert_eq!(array.data_type(), &DataType::ListView(field.clone())); + + builder.append_value([Some(4), Some(5)]); + let array = builder.finish(); + assert_eq!(array.data_type(), &DataType::ListView(field)); + assert_eq!(array.len(), 1); + } + + #[test] + #[should_panic( + expected = r#"Non-nullable field of ListViewArray \"item\" cannot contain nulls"# + )] + // If a non-nullable type is declared but a null value is used, it will be intercepted by the null check. + fn test_checks_nullability() { + let field = Arc::new(Field::new("item", DataType::Int32, false)); + let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone()); + builder.append_value([Some(1), None]); + builder.finish(); + } + + #[test] + #[should_panic(expected = "ListViewArray expected data type Int64 got Int32")] + // If the declared type does not match the actual appended type, it will be intercepted by type checking in the finish function. + fn test_checks_data_type() { + let field = Arc::new(Field::new("item", DataType::Int64, false)); + let mut builder = ListViewBuilder::new(Int32Builder::new()).with_field(field.clone()); + builder.append_value([Some(1)]); + builder.finish(); + } +} diff --git a/arrow-array/src/builder/map_builder.rs b/arrow-array/src/builder/map_builder.rs index 1d89d427aae1..012a454e76c9 100644 --- a/arrow-array/src/builder/map_builder.rs +++ b/arrow-array/src/builder/map_builder.rs @@ -61,6 +61,7 @@ pub struct MapBuilder { field_names: MapFieldNames, key_builder: K, value_builder: V, + key_field: Option, value_field: Option, } @@ -107,13 +108,27 @@ impl MapBuilder { field_names: field_names.unwrap_or_default(), key_builder, value_builder, + key_field: None, value_field: None, } } /// Override the field passed to [`MapBuilder::new`] /// - /// By default a nullable field is created with the name `values` + /// By default, a non-nullable field is created with the name `keys` + /// + /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the + /// field's data type does not match that of `K` or the field is nullable + pub fn with_keys_field(self, field: impl Into) -> Self { + Self { + key_field: Some(field.into()), + ..self + } + } + + /// Override the field passed to [`MapBuilder::new`] + /// + /// By default, a nullable field is created with the name `values` /// /// Note: [`Self::finish`] and [`Self::finish_cloned`] will panic if the /// field's data type does not match that of `V` @@ -194,11 +209,17 @@ impl MapBuilder { keys_arr.null_count() ); - let keys_field = Arc::new(Field::new( - self.field_names.key.as_str(), - keys_arr.data_type().clone(), - false, // always non-nullable - )); + let keys_field = match &self.key_field { + Some(f) => { + assert!(!f.is_nullable(), "Keys field must not be nullable"); + f.clone() + } + None => Arc::new(Field::new( + self.field_names.key.as_str(), + keys_arr.data_type().clone(), + false, // always non-nullable + )), + }; let values_field = match &self.value_field { Some(f) => f.clone(), None => Arc::new(Field::new( @@ -262,10 +283,10 @@ impl ArrayBuilder for MapBuilder { #[cfg(test)] mod tests { + use super::*; use crate::builder::{make_builder, Int32Builder, StringBuilder}; use crate::{Int32Array, StringArray}; - - use super::*; + use std::collections::HashMap; #[test] #[should_panic(expected = "Keys array must have no null values, found 1 null value(s)")] @@ -377,4 +398,67 @@ mod tests { ) ); } + + #[test] + fn test_with_keys_field() { + let mut key_metadata = HashMap::new(); + key_metadata.insert("foo".to_string(), "bar".to_string()); + let key_field = Arc::new( + Field::new("keys", DataType::Int32, false).with_metadata(key_metadata.clone()), + ); + let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new()) + .with_keys_field(key_field.clone()); + builder.keys().append_value(1); + builder.values().append_value(2); + builder.append(true).unwrap(); + let map = builder.finish(); + + assert_eq!(map.len(), 1); + assert_eq!( + map.data_type(), + &DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Arc::new( + Field::new("keys", DataType::Int32, false) + .with_metadata(key_metadata) + ), + Arc::new(Field::new("values", DataType::Int32, true)) + ] + .into() + ), + false, + )), + false + ) + ); + } + + #[test] + #[should_panic(expected = "Keys field must not be nullable")] + fn test_with_nullable_keys_field() { + let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new()) + .with_keys_field(Arc::new(Field::new("keys", DataType::Int32, true))); + + builder.keys().append_value(1); + builder.values().append_value(2); + builder.append(true).unwrap(); + + builder.finish(); + } + + #[test] + #[should_panic(expected = "Incorrect datatype")] + fn test_keys_field_type_mismatch() { + let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new()) + .with_keys_field(Arc::new(Field::new("keys", DataType::Utf8, false))); + + builder.keys().append_value(1); + builder.values().append_value(2); + builder.append(true).unwrap(); + + builder.finish(); + } } diff --git a/arrow-array/src/builder/mod.rs b/arrow-array/src/builder/mod.rs index 89a96280eb87..ea9c98f9b60e 100644 --- a/arrow-array/src/builder/mod.rs +++ b/arrow-array/src/builder/mod.rs @@ -78,6 +78,73 @@ //! )) //! ``` //! +//! # Using the [`Extend`] trait to append values from an iterable: +//! +//! ``` +//! # use arrow_array::{Array}; +//! # use arrow_array::builder::{ArrayBuilder, StringBuilder}; +//! +//! let mut builder = StringBuilder::new(); +//! builder.extend(vec![Some("🍐"), Some("🍎"), None]); +//! assert_eq!(builder.finish().len(), 3); +//! ``` +//! +//! # Using the [`Extend`] trait to write generic functions: +//! +//! ``` +//! # use arrow_array::{Array, ArrayRef, StringArray}; +//! # use arrow_array::builder::{ArrayBuilder, Int32Builder, ListBuilder, StringBuilder}; +//! +//! // For generic methods that fill a list of values for an [`ArrayBuilder`], use the [`Extend`] trait. +//! fn filter_and_fill>(builder: &mut impl Extend, values: I, filter: V) +//! where V: PartialEq +//! { +//! builder.extend(values.into_iter().filter(|v| *v == filter)); +//! } +//! let mut string_builder = StringBuilder::new(); +//! filter_and_fill( +//! &mut string_builder, +//! vec![Some("🍐"), Some("🍎"), None], +//! Some("🍎"), +//! ); +//! assert_eq!(string_builder.finish().len(), 1); +//! +//! let mut int_builder = Int32Builder::new(); +//! filter_and_fill( +//! &mut int_builder, +//! vec![Some(11), Some(42), None], +//! Some(42), +//! ); +//! assert_eq!(int_builder.finish().len(), 1); +//! +//! // For generic methods that fill lists-of-lists for an [`ArrayBuilder`], use the [`Extend`] trait. +//! fn filter_and_fill_if_contains>>( +//! list_builder: &mut impl Extend>, +//! values: I, +//! filter: Option, +//! ) where +//! T: PartialEq, +//! for<'a> &'a V: IntoIterator>, +//! { +//! list_builder.extend(values.into_iter().filter(|string: &Option| { +//! string +//! .as_ref() +//! .map(|str: &V| str.into_iter().any(|ch: &Option| ch == &filter)) +//! .unwrap_or(false) +//! })); +//! } +//! let builder = StringBuilder::new(); +//! let mut list_builder = ListBuilder::new(builder); +//! let pear_pear = vec![Some("🍐"),Some("🍐")]; +//! let pear_app = vec![Some("🍐"),Some("🍎")]; +//! filter_and_fill_if_contains( +//! &mut list_builder, +//! vec![Some(pear_pear), Some(pear_app), None], +//! Some("🍎"), +//! ); +//! assert_eq!(list_builder.finish().len(), 1); +//! ``` +//! //! # Custom Builders //! //! It is common to have a collection of statically defined Rust types that @@ -134,6 +201,8 @@ //! } //! } //! +//! /// For building arrays in generic code, use Extend instead of the append_* methods +//! /// e.g. append_value, append_option, append_null //! impl<'a> Extend<&'a MyRow> for MyRowBuilder { //! fn extend>(&mut self, iter: T) { //! iter.into_iter().for_each(|row| self.append(row)); @@ -147,8 +216,24 @@ //! RecordBatch::from(&builder.finish()) //! } //! ``` +//! +//! # Null / Validity Masks +//! +//! The [`NullBufferBuilder`] is optimized for creating the null mask for an array. +//! +//! ``` +//! # use arrow_array::builder::NullBufferBuilder; +//! let mut builder = NullBufferBuilder::new(8); +//! let mut builder = NullBufferBuilder::new(8); +//! builder.append_n_non_nulls(7); +//! builder.append_null(); +//! let buffer = builder.finish().unwrap(); +//! assert_eq!(buffer.len(), 8); +//! assert_eq!(buffer.iter().collect::>(), vec![true, true, true, true, true, true, true, false]); +//! ``` pub use arrow_buffer::BooleanBufferBuilder; +pub use arrow_buffer::NullBufferBuilder; mod boolean_builder; pub use boolean_builder::*; @@ -158,6 +243,8 @@ mod fixed_size_binary_builder; pub use fixed_size_binary_builder::*; mod fixed_size_list_builder; pub use fixed_size_list_builder::*; +mod fixed_size_binary_dictionary_builder; +pub use fixed_size_binary_dictionary_builder::*; mod generic_bytes_builder; pub use generic_bytes_builder::*; mod generic_list_builder; @@ -180,11 +267,15 @@ mod generic_byte_run_builder; pub use generic_byte_run_builder::*; mod generic_bytes_view_builder; pub use generic_bytes_view_builder::*; +mod generic_list_view_builder; +pub use generic_list_view_builder::*; mod union_builder; pub use union_builder::*; +use crate::types::{Int16Type, Int32Type, Int64Type, Int8Type}; use crate::ArrayRef; +use arrow_schema::{DataType, IntervalUnit, TimeUnit}; use std::any::Any; /// Trait for dealing with different array builders at runtime @@ -304,6 +395,12 @@ pub type ListBuilder = GenericListBuilder; /// Builder for [`LargeListArray`](crate::array::LargeListArray) pub type LargeListBuilder = GenericListBuilder; +/// Builder for [`ListViewArray`](crate::array::ListViewArray) +pub type ListViewBuilder = GenericListViewBuilder; + +/// Builder for [`LargeListViewArray`](crate::array::LargeListViewArray) +pub type LargeListViewBuilder = GenericListViewBuilder; + /// Builder for [`BinaryArray`](crate::array::BinaryArray) /// /// See examples on [`GenericBinaryBuilder`] @@ -323,3 +420,194 @@ pub type StringBuilder = GenericStringBuilder; /// /// See examples on [`GenericStringBuilder`] pub type LargeStringBuilder = GenericStringBuilder; + +/// Returns a builder with capacity for `capacity` elements of datatype +/// `DataType`. +/// +/// This function is useful to construct arrays from an arbitrary vectors with +/// known/expected schema. +/// +/// See comments on [StructBuilder] for retrieving collection builders built by +/// make_builder. +pub fn make_builder(datatype: &DataType, capacity: usize) -> Box { + use crate::builder::*; + match datatype { + DataType::Null => Box::new(NullBuilder::new()), + DataType::Boolean => Box::new(BooleanBuilder::with_capacity(capacity)), + DataType::Int8 => Box::new(Int8Builder::with_capacity(capacity)), + DataType::Int16 => Box::new(Int16Builder::with_capacity(capacity)), + DataType::Int32 => Box::new(Int32Builder::with_capacity(capacity)), + DataType::Int64 => Box::new(Int64Builder::with_capacity(capacity)), + DataType::UInt8 => Box::new(UInt8Builder::with_capacity(capacity)), + DataType::UInt16 => Box::new(UInt16Builder::with_capacity(capacity)), + DataType::UInt32 => Box::new(UInt32Builder::with_capacity(capacity)), + DataType::UInt64 => Box::new(UInt64Builder::with_capacity(capacity)), + DataType::Float16 => Box::new(Float16Builder::with_capacity(capacity)), + DataType::Float32 => Box::new(Float32Builder::with_capacity(capacity)), + DataType::Float64 => Box::new(Float64Builder::with_capacity(capacity)), + DataType::Binary => Box::new(BinaryBuilder::with_capacity(capacity, 1024)), + DataType::LargeBinary => Box::new(LargeBinaryBuilder::with_capacity(capacity, 1024)), + DataType::BinaryView => Box::new(BinaryViewBuilder::with_capacity(capacity)), + DataType::FixedSizeBinary(len) => { + Box::new(FixedSizeBinaryBuilder::with_capacity(capacity, *len)) + } + DataType::Decimal32(p, s) => Box::new( + Decimal32Builder::with_capacity(capacity).with_data_type(DataType::Decimal32(*p, *s)), + ), + DataType::Decimal64(p, s) => Box::new( + Decimal64Builder::with_capacity(capacity).with_data_type(DataType::Decimal64(*p, *s)), + ), + DataType::Decimal128(p, s) => Box::new( + Decimal128Builder::with_capacity(capacity).with_data_type(DataType::Decimal128(*p, *s)), + ), + DataType::Decimal256(p, s) => Box::new( + Decimal256Builder::with_capacity(capacity).with_data_type(DataType::Decimal256(*p, *s)), + ), + DataType::Utf8 => Box::new(StringBuilder::with_capacity(capacity, 1024)), + DataType::LargeUtf8 => Box::new(LargeStringBuilder::with_capacity(capacity, 1024)), + DataType::Utf8View => Box::new(StringViewBuilder::with_capacity(capacity)), + DataType::Date32 => Box::new(Date32Builder::with_capacity(capacity)), + DataType::Date64 => Box::new(Date64Builder::with_capacity(capacity)), + DataType::Time32(TimeUnit::Second) => { + Box::new(Time32SecondBuilder::with_capacity(capacity)) + } + DataType::Time32(TimeUnit::Millisecond) => { + Box::new(Time32MillisecondBuilder::with_capacity(capacity)) + } + DataType::Time64(TimeUnit::Microsecond) => { + Box::new(Time64MicrosecondBuilder::with_capacity(capacity)) + } + DataType::Time64(TimeUnit::Nanosecond) => { + Box::new(Time64NanosecondBuilder::with_capacity(capacity)) + } + DataType::Timestamp(TimeUnit::Second, tz) => Box::new( + TimestampSecondBuilder::with_capacity(capacity) + .with_data_type(DataType::Timestamp(TimeUnit::Second, tz.clone())), + ), + DataType::Timestamp(TimeUnit::Millisecond, tz) => Box::new( + TimestampMillisecondBuilder::with_capacity(capacity) + .with_data_type(DataType::Timestamp(TimeUnit::Millisecond, tz.clone())), + ), + DataType::Timestamp(TimeUnit::Microsecond, tz) => Box::new( + TimestampMicrosecondBuilder::with_capacity(capacity) + .with_data_type(DataType::Timestamp(TimeUnit::Microsecond, tz.clone())), + ), + DataType::Timestamp(TimeUnit::Nanosecond, tz) => Box::new( + TimestampNanosecondBuilder::with_capacity(capacity) + .with_data_type(DataType::Timestamp(TimeUnit::Nanosecond, tz.clone())), + ), + DataType::Interval(IntervalUnit::YearMonth) => { + Box::new(IntervalYearMonthBuilder::with_capacity(capacity)) + } + DataType::Interval(IntervalUnit::DayTime) => { + Box::new(IntervalDayTimeBuilder::with_capacity(capacity)) + } + DataType::Interval(IntervalUnit::MonthDayNano) => { + Box::new(IntervalMonthDayNanoBuilder::with_capacity(capacity)) + } + DataType::Duration(TimeUnit::Second) => { + Box::new(DurationSecondBuilder::with_capacity(capacity)) + } + DataType::Duration(TimeUnit::Millisecond) => { + Box::new(DurationMillisecondBuilder::with_capacity(capacity)) + } + DataType::Duration(TimeUnit::Microsecond) => { + Box::new(DurationMicrosecondBuilder::with_capacity(capacity)) + } + DataType::Duration(TimeUnit::Nanosecond) => { + Box::new(DurationNanosecondBuilder::with_capacity(capacity)) + } + DataType::List(field) => { + let builder = make_builder(field.data_type(), capacity); + Box::new(ListBuilder::with_capacity(builder, capacity).with_field(field.clone())) + } + DataType::LargeList(field) => { + let builder = make_builder(field.data_type(), capacity); + Box::new(LargeListBuilder::with_capacity(builder, capacity).with_field(field.clone())) + } + DataType::FixedSizeList(field, size) => { + let size = *size; + let values_builder_capacity = { + let size: usize = size.try_into().unwrap(); + capacity * size + }; + let builder = make_builder(field.data_type(), values_builder_capacity); + Box::new( + FixedSizeListBuilder::with_capacity(builder, size, capacity) + .with_field(field.clone()), + ) + } + DataType::ListView(field) => { + let builder = make_builder(field.data_type(), capacity); + Box::new(ListViewBuilder::with_capacity(builder, capacity).with_field(field.clone())) + } + DataType::LargeListView(field) => { + let builder = make_builder(field.data_type(), capacity); + Box::new( + LargeListViewBuilder::with_capacity(builder, capacity).with_field(field.clone()), + ) + } + DataType::Map(field, _) => match field.data_type() { + DataType::Struct(fields) => { + let map_field_names = MapFieldNames { + key: fields[0].name().clone(), + value: fields[1].name().clone(), + entry: field.name().clone(), + }; + let key_builder = make_builder(fields[0].data_type(), capacity); + let value_builder = make_builder(fields[1].data_type(), capacity); + Box::new( + MapBuilder::with_capacity( + Some(map_field_names), + key_builder, + value_builder, + capacity, + ) + .with_keys_field(fields[0].clone()) + .with_values_field(fields[1].clone()), + ) + } + t => panic!("The field of Map data type {t:?} should have a child Struct field"), + }, + DataType::Struct(fields) => Box::new(StructBuilder::from_fields(fields.clone(), capacity)), + t @ DataType::Dictionary(key_type, value_type) => { + macro_rules! dict_builder { + ($key_type:ty) => { + match &**value_type { + DataType::Utf8 => { + let dict_builder: StringDictionaryBuilder<$key_type> = + StringDictionaryBuilder::with_capacity(capacity, 256, 1024); + Box::new(dict_builder) + } + DataType::LargeUtf8 => { + let dict_builder: LargeStringDictionaryBuilder<$key_type> = + LargeStringDictionaryBuilder::with_capacity(capacity, 256, 1024); + Box::new(dict_builder) + } + DataType::Binary => { + let dict_builder: BinaryDictionaryBuilder<$key_type> = + BinaryDictionaryBuilder::with_capacity(capacity, 256, 1024); + Box::new(dict_builder) + } + DataType::LargeBinary => { + let dict_builder: LargeBinaryDictionaryBuilder<$key_type> = + LargeBinaryDictionaryBuilder::with_capacity(capacity, 256, 1024); + Box::new(dict_builder) + } + t => panic!("Dictionary value type {t:?} is not currently supported"), + } + }; + } + match &**key_type { + DataType::Int8 => dict_builder!(Int8Type), + DataType::Int16 => dict_builder!(Int16Type), + DataType::Int32 => dict_builder!(Int32Type), + DataType::Int64 => dict_builder!(Int64Type), + _ => { + panic!("Data type {t:?} with key type {key_type:?} is not currently supported") + } + } + } + t => panic!("Data type {t:?} is not currently supported"), + } +} diff --git a/arrow-array/src/builder/null_builder.rs b/arrow-array/src/builder/null_builder.rs index 59086dffa907..489822065b56 100644 --- a/arrow-array/src/builder/null_builder.rs +++ b/arrow-array/src/builder/null_builder.rs @@ -59,18 +59,6 @@ impl NullBuilder { Self { len: 0 } } - /// Creates a new null builder with space for `capacity` elements without re-allocating - #[deprecated = "there is no actual notion of capacity in the NullBuilder, so emulating it makes little sense"] - pub fn with_capacity(_capacity: usize) -> Self { - Self::new() - } - - /// Returns the capacity of this builder measured in slots of type `T` - #[deprecated = "there is no actual notion of capacity in the NullBuilder, so emulating it makes little sense"] - pub fn capacity(&self) -> usize { - self.len - } - /// Appends a null slot into the builder #[inline] pub fn append_null(&mut self) { diff --git a/arrow-array/src/builder/primitive_builder.rs b/arrow-array/src/builder/primitive_builder.rs index 3191fea6e407..7aca730ce192 100644 --- a/arrow-array/src/builder/primitive_builder.rs +++ b/arrow-array/src/builder/primitive_builder.rs @@ -17,7 +17,7 @@ use crate::builder::{ArrayBuilder, BufferBuilder}; use crate::types::*; -use crate::{ArrayRef, PrimitiveArray}; +use crate::{Array, ArrayRef, PrimitiveArray}; use arrow_buffer::NullBufferBuilder; use arrow_buffer::{Buffer, MutableBuffer}; use arrow_data::ArrayData; @@ -87,6 +87,10 @@ pub type DurationMicrosecondBuilder = PrimitiveBuilder; /// An elapsed time in nanoseconds array builder. pub type DurationNanosecondBuilder = PrimitiveBuilder; +/// A decimal 32 array builder +pub type Decimal32Builder = PrimitiveBuilder; +/// A decimal 64 array builder +pub type Decimal64Builder = PrimitiveBuilder; /// A decimal 128 array builder pub type Decimal128Builder = PrimitiveBuilder; /// A decimal 256 array builder @@ -175,7 +179,8 @@ impl PrimitiveBuilder { /// data type of the generated array. /// /// This method allows overriding the data type, to allow specifying timezones - /// for [`DataType::Timestamp`] or precision and scale for [`DataType::Decimal128`] and [`DataType::Decimal256`] + /// for [`DataType::Timestamp`] or precision and scale for [`DataType::Decimal32`], + /// [`DataType::Decimal64`], [`DataType::Decimal128`] and [`DataType::Decimal256`] /// /// # Panics /// @@ -255,6 +260,28 @@ impl PrimitiveBuilder { self.values_builder.append_slice(values); } + /// Appends array values and null to this builder as is + /// (this means that underlying null values are copied as is). + /// + /// # Panics + /// + /// Panics if `array` and `self` data types are different + #[inline] + pub fn append_array(&mut self, array: &PrimitiveArray) { + assert_eq!( + &self.data_type, + array.data_type(), + "array data type mismatch" + ); + + self.values_builder.append_slice(array.values()); + if let Some(null_buffer) = array.nulls() { + self.null_buffer_builder.append_buffer(null_buffer); + } else { + self.null_buffer_builder.append_n_non_nulls(array.len()); + } + } + /// Appends values from a trusted length iterator. /// /// # Safety @@ -366,6 +393,7 @@ impl Extend> for PrimitiveBuilder

{ #[cfg(test)] mod tests { use super::*; + use arrow_buffer::{NullBuffer, ScalarBuffer}; use arrow_schema::TimeUnit; use crate::array::Array; @@ -615,4 +643,63 @@ mod tests { let array = builder.finish(); assert_eq!(array.values(), &[1, 2, 3, 5, 2, 4, 4, 2, 4, 6, 2]); } + + #[test] + fn test_primitive_array_append_array() { + let input = vec![ + Some(1), + None, + Some(3), + None, + Some(5), + None, + None, + None, + Some(7), + Some(9), + Some(8), + Some(6), + Some(4), + ]; + let arr1 = Int32Array::from(input[..5].to_vec()); + let arr2 = Int32Array::from(input[5..8].to_vec()); + let arr3 = Int32Array::from(input[8..].to_vec()); + + let mut builder = Int32Array::builder(5); + builder.append_array(&arr1); + builder.append_array(&arr2); + builder.append_array(&arr3); + let actual = builder.finish(); + let expected = Int32Array::from(input); + + assert_eq!(actual, expected); + } + + #[test] + fn test_append_array_add_underlying_null_values() { + let array = Int32Array::new( + ScalarBuffer::from(vec![2, 3, 4, 5]), + Some(NullBuffer::from(&[true, true, false, false])), + ); + + let mut builder = Int32Array::builder(5); + builder.append_array(&array); + let actual = builder.finish(); + + assert_eq!(actual, array); + assert_eq!(actual.values(), array.values()) + } + + #[test] + #[should_panic(expected = "array data type mismatch")] + fn test_invalid_with_data_type_in_append_array() { + let array = { + let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2)); + builder.append_value(1); + builder.finish() + }; + + let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(2, 3)); + builder.append_array(&array) + } } diff --git a/arrow-array/src/builder/primitive_dictionary_builder.rs b/arrow-array/src/builder/primitive_dictionary_builder.rs index 282f0ae9d5b1..1d921c6df097 100644 --- a/arrow-array/src/builder/primitive_dictionary_builder.rs +++ b/arrow-array/src/builder/primitive_dictionary_builder.rs @@ -22,6 +22,7 @@ use crate::{ }; use arrow_buffer::{ArrowNativeType, ToByteSlice}; use arrow_schema::{ArrowError, DataType}; +use num::NumCast; use std::any::Any; use std::collections::HashMap; use std::sync::Arc; @@ -126,10 +127,11 @@ where keys_builder.is_empty() && values_builder.is_empty(), "keys and values builders must be empty" ); + let values_capacity = values_builder.capacity(); Self { keys_builder, values_builder, - map: HashMap::new(), + map: HashMap::with_capacity(values_capacity), } } @@ -168,6 +170,68 @@ where map: HashMap::with_capacity(values_capacity), } } + + /// Creates a new `PrimitiveDictionaryBuilder` from the existing builder with the same + /// keys and values, but with a new data type for the keys. + /// + /// # Example + /// ``` + /// # + /// # use arrow_array::builder::PrimitiveDictionaryBuilder; + /// # use arrow_array::types::{UInt8Type, UInt16Type, UInt64Type}; + /// # use arrow_array::UInt16Array; + /// # use arrow_schema::ArrowError; + /// + /// let mut u8_keyed_builder = PrimitiveDictionaryBuilder::::new(); + /// + /// // appending too many values causes the dictionary to overflow + /// for i in 0..256 { + /// u8_keyed_builder.append_value(i); + /// } + /// let result = u8_keyed_builder.append(256); + /// assert!(matches!(result, Err(ArrowError::DictionaryKeyOverflowError{}))); + /// + /// // we need to upgrade to a larger key type + /// let mut u16_keyed_builder = PrimitiveDictionaryBuilder::::try_new_from_builder(u8_keyed_builder).unwrap(); + /// let dictionary_array = u16_keyed_builder.finish(); + /// let keys = dictionary_array.keys(); + /// + /// assert_eq!(keys, &UInt16Array::from_iter(0..256)); + pub fn try_new_from_builder( + mut source: PrimitiveDictionaryBuilder, + ) -> Result + where + K::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast, + { + let map = source.map; + let values_builder = source.values_builder; + + let source_keys = source.keys_builder.finish(); + let new_keys: PrimitiveArray = source_keys.try_unary(|value| { + num::cast::cast::(value).ok_or_else(|| { + ArrowError::CastError(format!( + "Can't cast dictionary keys from source type {:?} to type {:?}", + K2::DATA_TYPE, + K::DATA_TYPE + )) + }) + })?; + + // drop source key here because currently source_keys and new_keys are holding reference to + // the same underlying null_buffer. Below we want to call new_keys.into_builder() it must + // be the only reference holder. + drop(source_keys); + + Ok(Self { + map, + keys_builder: new_keys + .into_builder() + .expect("underlying buffer has no references"), + values_builder, + }) + } } impl ArrayBuilder for PrimitiveDictionaryBuilder @@ -430,7 +494,11 @@ mod tests { use crate::array::{Int32Array, UInt32Array, UInt8Array}; use crate::builder::Decimal128Builder; use crate::cast::AsArray; - use crate::types::{Decimal128Type, Int32Type, UInt32Type, UInt8Type}; + use crate::types::{ + Date32Type, Decimal128Type, DurationNanosecondType, Float32Type, Float64Type, Int16Type, + Int32Type, Int64Type, Int8Type, TimestampNanosecondType, UInt16Type, UInt32Type, + UInt64Type, UInt8Type, + }; #[test] fn test_primitive_dictionary_builder() { @@ -633,4 +701,120 @@ mod tests { assert_eq!(values, [None, None]); } + + #[test] + fn creating_dictionary_from_builders_should_use_values_capacity_for_the_map() { + let builder = PrimitiveDictionaryBuilder::::new_from_empty_builders( + PrimitiveBuilder::with_capacity(1).with_data_type(DataType::Int32), + PrimitiveBuilder::with_capacity(2).with_data_type(DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+08:00".into()))), + ); + + assert!( + builder.map.capacity() >= builder.values_builder.capacity(), + "map capacity {} should be at least the values capacity {}", + builder.map.capacity(), + builder.values_builder.capacity() + ) + } + + fn _test_try_new_from_builder_generic_for_key_types(values: Vec) + where + K1: ArrowDictionaryKeyType, + K1::Native: NumCast, + K2: ArrowDictionaryKeyType, + K2::Native: NumCast + From, + V: ArrowPrimitiveType, + { + let mut source = PrimitiveDictionaryBuilder::::new(); + source.append(values[0]).unwrap(); + source.append_null(); + source.append(values[1]).unwrap(); + source.append(values[2]).unwrap(); + + let mut result = PrimitiveDictionaryBuilder::::try_new_from_builder(source).unwrap(); + let array = result.finish(); + + let mut expected_keys_builder = PrimitiveBuilder::::new(); + expected_keys_builder + .append_value(<::Native as From>::from(0u8)); + expected_keys_builder.append_null(); + expected_keys_builder + .append_value(<::Native as From>::from(1u8)); + expected_keys_builder + .append_value(<::Native as From>::from(2u8)); + let expected_keys = expected_keys_builder.finish(); + assert_eq!(array.keys(), &expected_keys); + + let av = array.values(); + let ava = av.as_any().downcast_ref::>().unwrap(); + assert_eq!(ava.value(0), values[0]); + assert_eq!(ava.value(1), values[1]); + assert_eq!(ava.value(2), values[2]); + } + + fn _test_try_new_from_builder_generic_for_value(values: Vec) + where + T: ArrowPrimitiveType, + { + // test cast to bigger size unsigned + _test_try_new_from_builder_generic_for_key_types::( + values.clone(), + ); + // test cast going to smaller size unsigned + _test_try_new_from_builder_generic_for_key_types::( + values.clone(), + ); + // test cast going to bigger size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test cast going to smaller size signed + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + // test going from signed to signed for different size changes + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + _test_try_new_from_builder_generic_for_key_types::(values.clone()); + } + + #[test] + fn test_try_new_from_builder() { + // test unsigned types + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + // test signed types + _test_try_new_from_builder_generic_for_value::(vec![-1, 0, 1]); + _test_try_new_from_builder_generic_for_value::(vec![-1, 0, 1]); + _test_try_new_from_builder_generic_for_value::(vec![-1, 0, 1]); + _test_try_new_from_builder_generic_for_value::(vec![-1, 0, 1]); + // test some date types + _test_try_new_from_builder_generic_for_value::(vec![5, 6, 7]); + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + _test_try_new_from_builder_generic_for_value::(vec![1, 2, 3]); + // test some floating point types + _test_try_new_from_builder_generic_for_value::(vec![0.1, 0.2, 0.3]); + _test_try_new_from_builder_generic_for_value::(vec![-0.1, 0.2, 0.3]); + } + + #[test] + fn test_try_new_from_builder_cast_fails() { + let mut source_builder = PrimitiveDictionaryBuilder::::new(); + for i in 0..257 { + source_builder.append_value(i); + } + + // there should be too many values that we can't downcast to the underlying type + // we have keys that wouldn't fit into UInt8Type + let result = PrimitiveDictionaryBuilder::::try_new_from_builder( + source_builder, + ); + assert!(result.is_err()); + if let Err(e) = result { + assert!(matches!(e, ArrowError::CastError(_))); + assert_eq!( + e.to_string(), + "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8" + ); + } + } } diff --git a/arrow-array/src/builder/struct_builder.rs b/arrow-array/src/builder/struct_builder.rs index 2b288445c74b..3afee5863f52 100644 --- a/arrow-array/src/builder/struct_builder.rs +++ b/arrow-array/src/builder/struct_builder.rs @@ -15,13 +15,10 @@ // specific language governing permissions and limitations // under the License. +use crate::builder::*; use crate::StructArray; -use crate::{ - builder::*, - types::{Int16Type, Int32Type, Int64Type, Int8Type}, -}; use arrow_buffer::NullBufferBuilder; -use arrow_schema::{DataType, Fields, IntervalUnit, SchemaBuilder, TimeUnit}; +use arrow_schema::{Fields, SchemaBuilder}; use std::sync::Arc; /// Builder for [`StructArray`] @@ -162,178 +159,6 @@ impl ArrayBuilder for StructBuilder { } } -/// Returns a builder with capacity for `capacity` elements of datatype -/// `DataType`. -/// -/// This function is useful to construct arrays from an arbitrary vectors with -/// known/expected schema. -/// -/// See comments on [StructBuilder] for retrieving collection builders built by -/// make_builder. -pub fn make_builder(datatype: &DataType, capacity: usize) -> Box { - use crate::builder::*; - match datatype { - DataType::Null => Box::new(NullBuilder::new()), - DataType::Boolean => Box::new(BooleanBuilder::with_capacity(capacity)), - DataType::Int8 => Box::new(Int8Builder::with_capacity(capacity)), - DataType::Int16 => Box::new(Int16Builder::with_capacity(capacity)), - DataType::Int32 => Box::new(Int32Builder::with_capacity(capacity)), - DataType::Int64 => Box::new(Int64Builder::with_capacity(capacity)), - DataType::UInt8 => Box::new(UInt8Builder::with_capacity(capacity)), - DataType::UInt16 => Box::new(UInt16Builder::with_capacity(capacity)), - DataType::UInt32 => Box::new(UInt32Builder::with_capacity(capacity)), - DataType::UInt64 => Box::new(UInt64Builder::with_capacity(capacity)), - DataType::Float16 => Box::new(Float16Builder::with_capacity(capacity)), - DataType::Float32 => Box::new(Float32Builder::with_capacity(capacity)), - DataType::Float64 => Box::new(Float64Builder::with_capacity(capacity)), - DataType::Binary => Box::new(BinaryBuilder::with_capacity(capacity, 1024)), - DataType::LargeBinary => Box::new(LargeBinaryBuilder::with_capacity(capacity, 1024)), - DataType::FixedSizeBinary(len) => { - Box::new(FixedSizeBinaryBuilder::with_capacity(capacity, *len)) - } - DataType::Decimal128(p, s) => Box::new( - Decimal128Builder::with_capacity(capacity).with_data_type(DataType::Decimal128(*p, *s)), - ), - DataType::Decimal256(p, s) => Box::new( - Decimal256Builder::with_capacity(capacity).with_data_type(DataType::Decimal256(*p, *s)), - ), - DataType::Utf8 => Box::new(StringBuilder::with_capacity(capacity, 1024)), - DataType::LargeUtf8 => Box::new(LargeStringBuilder::with_capacity(capacity, 1024)), - DataType::Date32 => Box::new(Date32Builder::with_capacity(capacity)), - DataType::Date64 => Box::new(Date64Builder::with_capacity(capacity)), - DataType::Time32(TimeUnit::Second) => { - Box::new(Time32SecondBuilder::with_capacity(capacity)) - } - DataType::Time32(TimeUnit::Millisecond) => { - Box::new(Time32MillisecondBuilder::with_capacity(capacity)) - } - DataType::Time64(TimeUnit::Microsecond) => { - Box::new(Time64MicrosecondBuilder::with_capacity(capacity)) - } - DataType::Time64(TimeUnit::Nanosecond) => { - Box::new(Time64NanosecondBuilder::with_capacity(capacity)) - } - DataType::Timestamp(TimeUnit::Second, tz) => Box::new( - TimestampSecondBuilder::with_capacity(capacity) - .with_data_type(DataType::Timestamp(TimeUnit::Second, tz.clone())), - ), - DataType::Timestamp(TimeUnit::Millisecond, tz) => Box::new( - TimestampMillisecondBuilder::with_capacity(capacity) - .with_data_type(DataType::Timestamp(TimeUnit::Millisecond, tz.clone())), - ), - DataType::Timestamp(TimeUnit::Microsecond, tz) => Box::new( - TimestampMicrosecondBuilder::with_capacity(capacity) - .with_data_type(DataType::Timestamp(TimeUnit::Microsecond, tz.clone())), - ), - DataType::Timestamp(TimeUnit::Nanosecond, tz) => Box::new( - TimestampNanosecondBuilder::with_capacity(capacity) - .with_data_type(DataType::Timestamp(TimeUnit::Nanosecond, tz.clone())), - ), - DataType::Interval(IntervalUnit::YearMonth) => { - Box::new(IntervalYearMonthBuilder::with_capacity(capacity)) - } - DataType::Interval(IntervalUnit::DayTime) => { - Box::new(IntervalDayTimeBuilder::with_capacity(capacity)) - } - DataType::Interval(IntervalUnit::MonthDayNano) => { - Box::new(IntervalMonthDayNanoBuilder::with_capacity(capacity)) - } - DataType::Duration(TimeUnit::Second) => { - Box::new(DurationSecondBuilder::with_capacity(capacity)) - } - DataType::Duration(TimeUnit::Millisecond) => { - Box::new(DurationMillisecondBuilder::with_capacity(capacity)) - } - DataType::Duration(TimeUnit::Microsecond) => { - Box::new(DurationMicrosecondBuilder::with_capacity(capacity)) - } - DataType::Duration(TimeUnit::Nanosecond) => { - Box::new(DurationNanosecondBuilder::with_capacity(capacity)) - } - DataType::List(field) => { - let builder = make_builder(field.data_type(), capacity); - Box::new(ListBuilder::with_capacity(builder, capacity).with_field(field.clone())) - } - DataType::LargeList(field) => { - let builder = make_builder(field.data_type(), capacity); - Box::new(LargeListBuilder::with_capacity(builder, capacity).with_field(field.clone())) - } - DataType::FixedSizeList(field, size) => { - let size = *size; - let values_builder_capacity = { - let size: usize = size.try_into().unwrap(); - capacity * size - }; - let builder = make_builder(field.data_type(), values_builder_capacity); - Box::new( - FixedSizeListBuilder::with_capacity(builder, size, capacity) - .with_field(field.clone()), - ) - } - DataType::Map(field, _) => match field.data_type() { - DataType::Struct(fields) => { - let map_field_names = MapFieldNames { - key: fields[0].name().clone(), - value: fields[1].name().clone(), - entry: field.name().clone(), - }; - let key_builder = make_builder(fields[0].data_type(), capacity); - let value_builder = make_builder(fields[1].data_type(), capacity); - Box::new( - MapBuilder::with_capacity( - Some(map_field_names), - key_builder, - value_builder, - capacity, - ) - .with_values_field(fields[1].clone()), - ) - } - t => panic!("The field of Map data type {t:?} should has a child Struct field"), - }, - DataType::Struct(fields) => Box::new(StructBuilder::from_fields(fields.clone(), capacity)), - t @ DataType::Dictionary(key_type, value_type) => { - macro_rules! dict_builder { - ($key_type:ty) => { - match &**value_type { - DataType::Utf8 => { - let dict_builder: StringDictionaryBuilder<$key_type> = - StringDictionaryBuilder::with_capacity(capacity, 256, 1024); - Box::new(dict_builder) - } - DataType::LargeUtf8 => { - let dict_builder: LargeStringDictionaryBuilder<$key_type> = - LargeStringDictionaryBuilder::with_capacity(capacity, 256, 1024); - Box::new(dict_builder) - } - DataType::Binary => { - let dict_builder: BinaryDictionaryBuilder<$key_type> = - BinaryDictionaryBuilder::with_capacity(capacity, 256, 1024); - Box::new(dict_builder) - } - DataType::LargeBinary => { - let dict_builder: LargeBinaryDictionaryBuilder<$key_type> = - LargeBinaryDictionaryBuilder::with_capacity(capacity, 256, 1024); - Box::new(dict_builder) - } - t => panic!("Dictionary value type {t:?} is not currently supported"), - } - }; - } - match &**key_type { - DataType::Int8 => dict_builder!(Int8Type), - DataType::Int16 => dict_builder!(Int16Type), - DataType::Int32 => dict_builder!(Int32Type), - DataType::Int64 => dict_builder!(Int64Type), - _ => { - panic!("Data type {t:?} with key type {key_type:?} is not currently supported") - } - } - } - t => panic!("Data type {t:?} is not currently supported"), - } -} - impl StructBuilder { /// Creates a new `StructBuilder` pub fn new(fields: impl Into, field_builders: Vec>) -> Self { @@ -361,6 +186,16 @@ impl StructBuilder { self.field_builders[i].as_any_mut().downcast_mut::() } + /// Returns a reference to field builders + pub fn field_builders(&self) -> &[Box] { + &self.field_builders + } + + /// Returns a mutable reference to field builders + pub fn field_builders_mut(&mut self) -> &mut [Box] { + &mut self.field_builders + } + /// Returns the number of fields for the struct this builder is building. pub fn num_fields(&self) -> usize { self.field_builders.len() @@ -379,6 +214,12 @@ impl StructBuilder { self.append(false) } + /// Appends `n` `null`s into the builder. + #[inline] + pub fn append_nulls(&mut self, n: usize) { + self.null_buffer_builder.append_slice(&vec![false; n]); + } + /// Builds the `StructArray` and reset this builder. pub fn finish(&mut self) -> StructArray { self.validate_content(); @@ -478,6 +319,8 @@ mod tests { string_builder.append_null(); string_builder.append_null(); string_builder.append_value("mark"); + string_builder.append_nulls(2); + string_builder.append_value("terry"); let int_builder = builder .field_builder::(1) @@ -486,35 +329,43 @@ mod tests { int_builder.append_value(2); int_builder.append_null(); int_builder.append_value(4); + int_builder.append_nulls(2); + int_builder.append_value(3); builder.append(true); builder.append(true); builder.append_null(); builder.append(true); + builder.append_nulls(2); + builder.append(true); + let struct_data = builder.finish().into_data(); - assert_eq!(4, struct_data.len()); - assert_eq!(1, struct_data.null_count()); - assert_eq!(&[11_u8], struct_data.nulls().unwrap().validity()); + assert_eq!(7, struct_data.len()); + assert_eq!(3, struct_data.null_count()); + assert_eq!(&[75_u8], struct_data.nulls().unwrap().validity()); let expected_string_data = ArrayData::builder(DataType::Utf8) - .len(4) - .null_bit_buffer(Some(Buffer::from(&[9_u8]))) - .add_buffer(Buffer::from_slice_ref([0, 3, 3, 3, 7])) - .add_buffer(Buffer::from_slice_ref(b"joemark")) + .len(7) + .null_bit_buffer(Some(Buffer::from(&[73_u8]))) + .add_buffer(Buffer::from_slice_ref([0, 3, 3, 3, 7, 7, 7, 12])) + .add_buffer(Buffer::from_slice_ref(b"joemarkterry")) .build() .unwrap(); let expected_int_data = ArrayData::builder(DataType::Int32) - .len(4) - .null_bit_buffer(Some(Buffer::from_slice_ref([11_u8]))) - .add_buffer(Buffer::from_slice_ref([1, 2, 0, 4])) + .len(7) + .null_bit_buffer(Some(Buffer::from_slice_ref([75_u8]))) + .add_buffer(Buffer::from_slice_ref([1, 2, 0, 4, 4, 4, 3])) .build() .unwrap(); assert_eq!(expected_string_data, struct_data.child_data()[0]); assert_eq!(expected_int_data, struct_data.child_data()[1]); + + assert!(struct_data.is_null(4)); + assert!(struct_data.is_null(5)); } #[test] diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs index fc657f94c6a6..41fffc4bc80c 100644 --- a/arrow-array/src/cast.rs +++ b/arrow-array/src/cast.rs @@ -60,6 +60,8 @@ macro_rules! repeat_pat { /// k.as_ref() => (dictionary_key_size_helper, u8), /// _ => unreachable!(), /// }, +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => u8::MAX, /// _ => u8::MAX, /// } /// } @@ -72,7 +74,7 @@ macro_rules! repeat_pat { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_integer { - ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => { + ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { match ($($data_type),+) { $crate::repeat_pat!($crate::cast::__private::DataType::Int8, $($data_type),+) => { $m!($crate::types::Int8Type $(, $args)*) @@ -98,7 +100,57 @@ macro_rules! downcast_integer { $crate::repeat_pat!($crate::cast::__private::DataType::UInt64, $($data_type),+) => { $m!($crate::types::UInt64Type $(, $args)*) } - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* + } + }; +} + +/// Given one or more expressions evaluating to an integer [`PrimitiveArray`] invokes the provided macro +/// with the corresponding array, along with match statements for any non integer array types +/// +/// ``` +/// # use arrow_array::{Array, downcast_integer_array, cast::as_string_array, cast::as_largestring_array}; +/// # use arrow_schema::DataType; +/// +/// fn print_integer(array: &dyn Array) { +/// downcast_integer_array!( +/// array => { +/// for v in array { +/// println!("{:?}", v); +/// } +/// } +/// DataType::Utf8 => { +/// for v in as_string_array(array) { +/// println!("{:?}", v); +/// } +/// } +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => { +/// for v in as_largestring_array(array) { +/// println!("{:?}", v); +/// } +/// } +/// t => println!("Unsupported datatype {}", t) +/// ) +/// } +/// ``` +/// +/// [`DataType`]: arrow_schema::DataType +#[macro_export] +macro_rules! downcast_integer_array { + ($values:ident => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_integer_array!($values => {$e} $($p $(if $pred)* => $fallback)*) + }; + (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_integer_array!($($values),+ => {$e} $($p $(if $pred)* => $fallback)*) + }; + ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_integer_array!(($($values),+) => $e $($p $(if $pred)* => $fallback)*) + }; + (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_integer!{ + $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e), + $($p $(if $pred)* => $fallback,)* } }; } @@ -123,6 +175,8 @@ macro_rules! downcast_integer { /// k.data_type() => (run_end_size_helper, u8), /// _ => unreachable!(), /// }, +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => u8::MAX, /// _ => u8::MAX, /// } /// } @@ -135,7 +189,7 @@ macro_rules! downcast_integer { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_run_end_index { - ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => { + ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { match ($($data_type),+) { $crate::repeat_pat!($crate::cast::__private::DataType::Int16, $($data_type),+) => { $m!($crate::types::Int16Type $(, $args)*) @@ -146,7 +200,7 @@ macro_rules! downcast_run_end_index { $crate::repeat_pat!($crate::cast::__private::DataType::Int64, $($data_type),+) => { $m!($crate::types::Int64Type $(, $args)*) } - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } }; } @@ -167,6 +221,8 @@ macro_rules! downcast_run_end_index { /// fn temporal_size(t: &DataType) -> u8 { /// downcast_temporal! { /// t => (temporal_size_helper, u8), +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => u8::MAX, /// _ => u8::MAX /// } /// } @@ -178,7 +234,7 @@ macro_rules! downcast_run_end_index { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_temporal { - ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => { + ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { match ($($data_type),+) { $crate::repeat_pat!($crate::cast::__private::DataType::Time32($crate::cast::__private::TimeUnit::Second), $($data_type),+) => { $m!($crate::types::Time32SecondType $(, $args)*) @@ -210,7 +266,7 @@ macro_rules! downcast_temporal { $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Nanosecond, _), $($data_type),+) => { $m!($crate::types::TimestampNanosecondType $(, $args)*) } - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } }; } @@ -219,7 +275,7 @@ macro_rules! downcast_temporal { /// accepts a number of subsequent patterns to match the data type /// /// ``` -/// # use arrow_array::{Array, downcast_temporal_array, cast::as_string_array}; +/// # use arrow_array::{Array, downcast_temporal_array, cast::as_string_array, cast::as_largestring_array}; /// # use arrow_schema::DataType; /// /// fn print_temporal(array: &dyn Array) { @@ -234,6 +290,12 @@ macro_rules! downcast_temporal { /// println!("{:?}", v); /// } /// } +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => { +/// for v in as_largestring_array(array) { +/// println!("{:?}", v); +/// } +/// } /// t => println!("Unsupported datatype {}", t) /// ) /// } @@ -242,19 +304,19 @@ macro_rules! downcast_temporal { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_temporal_array { - ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_temporal_array!($values => {$e} $($p => $fallback)*) + ($values:ident => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_temporal_array!($values => {$e} $($p $(if $pred)* => $fallback)*) }; - (($($values:ident),+) => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_temporal_array!($($values),+ => {$e} $($p => $fallback)*) + (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_temporal_array!($($values),+ => {$e} $($p $(if $pred)* => $fallback)*) }; - ($($values:ident),+ => $e:block $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_temporal_array!(($($values),+) => $e $($p => $fallback)*) + ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_temporal_array!(($($values),+) => $e $($p $(if $pred)* => $fallback)*) }; - (($($values:ident),+) => $e:block $($p:pat => $fallback:expr $(,)*)*) => { + (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { $crate::downcast_temporal!{ $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e), - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } }; } @@ -275,6 +337,8 @@ macro_rules! downcast_temporal_array { /// fn primitive_size(t: &DataType) -> u8 { /// downcast_primitive! { /// t => (primitive_size_helper, u8), +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => u8::MAX, /// _ => u8::MAX /// } /// } @@ -289,7 +353,7 @@ macro_rules! downcast_temporal_array { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_primitive { - ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => { + ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { $crate::downcast_integer! { $($data_type),+ => ($m $(, $args)*), $crate::repeat_pat!($crate::cast::__private::DataType::Float16, $($data_type),+) => { @@ -301,6 +365,12 @@ macro_rules! downcast_primitive { $crate::repeat_pat!($crate::cast::__private::DataType::Float64, $($data_type),+) => { $m!($crate::types::Float64Type $(, $args)*) } + $crate::repeat_pat!($crate::cast::__private::DataType::Decimal32(_, _), $($data_type),+) => { + $m!($crate::types::Decimal32Type $(, $args)*) + } + $crate::repeat_pat!($crate::cast::__private::DataType::Decimal64(_, _), $($data_type),+) => { + $m!($crate::types::Decimal64Type $(, $args)*) + } $crate::repeat_pat!($crate::cast::__private::DataType::Decimal128(_, _), $($data_type),+) => { $m!($crate::types::Decimal128Type $(, $args)*) } @@ -331,7 +401,7 @@ macro_rules! downcast_primitive { _ => { $crate::downcast_temporal! { $($data_type),+ => ($m $(, $args)*), - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } } } @@ -351,7 +421,7 @@ macro_rules! downcast_primitive_array_helper { /// accepts a number of subsequent patterns to match the data type /// /// ``` -/// # use arrow_array::{Array, downcast_primitive_array, cast::as_string_array}; +/// # use arrow_array::{Array, downcast_primitive_array, cast::as_string_array, cast::as_largestring_array}; /// # use arrow_schema::DataType; /// /// fn print_primitive(array: &dyn Array) { @@ -366,6 +436,12 @@ macro_rules! downcast_primitive_array_helper { /// println!("{:?}", v); /// } /// } +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => { +/// for v in as_largestring_array(array) { +/// println!("{:?}", v); +/// } +/// } /// t => println!("Unsupported datatype {}", t) /// ) /// } @@ -374,19 +450,19 @@ macro_rules! downcast_primitive_array_helper { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_primitive_array { - ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_primitive_array!($values => {$e} $($p => $fallback)*) + ($values:ident => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_primitive_array!($values => {$e} $($p $(if $pred)* => $fallback)*) }; - (($($values:ident),+) => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_primitive_array!($($values),+ => {$e} $($p => $fallback)*) + (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_primitive_array!($($values),+ => {$e} $($p $(if $pred)* => $fallback)*) }; - ($($values:ident),+ => $e:block $($p:pat => $fallback:expr $(,)*)*) => { - $crate::downcast_primitive_array!(($($values),+) => $e $($p => $fallback)*) + ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + $crate::downcast_primitive_array!(($($values),+) => $e $($p $(if $pred)* => $fallback)*) }; - (($($values:ident),+) => $e:block $($p:pat => $fallback:expr $(,)*)*) => { + (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { $crate::downcast_primitive!{ $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e), - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } }; } @@ -438,7 +514,7 @@ macro_rules! downcast_dictionary_array_helper { /// a number of subsequent patterns to match the data type /// /// ``` -/// # use arrow_array::{Array, StringArray, downcast_dictionary_array, cast::as_string_array}; +/// # use arrow_array::{Array, StringArray, downcast_dictionary_array, cast::as_string_array, cast::as_largestring_array}; /// # use arrow_schema::DataType; /// /// fn print_strings(array: &dyn Array) { @@ -456,6 +532,12 @@ macro_rules! downcast_dictionary_array_helper { /// println!("{:?}", v); /// } /// } +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => { +/// for v in as_largestring_array(array) { +/// println!("{:?}", v); +/// } +/// } /// t => println!("Unsupported datatype {}", t) /// ) /// } @@ -464,11 +546,11 @@ macro_rules! downcast_dictionary_array_helper { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_dictionary_array { - ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - downcast_dictionary_array!($values => {$e} $($p => $fallback)*) + ($values:ident => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + downcast_dictionary_array!($values => {$e} $($p $(if $pred)* => $fallback)*) }; - ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => { + ($values:ident => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { match $values.data_type() { $crate::cast::__private::DataType::Dictionary(k, _) => { $crate::downcast_integer! { @@ -476,7 +558,7 @@ macro_rules! downcast_dictionary_array { k => unreachable!("unsupported dictionary key type: {}", k) } } - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } } } @@ -540,7 +622,7 @@ macro_rules! downcast_run_array_helper { /// a number of subsequent patterns to match the data type /// /// ``` -/// # use arrow_array::{Array, StringArray, downcast_run_array, cast::as_string_array}; +/// # use arrow_array::{Array, StringArray, downcast_run_array, cast::as_string_array, cast::as_largestring_array}; /// # use arrow_schema::DataType; /// /// fn print_strings(array: &dyn Array) { @@ -558,6 +640,12 @@ macro_rules! downcast_run_array_helper { /// println!("{:?}", v); /// } /// } +/// // You can also add a guard to the pattern +/// DataType::LargeUtf8 if true => { +/// for v in as_largestring_array(array) { +/// println!("{:?}", v); +/// } +/// } /// t => println!("Unsupported datatype {}", t) /// ) /// } @@ -566,11 +654,11 @@ macro_rules! downcast_run_array_helper { /// [`DataType`]: arrow_schema::DataType #[macro_export] macro_rules! downcast_run_array { - ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { - downcast_run_array!($values => {$e} $($p => $fallback)*) + ($values:ident => $e:expr, $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { + downcast_run_array!($values => {$e} $($p $(if $pred)* => $fallback)*) }; - ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => { + ($values:ident => $e:block $($p:pat $(if $pred:expr)* => $fallback:expr $(,)*)*) => { match $values.data_type() { $crate::cast::__private::DataType::RunEndEncoded(k, _) => { $crate::downcast_run_end_index! { @@ -578,7 +666,7 @@ macro_rules! downcast_run_array { k => unreachable!("unsupported run end index type: {}", k) } } - $($p => $fallback,)* + $($p $(if $pred)* => $fallback,)* } } } @@ -832,6 +920,14 @@ pub trait AsArray: private::Sealed { self.as_list_opt().expect("list array") } + /// Downcast this to a [`GenericListViewArray`] returning `None` if not possible + fn as_list_view_opt(&self) -> Option<&GenericListViewArray>; + + /// Downcast this to a [`GenericListViewArray`] panicking if not possible + fn as_list_view(&self) -> &GenericListViewArray { + self.as_list_view_opt().expect("list view array") + } + /// Downcast this to a [`FixedSizeBinaryArray`] returning `None` if not possible fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray>; @@ -866,6 +962,14 @@ pub trait AsArray: private::Sealed { self.as_dictionary_opt().expect("dictionary array") } + /// Downcast this to a [`RunArray`] returning `None` if not possible + fn as_run_opt(&self) -> Option<&RunArray>; + + /// Downcast this to a [`RunArray`] panicking if not possible + fn as_run(&self) -> &RunArray { + self.as_run_opt().expect("run array") + } + /// Downcasts this to a [`AnyDictionaryArray`] returning `None` if not possible fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray>; @@ -905,6 +1009,10 @@ impl AsArray for dyn Array + '_ { self.as_any().downcast_ref() } + fn as_list_view_opt(&self) -> Option<&GenericListViewArray> { + self.as_any().downcast_ref() + } + fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> { self.as_any().downcast_ref() } @@ -921,6 +1029,10 @@ impl AsArray for dyn Array + '_ { self.as_any().downcast_ref() } + fn as_run_opt(&self) -> Option<&RunArray> { + self.as_any().downcast_ref() + } + fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> { let array = self; downcast_dictionary_array! { @@ -960,6 +1072,10 @@ impl AsArray for ArrayRef { self.as_ref().as_list_opt() } + fn as_list_view_opt(&self) -> Option<&GenericListViewArray> { + self.as_ref().as_list_view_opt() + } + fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> { self.as_ref().as_fixed_size_binary_opt() } @@ -979,15 +1095,23 @@ impl AsArray for ArrayRef { fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> { self.as_ref().as_any_dictionary_opt() } + + fn as_run_opt(&self) -> Option<&RunArray> { + self.as_ref().as_run_opt() + } + + fn as_string_opt(&self) -> Option<&GenericStringArray> { + self.as_ref().as_string_opt() + } } #[cfg(test)] mod tests { + use super::*; use arrow_buffer::i256; + use arrow_schema::DataType; use std::sync::Arc; - use super::*; - #[test] fn test_as_primitive_array_ref() { let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect(); @@ -1019,4 +1143,46 @@ mod tests { let a = Decimal256Array::from_iter_values([1, 2, 4, 5].into_iter().map(i256::from_i128)); assert!(!as_primitive_array::(&a).is_empty()); } + + #[test] + fn downcast_integer_array_should_match_only_integers() { + let i32_array: ArrayRef = Arc::new(Int32Array::new_null(1)); + let i32_array_ref = &i32_array; + downcast_integer_array!( + i32_array_ref => { + assert_eq!(i32_array_ref.null_count(), 1); + }, + _ => panic!("unexpected data type") + ); + } + + #[test] + fn downcast_integer_array_should_not_match_primitive_that_are_not_integers() { + let array: ArrayRef = Arc::new(Float32Array::new_null(1)); + let array_ref = &array; + downcast_integer_array!( + array_ref => { + panic!("unexpected data type {}", array_ref.data_type()) + }, + DataType::Float32 => { + assert_eq!(array_ref.null_count(), 1); + }, + _ => panic!("unexpected data type") + ); + } + + #[test] + fn downcast_integer_array_should_not_match_non_primitive() { + let array: ArrayRef = Arc::new(StringArray::new_null(1)); + let array_ref = &array; + downcast_integer_array!( + array_ref => { + panic!("unexpected data type {}", array_ref.data_type()) + }, + DataType::Utf8 => { + assert_eq!(array_ref.null_count(), 1); + }, + _ => panic!("unexpected data type") + ); + } } diff --git a/arrow-array/src/ffi.rs b/arrow-array/src/ffi.rs index 144f2a21afec..f3c34f6ccd13 100644 --- a/arrow-array/src/ffi.rs +++ b/arrow-array/src/ffi.rs @@ -1298,12 +1298,12 @@ mod tests_to_then_from_ffi { mod tests_from_ffi { use std::sync::Arc; - use arrow_buffer::{bit_util, buffer::Buffer, MutableBuffer, OffsetBuffer}; + use arrow_buffer::{bit_util, buffer::Buffer}; use arrow_data::transform::MutableArrayData; use arrow_data::ArrayData; use arrow_schema::{DataType, Field}; - use super::{ImportedArrowArray, Result}; + use super::Result; use crate::builder::GenericByteViewBuilder; use crate::types::{BinaryViewType, ByteViewType, Int32Type, StringViewType}; use crate::{ @@ -1507,7 +1507,11 @@ mod tests_from_ffi { } #[test] + #[cfg(not(feature = "force_validate"))] fn test_empty_string_with_non_zero_offset() -> Result<()> { + use super::ImportedArrowArray; + use arrow_buffer::{MutableBuffer, OffsetBuffer}; + // Simulate an empty string array with a non-zero offset from a producer let data: Buffer = MutableBuffer::new(0).into(); let offsets = OffsetBuffer::new(vec![123].into()); @@ -1572,7 +1576,7 @@ mod tests_from_ffi { let mut strings = vec![]; for i in 0..1000 { - strings.push(format!("string: {}", i)); + strings.push(format!("string: {i}")); } let string_array = StringArray::from(strings); diff --git a/arrow-array/src/lib.rs b/arrow-array/src/lib.rs index 0fc9d30ab6e3..91696540d219 100644 --- a/arrow-array/src/lib.rs +++ b/arrow-array/src/lib.rs @@ -221,6 +221,11 @@ //! [DataFusion]: https://github.com/apache/arrow-datafusion //! [RecordBatchStream]: https://docs.rs/datafusion/latest/datafusion/execution/trait.RecordBatchStream.html +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![deny(rustdoc::broken_intra_doc_links)] #![warn(missing_docs)] diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs index 8958ca6fae62..c1023b739081 100644 --- a/arrow-array/src/record_batch.rs +++ b/arrow-array/src/record_batch.rs @@ -18,8 +18,9 @@ //! A two-dimensional batch of column-oriented data with a defined //! [schema](arrow_schema::Schema). +use crate::cast::AsArray; use crate::{new_empty_array, Array, ArrayRef, StructArray}; -use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaBuilder, SchemaRef}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaBuilder, SchemaRef}; use std::ops::Index; use std::sync::Arc; @@ -64,7 +65,7 @@ pub trait RecordBatchWriter { /// Support for limited data types is available. The macro will return a compile error if an unsupported data type is used. /// Presently supported data types are: /// - `Boolean`, `Null` -/// - `Decimal128`, `Decimal256` +/// - `Decimal32`, `Decimal64`, `Decimal128`, `Decimal256` /// - `Float16`, `Float32`, `Float64` /// - `Int8`, `Int16`, `Int32`, `Int64` /// - `UInt8`, `UInt16`, `UInt32`, `UInt64` @@ -106,6 +107,8 @@ macro_rules! create_array { (@from DurationMillisecond) => { $crate::DurationMillisecondArray }; (@from DurationMicrosecond) => { $crate::DurationMicrosecondArray }; (@from DurationNanosecond) => { $crate::DurationNanosecondArray }; + (@from Decimal32) => { $crate::Decimal32Array }; + (@from Decimal64) => { $crate::Decimal64Array }; (@from Decimal128) => { $crate::Decimal128Array }; (@from Decimal256) => { $crate::Decimal256Array }; (@from TimestampSecond) => { $crate::TimestampSecondArray }; @@ -210,10 +213,11 @@ impl RecordBatch { /// Creates a `RecordBatch` from a schema and columns. /// /// Expects the following: - /// * the vec of columns to not be empty - /// * the schema and column data types to have equal lengths - /// and match - /// * each array in columns to have the same length + /// + /// * `!columns.is_empty()` + /// * `schema.fields.len() == columns.len()` + /// * `schema.fields[i].data_type() == columns[i].data_type()` + /// * `columns[i].len() == columns[j].len()` /// /// If the conditions are not met, an error is returned. /// @@ -239,6 +243,33 @@ impl RecordBatch { Self::try_new_impl(schema, columns, &options) } + /// Creates a `RecordBatch` from a schema and columns, without validation. + /// + /// See [`Self::try_new`] for the checked version. + /// + /// # Safety + /// + /// Expects the following: + /// + /// * `schema.fields.len() == columns.len()` + /// * `schema.fields[i].data_type() == columns[i].data_type()` + /// * `columns[i].len() == row_count` + /// + /// Note: if the schema does not match the underlying data exactly, it can lead to undefined + /// behavior, for example, via conversion to a `StructArray`, which in turn could lead + /// to incorrect access. + pub unsafe fn new_unchecked( + schema: SchemaRef, + columns: Vec>, + row_count: usize, + ) -> Self { + Self { + schema, + columns, + row_count, + } + } + /// Creates a `RecordBatch` from a schema and columns, with additional options, /// such as whether to strictly validate field names. /// @@ -339,10 +370,17 @@ impl RecordBatch { }) } + /// Return the schema, columns and row count of this [`RecordBatch`] + pub fn into_parts(self) -> (SchemaRef, Vec, usize) { + (self.schema, self.columns, self.row_count) + } + /// Override the schema of this [`RecordBatch`] /// /// Returns an error if `schema` is not a superset of the current schema /// as determined by [`Schema::contains`] + /// + /// See also [`Self::schema_metadata_mut`]. pub fn with_schema(self, schema: SchemaRef) -> Result { if !schema.contains(self.schema.as_ref()) { return Err(ArrowError::SchemaError(format!( @@ -368,6 +406,28 @@ impl RecordBatch { &self.schema } + /// Mutable access to the metadata of the schema. + /// + /// This allows you to modify [`Schema::metadata`] of [`Self::schema`] in a convenient and fast way. + /// + /// Note this will clone the entire underlying `Schema` object if it is currently shared + /// + /// # Example + /// ``` + /// # use std::sync::Arc; + /// # use arrow_array::{record_batch, RecordBatch}; + /// let mut batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap(); + /// // Initially, the metadata is empty + /// assert!(batch.schema().metadata().get("key").is_none()); + /// // Insert a key-value pair into the metadata + /// batch.schema_metadata_mut().insert("key".into(), "value".into()); + /// assert_eq!(batch.schema().metadata().get("key"), Some(&String::from("value"))); + /// ``` + pub fn schema_metadata_mut(&mut self) -> &mut std::collections::HashMap { + let schema = Arc::make_mut(&mut self.schema); + &mut schema.metadata + } + /// Projects the schema onto the specified columns pub fn project(&self, indices: &[usize]) -> Result { let projected_schema = self.schema.project(indices)?; @@ -394,6 +454,108 @@ impl RecordBatch { ) } + /// Normalize a semi-structured [`RecordBatch`] into a flat table. + /// + /// Nested [`Field`]s will generate names separated by `separator`, up to a depth of `max_level` + /// (unlimited if `None`). + /// + /// e.g. given a [`RecordBatch`] with schema: + /// + /// ```text + /// "foo": StructArray<"bar": Utf8> + /// ``` + /// + /// A separator of `"."` would generate a batch with the schema: + /// + /// ```text + /// "foo.bar": Utf8 + /// ``` + /// + /// Note that giving a depth of `Some(0)` to `max_level` is the same as passing in `None`; + /// it will be treated as unlimited. + /// + /// # Example + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow_array::{ArrayRef, Int64Array, StringArray, StructArray, RecordBatch}; + /// # use arrow_schema::{DataType, Field, Fields, Schema}; + /// # + /// let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""])); + /// let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)])); + /// + /// let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true)); + /// let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true)); + /// + /// let a = Arc::new(StructArray::from(vec![ + /// (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef), + /// (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef), + /// ])); + /// + /// let schema = Schema::new(vec![ + /// Field::new( + /// "a", + /// DataType::Struct(Fields::from(vec![animals_field, n_legs_field])), + /// false, + /// ) + /// ]); + /// + /// let normalized = RecordBatch::try_new(Arc::new(schema), vec![a]) + /// .expect("valid conversion") + /// .normalize(".", None) + /// .expect("valid normalization"); + /// + /// let expected = RecordBatch::try_from_iter_with_nullable(vec![ + /// ("a.animals", animals.clone(), true), + /// ("a.n_legs", n_legs.clone(), true), + /// ]) + /// .expect("valid conversion"); + /// + /// assert_eq!(expected, normalized); + /// ``` + pub fn normalize(&self, separator: &str, max_level: Option) -> Result { + let max_level = match max_level.unwrap_or(usize::MAX) { + 0 => usize::MAX, + val => val, + }; + let mut stack: Vec<(usize, &ArrayRef, Vec<&str>, &FieldRef)> = self + .columns + .iter() + .zip(self.schema.fields()) + .rev() + .map(|(c, f)| { + let name_vec: Vec<&str> = vec![f.name()]; + (0, c, name_vec, f) + }) + .collect(); + let mut columns: Vec = Vec::new(); + let mut fields: Vec = Vec::new(); + + while let Some((depth, c, name, field_ref)) = stack.pop() { + match field_ref.data_type() { + DataType::Struct(ff) if depth < max_level => { + // Need to zip these in reverse to maintain original order + for (cff, fff) in c.as_struct().columns().iter().zip(ff.into_iter()).rev() { + let mut name = name.clone(); + name.push(separator); + name.push(fff.name()); + stack.push((depth + 1, cff, name, fff)) + } + } + _ => { + let updated_field = Field::new( + name.concat(), + field_ref.data_type().clone(), + field_ref.is_nullable(), + ); + columns.push(c.clone()); + fields.push(Arc::new(updated_field)); + } + } + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) + } + /// Returns the number of columns in the record batch. /// /// # Example @@ -768,8 +930,6 @@ where #[cfg(test)] mod tests { - use std::collections::HashMap; - use super::*; use crate::{ BooleanArray, Int32Array, Int64Array, Int8Array, ListArray, StringArray, StringViewArray, @@ -777,6 +937,7 @@ mod tests { use arrow_buffer::{Buffer, ToByteSlice}; use arrow_data::{ArrayData, ArrayDataBuilder}; use arrow_schema::Fields; + use std::collections::HashMap; #[test] fn create_record_batch() { @@ -938,8 +1099,8 @@ mod tests { let a = Int64Array::from(vec![1, 2, 3, 4, 5]); - let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]); - assert!(batch.is_err()); + let err = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap_err(); + assert_eq!(err.to_string(), "Invalid argument error: column types must match schema types, expected Int32 but found Int64 at column index 0"); } #[test] @@ -1197,6 +1358,181 @@ mod tests { assert_ne!(batch1, batch2); } + #[test] + fn normalize_simple() { + let animals: ArrayRef = Arc::new(StringArray::from(vec!["Parrot", ""])); + let n_legs: ArrayRef = Arc::new(Int64Array::from(vec![Some(2), Some(4)])); + let year: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(2022)])); + + let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true)); + let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true)); + let year_field = Arc::new(Field::new("year", DataType::Int64, true)); + + let a = Arc::new(StructArray::from(vec![ + (animals_field.clone(), Arc::new(animals.clone()) as ArrayRef), + (n_legs_field.clone(), Arc::new(n_legs.clone()) as ArrayRef), + (year_field.clone(), Arc::new(year.clone()) as ArrayRef), + ])); + + let month = Arc::new(Int64Array::from(vec![Some(4), Some(6)])); + + let schema = Schema::new(vec![ + Field::new( + "a", + DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])), + false, + ), + Field::new("month", DataType::Int64, true), + ]); + + let normalized = + RecordBatch::try_new(Arc::new(schema.clone()), vec![a.clone(), month.clone()]) + .expect("valid conversion") + .normalize(".", Some(0)) + .expect("valid normalization"); + + let expected = RecordBatch::try_from_iter_with_nullable(vec![ + ("a.animals", animals.clone(), true), + ("a.n_legs", n_legs.clone(), true), + ("a.year", year.clone(), true), + ("month", month.clone(), true), + ]) + .expect("valid conversion"); + + assert_eq!(expected, normalized); + + // check 0 and None have the same effect + let normalized = RecordBatch::try_new(Arc::new(schema), vec![a, month.clone()]) + .expect("valid conversion") + .normalize(".", None) + .expect("valid normalization"); + + assert_eq!(expected, normalized); + } + + #[test] + fn normalize_nested() { + // Initialize schema + let a = Arc::new(Field::new("a", DataType::Int64, true)); + let b = Arc::new(Field::new("b", DataType::Int64, false)); + let c = Arc::new(Field::new("c", DataType::Int64, true)); + + let one = Arc::new(Field::new( + "1", + DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])), + false, + )); + let two = Arc::new(Field::new( + "2", + DataType::Struct(Fields::from(vec![a.clone(), b.clone(), c.clone()])), + true, + )); + + let exclamation = Arc::new(Field::new( + "!", + DataType::Struct(Fields::from(vec![one.clone(), two.clone()])), + false, + )); + + let schema = Schema::new(vec![exclamation.clone()]); + + // Initialize fields + let a_field = Int64Array::from(vec![Some(0), Some(1)]); + let b_field = Int64Array::from(vec![Some(2), Some(3)]); + let c_field = Int64Array::from(vec![None, Some(4)]); + + let one_field = StructArray::from(vec![ + (a.clone(), Arc::new(a_field.clone()) as ArrayRef), + (b.clone(), Arc::new(b_field.clone()) as ArrayRef), + (c.clone(), Arc::new(c_field.clone()) as ArrayRef), + ]); + let two_field = StructArray::from(vec![ + (a.clone(), Arc::new(a_field.clone()) as ArrayRef), + (b.clone(), Arc::new(b_field.clone()) as ArrayRef), + (c.clone(), Arc::new(c_field.clone()) as ArrayRef), + ]); + + let exclamation_field = Arc::new(StructArray::from(vec![ + (one.clone(), Arc::new(one_field) as ArrayRef), + (two.clone(), Arc::new(two_field) as ArrayRef), + ])); + + // Normalize top level + let normalized = + RecordBatch::try_new(Arc::new(schema.clone()), vec![exclamation_field.clone()]) + .expect("valid conversion") + .normalize(".", Some(1)) + .expect("valid normalization"); + + let expected = RecordBatch::try_from_iter_with_nullable(vec![ + ( + "!.1", + Arc::new(StructArray::from(vec![ + (a.clone(), Arc::new(a_field.clone()) as ArrayRef), + (b.clone(), Arc::new(b_field.clone()) as ArrayRef), + (c.clone(), Arc::new(c_field.clone()) as ArrayRef), + ])) as ArrayRef, + false, + ), + ( + "!.2", + Arc::new(StructArray::from(vec![ + (a.clone(), Arc::new(a_field.clone()) as ArrayRef), + (b.clone(), Arc::new(b_field.clone()) as ArrayRef), + (c.clone(), Arc::new(c_field.clone()) as ArrayRef), + ])) as ArrayRef, + true, + ), + ]) + .expect("valid conversion"); + + assert_eq!(expected, normalized); + + // Normalize all levels + let normalized = RecordBatch::try_new(Arc::new(schema), vec![exclamation_field]) + .expect("valid conversion") + .normalize(".", None) + .expect("valid normalization"); + + let expected = RecordBatch::try_from_iter_with_nullable(vec![ + ("!.1.a", Arc::new(a_field.clone()) as ArrayRef, true), + ("!.1.b", Arc::new(b_field.clone()) as ArrayRef, false), + ("!.1.c", Arc::new(c_field.clone()) as ArrayRef, true), + ("!.2.a", Arc::new(a_field.clone()) as ArrayRef, true), + ("!.2.b", Arc::new(b_field.clone()) as ArrayRef, false), + ("!.2.c", Arc::new(c_field.clone()) as ArrayRef, true), + ]) + .expect("valid conversion"); + + assert_eq!(expected, normalized); + } + + #[test] + fn normalize_empty() { + let animals_field = Arc::new(Field::new("animals", DataType::Utf8, true)); + let n_legs_field = Arc::new(Field::new("n_legs", DataType::Int64, true)); + let year_field = Arc::new(Field::new("year", DataType::Int64, true)); + + let schema = Schema::new(vec![ + Field::new( + "a", + DataType::Struct(Fields::from(vec![animals_field, n_legs_field, year_field])), + false, + ), + Field::new("month", DataType::Int64, true), + ]); + + let normalized = RecordBatch::new_empty(Arc::new(schema.clone())) + .normalize(".", Some(0)) + .expect("valid normalization"); + + let expected = RecordBatch::new_empty(Arc::new( + schema.normalize(".", Some(0)).expect("valid normalization"), + )); + + assert_eq!(expected, normalized); + } + #[test] fn project() { let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); diff --git a/arrow-array/src/run_iterator.rs b/arrow-array/src/run_iterator.rs index 2922bf04dd2f..4fb0eef32eca 100644 --- a/arrow-array/src/run_iterator.rs +++ b/arrow-array/src/run_iterator.rs @@ -172,7 +172,7 @@ where #[cfg(test)] mod tests { - use rand::{seq::SliceRandom, thread_rng, Rng}; + use rand::{rng, seq::SliceRandom, Rng}; use crate::{ array::{Int32Array, StringArray}, @@ -200,7 +200,7 @@ mod tests { ]; let mut result: Vec> = Vec::with_capacity(size); let mut ix = 0; - let mut rng = thread_rng(); + let mut rng = rng(); // run length can go up to 8. Cap the max run length for smaller arrays to size / 2. let max_run_length = 8_usize.min(1_usize.max(size / 2)); while result.len() < size { @@ -209,7 +209,7 @@ mod tests { seed.shuffle(&mut rng); } // repeat the items between 1 and 8 times. Cap the length for smaller sized arrays - let num = max_run_length.min(rand::thread_rng().gen_range(1..=max_run_length)); + let num = max_run_length.min(rng.random_range(1..=max_run_length)); for _ in 0..num { result.push(seed[ix]); } diff --git a/arrow-array/src/temporal_conversions.rs b/arrow-array/src/temporal_conversions.rs index 23f950d55048..7a4c67602932 100644 --- a/arrow-array/src/temporal_conversions.rs +++ b/arrow-array/src/temporal_conversions.rs @@ -217,14 +217,28 @@ pub(crate) fn split_second(v: i64, base: i64) -> (i64, u32) { /// converts a `i64` representing a `duration(s)` to [`Duration`] #[inline] +#[deprecated(since = "55.2.0", note = "Use `try_duration_s_to_duration` instead")] pub fn duration_s_to_duration(v: i64) -> Duration { Duration::try_seconds(v).unwrap() } +/// converts a `i64` representing a `duration(s)` to [`Option`] +#[inline] +pub fn try_duration_s_to_duration(v: i64) -> Option { + Duration::try_seconds(v) +} + /// converts a `i64` representing a `duration(ms)` to [`Duration`] #[inline] +#[deprecated(since = "55.2.0", note = "Use `try_duration_ms_to_duration` instead")] pub fn duration_ms_to_duration(v: i64) -> Duration { - Duration::try_milliseconds(v).unwrap() + Duration::try_seconds(v).unwrap() +} + +/// converts a `i64` representing a `duration(ms)` to [`Option`] +#[inline] +pub fn try_duration_ms_to_duration(v: i64) -> Option { + Duration::try_milliseconds(v) } /// converts a `i64` representing a `duration(us)` to [`Duration`] @@ -296,8 +310,8 @@ pub fn as_time(v: i64) -> Option { pub fn as_duration(v: i64) -> Option { match T::DATA_TYPE { DataType::Duration(unit) => match unit { - TimeUnit::Second => Some(duration_s_to_duration(v)), - TimeUnit::Millisecond => Some(duration_ms_to_duration(v)), + TimeUnit::Second => try_duration_s_to_duration(v), + TimeUnit::Millisecond => try_duration_ms_to_duration(v), TimeUnit::Microsecond => Some(duration_us_to_duration(v)), TimeUnit::Nanosecond => Some(duration_ns_to_duration(v)), }, diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs index 3d8cfcdb112b..96c496a536bb 100644 --- a/arrow-array/src/types.rs +++ b/arrow-array/src/types.rs @@ -25,13 +25,16 @@ use crate::timezone::Tz; use crate::{ArrowNativeTypeOp, OffsetSizeTrait}; use arrow_buffer::{i256, Buffer, OffsetBuffer}; use arrow_data::decimal::{ - is_validate_decimal256_precision, is_validate_decimal_precision, validate_decimal256_precision, - validate_decimal_precision, + is_validate_decimal256_precision, is_validate_decimal32_precision, + is_validate_decimal64_precision, is_validate_decimal_precision, validate_decimal256_precision, + validate_decimal32_precision, validate_decimal64_precision, validate_decimal_precision, }; use arrow_data::{validate_binary_view, validate_string_view}; use arrow_schema::{ ArrowError, DataType, IntervalUnit, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, - DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DECIMAL_DEFAULT_SCALE, + DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DECIMAL32_DEFAULT_SCALE, + DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, DECIMAL64_DEFAULT_SCALE, DECIMAL64_MAX_PRECISION, + DECIMAL64_MAX_SCALE, DECIMAL_DEFAULT_SCALE, }; use chrono::{Duration, NaiveDate, NaiveDateTime}; use half::f16; @@ -68,12 +71,6 @@ pub trait ArrowPrimitiveType: primitive::PrimitiveTypeSealed + 'static { /// the corresponding Arrow data type of this primitive type. const DATA_TYPE: DataType; - /// Returns the byte width of this primitive type. - #[deprecated(since = "52.0.0", note = "Use ArrowNativeType::get_byte_width")] - fn get_byte_width() -> usize { - std::mem::size_of::() - } - /// Returns a default value of this primitive type. /// /// This is useful for aggregate array ops like `sum()`, `mean()`. @@ -1031,9 +1028,25 @@ impl Date64Type { /// # Arguments /// /// * `i` - The Date64Type to convert + #[deprecated(since = "56.0.0", note = "Use to_naive_date_opt instead.")] pub fn to_naive_date(i: ::Native) -> NaiveDate { + Self::to_naive_date_opt(i) + .unwrap_or_else(|| panic!("Date64Type::to_naive_date overflowed for date: {i}",)) + } + + /// Converts an arrow Date64Type into a chrono::NaiveDateTime if it fits in the range that chrono::NaiveDateTime can represent. + /// Returns `None` if the calculation would overflow or underflow. + /// + /// This function is able to handle dates ranging between 1677-09-21 (-9,223,372,800,000) and 2262-04-11 (9,223,286,400,000). + /// + /// # Arguments + /// + /// * `i` - The Date64Type to convert + /// + /// Returns `Some(NaiveDateTime)` if it fits, `None` otherwise. + pub fn to_naive_date_opt(i: ::Native) -> Option { let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); - epoch.add(Duration::try_milliseconds(i).unwrap()) + Duration::try_milliseconds(i).and_then(|d| epoch.checked_add_signed(d)) } /// Converts a chrono::NaiveDate into an arrow Date64Type @@ -1052,14 +1065,35 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to add + #[deprecated( + since = "56.0.0", + note = "Use `add_year_months_opt` instead, which returns an Option to handle overflow." + )] pub fn add_year_months( date: ::Native, delta: ::Native, ) -> ::Native { - let prior = Date64Type::to_naive_date(date); + Self::add_year_months_opt(date, delta).unwrap_or_else(|| { + panic!("Date64Type::add_year_months overflowed for date: {date}, delta: {delta}",) + }) + } + + /// Adds the given IntervalYearMonthType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to add + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn add_year_months_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { + let prior = Date64Type::to_naive_date_opt(date)?; let months = IntervalYearMonthType::to_months(delta); let posterior = shift_months(prior, months); - Date64Type::from_naive_date(posterior) + Some(Date64Type::from_naive_date(posterior)) } /// Adds the given IntervalDayTimeType to an arrow Date64Type @@ -1068,15 +1102,36 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to add + #[deprecated( + since = "56.0.0", + note = "Use `add_day_time_opt` instead, which returns an Option to handle overflow." + )] pub fn add_day_time( date: ::Native, delta: ::Native, ) -> ::Native { + Self::add_day_time_opt(date, delta).unwrap_or_else(|| { + panic!("Date64Type::add_day_time overflowed for date: {date}, delta: {delta:?}",) + }) + } + + /// Adds the given IntervalDayTimeType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to add + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn add_day_time_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { let (days, ms) = IntervalDayTimeType::to_parts(delta); - let res = Date64Type::to_naive_date(date); - let res = res.add(Duration::try_days(days as i64).unwrap()); - let res = res.add(Duration::try_milliseconds(ms as i64).unwrap()); - Date64Type::from_naive_date(res) + let res = Date64Type::to_naive_date_opt(date)?; + let res = res.checked_add_signed(Duration::try_days(days as i64)?)?; + let res = res.checked_add_signed(Duration::try_milliseconds(ms as i64)?)?; + Some(Date64Type::from_naive_date(res)) } /// Adds the given IntervalMonthDayNanoType to an arrow Date64Type @@ -1085,16 +1140,37 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to add + #[deprecated( + since = "56.0.0", + note = "Use `add_month_day_nano_opt` instead, which returns an Option to handle overflow." + )] pub fn add_month_day_nano( date: ::Native, delta: ::Native, ) -> ::Native { + Self::add_month_day_nano_opt(date, delta).unwrap_or_else(|| { + panic!("Date64Type::add_month_day_nano overflowed for date: {date}, delta: {delta:?}",) + }) + } + + /// Adds the given IntervalMonthDayNanoType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to add + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn add_month_day_nano_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta); - let res = Date64Type::to_naive_date(date); + let res = Date64Type::to_naive_date_opt(date)?; let res = shift_months(res, months); - let res = res.add(Duration::try_days(days as i64).unwrap()); - let res = res.add(Duration::nanoseconds(nanos)); - Date64Type::from_naive_date(res) + let res = res.checked_add_signed(Duration::try_days(days as i64)?)?; + let res = res.checked_add_signed(Duration::nanoseconds(nanos))?; + Some(Date64Type::from_naive_date(res)) } /// Subtract the given IntervalYearMonthType to an arrow Date64Type @@ -1103,14 +1179,35 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to subtract + #[deprecated( + since = "56.0.0", + note = "Use `subtract_year_months_opt` instead, which returns an Option to handle overflow." + )] pub fn subtract_year_months( date: ::Native, delta: ::Native, ) -> ::Native { - let prior = Date64Type::to_naive_date(date); + Self::subtract_year_months_opt(date, delta).unwrap_or_else(|| { + panic!("Date64Type::subtract_year_months overflowed for date: {date}, delta: {delta}",) + }) + } + + /// Subtract the given IntervalYearMonthType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to subtract + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn subtract_year_months_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { + let prior = Date64Type::to_naive_date_opt(date)?; let months = IntervalYearMonthType::to_months(-delta); let posterior = shift_months(prior, months); - Date64Type::from_naive_date(posterior) + Some(Date64Type::from_naive_date(posterior)) } /// Subtract the given IntervalDayTimeType to an arrow Date64Type @@ -1119,15 +1216,36 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to subtract + #[deprecated( + since = "56.0.0", + note = "Use `subtract_day_time_opt` instead, which returns an Option to handle overflow." + )] pub fn subtract_day_time( date: ::Native, delta: ::Native, ) -> ::Native { + Self::subtract_day_time_opt(date, delta).unwrap_or_else(|| { + panic!("Date64Type::subtract_day_time overflowed for date: {date}, delta: {delta:?}",) + }) + } + + /// Subtract the given IntervalDayTimeType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to subtract + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn subtract_day_time_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { let (days, ms) = IntervalDayTimeType::to_parts(delta); - let res = Date64Type::to_naive_date(date); - let res = res.sub(Duration::try_days(days as i64).unwrap()); - let res = res.sub(Duration::try_milliseconds(ms as i64).unwrap()); - Date64Type::from_naive_date(res) + let res = Date64Type::to_naive_date_opt(date)?; + let res = res.checked_sub_signed(Duration::try_days(days as i64)?)?; + let res = res.checked_sub_signed(Duration::try_milliseconds(ms as i64)?)?; + Some(Date64Type::from_naive_date(res)) } /// Subtract the given IntervalMonthDayNanoType to an arrow Date64Type @@ -1136,16 +1254,39 @@ impl Date64Type { /// /// * `date` - The date on which to perform the operation /// * `delta` - The interval to subtract + #[deprecated( + since = "56.0.0", + note = "Use `subtract_month_day_nano_opt` instead, which returns an Option to handle overflow." + )] pub fn subtract_month_day_nano( date: ::Native, delta: ::Native, ) -> ::Native { + Self::subtract_month_day_nano_opt(date, delta).unwrap_or_else(|| { + panic!( + "Date64Type::subtract_month_day_nano overflowed for date: {date}, delta: {delta:?}", + ) + }) + } + + /// Subtract the given IntervalMonthDayNanoType to an arrow Date64Type + /// + /// # Arguments + /// + /// * `date` - The date on which to perform the operation + /// * `delta` - The interval to subtract + /// + /// Returns `Some(Date64Type)` if it fits, `None` otherwise. + pub fn subtract_month_day_nano_opt( + date: ::Native, + delta: ::Native, + ) -> Option<::Native> { let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(delta); - let res = Date64Type::to_naive_date(date); + let res = Date64Type::to_naive_date_opt(date)?; let res = shift_months(res, -months); - let res = res.sub(Duration::try_days(days as i64).unwrap()); - let res = res.sub(Duration::nanoseconds(nanos)); - Date64Type::from_naive_date(res) + let res = res.checked_sub_signed(Duration::try_days(days as i64)?)?; + let res = res.checked_sub_signed(Duration::nanoseconds(nanos))?; + Some(Date64Type::from_naive_date(res)) } } @@ -1156,6 +1297,8 @@ mod decimal { use super::*; pub trait DecimalTypeSealed {} + impl DecimalTypeSealed for Decimal32Type {} + impl DecimalTypeSealed for Decimal64Type {} impl DecimalTypeSealed for Decimal128Type {} impl DecimalTypeSealed for Decimal256Type {} } @@ -1163,10 +1306,12 @@ mod decimal { /// A trait over the decimal types, used by [`PrimitiveArray`] to provide a generic /// implementation across the various decimal types /// -/// Implemented by [`Decimal128Type`] and [`Decimal256Type`] for [`Decimal128Array`] -/// and [`Decimal256Array`] respectively +/// Implemented by [`Decimal32Type`], [`Decimal64Type`], [`Decimal128Type`] and [`Decimal256Type`] +/// for [`Decimal32Array`], [`Decimal64Array`], [`Decimal128Array`] and [`Decimal256Array`] respectively /// /// [`PrimitiveArray`]: crate::array::PrimitiveArray +/// [`Decimal32Array`]: crate::array::Decimal32Array +/// [`Decimal64Array`]: crate::array::Decimal64Array /// [`Decimal128Array`]: crate::array::Decimal128Array /// [`Decimal256Array`]: crate::array::Decimal256Array pub trait DecimalType: @@ -1183,7 +1328,7 @@ pub trait DecimalType: /// Default values for [`DataType`] const DEFAULT_TYPE: DataType; - /// "Decimal128" or "Decimal256", for use in error messages + /// "Decimal32", "Decimal64", "Decimal128" or "Decimal256", for use in error messages const PREFIX: &'static str; /// Formats the decimal value with the provided precision and scale @@ -1236,6 +1381,74 @@ pub fn validate_decimal_precision_and_scale( Ok(()) } +/// The decimal type for a Decimal32Array +#[derive(Debug)] +pub struct Decimal32Type {} + +impl DecimalType for Decimal32Type { + const BYTE_LENGTH: usize = 4; + const MAX_PRECISION: u8 = DECIMAL32_MAX_PRECISION; + const MAX_SCALE: i8 = DECIMAL32_MAX_SCALE; + const TYPE_CONSTRUCTOR: fn(u8, i8) -> DataType = DataType::Decimal32; + const DEFAULT_TYPE: DataType = + DataType::Decimal32(DECIMAL32_MAX_PRECISION, DECIMAL32_DEFAULT_SCALE); + const PREFIX: &'static str = "Decimal32"; + + fn format_decimal(value: Self::Native, precision: u8, scale: i8) -> String { + format_decimal_str(&value.to_string(), precision as usize, scale) + } + + fn validate_decimal_precision(num: i32, precision: u8) -> Result<(), ArrowError> { + validate_decimal32_precision(num, precision) + } + + fn is_valid_decimal_precision(value: Self::Native, precision: u8) -> bool { + is_validate_decimal32_precision(value, precision) + } +} + +impl ArrowPrimitiveType for Decimal32Type { + type Native = i32; + + const DATA_TYPE: DataType = ::DEFAULT_TYPE; +} + +impl primitive::PrimitiveTypeSealed for Decimal32Type {} + +/// The decimal type for a Decimal64Array +#[derive(Debug)] +pub struct Decimal64Type {} + +impl DecimalType for Decimal64Type { + const BYTE_LENGTH: usize = 8; + const MAX_PRECISION: u8 = DECIMAL64_MAX_PRECISION; + const MAX_SCALE: i8 = DECIMAL64_MAX_SCALE; + const TYPE_CONSTRUCTOR: fn(u8, i8) -> DataType = DataType::Decimal64; + const DEFAULT_TYPE: DataType = + DataType::Decimal64(DECIMAL64_MAX_PRECISION, DECIMAL64_DEFAULT_SCALE); + const PREFIX: &'static str = "Decimal64"; + + fn format_decimal(value: Self::Native, precision: u8, scale: i8) -> String { + format_decimal_str(&value.to_string(), precision as usize, scale) + } + + fn validate_decimal_precision(num: i64, precision: u8) -> Result<(), ArrowError> { + validate_decimal64_precision(num, precision) + } + + fn is_valid_decimal_precision(value: Self::Native, precision: u8) -> bool { + is_validate_decimal64_precision(value, precision) + } +} + +impl ArrowPrimitiveType for Decimal64Type { + type Native = i64; + + const DATA_TYPE: DataType = ::DEFAULT_TYPE; +} + +impl primitive::PrimitiveTypeSealed for Decimal64Type {} + /// The decimal type for a Decimal128Array #[derive(Debug)] pub struct Decimal128Type {} diff --git a/arrow-avro/Cargo.toml b/arrow-avro/Cargo.toml index c103c2ecc0f3..e2280b251ff6 100644 --- a/arrow-avro/Cargo.toml +++ b/arrow-avro/Cargo.toml @@ -30,13 +30,16 @@ rust-version = { workspace = true } [lib] name = "arrow_avro" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + [features] -default = ["deflate", "snappy", "zstd"] +default = ["deflate", "snappy", "zstd", "bzip2", "xz"] deflate = ["flate2"] snappy = ["snap", "crc"] +canonical_extension_types = ["arrow-schema/canonical_extension_types"] [dependencies] arrow-schema = { workspace = true } @@ -44,12 +47,30 @@ arrow-buffer = { workspace = true } arrow-array = { workspace = true } serde_json = { version = "1.0", default-features = false, features = ["std"] } serde = { version = "1.0.188", features = ["derive"] } -flate2 = { version = "1.0", default-features = false, features = ["rust_backend"], optional = true } +flate2 = { version = "1.0", default-features = false, features = [ + "rust_backend", +], optional = true } 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 } - +uuid = "1.17" [dev-dependencies] -rand = { version = "0.8", default-features = false, features = ["std", "std_rng"] } +arrow-data = { workspace = true } +rand = { version = "0.9.1", default-features = false, features = [ + "std", + "std_rng", + "thread_rng", +] } +criterion = { version = "0.6.0", default-features = false } +tempfile = "3.3" +arrow = { workspace = true } +futures = "0.3.31" +bytes = "1.10.1" +async-stream = "0.3.6" +[[bench]] +name = "avro_reader" +harness = false diff --git a/arrow-avro/benches/avro_reader.rs b/arrow-avro/benches/avro_reader.rs new file mode 100644 index 000000000000..2f2a3a10dbf3 --- /dev/null +++ b/arrow-avro/benches/avro_reader.rs @@ -0,0 +1,268 @@ +// 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. + +//! Comprehensive benchmarks comparing StringArray vs StringViewArray performance +//! +//! This benchmark suite compares the performance characteristics of StringArray vs +//! StringViewArray across three key dimensions: +//! 1. Array creation performance +//! 2. String value access operations +//! 3. Avro file reading with each array type + +use std::fs::File; +use std::io::{BufReader, Read, Write}; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::RecordBatch; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow_array::{ArrayRef, Int32Array, StringArray, StringViewArray}; +use arrow_schema::ArrowError; +use criterion::*; +use tempfile::NamedTempFile; + +fn create_test_data(count: usize, str_length: usize) -> Vec { + (0..count) + .map(|i| format!("str_{i}") + &"a".repeat(str_length)) + .collect() +} + +fn create_avro_test_file(row_count: usize, str_length: usize) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("string_field", DataType::Utf8, false), + Field::new("int_field", DataType::Int32, false), + ])); + + let strings = create_test_data(row_count, str_length); + let string_array = StringArray::from_iter(strings.iter().map(|s| Some(s.as_str()))); + let int_array = Int32Array::from_iter_values(0..row_count as i32); + let _batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(string_array) as ArrayRef, + Arc::new(int_array) as ArrayRef, + ], + )?; + + let temp_file = NamedTempFile::new()?; + + let mut file = temp_file.reopen()?; + + file.write_all(b"AVRO")?; + + for (i, string) in strings.iter().enumerate().take(row_count) { + let s = string.as_bytes(); + let len = s.len() as u32; + file.write_all(&len.to_le_bytes())?; + file.write_all(s)?; + file.write_all(&(i as i32).to_le_bytes())?; + } + + file.flush()?; + Ok(temp_file) +} + +fn read_avro_test_file( + file_path: &std::path::Path, + use_utf8view: bool, +) -> Result { + let file = File::open(file_path)?; + let mut reader = BufReader::new(file); + + let mut header = [0u8; 4]; + reader.read_exact(&mut header)?; + + let mut strings = Vec::new(); + let mut ints = Vec::new(); + + loop { + let mut len_bytes = [0u8; 4]; + if reader.read_exact(&mut len_bytes).is_err() { + break; // End of file + } + + let len = u32::from_le_bytes(len_bytes) as usize; + let mut buf = vec![0u8; len]; + reader.read_exact(&mut buf)?; + + let s = String::from_utf8(buf) + .map_err(|e| ArrowError::ParseError(format!("Invalid UTF-8: {e}")))?; + + strings.push(s); + + let mut int_bytes = [0u8; 4]; + reader.read_exact(&mut int_bytes)?; + ints.push(i32::from_le_bytes(int_bytes)); + } + + let string_array: ArrayRef = if use_utf8view { + Arc::new(StringViewArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + )) + } else { + Arc::new(StringArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + )) + }; + + let int_array: ArrayRef = Arc::new(Int32Array::from(ints)); + + let schema = Arc::new(Schema::new(vec![ + if use_utf8view { + Field::new("string_field", DataType::Utf8View, false) + } else { + Field::new("string_field", DataType::Utf8, false) + }, + Field::new("int_field", DataType::Int32, false), + ])); + + RecordBatch::try_new(schema, vec![string_array, int_array]) +} + +fn bench_array_creation(c: &mut Criterion) { + let mut group = c.benchmark_group("array_creation"); + group.sample_size(20); + group.measurement_time(Duration::from_secs(5)); + + for &str_length in &[10, 100, 1000] { + let data = create_test_data(10000, str_length); + let row_count = 1000; + + group.bench_function(format!("string_array_{str_length}_chars"), |b| { + b.iter(|| { + let string_array = + StringArray::from_iter(data[0..row_count].iter().map(|s| Some(s.as_str()))); + let int_array = Int32Array::from_iter_values(0..row_count as i32); + + let schema = Arc::new(Schema::new(vec![ + Field::new("string_field", DataType::Utf8, false), + Field::new("int_field", DataType::Int32, false), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(string_array) as ArrayRef, + Arc::new(int_array) as ArrayRef, + ], + ) + .unwrap(); + + std::hint::black_box(batch) + }) + }); + + group.bench_function(format!("string_view_{str_length}_chars"), |b| { + b.iter(|| { + let string_array = + StringViewArray::from_iter(data[0..row_count].iter().map(|s| Some(s.as_str()))); + let int_array = Int32Array::from_iter_values(0..row_count as i32); + + let schema = Arc::new(Schema::new(vec![ + Field::new("string_field", DataType::Utf8View, false), + Field::new("int_field", DataType::Int32, false), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(string_array) as ArrayRef, + Arc::new(int_array) as ArrayRef, + ], + ) + .unwrap(); + + std::hint::black_box(batch) + }) + }); + } + + group.finish(); +} + +fn bench_string_operations(c: &mut Criterion) { + let mut group = c.benchmark_group("string_operations"); + group.sample_size(20); + group.measurement_time(Duration::from_secs(5)); + + for &str_length in &[10, 100, 1000] { + let data = create_test_data(10000, str_length); + let rows = 1000; + + let string_array = StringArray::from_iter(data[0..rows].iter().map(|s| Some(s.as_str()))); + let string_view_array = + StringViewArray::from_iter(data[0..rows].iter().map(|s| Some(s.as_str()))); + + group.bench_function(format!("string_array_value_{str_length}_chars"), |b| { + b.iter(|| { + let mut sum_len = 0; + for i in 0..rows { + sum_len += string_array.value(i).len(); + } + std::hint::black_box(sum_len) + }) + }); + + group.bench_function(format!("string_view_value_{str_length}_chars"), |b| { + b.iter(|| { + let mut sum_len = 0; + for i in 0..rows { + sum_len += string_view_array.value(i).len(); + } + std::hint::black_box(sum_len) + }) + }); + } + + group.finish(); +} + +fn bench_avro_reader(c: &mut Criterion) { + let mut group = c.benchmark_group("avro_reader"); + group.sample_size(20); + group.measurement_time(Duration::from_secs(5)); + + for &str_length in &[10, 100, 1000] { + let row_count = 1000; + let temp_file = create_avro_test_file(row_count, str_length).unwrap(); + let file_path = temp_file.path(); + + group.bench_function(format!("string_array_{str_length}_chars"), |b| { + b.iter(|| { + let batch = read_avro_test_file(file_path, false).unwrap(); + std::hint::black_box(batch) + }) + }); + + group.bench_function(format!("string_view_{str_length}_chars"), |b| { + b.iter(|| { + let batch = read_avro_test_file(file_path, true).unwrap(); + std::hint::black_box(batch) + }) + }); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_array_creation, + bench_string_operations, + bench_avro_reader +); +criterion_main!(benches); diff --git a/arrow-avro/examples/read_with_utf8view.rs b/arrow-avro/examples/read_with_utf8view.rs new file mode 100644 index 000000000000..707be575168a --- /dev/null +++ b/arrow-avro/examples/read_with_utf8view.rs @@ -0,0 +1,108 @@ +// 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. + +//! This example demonstrates how to use Utf8View support in the Arrow Avro reader +//! +//! It reads an Avro file with string data twice - once with regular StringArray +//! and once with StringViewArray - and compares the performance. + +use std::env; +use std::fs::File; +use std::io::BufReader; +use std::time::Instant; + +use arrow_array::{RecordBatch, StringArray, StringViewArray}; +use arrow_avro::reader::ReaderBuilder; + +fn main() -> Result<(), Box> { + let args: Vec = env::args().collect(); + let file_path = if args.len() > 1 { + &args[1] + } else { + eprintln!("No file specified, please provide an Avro file path"); + eprintln!("Usage: {} ", args[0]); + return Ok(()); + }; + + let file = File::open(file_path)?; + let file_for_view = file.try_clone()?; + + let start = Instant::now(); + let reader = BufReader::new(file); + let avro_reader = ReaderBuilder::new().build(reader)?; + let schema = avro_reader.schema(); + let batches: Vec = avro_reader.collect::>()?; + let regular_duration = start.elapsed(); + + let start = Instant::now(); + let reader_view = BufReader::new(file_for_view); + let avro_reader_view = ReaderBuilder::new() + .with_utf8_view(true) + .build(reader_view)?; + let batches_view: Vec = avro_reader_view.collect::>()?; + let view_duration = start.elapsed(); + + let num_rows = batches.iter().map(|b| b.num_rows()).sum::(); + + println!("Read {num_rows} rows from {file_path}"); + println!("Reading with StringArray: {regular_duration:?}"); + println!("Reading with StringViewArray: {view_duration:?}"); + + if regular_duration > view_duration { + println!( + "StringViewArray was {:.2}x faster", + regular_duration.as_secs_f64() / view_duration.as_secs_f64() + ); + } else { + println!( + "StringArray was {:.2}x faster", + view_duration.as_secs_f64() / regular_duration.as_secs_f64() + ); + } + + if batches.is_empty() { + println!("No data read from file."); + return Ok(()); + } + + // Inspect the first batch from each run to show the array types + let batch = &batches[0]; + let batch_view = &batches_view[0]; + + for (i, field) in schema.fields().iter().enumerate() { + let col = batch.column(i); + let col_view = batch_view.column(i); + + if col.as_any().is::() { + println!( + "Column {} '{}' is StringArray in regular version", + i, + field.name() + ); + } + + if col_view.as_any().is::() { + println!( + "Column {} '{}' is StringViewArray in utf8view version", + i, + field.name() + ); + } + } + + Ok(()) +} diff --git a/arrow-avro/src/codec.rs b/arrow-avro/src/codec.rs index 2ac1ad038bd7..bd265503d755 100644 --- a/arrow-avro/src/codec.rs +++ b/arrow-avro/src/codec.rs @@ -17,7 +17,8 @@ use crate::schema::{Attributes, ComplexType, PrimitiveType, Record, Schema, TypeName}; use arrow_schema::{ - ArrowError, DataType, Field, FieldRef, IntervalUnit, SchemaBuilder, SchemaRef, TimeUnit, + ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, SchemaBuilder, SchemaRef, + TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, }; use std::borrow::Cow; use std::collections::HashMap; @@ -36,6 +37,14 @@ pub enum Nullability { NullSecond, } +#[cfg(feature = "canonical_extension_types")] +fn with_extension_type(codec: &Codec, field: Field) -> Field { + match codec { + Codec::Uuid => field.with_extension_type(arrow_schema::extension::Uuid), + _ => field, + } +} + /// An Avro datatype mapped to the arrow data model #[derive(Debug, Clone)] pub struct AvroDataType { @@ -45,16 +54,45 @@ pub struct AvroDataType { } impl AvroDataType { + /// Create a new [`AvroDataType`] with the given parts. + pub fn new( + codec: Codec, + metadata: HashMap, + nullability: Option, + ) -> Self { + AvroDataType { + codec, + metadata, + nullability, + } + } + /// Returns an arrow [`Field`] with the given name pub fn field_with_name(&self, name: &str) -> Field { - let d = self.codec.data_type(); - Field::new(name, d, self.nullability.is_some()).with_metadata(self.metadata.clone()) + let nullable = self.nullability.is_some(); + let data_type = self.codec.data_type(); + let field = Field::new(name, data_type, nullable).with_metadata(self.metadata.clone()); + #[cfg(feature = "canonical_extension_types")] + return with_extension_type(&self.codec, field); + #[cfg(not(feature = "canonical_extension_types"))] + field } + /// Returns a reference to the codec used by this data type + /// + /// The codec determines how Avro data is encoded and mapped to Arrow data types. + /// This is useful when we need to inspect or use the specific encoding of a field. pub fn codec(&self) -> &Codec { &self.codec } + /// Returns the nullability status of this data type + /// + /// In Avro, nullability is represented through unions with null types. + /// The returned value indicates how nulls are encoded in the Avro format: + /// - `Some(Nullability::NullFirst)` - Nulls are encoded as the first union variant + /// - `Some(Nullability::NullSecond)` - Nulls are encoded as the second union variant + /// - `None` - The type is not nullable pub fn nullability(&self) -> Option { self.nullability } @@ -78,6 +116,26 @@ impl AvroField { &self.data_type } + /// Returns a new [`AvroField`] with Utf8View support enabled + /// + /// This will convert any Utf8 codecs to Utf8View codecs. This method is used to + /// enable potential performance optimizations in string-heavy workloads by using + /// Arrow's StringViewArray data structure. + /// + /// Returns a new `AvroField` with the same structure, but with string types + /// converted to use `Utf8View` instead of `Utf8`. + pub fn with_utf8view(&self) -> Self { + let mut field = self.clone(); + if let Codec::Utf8 = field.data_type.codec { + field.data_type.codec = Codec::Utf8View; + } + field + } + + /// Returns the name of this Avro field + /// + /// This is the field name as defined in the Avro schema. + /// It's used to identify fields within a record structure. pub fn name(&self) -> &str { &self.name } @@ -90,7 +148,7 @@ impl<'a> TryFrom<&Schema<'a>> for AvroField { match schema { Schema::Complex(ComplexType::Record(r)) => { let mut resolver = Resolver::default(); - let data_type = make_data_type(schema, None, &mut resolver)?; + let data_type = make_data_type(schema, None, &mut resolver, false, false)?; Ok(AvroField { data_type, name: r.name.to_string(), @@ -103,29 +161,125 @@ impl<'a> TryFrom<&Schema<'a>> for AvroField { } } +/// Builder for an [`AvroField`] +#[derive(Debug)] +pub struct AvroFieldBuilder<'a> { + schema: &'a Schema<'a>, + use_utf8view: bool, + strict_mode: bool, +} + +impl<'a> AvroFieldBuilder<'a> { + /// Creates a new [`AvroFieldBuilder`] + pub fn new(schema: &'a Schema<'a>) -> Self { + Self { + schema, + use_utf8view: false, + strict_mode: false, + } + } + + /// Enable or disable Utf8View support + pub fn with_utf8view(mut self, use_utf8view: bool) -> Self { + self.use_utf8view = use_utf8view; + self + } + + /// Enable or disable strict mode. + pub fn with_strict_mode(mut self, strict_mode: bool) -> Self { + self.strict_mode = strict_mode; + self + } + + /// Build an [`AvroField`] from the builder + pub fn build(self) -> Result { + match self.schema { + Schema::Complex(ComplexType::Record(r)) => { + let mut resolver = Resolver::default(); + let data_type = make_data_type( + self.schema, + None, + &mut resolver, + self.use_utf8view, + self.strict_mode, + )?; + Ok(AvroField { + name: r.name.to_string(), + data_type, + }) + } + _ => Err(ArrowError::ParseError(format!( + "Expected a Record schema to build an AvroField, but got {:?}", + self.schema + ))), + } + } +} /// An Avro encoding /// /// #[derive(Debug, Clone)] pub enum Codec { + /// Represents Avro null type, maps to Arrow's Null data type Null, + /// Represents Avro boolean type, maps to Arrow's Boolean data type Boolean, + /// Represents Avro int type, maps to Arrow's Int32 data type Int32, + /// Represents Avro long type, maps to Arrow's Int64 data type Int64, + /// Represents Avro float type, maps to Arrow's Float32 data type Float32, + /// Represents Avro double type, maps to Arrow's Float64 data type Float64, + /// Represents Avro bytes type, maps to Arrow's Binary data type Binary, + /// String data represented as UTF-8 encoded bytes, corresponding to Arrow's StringArray Utf8, + /// String data represented as UTF-8 encoded bytes with an optimized view representation, + /// corresponding to Arrow's StringViewArray which provides better performance for string operations + /// + /// The Utf8View option can be enabled via `ReadOptions::use_utf8view`. + Utf8View, + /// Represents Avro date logical type, maps to Arrow's Date32 data type Date32, + /// Represents Avro time-millis logical type, maps to Arrow's Time32(TimeUnit::Millisecond) data type TimeMillis, + /// Represents Avro time-micros logical type, maps to Arrow's Time64(TimeUnit::Microsecond) data type TimeMicros, - /// TimestampMillis(is_utc) + /// Represents Avro timestamp-millis or local-timestamp-millis logical type + /// + /// Maps to Arrow's Timestamp(TimeUnit::Millisecond) data type + /// The boolean parameter indicates whether the timestamp has a UTC timezone (true) or is local time (false) TimestampMillis(bool), - /// TimestampMicros(is_utc) + /// Represents Avro timestamp-micros or local-timestamp-micros logical type + /// + /// Maps to Arrow's Timestamp(TimeUnit::Microsecond) data type + /// The boolean parameter indicates whether the timestamp has a UTC timezone (true) or is local time (false) TimestampMicros(bool), + /// Represents Avro fixed type, maps to Arrow's FixedSizeBinary data type + /// The i32 parameter indicates the fixed binary size Fixed(i32), + /// Represents Avro decimal type, maps to Arrow's Decimal128 or Decimal256 data types + /// + /// The fields are `(precision, scale, fixed_size)`. + /// - `precision` (`usize`): Total number of digits. + /// - `scale` (`Option`): Number of fractional digits. + /// - `fixed_size` (`Option`): Size in bytes if backed by a `fixed` type, otherwise `None`. + Decimal(usize, Option, Option), + /// Represents Avro Uuid type, a FixedSizeBinary with a length of 16. + Uuid, + /// Represents an Avro enum, maps to Arrow's Dictionary(Int32, Utf8) type. + /// + /// The enclosed value contains the enum's symbols. + Enum(Arc<[String]>), + /// Represents Avro array type, maps to Arrow's List data type List(Arc), + /// Represents Avro record type, maps to Arrow's Struct data type Struct(Arc<[AvroField]>), + /// Represents Avro map type, maps to Arrow's Map data type + Map(Arc), + /// Represents Avro duration logical type, maps to Arrow's Interval(IntervalUnit::MonthDayNano) data type Interval, } @@ -140,6 +294,7 @@ impl Codec { Self::Float64 => DataType::Float64, Self::Binary => DataType::Binary, Self::Utf8 => DataType::Utf8, + Self::Utf8View => DataType::Utf8View, Self::Date32 => DataType::Date32, Self::TimeMillis => DataType::Time32(TimeUnit::Millisecond), Self::TimeMicros => DataType::Time64(TimeUnit::Microsecond), @@ -151,10 +306,46 @@ impl Codec { } Self::Interval => DataType::Interval(IntervalUnit::MonthDayNano), Self::Fixed(size) => DataType::FixedSizeBinary(*size), + Self::Decimal(precision, scale, size) => { + let p = *precision as u8; + let s = scale.unwrap_or(0) as i8; + let too_large_for_128 = match *size { + Some(sz) => sz > 16, + None => { + (p as usize) > DECIMAL128_MAX_PRECISION as usize + || (s as usize) > DECIMAL128_MAX_SCALE as usize + } + }; + if too_large_for_128 { + DataType::Decimal256(p, s) + } else { + DataType::Decimal128(p, s) + } + } + Self::Uuid => DataType::FixedSizeBinary(16), + Self::Enum(_) => { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } Self::List(f) => { DataType::List(Arc::new(f.field_with_name(Field::LIST_FIELD_DEFAULT_NAME))) } Self::Struct(f) => DataType::Struct(f.iter().map(|x| x.field()).collect()), + Self::Map(value_type) => { + let val_dt = value_type.codec.data_type(); + let val_field = Field::new("value", val_dt, value_type.nullability.is_some()) + .with_metadata(value_type.metadata.clone()); + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + val_field, + ])), + false, + )), + false, + ) + } } } } @@ -174,6 +365,62 @@ impl From for Codec { } } +fn parse_decimal_attributes( + attributes: &Attributes, + fallback_size: Option, + precision_required: bool, +) -> Result<(usize, usize, Option), ArrowError> { + let precision = attributes + .additional + .get("precision") + .and_then(|v| v.as_u64()) + .or(if precision_required { None } else { Some(10) }) + .ok_or_else(|| ArrowError::ParseError("Decimal requires precision".to_string()))? + as usize; + let scale = attributes + .additional + .get("scale") + .and_then(|v| v.as_u64()) + .unwrap_or(0) as usize; + let size = attributes + .additional + .get("size") + .and_then(|v| v.as_u64()) + .map(|s| s as usize) + .or(fallback_size); + Ok((precision, scale, size)) +} + +impl Codec { + /// Converts a string codec to use Utf8View if requested + /// + /// The conversion only happens if both: + /// 1. `use_utf8view` is true + /// 2. The codec is currently `Utf8` + /// + /// # Example + /// ``` + /// # use arrow_avro::codec::Codec; + /// let utf8_codec1 = Codec::Utf8; + /// let utf8_codec2 = Codec::Utf8; + /// + /// // Convert to Utf8View + /// let view_codec = utf8_codec1.with_utf8view(true); + /// assert!(matches!(view_codec, Codec::Utf8View)); + /// + /// // Don't convert if use_utf8view is false + /// let unchanged_codec = utf8_codec2.with_utf8view(false); + /// assert!(matches!(unchanged_codec, Codec::Utf8)); + /// ``` + pub fn with_utf8view(self, use_utf8view: bool) -> Self { + if use_utf8view && matches!(self, Self::Utf8) { + Self::Utf8View + } else { + self + } + } +} + /// Resolves Avro type names to [`AvroDataType`] /// /// See @@ -203,19 +450,31 @@ impl<'a> Resolver<'a> { /// /// `name`: is name used to refer to `schema` in its parent /// `namespace`: an optional qualifier used as part of a type hierarchy +/// If the data type is a string, convert to use Utf8View if requested +/// +/// This function is used during the schema conversion process to determine whether +/// string data should be represented as StringArray (default) or StringViewArray. +/// +/// `use_utf8view`: if true, use Utf8View instead of Utf8 for string types /// /// See [`Resolver`] for more information fn make_data_type<'a>( schema: &Schema<'a>, namespace: Option<&'a str>, resolver: &mut Resolver<'a>, + use_utf8view: bool, + strict_mode: bool, ) -> Result { match schema { - Schema::TypeName(TypeName::Primitive(p)) => Ok(AvroDataType { - nullability: None, - metadata: Default::default(), - codec: (*p).into(), - }), + Schema::TypeName(TypeName::Primitive(p)) => { + let codec: Codec = (*p).into(); + let codec = codec.with_utf8view(use_utf8view); + Ok(AvroDataType { + nullability: None, + metadata: Default::default(), + codec, + }) + } Schema::TypeName(TypeName::Ref(name)) => resolver.resolve(name, namespace), Schema::Union(f) => { // Special case the common case of nullable primitives @@ -224,12 +483,20 @@ fn make_data_type<'a>( .position(|x| x == &Schema::TypeName(TypeName::Primitive(PrimitiveType::Null))); match (f.len() == 2, null) { (true, Some(0)) => { - let mut field = make_data_type(&f[1], namespace, resolver)?; + let mut field = + make_data_type(&f[1], namespace, resolver, use_utf8view, strict_mode)?; field.nullability = Some(Nullability::NullFirst); Ok(field) } (true, Some(1)) => { - let mut field = make_data_type(&f[0], namespace, resolver)?; + if strict_mode { + return Err(ArrowError::SchemaError( + "Found Avro union of the form ['T','null'], which is disallowed in strict_mode" + .to_string(), + )); + } + let mut field = + make_data_type(&f[0], namespace, resolver, use_utf8view, strict_mode)?; field.nullability = Some(Nullability::NullSecond); Ok(field) } @@ -247,11 +514,16 @@ fn make_data_type<'a>( .map(|field| { Ok(AvroField { name: field.name.to_string(), - data_type: make_data_type(&field.r#type, namespace, resolver)?, + data_type: make_data_type( + &field.r#type, + namespace, + resolver, + use_utf8view, + strict_mode, + )?, }) }) .collect::>()?; - let field = AvroDataType { nullability: None, codec: Codec::Struct(fields), @@ -261,7 +533,13 @@ fn make_data_type<'a>( Ok(field) } ComplexType::Array(a) => { - let mut field = make_data_type(a.items.as_ref(), namespace, resolver)?; + let mut field = make_data_type( + a.items.as_ref(), + namespace, + resolver, + use_utf8view, + strict_mode, + )?; Ok(AvroDataType { nullability: None, metadata: a.attributes.field_metadata(), @@ -272,32 +550,83 @@ fn make_data_type<'a>( let size = f.size.try_into().map_err(|e| { ArrowError::ParseError(format!("Overflow converting size to i32: {e}")) })?; + let md = f.attributes.field_metadata(); + let field = match f.attributes.logical_type { + Some("decimal") => { + let (precision, scale, _) = + parse_decimal_attributes(&f.attributes, Some(size as usize), true)?; + AvroDataType { + nullability: None, + metadata: md, + codec: Codec::Decimal(precision, Some(scale), Some(size as usize)), + } + } + Some("duration") => { + if size != 12 { + return Err(ArrowError::ParseError(format!( + "Invalid fixed size for Duration: {size}, must be 12" + ))); + }; + AvroDataType { + nullability: None, + metadata: md, + codec: Codec::Interval, + } + } + _ => AvroDataType { + nullability: None, + metadata: md, + codec: Codec::Fixed(size), + }, + }; + resolver.register(f.name, namespace, field.clone()); + Ok(field) + } + ComplexType::Enum(e) => { + let namespace = e.namespace.or(namespace); + let symbols = e + .symbols + .iter() + .map(|s| s.to_string()) + .collect::>(); + let mut metadata = e.attributes.field_metadata(); + let symbols_json = serde_json::to_string(&e.symbols).map_err(|e| { + ArrowError::ParseError(format!("Failed to serialize enum symbols: {e}")) + })?; + metadata.insert("avro.enum.symbols".to_string(), symbols_json); let field = AvroDataType { nullability: None, - metadata: f.attributes.field_metadata(), - codec: Codec::Fixed(size), + metadata, + codec: Codec::Enum(symbols), }; - resolver.register(f.name, namespace, field.clone()); + resolver.register(e.name, namespace, field.clone()); Ok(field) } - ComplexType::Enum(e) => Err(ArrowError::NotYetImplemented(format!( - "Enum of {e:?} not currently supported" - ))), - ComplexType::Map(m) => Err(ArrowError::NotYetImplemented(format!( - "Map of {m:?} not currently supported" - ))), + ComplexType::Map(m) => { + let val = + make_data_type(&m.values, namespace, resolver, use_utf8view, strict_mode)?; + Ok(AvroDataType { + nullability: None, + metadata: m.attributes.field_metadata(), + codec: Codec::Map(Arc::new(val)), + }) + } }, Schema::Type(t) => { - let mut field = - make_data_type(&Schema::TypeName(t.r#type.clone()), namespace, resolver)?; + let mut field = make_data_type( + &Schema::TypeName(t.r#type.clone()), + namespace, + resolver, + use_utf8view, + strict_mode, + )?; // https://avro.apache.org/docs/1.11.1/specification/#logical-types match (t.attributes.logical_type, &mut field.codec) { - (Some("decimal"), c @ Codec::Fixed(_)) => { - return Err(ArrowError::NotYetImplemented( - "Decimals are not currently supported".to_string(), - )) + (Some("decimal"), c @ Codec::Binary) => { + let (prec, sc, _) = parse_decimal_attributes(&t.attributes, None, false)?; + *c = Codec::Decimal(prec, Some(sc), None); } (Some("date"), c @ Codec::Int32) => *c = Codec::Date32, (Some("time-millis"), c @ Codec::Int32) => *c = Codec::TimeMillis, @@ -310,7 +639,7 @@ fn make_data_type<'a>( (Some("local-timestamp-micros"), c @ Codec::Int64) => { *c = Codec::TimestampMicros(false) } - (Some("duration"), c @ Codec::Fixed(12)) => *c = Codec::Interval, + (Some("uuid"), c @ Codec::Utf8) => *c = Codec::Uuid, (Some(logical), _) => { // Insert unrecognized logical type into metadata map field.metadata.insert("logicalType".into(), logical.into()); @@ -327,3 +656,245 @@ fn make_data_type<'a>( } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema::{ + Attributes, ComplexType, Fixed, PrimitiveType, Record, Schema, Type, TypeName, + }; + use serde_json; + use std::collections::HashMap; + + fn create_schema_with_logical_type( + primitive_type: PrimitiveType, + logical_type: &'static str, + ) -> Schema<'static> { + let attributes = Attributes { + logical_type: Some(logical_type), + additional: Default::default(), + }; + + Schema::Type(Type { + r#type: TypeName::Primitive(primitive_type), + attributes, + }) + } + + fn create_fixed_schema(size: usize, logical_type: &'static str) -> Schema<'static> { + let attributes = Attributes { + logical_type: Some(logical_type), + additional: Default::default(), + }; + + Schema::Complex(ComplexType::Fixed(Fixed { + name: "fixed_type", + namespace: None, + aliases: Vec::new(), + size, + attributes, + })) + } + + #[test] + fn test_date_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Int, "date"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::Date32)); + } + + #[test] + fn test_time_millis_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Int, "time-millis"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimeMillis)); + } + + #[test] + fn test_time_micros_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Long, "time-micros"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimeMicros)); + } + + #[test] + fn test_timestamp_millis_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Long, "timestamp-millis"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimestampMillis(true))); + } + + #[test] + fn test_timestamp_micros_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Long, "timestamp-micros"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimestampMicros(true))); + } + + #[test] + fn test_local_timestamp_millis_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Long, "local-timestamp-millis"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimestampMillis(false))); + } + + #[test] + fn test_local_timestamp_micros_logical_type() { + let schema = create_schema_with_logical_type(PrimitiveType::Long, "local-timestamp-micros"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::TimestampMicros(false))); + } + + #[test] + fn test_uuid_type() { + let mut codec = Codec::Fixed(16); + + if let c @ Codec::Fixed(16) = &mut codec { + *c = Codec::Uuid; + } + + assert!(matches!(codec, Codec::Uuid)); + } + + #[test] + fn test_duration_logical_type() { + let mut codec = Codec::Fixed(12); + + if let c @ Codec::Fixed(12) = &mut codec { + *c = Codec::Interval; + } + + assert!(matches!(codec, Codec::Interval)); + } + + #[test] + fn test_decimal_logical_type_not_implemented() { + let mut codec = Codec::Fixed(16); + + let process_decimal = || -> Result<(), ArrowError> { + if let Codec::Fixed(_) = codec { + return Err(ArrowError::NotYetImplemented( + "Decimals are not currently supported".to_string(), + )); + } + Ok(()) + }; + + let result = process_decimal(); + + assert!(result.is_err()); + if let Err(ArrowError::NotYetImplemented(msg)) = result { + assert!(msg.contains("Decimals are not currently supported")); + } else { + panic!("Expected NotYetImplemented error"); + } + } + + #[test] + fn test_unknown_logical_type_added_to_metadata() { + let schema = create_schema_with_logical_type(PrimitiveType::Int, "custom-type"); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert_eq!( + result.metadata.get("logicalType"), + Some(&"custom-type".to_string()) + ); + } + + #[test] + fn test_string_with_utf8view_enabled() { + let schema = Schema::TypeName(TypeName::Primitive(PrimitiveType::String)); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, true, false).unwrap(); + + assert!(matches!(result.codec, Codec::Utf8View)); + } + + #[test] + fn test_string_without_utf8view_enabled() { + let schema = Schema::TypeName(TypeName::Primitive(PrimitiveType::String)); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, false).unwrap(); + + assert!(matches!(result.codec, Codec::Utf8)); + } + + #[test] + fn test_record_with_string_and_utf8view_enabled() { + let field_schema = Schema::TypeName(TypeName::Primitive(PrimitiveType::String)); + + let avro_field = crate::schema::Field { + name: "string_field", + r#type: field_schema, + default: None, + doc: None, + }; + + let record = Record { + name: "test_record", + namespace: None, + aliases: vec![], + doc: None, + fields: vec![avro_field], + attributes: Attributes::default(), + }; + + let schema = Schema::Complex(ComplexType::Record(record)); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, true, false).unwrap(); + + if let Codec::Struct(fields) = &result.codec { + let first_field_codec = &fields[0].data_type().codec; + assert!(matches!(first_field_codec, Codec::Utf8View)); + } else { + panic!("Expected Struct codec"); + } + } + + #[test] + fn test_union_with_strict_mode() { + let schema = Schema::Union(vec![ + Schema::TypeName(TypeName::Primitive(PrimitiveType::String)), + Schema::TypeName(TypeName::Primitive(PrimitiveType::Null)), + ]); + + let mut resolver = Resolver::default(); + let result = make_data_type(&schema, None, &mut resolver, false, true); + + assert!(result.is_err()); + match result { + Err(ArrowError::SchemaError(msg)) => { + assert!(msg.contains( + "Found Avro union of the form ['T','null'], which is disallowed in strict_mode" + )); + } + _ => panic!("Expected SchemaError"), + } + } +} diff --git a/arrow-avro/src/compression.rs b/arrow-avro/src/compression.rs index f29b8dd07606..1e1960dc841f 100644 --- a/arrow-avro/src/compression.rs +++ b/arrow-avro/src/compression.rs @@ -23,10 +23,21 @@ use std::io::Read; pub const CODEC_METADATA_KEY: &str = "avro.codec"; #[derive(Debug, Copy, Clone, Eq, PartialEq)] +/// Supported compression codecs for Avro data +/// +/// Avro supports multiple compression formats for data blocks. +/// This enum represents the compression codecs available in this implementation. pub enum CompressionCodec { + /// Deflate compression (RFC 1951) Deflate, + /// Snappy compression Snappy, + /// ZStandard compression ZStandard, + /// Bzip2 compression + Bzip2, + /// Xz compression + Xz, } impl CompressionCodec { @@ -77,6 +88,28 @@ impl CompressionCodec { CompressionCodec::ZStandard => Err(ArrowError::ParseError( "ZStandard codec requires zstd feature".to_string(), )), + #[cfg(feature = "bzip2")] + CompressionCodec::Bzip2 => { + let mut decoder = bzip2::read::BzDecoder::new(block); + let mut out = Vec::new(); + decoder.read_to_end(&mut out)?; + Ok(out) + } + #[cfg(not(feature = "bzip2"))] + CompressionCodec::Bzip2 => Err(ArrowError::ParseError( + "Bzip2 codec requires bzip2 feature".to_string(), + )), + #[cfg(feature = "xz")] + CompressionCodec::Xz => { + let mut decoder = xz::read::XzDecoder::new(block); + let mut out = Vec::new(); + decoder.read_to_end(&mut out)?; + Ok(out) + } + #[cfg(not(feature = "xz"))] + CompressionCodec::Xz => Err(ArrowError::ParseError( + "XZ codec requires xz feature".to_string(), + )), } } } diff --git a/arrow-avro/src/lib.rs b/arrow-avro/src/lib.rs index d01d681b7af0..ae13c3861842 100644 --- a/arrow-avro/src/lib.rs +++ b/arrow-avro/src/lib.rs @@ -20,15 +20,51 @@ //! [Apache Arrow]: https://arrow.apache.org //! [Apache Avro]: https://avro.apache.org/ +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] #![allow(unused)] // Temporary +/// Core functionality for reading Avro data into Arrow arrays +/// +/// Implements the primary reader interface and record decoding logic. pub mod reader; + +// Avro schema parsing and representation +// +// Provides types for parsing and representing Avro schema definitions. mod schema; -mod compression; +/// Compression codec implementations for Avro +/// +/// Provides support for various compression algorithms used in Avro files, +/// including Deflate, Snappy, and ZStandard. +pub mod compression; -mod codec; +/// Data type conversions between Avro and Arrow types +/// +/// This module contains the necessary types and functions to convert between +/// Avro data types and Arrow data types. +pub mod codec; + +/// Extension trait for AvroField to add Utf8View support +/// +/// This trait adds methods for working with Utf8View support to the AvroField struct. +pub trait AvroFieldExt { + /// Returns a new field with Utf8View support enabled for string data + /// + /// This will convert any string data to use StringViewArray instead of StringArray. + fn with_utf8view(&self) -> Self; +} + +impl AvroFieldExt for codec::AvroField { + fn with_utf8view(&self) -> Self { + codec::AvroField::with_utf8view(self) + } +} #[cfg(test)] mod test_util { diff --git a/arrow-avro/src/reader/cursor.rs b/arrow-avro/src/reader/cursor.rs index 4b6a5a4d65db..1b89ff86c38c 100644 --- a/arrow-avro/src/reader/cursor.rs +++ b/arrow-avro/src/reader/cursor.rs @@ -118,4 +118,16 @@ impl<'a> AvroCursor<'a> { self.buf = &self.buf[8..]; Ok(ret) } + + /// Read exactly `n` bytes from the buffer (e.g. for Avro `fixed`). + pub(crate) fn get_fixed(&mut self, n: usize) -> Result<&'a [u8], ArrowError> { + if self.buf.len() < n { + return Err(ArrowError::ParseError( + "Unexpected EOF reading fixed".to_string(), + )); + } + let ret = &self.buf[..n]; + self.buf = &self.buf[n..]; + Ok(ret) + } } diff --git a/arrow-avro/src/reader/header.rs b/arrow-avro/src/reader/header.rs index 98c285171bf3..0f7ffd3f8d6e 100644 --- a/arrow-avro/src/reader/header.rs +++ b/arrow-avro/src/reader/header.rs @@ -77,12 +77,13 @@ impl Header { /// Returns the [`CompressionCodec`] if any pub fn compression(&self) -> Result, ArrowError> { let v = self.get(CODEC_METADATA_KEY); - match v { None | Some(b"null") => Ok(None), Some(b"deflate") => Ok(Some(CompressionCodec::Deflate)), Some(b"snappy") => Ok(Some(CompressionCodec::Snappy)), Some(b"zstandard") => Ok(Some(CompressionCodec::ZStandard)), + Some(b"bzip2") => Ok(Some(CompressionCodec::Bzip2)), + Some(b"xz") => Ok(Some(CompressionCodec::Xz)), Some(v) => Err(ArrowError::ParseError(format!( "Unrecognized compression codec \'{}\'", String::from_utf8_lossy(v) diff --git a/arrow-avro/src/reader/mod.rs b/arrow-avro/src/reader/mod.rs index 12fa67d9c8e3..02d3f49aa10c 100644 --- a/arrow-avro/src/reader/mod.rs +++ b/arrow-avro/src/reader/mod.rs @@ -15,22 +15,93 @@ // specific language governing permissions and limitations // under the License. -//! Read Avro data to Arrow +//! Avro reader +//! +//! This module provides facilities to read Apache Avro-encoded files or streams +//! into Arrow's `RecordBatch` format. In particular, it introduces: +//! +//! * `ReaderBuilder`: Configures Avro reading, e.g., batch size +//! * `Reader`: Yields `RecordBatch` values, implementing `Iterator` +//! * `Decoder`: A low-level push-based decoder for Avro records +//! +//! # Basic Usage +//! +//! `Reader` can be used directly with synchronous data sources, such as [`std::fs::File`]. +//! +//! ## Reading a Single Batch +//! +//! ``` +//! # use std::fs::File; +//! # use std::io::BufReader; +//! # use arrow_avro::reader::ReaderBuilder; +//! +//! let file = File::open("../testing/data/avro/alltypes_plain.avro").unwrap(); +//! let mut avro = ReaderBuilder::new().build(BufReader::new(file)).unwrap(); +//! let batch = avro.next().unwrap(); +//! ``` +//! +//! # Async Usage +//! +//! The lower-level `Decoder` can be integrated with various forms of async data streams, +//! and is designed to be agnostic to different async IO primitives within +//! the Rust ecosystem. It works by incrementally decoding Avro data from byte slices. +//! +//! For example, see below for how it could be used with an arbitrary `Stream` of `Bytes`: +//! +//! ``` +//! # use std::task::{Poll, ready}; +//! # use bytes::{Buf, Bytes}; +//! # use arrow_schema::ArrowError; +//! # use futures::stream::{Stream, StreamExt}; +//! # use arrow_array::RecordBatch; +//! # use arrow_avro::reader::Decoder; +//! +//! fn decode_stream + Unpin>( +//! mut decoder: Decoder, +//! mut input: S, +//! ) -> impl Stream> { +//! let mut buffered = Bytes::new(); +//! futures::stream::poll_fn(move |cx| { +//! loop { +//! if buffered.is_empty() { +//! buffered = match ready!(input.poll_next_unpin(cx)) { +//! Some(b) => b, +//! None => break, +//! }; +//! } +//! let decoded = match decoder.decode(buffered.as_ref()) { +//! Ok(decoded) => decoded, +//! Err(e) => return Poll::Ready(Some(Err(e))), +//! }; +//! let read = buffered.len(); +//! buffered.advance(decoded); +//! if decoded != read { +//! break +//! } +//! } +//! // Convert any fully-decoded rows to a RecordBatch, if available +//! Poll::Ready(decoder.flush().transpose()) +//! }) +//! } +//! ``` +//! -use crate::reader::block::{Block, BlockDecoder}; -use crate::reader::header::{Header, HeaderDecoder}; -use arrow_schema::ArrowError; +use crate::codec::AvroFieldBuilder; +use crate::schema::Schema as AvroSchema; +use arrow_array::{RecordBatch, RecordBatchReader}; +use arrow_schema::{ArrowError, SchemaRef}; +use block::BlockDecoder; +use header::{Header, HeaderDecoder}; +use record::RecordDecoder; use std::io::BufRead; -mod header; - mod block; - mod cursor; +mod header; mod record; mod vlq; -/// Read a [`Header`] from the provided [`BufRead`] +/// Read the Avro file header (magic, metadata, sync marker) from `reader`. fn read_header(mut reader: R) -> Result { let mut decoder = HeaderDecoder::default(); loop { @@ -45,75 +116,407 @@ fn read_header(mut reader: R) -> Result { break; } } + decoder.flush().ok_or_else(|| { + ArrowError::ParseError("Unexpected EOF while reading Avro header".to_string()) + }) +} + +/// A low-level interface for decoding Avro-encoded bytes into Arrow `RecordBatch`. +#[derive(Debug)] +pub struct Decoder { + record_decoder: RecordDecoder, + batch_size: usize, + decoded_rows: usize, +} + +impl Decoder { + fn new(record_decoder: RecordDecoder, batch_size: usize) -> Self { + Self { + record_decoder, + batch_size, + decoded_rows: 0, + } + } + + /// Return the Arrow schema for the rows decoded by this decoder + pub fn schema(&self) -> SchemaRef { + self.record_decoder.schema().clone() + } + + /// Return the configured maximum number of rows per batch + pub fn batch_size(&self) -> usize { + self.batch_size + } + + /// Feed `data` into the decoder row by row until we either: + /// - consume all bytes in `data`, or + /// - reach `batch_size` decoded rows. + /// + /// Returns the number of bytes consumed. + pub fn decode(&mut self, data: &[u8]) -> Result { + let mut total_consumed = 0usize; + while total_consumed < data.len() && self.decoded_rows < self.batch_size { + let consumed = self.record_decoder.decode(&data[total_consumed..], 1)?; + // A successful call to record_decoder.decode means one row was decoded. + // If `consumed` is 0 on a non-empty buffer, it implies a valid zero-byte record. + // We increment `decoded_rows` to mark progress and avoid an infinite loop. + // We add `consumed` (which can be 0) to `total_consumed`. + total_consumed += consumed; + self.decoded_rows += 1; + } + Ok(total_consumed) + } + + /// Produce a `RecordBatch` if at least one row is fully decoded, returning + /// `Ok(None)` if no new rows are available. + pub fn flush(&mut self) -> Result, ArrowError> { + if self.decoded_rows == 0 { + Ok(None) + } else { + let batch = self.record_decoder.flush()?; + self.decoded_rows = 0; + Ok(Some(batch)) + } + } - decoder - .flush() - .ok_or_else(|| ArrowError::ParseError("Unexpected EOF".to_string())) + /// Returns the number of rows that can be added to this decoder before it is full. + pub fn capacity(&self) -> usize { + self.batch_size.saturating_sub(self.decoded_rows) + } + + /// Returns true if the decoder has reached its capacity for the current batch. + pub fn batch_is_full(&self) -> bool { + self.capacity() == 0 + } +} + +/// A builder to create an [`Avro Reader`](Reader) that reads Avro data +/// into Arrow `RecordBatch`. +#[derive(Debug)] +pub struct ReaderBuilder { + batch_size: usize, + strict_mode: bool, + utf8_view: bool, + schema: Option>, } -/// Return an iterator of [`Block`] from the provided [`BufRead`] -fn read_blocks(mut reader: R) -> impl Iterator> { - let mut decoder = BlockDecoder::default(); +impl Default for ReaderBuilder { + fn default() -> Self { + Self { + batch_size: 1024, + strict_mode: false, + utf8_view: false, + schema: None, + } + } +} + +impl ReaderBuilder { + /// Creates a new [`ReaderBuilder`] with default settings: + /// - `batch_size` = 1024 + /// - `strict_mode` = false + /// - `utf8_view` = false + /// - `schema` = None + pub fn new() -> Self { + Self::default() + } + + fn make_record_decoder(&self, schema: &AvroSchema<'_>) -> Result { + let root_field = AvroFieldBuilder::new(schema) + .with_utf8view(self.utf8_view) + .with_strict_mode(self.strict_mode) + .build()?; + RecordDecoder::try_new_with_options(root_field.data_type(), self.utf8_view) + } + + fn build_impl(self, reader: &mut R) -> Result<(Header, Decoder), ArrowError> { + let header = read_header(reader)?; + let record_decoder = if let Some(schema) = &self.schema { + self.make_record_decoder(schema)? + } else { + let avro_schema: Option> = header + .schema() + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + let avro_schema = avro_schema.ok_or_else(|| { + ArrowError::ParseError("No Avro schema present in file header".to_string()) + })?; + self.make_record_decoder(&avro_schema)? + }; + let decoder = Decoder::new(record_decoder, self.batch_size); + Ok((header, decoder)) + } + + /// Sets the row-based batch size + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + /// Set whether to use StringViewArray for string data + /// + /// When enabled, string data from Avro files will be loaded into + /// Arrow's StringViewArray instead of the standard StringArray. + pub fn with_utf8_view(mut self, utf8_view: bool) -> Self { + self.utf8_view = utf8_view; + self + } + + /// Get whether StringViewArray is enabled for string data + pub fn use_utf8view(&self) -> bool { + self.utf8_view + } - let mut try_next = move || { - loop { - let buf = reader.fill_buf()?; - if buf.is_empty() { - break; + /// Controls whether certain Avro unions of the form `[T, "null"]` should produce an error. + pub fn with_strict_mode(mut self, strict_mode: bool) -> Self { + self.strict_mode = strict_mode; + self + } + + /// Sets the Avro schema. + /// + /// If a schema is not provided, the schema will be read from the Avro file header. + pub fn with_schema(mut self, schema: AvroSchema<'static>) -> Self { + self.schema = Some(schema); + self + } + + /// Create a [`Reader`] from this builder and a `BufRead` + pub fn build(self, mut reader: R) -> Result, ArrowError> { + let (header, decoder) = self.build_impl(&mut reader)?; + Ok(Reader { + reader, + header, + decoder, + block_decoder: BlockDecoder::default(), + block_data: Vec::new(), + block_cursor: 0, + finished: false, + }) + } + + /// Create a [`Decoder`] from this builder and a `BufRead` by + /// reading and parsing the Avro file's header. This will + /// not create a full [`Reader`]. + pub fn build_decoder(self, mut reader: R) -> Result { + match self.schema { + Some(ref schema) => { + let record_decoder = self.make_record_decoder(schema)?; + Ok(Decoder::new(record_decoder, self.batch_size)) } - let read = buf.len(); - let decoded = decoder.decode(buf)?; - reader.consume(decoded); - if decoded != read { - break; + None => { + let (_, decoder) = self.build_impl(&mut reader)?; + Ok(decoder) } } - Ok(decoder.flush()) - }; - std::iter::from_fn(move || try_next().transpose()) + } +} + +/// A high-level Avro `Reader` that reads container-file blocks +/// and feeds them into a row-level [`Decoder`]. +#[derive(Debug)] +pub struct Reader { + reader: R, + header: Header, + decoder: Decoder, + block_decoder: BlockDecoder, + block_data: Vec, + block_cursor: usize, + finished: bool, +} + +impl Reader { + /// Return the Arrow schema discovered from the Avro file header + pub fn schema(&self) -> SchemaRef { + self.decoder.schema() + } + + /// Return the Avro container-file header + pub fn avro_header(&self) -> &Header { + &self.header + } + + /// Reads the next [`RecordBatch`] from the Avro file or `Ok(None)` on EOF + fn read(&mut self) -> Result, ArrowError> { + 'outer: while !self.finished && !self.decoder.batch_is_full() { + while self.block_cursor == self.block_data.len() { + let buf = self.reader.fill_buf()?; + if buf.is_empty() { + self.finished = true; + break 'outer; + } + // Try to decode another block from the buffered reader. + let consumed = self.block_decoder.decode(buf)?; + self.reader.consume(consumed); + if let Some(block) = self.block_decoder.flush() { + // Successfully decoded a block. + let block_data = if let Some(ref codec) = self.header.compression()? { + codec.decompress(&block.data)? + } else { + block.data + }; + self.block_data = block_data; + self.block_cursor = 0; + } else if consumed == 0 { + // The block decoder made no progress on a non-empty buffer. + return Err(ArrowError::ParseError( + "Could not decode next Avro block from partial data".to_string(), + )); + } + } + // Try to decode more rows from the current block. + let consumed = self.decoder.decode(&self.block_data[self.block_cursor..])?; + self.block_cursor += consumed; + } + self.decoder.flush() + } +} + +impl Iterator for Reader { + type Item = Result; + + fn next(&mut self) -> Option { + self.read().transpose() + } +} + +impl RecordBatchReader for Reader { + fn schema(&self) -> SchemaRef { + self.schema() + } } #[cfg(test)] mod test { - use crate::codec::AvroField; + use crate::codec::{AvroDataType, AvroField, Codec}; use crate::compression::CompressionCodec; use crate::reader::record::RecordDecoder; - use crate::reader::{read_blocks, read_header}; + use crate::reader::vlq::VLQDecoder; + use crate::reader::{read_header, Decoder, Reader, ReaderBuilder}; use crate::test_util::arrow_test_data; + use arrow::array::ArrayDataBuilder; + use arrow_array::builder::{ + ArrayBuilder, BooleanBuilder, Float32Builder, Float64Builder, Int32Builder, Int64Builder, + ListBuilder, MapBuilder, StringBuilder, StructBuilder, + }; + use arrow_array::types::{Int32Type, IntervalMonthDayNanoType}; use arrow_array::*; + use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema}; + use bytes::{Buf, BufMut, Bytes}; + use futures::executor::block_on; + use futures::{stream, Stream, StreamExt, TryStreamExt}; + use std::collections::HashMap; + use std::fs; use std::fs::File; - use std::io::BufReader; + use std::io::{BufReader, Cursor, Read}; use std::sync::Arc; + use std::task::{ready, Poll}; - fn read_file(file: &str, batch_size: usize) -> RecordBatch { - let file = File::open(file).unwrap(); - let mut reader = BufReader::new(file); - let header = read_header(&mut reader).unwrap(); - let compression = header.compression().unwrap(); - let schema = header.schema().unwrap().unwrap(); - let root = AvroField::try_from(&schema).unwrap(); - let mut decoder = RecordDecoder::try_new(root.data_type()).unwrap(); - - for result in read_blocks(reader) { - let block = result.unwrap(); - assert_eq!(block.sync, header.sync()); - if let Some(c) = compression { - let decompressed = c.decompress(&block.data).unwrap(); - - let mut offset = 0; - let mut remaining = block.count; - while remaining > 0 { - let to_read = remaining.max(batch_size); - offset += decoder - .decode(&decompressed[offset..], block.count) - .unwrap(); - - remaining -= to_read; + fn read_file(path: &str, batch_size: usize, utf8_view: bool) -> RecordBatch { + let file = File::open(path).unwrap(); + let reader = ReaderBuilder::new() + .with_batch_size(batch_size) + .with_utf8_view(utf8_view) + .build(BufReader::new(file)) + .unwrap(); + let schema = reader.schema(); + let batches = reader.collect::, _>>().unwrap(); + arrow::compute::concat_batches(&schema, &batches).unwrap() + } + + fn read_file_strict( + path: &str, + batch_size: usize, + utf8_view: bool, + ) -> Result>, ArrowError> { + let file = File::open(path).unwrap(); + ReaderBuilder::new() + .with_batch_size(batch_size) + .with_utf8_view(utf8_view) + .with_strict_mode(true) + .build(BufReader::new(file)) + } + + fn decode_stream + Unpin>( + mut decoder: Decoder, + mut input: S, + ) -> impl Stream> { + async_stream::try_stream! { + if let Some(data) = input.next().await { + let consumed = decoder.decode(&data)?; + if consumed < data.len() { + Err(ArrowError::ParseError( + "did not consume all bytes".to_string(), + ))?; } - assert_eq!(offset, decompressed.len()); + } + if let Some(batch) = decoder.flush()? { + yield batch } } - decoder.flush().unwrap() + } + + #[test] + fn test_utf8view_support() { + let schema_json = r#"{ + "type": "record", + "name": "test", + "fields": [{ + "name": "str_field", + "type": "string" + }] + }"#; + + let schema: crate::schema::Schema = serde_json::from_str(schema_json).unwrap(); + let avro_field = AvroField::try_from(&schema).unwrap(); + + let data_type = avro_field.data_type(); + + struct TestHelper; + impl TestHelper { + fn with_utf8view(field: &Field) -> Field { + match field.data_type() { + DataType::Utf8 => { + Field::new(field.name(), DataType::Utf8View, field.is_nullable()) + .with_metadata(field.metadata().clone()) + } + _ => field.clone(), + } + } + } + + let field = TestHelper::with_utf8view(&Field::new("str_field", DataType::Utf8, false)); + + assert_eq!(field.data_type(), &DataType::Utf8View); + + let array = StringViewArray::from(vec!["test1", "test2"]); + let batch = + RecordBatch::try_from_iter(vec![("str_field", Arc::new(array) as ArrayRef)]).unwrap(); + + assert!(batch.column(0).as_any().is::()); + } + + #[test] + fn test_read_zero_byte_avro_file() { + let batch = read_file("test/data/zero_byte.avro", 3, false); + let schema = batch.schema(); + assert_eq!(schema.fields().len(), 1); + let field = schema.field(0); + assert_eq!(field.name(), "data"); + assert_eq!(field.data_type(), &DataType::Binary); + assert!(field.is_nullable()); + assert_eq!(batch.num_rows(), 3); + assert_eq!(batch.num_columns(), 1); + let binary_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(binary_array.is_null(0)); + assert!(binary_array.is_valid(1)); + assert_eq!(binary_array.value(1), b""); + assert!(binary_array.is_valid(2)); + assert_eq!(binary_array.value(2), b"some bytes"); } #[test] @@ -122,6 +525,8 @@ mod test { "avro/alltypes_plain.avro", "avro/alltypes_plain.snappy.avro", "avro/alltypes_plain.zstandard.avro", + "avro/alltypes_plain.bzip2.avro", + "avro/alltypes_plain.xz.avro", ]; let expected = RecordBatch::try_from_iter_with_nullable([ @@ -211,8 +616,1237 @@ mod test { for file in files { let file = arrow_test_data(file); - assert_eq!(read_file(&file, 8), expected); - assert_eq!(read_file(&file, 3), expected); + assert_eq!(read_file(&file, 8, false), expected); + assert_eq!(read_file(&file, 3, false), expected); + } + } + + #[test] + fn test_alltypes_dictionary() { + let file = "avro/alltypes_dictionary.avro"; + let expected = RecordBatch::try_from_iter_with_nullable([ + ("id", Arc::new(Int32Array::from(vec![0, 1])) as _, true), + ( + "bool_col", + Arc::new(BooleanArray::from(vec![Some(true), Some(false)])) as _, + true, + ), + ( + "tinyint_col", + Arc::new(Int32Array::from(vec![0, 1])) as _, + true, + ), + ( + "smallint_col", + Arc::new(Int32Array::from(vec![0, 1])) as _, + true, + ), + ("int_col", Arc::new(Int32Array::from(vec![0, 1])) as _, true), + ( + "bigint_col", + Arc::new(Int64Array::from(vec![0, 10])) as _, + true, + ), + ( + "float_col", + Arc::new(Float32Array::from(vec![0.0, 1.1])) as _, + true, + ), + ( + "double_col", + Arc::new(Float64Array::from(vec![0.0, 10.1])) as _, + true, + ), + ( + "date_string_col", + Arc::new(BinaryArray::from_iter_values([b"01/01/09", b"01/01/09"])) as _, + true, + ), + ( + "string_col", + Arc::new(BinaryArray::from_iter_values([b"0", b"1"])) as _, + true, + ), + ( + "timestamp_col", + Arc::new( + TimestampMicrosecondArray::from_iter_values([ + 1230768000000000, // 2009-01-01T00:00:00.000 + 1230768060000000, // 2009-01-01T00:01:00.000 + ]) + .with_timezone("+00:00"), + ) as _, + true, + ), + ]) + .unwrap(); + let file_path = arrow_test_data(file); + let batch_large = read_file(&file_path, 8, false); + assert_eq!( + batch_large, expected, + "Decoded RecordBatch does not match for file {file}" + ); + let batch_small = read_file(&file_path, 3, false); + assert_eq!( + batch_small, expected, + "Decoded RecordBatch (batch size 3) does not match for file {file}" + ); + } + + #[test] + fn test_alltypes_nulls_plain() { + let file = "avro/alltypes_nulls_plain.avro"; + let expected = RecordBatch::try_from_iter_with_nullable([ + ( + "string_col", + Arc::new(StringArray::from(vec![None::<&str>])) as _, + true, + ), + ("int_col", Arc::new(Int32Array::from(vec![None])) as _, true), + ( + "bool_col", + Arc::new(BooleanArray::from(vec![None])) as _, + true, + ), + ( + "bigint_col", + Arc::new(Int64Array::from(vec![None])) as _, + true, + ), + ( + "float_col", + Arc::new(Float32Array::from(vec![None])) as _, + true, + ), + ( + "double_col", + Arc::new(Float64Array::from(vec![None])) as _, + true, + ), + ( + "bytes_col", + Arc::new(BinaryArray::from(vec![None::<&[u8]>])) as _, + true, + ), + ]) + .unwrap(); + let file_path = arrow_test_data(file); + let batch_large = read_file(&file_path, 8, false); + assert_eq!( + batch_large, expected, + "Decoded RecordBatch does not match for file {file}" + ); + let batch_small = read_file(&file_path, 3, false); + assert_eq!( + batch_small, expected, + "Decoded RecordBatch (batch size 3) does not match for file {file}" + ); + } + + #[test] + fn test_binary() { + let file = arrow_test_data("avro/binary.avro"); + let batch = read_file(&file, 8, false); + let expected = RecordBatch::try_from_iter_with_nullable([( + "foo", + Arc::new(BinaryArray::from_iter_values(vec![ + b"\x00".as_ref(), + b"\x01".as_ref(), + b"\x02".as_ref(), + b"\x03".as_ref(), + b"\x04".as_ref(), + b"\x05".as_ref(), + b"\x06".as_ref(), + b"\x07".as_ref(), + b"\x08".as_ref(), + b"\t".as_ref(), + b"\n".as_ref(), + b"\x0b".as_ref(), + ])) as Arc, + true, + )]) + .unwrap(); + assert_eq!(batch, expected); + } + + #[test] + fn test_decode_stream_with_schema() { + struct TestCase<'a> { + name: &'a str, + schema: &'a str, + expected_error: Option<&'a str>, + } + let tests = vec![ + TestCase { + name: "success", + schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"string"}]}"#, + expected_error: None, + }, + TestCase { + name: "valid schema invalid data", + schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"long"}]}"#, + expected_error: Some("did not consume all bytes"), + }, + ]; + for test in tests { + let schema_s2: crate::schema::Schema = serde_json::from_str(test.schema).unwrap(); + let record_val = "some_string"; + let mut body = vec![]; + body.push((record_val.len() as u8) << 1); + body.extend_from_slice(record_val.as_bytes()); + let mut reader_placeholder = Cursor::new(&[] as &[u8]); + let builder = ReaderBuilder::new() + .with_batch_size(1) + .with_schema(schema_s2); + let decoder_result = builder.build_decoder(&mut reader_placeholder); + let decoder = match decoder_result { + Ok(decoder) => decoder, + Err(e) => { + if let Some(expected) = test.expected_error { + assert!( + e.to_string().contains(expected), + "Test '{}' failed: unexpected error message at build.\nExpected to contain: '{expected}'\nActual: '{e}'", + test.name, + ); + continue; + } else { + panic!("Test '{}' failed at decoder build: {e}", test.name); + } + } + }; + let stream = Box::pin(stream::once(async { Bytes::from(body) })); + let decoded_stream = decode_stream(decoder, stream); + let batches_result: Result, ArrowError> = + block_on(decoded_stream.try_collect()); + match (batches_result, test.expected_error) { + (Ok(batches), None) => { + let batch = + arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap(); + let expected_field = Field::new("f2", DataType::Utf8, false); + let expected_schema = Arc::new(Schema::new(vec![expected_field])); + let expected_array = Arc::new(StringArray::from(vec![record_val])); + let expected_batch = + RecordBatch::try_new(expected_schema, vec![expected_array]).unwrap(); + assert_eq!(batch, expected_batch, "Test '{}' failed", test.name); + assert_eq!( + batch.schema().field(0).name(), + "f2", + "Test '{}' failed", + test.name + ); + } + (Err(e), Some(expected)) => { + assert!( + e.to_string().contains(expected), + "Test '{}' failed: unexpected error message at decode.\nExpected to contain: '{expected}'\nActual: '{e}'", + test.name, + ); + } + (Ok(batches), Some(expected)) => { + panic!( + "Test '{}' was expected to fail with '{expected}', but it succeeded with: {:?}", + test.name, batches + ); + } + (Err(e), None) => { + panic!( + "Test '{}' was not expected to fail, but it did with '{e}'", + test.name + ); + } + } + } + } + + #[test] + fn test_decimal() { + let files = [ + ("avro/fixed_length_decimal.avro", 25, 2), + ("avro/fixed_length_decimal_legacy.avro", 13, 2), + ("avro/int32_decimal.avro", 4, 2), + ("avro/int64_decimal.avro", 10, 2), + ]; + let decimal_values: Vec = (1..=24).map(|n| n as i128 * 100).collect(); + for (file, precision, scale) in files { + let file_path = arrow_test_data(file); + let actual_batch = read_file(&file_path, 8, false); + let expected_array = Decimal128Array::from_iter_values(decimal_values.clone()) + .with_precision_and_scale(precision, scale) + .unwrap(); + let mut meta = HashMap::new(); + meta.insert("precision".to_string(), precision.to_string()); + meta.insert("scale".to_string(), scale.to_string()); + let field_with_meta = Field::new("value", DataType::Decimal128(precision, scale), true) + .with_metadata(meta); + let expected_schema = Arc::new(Schema::new(vec![field_with_meta])); + let expected_batch = + RecordBatch::try_new(expected_schema.clone(), vec![Arc::new(expected_array)]) + .expect("Failed to build expected RecordBatch"); + assert_eq!( + actual_batch, expected_batch, + "Decoded RecordBatch does not match the expected Decimal128 data for file {file}" + ); + let actual_batch_small = read_file(&file_path, 3, false); + assert_eq!( + actual_batch_small, + expected_batch, + "Decoded RecordBatch does not match the expected Decimal128 data for file {file} with batch size 3" + ); } } + + #[test] + fn test_dict_pages_offset_zero() { + let file = arrow_test_data("avro/dict-page-offset-zero.avro"); + let batch = read_file(&file, 32, false); + let num_rows = batch.num_rows(); + let expected_field = Int32Array::from(vec![Some(1552); num_rows]); + let expected = RecordBatch::try_from_iter_with_nullable([( + "l_partkey", + Arc::new(expected_field) as Arc, + true, + )]) + .unwrap(); + assert_eq!(batch, expected); + } + + #[test] + fn test_list_columns() { + let file = arrow_test_data("avro/list_columns.avro"); + let mut int64_list_builder = ListBuilder::new(Int64Builder::new()); + { + { + let values = int64_list_builder.values(); + values.append_value(1); + values.append_value(2); + values.append_value(3); + } + int64_list_builder.append(true); + } + { + { + let values = int64_list_builder.values(); + values.append_null(); + values.append_value(1); + } + int64_list_builder.append(true); + } + { + { + let values = int64_list_builder.values(); + values.append_value(4); + } + int64_list_builder.append(true); + } + let int64_list = int64_list_builder.finish(); + let mut utf8_list_builder = ListBuilder::new(StringBuilder::new()); + { + { + let values = utf8_list_builder.values(); + values.append_value("abc"); + values.append_value("efg"); + values.append_value("hij"); + } + utf8_list_builder.append(true); + } + { + utf8_list_builder.append(false); + } + { + { + let values = utf8_list_builder.values(); + values.append_value("efg"); + values.append_null(); + values.append_value("hij"); + values.append_value("xyz"); + } + utf8_list_builder.append(true); + } + let utf8_list = utf8_list_builder.finish(); + let expected = RecordBatch::try_from_iter_with_nullable([ + ("int64_list", Arc::new(int64_list) as Arc, true), + ("utf8_list", Arc::new(utf8_list) as Arc, true), + ]) + .unwrap(); + let batch = read_file(&file, 8, false); + assert_eq!(batch, expected); + } + + #[test] + fn test_nested_lists() { + use arrow_data::ArrayDataBuilder; + let file = arrow_test_data("avro/nested_lists.snappy.avro"); + let inner_values = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + Some("f"), + ]); + let inner_offsets = Buffer::from_slice_ref([0, 2, 3, 3, 4, 6, 8, 8, 9, 11, 13, 14, 14, 15]); + let inner_validity = [ + true, true, false, true, true, true, false, true, true, true, true, false, true, + ]; + let inner_null_buffer = Buffer::from_iter(inner_validity.iter().copied()); + let inner_field = Field::new("item", DataType::Utf8, true); + let inner_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(inner_field))) + .len(13) + .add_buffer(inner_offsets) + .add_child_data(inner_values.to_data()) + .null_bit_buffer(Some(inner_null_buffer)) + .build() + .unwrap(); + let inner_list_array = ListArray::from(inner_list_data); + let middle_offsets = Buffer::from_slice_ref([0, 2, 4, 6, 8, 11, 13]); + let middle_validity = [true; 6]; + let middle_null_buffer = Buffer::from_iter(middle_validity.iter().copied()); + let middle_field = Field::new("item", inner_list_array.data_type().clone(), true); + let middle_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(middle_field))) + .len(6) + .add_buffer(middle_offsets) + .add_child_data(inner_list_array.to_data()) + .null_bit_buffer(Some(middle_null_buffer)) + .build() + .unwrap(); + let middle_list_array = ListArray::from(middle_list_data); + let outer_offsets = Buffer::from_slice_ref([0, 2, 4, 6]); + let outer_null_buffer = Buffer::from_slice_ref([0b111]); // all 3 rows valid + let outer_field = Field::new("item", middle_list_array.data_type().clone(), true); + let outer_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(outer_field))) + .len(3) + .add_buffer(outer_offsets) + .add_child_data(middle_list_array.to_data()) + .null_bit_buffer(Some(outer_null_buffer)) + .build() + .unwrap(); + let a_expected = ListArray::from(outer_list_data); + let b_expected = Int32Array::from(vec![1, 1, 1]); + let expected = RecordBatch::try_from_iter_with_nullable([ + ("a", Arc::new(a_expected) as Arc, true), + ("b", Arc::new(b_expected) as Arc, true), + ]) + .unwrap(); + let left = read_file(&file, 8, false); + assert_eq!(left, expected, "Mismatch for batch size=8"); + let left_small = read_file(&file, 3, false); + assert_eq!(left_small, expected, "Mismatch for batch size=3"); + } + + #[test] + fn test_simple() { + let tests = [ + ("avro/simple_enum.avro", 4, build_expected_enum(), 2), + ("avro/simple_fixed.avro", 2, build_expected_fixed(), 1), + ]; + + fn build_expected_enum() -> RecordBatch { + // Build the DictionaryArrays for f1, f2, f3 + let keys_f1 = Int32Array::from(vec![0, 1, 2, 3]); + let vals_f1 = StringArray::from(vec!["a", "b", "c", "d"]); + let f1_dict = + DictionaryArray::::try_new(keys_f1, Arc::new(vals_f1)).unwrap(); + let keys_f2 = Int32Array::from(vec![2, 3, 0, 1]); + let vals_f2 = StringArray::from(vec!["e", "f", "g", "h"]); + let f2_dict = + DictionaryArray::::try_new(keys_f2, Arc::new(vals_f2)).unwrap(); + let keys_f3 = Int32Array::from(vec![Some(1), Some(2), None, Some(0)]); + let vals_f3 = StringArray::from(vec!["i", "j", "k"]); + let f3_dict = + DictionaryArray::::try_new(keys_f3, Arc::new(vals_f3)).unwrap(); + let dict_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let mut md_f1 = HashMap::new(); + md_f1.insert( + "avro.enum.symbols".to_string(), + r#"["a","b","c","d"]"#.to_string(), + ); + let f1_field = Field::new("f1", dict_type.clone(), false).with_metadata(md_f1); + let mut md_f2 = HashMap::new(); + md_f2.insert( + "avro.enum.symbols".to_string(), + r#"["e","f","g","h"]"#.to_string(), + ); + let f2_field = Field::new("f2", dict_type.clone(), false).with_metadata(md_f2); + let mut md_f3 = HashMap::new(); + md_f3.insert( + "avro.enum.symbols".to_string(), + r#"["i","j","k"]"#.to_string(), + ); + let f3_field = Field::new("f3", dict_type.clone(), true).with_metadata(md_f3); + let expected_schema = Arc::new(Schema::new(vec![f1_field, f2_field, f3_field])); + RecordBatch::try_new( + expected_schema, + vec![ + Arc::new(f1_dict) as Arc, + Arc::new(f2_dict) as Arc, + Arc::new(f3_dict) as Arc, + ], + ) + .unwrap() + } + + fn build_expected_fixed() -> RecordBatch { + let f1 = + FixedSizeBinaryArray::try_from_iter(vec![b"abcde", b"12345"].into_iter()).unwrap(); + let f2 = + FixedSizeBinaryArray::try_from_iter(vec![b"fghijklmno", b"1234567890"].into_iter()) + .unwrap(); + let f3 = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + vec![Some(b"ABCDEF" as &[u8]), None].into_iter(), + 6, + ) + .unwrap(); + let expected_schema = Arc::new(Schema::new(vec![ + Field::new("f1", DataType::FixedSizeBinary(5), false), + Field::new("f2", DataType::FixedSizeBinary(10), false), + Field::new("f3", DataType::FixedSizeBinary(6), true), + ])); + RecordBatch::try_new( + expected_schema, + vec![ + Arc::new(f1) as Arc, + Arc::new(f2) as Arc, + Arc::new(f3) as Arc, + ], + ) + .unwrap() + } + for (file_name, batch_size, expected, alt_batch_size) in tests { + let file = arrow_test_data(file_name); + let actual = read_file(&file, batch_size, false); + assert_eq!(actual, expected); + let actual2 = read_file(&file, alt_batch_size, false); + assert_eq!(actual2, expected); + } + } + + #[test] + fn test_single_nan() { + let file = arrow_test_data("avro/single_nan.avro"); + let actual = read_file(&file, 1, false); + use arrow_array::Float64Array; + let schema = Arc::new(Schema::new(vec![Field::new( + "mycol", + DataType::Float64, + true, + )])); + let col = Float64Array::from(vec![None]); + let expected = RecordBatch::try_new(schema, vec![Arc::new(col)]).unwrap(); + assert_eq!(actual, expected); + let actual2 = read_file(&file, 2, false); + assert_eq!(actual2, expected); + } + + #[test] + fn test_duration_uuid() { + let batch = read_file("test/data/duration_uuid.avro", 4, false); + let schema = batch.schema(); + let fields = schema.fields(); + assert_eq!(fields.len(), 2); + assert_eq!(fields[0].name(), "duration_field"); + assert_eq!( + fields[0].data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + assert_eq!(fields[1].name(), "uuid_field"); + assert_eq!(fields[1].data_type(), &DataType::FixedSizeBinary(16)); + assert_eq!(batch.num_rows(), 4); + assert_eq!(batch.num_columns(), 2); + let duration_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let expected_duration_array: IntervalMonthDayNanoArray = [ + Some(IntervalMonthDayNanoType::make_value(1, 15, 500_000_000)), + Some(IntervalMonthDayNanoType::make_value(0, 5, 2_500_000_000)), + Some(IntervalMonthDayNanoType::make_value(2, 0, 0)), + Some(IntervalMonthDayNanoType::make_value(12, 31, 999_000_000)), + ] + .iter() + .copied() + .collect(); + assert_eq!(&expected_duration_array, duration_array); + let uuid_array = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let expected_uuid_array = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [ + Some([ + 0xfe, 0x7b, 0xc3, 0x0b, 0x4c, 0xe8, 0x4c, 0x5e, 0xb6, 0x7c, 0x22, 0x34, 0xa2, + 0xd3, 0x8e, 0x66, + ]), + Some([ + 0xb3, 0x3f, 0x2a, 0xd7, 0x97, 0xb4, 0x4d, 0xe1, 0x8b, 0xfe, 0x94, 0x94, 0x1d, + 0x60, 0x15, 0x6e, + ]), + Some([ + 0x5f, 0x74, 0x92, 0x64, 0x07, 0x4b, 0x40, 0x05, 0x84, 0xbf, 0x11, 0x5e, 0xa8, + 0x4e, 0xd2, 0x0a, + ]), + Some([ + 0x08, 0x26, 0xcc, 0x06, 0xd2, 0xe3, 0x45, 0x99, 0xb4, 0xad, 0xaf, 0x5f, 0xa6, + 0x90, 0x5c, 0xdb, + ]), + ] + .into_iter(), + 16, + ) + .unwrap(); + assert_eq!(&expected_uuid_array, uuid_array); + } + + #[test] + fn test_datapage_v2() { + let file = arrow_test_data("avro/datapage_v2.snappy.avro"); + let batch = read_file(&file, 8, false); + let a = StringArray::from(vec![ + Some("abc"), + Some("abc"), + Some("abc"), + None, + Some("abc"), + ]); + let b = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); + let c = Float64Array::from(vec![Some(2.0), Some(3.0), Some(4.0), Some(5.0), Some(2.0)]); + let d = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(false), + Some(true), + ]); + let e_values = Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(1), + Some(2), + Some(3), + Some(1), + Some(2), + ]); + let e_offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 3, 3, 6, 8])); + let e_validity = Some(NullBuffer::from(vec![true, false, false, true, true])); + let field_e = Arc::new(Field::new("item", DataType::Int32, true)); + let e = ListArray::new(field_e, e_offsets, Arc::new(e_values), e_validity); + let expected = RecordBatch::try_from_iter_with_nullable([ + ("a", Arc::new(a) as Arc, true), + ("b", Arc::new(b) as Arc, true), + ("c", Arc::new(c) as Arc, true), + ("d", Arc::new(d) as Arc, true), + ("e", Arc::new(e) as Arc, true), + ]) + .unwrap(); + assert_eq!(batch, expected); + } + + #[test] + fn test_nested_records() { + let f1_f1_1 = StringArray::from(vec!["aaa", "bbb"]); + let f1_f1_2 = Int32Array::from(vec![10, 20]); + let rounded_pi = (std::f64::consts::PI * 100.0).round() / 100.0; + let f1_f1_3_1 = Float64Array::from(vec![rounded_pi, rounded_pi]); + let f1_f1_3 = StructArray::from(vec![( + Arc::new(Field::new("f1_3_1", DataType::Float64, false)), + Arc::new(f1_f1_3_1) as Arc, + )]); + let f1_expected = StructArray::from(vec![ + ( + Arc::new(Field::new("f1_1", DataType::Utf8, false)), + Arc::new(f1_f1_1) as Arc, + ), + ( + Arc::new(Field::new("f1_2", DataType::Int32, false)), + Arc::new(f1_f1_2) as Arc, + ), + ( + Arc::new(Field::new( + "f1_3", + DataType::Struct(Fields::from(vec![Field::new( + "f1_3_1", + DataType::Float64, + false, + )])), + false, + )), + Arc::new(f1_f1_3) as Arc, + ), + ]); + + let f2_fields = vec![ + Field::new("f2_1", DataType::Boolean, false), + Field::new("f2_2", DataType::Float32, false), + ]; + let f2_struct_builder = StructBuilder::new( + f2_fields + .iter() + .map(|f| Arc::new(f.clone())) + .collect::>>(), + vec![ + Box::new(BooleanBuilder::new()) as Box, + Box::new(Float32Builder::new()) as Box, + ], + ); + let mut f2_list_builder = ListBuilder::new(f2_struct_builder); + { + let struct_builder = f2_list_builder.values(); + struct_builder.append(true); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_value(true); + } + { + let b = struct_builder.field_builder::(1).unwrap(); + b.append_value(1.2_f32); + } + struct_builder.append(true); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_value(true); + } + { + let b = struct_builder.field_builder::(1).unwrap(); + b.append_value(2.2_f32); + } + f2_list_builder.append(true); + } + { + let struct_builder = f2_list_builder.values(); + struct_builder.append(true); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_value(false); + } + { + let b = struct_builder.field_builder::(1).unwrap(); + b.append_value(10.2_f32); + } + f2_list_builder.append(true); + } + + let list_array_with_nullable_items = f2_list_builder.finish(); + + let item_field = Arc::new(Field::new( + "item", + list_array_with_nullable_items.values().data_type().clone(), + false, + )); + let list_data_type = DataType::List(item_field); + + let f2_array_data = list_array_with_nullable_items + .to_data() + .into_builder() + .data_type(list_data_type) + .build() + .unwrap(); + let f2_expected = ListArray::from(f2_array_data); + + let mut f3_struct_builder = StructBuilder::new( + vec![Arc::new(Field::new("f3_1", DataType::Utf8, false))], + vec![Box::new(StringBuilder::new()) as Box], + ); + f3_struct_builder.append(true); + { + let b = f3_struct_builder.field_builder::(0).unwrap(); + b.append_value("xyz"); + } + f3_struct_builder.append(false); + { + let b = f3_struct_builder.field_builder::(0).unwrap(); + b.append_null(); + } + let f3_expected = f3_struct_builder.finish(); + let f4_fields = [Field::new("f4_1", DataType::Int64, false)]; + let f4_struct_builder = StructBuilder::new( + f4_fields + .iter() + .map(|f| Arc::new(f.clone())) + .collect::>>(), + vec![Box::new(Int64Builder::new()) as Box], + ); + let mut f4_list_builder = ListBuilder::new(f4_struct_builder); + { + let struct_builder = f4_list_builder.values(); + struct_builder.append(true); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_value(200); + } + struct_builder.append(false); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_null(); + } + f4_list_builder.append(true); + } + { + let struct_builder = f4_list_builder.values(); + struct_builder.append(false); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_null(); + } + struct_builder.append(true); + { + let b = struct_builder.field_builder::(0).unwrap(); + b.append_value(300); + } + f4_list_builder.append(true); + } + let f4_expected = f4_list_builder.finish(); + + let expected = RecordBatch::try_from_iter_with_nullable([ + ("f1", Arc::new(f1_expected) as Arc, false), + ("f2", Arc::new(f2_expected) as Arc, false), + ("f3", Arc::new(f3_expected) as Arc, true), + ("f4", Arc::new(f4_expected) as Arc, false), + ]) + .unwrap(); + + let file = arrow_test_data("avro/nested_records.avro"); + let batch_large = read_file(&file, 8, false); + assert_eq!( + batch_large, expected, + "Decoded RecordBatch does not match expected data for nested records (batch size 8)" + ); + let batch_small = read_file(&file, 3, false); + assert_eq!( + batch_small, expected, + "Decoded RecordBatch does not match expected data for nested records (batch size 3)" + ); + } + + #[test] + fn test_repeated_no_annotation() { + let file = arrow_test_data("avro/repeated_no_annotation.avro"); + let batch_large = read_file(&file, 8, false); + use arrow_array::{Int32Array, Int64Array, ListArray, StringArray, StructArray}; + use arrow_buffer::Buffer; + use arrow_schema::{DataType, Field, Fields}; + let id_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let number_array = Int64Array::from(vec![ + Some(5555555555), + Some(1111111111), + Some(1111111111), + Some(2222222222), + Some(3333333333), + ]); + let kind_array = + StringArray::from(vec![None, Some("home"), Some("home"), None, Some("mobile")]); + let phone_fields = Fields::from(vec![ + Field::new("number", DataType::Int64, true), + Field::new("kind", DataType::Utf8, true), + ]); + let phone_struct_data = ArrayDataBuilder::new(DataType::Struct(phone_fields)) + .len(5) + .child_data(vec![number_array.into_data(), kind_array.into_data()]) + .build() + .unwrap(); + let phone_struct_array = StructArray::from(phone_struct_data); + let phone_list_offsets = Buffer::from_slice_ref([0, 0, 0, 0, 1, 2, 5]); + let phone_list_validity = Buffer::from_iter([false, false, true, true, true, true]); + let phone_item_field = Field::new("item", phone_struct_array.data_type().clone(), true); + let phone_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(phone_item_field))) + .len(6) + .add_buffer(phone_list_offsets) + .null_bit_buffer(Some(phone_list_validity)) + .child_data(vec![phone_struct_array.into_data()]) + .build() + .unwrap(); + let phone_list_array = ListArray::from(phone_list_data); + let phone_numbers_validity = Buffer::from_iter([false, false, true, true, true, true]); + let phone_numbers_field = Field::new("phone", phone_list_array.data_type().clone(), true); + let phone_numbers_struct_data = + ArrayDataBuilder::new(DataType::Struct(Fields::from(vec![phone_numbers_field]))) + .len(6) + .null_bit_buffer(Some(phone_numbers_validity)) + .child_data(vec![phone_list_array.into_data()]) + .build() + .unwrap(); + let phone_numbers_struct_array = StructArray::from(phone_numbers_struct_data); + let expected = arrow_array::RecordBatch::try_from_iter_with_nullable([ + ("id", Arc::new(id_array) as _, true), + ( + "phoneNumbers", + Arc::new(phone_numbers_struct_array) as _, + true, + ), + ]) + .unwrap(); + assert_eq!(batch_large, expected, "Mismatch for batch_size=8"); + let batch_small = read_file(&file, 3, false); + assert_eq!(batch_small, expected, "Mismatch for batch_size=3"); + } + + #[test] + fn test_nonnullable_impala() { + let file = arrow_test_data("avro/nonnullable.impala.avro"); + let id = Int64Array::from(vec![Some(8)]); + let mut int_array_builder = ListBuilder::new(Int32Builder::new()); + { + let vb = int_array_builder.values(); + vb.append_value(-1); + } + int_array_builder.append(true); // finalize one sub-list + let int_array = int_array_builder.finish(); + let mut iaa_builder = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + { + let inner_list_builder = iaa_builder.values(); + { + let vb = inner_list_builder.values(); + vb.append_value(-1); + vb.append_value(-2); + } + inner_list_builder.append(true); + inner_list_builder.append(true); + } + iaa_builder.append(true); + let int_array_array = iaa_builder.finish(); + use arrow_array::builder::MapFieldNames; + let field_names = MapFieldNames { + entry: "entries".to_string(), + key: "key".to_string(), + value: "value".to_string(), + }; + let mut int_map_builder = + MapBuilder::new(Some(field_names), StringBuilder::new(), Int32Builder::new()); + { + let (keys, vals) = int_map_builder.entries(); + keys.append_value("k1"); + vals.append_value(-1); + } + int_map_builder.append(true).unwrap(); // finalize map for row 0 + let int_map = int_map_builder.finish(); + let field_names2 = MapFieldNames { + entry: "entries".to_string(), + key: "key".to_string(), + value: "value".to_string(), + }; + let mut ima_builder = ListBuilder::new(MapBuilder::new( + Some(field_names2), + StringBuilder::new(), + Int32Builder::new(), + )); + { + let map_builder = ima_builder.values(); + map_builder.append(true).unwrap(); + { + let (keys, vals) = map_builder.entries(); + keys.append_value("k1"); + vals.append_value(1); + } + map_builder.append(true).unwrap(); + map_builder.append(true).unwrap(); + map_builder.append(true).unwrap(); + } + ima_builder.append(true); + let int_map_array_ = ima_builder.finish(); + let mut nested_sb = StructBuilder::new( + vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new( + "B", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )), + Arc::new(Field::new( + "c", + DataType::Struct( + vec![Field::new( + "D", + DataType::List(Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct( + vec![ + Field::new("e", DataType::Int32, true), + Field::new("f", DataType::Utf8, true), + ] + .into(), + ), + true, + ))), + true, + ))), + true, + )] + .into(), + ), + true, + )), + Arc::new(Field::new( + "G", + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new( + "value", + DataType::Struct( + vec![Field::new( + "h", + DataType::Struct( + vec![Field::new( + "i", + DataType::List(Arc::new(Field::new( + "item", + DataType::Float64, + true, + ))), + true, + )] + .into(), + ), + true, + )] + .into(), + ), + true, + ), + ] + .into(), + ), + false, + )), + false, + ), + true, + )), + ], + vec![ + Box::new(Int32Builder::new()), + Box::new(ListBuilder::new(Int32Builder::new())), + { + let d_field = Field::new( + "D", + DataType::List(Arc::new(Field::new( + "item", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct( + vec![ + Field::new("e", DataType::Int32, true), + Field::new("f", DataType::Utf8, true), + ] + .into(), + ), + true, + ))), + true, + ))), + true, + ); + Box::new(StructBuilder::new( + vec![Arc::new(d_field)], + vec![Box::new({ + let ef_struct_builder = StructBuilder::new( + vec![ + Arc::new(Field::new("e", DataType::Int32, true)), + Arc::new(Field::new("f", DataType::Utf8, true)), + ], + vec![ + Box::new(Int32Builder::new()), + Box::new(StringBuilder::new()), + ], + ); + let list_of_ef = ListBuilder::new(ef_struct_builder); + ListBuilder::new(list_of_ef) + })], + )) + }, + { + let map_field_names = MapFieldNames { + entry: "entries".to_string(), + key: "key".to_string(), + value: "value".to_string(), + }; + let i_list_builder = ListBuilder::new(Float64Builder::new()); + let h_struct = StructBuilder::new( + vec![Arc::new(Field::new( + "i", + DataType::List(Arc::new(Field::new("item", DataType::Float64, true))), + true, + ))], + vec![Box::new(i_list_builder)], + ); + let g_value_builder = StructBuilder::new( + vec![Arc::new(Field::new( + "h", + DataType::Struct( + vec![Field::new( + "i", + DataType::List(Arc::new(Field::new( + "item", + DataType::Float64, + true, + ))), + true, + )] + .into(), + ), + true, + ))], + vec![Box::new(h_struct)], + ); + Box::new(MapBuilder::new( + Some(map_field_names), + StringBuilder::new(), + g_value_builder, + )) + }, + ], + ); + nested_sb.append(true); + { + let a_builder = nested_sb.field_builder::(0).unwrap(); + a_builder.append_value(-1); + } + { + let b_builder = nested_sb + .field_builder::>(1) + .unwrap(); + { + let vb = b_builder.values(); + vb.append_value(-1); + } + b_builder.append(true); + } + { + let c_struct_builder = nested_sb.field_builder::(2).unwrap(); + c_struct_builder.append(true); + let d_list_builder = c_struct_builder + .field_builder::>>(0) + .unwrap(); + { + let sub_list_builder = d_list_builder.values(); + { + let ef_struct = sub_list_builder.values(); + ef_struct.append(true); + { + let e_b = ef_struct.field_builder::(0).unwrap(); + e_b.append_value(-1); + let f_b = ef_struct.field_builder::(1).unwrap(); + f_b.append_value("nonnullable"); + } + sub_list_builder.append(true); + } + d_list_builder.append(true); + } + } + { + let g_map_builder = nested_sb + .field_builder::>(3) + .unwrap(); + g_map_builder.append(true).unwrap(); + } + let nested_struct = nested_sb.finish(); + let expected = RecordBatch::try_from_iter_with_nullable([ + ("ID", Arc::new(id) as Arc, true), + ("Int_Array", Arc::new(int_array), true), + ("int_array_array", Arc::new(int_array_array), true), + ("Int_Map", Arc::new(int_map), true), + ("int_map_array", Arc::new(int_map_array_), true), + ("nested_Struct", Arc::new(nested_struct), true), + ]) + .unwrap(); + let batch_large = read_file(&file, 8, false); + assert_eq!(batch_large, expected, "Mismatch for batch_size=8"); + let batch_small = read_file(&file, 3, false); + assert_eq!(batch_small, expected, "Mismatch for batch_size=3"); + } + + #[test] + fn test_nonnullable_impala_strict() { + let file = arrow_test_data("avro/nonnullable.impala.avro"); + let err = read_file_strict(&file, 8, false).unwrap_err(); + assert!(err.to_string().contains( + "Found Avro union of the form ['T','null'], which is disallowed in strict_mode" + )); + } + + #[test] + fn test_nullable_impala() { + let file = arrow_test_data("avro/nullable.impala.avro"); + let batch1 = read_file(&file, 3, false); + let batch2 = read_file(&file, 8, false); + assert_eq!(batch1, batch2); + let batch = batch1; + assert_eq!(batch.num_rows(), 7); + let id_array = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("id column should be an Int64Array"); + let expected_ids = [1, 2, 3, 4, 5, 6, 7]; + for (i, &expected_id) in expected_ids.iter().enumerate() { + assert_eq!(id_array.value(i), expected_id, "Mismatch in id at row {i}",); + } + let int_array = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("int_array column should be a ListArray"); + { + let offsets = int_array.value_offsets(); + let start = offsets[0] as usize; + let end = offsets[1] as usize; + let values = int_array + .values() + .as_any() + .downcast_ref::() + .expect("Values of int_array should be an Int32Array"); + let row0: Vec> = (start..end).map(|i| Some(values.value(i))).collect(); + assert_eq!( + row0, + vec![Some(1), Some(2), Some(3)], + "Mismatch in int_array row 0" + ); + } + let nested_struct = batch + .column(5) + .as_any() + .downcast_ref::() + .expect("nested_struct column should be a StructArray"); + let a_array = nested_struct + .column_by_name("A") + .expect("Field A should exist in nested_struct") + .as_any() + .downcast_ref::() + .expect("Field A should be an Int32Array"); + assert_eq!(a_array.value(0), 1, "Mismatch in nested_struct.A at row 0"); + assert!( + !a_array.is_valid(1), + "Expected null in nested_struct.A at row 1" + ); + assert!( + !a_array.is_valid(3), + "Expected null in nested_struct.A at row 3" + ); + assert_eq!(a_array.value(6), 7, "Mismatch in nested_struct.A at row 6"); + } + + #[test] + fn test_nullable_impala_strict() { + let file = arrow_test_data("avro/nullable.impala.avro"); + let err = read_file_strict(&file, 8, false).unwrap_err(); + assert!(err.to_string().contains( + "Found Avro union of the form ['T','null'], which is disallowed in strict_mode" + )); + } } diff --git a/arrow-avro/src/reader/record.rs b/arrow-avro/src/reader/record.rs index 52a58cf63303..180afcd2d8c3 100644 --- a/arrow-avro/src/reader/record.rs +++ b/arrow-avro/src/reader/record.rs @@ -20,28 +20,92 @@ use crate::reader::block::{Block, BlockDecoder}; use crate::reader::cursor::AvroCursor; use crate::reader::header::Header; use crate::schema::*; +use arrow_array::builder::{ + ArrayBuilder, Decimal128Builder, Decimal256Builder, IntervalMonthDayNanoBuilder, + PrimitiveBuilder, +}; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::*; use arrow_schema::{ - ArrowError, DataType, Field as ArrowField, FieldRef, Fields, Schema as ArrowSchema, SchemaRef, + ArrowError, DataType, Field as ArrowField, FieldRef, Fields, IntervalUnit, + Schema as ArrowSchema, SchemaRef, DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, }; +use std::cmp::Ordering; use std::collections::HashMap; use std::io::Read; use std::sync::Arc; +use uuid::Uuid; + +const DEFAULT_CAPACITY: usize = 1024; + +#[derive(Debug)] +pub(crate) struct RecordDecoderBuilder<'a> { + data_type: &'a AvroDataType, + use_utf8view: bool, +} + +impl<'a> RecordDecoderBuilder<'a> { + pub(crate) fn new(data_type: &'a AvroDataType) -> Self { + Self { + data_type, + use_utf8view: false, + } + } + + pub(crate) fn with_utf8_view(mut self, use_utf8view: bool) -> Self { + self.use_utf8view = use_utf8view; + self + } + + /// Builds the `RecordDecoder`. + pub(crate) fn build(self) -> Result { + RecordDecoder::try_new_with_options(self.data_type, self.use_utf8view) + } +} /// Decodes avro encoded data into [`RecordBatch`] -pub struct RecordDecoder { +#[derive(Debug)] +pub(crate) struct RecordDecoder { schema: SchemaRef, fields: Vec, + use_utf8view: bool, } impl RecordDecoder { - pub fn try_new(data_type: &AvroDataType) -> Result { + /// Creates a new `RecordDecoderBuilder` for configuring a `RecordDecoder`. + pub(crate) fn new(data_type: &'_ AvroDataType) -> Self { + RecordDecoderBuilder::new(data_type).build().unwrap() + } + + /// Create a new [`RecordDecoder`] from the provided [`AvroDataType`] with default options + pub(crate) fn try_new(data_type: &AvroDataType) -> Result { + RecordDecoderBuilder::new(data_type) + .with_utf8_view(true) + .build() + } + + /// Creates a new [`RecordDecoder`] from the provided [`AvroDataType`] with additional options. + /// + /// This method allows you to customize how the Avro data is decoded into Arrow arrays. + /// + /// # Arguments + /// * `data_type` - The Avro data type to decode. + /// * `use_utf8view` - A flag indicating whether to use `Utf8View` for string types. + /// * `strict_mode` - A flag to enable strict decoding, returning an error if the data + /// does not conform to the schema. + /// + /// # Errors + /// This function will return an error if the provided `data_type` is not a `Record`. + pub(crate) fn try_new_with_options( + data_type: &AvroDataType, + use_utf8view: bool, + ) -> Result { match Decoder::try_new(data_type)? { Decoder::Record(fields, encodings) => Ok(Self { schema: Arc::new(ArrowSchema::new(fields)), fields: encodings, + use_utf8view, }), encoding => Err(ArrowError::ParseError(format!( "Expected record got {encoding:?}" @@ -49,12 +113,13 @@ impl RecordDecoder { } } - pub fn schema(&self) -> &SchemaRef { + /// Returns the decoder's `SchemaRef` + pub(crate) fn schema(&self) -> &SchemaRef { &self.schema } /// Decode `count` records from `buf` - pub fn decode(&mut self, buf: &[u8], count: usize) -> Result { + pub(crate) fn decode(&mut self, buf: &[u8], count: usize) -> Result { let mut cursor = AvroCursor::new(buf); for _ in 0..count { for field in &mut self.fields { @@ -65,7 +130,7 @@ impl RecordDecoder { } /// Flush the decoded records into a [`RecordBatch`] - pub fn flush(&mut self) -> Result { + pub(crate) fn flush(&mut self) -> Result { let arrays = self .fields .iter_mut() @@ -90,16 +155,30 @@ enum Decoder { TimestampMillis(bool, Vec), TimestampMicros(bool, Vec), Binary(OffsetBufferBuilder, Vec), + /// String data encoded as UTF-8 bytes, mapped to Arrow's StringArray String(OffsetBufferBuilder, Vec), - List(FieldRef, OffsetBufferBuilder, Box), + /// String data encoded as UTF-8 bytes, but mapped to Arrow's StringViewArray + StringView(OffsetBufferBuilder, Vec), + Array(FieldRef, OffsetBufferBuilder, Box), Record(Fields, Vec), + Map( + FieldRef, + OffsetBufferBuilder, + OffsetBufferBuilder, + Vec, + Box, + ), + Fixed(i32, Vec), + Enum(Vec, Arc<[String]>), + Duration(IntervalMonthDayNanoBuilder), + Uuid(Vec), + Decimal128(usize, Option, Option, Decimal128Builder), + Decimal256(usize, Option, Option, Decimal256Builder), Nullable(Nullability, NullBufferBuilder, Box), } impl Decoder { fn try_new(data_type: &AvroDataType) -> Result { - let nyi = |s: &str| Err(ArrowError::NotYetImplemented(s.to_string())); - let decoder = match data_type.codec() { Codec::Null => Self::Null(0), Codec::Boolean => Self::Boolean(BooleanBufferBuilder::new(DEFAULT_CAPACITY)), @@ -115,6 +194,10 @@ impl Decoder { OffsetBufferBuilder::new(DEFAULT_CAPACITY), Vec::with_capacity(DEFAULT_CAPACITY), ), + Codec::Utf8View => Self::StringView( + OffsetBufferBuilder::new(DEFAULT_CAPACITY), + Vec::with_capacity(DEFAULT_CAPACITY), + ), Codec::Date32 => Self::Date32(Vec::with_capacity(DEFAULT_CAPACITY)), Codec::TimeMillis => Self::TimeMillis(Vec::with_capacity(DEFAULT_CAPACITY)), Codec::TimeMicros => Self::TimeMicros(Vec::with_capacity(DEFAULT_CAPACITY)), @@ -124,16 +207,58 @@ impl Decoder { Codec::TimestampMicros(is_utc) => { Self::TimestampMicros(*is_utc, Vec::with_capacity(DEFAULT_CAPACITY)) } - Codec::Fixed(_) => return nyi("decoding fixed"), - Codec::Interval => return nyi("decoding interval"), + Codec::Fixed(sz) => Self::Fixed(*sz, Vec::with_capacity(DEFAULT_CAPACITY)), + Codec::Decimal(precision, scale, size) => { + let p = *precision; + let s = *scale; + let sz = *size; + let prec = p as u8; + let scl = s.unwrap_or(0) as i8; + match (sz, p) { + (Some(fixed_size), _) if fixed_size <= 16 => { + let builder = + Decimal128Builder::new().with_precision_and_scale(prec, scl)?; + Self::Decimal128(p, s, sz, builder) + } + (Some(fixed_size), _) if fixed_size <= 32 => { + let builder = + Decimal256Builder::new().with_precision_and_scale(prec, scl)?; + Self::Decimal256(p, s, sz, builder) + } + (Some(fixed_size), _) => { + return Err(ArrowError::ParseError(format!( + "Unsupported decimal size: {fixed_size:?}" + ))); + } + (None, p) if p <= DECIMAL128_MAX_PRECISION as usize => { + let builder = + Decimal128Builder::new().with_precision_and_scale(prec, scl)?; + Self::Decimal128(p, s, sz, builder) + } + (None, p) if p <= DECIMAL256_MAX_PRECISION as usize => { + let builder = + Decimal256Builder::new().with_precision_and_scale(prec, scl)?; + Self::Decimal256(p, s, sz, builder) + } + (None, _) => { + return Err(ArrowError::ParseError(format!( + "Decimal precision {p} exceeds maximum supported" + ))); + } + } + } + Codec::Interval => Self::Duration(IntervalMonthDayNanoBuilder::new()), Codec::List(item) => { let decoder = Self::try_new(item)?; - Self::List( + Self::Array( Arc::new(item.field_with_name("item")), OffsetBufferBuilder::new(DEFAULT_CAPACITY), Box::new(decoder), ) } + Codec::Enum(symbols) => { + Self::Enum(Vec::with_capacity(DEFAULT_CAPACITY), symbols.clone()) + } Codec::Struct(fields) => { let mut arrow_fields = Vec::with_capacity(fields.len()); let mut encodings = Vec::with_capacity(fields.len()); @@ -144,8 +269,27 @@ impl Decoder { } Self::Record(arrow_fields.into(), encodings) } + Codec::Map(child) => { + let val_field = child.field_with_name("value").with_nullable(true); + let map_field = Arc::new(ArrowField::new( + "entries", + DataType::Struct(Fields::from(vec![ + ArrowField::new("key", DataType::Utf8, false), + val_field, + ])), + false, + )); + let val_dec = Self::try_new(child)?; + Self::Map( + map_field, + OffsetBufferBuilder::new(DEFAULT_CAPACITY), + OffsetBufferBuilder::new(DEFAULT_CAPACITY), + Vec::with_capacity(DEFAULT_CAPACITY), + Box::new(val_dec), + ) + } + Codec::Uuid => Self::Uuid(Vec::with_capacity(DEFAULT_CAPACITY)), }; - Ok(match data_type.nullability() { Some(nullability) => Self::Nullable( nullability, @@ -168,13 +312,30 @@ impl Decoder { | Self::TimestampMicros(_, v) => v.push(0), Self::Float32(v) => v.push(0.), Self::Float64(v) => v.push(0.), - Self::Binary(offsets, _) | Self::String(offsets, _) => offsets.push_length(0), - Self::List(_, offsets, e) => { + Self::Binary(offsets, _) | Self::String(offsets, _) | Self::StringView(offsets, _) => { + offsets.push_length(0); + } + Self::Uuid(v) => { + v.extend([0; 16]); + } + Self::Array(_, offsets, e) => { offsets.push_length(0); - e.append_null(); } Self::Record(_, e) => e.iter_mut().for_each(|e| e.append_null()), - Self::Nullable(_, _, _) => unreachable!("Nulls cannot be nested"), + Self::Map(_, _koff, moff, _, _) => { + moff.push_length(0); + } + Self::Fixed(sz, accum) => { + accum.extend(std::iter::repeat_n(0u8, *sz as usize)); + } + Self::Decimal128(_, _, _, builder) => builder.append_value(0), + Self::Decimal256(_, _, _, builder) => builder.append_value(i256::ZERO), + Self::Enum(indices, _) => indices.push(0), + Self::Duration(builder) => builder.append_null(), + Self::Nullable(_, null_buffer, inner) => { + null_buffer.append(false); + inner.append_null(); + } } } @@ -192,27 +353,86 @@ impl Decoder { | Self::TimestampMicros(_, values) => values.push(buf.get_long()?), Self::Float32(values) => values.push(buf.get_float()?), Self::Float64(values) => values.push(buf.get_double()?), - Self::Binary(offsets, values) | Self::String(offsets, values) => { + Self::Binary(offsets, values) + | Self::String(offsets, values) + | Self::StringView(offsets, values) => { let data = buf.get_bytes()?; offsets.push_length(data.len()); values.extend_from_slice(data); } - Self::List(_, _, _) => { - return Err(ArrowError::NotYetImplemented( - "Decoding ListArray".to_string(), - )) + Self::Uuid(values) => { + let s_bytes = buf.get_bytes()?; + let s = std::str::from_utf8(s_bytes).map_err(|e| { + ArrowError::ParseError(format!("UUID bytes are not valid UTF-8: {e}")) + })?; + let uuid = Uuid::try_parse(s) + .map_err(|e| ArrowError::ParseError(format!("Failed to parse uuid: {e}")))?; + values.extend_from_slice(uuid.as_bytes()); + } + Self::Array(_, off, encoding) => { + let total_items = read_blocks(buf, |cursor| encoding.decode(cursor))?; + off.push_length(total_items); } Self::Record(_, encodings) => { for encoding in encodings { encoding.decode(buf)?; } } - Self::Nullable(nullability, nulls, e) => { - let is_valid = buf.get_bool()? == matches!(nullability, Nullability::NullFirst); - nulls.append(is_valid); - match is_valid { - true => e.decode(buf)?, - false => e.append_null(), + Self::Map(_, koff, moff, kdata, valdec) => { + let newly_added = read_blocks(buf, |cur| { + let kb = cur.get_bytes()?; + koff.push_length(kb.len()); + kdata.extend_from_slice(kb); + valdec.decode(cur) + })?; + moff.push_length(newly_added); + } + Self::Fixed(sz, accum) => { + let fx = buf.get_fixed(*sz as usize)?; + accum.extend_from_slice(fx); + } + Self::Decimal128(_, _, size, builder) => { + let raw = if let Some(s) = size { + buf.get_fixed(*s)? + } else { + buf.get_bytes()? + }; + let ext = sign_extend_to::<16>(raw)?; + let val = i128::from_be_bytes(ext); + builder.append_value(val); + } + Self::Decimal256(_, _, size, builder) => { + let raw = if let Some(s) = size { + buf.get_fixed(*s)? + } else { + buf.get_bytes()? + }; + let ext = sign_extend_to::<32>(raw)?; + let val = i256::from_be_bytes(ext); + builder.append_value(val); + } + Self::Enum(indices, _) => { + indices.push(buf.get_int()?); + } + Self::Duration(builder) => { + let b = buf.get_fixed(12)?; + let months = u32::from_le_bytes(b[0..4].try_into().unwrap()); + let days = u32::from_le_bytes(b[4..8].try_into().unwrap()); + let millis = u32::from_le_bytes(b[8..12].try_into().unwrap()); + let nanos = (millis as i64) * 1_000_000; + builder.append_value(IntervalMonthDayNano::new(months as i32, days as i32, nanos)); + } + Self::Nullable(order, nb, encoding) => { + let branch = buf.read_vlq()?; + let is_not_null = match *order { + Nullability::NullFirst => branch != 0, + Nullability::NullSecond => branch == 0, + }; + nb.append(is_not_null); + if is_not_null { + encoding.decode(buf)?; + } else { + encoding.append_null(); } } } @@ -244,7 +464,6 @@ impl Decoder { ), Self::Float32(values) => Arc::new(flush_primitive::(values, nulls)), Self::Float64(values) => Arc::new(flush_primitive::(values, nulls)), - Self::Binary(offsets, values) => { let offsets = flush_offsets(offsets); let values = flush_values(values).into(); @@ -255,7 +474,22 @@ impl Decoder { let values = flush_values(values).into(); Arc::new(StringArray::new(offsets, values, nulls)) } - Self::List(field, offsets, values) => { + Self::StringView(offsets, values) => { + let offsets = flush_offsets(offsets); + let values = flush_values(values); + let array = StringArray::new(offsets, values.into(), nulls.clone()); + let values: Vec<&str> = (0..array.len()) + .map(|i| { + if array.is_valid(i) { + array.value(i) + } else { + "" + } + }) + .collect(); + Arc::new(StringViewArray::from(values)) + } + Self::Array(field, offsets, values) => { let values = values.flush(None)?; let offsets = flush_offsets(offsets); Arc::new(ListArray::new(field.clone(), offsets, values, nulls)) @@ -267,10 +501,131 @@ impl Decoder { .collect::, _>>()?; Arc::new(StructArray::new(fields.clone(), arrays, nulls)) } + Self::Map(map_field, k_off, m_off, kdata, valdec) => { + let moff = flush_offsets(m_off); + let koff = flush_offsets(k_off); + let kd = flush_values(kdata).into(); + let val_arr = valdec.flush(None)?; + let key_arr = StringArray::new(koff, kd, None); + if key_arr.len() != val_arr.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Map keys length ({}) != map values length ({})", + key_arr.len(), + val_arr.len() + ))); + } + let final_len = moff.len() - 1; + if let Some(n) = &nulls { + if n.len() != final_len { + return Err(ArrowError::InvalidArgumentError(format!( + "Map array null buffer length {} != final map length {final_len}", + n.len() + ))); + } + } + let entries_struct = StructArray::new( + Fields::from(vec![ + Arc::new(ArrowField::new("key", DataType::Utf8, false)), + Arc::new(ArrowField::new("value", val_arr.data_type().clone(), true)), + ]), + vec![Arc::new(key_arr), val_arr], + None, + ); + let map_arr = MapArray::new(map_field.clone(), moff, entries_struct, nulls, false); + Arc::new(map_arr) + } + Self::Fixed(sz, accum) => { + let b: Buffer = flush_values(accum).into(); + let arr = FixedSizeBinaryArray::try_new(*sz, b, nulls) + .map_err(|e| ArrowError::ParseError(e.to_string()))?; + Arc::new(arr) + } + Self::Uuid(values) => { + let arr = FixedSizeBinaryArray::try_new(16, std::mem::take(values).into(), nulls) + .map_err(|e| ArrowError::ParseError(e.to_string()))?; + Arc::new(arr) + } + Self::Decimal128(precision, scale, _, builder) => { + let (_, vals, _) = builder.finish().into_parts(); + let scl = scale.unwrap_or(0); + let dec = Decimal128Array::new(vals, nulls) + .with_precision_and_scale(*precision as u8, scl as i8) + .map_err(|e| ArrowError::ParseError(e.to_string()))?; + Arc::new(dec) + } + Self::Decimal256(precision, scale, _, builder) => { + let (_, vals, _) = builder.finish().into_parts(); + let scl = scale.unwrap_or(0); + let dec = Decimal256Array::new(vals, nulls) + .with_precision_and_scale(*precision as u8, scl as i8) + .map_err(|e| ArrowError::ParseError(e.to_string()))?; + Arc::new(dec) + } + Self::Enum(indices, symbols) => { + let keys = flush_primitive::(indices, nulls); + let values = Arc::new(StringArray::from( + symbols.iter().map(|s| s.as_str()).collect::>(), + )); + Arc::new(DictionaryArray::try_new(keys, values)?) + } + Self::Duration(builder) => { + let (_, vals, _) = builder.finish().into_parts(); + let vals = IntervalMonthDayNanoArray::try_new(vals, nulls) + .map_err(|e| ArrowError::ParseError(e.to_string()))?; + Arc::new(vals) + } }) } } +#[inline] +fn read_blocks( + buf: &mut AvroCursor, + decode_entry: impl FnMut(&mut AvroCursor) -> Result<(), ArrowError>, +) -> Result { + read_blockwise_items(buf, true, decode_entry) +} + +#[inline] +fn read_blockwise_items( + buf: &mut AvroCursor, + read_size_after_negative: bool, + mut decode_fn: impl FnMut(&mut AvroCursor) -> Result<(), ArrowError>, +) -> Result { + let mut total = 0usize; + loop { + // Read the block count + // positive = that many items + // negative = that many items + read block size + // See: https://avro.apache.org/docs/1.11.1/specification/#maps + let block_count = buf.get_long()?; + match block_count.cmp(&0) { + Ordering::Equal => break, + Ordering::Less => { + // If block_count is negative, read the absolute value of count, + // then read the block size as a long and discard + let count = (-block_count) as usize; + if read_size_after_negative { + let _size_in_bytes = buf.get_long()?; + } + for _ in 0..count { + decode_fn(buf)?; + } + total += count; + } + Ordering::Greater => { + // If block_count is positive, decode that many items + let count = block_count as usize; + for _i in 0..count { + decode_fn(buf)?; + } + total += count; + } + } + } + Ok(total) +} + #[inline] fn flush_values(values: &mut Vec) -> Vec { std::mem::replace(values, Vec::with_capacity(DEFAULT_CAPACITY)) @@ -289,4 +644,548 @@ fn flush_primitive( PrimitiveArray::new(flush_values(values).into(), nulls) } -const DEFAULT_CAPACITY: usize = 1024; +/// Sign extends a byte slice to a fixed-size array of N bytes. +/// This is done by filling the leading bytes with 0x00 for positive numbers +/// or 0xFF for negative numbers. +#[inline] +fn sign_extend_to(raw: &[u8]) -> Result<[u8; N], ArrowError> { + if raw.len() > N { + return Err(ArrowError::ParseError(format!( + "Cannot extend a slice of length {} to {} bytes.", + raw.len(), + N + ))); + } + let mut arr = [0u8; N]; + let pad_len = N - raw.len(); + // Determine the byte to use for padding based on the sign bit of the raw data. + let extension_byte = if raw.is_empty() || (raw[0] & 0x80 == 0) { + 0x00 + } else { + 0xFF + }; + arr[..pad_len].fill(extension_byte); + arr[pad_len..].copy_from_slice(raw); + Ok(arr) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ + cast::AsArray, Array, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, + IntervalMonthDayNanoArray, ListArray, MapArray, StringArray, StructArray, + }; + + fn encode_avro_int(value: i32) -> Vec { + let mut buf = Vec::new(); + let mut v = (value << 1) ^ (value >> 31); + while v & !0x7F != 0 { + buf.push(((v & 0x7F) | 0x80) as u8); + v >>= 7; + } + buf.push(v as u8); + buf + } + + fn encode_avro_long(value: i64) -> Vec { + let mut buf = Vec::new(); + let mut v = (value << 1) ^ (value >> 63); + while v & !0x7F != 0 { + buf.push(((v & 0x7F) | 0x80) as u8); + v >>= 7; + } + buf.push(v as u8); + buf + } + + fn encode_avro_bytes(bytes: &[u8]) -> Vec { + let mut buf = encode_avro_long(bytes.len() as i64); + buf.extend_from_slice(bytes); + buf + } + + fn avro_from_codec(codec: Codec) -> AvroDataType { + AvroDataType::new(codec, Default::default(), None) + } + + #[test] + fn test_map_decoding_one_entry() { + let value_type = avro_from_codec(Codec::Utf8); + let map_type = avro_from_codec(Codec::Map(Arc::new(value_type))); + let mut decoder = Decoder::try_new(&map_type).unwrap(); + // Encode a single map with one entry: {"hello": "world"} + let mut data = Vec::new(); + data.extend_from_slice(&encode_avro_long(1)); + data.extend_from_slice(&encode_avro_bytes(b"hello")); // key + data.extend_from_slice(&encode_avro_bytes(b"world")); // value + data.extend_from_slice(&encode_avro_long(0)); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + let array = decoder.flush(None).unwrap(); + let map_arr = array.as_any().downcast_ref::().unwrap(); + assert_eq!(map_arr.len(), 1); // one map + assert_eq!(map_arr.value_length(0), 1); + let entries = map_arr.value(0); + let struct_entries = entries.as_any().downcast_ref::().unwrap(); + assert_eq!(struct_entries.len(), 1); + let key_arr = struct_entries + .column_by_name("key") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let val_arr = struct_entries + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(key_arr.value(0), "hello"); + assert_eq!(val_arr.value(0), "world"); + } + + #[test] + fn test_map_decoding_empty() { + let value_type = avro_from_codec(Codec::Utf8); + let map_type = avro_from_codec(Codec::Map(Arc::new(value_type))); + let mut decoder = Decoder::try_new(&map_type).unwrap(); + let data = encode_avro_long(0); + decoder.decode(&mut AvroCursor::new(&data)).unwrap(); + let array = decoder.flush(None).unwrap(); + let map_arr = array.as_any().downcast_ref::().unwrap(); + assert_eq!(map_arr.len(), 1); + assert_eq!(map_arr.value_length(0), 0); + } + + #[test] + fn test_fixed_decoding() { + let avro_type = avro_from_codec(Codec::Fixed(3)); + let mut decoder = Decoder::try_new(&avro_type).expect("Failed to create decoder"); + + let data1 = [1u8, 2, 3]; + let mut cursor1 = AvroCursor::new(&data1); + decoder + .decode(&mut cursor1) + .expect("Failed to decode data1"); + assert_eq!(cursor1.position(), 3, "Cursor should advance by fixed size"); + let data2 = [4u8, 5, 6]; + let mut cursor2 = AvroCursor::new(&data2); + decoder + .decode(&mut cursor2) + .expect("Failed to decode data2"); + assert_eq!(cursor2.position(), 3, "Cursor should advance by fixed size"); + let array = decoder.flush(None).expect("Failed to flush decoder"); + assert_eq!(array.len(), 2, "Array should contain two items"); + let fixed_size_binary_array = array + .as_any() + .downcast_ref::() + .expect("Failed to downcast to FixedSizeBinaryArray"); + assert_eq!( + fixed_size_binary_array.value_length(), + 3, + "Fixed size of binary values should be 3" + ); + assert_eq!( + fixed_size_binary_array.value(0), + &[1, 2, 3], + "First item mismatch" + ); + assert_eq!( + fixed_size_binary_array.value(1), + &[4, 5, 6], + "Second item mismatch" + ); + } + + #[test] + fn test_fixed_decoding_empty() { + let avro_type = avro_from_codec(Codec::Fixed(5)); + let mut decoder = Decoder::try_new(&avro_type).expect("Failed to create decoder"); + + let array = decoder + .flush(None) + .expect("Failed to flush decoder for empty input"); + + assert_eq!(array.len(), 0, "Array should be empty"); + let fixed_size_binary_array = array + .as_any() + .downcast_ref::() + .expect("Failed to downcast to FixedSizeBinaryArray for empty array"); + + assert_eq!( + fixed_size_binary_array.value_length(), + 5, + "Fixed size of binary values should be 5 as per type" + ); + } + + #[test] + fn test_uuid_decoding() { + let avro_type = avro_from_codec(Codec::Uuid); + let mut decoder = Decoder::try_new(&avro_type).expect("Failed to create decoder"); + let uuid_str = "f81d4fae-7dec-11d0-a765-00a0c91e6bf6"; + let data = encode_avro_bytes(uuid_str.as_bytes()); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).expect("Failed to decode data"); + assert_eq!( + cursor.position(), + data.len(), + "Cursor should advance by varint size + data size" + ); + let array = decoder.flush(None).expect("Failed to flush decoder"); + let fixed_size_binary_array = array + .as_any() + .downcast_ref::() + .expect("Array should be a FixedSizeBinaryArray"); + assert_eq!(fixed_size_binary_array.len(), 1); + assert_eq!(fixed_size_binary_array.value_length(), 16); + let expected_bytes = [ + 0xf8, 0x1d, 0x4f, 0xae, 0x7d, 0xec, 0x11, 0xd0, 0xa7, 0x65, 0x00, 0xa0, 0xc9, 0x1e, + 0x6b, 0xf6, + ]; + assert_eq!(fixed_size_binary_array.value(0), &expected_bytes); + } + + #[test] + fn test_array_decoding() { + let item_dt = avro_from_codec(Codec::Int32); + let list_dt = avro_from_codec(Codec::List(Arc::new(item_dt))); + let mut decoder = Decoder::try_new(&list_dt).unwrap(); + let mut row1 = Vec::new(); + row1.extend_from_slice(&encode_avro_long(2)); + row1.extend_from_slice(&encode_avro_int(10)); + row1.extend_from_slice(&encode_avro_int(20)); + row1.extend_from_slice(&encode_avro_long(0)); + let row2 = encode_avro_long(0); + let mut cursor = AvroCursor::new(&row1); + decoder.decode(&mut cursor).unwrap(); + let mut cursor2 = AvroCursor::new(&row2); + decoder.decode(&mut cursor2).unwrap(); + let array = decoder.flush(None).unwrap(); + let list_arr = array.as_any().downcast_ref::().unwrap(); + assert_eq!(list_arr.len(), 2); + let offsets = list_arr.value_offsets(); + assert_eq!(offsets, &[0, 2, 2]); + let values = list_arr.values(); + let int_arr = values.as_primitive::(); + assert_eq!(int_arr.len(), 2); + assert_eq!(int_arr.value(0), 10); + assert_eq!(int_arr.value(1), 20); + } + + #[test] + fn test_array_decoding_with_negative_block_count() { + let item_dt = avro_from_codec(Codec::Int32); + let list_dt = avro_from_codec(Codec::List(Arc::new(item_dt))); + let mut decoder = Decoder::try_new(&list_dt).unwrap(); + let mut data = encode_avro_long(-3); + data.extend_from_slice(&encode_avro_long(12)); + data.extend_from_slice(&encode_avro_int(1)); + data.extend_from_slice(&encode_avro_int(2)); + data.extend_from_slice(&encode_avro_int(3)); + data.extend_from_slice(&encode_avro_long(0)); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + let array = decoder.flush(None).unwrap(); + let list_arr = array.as_any().downcast_ref::().unwrap(); + assert_eq!(list_arr.len(), 1); + assert_eq!(list_arr.value_length(0), 3); + let values = list_arr.values().as_primitive::(); + assert_eq!(values.len(), 3); + assert_eq!(values.value(0), 1); + assert_eq!(values.value(1), 2); + assert_eq!(values.value(2), 3); + } + + #[test] + fn test_nested_array_decoding() { + let inner_ty = avro_from_codec(Codec::List(Arc::new(avro_from_codec(Codec::Int32)))); + let nested_ty = avro_from_codec(Codec::List(Arc::new(inner_ty.clone()))); + let mut decoder = Decoder::try_new(&nested_ty).unwrap(); + let mut buf = Vec::new(); + buf.extend(encode_avro_long(1)); + buf.extend(encode_avro_long(2)); + buf.extend(encode_avro_int(5)); + buf.extend(encode_avro_int(6)); + buf.extend(encode_avro_long(0)); + buf.extend(encode_avro_long(0)); + let mut cursor = AvroCursor::new(&buf); + decoder.decode(&mut cursor).unwrap(); + let arr = decoder.flush(None).unwrap(); + let outer = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(outer.len(), 1); + assert_eq!(outer.value_length(0), 1); + let inner = outer.values().as_any().downcast_ref::().unwrap(); + assert_eq!(inner.len(), 1); + assert_eq!(inner.value_length(0), 2); + let values = inner + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values(), &[5, 6]); + } + + #[test] + fn test_array_decoding_empty_array() { + let value_type = avro_from_codec(Codec::Utf8); + let map_type = avro_from_codec(Codec::List(Arc::new(value_type))); + let mut decoder = Decoder::try_new(&map_type).unwrap(); + let data = encode_avro_long(0); + decoder.decode(&mut AvroCursor::new(&data)).unwrap(); + let array = decoder.flush(None).unwrap(); + let list_arr = array.as_any().downcast_ref::().unwrap(); + assert_eq!(list_arr.len(), 1); + assert_eq!(list_arr.value_length(0), 0); + } + + #[test] + fn test_decimal_decoding_fixed256() { + let dt = avro_from_codec(Codec::Decimal(5, Some(2), Some(32))); + let mut decoder = Decoder::try_new(&dt).unwrap(); + let row1 = [ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x30, 0x39, + ]; + let row2 = [ + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0xFF, 0xFF, 0x85, + ]; + let mut data = Vec::new(); + data.extend_from_slice(&row1); + data.extend_from_slice(&row2); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let arr = decoder.flush(None).unwrap(); + let dec = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(dec.len(), 2); + assert_eq!(dec.value_as_string(0), "123.45"); + assert_eq!(dec.value_as_string(1), "-1.23"); + } + + #[test] + fn test_decimal_decoding_fixed128() { + let dt = avro_from_codec(Codec::Decimal(5, Some(2), Some(16))); + let mut decoder = Decoder::try_new(&dt).unwrap(); + let row1 = [ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x30, 0x39, + ]; + let row2 = [ + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0x85, + ]; + let mut data = Vec::new(); + data.extend_from_slice(&row1); + data.extend_from_slice(&row2); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let arr = decoder.flush(None).unwrap(); + let dec = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(dec.len(), 2); + assert_eq!(dec.value_as_string(0), "123.45"); + assert_eq!(dec.value_as_string(1), "-1.23"); + } + + #[test] + fn test_decimal_decoding_bytes_with_nulls() { + let dt = avro_from_codec(Codec::Decimal(4, Some(1), None)); + let inner = Decoder::try_new(&dt).unwrap(); + let mut decoder = Decoder::Nullable( + Nullability::NullSecond, + NullBufferBuilder::new(DEFAULT_CAPACITY), + Box::new(inner), + ); + let mut data = Vec::new(); + data.extend_from_slice(&encode_avro_int(0)); + data.extend_from_slice(&encode_avro_bytes(&[0x04, 0xD2])); + data.extend_from_slice(&encode_avro_int(1)); + data.extend_from_slice(&encode_avro_int(0)); + data.extend_from_slice(&encode_avro_bytes(&[0xFB, 0x2E])); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); // row1 + decoder.decode(&mut cursor).unwrap(); // row2 + decoder.decode(&mut cursor).unwrap(); // row3 + let arr = decoder.flush(None).unwrap(); + let dec_arr = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(dec_arr.len(), 3); + assert!(dec_arr.is_valid(0)); + assert!(!dec_arr.is_valid(1)); + assert!(dec_arr.is_valid(2)); + assert_eq!(dec_arr.value_as_string(0), "123.4"); + assert_eq!(dec_arr.value_as_string(2), "-123.4"); + } + + #[test] + fn test_decimal_decoding_bytes_with_nulls_fixed_size() { + let dt = avro_from_codec(Codec::Decimal(6, Some(2), Some(16))); + let inner = Decoder::try_new(&dt).unwrap(); + let mut decoder = Decoder::Nullable( + Nullability::NullSecond, + NullBufferBuilder::new(DEFAULT_CAPACITY), + Box::new(inner), + ); + let row1 = [ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, + 0xE2, 0x40, + ]; + let row3 = [ + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFE, + 0x1D, 0xC0, + ]; + let mut data = Vec::new(); + data.extend_from_slice(&encode_avro_int(0)); + data.extend_from_slice(&row1); + data.extend_from_slice(&encode_avro_int(1)); + data.extend_from_slice(&encode_avro_int(0)); + data.extend_from_slice(&row3); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let arr = decoder.flush(None).unwrap(); + let dec_arr = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(dec_arr.len(), 3); + assert!(dec_arr.is_valid(0)); + assert!(!dec_arr.is_valid(1)); + assert!(dec_arr.is_valid(2)); + assert_eq!(dec_arr.value_as_string(0), "1234.56"); + assert_eq!(dec_arr.value_as_string(2), "-1234.56"); + } + + #[test] + fn test_enum_decoding() { + let symbols: Arc<[String]> = vec!["A", "B", "C"].into_iter().map(String::from).collect(); + let avro_type = avro_from_codec(Codec::Enum(symbols.clone())); + let mut decoder = Decoder::try_new(&avro_type).unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&encode_avro_int(2)); + data.extend_from_slice(&encode_avro_int(0)); + data.extend_from_slice(&encode_avro_int(1)); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let array = decoder.flush(None).unwrap(); + let dict_array = array + .as_any() + .downcast_ref::>() + .unwrap(); + + assert_eq!(dict_array.len(), 3); + let values = dict_array + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), "A"); + assert_eq!(values.value(1), "B"); + assert_eq!(values.value(2), "C"); + assert_eq!(dict_array.keys().values(), &[2, 0, 1]); + } + + #[test] + fn test_enum_decoding_with_nulls() { + let symbols: Arc<[String]> = vec!["X", "Y"].into_iter().map(String::from).collect(); + let enum_codec = Codec::Enum(symbols.clone()); + let avro_type = + AvroDataType::new(enum_codec, Default::default(), Some(Nullability::NullFirst)); + let mut decoder = Decoder::try_new(&avro_type).unwrap(); + let mut data = Vec::new(); + data.extend_from_slice(&encode_avro_long(1)); + data.extend_from_slice(&encode_avro_int(1)); + data.extend_from_slice(&encode_avro_long(0)); + data.extend_from_slice(&encode_avro_long(1)); + data.extend_from_slice(&encode_avro_int(0)); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let array = decoder.flush(None).unwrap(); + let dict_array = array + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(dict_array.len(), 3); + assert!(dict_array.is_valid(0)); + assert!(dict_array.is_null(1)); + assert!(dict_array.is_valid(2)); + let expected_keys = Int32Array::from(vec![Some(1), None, Some(0)]); + assert_eq!(dict_array.keys(), &expected_keys); + let values = dict_array + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.value(0), "X"); + assert_eq!(values.value(1), "Y"); + } + + #[test] + fn test_duration_decoding_with_nulls() { + let duration_codec = Codec::Interval; + let avro_type = AvroDataType::new( + duration_codec, + Default::default(), + Some(Nullability::NullFirst), + ); + let mut decoder = Decoder::try_new(&avro_type).unwrap(); + let mut data = Vec::new(); + // First value: 1 month, 2 days, 3 millis + data.extend_from_slice(&encode_avro_long(1)); // not null + let mut duration1 = Vec::new(); + duration1.extend_from_slice(&1u32.to_le_bytes()); + duration1.extend_from_slice(&2u32.to_le_bytes()); + duration1.extend_from_slice(&3u32.to_le_bytes()); + data.extend_from_slice(&duration1); + // Second value: null + data.extend_from_slice(&encode_avro_long(0)); // null + data.extend_from_slice(&encode_avro_long(1)); // not null + let mut duration2 = Vec::new(); + duration2.extend_from_slice(&4u32.to_le_bytes()); + duration2.extend_from_slice(&5u32.to_le_bytes()); + duration2.extend_from_slice(&6u32.to_le_bytes()); + data.extend_from_slice(&duration2); + let mut cursor = AvroCursor::new(&data); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + decoder.decode(&mut cursor).unwrap(); + let array = decoder.flush(None).unwrap(); + let interval_array = array + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(interval_array.len(), 3); + assert!(interval_array.is_valid(0)); + assert!(interval_array.is_null(1)); + assert!(interval_array.is_valid(2)); + let expected = IntervalMonthDayNanoArray::from(vec![ + Some(IntervalMonthDayNano { + months: 1, + days: 2, + nanoseconds: 3_000_000, + }), + None, + Some(IntervalMonthDayNano { + months: 4, + days: 5, + nanoseconds: 6_000_000, + }), + ]); + assert_eq!(interval_array, &expected); + } + + #[test] + fn test_duration_decoding_empty() { + let duration_codec = Codec::Interval; + let avro_type = AvroDataType::new(duration_codec, Default::default(), None); + let mut decoder = Decoder::try_new(&avro_type).unwrap(); + let array = decoder.flush(None).unwrap(); + assert_eq!(array.len(), 0); + } +} diff --git a/arrow-avro/src/schema.rs b/arrow-avro/src/schema.rs index a9d91e47948b..c3e4549c8c38 100644 --- a/arrow-avro/src/schema.rs +++ b/arrow-avro/src/schema.rs @@ -26,8 +26,13 @@ pub const SCHEMA_METADATA_KEY: &str = "avro.schema"; /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(untagged)] +/// A type name in an Avro schema +/// +/// This represents the different ways a type can be referenced in an Avro schema. pub enum TypeName<'a> { + /// A primitive type like null, boolean, int, etc. Primitive(PrimitiveType), + /// A reference to another named type Ref(&'a str), } @@ -37,13 +42,21 @@ pub enum TypeName<'a> { #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub enum PrimitiveType { + /// null: no value Null, + /// boolean: a binary value Boolean, + /// int: 32-bit signed integer Int, + /// long: 64-bit signed integer Long, + /// float: single precision (32-bit) IEEE 754 floating-point number Float, + /// double: double precision (64-bit) IEEE 754 floating-point number Double, + /// bytes: sequence of 8-bit unsigned bytes Bytes, + /// string: Unicode character sequence String, } @@ -78,22 +91,31 @@ impl Attributes<'_> { #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub struct Type<'a> { + /// The type of this Avro data structure #[serde(borrow)] pub r#type: TypeName<'a>, + /// Additional attributes associated with this type #[serde(flatten)] pub attributes: Attributes<'a>, } /// An Avro schema +/// +/// This represents the different shapes of Avro schemas as defined in the specification. +/// See for more details. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(untagged)] pub enum Schema<'a> { + /// A direct type name (primitive or reference) #[serde(borrow)] TypeName(TypeName<'a>), + /// A union of multiple schemas (e.g., ["null", "string"]) #[serde(borrow)] Union(Vec>), + /// A complex type such as record, array, map, etc. #[serde(borrow)] Complex(ComplexType<'a>), + /// A type with attributes #[serde(borrow)] Type(Type<'a>), } @@ -104,14 +126,19 @@ pub enum Schema<'a> { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "camelCase")] pub enum ComplexType<'a> { + /// Record type: a sequence of fields with names and types #[serde(borrow)] Record(Record<'a>), + /// Enum type: a set of named values #[serde(borrow)] Enum(Enum<'a>), + /// Array type: a sequence of values of the same type #[serde(borrow)] Array(Array<'a>), + /// Map type: a mapping from strings to values of the same type #[serde(borrow)] Map(Map<'a>), + /// Fixed type: a fixed-size byte array #[serde(borrow)] Fixed(Fixed<'a>), } @@ -121,16 +148,22 @@ pub enum ComplexType<'a> { /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Record<'a> { + /// Name of the record #[serde(borrow)] pub name: &'a str, + /// Optional namespace for the record, provides a way to organize names #[serde(borrow, default)] pub namespace: Option<&'a str>, + /// Optional documentation string for the record #[serde(borrow, default)] pub doc: Option<&'a str>, + /// Alternative names for this record #[serde(borrow, default)] pub aliases: Vec<&'a str>, + /// The fields contained in this record #[serde(borrow)] pub fields: Vec>, + /// Additional attributes for this record #[serde(flatten)] pub attributes: Attributes<'a>, } @@ -138,12 +171,16 @@ pub struct Record<'a> { /// A field within a [`Record`] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Field<'a> { + /// Name of the field within the record #[serde(borrow)] pub name: &'a str, + /// Optional documentation for this field #[serde(borrow, default)] pub doc: Option<&'a str>, + /// The field's type definition #[serde(borrow)] pub r#type: Schema<'a>, + /// Optional default value for this field #[serde(borrow, default)] pub default: Option<&'a str>, } @@ -153,18 +190,25 @@ pub struct Field<'a> { /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Enum<'a> { + /// Name of the enum #[serde(borrow)] pub name: &'a str, + /// Optional namespace for the enum, provides organizational structure #[serde(borrow, default)] pub namespace: Option<&'a str>, + /// Optional documentation string describing the enum #[serde(borrow, default)] pub doc: Option<&'a str>, + /// Alternative names for this enum #[serde(borrow, default)] pub aliases: Vec<&'a str>, + /// The symbols (values) that this enum can have #[serde(borrow)] pub symbols: Vec<&'a str>, + /// Optional default value for this enum #[serde(borrow, default)] pub default: Option<&'a str>, + /// Additional attributes for this enum #[serde(flatten)] pub attributes: Attributes<'a>, } @@ -174,8 +218,10 @@ pub struct Enum<'a> { /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Array<'a> { + /// The schema for items in this array #[serde(borrow)] pub items: Box>, + /// Additional attributes for this array #[serde(flatten)] pub attributes: Attributes<'a>, } @@ -185,8 +231,10 @@ pub struct Array<'a> { /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Map<'a> { + /// The schema for values in this map #[serde(borrow)] pub values: Box>, + /// Additional attributes for this map #[serde(flatten)] pub attributes: Attributes<'a>, } @@ -196,13 +244,18 @@ pub struct Map<'a> { /// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Fixed<'a> { + /// Name of the fixed type #[serde(borrow)] pub name: &'a str, + /// Optional namespace for the fixed type #[serde(borrow, default)] pub namespace: Option<&'a str>, + /// Alternative names for this fixed type #[serde(borrow, default)] pub aliases: Vec<&'a str>, + /// The number of bytes in this fixed type pub size: usize, + /// Additional attributes for this fixed type #[serde(flatten)] pub attributes: Attributes<'a>, } diff --git a/arrow-avro/test/data/duration_uuid.avro b/arrow-avro/test/data/duration_uuid.avro new file mode 100644 index 000000000000..09dd67b7807a Binary files /dev/null and b/arrow-avro/test/data/duration_uuid.avro differ diff --git a/arrow-avro/test/data/zero_byte.avro b/arrow-avro/test/data/zero_byte.avro new file mode 100644 index 000000000000..f7ffd29b6890 Binary files /dev/null and b/arrow-avro/test/data/zero_byte.avro differ diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index 68bfe8ddf732..21ed4212da65 100644 --- a/arrow-buffer/Cargo.toml +++ b/arrow-buffer/Cargo.toml @@ -30,9 +30,14 @@ rust-version = { workspace = true } [lib] name = "arrow_buffer" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + +[features] +pool = [] + [dependencies] bytes = { version = "1.4" } num = { version = "0.4", default-features = false, features = ["std"] } @@ -40,9 +45,7 @@ half = { version = "2.1", default-features = false } [dev-dependencies] criterion = { version = "0.5", default-features = false } -rand = { version = "0.8", default-features = false, features = ["std", "std_rng"] } - -[build-dependencies] +rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } [[bench]] name = "bit_mask" diff --git a/arrow-buffer/benches/bit_mask.rs b/arrow-buffer/benches/bit_mask.rs index 6907e336a418..545528724e5d 100644 --- a/arrow-buffer/benches/bit_mask.rs +++ b/arrow-buffer/benches/bit_mask.rs @@ -16,7 +16,8 @@ // under the License. use arrow_buffer::bit_mask::set_bits; -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("bit_mask"); @@ -38,11 +39,11 @@ fn criterion_benchmark(c: &mut Criterion) { |b, &x| { b.iter(|| { set_bits( - black_box(&mut [0u8; 9]), - black_box(&[x.3; 9]), - black_box(x.0), - black_box(x.1), - black_box(x.2), + hint::black_box(&mut [0u8; 9]), + hint::black_box(&[x.3; 9]), + hint::black_box(x.0), + hint::black_box(x.1), + hint::black_box(x.2), ) }); }, diff --git a/arrow-buffer/benches/i256.rs b/arrow-buffer/benches/i256.rs index ebb45e793bd0..7dec226bbc08 100644 --- a/arrow-buffer/benches/i256.rs +++ b/arrow-buffer/benches/i256.rs @@ -19,7 +19,7 @@ use arrow_buffer::i256; use criterion::*; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -use std::str::FromStr; +use std::{hint, str::FromStr}; const SIZE: usize = 1024; @@ -37,7 +37,7 @@ fn criterion_benchmark(c: &mut Criterion) { ]; for number in numbers { - let t = black_box(number.to_string()); + let t = hint::black_box(number.to_string()); c.bench_function(&format!("i256_parse({t})"), |b| { b.iter(|| i256::from_str(&t).unwrap()); }); @@ -47,8 +47,8 @@ fn criterion_benchmark(c: &mut Criterion) { let numerators: Vec<_> = (0..SIZE) .map(|_| { - let high = rng.gen_range(1000..i128::MAX); - let low = rng.gen(); + let high = rng.random_range(1000..i128::MAX); + let low = rng.random(); i256::from_parts(low, high) }) .collect(); @@ -56,7 +56,7 @@ fn criterion_benchmark(c: &mut Criterion) { let divisors: Vec<_> = numerators .iter() .map(|n| { - let quotient = rng.gen_range(1..100_i32); + let quotient = rng.random_range(1..100_i32); n.wrapping_div(i256::from(quotient)) }) .collect(); @@ -64,19 +64,19 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function("i256_div_rem small quotient", |b| { b.iter(|| { for (n, d) in numerators.iter().zip(&divisors) { - black_box(n.wrapping_div(*d)); + hint::black_box(n.wrapping_div(*d)); } }); }); let divisors: Vec<_> = (0..SIZE) - .map(|_| i256::from(rng.gen_range(1..100_i32))) + .map(|_| i256::from(rng.random_range(1..100_i32))) .collect(); c.bench_function("i256_div_rem small divisor", |b| { b.iter(|| { for (n, d) in numerators.iter().zip(&divisors) { - black_box(n.wrapping_div(*d)); + hint::black_box(n.wrapping_div(*d)); } }); }); diff --git a/arrow-buffer/benches/offset.rs b/arrow-buffer/benches/offset.rs index 1aea5024fbd1..0d6128b69288 100644 --- a/arrow-buffer/benches/offset.rs +++ b/arrow-buffer/benches/offset.rs @@ -19,12 +19,13 @@ use arrow_buffer::{OffsetBuffer, OffsetBufferBuilder}; use criterion::*; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use std::hint; const SIZE: usize = 1024; fn criterion_benchmark(c: &mut Criterion) { let mut rng = StdRng::seed_from_u64(42); - let lengths: Vec = black_box((0..SIZE).map(|_| rng.gen_range(0..40)).collect()); + let lengths: Vec = hint::black_box((0..SIZE).map(|_| rng.random_range(0..40)).collect()); c.bench_function("OffsetBuffer::from_lengths", |b| { b.iter(|| OffsetBuffer::::from_lengths(lengths.iter().copied())); @@ -41,7 +42,7 @@ fn criterion_benchmark(c: &mut Criterion) { let offsets = OffsetBuffer::::from_lengths(lengths.iter().copied()).into_inner(); c.bench_function("OffsetBuffer::new", |b| { - b.iter(|| OffsetBuffer::new(black_box(offsets.clone()))); + b.iter(|| OffsetBuffer::new(hint::black_box(offsets.clone()))); }); } diff --git a/arrow-buffer/src/bigint/div.rs b/arrow-buffer/src/bigint/div.rs index 8a75dad0ffd8..f094f662ccf8 100644 --- a/arrow-buffer/src/bigint/div.rs +++ b/arrow-buffer/src/bigint/div.rs @@ -39,8 +39,8 @@ pub fn div_rem(numerator: &[u64; N], divisor: &[u64; N]) -> ([u6 return div_rem_small(numerator, divisor[0]); } - let numerator_words = (numerator_bits + 63) / 64; - let divisor_words = (divisor_bits + 63) / 64; + let numerator_words = numerator_bits.div_ceil(64); + let divisor_words = divisor_bits.div_ceil(64); let n = divisor_words; let m = numerator_words - divisor_words; @@ -258,7 +258,7 @@ fn full_shl(v: &[u64; N], shift: u32) -> ArrayPlusOne { let mut out = [0u64; N]; out[0] = v[0] << shift; for i in 1..N { - out[i] = v[i - 1] >> (64 - shift) | v[i] << shift + out[i] = (v[i - 1] >> (64 - shift)) | (v[i] << shift) } let carry = v[N - 1] >> (64 - shift); ArrayPlusOne(out, carry) @@ -272,7 +272,7 @@ fn full_shr(a: &ArrayPlusOne, shift: u32) -> [u64; N] { } let mut out = [0; N]; for i in 0..N - 1 { - out[i] = a[i] >> shift | a[i + 1] << (64 - shift) + out[i] = (a[i] >> shift) | (a[i + 1] << (64 - shift)) } out[N - 1] = a[N - 1] >> shift; out diff --git a/arrow-buffer/src/bigint/mod.rs b/arrow-buffer/src/bigint/mod.rs index f5fab75dc5ef..9868ab55cc11 100644 --- a/arrow-buffer/src/bigint/mod.rs +++ b/arrow-buffer/src/bigint/mod.rs @@ -475,8 +475,8 @@ impl i256 { /// Interpret 4 `u64` digits, least significant first, as a [`i256`] fn from_digits(digits: [u64; 4]) -> Self { Self::from_parts( - digits[0] as u128 | (digits[1] as u128) << 64, - digits[2] as i128 | (digits[3] as i128) << 64, + digits[0] as u128 | ((digits[1] as u128) << 64), + digits[2] as i128 | ((digits[3] as i128) << 64), ) } @@ -746,7 +746,7 @@ impl Shl for i256 { self } else if rhs < 128 { Self { - high: self.high << rhs | (self.low >> (128 - rhs)) as i128, + high: (self.high << rhs) | (self.low >> (128 - rhs)) as i128, low: self.low << rhs, } } else { @@ -768,7 +768,7 @@ impl Shr for i256 { } else if rhs < 128 { Self { high: self.high >> rhs, - low: self.low >> rhs | ((self.high as u128) << (128 - rhs)), + low: (self.low >> rhs) | ((self.high as u128) << (128 - rhs)), } } else { Self { @@ -840,7 +840,7 @@ impl ToPrimitive for i256 { mod tests { use super::*; use num::Signed; - use rand::{thread_rng, Rng}; + use rand::{rng, Rng}; #[test] fn test_signed_cmp() { @@ -1091,16 +1091,16 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_i256_fuzz() { - let mut rng = thread_rng(); + let mut rng = rng(); for _ in 0..1000 { let mut l = [0_u8; 32]; - let len = rng.gen_range(0..32); - l.iter_mut().take(len).for_each(|x| *x = rng.gen()); + let len = rng.random_range(0..32); + l.iter_mut().take(len).for_each(|x| *x = rng.random()); let mut r = [0_u8; 32]; - let len = rng.gen_range(0..32); - r.iter_mut().take(len).for_each(|x| *x = rng.gen()); + let len = rng.random_range(0..32); + r.iter_mut().take(len).for_each(|x| *x = rng.random()); test_ops(i256::from_le_bytes(l), i256::from_le_bytes(r)) } diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index aaa86832f692..c8e5144c14cb 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -25,6 +25,14 @@ use crate::{ use std::ops::{BitAnd, BitOr, BitXor, Not}; /// A slice-able [`Buffer`] containing bit-packed booleans +/// +/// `BooleanBuffer`s can be creating using [`BooleanBufferBuilder`] +/// +/// # See Also +/// +/// * [`NullBuffer`] for representing null values in Arrow arrays +/// +/// [`NullBuffer`]: crate::NullBuffer #[derive(Debug, Clone, Eq)] pub struct BooleanBuffer { buffer: Buffer, diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index d0c8ffa39783..2b55bf6604e6 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -25,11 +25,49 @@ use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; use crate::BufferBuilder; use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; +#[cfg(feature = "pool")] +use crate::pool::MemoryPool; + use super::ops::bitwise_unary_op_helper; use super::{MutableBuffer, ScalarBuffer}; -/// Buffer represents a contiguous memory region that can be shared with other buffers and across -/// thread boundaries. +/// A contiguous memory region that can be shared with other buffers and across +/// thread boundaries that stores Arrow data. +/// +/// `Buffer`s can be sliced and cloned without copying the underlying data and can +/// be created from memory allocated by non-Rust sources such as C/C++. +/// +/// # Example: Create a `Buffer` from a `Vec` (without copying) +/// ``` +/// # use arrow_buffer::Buffer; +/// let vec: Vec = vec![1, 2, 3]; +/// let buffer = Buffer::from(vec); +/// ``` +/// +/// # Example: Convert a `Buffer` to a `Vec` (without copying) +/// +/// Use [`Self::into_vec`] to convert a `Buffer` back into a `Vec` if there are +/// no other references and the types are aligned correctly. +/// ``` +/// # use arrow_buffer::Buffer; +/// # let vec: Vec = vec![1, 2, 3]; +/// # let buffer = Buffer::from(vec); +/// // convert the buffer back into a Vec of u32 +/// // note this will fail if the buffer is shared or not aligned correctly +/// let vec: Vec = buffer.into_vec().unwrap(); +/// ``` +/// +/// # Example: Create a `Buffer` from a [`bytes::Bytes`] (without copying) +/// +/// [`bytes::Bytes`] is a common type in the Rust ecosystem for shared memory +/// regions. You can create a buffer from a `Bytes` instance using the `From` +/// implementation, also without copying. +/// +/// ``` +/// # use arrow_buffer::Buffer; +/// let bytes = bytes::Bytes::from("hello"); +/// let buffer = Buffer::from(bytes); +///``` #[derive(Clone, Debug)] pub struct Buffer { /// the internal byte buffer. @@ -47,6 +85,13 @@ pub struct Buffer { length: usize, } +impl Default for Buffer { + #[inline] + fn default() -> Self { + MutableBuffer::default().into() + } +} + impl PartialEq for Buffer { fn eq(&self, other: &Self) -> bool { self.as_slice().eq(other.as_slice()) @@ -59,16 +104,15 @@ unsafe impl Send for Buffer where Bytes: Send {} unsafe impl Sync for Buffer where Bytes: Sync {} impl Buffer { - /// Auxiliary method to create a new Buffer - #[inline] + /// Create a new Buffer from a (internal) `Bytes` + /// + /// NOTE despite the same name, `Bytes` is an internal struct in arrow-rs + /// and is different than [`bytes::Bytes`]. + /// + /// See examples on [`Buffer`] for ways to create a buffer from a [`bytes::Bytes`]. + #[deprecated(since = "54.1.0", note = "Use Buffer::from instead")] pub fn from_bytes(bytes: Bytes) -> Self { - let length = bytes.len(); - let ptr = bytes.as_ptr(); - Buffer { - data: Arc::new(bytes), - ptr, - length, - } + Self::from(bytes) } /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` @@ -84,6 +128,15 @@ impl Buffer { self.data.ptr() } + /// Returns the number of strong references to the buffer. + /// + /// This method is safe but if the buffer is shared across multiple threads + /// the underlying value could change between calling this method and using + /// the result. + pub fn strong_count(&self) -> usize { + Arc::strong_count(&self.data) + } + /// Create a [`Buffer`] from the provided [`Vec`] without copying #[inline] pub fn from_vec(vec: Vec) -> Self { @@ -99,8 +152,11 @@ impl Buffer { buffer.into() } - /// Creates a buffer from an existing memory region. Ownership of the memory is tracked via reference counting - /// and the memory will be freed using the `drop` method of [crate::alloc::Allocation] when the reference count reaches zero. + /// Creates a buffer from an existing memory region. + /// + /// Ownership of the memory is tracked via reference counting + /// and the memory will be freed using the `drop` method of + /// [crate::alloc::Allocation] when the reference count reaches zero. /// /// # Arguments /// @@ -147,7 +203,7 @@ impl Buffer { self.data.capacity() } - /// Tried to shrink the capacity of the buffer as much as possible, freeing unused memory. + /// Tries to shrink the capacity of the buffer as much as possible, freeing unused memory. /// /// If the buffer is shared, this is a no-op. /// @@ -182,7 +238,7 @@ impl Buffer { } } - /// Returns whether the buffer is empty. + /// Returns true if the buffer is empty. #[inline] pub fn is_empty(&self) -> bool { self.length == 0 @@ -198,7 +254,9 @@ impl Buffer { } /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`. - /// Doing so allows the same memory region to be shared between buffers. + /// + /// This function is `O(1)` and does not copy any data, allowing the + /// same memory region to be shared between buffers. /// /// # Panics /// @@ -232,7 +290,10 @@ impl Buffer { /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`, /// with `length` bytes. - /// Doing so allows the same memory region to be shared between buffers. + /// + /// This function is `O(1)` and does not copy any data, allowing the same + /// memory region to be shared between buffers. + /// /// # Panics /// Panics iff `(offset + length)` is larger than the existing length. pub fn slice_with_length(&self, offset: usize, length: usize) -> Self { @@ -320,10 +381,16 @@ impl Buffer { }) } - /// Returns `Vec` for mutating the buffer + /// Converts self into a `Vec`, if possible. + /// + /// This can be used to reuse / mutate the underlying data. /// - /// Returns `Err(self)` if this buffer does not have the same [`Layout`] as - /// the destination Vec or contains a non-zero offset + /// # Errors + /// + /// Returns `Err(self)` if + /// 1. this buffer does not have the same [`Layout`] as the destination Vec + /// 2. contains a non-zero offset + /// 3. The buffer is shared pub fn into_vec(self) -> Result, Self> { let layout = match self.data.deallocation() { Deallocation::Standard(l) => l, @@ -366,6 +433,17 @@ impl Buffer { pub fn ptr_eq(&self, other: &Self) -> bool { self.ptr == other.ptr && self.length == other.length } + + /// Register this [`Buffer`] with the provided [`MemoryPool`] + /// + /// This claims the memory used by this buffer in the pool, allowing for + /// accurate accounting of memory usage. Any prior reservation will be + /// released so this works well when the buffer is being shared among + /// multiple arrays. + #[cfg(feature = "pool")] + pub fn claim(&self, pool: &dyn MemoryPool) { + self.data.claim(pool) + } } /// Note that here we deliberately do not implement @@ -406,7 +484,29 @@ impl From> for Buffer { } } -/// Creating a `Buffer` instance by storing the boolean values into the buffer +/// Convert from internal `Bytes` (not [`bytes::Bytes`]) to `Buffer` +impl From for Buffer { + #[inline] + fn from(bytes: Bytes) -> Self { + let length = bytes.len(); + let ptr = bytes.as_ptr(); + Self { + data: Arc::new(bytes), + ptr, + length, + } + } +} + +/// Convert from [`bytes::Bytes`], not internal `Bytes` to `Buffer` +impl From for Buffer { + fn from(bytes: bytes::Bytes) -> Self { + let bytes: Bytes = bytes.into(); + Self::from(bytes) + } +} + +/// Create a `Buffer` instance by storing the boolean values into the buffer impl FromIterator for Buffer { fn from_iter(iter: I) -> Self where @@ -439,7 +539,9 @@ impl From> for Buffer { impl Buffer { /// Creates a [`Buffer`] from an [`Iterator`] with a trusted (upper) length. + /// /// Prefer this to `collect` whenever possible, as it is ~60% faster. + /// /// # Example /// ``` /// # use arrow_buffer::buffer::Buffer; @@ -895,13 +997,13 @@ mod tests { #[should_panic(expected = "capacity overflow")] fn test_from_iter_overflow() { let iter_len = usize::MAX / std::mem::size_of::() + 1; - let _ = Buffer::from_iter(std::iter::repeat(0_u64).take(iter_len)); + let _ = Buffer::from_iter(std::iter::repeat_n(0_u64, iter_len)); } #[test] fn bit_slice_length_preserved() { // Create a boring buffer - let buf = Buffer::from_iter(std::iter::repeat(true).take(64)); + let buf = Buffer::from_iter(std::iter::repeat_n(true, 64)); let assert_preserved = |offset: usize, len: usize| { let new_buf = buf.bit_slice(offset, len); @@ -930,4 +1032,31 @@ mod tests { } } } + + #[test] + fn test_strong_count() { + let buffer = Buffer::from_iter(std::iter::repeat_n(0_u8, 100)); + assert_eq!(buffer.strong_count(), 1); + + let buffer2 = buffer.clone(); + assert_eq!(buffer.strong_count(), 2); + + let buffer3 = buffer2.clone(); + assert_eq!(buffer.strong_count(), 3); + + drop(buffer); + assert_eq!(buffer2.strong_count(), 2); + assert_eq!(buffer3.strong_count(), 2); + + // Strong count does not increase on move + let capture = move || { + assert_eq!(buffer3.strong_count(), 2); + }; + + capture(); + assert_eq!(buffer2.strong_count(), 2); + + drop(capture); + assert_eq!(buffer2.strong_count(), 1); + } } diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index c4315a1d64cd..63fdbf598bdb 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -26,6 +26,11 @@ use crate::{ util::bit_util, }; +#[cfg(feature = "pool")] +use crate::pool::{MemoryPool, MemoryReservation}; +#[cfg(feature = "pool")] +use std::sync::Mutex; + use super::Buffer; /// A [`MutableBuffer`] is Arrow's interface to build a [`Buffer`] out of items or slices of items. @@ -57,6 +62,10 @@ pub struct MutableBuffer { // invariant: len <= capacity len: usize, layout: Layout, + + /// Memory reservation for tracking memory usage + #[cfg(feature = "pool")] + reservation: Mutex>>, } impl MutableBuffer { @@ -91,6 +100,8 @@ impl MutableBuffer { data, len: 0, layout, + #[cfg(feature = "pool")] + reservation: std::sync::Mutex::new(None), } } @@ -115,7 +126,13 @@ impl MutableBuffer { NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout)) } }; - Self { data, len, layout } + Self { + data, + len, + layout, + #[cfg(feature = "pool")] + reservation: std::sync::Mutex::new(None), + } } /// Allocates a new [MutableBuffer] from given `Bytes`. @@ -127,9 +144,17 @@ impl MutableBuffer { let len = bytes.len(); let data = bytes.ptr(); + #[cfg(feature = "pool")] + let reservation = bytes.reservation.lock().unwrap().take(); mem::forget(bytes); - Ok(Self { data, len, layout }) + Ok(Self { + data, + len, + layout, + #[cfg(feature = "pool")] + reservation: Mutex::new(reservation), + }) } /// creates a new [MutableBuffer] with capacity and length capable of holding `len` bits. @@ -217,6 +242,12 @@ impl MutableBuffer { }; self.data = NonNull::new(data).unwrap_or_else(|| handle_alloc_error(new_layout)); self.layout = new_layout; + #[cfg(feature = "pool")] + { + if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { + reservation.resize(self.layout.size()); + } + } } /// Truncates this buffer to `len` bytes @@ -228,6 +259,12 @@ impl MutableBuffer { return; } self.len = len; + #[cfg(feature = "pool")] + { + if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { + reservation.resize(self.len); + } + } } /// Resizes the buffer, either truncating its contents (with no change in capacity), or @@ -251,6 +288,12 @@ impl MutableBuffer { } // this truncates the buffer when new_len < self.len self.len = new_len; + #[cfg(feature = "pool")] + { + if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { + reservation.resize(self.len); + } + } } /// Shrinks the capacity of the buffer as much as possible. @@ -288,7 +331,8 @@ impl MutableBuffer { self.len } - /// Returns the total capacity in this buffer. + /// Returns the total capacity in this buffer, in bytes. + /// /// The invariant `buffer.len() <= buffer.capacity()` is always upheld. #[inline] pub const fn capacity(&self) -> usize { @@ -327,8 +371,13 @@ impl MutableBuffer { #[inline] pub(super) fn into_buffer(self) -> Buffer { let bytes = unsafe { Bytes::new(self.data, self.len, Deallocation::Standard(self.layout)) }; + #[cfg(feature = "pool")] + { + let reservation = self.reservation.lock().unwrap().take(); + *bytes.reservation.lock().unwrap() = reservation; + } std::mem::forget(self); - Buffer::from_bytes(bytes) + Buffer::from(bytes) } /// View this buffer as a mutable slice of a specific type. @@ -465,6 +514,17 @@ impl MutableBuffer { buffer.truncate(bit_util::ceil(len, 8)); buffer } + + /// Register this [`MutableBuffer`] with the provided [`MemoryPool`] + /// + /// This claims the memory used by this buffer in the pool, allowing for + /// accurate accounting of memory usage. Any prior reservation will be + /// released so this works well when the buffer is being shared among + /// multiple arrays. + #[cfg(feature = "pool")] + pub fn claim(&self, pool: &dyn MemoryPool) { + *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity())); + } } /// Creates a non-null pointer with alignment of [`ALIGNMENT`] @@ -505,7 +565,13 @@ impl From> for MutableBuffer { // This is based on `RawVec::current_memory` let layout = unsafe { Layout::array::(value.capacity()).unwrap_unchecked() }; mem::forget(value); - Self { data, len, layout } + Self { + data, + len, + layout, + #[cfg(feature = "pool")] + reservation: std::sync::Mutex::new(None), + } } } @@ -1012,4 +1078,108 @@ mod tests { let max_capacity = isize::MAX as usize - (isize::MAX as usize % ALIGNMENT); let _ = MutableBuffer::with_capacity(max_capacity + 1); } + + #[cfg(feature = "pool")] + mod pool_tests { + use super::*; + use crate::pool::{MemoryPool, TrackingMemoryPool}; + + #[test] + fn test_reallocate_with_pool() { + let pool = TrackingMemoryPool::default(); + let mut buffer = MutableBuffer::with_capacity(100); + buffer.claim(&pool); + + // Initial capacity should be 128 (multiple of 64) + assert_eq!(buffer.capacity(), 128); + assert_eq!(pool.used(), 128); + + // Reallocate to a larger size + buffer.reallocate(200); + + // The capacity is exactly the requested size, not rounded up + assert_eq!(buffer.capacity(), 200); + assert_eq!(pool.used(), 200); + + // Reallocate to a smaller size + buffer.reallocate(50); + + // The capacity is exactly the requested size, not rounded up + assert_eq!(buffer.capacity(), 50); + assert_eq!(pool.used(), 50); + } + + #[test] + fn test_truncate_with_pool() { + let pool = TrackingMemoryPool::default(); + let mut buffer = MutableBuffer::with_capacity(100); + + // Fill buffer with some data + buffer.resize(80, 1); + assert_eq!(buffer.len(), 80); + + buffer.claim(&pool); + assert_eq!(pool.used(), 128); + + // Truncate buffer + buffer.truncate(40); + assert_eq!(buffer.len(), 40); + assert_eq!(pool.used(), 40); + + // Truncate to zero + buffer.truncate(0); + assert_eq!(buffer.len(), 0); + assert_eq!(pool.used(), 0); + } + + #[test] + fn test_resize_with_pool() { + let pool = TrackingMemoryPool::default(); + let mut buffer = MutableBuffer::with_capacity(100); + buffer.claim(&pool); + + // Initial state + assert_eq!(buffer.len(), 0); + assert_eq!(pool.used(), 128); + + // Resize to increase length + buffer.resize(50, 1); + assert_eq!(buffer.len(), 50); + assert_eq!(pool.used(), 50); + + // Resize to increase length beyond capacity + buffer.resize(150, 1); + assert_eq!(buffer.len(), 150); + assert_eq!(buffer.capacity(), 256); + assert_eq!(pool.used(), 150); + + // Resize to decrease length + buffer.resize(30, 1); + assert_eq!(buffer.len(), 30); + assert_eq!(pool.used(), 30); + } + + #[test] + fn test_buffer_lifecycle_with_pool() { + let pool = TrackingMemoryPool::default(); + + // Create a buffer with memory reservation + let mut mutable = MutableBuffer::with_capacity(100); + mutable.resize(80, 1); + mutable.claim(&pool); + + // Memory reservation is based on capacity when using claim() + assert_eq!(pool.used(), 128); + + // Convert to immutable Buffer + let buffer = mutable.into_buffer(); + + // Memory reservation should be preserved + assert_eq!(pool.used(), 128); + + // Drop the buffer and the reservation should be released + drop(buffer); + assert_eq!(pool.used(), 0); + } + } } diff --git a/arrow-buffer/src/buffer/null.rs b/arrow-buffer/src/buffer/null.rs index ec12b885eb5a..e5e3a610ead2 100644 --- a/arrow-buffer/src/buffer/null.rs +++ b/arrow-buffer/src/buffer/null.rs @@ -19,13 +19,16 @@ use crate::bit_iterator::{BitIndexIterator, BitIterator, BitSliceIterator}; use crate::buffer::BooleanBuffer; use crate::{Buffer, MutableBuffer}; -/// A [`BooleanBuffer`] used to encode validity for arrow arrays +/// A [`BooleanBuffer`] used to encode validity for Arrow arrays /// -/// As per the [Arrow specification], array validity is encoded in a packed bitmask with a +/// In the [Arrow specification], array validity is encoded in a packed bitmask with a /// `true` value indicating the corresponding slot is not null, and `false` indicating /// that it is null. /// +/// `NullBuffer`s can be creating using [`NullBufferBuilder`] +/// /// [Arrow specification]: https://arrow.apache.org/docs/format/Columnar.html#validity-bitmaps +/// [`NullBufferBuilder`]: crate::NullBufferBuilder #[derive(Debug, Clone, Eq, PartialEq)] pub struct NullBuffer { buffer: BooleanBuffer, @@ -49,7 +52,8 @@ impl NullBuffer { /// Create a new [`NullBuffer`] of length `len` where all values are valid /// - /// Note: it is more efficient to not set the null buffer if it is known to be all valid + /// Note: it is more efficient to not set the null buffer if it is known to + /// be all valid (aka all values are not null) pub fn new_valid(len: usize) -> Self { Self { buffer: BooleanBuffer::new_set(len), @@ -112,7 +116,7 @@ impl NullBuffer { } } - /// Returns the length of this [`NullBuffer`] + /// Returns the length of this [`NullBuffer`] in bits #[inline] pub fn len(&self) -> usize { self.buffer.len() diff --git a/arrow-buffer/src/buffer/offset.rs b/arrow-buffer/src/buffer/offset.rs index a6be2b67af84..fe3a57a38248 100644 --- a/arrow-buffer/src/buffer/offset.rs +++ b/arrow-buffer/src/buffer/offset.rs @@ -55,7 +55,7 @@ use std::ops::Deref; /// (offsets[i], /// offsets[i+1]) /// ``` -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct OffsetBuffer(ScalarBuffer); impl OffsetBuffer { @@ -133,6 +133,38 @@ impl OffsetBuffer { Self(out.into()) } + /// Get an Iterator over the lengths of this [`OffsetBuffer`] + /// + /// ``` + /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + /// let offsets = OffsetBuffer::<_>::new(ScalarBuffer::::from(vec![0, 1, 4, 9])); + /// assert_eq!(offsets.lengths().collect::>(), vec![1, 3, 5]); + /// ``` + /// + /// Empty [`OffsetBuffer`] will return an empty iterator + /// ``` + /// # use arrow_buffer::OffsetBuffer; + /// let offsets = OffsetBuffer::::new_empty(); + /// assert_eq!(offsets.lengths().count(), 0); + /// ``` + /// + /// This can be used to merge multiple [`OffsetBuffer`]s to one + /// ``` + /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + /// + /// let buffer1 = OffsetBuffer::::from_lengths([2, 6, 3, 7, 2]); + /// let buffer2 = OffsetBuffer::::from_lengths([1, 3, 5, 7, 9]); + /// + /// let merged = OffsetBuffer::::from_lengths( + /// vec![buffer1, buffer2].iter().flat_map(|x| x.lengths()) + /// ); + /// + /// assert_eq!(merged.lengths().collect::>(), &[2, 6, 3, 7, 2, 1, 3, 5, 7, 9]); + /// ``` + pub fn lengths(&self) -> impl ExactSizeIterator + '_ { + self.0.windows(2).map(|x| x[1].as_usize() - x[0].as_usize()) + } + /// Free up unused memory. pub fn shrink_to_fit(&mut self) { self.0.shrink_to_fit(); @@ -184,6 +216,12 @@ impl From> for OffsetBuffer { } } +impl Default for OffsetBuffer { + fn default() -> Self { + Self::new_empty() + } +} + #[cfg(test)] mod tests { use super::*; @@ -244,4 +282,45 @@ mod tests { fn from_lengths_usize_overflow() { OffsetBuffer::::from_lengths([usize::MAX, 1]); } + + #[test] + fn get_lengths() { + let offsets = OffsetBuffer::::new(ScalarBuffer::::from(vec![0, 1, 4, 9])); + assert_eq!(offsets.lengths().collect::>(), vec![1, 3, 5]); + } + + #[test] + fn get_lengths_should_be_with_fixed_size() { + let offsets = OffsetBuffer::::new(ScalarBuffer::::from(vec![0, 1, 4, 9])); + let iter = offsets.lengths(); + assert_eq!(iter.size_hint(), (3, Some(3))); + assert_eq!(iter.len(), 3); + } + + #[test] + fn get_lengths_from_empty_offset_buffer_should_be_empty_iterator() { + let offsets = OffsetBuffer::::new_empty(); + assert_eq!(offsets.lengths().collect::>(), vec![]); + } + + #[test] + fn impl_eq() { + fn are_equal(a: &T, b: &T) -> bool { + a.eq(b) + } + + assert!( + are_equal( + &OffsetBuffer::new(ScalarBuffer::::from(vec![0, 1, 4, 9])), + &OffsetBuffer::new(ScalarBuffer::::from(vec![0, 1, 4, 9])) + ), + "OffsetBuffer should implement Eq." + ); + } + + #[test] + fn impl_default() { + let default = OffsetBuffer::::default(); + assert_eq!(default.as_ref(), &[0]); + } } diff --git a/arrow-buffer/src/buffer/scalar.rs b/arrow-buffer/src/buffer/scalar.rs index ab6c87168e5c..6c66060fb95f 100644 --- a/arrow-buffer/src/buffer/scalar.rs +++ b/arrow-buffer/src/buffer/scalar.rs @@ -41,7 +41,7 @@ use std::ops::Deref; /// let sliced = buffer.slice(1, 2); /// assert_eq!(&sliced, &[2, 3]); /// ``` -#[derive(Clone)] +#[derive(Clone, Default)] pub struct ScalarBuffer { /// Underlying data buffer buffer: Buffer, @@ -182,6 +182,7 @@ impl From> for ScalarBuffer { } impl FromIterator for ScalarBuffer { + #[inline] fn from_iter>(iter: I) -> Self { iter.into_iter().collect::>().into() } @@ -220,6 +221,9 @@ impl PartialEq> for Vec { } } +/// If T implements Eq, then so does ScalarBuffer. +impl Eq for ScalarBuffer {} + #[cfg(test)] mod tests { use std::{ptr::NonNull, sync::Arc}; @@ -341,4 +345,19 @@ mod tests { assert_eq!(vec, input.as_slice()); assert_ne!(vec.as_ptr(), input.as_ptr()); } + + #[test] + fn scalar_buffer_impl_eq() { + fn are_equal(a: &T, b: &T) -> bool { + a.eq(b) + } + + assert!( + are_equal( + &ScalarBuffer::::from(vec![23]), + &ScalarBuffer::::from(vec![23]) + ), + "ScalarBuffer should implement Eq if the inner type does" + ); + } } diff --git a/arrow-buffer/src/builder/boolean.rs b/arrow-buffer/src/builder/boolean.rs index ca178ae5ce4e..bdcc3a55dbf2 100644 --- a/arrow-buffer/src/builder/boolean.rs +++ b/arrow-buffer/src/builder/boolean.rs @@ -19,6 +19,12 @@ use crate::{bit_mask, bit_util, BooleanBuffer, Buffer, MutableBuffer}; use std::ops::Range; /// Builder for [`BooleanBuffer`] +/// +/// # See Also +/// +/// * [`NullBuffer`] for building [`BooleanBuffer`]s for representing nulls +/// +/// [`NullBuffer`]: crate::NullBuffer #[derive(Debug)] pub struct BooleanBufferBuilder { buffer: MutableBuffer, @@ -26,7 +32,11 @@ pub struct BooleanBufferBuilder { } impl BooleanBufferBuilder { - /// Creates a new `BooleanBufferBuilder` + /// Creates a new `BooleanBufferBuilder` with sufficient space for + /// `capacity` bits (not bytes). + /// + /// The capacity is rounded up to the nearest multiple of 8 for the + /// allocation. #[inline] pub fn new(capacity: usize) -> Self { let byte_capacity = bit_util::ceil(capacity, 8); @@ -73,7 +83,24 @@ impl BooleanBufferBuilder { self.len == 0 } - /// Returns the capacity of the buffer + /// Returns the capacity of the buffer, in bits (not bytes) + /// + /// Note this + /// + /// # Example + /// ``` + /// # use arrow_buffer::builder::BooleanBufferBuilder; + /// // empty requires 0 bytes + /// let b = BooleanBufferBuilder::new(0); + /// assert_eq!(0, b.capacity()); + /// // Creating space for 1 bit results in 64 bytes (space for 512 bits) + /// // (64 is the minimum allocation size for 64 bit architectures) + /// let mut b = BooleanBufferBuilder::new(1); + /// assert_eq!(512, b.capacity()); + /// // 1000 bits requires 128 bytes (space for 1024 bits) + /// b.append_n(1000, true); + /// assert_eq!(1024, b.capacity()); + /// ``` #[inline] pub fn capacity(&self) -> usize { self.buffer.capacity() * 8 @@ -389,7 +416,7 @@ mod tests { let mut buffer = BooleanBufferBuilder::new(12); let mut all_bools = vec![]; - let mut rng = rand::thread_rng(); + let mut rng = rand::rng(); let src_len = 32; let (src, compacted_src) = { diff --git a/arrow-buffer/src/builder/mod.rs b/arrow-buffer/src/builder/mod.rs index f7e0e29dace4..abe510bdabc6 100644 --- a/arrow-buffer/src/builder/mod.rs +++ b/arrow-buffer/src/builder/mod.rs @@ -26,7 +26,7 @@ pub use null::*; pub use offset::*; use crate::{ArrowNativeType, Buffer, MutableBuffer}; -use std::{iter, marker::PhantomData}; +use std::marker::PhantomData; /// Builder for creating a [Buffer] object. /// @@ -214,7 +214,7 @@ impl BufferBuilder { #[inline] pub fn append_n(&mut self, n: usize, v: T) { self.reserve(n); - self.extend(iter::repeat(v).take(n)) + self.extend(std::iter::repeat_n(v, n)) } /// Appends `n`, zero-initialized values diff --git a/arrow-buffer/src/builder/null.rs b/arrow-buffer/src/builder/null.rs index 298b479e87df..e6f426615be5 100644 --- a/arrow-buffer/src/builder/null.rs +++ b/arrow-buffer/src/builder/null.rs @@ -17,23 +17,55 @@ use crate::{BooleanBufferBuilder, MutableBuffer, NullBuffer}; -/// Builder for creating the null bit buffer. +/// Builder for creating [`NullBuffer`] +/// +/// # Performance /// /// This builder only materializes the buffer when we append `false`. /// If you only append `true`s to the builder, what you get will be /// `None` when calling [`finish`](#method.finish). -/// This optimization is **very** important for the performance. +/// +/// This optimization is **very** important for the performance as it avoids +/// allocating memory for the null buffer when there are no nulls. +/// +/// See [`Self::allocated_size`] to get the current memory allocated by the builder. +/// +/// # Example +/// ``` +/// # use arrow_buffer::NullBufferBuilder; +/// let mut builder = NullBufferBuilder::new(8); +/// builder.append_n_non_nulls(8); +/// // If no non null values are appended, the null buffer is not created +/// let buffer = builder.finish(); +/// assert!(buffer.is_none()); +/// // however, if a null value is appended, the null buffer is created +/// let mut builder = NullBufferBuilder::new(8); +/// builder.append_n_non_nulls(7); +/// builder.append_null(); +/// let buffer = builder.finish().unwrap(); +/// assert_eq!(buffer.len(), 8); +/// assert_eq!(buffer.iter().collect::>(), vec![true, true, true, true, true, true, true, false]); +/// ``` #[derive(Debug)] pub struct NullBufferBuilder { + /// The bitmap builder to store the null buffer: + /// * `Some` if any nulls have been appended ("materialized") + /// * `None` if no nulls have been appended. bitmap_builder: Option, - /// Store the length of the buffer before materializing. + /// Length of the buffer before materializing. + /// + /// if `bitmap_buffer` buffer is `Some`, this value is not used. len: usize, + /// Initial capacity of the `bitmap_builder`, when it is materialized. capacity: usize, } impl NullBufferBuilder { /// Creates a new empty builder. - /// `capacity` is the number of bits in the null buffer. + /// + /// Note that this method does not allocate any memory, regardless of the + /// `capacity` parameter. If an allocation is required, `capacity` is the + /// size in bits (not bytes) that will be allocated at minimum. pub fn new(capacity: usize) -> Self { Self { bitmap_builder: None, @@ -54,7 +86,6 @@ impl NullBufferBuilder { /// Creates a new builder from a `MutableBuffer`. pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self { let capacity = buffer.len() * 8; - assert!(len <= capacity); let bitmap_builder = Some(BooleanBufferBuilder::new_from_buffer(buffer, len)); @@ -113,6 +144,28 @@ impl NullBufferBuilder { } } + /// Gets a bit in the buffer at `index` + #[inline] + pub fn is_valid(&self, index: usize) -> bool { + if let Some(ref buf) = self.bitmap_builder { + buf.get_bit(index) + } else { + true + } + } + + /// Truncates the builder to the given length + /// + /// If `len` is greater than the buffer's current length, this has no effect + #[inline] + pub fn truncate(&mut self, len: usize) { + if let Some(buf) = self.bitmap_builder.as_mut() { + buf.truncate(len); + } else if len <= self.len { + self.len = len + } + } + /// Appends a boolean slice into the builder /// to indicate the validations of these items. pub fn append_slice(&mut self, slice: &[bool]) { @@ -126,6 +179,20 @@ impl NullBufferBuilder { } } + /// Append [`NullBuffer`] to this [`NullBufferBuilder`] + /// + /// This is useful when you want to concatenate two null buffers. + pub fn append_buffer(&mut self, buffer: &NullBuffer) { + if buffer.null_count() > 0 { + self.materialize_if_needed(); + } + if let Some(buf) = self.bitmap_builder.as_mut() { + buf.append_buffer(buffer.inner()) + } else { + self.len += buffer.len(); + } + } + /// Builds the null buffer and resets the builder. /// Returns `None` if the builder only contains `true`s. pub fn finish(&mut self) -> Option { @@ -168,7 +235,7 @@ impl NullBufferBuilder { pub fn allocated_size(&self) -> usize { self.bitmap_builder .as_ref() - .map(|b| b.capacity()) + .map(|b| b.capacity() / 8) .unwrap_or(0) } } @@ -197,6 +264,7 @@ mod tests { builder.append_n_nulls(2); builder.append_n_non_nulls(2); assert_eq!(6, builder.len()); + assert_eq!(64, builder.allocated_size()); let buf = builder.finish().unwrap(); assert_eq!(&[0b110010_u8], buf.validity()); @@ -209,6 +277,7 @@ mod tests { builder.append_n_nulls(2); builder.append_slice(&[false, false, false]); assert_eq!(6, builder.len()); + assert_eq!(64, builder.allocated_size()); let buf = builder.finish().unwrap(); assert_eq!(&[0b0_u8], buf.validity()); @@ -221,6 +290,7 @@ mod tests { builder.append_n_non_nulls(2); builder.append_slice(&[true, true, true]); assert_eq!(6, builder.len()); + assert_eq!(0, builder.allocated_size()); let buf = builder.finish(); assert!(buf.is_none()); @@ -242,4 +312,94 @@ mod tests { let buf = builder.finish().unwrap(); assert_eq!(&[0b1011_u8], buf.validity()); } + + #[test] + fn test_null_buffer_builder_is_valid() { + let mut builder = NullBufferBuilder::new(0); + builder.append_n_non_nulls(6); + assert!(builder.is_valid(0)); + + builder.append_null(); + assert!(!builder.is_valid(6)); + + builder.append_non_null(); + assert!(builder.is_valid(7)); + } + + #[test] + fn test_null_buffer_builder_truncate() { + let mut builder = NullBufferBuilder::new(10); + builder.append_n_non_nulls(16); + assert_eq!(builder.as_slice(), None); + builder.truncate(20); + assert_eq!(builder.as_slice(), None); + assert_eq!(builder.len(), 16); + assert_eq!(builder.allocated_size(), 0); + builder.truncate(14); + assert_eq!(builder.as_slice(), None); + assert_eq!(builder.len(), 14); + builder.append_null(); + builder.append_non_null(); + assert_eq!(builder.as_slice().unwrap(), &[0xFF, 0b10111111]); + assert_eq!(builder.allocated_size(), 64); + } + + #[test] + fn test_null_buffer_builder_truncate_never_materialized() { + let mut builder = NullBufferBuilder::new(0); + assert_eq!(builder.len(), 0); + builder.append_n_nulls(2); // doesn't materialize + assert_eq!(builder.len(), 2); + builder.truncate(1); + assert_eq!(builder.len(), 1); + } + + #[test] + fn test_append_buffers() { + let mut builder = NullBufferBuilder::new(0); + let buffer1 = NullBuffer::from(&[true, true]); + let buffer2 = NullBuffer::from(&[true, true, false]); + + builder.append_buffer(&buffer1); + builder.append_buffer(&buffer2); + + assert_eq!(builder.as_slice().unwrap(), &[0b01111_u8]); + } + + #[test] + fn test_append_buffers_with_unaligned_length() { + let mut builder = NullBufferBuilder::new(0); + let buffer = NullBuffer::from(&[true, true, false, true, false]); + builder.append_buffer(&buffer); + assert_eq!(builder.as_slice().unwrap(), &[0b01011_u8]); + + let buffer = NullBuffer::from(&[false, false, true, true, true, false, false]); + builder.append_buffer(&buffer); + assert_eq!(builder.as_slice().unwrap(), &[0b10001011_u8, 0b0011_u8]); + } + + #[test] + fn test_append_empty_buffer() { + let mut builder = NullBufferBuilder::new(0); + let buffer = NullBuffer::from(&[true, true, false, true]); + builder.append_buffer(&buffer); + assert_eq!(builder.as_slice().unwrap(), &[0b1011_u8]); + + let buffer = NullBuffer::from(&[]); + builder.append_buffer(&buffer); + + assert_eq!(builder.as_slice().unwrap(), &[0b1011_u8]); + } + + #[test] + fn test_should_not_materialize_when_appending_all_valid_buffers() { + let mut builder = NullBufferBuilder::new(0); + let buffer = NullBuffer::from(&[true; 10]); + builder.append_buffer(&buffer); + + let buffer = NullBuffer::from(&[true; 2]); + builder.append_buffer(&buffer); + + assert_eq!(builder.finish(), None); + } } diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs index 77724137aef7..8f912b807da5 100644 --- a/arrow-buffer/src/bytes.rs +++ b/arrow-buffer/src/bytes.rs @@ -26,16 +26,25 @@ use std::{fmt::Debug, fmt::Formatter}; use crate::alloc::Deallocation; use crate::buffer::dangling_ptr; +#[cfg(feature = "pool")] +use crate::pool::{MemoryPool, MemoryReservation}; +#[cfg(feature = "pool")] +use std::sync::Mutex; + /// A continuous, fixed-size, immutable memory region that knows how to de-allocate itself. /// -/// This structs' API is inspired by the `bytes::Bytes`, but it is not limited to using rust's -/// global allocator nor u8 alignment. +/// Note that this structure is an internal implementation detail of the +/// arrow-rs crate. While it has the same name and similar API as +/// [`bytes::Bytes`] it is not limited to rust's global allocator nor u8 +/// alignment. It is possible to create a `Bytes` from `bytes::Bytes` using the +/// `From` implementation. /// /// In the most common case, this buffer is allocated using [`alloc`](std::alloc::alloc) /// with an alignment of [`ALIGNMENT`](crate::alloc::ALIGNMENT) /// /// When the region is allocated by a different allocator, [Deallocation::Custom], this calls the /// custom deallocator to deallocate the region when it is no longer needed. +/// pub struct Bytes { /// The raw pointer to be beginning of the region ptr: NonNull, @@ -45,6 +54,10 @@ pub struct Bytes { /// how to deallocate this region deallocation: Deallocation, + + /// Memory reservation for tracking memory usage + #[cfg(feature = "pool")] + pub(super) reservation: Mutex>>, } impl Bytes { @@ -66,6 +79,8 @@ impl Bytes { ptr, len, deallocation, + #[cfg(feature = "pool")] + reservation: Mutex::new(None), } } @@ -97,6 +112,27 @@ impl Bytes { } } + /// Register this [`Bytes`] with the provided [`MemoryPool`], replacing any prior reservation. + #[cfg(feature = "pool")] + pub fn claim(&self, pool: &dyn MemoryPool) { + *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity())); + } + + /// Resize the memory reservation of this buffer + /// + /// This is a no-op if this buffer doesn't have a reservation. + #[cfg(feature = "pool")] + fn resize_reservation(&self, new_size: usize) { + let mut guard = self.reservation.lock().unwrap(); + if let Some(mut reservation) = guard.take() { + // Resize the reservation + reservation.resize(new_size); + + // Put it back + *guard = Some(reservation); + } + } + /// Try to reallocate the underlying memory region to a new size (smaller or larger). /// /// Only works for bytes allocated with the standard allocator. @@ -131,6 +167,13 @@ impl Bytes { self.ptr = ptr; self.len = new_len; self.deallocation = Deallocation::Standard(new_layout); + + #[cfg(feature = "pool")] + { + // Resize reservation + self.resize_reservation(new_len); + } + return Ok(()); } } @@ -195,6 +238,8 @@ impl From for Bytes { len, ptr: NonNull::new(value.as_ptr() as _).unwrap(), deallocation: Deallocation::Custom(std::sync::Arc::new(value), len), + #[cfg(feature = "pool")] + reservation: Mutex::new(None), } } } @@ -205,14 +250,83 @@ mod tests { #[test] fn test_from_bytes() { - let bytes = bytes::Bytes::from(vec![1, 2, 3, 4]); - let arrow_bytes: Bytes = bytes.clone().into(); + let message = b"hello arrow"; + + // we can create a Bytes from bytes::Bytes (created from slices) + let c_bytes: bytes::Bytes = message.as_ref().into(); + let a_bytes: Bytes = c_bytes.into(); + assert_eq!(a_bytes.as_slice(), message); + + // we can create a Bytes from bytes::Bytes (created from Vec) + let c_bytes: bytes::Bytes = bytes::Bytes::from(message.to_vec()); + let a_bytes: Bytes = c_bytes.into(); + assert_eq!(a_bytes.as_slice(), message); + } - assert_eq!(bytes.as_ptr(), arrow_bytes.as_ptr()); + #[cfg(feature = "pool")] + mod pool_tests { + use super::*; + + use crate::pool::TrackingMemoryPool; + + #[test] + fn test_bytes_with_pool() { + // Create a standard allocation + let buffer = unsafe { + let layout = + std::alloc::Layout::from_size_align(1024, crate::alloc::ALIGNMENT).unwrap(); + let ptr = std::alloc::alloc(layout); + assert!(!ptr.is_null()); + + Bytes::new( + NonNull::new(ptr).unwrap(), + 1024, + Deallocation::Standard(layout), + ) + }; + + // Create a memory pool + let pool = TrackingMemoryPool::default(); + assert_eq!(pool.used(), 0); + + // Reserve memory and assign to buffer. Claim twice. + buffer.claim(&pool); + assert_eq!(pool.used(), 1024); + buffer.claim(&pool); + assert_eq!(pool.used(), 1024); + + // Memory should be released when buffer is dropped + drop(buffer); + assert_eq!(pool.used(), 0); + } - drop(bytes); - drop(arrow_bytes); + #[test] + fn test_bytes_drop_releases_pool() { + let pool = TrackingMemoryPool::default(); - let _ = Bytes::from(bytes::Bytes::new()); + { + // Create a buffer with pool + let _buffer = unsafe { + let layout = + std::alloc::Layout::from_size_align(1024, crate::alloc::ALIGNMENT).unwrap(); + let ptr = std::alloc::alloc(layout); + assert!(!ptr.is_null()); + + let bytes = Bytes::new( + NonNull::new(ptr).unwrap(), + 1024, + Deallocation::Standard(layout), + ); + + bytes.claim(&pool); + bytes + }; + + assert_eq!(pool.used(), 1024); + } + + // Buffer has been dropped, memory should be released + assert_eq!(pool.used(), 0); + } } } diff --git a/arrow-buffer/src/lib.rs b/arrow-buffer/src/lib.rs index 34e432208ada..1090146f3636 100644 --- a/arrow-buffer/src/lib.rs +++ b/arrow-buffer/src/lib.rs @@ -17,6 +17,11 @@ //! Low-level buffer abstractions for [Apache Arrow Rust](https://docs.rs/arrow) +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // used by [`buffer::mutable::dangling_ptr`] #![cfg_attr(miri, feature(strict_provenance))] #![warn(missing_docs)] @@ -43,3 +48,8 @@ mod interval; pub use interval::*; mod arith; + +#[cfg(feature = "pool")] +mod pool; +#[cfg(feature = "pool")] +pub use pool::*; diff --git a/arrow-buffer/src/pool.rs b/arrow-buffer/src/pool.rs new file mode 100644 index 000000000000..bf22d433d615 --- /dev/null +++ b/arrow-buffer/src/pool.rs @@ -0,0 +1,189 @@ +// 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. + +//! This module contains traits for memory pool traits and an implementation +//! for tracking memory usage. +//! +//! The basic traits are [`MemoryPool`] and [`MemoryReservation`]. And default +//! implementation of [`MemoryPool`] is [`TrackingMemoryPool`]. Their relationship +//! is as follows: +//! +//! ```text +//! (pool tracker) (resizable) +//! ┌──────────────────┐ fn reserve() ┌─────────────────────────┐ +//! │ trait MemoryPool │─────────────►│ trait MemoryReservation │ +//! └──────────────────┘ └─────────────────────────┘ +//! ``` + +use std::fmt::Debug; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +/// A memory reservation within a [`MemoryPool`] that is freed on drop +pub trait MemoryReservation: Debug + Send + Sync { + /// Returns the size of this reservation in bytes. + fn size(&self) -> usize; + + /// Resize this reservation to a new size in bytes. + fn resize(&mut self, new_size: usize); +} + +/// A pool of memory that can be reserved and released. +/// +/// This is used to accurately track memory usage when buffers are shared +/// between multiple arrays or other data structures. +/// +/// For example, assume we have two arrays that share underlying buffer. +/// It's hard to tell how much memory is used by them because we can't +/// tell if the buffer is shared or not. +/// +/// ```text +/// Array A Array B +/// ┌────────────┐ ┌────────────┐ +/// │ slices... │ │ slices... │ +/// │────────────│ │────────────│ +/// │ Arc │ │ Arc │ (shared buffer) +/// └─────▲──────┘ └───────▲────┘ +/// │ │ +/// │ Bytes │ +/// │ ┌─────────────┐ │ +/// │ │ data... │ │ +/// │ │─────────────│ │ +/// └──│ Memory │──┘ (tracked with a memory pool) +/// │ Reservation │ +/// └─────────────┘ +/// ``` +/// +/// With a memory pool, we can count the memory usage by the shared buffer +/// directly. +pub trait MemoryPool: Debug + Send + Sync { + /// Reserves memory from the pool. Infallible. + /// + /// Returns a reservation of the requested size. + fn reserve(&self, size: usize) -> Box; + + /// Returns the current available memory in the pool. + /// + /// The pool may be overfilled, so this method might return a negative value. + fn available(&self) -> isize; + + /// Returns the current used memory from the pool. + fn used(&self) -> usize; + + /// Returns the maximum memory that can be reserved from the pool. + fn capacity(&self) -> usize; +} + +/// A simple [`MemoryPool`] that reports the total memory usage +#[derive(Debug, Default)] +pub struct TrackingMemoryPool(Arc); + +impl TrackingMemoryPool { + /// Returns the total allocated size + pub fn allocated(&self) -> usize { + self.0.load(Ordering::Relaxed) + } +} + +impl MemoryPool for TrackingMemoryPool { + fn reserve(&self, size: usize) -> Box { + self.0.fetch_add(size, Ordering::Relaxed); + Box::new(Tracker { + size, + shared: Arc::clone(&self.0), + }) + } + + fn available(&self) -> isize { + isize::MAX - self.used() as isize + } + + fn used(&self) -> usize { + self.0.load(Ordering::Relaxed) + } + + fn capacity(&self) -> usize { + usize::MAX + } +} + +#[derive(Debug)] +struct Tracker { + size: usize, + shared: Arc, +} + +impl Drop for Tracker { + fn drop(&mut self) { + self.shared.fetch_sub(self.size, Ordering::Relaxed); + } +} + +impl MemoryReservation for Tracker { + fn size(&self) -> usize { + self.size + } + + fn resize(&mut self, new: usize) { + match self.size < new { + true => self.shared.fetch_add(new - self.size, Ordering::Relaxed), + false => self.shared.fetch_sub(self.size - new, Ordering::Relaxed), + }; + self.size = new; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tracking_memory_pool() { + let pool = TrackingMemoryPool::default(); + + // Reserve 512 bytes + let reservation = pool.reserve(512); + assert_eq!(reservation.size(), 512); + assert_eq!(pool.used(), 512); + assert_eq!(pool.available(), isize::MAX - 512); + + // Reserve another 256 bytes + let reservation2 = pool.reserve(256); + assert_eq!(reservation2.size(), 256); + assert_eq!(pool.used(), 768); + assert_eq!(pool.available(), isize::MAX - 768); + + // Test resize to increase + let mut reservation_mut = reservation; + reservation_mut.resize(600); + assert_eq!(reservation_mut.size(), 600); + assert_eq!(pool.used(), 856); // 600 + 256 + + // Test resize to decrease + reservation_mut.resize(400); + assert_eq!(reservation_mut.size(), 400); + assert_eq!(pool.used(), 656); // 400 + 256 + + // Drop the first reservation + drop(reservation_mut); + assert_eq!(pool.used(), 256); + + // Drop the second reservation + drop(reservation2); + assert_eq!(pool.used(), 0); + } +} diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs b/arrow-buffer/src/util/bit_chunk_iterator.rs index 54995314c49b..ea8e8f472ace 100644 --- a/arrow-buffer/src/util/bit_chunk_iterator.rs +++ b/arrow-buffer/src/util/bit_chunk_iterator.rs @@ -53,7 +53,7 @@ impl<'a> UnalignedBitChunk<'a> { let byte_offset = offset / 8; let offset_padding = offset % 8; - let bytes_len = (len + offset_padding + 7) / 8; + let bytes_len = (len + offset_padding).div_ceil(8); let buffer = &buffer[byte_offset..byte_offset + bytes_len]; let prefix_mask = compute_prefix_mask(offset_padding); @@ -371,7 +371,10 @@ impl ExactSizeIterator for BitChunkIterator<'_> { #[cfg(test)] mod tests { + use rand::distr::uniform::UniformSampler; + use rand::distr::uniform::UniformUsize; use rand::prelude::*; + use rand::rng; use crate::buffer::Buffer; use crate::util::bit_chunk_iterator::UnalignedBitChunk; @@ -624,21 +627,25 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn fuzz_unaligned_bit_chunk_iterator() { - let mut rng = thread_rng(); + let mut rng = rng(); + let uusize = UniformUsize::new(usize::MIN, usize::MAX).unwrap(); for _ in 0..100 { - let mask_len = rng.gen_range(0..1024); - let bools: Vec<_> = std::iter::from_fn(|| Some(rng.gen())) + let mask_len = rng.random_range(0..1024); + let bools: Vec<_> = std::iter::from_fn(|| Some(rng.random())) .take(mask_len) .collect(); let buffer = Buffer::from_iter(bools.iter().cloned()); let max_offset = 64.min(mask_len); - let offset = rng.gen::().checked_rem(max_offset).unwrap_or(0); + let offset = uusize.sample(&mut rng).checked_rem(max_offset).unwrap_or(0); let max_truncate = 128.min(mask_len - offset); - let truncate = rng.gen::().checked_rem(max_truncate).unwrap_or(0); + let truncate = uusize + .sample(&mut rng) + .checked_rem(max_truncate) + .unwrap_or(0); let unaligned = UnalignedBitChunk::new(buffer.as_slice(), offset, mask_len - offset - truncate); diff --git a/arrow-buffer/src/util/bit_iterator.rs b/arrow-buffer/src/util/bit_iterator.rs index f667ab1e7b9d..6a783138884b 100644 --- a/arrow-buffer/src/util/bit_iterator.rs +++ b/arrow-buffer/src/util/bit_iterator.rs @@ -93,6 +93,8 @@ impl DoubleEndedIterator for BitIterator<'_> { /// Returns `(usize, usize)` each representing an interval where the corresponding /// bits in the provides mask are set /// +/// the first value is the start of the range (inclusive) and the second value is the end of the range (exclusive) +/// #[derive(Debug)] pub struct BitSliceIterator<'a> { iter: UnalignedBitChunkIterator<'a>, @@ -214,6 +216,7 @@ impl<'a> BitIndexIterator<'a> { impl Iterator for BitIndexIterator<'_> { type Item = usize; + #[inline] fn next(&mut self) -> Option { loop { if self.current_chunk != 0 { diff --git a/arrow-buffer/src/util/bit_mask.rs b/arrow-buffer/src/util/bit_mask.rs index 97be7e006dec..6030cb4b1b8c 100644 --- a/arrow-buffer/src/util/bit_mask.rs +++ b/arrow-buffer/src/util/bit_mask.rs @@ -164,7 +164,7 @@ mod tests { use super::*; use crate::bit_util::{get_bit, set_bit, unset_bit}; use rand::prelude::StdRng; - use rand::{Fill, Rng, SeedableRng}; + use rand::{Rng, SeedableRng, TryRngCore}; use std::fmt::Display; #[test] @@ -278,7 +278,7 @@ mod tests { impl Display for BinaryFormatter<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { for byte in self.0 { - write!(f, "{:08b} ", byte)?; + write!(f, "{byte:08b} ")?; } write!(f, " ")?; Ok(()) @@ -322,20 +322,20 @@ mod tests { // -------------------+-----------------+------- // length of data to copy - let len = rng.gen_range(0..=200); + let len = rng.random_range(0..=200); // randomly pick where we will write to - let offset_write_bits = rng.gen_range(0..=200); + let offset_write_bits = rng.random_range(0..=200); let offset_write_bytes = if offset_write_bits % 8 == 0 { offset_write_bits / 8 } else { (offset_write_bits / 8) + 1 }; - let extra_write_data_bytes = rng.gen_range(0..=5); // ensure 0 shows up often + let extra_write_data_bytes = rng.random_range(0..=5); // ensure 0 shows up often // randomly decide where we will read from - let extra_read_data_bytes = rng.gen_range(0..=5); // make sure 0 shows up often - let offset_read_bits = rng.gen_range(0..=200); + let extra_read_data_bytes = rng.random_range(0..=5); // make sure 0 shows up often + let offset_read_bits = rng.random_range(0..=200); let offset_read_bytes = if offset_read_bits % 8 != 0 { (offset_read_bits / 8) + 1 } else { @@ -356,7 +356,7 @@ mod tests { self.data .resize(offset_read_bytes + len + extra_read_data_bytes, 0); // fill source data with random bytes - self.data.try_fill(rng).unwrap(); + rng.try_fill_bytes(self.data.as_mut_slice()).unwrap(); self.offset_read = offset_read_bits; self.len = len; @@ -389,8 +389,8 @@ mod tests { self.len, ); - assert_eq!(actual, self.expected_data, "self: {}", self); - assert_eq!(null_count, self.expected_null_count, "self: {}", self); + assert_eq!(actual, self.expected_data, "self: {self}"); + assert_eq!(null_count, self.expected_null_count, "self: {self}"); } } diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index ed5d363d607f..c297321bdcf9 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -20,7 +20,8 @@ /// Returns the nearest number that is `>=` than `num` and is a multiple of 64 #[inline] pub fn round_upto_multiple_of_64(num: usize) -> usize { - round_upto_power_of_2(num, 64) + num.checked_next_multiple_of(64) + .expect("failed to round upto multiple of 64") } /// Returns the nearest multiple of `factor` that is `>=` than `num`. Here `factor` must @@ -86,9 +87,7 @@ pub unsafe fn unset_bit_raw(data: *mut u8, i: usize) { /// Returns the ceil of `value`/`divisor` #[inline] pub fn ceil(value: usize, divisor: usize) -> usize { - // Rewrite as `value.div_ceil(&divisor)` after - // https://github.com/rust-lang/rust/issues/88581 is merged. - value / divisor + (0 != value % divisor) as usize + value.div_ceil(divisor) } #[cfg(test)] @@ -109,6 +108,12 @@ mod tests { assert_eq!(192, round_upto_multiple_of_64(129)); } + #[test] + #[should_panic(expected = "failed to round upto multiple of 64")] + fn test_round_upto_multiple_of_64_panic() { + let _ = round_upto_multiple_of_64(usize::MAX); + } + #[test] #[should_panic(expected = "failed to round to next highest power of 2")] fn test_round_upto_panic() { @@ -153,7 +158,7 @@ mod tests { let mut expected = vec![]; let mut rng = seedable_rng(); for i in 0..8 * NUM_BYTE { - let b = rng.gen_bool(0.5); + let b = rng.random_bool(0.5); expected.push(b); if b { set_bit(&mut buf[..], i) @@ -197,7 +202,7 @@ mod tests { let mut expected = vec![]; let mut rng = seedable_rng(); for i in 0..8 * NUM_BYTE { - let b = rng.gen_bool(0.5); + let b = rng.random_bool(0.5); expected.push(b); if b { unsafe { @@ -221,7 +226,7 @@ mod tests { let mut expected = vec![]; let mut rng = seedable_rng(); for i in 0..8 * NUM_BYTE { - let b = rng.gen_bool(0.5); + let b = rng.random_bool(0.5); expected.push(b); if !b { unsafe { @@ -247,7 +252,7 @@ mod tests { let mut v = HashSet::new(); let mut rng = seedable_rng(); for _ in 0..NUM_SETS { - let offset = rng.gen_range(0..8 * NUM_BYTES); + let offset = rng.random_range(0..8 * NUM_BYTES); v.insert(offset); set_bit(&mut buffer[..], offset); } diff --git a/arrow-cast/Cargo.toml b/arrow-cast/Cargo.toml index 4046f5226094..49145cf987f9 100644 --- a/arrow-cast/Cargo.toml +++ b/arrow-cast/Cargo.toml @@ -30,11 +30,10 @@ rust-version = { workspace = true } [lib] name = "arrow_cast" -path = "src/lib.rs" bench = false [package.metadata.docs.rs] -features = ["prettyprint"] +all-features = true [features] prettyprint = ["comfy-table"] @@ -58,9 +57,7 @@ ryu = "1.0.16" [dev-dependencies] criterion = { version = "0.5", default-features = false } half = { version = "2.1", default-features = false } -rand = "0.8" - -[build-dependencies] +rand = "0.9" [[bench]] name = "parse_timestamp" diff --git a/arrow-cast/benches/parse_date.rs b/arrow-cast/benches/parse_date.rs index e05d38d2f853..0fdc13911837 100644 --- a/arrow-cast/benches/parse_date.rs +++ b/arrow-cast/benches/parse_date.rs @@ -18,12 +18,13 @@ use arrow_array::types::Date32Type; use arrow_cast::parse::Parser; use criterion::*; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let timestamps = ["2020-09-08", "2020-9-8", "2020-09-8", "2020-9-08"]; for timestamp in timestamps { - let t = black_box(timestamp); + let t = hint::black_box(timestamp); c.bench_function(t, |b| { b.iter(|| Date32Type::parse(t).unwrap()); }); diff --git a/arrow-cast/benches/parse_decimal.rs b/arrow-cast/benches/parse_decimal.rs index 5682859dd25a..28364b733ce4 100644 --- a/arrow-cast/benches/parse_decimal.rs +++ b/arrow-cast/benches/parse_decimal.rs @@ -18,6 +18,7 @@ use arrow_array::types::Decimal256Type; use arrow_cast::parse::parse_decimal; use criterion::*; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let decimals = [ @@ -45,7 +46,7 @@ fn criterion_benchmark(c: &mut Criterion) { ]; for decimal in decimals { - let d = black_box(decimal); + let d = hint::black_box(decimal); c.bench_function(d, |b| { b.iter(|| parse_decimal::(d, 20, 3).unwrap()); }); diff --git a/arrow-cast/benches/parse_time.rs b/arrow-cast/benches/parse_time.rs index d28b9c7c613d..8d3915bdafd6 100644 --- a/arrow-cast/benches/parse_time.rs +++ b/arrow-cast/benches/parse_time.rs @@ -17,6 +17,7 @@ use arrow_cast::parse::string_to_time_nanoseconds; use criterion::*; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let timestamps = [ @@ -31,7 +32,7 @@ fn criterion_benchmark(c: &mut Criterion) { ]; for timestamp in timestamps { - let t = black_box(timestamp); + let t = hint::black_box(timestamp); c.bench_function(t, |b| { b.iter(|| string_to_time_nanoseconds(t).unwrap()); }); diff --git a/arrow-cast/benches/parse_timestamp.rs b/arrow-cast/benches/parse_timestamp.rs index d3ab41863e70..12cad5b78e9d 100644 --- a/arrow-cast/benches/parse_timestamp.rs +++ b/arrow-cast/benches/parse_timestamp.rs @@ -17,6 +17,7 @@ use arrow_cast::parse::string_to_timestamp_nanos; use criterion::*; +use std::hint; fn criterion_benchmark(c: &mut Criterion) { let timestamps = [ @@ -33,7 +34,7 @@ fn criterion_benchmark(c: &mut Criterion) { ]; for timestamp in timestamps { - let t = black_box(timestamp); + let t = hint::black_box(timestamp); c.bench_function(t, |b| { b.iter(|| string_to_timestamp_nanos(t).unwrap()); }); diff --git a/arrow-cast/src/base64.rs b/arrow-cast/src/base64.rs index 534b21878c56..e7bb84ebe24c 100644 --- a/arrow-cast/src/base64.rs +++ b/arrow-cast/src/base64.rs @@ -90,7 +90,7 @@ pub fn b64_decode( mod tests { use super::*; use arrow_array::BinaryArray; - use rand::{thread_rng, Rng}; + use rand::{rng, Rng}; fn test_engine(e: &E, a: &BinaryArray) { let encoded = b64_encode(e, a); @@ -105,12 +105,12 @@ mod tests { #[test] fn test_b64() { - let mut rng = thread_rng(); - let len = rng.gen_range(1024..1050); + let mut rng = rng(); + let len = rng.random_range(1024..1050); let data: BinaryArray = (0..len) .map(|_| { - let len = rng.gen_range(0..16); - Some((0..len).map(|_| rng.gen()).collect::>()) + let len = rng.random_range(0..16); + Some((0..len).map(|_| rng.random()).collect::>()) }) .collect(); diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index ba82ca9040c7..597f384fa452 100644 --- a/arrow-cast/src/cast/decimal.rs +++ b/arrow-cast/src/cast/decimal.rs @@ -25,6 +25,8 @@ pub(crate) trait DecimalCast: Sized { fn to_i256(self) -> Option; fn from_decimal(n: T) -> Option; + + fn from_f64(n: f64) -> Option; } impl DecimalCast for i128 { @@ -39,6 +41,10 @@ impl DecimalCast for i128 { fn from_decimal(n: T) -> Option { n.to_i128() } + + fn from_f64(n: f64) -> Option { + n.to_i128() + } } impl DecimalCast for i256 { @@ -53,6 +59,10 @@ impl DecimalCast for i256 { fn from_decimal(n: T) -> Option { n.to_i256() } + + fn from_f64(n: f64) -> Option { + i256::from_f64(n) + } } pub(crate) fn cast_decimal_to_decimal_error( @@ -78,6 +88,7 @@ where pub(crate) fn convert_to_smaller_scale_decimal( array: &PrimitiveArray, + input_precision: u8, input_scale: i8, output_precision: u8, output_scale: i8, @@ -90,9 +101,22 @@ where O::Native: DecimalCast + ArrowNativeTypeOp, { let error = cast_decimal_to_decimal_error::(output_precision, output_scale); + let delta_scale = input_scale - output_scale; + // if the reduction of the input number through scaling (dividing) is greater + // than a possible precision loss (plus potential increase via rounding) + // every input number will fit into the output type + // Example: If we are starting with any number of precision 5 [xxxxx], + // then and decrease the scale by 3 will have the following effect on the representation: + // [xxxxx] -> [xx] (+ 1 possibly, due to rounding). + // The rounding may add an additional digit, so the cast to be infallible, + // the output type needs to have at least 3 digits of precision. + // e.g. Decimal(5, 3) 99.999 to Decimal(3, 0) will result in 100: + // [99999] -> [99] + 1 = [100], a cast to Decimal(2, 0) would not be possible + let is_infallible_cast = (input_precision as i8) - delta_scale < (output_precision as i8); + let div = I::Native::from_decimal(10_i128) .unwrap() - .pow_checked((input_scale - output_scale) as u32)?; + .pow_checked(delta_scale as u32)?; let half = div.div_wrapping(I::Native::from_usize(2).unwrap()); let half_neg = half.neg_wrapping(); @@ -111,7 +135,13 @@ where O::Native::from_decimal(adjusted) }; - Ok(if cast_options.safe { + Ok(if is_infallible_cast { + // make sure we don't perform calculations that don't make sense w/o validation + validate_decimal_precision_and_scale::(output_precision, output_scale)?; + let g = |x: I::Native| f(x).unwrap(); // unwrapping is safe since the result is guaranteed + // to fit into the target type + array.unary(g) + } else if cast_options.safe { array.unary_opt(|x| f(x).filter(|v| O::is_valid_decimal_precision(*v, output_precision))) } else { array.try_unary(|x| { @@ -123,6 +153,7 @@ where pub(crate) fn convert_to_bigger_or_equal_scale_decimal( array: &PrimitiveArray, + input_precision: u8, input_scale: i8, output_precision: u8, output_scale: i8, @@ -135,13 +166,27 @@ where O::Native: DecimalCast + ArrowNativeTypeOp, { let error = cast_decimal_to_decimal_error::(output_precision, output_scale); + let delta_scale = output_scale - input_scale; let mul = O::Native::from_decimal(10_i128) .unwrap() - .pow_checked((output_scale - input_scale) as u32)?; - + .pow_checked(delta_scale as u32)?; + + // if the gain in precision (digits) is greater than the multiplication due to scaling + // every number will fit into the output type + // Example: If we are starting with any number of precision 5 [xxxxx], + // then an increase of scale by 3 will have the following effect on the representation: + // [xxxxx] -> [xxxxx000], so for the cast to be infallible, the output type + // needs to provide at least 8 digits precision + let is_infallible_cast = (input_precision as i8) + delta_scale <= (output_precision as i8); let f = |x| O::Native::from_decimal(x).and_then(|x| x.mul_checked(mul).ok()); - Ok(if cast_options.safe { + Ok(if is_infallible_cast { + // make sure we don't perform calculations that don't make sense w/o validation + validate_decimal_precision_and_scale::(output_precision, output_scale)?; + // unwrapping is safe since the result is guaranteed to fit into the target type + let f = |x| O::Native::from_decimal(x).unwrap().mul_wrapping(mul); + array.unary(f) + } else if cast_options.safe { array.unary_opt(|x| f(x).filter(|v| O::is_valid_decimal_precision(*v, output_precision))) } else { array.try_unary(|x| { @@ -167,18 +212,20 @@ where let array: PrimitiveArray = if input_scale == output_scale && input_precision <= output_precision { array.clone() - } else if input_scale < output_scale { - // the scale doesn't change, but precision may change and cause overflow + } else if input_scale <= output_scale { convert_to_bigger_or_equal_scale_decimal::( array, + input_precision, input_scale, output_precision, output_scale, cast_options, )? } else { + // input_scale > output_scale convert_to_smaller_scale_decimal::( array, + input_precision, input_scale, output_precision, output_scale, @@ -195,6 +242,7 @@ where // Support two different types of decimal cast operations pub(crate) fn cast_decimal_to_decimal( array: &PrimitiveArray, + input_precision: u8, input_scale: i8, output_precision: u8, output_scale: i8, @@ -209,6 +257,7 @@ where let array: PrimitiveArray = if input_scale > output_scale { convert_to_smaller_scale_decimal::( array, + input_precision, input_scale, output_precision, output_scale, @@ -217,6 +266,7 @@ where } else { convert_to_bigger_or_equal_scale_decimal::( array, + input_precision, input_scale, output_precision, output_scale, @@ -455,8 +505,7 @@ where )?, other => { return Err(ArrowError::ComputeError(format!( - "Cannot cast {:?} to decimal", - other + "Cannot cast {other:?} to decimal", ))) } }; @@ -464,52 +513,7 @@ where Ok(Arc::new(result)) } -pub(crate) fn cast_floating_point_to_decimal128( - array: &PrimitiveArray, - precision: u8, - scale: i8, - cast_options: &CastOptions, -) -> Result -where - ::Native: AsPrimitive, -{ - let mul = 10_f64.powi(scale as i32); - - if cast_options.safe { - array - .unary_opt::<_, Decimal128Type>(|v| { - (mul * v.as_()) - .round() - .to_i128() - .filter(|v| Decimal128Type::is_valid_decimal_precision(*v, precision)) - }) - .with_precision_and_scale(precision, scale) - .map(|a| Arc::new(a) as ArrayRef) - } else { - array - .try_unary::<_, Decimal128Type, _>(|v| { - (mul * v.as_()) - .round() - .to_i128() - .ok_or_else(|| { - ArrowError::CastError(format!( - "Cannot cast to {}({}, {}). Overflowing on {:?}", - Decimal128Type::PREFIX, - precision, - scale, - v - )) - }) - .and_then(|v| { - Decimal128Type::validate_decimal_precision(v, precision).map(|_| v) - }) - })? - .with_precision_and_scale(precision, scale) - .map(|a| Arc::new(a) as ArrayRef) - } -} - -pub(crate) fn cast_floating_point_to_decimal256( +pub(crate) fn cast_floating_point_to_decimal( array: &PrimitiveArray, precision: u8, scale: i8, @@ -517,33 +521,33 @@ pub(crate) fn cast_floating_point_to_decimal256( ) -> Result where ::Native: AsPrimitive, + D: DecimalType + ArrowPrimitiveType, + ::Native: DecimalCast, { let mul = 10_f64.powi(scale as i32); if cast_options.safe { array - .unary_opt::<_, Decimal256Type>(|v| { - i256::from_f64((v.as_() * mul).round()) - .filter(|v| Decimal256Type::is_valid_decimal_precision(*v, precision)) + .unary_opt::<_, D>(|v| { + D::Native::from_f64((mul * v.as_()).round()) + .filter(|v| D::is_valid_decimal_precision(*v, precision)) }) .with_precision_and_scale(precision, scale) .map(|a| Arc::new(a) as ArrayRef) } else { array - .try_unary::<_, Decimal256Type, _>(|v| { - i256::from_f64((v.as_() * mul).round()) + .try_unary::<_, D, _>(|v| { + D::Native::from_f64((mul * v.as_()).round()) .ok_or_else(|| { ArrowError::CastError(format!( "Cannot cast to {}({}, {}). Overflowing on {:?}", - Decimal256Type::PREFIX, + D::PREFIX, precision, scale, v )) }) - .and_then(|v| { - Decimal256Type::validate_decimal_precision(v, precision).map(|_| v) - }) + .and_then(|v| D::validate_decimal_precision(v, precision).map(|_| v)) })? .with_precision_and_scale(precision, scale) .map(|a| Arc::new(a) as ArrayRef) @@ -610,7 +614,11 @@ where Ok(Arc::new(value_builder.finish())) } -// Cast the decimal array to floating-point array +/// Cast a decimal array to a floating point array. +/// +/// Conversion is lossy and follows standard floating point semantics. Values +/// that exceed the representable range become `INFINITY` or `-INFINITY` without +/// returning an error. pub(crate) fn cast_decimal_to_float( array: &dyn Array, op: F, diff --git a/arrow-cast/src/cast/dictionary.rs b/arrow-cast/src/cast/dictionary.rs index ec0ab346f997..43a67a7d9a2d 100644 --- a/arrow-cast/src/cast/dictionary.rs +++ b/arrow-cast/src/cast/dictionary.rs @@ -214,50 +214,34 @@ pub(crate) fn cast_to_dictionary( UInt16 => pack_numeric_to_dictionary::(array, dict_value_type, cast_options), UInt32 => pack_numeric_to_dictionary::(array, dict_value_type, cast_options), UInt64 => pack_numeric_to_dictionary::(array, dict_value_type, cast_options), - Decimal128(p, s) => { - let dict = pack_numeric_to_dictionary::( - array, - dict_value_type, - cast_options, - )?; - let dict = dict - .as_dictionary::() - .downcast_dict::() - .ok_or_else(|| { - ArrowError::ComputeError( - "Internal Error: Cannot cast dict to Decimal128Array".to_string(), - ) - })?; - let value = dict.values().clone(); - // Set correct precision/scale - let value = value.with_precision_and_scale(p, s)?; - Ok(Arc::new(DictionaryArray::::try_new( - dict.keys().clone(), - Arc::new(value), - )?)) - } - Decimal256(p, s) => { - let dict = pack_numeric_to_dictionary::( - array, - dict_value_type, - cast_options, - )?; - let dict = dict - .as_dictionary::() - .downcast_dict::() - .ok_or_else(|| { - ArrowError::ComputeError( - "Internal Error: Cannot cast dict to Decimal256Array".to_string(), - ) - })?; - let value = dict.values().clone(); - // Set correct precision/scale - let value = value.with_precision_and_scale(p, s)?; - Ok(Arc::new(DictionaryArray::::try_new( - dict.keys().clone(), - Arc::new(value), - )?)) - } + Decimal32(p, s) => pack_decimal_to_dictionary::( + array, + dict_value_type, + p, + s, + cast_options, + ), + Decimal64(p, s) => pack_decimal_to_dictionary::( + array, + dict_value_type, + p, + s, + cast_options, + ), + Decimal128(p, s) => pack_decimal_to_dictionary::( + array, + dict_value_type, + p, + s, + cast_options, + ), + Decimal256(p, s) => pack_decimal_to_dictionary::( + array, + dict_value_type, + p, + s, + cast_options, + ), Float16 => { pack_numeric_to_dictionary::(array, dict_value_type, cast_options) } @@ -325,6 +309,9 @@ pub(crate) fn cast_to_dictionary( } pack_byte_to_dictionary::>(array, cast_options) } + FixedSizeBinary(byte_size) => { + pack_byte_to_fixed_size_dictionary::(array, cast_options, byte_size) + } _ => Err(ArrowError::CastError(format!( "Unsupported output type for dictionary packing: {dict_value_type:?}" ))), @@ -359,6 +346,36 @@ where Ok(Arc::new(b.finish())) } +pub(crate) fn pack_decimal_to_dictionary( + array: &dyn Array, + dict_value_type: &DataType, + precision: u8, + scale: i8, + cast_options: &CastOptions, +) -> Result +where + K: ArrowDictionaryKeyType, + D: DecimalType + ArrowPrimitiveType, +{ + let dict = pack_numeric_to_dictionary::(array, dict_value_type, cast_options)?; + let dict = dict + .as_dictionary::() + .downcast_dict::>() + .ok_or_else(|| { + ArrowError::ComputeError(format!( + "Internal Error: Cannot cast dict to {}Array", + D::PREFIX + )) + })?; + let value = dict.values().clone(); + // Set correct precision/scale + let value = value.with_precision_and_scale(precision, scale)?; + Ok(Arc::new(DictionaryArray::::try_new( + dict.keys().clone(), + Arc::new(value), + )?)) +} + pub(crate) fn string_view_to_dictionary( array: &dyn Array, ) -> Result @@ -450,3 +467,34 @@ where } Ok(Arc::new(b.finish())) } + +// Packs the data as a GenericByteDictionaryBuilder, if possible, with the +// key types of K +pub(crate) fn pack_byte_to_fixed_size_dictionary( + array: &dyn Array, + cast_options: &CastOptions, + byte_width: i32, +) -> Result +where + K: ArrowDictionaryKeyType, +{ + let cast_values = + cast_with_options(array, &DataType::FixedSizeBinary(byte_width), cast_options)?; + let values = cast_values + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::ComputeError("Internal Error: Cannot cast to GenericByteArray".to_string()) + })?; + let mut b = FixedSizeBinaryDictionaryBuilder::::with_capacity(1024, 1024, byte_width); + + // copy each element one at a time + for i in 0..values.len() { + if values.is_null(i) { + b.append_null(); + } else { + b.append(values.value(i))?; + } + } + Ok(Arc::new(b.finish())) +} diff --git a/arrow-cast/src/cast/list.rs b/arrow-cast/src/cast/list.rs index ec7a5c57d504..1728cc4061a8 100644 --- a/arrow-cast/src/cast/list.rs +++ b/arrow-cast/src/cast/list.rs @@ -24,7 +24,7 @@ pub(crate) fn cast_values_to_list( cast_options: &CastOptions, ) -> Result { let values = cast_with_options(array, to.data_type(), cast_options)?; - let offsets = OffsetBuffer::from_lengths(std::iter::repeat(1).take(values.len())); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(1, values.len())); let list = GenericListArray::::new(to.clone(), offsets, values, None); Ok(Arc::new(list)) } @@ -88,6 +88,17 @@ where let mut mutable = MutableArrayData::new(vec![&values], nullable, cap); // The end position in values of the last incorrectly-sized list slice let mut last_pos = 0; + + // Need to flag when previous vector(s) are empty/None to distinguish from 'All slices were correct length' cases. + let is_prev_empty = if array.offsets().len() < 2 { + false + } else { + let first_offset = array.offsets()[0].as_usize(); + let second_offset = array.offsets()[1].as_usize(); + + first_offset == 0 && second_offset == 0 + }; + for (idx, w) in array.offsets().windows(2).enumerate() { let start_pos = w[0].as_usize(); let end_pos = w[1].as_usize(); @@ -113,7 +124,7 @@ where } let values = match last_pos { - 0 => array.values().slice(0, cap), // All slices were the correct length + 0 if !is_prev_empty => array.values().slice(0, cap), // All slices were the correct length _ => { if mutable.len() != cap { // Remaining slices were all correct length diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index ba470635c6cd..dbe4401c7863 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -213,7 +213,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { (Binary, LargeBinary | Utf8 | LargeUtf8 | FixedSizeBinary(_) | BinaryView | Utf8View ) => true, (LargeBinary, Binary | Utf8 | LargeUtf8 | FixedSizeBinary(_) | BinaryView | Utf8View ) => true, - (FixedSizeBinary(_), Binary | LargeBinary) => true, + (FixedSizeBinary(_), Binary | LargeBinary | BinaryView) => true, ( Utf8 | LargeUtf8 | Utf8View, Binary @@ -234,7 +234,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { | BinaryView, ) => true, (Utf8 | LargeUtf8, Utf8View) => true, - (BinaryView, Binary | LargeBinary | Utf8 | LargeUtf8 | Utf8View ) => true, + (BinaryView, Binary | LargeBinary | Utf8 | LargeUtf8 | Utf8View) => true, (Utf8View | Utf8 | LargeUtf8, _) => to_type.is_numeric() && to_type != &Float16, (_, Utf8 | LargeUtf8) => from_type.is_primitive(), (_, Utf8View) => from_type.is_numeric(), @@ -265,8 +265,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool { } (Timestamp(_, _), _) if to_type.is_numeric() => true, (_, Timestamp(_, _)) if from_type.is_numeric() => true, - (Date64, Timestamp(_, None)) => true, - (Date32, Timestamp(_, None)) => true, + (Date64, Timestamp(_, _)) => true, + (Date32, Timestamp(_, _)) => true, ( Timestamp(_, _), Timestamp(_, _) @@ -603,6 +603,8 @@ fn timestamp_to_date32( /// * Temporal to/from backing Primitive: zero-copy with data type change /// * `Float32/Float64` to `Decimal(precision, scale)` rounds to the `scale` decimals /// (i.e. casting `6.4999` to `Decimal(10, 1)` becomes `6.5`). +/// * `Decimal` to `Float32/Float64` is lossy and values outside the representable +/// range become `INFINITY` or `-INFINITY` without error. /// /// Unsupported Casts (check with `can_cast_types` before calling): /// * To or from `StructArray` @@ -830,6 +832,7 @@ pub fn cast_with_options( (Map(_, ordered1), Map(_, ordered2)) if ordered1 == ordered2 => { cast_map_values(array.as_map(), to_type, cast_options, ordered1.to_owned()) } + // Decimal to decimal, same width (Decimal128(p1, s1), Decimal128(p2, s2)) => { cast_decimal_to_decimal_same_type::( array.as_primitive(), @@ -850,333 +853,72 @@ pub fn cast_with_options( cast_options, ) } - (Decimal128(_, s1), Decimal256(p2, s2)) => { + // Decimal to decimal, different width + (Decimal128(p1, s1), Decimal256(p2, s2)) => { cast_decimal_to_decimal::( array.as_primitive(), + *p1, *s1, *p2, *s2, cast_options, ) } - (Decimal256(_, s1), Decimal128(p2, s2)) => { + (Decimal256(p1, s1), Decimal128(p2, s2)) => { cast_decimal_to_decimal::( array.as_primitive(), + *p1, *s1, *p2, *s2, cast_options, ) } + // Decimal to non-decimal (Decimal128(_, scale), _) if !to_type.is_temporal() => { - // cast decimal to other type - match to_type { - UInt8 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - UInt16 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - UInt32 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - UInt64 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - Int8 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - Int16 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - Int32 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - Int64 => cast_decimal_to_integer::( - array, - 10_i128, - *scale, - cast_options, - ), - Float32 => cast_decimal_to_float::(array, |x| { - (x as f64 / 10_f64.powi(*scale as i32)) as f32 - }), - Float64 => cast_decimal_to_float::(array, |x| { - x as f64 / 10_f64.powi(*scale as i32) - }), - Utf8View => value_to_string_view(array, cast_options), - Utf8 => value_to_string::(array, cast_options), - LargeUtf8 => value_to_string::(array, cast_options), - Null => Ok(new_null_array(to_type, array.len())), - _ => Err(ArrowError::CastError(format!( - "Casting from {from_type:?} to {to_type:?} not supported" - ))), - } + cast_from_decimal::( + array, + 10_i128, + scale, + from_type, + to_type, + |x: i128| x as f64, + cast_options, + ) } (Decimal256(_, scale), _) if !to_type.is_temporal() => { - // cast decimal to other type - match to_type { - UInt8 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - UInt16 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - UInt32 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - UInt64 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - Int8 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - Int16 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - Int32 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - Int64 => cast_decimal_to_integer::( - array, - i256::from_i128(10_i128), - *scale, - cast_options, - ), - Float32 => cast_decimal_to_float::(array, |x| { - (x.to_f64().unwrap() / 10_f64.powi(*scale as i32)) as f32 - }), - Float64 => cast_decimal_to_float::(array, |x| { - x.to_f64().unwrap() / 10_f64.powi(*scale as i32) - }), - Utf8View => value_to_string_view(array, cast_options), - Utf8 => value_to_string::(array, cast_options), - LargeUtf8 => value_to_string::(array, cast_options), - Null => Ok(new_null_array(to_type, array.len())), - _ => Err(ArrowError::CastError(format!( - "Casting from {from_type:?} to {to_type:?} not supported" - ))), - } + cast_from_decimal::( + array, + i256::from_i128(10_i128), + scale, + from_type, + to_type, + |x: i256| decimal256_to_f64(x), + cast_options, + ) } + // Non-decimal to decimal (_, Decimal128(precision, scale)) if !from_type.is_temporal() => { - // cast data to decimal - match from_type { - UInt8 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - UInt16 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - UInt32 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - UInt64 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - Int8 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - Int16 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - Int32 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - Int64 => cast_integer_to_decimal::<_, Decimal128Type, _>( - array.as_primitive::(), - *precision, - *scale, - 10_i128, - cast_options, - ), - Float32 => cast_floating_point_to_decimal128( - array.as_primitive::(), - *precision, - *scale, - cast_options, - ), - Float64 => cast_floating_point_to_decimal128( - array.as_primitive::(), - *precision, - *scale, - cast_options, - ), - Utf8View | Utf8 => cast_string_to_decimal::( - array, - *precision, - *scale, - cast_options, - ), - LargeUtf8 => cast_string_to_decimal::( - array, - *precision, - *scale, - cast_options, - ), - Null => Ok(new_null_array(to_type, array.len())), - _ => Err(ArrowError::CastError(format!( - "Casting from {from_type:?} to {to_type:?} not supported" - ))), - } + cast_to_decimal::( + array, + 10_i128, + precision, + scale, + from_type, + to_type, + cast_options, + ) } (_, Decimal256(precision, scale)) if !from_type.is_temporal() => { - // cast data to decimal - match from_type { - UInt8 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - UInt16 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - UInt32 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - UInt64 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - Int8 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - Int16 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - Int32 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - Int64 => cast_integer_to_decimal::<_, Decimal256Type, _>( - array.as_primitive::(), - *precision, - *scale, - i256::from_i128(10_i128), - cast_options, - ), - Float32 => cast_floating_point_to_decimal256( - array.as_primitive::(), - *precision, - *scale, - cast_options, - ), - Float64 => cast_floating_point_to_decimal256( - array.as_primitive::(), - *precision, - *scale, - cast_options, - ), - Utf8View | Utf8 => cast_string_to_decimal::( - array, - *precision, - *scale, - cast_options, - ), - LargeUtf8 => cast_string_to_decimal::( - array, - *precision, - *scale, - cast_options, - ), - Null => Ok(new_null_array(to_type, array.len())), - _ => Err(ArrowError::CastError(format!( - "Casting from {from_type:?} to {to_type:?} not supported" - ))), - } + cast_to_decimal::( + array, + i256::from_i128(10_i128), + precision, + scale, + from_type, + to_type, + cast_options, + ) } (Struct(_), Struct(to_fields)) => { let array = array.as_struct(); @@ -1189,12 +931,12 @@ pub fn cast_with_options( let array = StructArray::try_new(to_fields.clone(), fields, array.nulls().cloned())?; Ok(Arc::new(array) as ArrayRef) } - (Struct(_), _) => Err(ArrowError::CastError( - "Cannot cast from struct to other types except struct".to_string(), - )), - (_, Struct(_)) => Err(ArrowError::CastError( - "Cannot cast to struct from other types except struct".to_string(), - )), + (Struct(_), _) => Err(ArrowError::CastError(format!( + "Casting from {from_type:?} to {to_type:?} not supported" + ))), + (_, Struct(_)) => Err(ArrowError::CastError(format!( + "Casting from {from_type:?} to {to_type:?} not supported" + ))), (_, Boolean) => match from_type { UInt8 => cast_numeric_to_bool::(array), UInt16 => cast_numeric_to_bool::(array), @@ -1452,6 +1194,7 @@ pub fn cast_with_options( (FixedSizeBinary(size), _) => match to_type { Binary => cast_fixed_size_binary_to_binary::(array, *size), LargeBinary => cast_fixed_size_binary_to_binary::(array, *size), + BinaryView => cast_fixed_size_binary_to_binary_view(array, *size), _ => Err(ArrowError::CastError(format!( "Casting from {from_type:?} to {to_type:?} not supported", ))), @@ -2068,44 +1811,63 @@ pub fn cast_with_options( })?, )) } - (Date64, Timestamp(TimeUnit::Second, None)) => Ok(Arc::new( - array + (Date64, Timestamp(TimeUnit::Second, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampSecondType>(|x| x / MILLISECONDS), - )), - (Date64, Timestamp(TimeUnit::Millisecond, None)) => { - cast_reinterpret_arrays::(array) + .unary::<_, TimestampSecondType>(|x| x / MILLISECONDS); + + cast_with_options(&array, to_type, cast_options) } - (Date64, Timestamp(TimeUnit::Microsecond, None)) => Ok(Arc::new( - array + (Date64, Timestamp(TimeUnit::Millisecond, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampMicrosecondType>(|x| x * (MICROSECONDS / MILLISECONDS)), - )), - (Date64, Timestamp(TimeUnit::Nanosecond, None)) => Ok(Arc::new( - array + .reinterpret_cast::(); + + cast_with_options(&array, to_type, cast_options) + } + + (Date64, Timestamp(TimeUnit::Microsecond, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampNanosecondType>(|x| x * (NANOSECONDS / MILLISECONDS)), - )), - (Date32, Timestamp(TimeUnit::Second, None)) => Ok(Arc::new( - array + .unary::<_, TimestampMicrosecondType>(|x| x * (MICROSECONDS / MILLISECONDS)); + + cast_with_options(&array, to_type, cast_options) + } + (Date64, Timestamp(TimeUnit::Nanosecond, _)) => { + let array = array + .as_primitive::() + .unary::<_, TimestampNanosecondType>(|x| x * (NANOSECONDS / MILLISECONDS)); + + cast_with_options(&array, to_type, cast_options) + } + (Date32, Timestamp(TimeUnit::Second, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampSecondType>(|x| (x as i64) * SECONDS_IN_DAY), - )), - (Date32, Timestamp(TimeUnit::Millisecond, None)) => Ok(Arc::new( - array + .unary::<_, TimestampSecondType>(|x| (x as i64) * SECONDS_IN_DAY); + + cast_with_options(&array, to_type, cast_options) + } + (Date32, Timestamp(TimeUnit::Millisecond, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampMillisecondType>(|x| (x as i64) * MILLISECONDS_IN_DAY), - )), - (Date32, Timestamp(TimeUnit::Microsecond, None)) => Ok(Arc::new( - array + .unary::<_, TimestampMillisecondType>(|x| (x as i64) * MILLISECONDS_IN_DAY); + + cast_with_options(&array, to_type, cast_options) + } + (Date32, Timestamp(TimeUnit::Microsecond, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampMicrosecondType>(|x| (x as i64) * MICROSECONDS_IN_DAY), - )), - (Date32, Timestamp(TimeUnit::Nanosecond, None)) => Ok(Arc::new( - array + .unary::<_, TimestampMicrosecondType>(|x| (x as i64) * MICROSECONDS_IN_DAY); + + cast_with_options(&array, to_type, cast_options) + } + (Date32, Timestamp(TimeUnit::Nanosecond, _)) => { + let array = array .as_primitive::() - .unary::<_, TimestampNanosecondType>(|x| (x as i64) * NANOSECONDS_IN_DAY), - )), + .unary::<_, TimestampNanosecondType>(|x| (x as i64) * NANOSECONDS_IN_DAY); + + cast_with_options(&array, to_type, cast_options) + } (_, Duration(unit)) if from_type.is_numeric() => { let array = cast_with_options(array, &Int64, cast_options)?; @@ -2192,6 +1954,161 @@ pub fn cast_with_options( } } +fn cast_from_decimal( + array: &dyn Array, + base: D::Native, + scale: &i8, + from_type: &DataType, + to_type: &DataType, + as_float: F, + cast_options: &CastOptions, +) -> Result +where + D: DecimalType + ArrowPrimitiveType, + ::Native: ArrowNativeTypeOp + ToPrimitive, + F: Fn(D::Native) -> f64, +{ + use DataType::*; + // cast decimal to other type + match to_type { + UInt8 => cast_decimal_to_integer::(array, base, *scale, cast_options), + UInt16 => cast_decimal_to_integer::(array, base, *scale, cast_options), + UInt32 => cast_decimal_to_integer::(array, base, *scale, cast_options), + UInt64 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Int8 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Int16 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Int32 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Int64 => cast_decimal_to_integer::(array, base, *scale, cast_options), + Float32 => cast_decimal_to_float::(array, |x| { + (as_float(x) / 10_f64.powi(*scale as i32)) as f32 + }), + Float64 => cast_decimal_to_float::(array, |x| { + as_float(x) / 10_f64.powi(*scale as i32) + }), + Utf8View => value_to_string_view(array, cast_options), + Utf8 => value_to_string::(array, cast_options), + LargeUtf8 => value_to_string::(array, cast_options), + Null => Ok(new_null_array(to_type, array.len())), + _ => Err(ArrowError::CastError(format!( + "Casting from {from_type:?} to {to_type:?} not supported" + ))), + } +} + +/// Convert a [`i256`] to `f64` saturating to infinity on overflow. +fn decimal256_to_f64(v: i256) -> f64 { + v.to_f64().unwrap_or_else(|| { + if v.is_negative() { + f64::NEG_INFINITY + } else { + f64::INFINITY + } + }) +} + +fn cast_to_decimal( + array: &dyn Array, + base: M, + precision: &u8, + scale: &i8, + from_type: &DataType, + to_type: &DataType, + cast_options: &CastOptions, +) -> Result +where + D: DecimalType + ArrowPrimitiveType, + M: ArrowNativeTypeOp + DecimalCast, + u8: num::traits::AsPrimitive, + u16: num::traits::AsPrimitive, + u32: num::traits::AsPrimitive, + u64: num::traits::AsPrimitive, + i8: num::traits::AsPrimitive, + i16: num::traits::AsPrimitive, + i32: num::traits::AsPrimitive, + i64: num::traits::AsPrimitive, +{ + use DataType::*; + // cast data to decimal + match from_type { + UInt8 => cast_integer_to_decimal::<_, D, M>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + UInt16 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + UInt32 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + UInt64 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + Int8 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + Int16 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + Int32 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + Int64 => cast_integer_to_decimal::<_, D, _>( + array.as_primitive::(), + *precision, + *scale, + base, + cast_options, + ), + Float32 => cast_floating_point_to_decimal::<_, D>( + array.as_primitive::(), + *precision, + *scale, + cast_options, + ), + Float64 => cast_floating_point_to_decimal::<_, D>( + array.as_primitive::(), + *precision, + *scale, + cast_options, + ), + Utf8View | Utf8 => { + cast_string_to_decimal::(array, *precision, *scale, cast_options) + } + LargeUtf8 => cast_string_to_decimal::(array, *precision, *scale, cast_options), + Null => Ok(new_null_array(to_type, array.len())), + _ => Err(ArrowError::CastError(format!( + "Casting from {from_type:?} to {to_type:?} not supported" + ))), + } +} + /// Get the time unit as a multiple of a second const fn time_unit_multiple(unit: &TimeUnit) -> i64 { match unit { @@ -2263,7 +2180,7 @@ fn cast_numeric_to_binary( ) -> Result { let array = array.as_primitive::(); let size = std::mem::size_of::(); - let offsets = OffsetBuffer::from_lengths(std::iter::repeat(size).take(array.len())); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(size, array.len())); Ok(Arc::new(GenericBinaryArray::::new( offsets, array.values().inner().clone(), @@ -2424,6 +2341,27 @@ fn cast_fixed_size_binary_to_binary( Ok(Arc::new(builder.finish())) } +fn cast_fixed_size_binary_to_binary_view( + array: &dyn Array, + _byte_width: i32, +) -> Result { + let array = array + .as_any() + .downcast_ref::() + .unwrap(); + + let mut builder = BinaryViewBuilder::with_capacity(array.len()); + for i in 0..array.len() { + if array.is_null(i) { + builder.append_null(); + } else { + builder.append_value(array.value(i)); + } + } + + Ok(Arc::new(builder.finish())) +} + /// Helper function to cast from one `ByteArrayType` to another and vice versa. /// If the target one (e.g., `LargeUtf8`) is too large for the source array it will return an Error. fn cast_byte_container(array: &dyn Array) -> Result @@ -2505,6 +2443,19 @@ mod tests { use chrono::NaiveDate; use half::f16; + #[derive(Clone)] + struct DecimalCastTestConfig { + input_prec: u8, + input_scale: i8, + input_repr: i128, + output_prec: u8, + output_scale: i8, + expected_output_repr: Result, // the error variant can contain a string + // template where the "{}" will be + // replaced with the decimal type name + // (e.g. Decimal128) + } + macro_rules! generate_cast_test_case { ($INPUT_ARRAY: expr, $OUTPUT_TYPE_ARRAY: ident, $OUTPUT_TYPE: expr, $OUTPUT_VALUES: expr) => { let output = @@ -2527,7 +2478,49 @@ mod tests { }; } - fn create_decimal_array( + fn run_decimal_cast_test_case(t: DecimalCastTestConfig) + where + I: DecimalType, + O: DecimalType, + I::Native: DecimalCast, + O::Native: DecimalCast, + { + let array = vec![I::Native::from_decimal(t.input_repr)]; + let array = array + .into_iter() + .collect::>() + .with_precision_and_scale(t.input_prec, t.input_scale) + .unwrap(); + let input_type = array.data_type(); + let output_type = O::TYPE_CONSTRUCTOR(t.output_prec, t.output_scale); + assert!(can_cast_types(input_type, &output_type)); + + let options = CastOptions { + safe: false, + ..Default::default() + }; + let result = cast_with_options(&array, &output_type, &options); + + match t.expected_output_repr { + Ok(v) => { + let expected_array = vec![O::Native::from_decimal(v)]; + let expected_array = expected_array + .into_iter() + .collect::>() + .with_precision_and_scale(t.output_prec, t.output_scale) + .unwrap(); + assert_eq!(*result.unwrap(), expected_array); + } + Err(expected_output_message_template) => { + assert!(result.is_err()); + let expected_error_message = + expected_output_message_template.replace("{}", O::PREFIX); + assert_eq!(result.unwrap_err().to_string(), expected_error_message); + } + } + } + + fn create_decimal128_array( array: Vec>, precision: u8, scale: i8, @@ -2596,7 +2589,7 @@ mod tests { Some(-3123456), None, ]; - let array = create_decimal_array(array, 20, 4).unwrap(); + let array = create_decimal128_array(array, 20, 4).unwrap(); // decimal128 to decimal128 let input_type = DataType::Decimal128(20, 4); let output_type = DataType::Decimal128(20, 3); @@ -2681,7 +2674,7 @@ mod tests { let output_type = DataType::Decimal128(20, 4); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(1123456), Some(2123456), Some(3123456), None]; - let array = create_decimal_array(array, 20, 3).unwrap(); + let array = create_decimal128_array(array, 20, 3).unwrap(); generate_cast_test_case!( &array, Decimal128Array, @@ -2695,7 +2688,7 @@ mod tests { ); // negative test let array = vec![Some(123456), None]; - let array = create_decimal_array(array, 10, 0).unwrap(); + let array = create_decimal128_array(array, 10, 0).unwrap(); let result_safe = cast(&array, &DataType::Decimal128(2, 2)); assert!(result_safe.is_ok()); let options = CastOptions { @@ -2719,7 +2712,7 @@ mod tests { ); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(1123456), Some(2123456), Some(3123456), None]; - let array = create_decimal_array(array, p, s).unwrap(); + let array = create_decimal128_array(array, p, s).unwrap(); let cast_array = cast_with_options(&array, &output_type, &CastOptions::default()).unwrap(); assert_eq!(cast_array.data_type(), &output_type); } @@ -2735,7 +2728,7 @@ mod tests { ); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(1123456), Some(2123456), Some(3123456), None]; - let array = create_decimal_array(array, p, s).unwrap(); + let array = create_decimal128_array(array, p, s).unwrap(); let cast_array = cast_with_options(&array, &output_type, &CastOptions::default()).unwrap(); assert_eq!(cast_array.data_type(), &output_type); } @@ -2747,7 +2740,7 @@ mod tests { assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(i128::MAX)]; - let array = create_decimal_array(array, 38, 3).unwrap(); + let array = create_decimal128_array(array, 38, 3).unwrap(); let result = cast_with_options( &array, &output_type, @@ -2767,7 +2760,7 @@ mod tests { assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(i128::MAX)]; - let array = create_decimal_array(array, 38, 3).unwrap(); + let array = create_decimal128_array(array, 38, 3).unwrap(); let result = cast_with_options( &array, &output_type, @@ -2786,7 +2779,7 @@ mod tests { let output_type = DataType::Decimal256(20, 4); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(1123456), Some(2123456), Some(3123456), None]; - let array = create_decimal_array(array, 20, 3).unwrap(); + let array = create_decimal128_array(array, 20, 3).unwrap(); generate_cast_test_case!( &array, Decimal256Array, @@ -2888,69 +2881,69 @@ mod tests { ); } - #[test] - fn test_cast_decimal_to_numeric() { - let value_array: Vec> = vec![Some(125), Some(225), Some(325), None, Some(525)]; - let array = create_decimal_array(value_array, 38, 2).unwrap(); + fn generate_decimal_to_numeric_cast_test_case(array: &PrimitiveArray) + where + T: ArrowPrimitiveType + DecimalType, + { // u8 generate_cast_test_case!( - &array, + array, UInt8Array, &DataType::UInt8, vec![Some(1_u8), Some(2_u8), Some(3_u8), None, Some(5_u8)] ); // u16 generate_cast_test_case!( - &array, + array, UInt16Array, &DataType::UInt16, vec![Some(1_u16), Some(2_u16), Some(3_u16), None, Some(5_u16)] ); // u32 generate_cast_test_case!( - &array, + array, UInt32Array, &DataType::UInt32, vec![Some(1_u32), Some(2_u32), Some(3_u32), None, Some(5_u32)] ); // u64 generate_cast_test_case!( - &array, + array, UInt64Array, &DataType::UInt64, vec![Some(1_u64), Some(2_u64), Some(3_u64), None, Some(5_u64)] ); // i8 generate_cast_test_case!( - &array, + array, Int8Array, &DataType::Int8, vec![Some(1_i8), Some(2_i8), Some(3_i8), None, Some(5_i8)] ); // i16 generate_cast_test_case!( - &array, + array, Int16Array, &DataType::Int16, vec![Some(1_i16), Some(2_i16), Some(3_i16), None, Some(5_i16)] ); // i32 generate_cast_test_case!( - &array, + array, Int32Array, &DataType::Int32, vec![Some(1_i32), Some(2_i32), Some(3_i32), None, Some(5_i32)] ); // i64 generate_cast_test_case!( - &array, + array, Int64Array, &DataType::Int64, vec![Some(1_i64), Some(2_i64), Some(3_i64), None, Some(5_i64)] ); // f32 generate_cast_test_case!( - &array, + array, Float32Array, &DataType::Float32, vec![ @@ -2963,7 +2956,7 @@ mod tests { ); // f64 generate_cast_test_case!( - &array, + array, Float64Array, &DataType::Float64, vec![ @@ -2974,10 +2967,18 @@ mod tests { Some(5.25_f64) ] ); + } + + #[test] + fn test_cast_decimal128_to_numeric() { + let value_array: Vec> = vec![Some(125), Some(225), Some(325), None, Some(525)]; + let array = create_decimal128_array(value_array, 38, 2).unwrap(); + + generate_decimal_to_numeric_cast_test_case(&array); // overflow test: out of range of max u8 let value_array: Vec> = vec![Some(51300)]; - let array = create_decimal_array(value_array, 38, 2).unwrap(); + let array = create_decimal128_array(value_array, 38, 2).unwrap(); let casted_array = cast_with_options( &array, &DataType::UInt8, @@ -3004,7 +3005,7 @@ mod tests { // overflow test: out of range of max i8 let value_array: Vec> = vec![Some(24400)]; - let array = create_decimal_array(value_array, 38, 2).unwrap(); + let array = create_decimal128_array(value_array, 38, 2).unwrap(); let casted_array = cast_with_options( &array, &DataType::Int8, @@ -3041,7 +3042,7 @@ mod tests { Some(112345678), Some(112345679), ]; - let array = create_decimal_array(value_array, 38, 2).unwrap(); + let array = create_decimal128_array(value_array, 38, 2).unwrap(); generate_cast_test_case!( &array, Float32Array, @@ -3068,7 +3069,7 @@ mod tests { Some(112345678901234568), Some(112345678901234560), ]; - let array = create_decimal_array(value_array, 38, 2).unwrap(); + let array = create_decimal128_array(value_array, 38, 2).unwrap(); generate_cast_test_case!( &array, Float64Array, @@ -3767,7 +3768,6 @@ mod tests { Arc::new(StringViewArray::from(vec![Some("1.5"), Some("2.5"), None])); for array in inputs { - println!("type: {}", array.data_type()); assert!(can_cast_types(array.data_type(), &DataType::Utf8View)); let arr = cast(&array, &DataType::Utf8View).unwrap(); assert_eq!(expected.as_ref(), arr.as_ref()); @@ -4340,6 +4340,48 @@ mod tests { } } + #[test] + fn test_cast_string_with_large_date_to_date32() { + let array = Arc::new(StringArray::from(vec![ + Some("+10999-12-31"), + Some("-0010-02-28"), + Some("0010-02-28"), + Some("0000-01-01"), + Some("-0000-01-01"), + Some("-0001-01-01"), + ])) as ArrayRef; + let to_type = DataType::Date32; + let options = CastOptions { + safe: false, + format_options: FormatOptions::default(), + }; + let b = cast_with_options(&array, &to_type, &options).unwrap(); + let c = b.as_primitive::(); + assert_eq!(3298139, c.value(0)); // 10999-12-31 + assert_eq!(-723122, c.value(1)); // -0010-02-28 + assert_eq!(-715817, c.value(2)); // 0010-02-28 + assert_eq!(c.value(3), c.value(4)); // Expect 0000-01-01 and -0000-01-01 to be parsed the same + assert_eq!(-719528, c.value(3)); // 0000-01-01 + assert_eq!(-719528, c.value(4)); // -0000-01-01 + assert_eq!(-719893, c.value(5)); // -0001-01-01 + } + + #[test] + fn test_cast_invalid_string_with_large_date_to_date32() { + // Large dates need to be prefixed with a + or - sign, otherwise they are not parsed correctly + let array = Arc::new(StringArray::from(vec![Some("10999-12-31")])) as ArrayRef; + let to_type = DataType::Date32; + let options = CastOptions { + safe: false, + format_options: FormatOptions::default(), + }; + let err = cast_with_options(&array, &to_type, &options).unwrap_err(); + assert_eq!( + err.to_string(), + "Cast error: Cannot cast string '10999-12-31' to value of Date32 type" + ); + } + #[test] fn test_cast_string_format_yyyymmdd_to_date32() { let a0 = Arc::new(StringViewArray::from(vec![ @@ -4840,6 +4882,73 @@ mod tests { assert_eq!(bytes_1, down_cast.value(0)); assert_eq!(bytes_2, down_cast.value(1)); assert!(down_cast.is_null(2)); + + let array_ref = cast(&a1, &DataType::BinaryView).unwrap(); + let down_cast = array_ref.as_binary_view(); + assert_eq!(bytes_1, down_cast.value(0)); + assert_eq!(bytes_2, down_cast.value(1)); + assert!(down_cast.is_null(2)); + } + + #[test] + fn test_fixed_size_binary_to_dictionary() { + let bytes_1 = "Hiiii".as_bytes(); + let bytes_2 = "Hello".as_bytes(); + + let binary_data = vec![Some(bytes_1), Some(bytes_2), Some(bytes_1), None]; + let a1 = Arc::new(FixedSizeBinaryArray::from(binary_data.clone())) as ArrayRef; + + let cast_type = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::FixedSizeBinary(5)), + ); + let cast_array = cast(&a1, &cast_type).unwrap(); + assert_eq!(cast_array.data_type(), &cast_type); + assert_eq!( + array_to_strings(&cast_array), + vec!["4869696969", "48656c6c6f", "4869696969", "null"] + ); + // dictionary should only have two distinct values + let dict_array = cast_array + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(dict_array.values().len(), 2); + } + + #[test] + fn test_binary_to_dictionary() { + let mut builder = GenericBinaryBuilder::::new(); + builder.append_value(b"hello"); + builder.append_value(b"hiiii"); + builder.append_value(b"hiiii"); // duplicate + builder.append_null(); + builder.append_value(b"rustt"); + + let a1 = builder.finish(); + + let cast_type = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::FixedSizeBinary(5)), + ); + let cast_array = cast(&a1, &cast_type).unwrap(); + assert_eq!(cast_array.data_type(), &cast_type); + assert_eq!( + array_to_strings(&cast_array), + vec![ + "68656c6c6f", + "6869696969", + "6869696969", + "null", + "7275737474" + ] + ); + // dictionary should only have three distinct values + let dict_array = cast_array + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(dict_array.values().len(), 3); } #[test] @@ -5286,6 +5395,197 @@ mod tests { }}; } + #[test] + fn test_cast_date32_to_timestamp_and_timestamp_with_timezone() { + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let a = Date32Array::from(vec![Some(18628), None, None]); // 2021-1-1, 2022-1-1 + let array = Arc::new(a) as ArrayRef; + + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Second, Some(tz.into())), + ) + .unwrap(); + let c = b.as_primitive::(); + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00+05:45", result.value(0)); + + let b = cast(&array, &DataType::Timestamp(TimeUnit::Second, None)).unwrap(); + let c = b.as_primitive::(); + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00", result.value(0)); + } + + #[test] + fn test_cast_date32_to_timestamp_with_timezone() { + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1 + let array = Arc::new(a) as ArrayRef; + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Second, Some(tz.into())), + ) + .unwrap(); + let c = b.as_primitive::(); + assert_eq!(1609438500, c.value(0)); + assert_eq!(1640974500, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00+05:45", result.value(0)); + assert_eq!("2022-01-01T00:00:00+05:45", result.value(1)); + } + + #[test] + fn test_cast_date32_to_timestamp_with_timezone_ms() { + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1 + let array = Arc::new(a) as ArrayRef; + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Millisecond, Some(tz.into())), + ) + .unwrap(); + let c = b.as_primitive::(); + assert_eq!(1609438500000, c.value(0)); + assert_eq!(1640974500000, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00+05:45", result.value(0)); + assert_eq!("2022-01-01T00:00:00+05:45", result.value(1)); + } + + #[test] + fn test_cast_date32_to_timestamp_with_timezone_us() { + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1 + let array = Arc::new(a) as ArrayRef; + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Microsecond, Some(tz.into())), + ) + .unwrap(); + let c = b.as_primitive::(); + assert_eq!(1609438500000000, c.value(0)); + assert_eq!(1640974500000000, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00+05:45", result.value(0)); + assert_eq!("2022-01-01T00:00:00+05:45", result.value(1)); + } + + #[test] + fn test_cast_date32_to_timestamp_with_timezone_ns() { + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let a = Date32Array::from(vec![Some(18628), Some(18993), None]); // 2021-1-1, 2022-1-1 + let array = Arc::new(a) as ArrayRef; + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Nanosecond, Some(tz.into())), + ) + .unwrap(); + let c = b.as_primitive::(); + assert_eq!(1609438500000000000, c.value(0)); + assert_eq!(1640974500000000000, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("2021-01-01T00:00:00+05:45", result.value(0)); + assert_eq!("2022-01-01T00:00:00+05:45", result.value(1)); + } + + #[test] + fn test_cast_date64_to_timestamp_with_timezone() { + let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]); + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Second, Some(tz.into())), + ) + .unwrap(); + + let c = b.as_primitive::(); + assert_eq!(863979300, c.value(0)); + assert_eq!(1545675300, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("1997-05-19T00:00:00+05:45", result.value(0)); + assert_eq!("2018-12-25T00:00:00+05:45", result.value(1)); + } + + #[test] + fn test_cast_date64_to_timestamp_with_timezone_ms() { + let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]); + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Millisecond, Some(tz.into())), + ) + .unwrap(); + + let c = b.as_primitive::(); + assert_eq!(863979300005, c.value(0)); + assert_eq!(1545675300001, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0)); + assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1)); + } + + #[test] + fn test_cast_date64_to_timestamp_with_timezone_us() { + let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]); + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Microsecond, Some(tz.into())), + ) + .unwrap(); + + let c = b.as_primitive::(); + assert_eq!(863979300005000, c.value(0)); + assert_eq!(1545675300001000, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0)); + assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1)); + } + + #[test] + fn test_cast_date64_to_timestamp_with_timezone_ns() { + let array = Date64Array::from(vec![Some(864000000005), Some(1545696000001), None]); + let tz = "+0545"; // UTC + 0545 is Asia/Kathmandu + let b = cast( + &array, + &DataType::Timestamp(TimeUnit::Nanosecond, Some(tz.into())), + ) + .unwrap(); + + let c = b.as_primitive::(); + assert_eq!(863979300005000000, c.value(0)); + assert_eq!(1545675300001000000, c.value(1)); + assert!(c.is_null(2)); + + let string_array = cast(&c, &DataType::Utf8).unwrap(); + let result = string_array.as_string::(); + assert_eq!("1997-05-19T00:00:00.005+05:45", result.value(0)); + assert_eq!("2018-12-25T00:00:00.001+05:45", result.value(1)); + } + #[test] fn test_cast_timestamp_to_strings() { // "2018-12-25T00:00:02.001", "1997-05-19T00:00:03.005", None @@ -7408,13 +7708,11 @@ mod tests { ); let list_array = cast(&array, expected.data_type()) - .unwrap_or_else(|_| panic!("Failed to cast {:?} to {:?}", array, expected)); + .unwrap_or_else(|_| panic!("Failed to cast {array:?} to {expected:?}")); assert_eq!( list_array.as_ref(), &expected, - "Incorrect result from casting {:?} to {:?}", - array, - expected + "Incorrect result from casting {array:?} to {expected:?}", ); } } @@ -7648,7 +7946,7 @@ mod tests { }, ); assert!(res.is_err()); - assert!(format!("{:?}", res) + assert!(format!("{res:?}") .contains("Cannot cast to FixedSizeList(3): value at index 1 has length 2")); // When safe=true (default), the cast will fill nulls for lists that are @@ -7739,7 +8037,7 @@ mod tests { }, ); assert!(res.is_err()); - assert!(format!("{:?}", res).contains("Can't cast value 2147483647 to type Int16")); + assert!(format!("{res:?}").contains("Can't cast value 2147483647 to type Int16")); } #[test] @@ -8375,6 +8673,28 @@ mod tests { "did not find expected error '{expected_error}' in actual error '{err}'" ); } + #[test] + fn test_cast_decimal256_to_f64_overflow() { + // Test positive overflow (positive infinity) + let array = vec![Some(i256::MAX)]; + let array = create_decimal256_array(array, 76, 2).unwrap(); + let array = Arc::new(array) as ArrayRef; + + let result = cast(&array, &DataType::Float64).unwrap(); + let result = result.as_primitive::(); + assert!(result.value(0).is_infinite()); + assert!(result.value(0) > 0.0); // Positive infinity + + // Test negative overflow (negative infinity) + let array = vec![Some(i256::MIN)]; + let array = create_decimal256_array(array, 76, 2).unwrap(); + let array = Arc::new(array) as ArrayRef; + + let result = cast(&array, &DataType::Float64).unwrap(); + let result = result.as_primitive::(); + assert!(result.value(0).is_infinite()); + assert!(result.value(0) < 0.0); // Negative infinity + } #[test] fn test_cast_decimal128_to_decimal128_negative_scale() { @@ -8382,7 +8702,7 @@ mod tests { let output_type = DataType::Decimal128(20, -1); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(1123450), Some(2123455), Some(3123456), None]; - let input_decimal_array = create_decimal_array(array, 20, 0).unwrap(); + let input_decimal_array = create_decimal128_array(array, 20, 0).unwrap(); let array = Arc::new(input_decimal_array) as ArrayRef; generate_cast_test_case!( &array, @@ -8440,7 +8760,7 @@ mod tests { let output_type = DataType::Decimal128(10, -2); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(123)]; - let input_decimal_array = create_decimal_array(array, 10, -1).unwrap(); + let input_decimal_array = create_decimal128_array(array, 10, -1).unwrap(); let array = Arc::new(input_decimal_array) as ArrayRef; generate_cast_test_case!(&array, Decimal128Array, &output_type, vec![Some(12_i128),]); @@ -8450,7 +8770,7 @@ mod tests { assert_eq!("1200", decimal_arr.value_as_string(0)); let array = vec![Some(125)]; - let input_decimal_array = create_decimal_array(array, 10, -1).unwrap(); + let input_decimal_array = create_decimal128_array(array, 10, -1).unwrap(); let array = Arc::new(input_decimal_array) as ArrayRef; generate_cast_test_case!(&array, Decimal128Array, &output_type, vec![Some(13_i128),]); @@ -8466,7 +8786,7 @@ mod tests { let output_type = DataType::Decimal256(10, 5); assert!(can_cast_types(&input_type, &output_type)); let array = vec![Some(123456), Some(-123456)]; - let input_decimal_array = create_decimal_array(array, 10, 3).unwrap(); + let input_decimal_array = create_decimal128_array(array, 10, 3).unwrap(); let array = Arc::new(input_decimal_array) as ArrayRef; let hundred = i256::from_i128(100); @@ -8803,7 +9123,7 @@ mod tests { Some(array.value_as_string(i)) }; let actual = actual.as_ref().map(|s| s.as_ref()); - assert_eq!(*expected, actual, "Expected at position {}", i); + assert_eq!(*expected, actual, "Expected at position {i}"); } } @@ -9288,15 +9608,15 @@ mod tests { test_decimal_to_string::( DataType::Utf8View, - create_decimal_array(array128.clone(), 7, 3).unwrap(), + create_decimal128_array(array128.clone(), 7, 3).unwrap(), ); test_decimal_to_string::( DataType::Utf8, - create_decimal_array(array128.clone(), 7, 3).unwrap(), + create_decimal128_array(array128.clone(), 7, 3).unwrap(), ); test_decimal_to_string::( DataType::LargeUtf8, - create_decimal_array(array128, 7, 3).unwrap(), + create_decimal128_array(array128, 7, 3).unwrap(), ); test_decimal_to_string::( @@ -9941,11 +10261,333 @@ mod tests { ); } + #[test] + fn test_cast_struct_to_non_struct() { + let boolean = Arc::new(BooleanArray::from(vec![true, false])); + let struct_array = StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Boolean, false)), + boolean.clone() as ArrayRef, + )]); + let to_type = DataType::Utf8; + let result = cast(&struct_array, &to_type); + assert_eq!( + r#"Cast error: Casting from Struct([Field { name: "a", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]) to Utf8 not supported"#, + result.unwrap_err().to_string() + ); + } + + #[test] + fn test_cast_non_struct_to_struct() { + let array = StringArray::from(vec!["a", "b"]); + let to_type = DataType::Struct(vec![Field::new("a", DataType::Boolean, false)].into()); + let result = cast(&array, &to_type); + assert_eq!( + r#"Cast error: Casting from Utf8 to Struct([Field { name: "a", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]) not supported"#, + result.unwrap_err().to_string() + ); + } + + fn run_decimal_cast_test_case_between_multiple_types(t: DecimalCastTestConfig) { + run_decimal_cast_test_case::(t.clone()); + run_decimal_cast_test_case::(t.clone()); + run_decimal_cast_test_case::(t.clone()); + run_decimal_cast_test_case::(t.clone()); + } + + #[test] + fn test_decimal_to_decimal_coverage() { + let test_cases = [ + // increase precision, increase scale, infallible + DecimalCastTestConfig { + input_prec: 5, + input_scale: 1, + input_repr: 99999, // 9999.9 + output_prec: 10, + output_scale: 6, + expected_output_repr: Ok(9999900000), // 9999.900000 + }, + // increase precision, increase scale, fallible, safe + DecimalCastTestConfig { + input_prec: 5, + input_scale: 1, + input_repr: 99, // 9999.9 + output_prec: 7, + output_scale: 6, + expected_output_repr: Ok(9900000), // 9.900000 + }, + // increase precision, increase scale, fallible, unsafe + DecimalCastTestConfig { + input_prec: 5, + input_scale: 1, + input_repr: 99999, // 9999.9 + output_prec: 7, + output_scale: 6, + expected_output_repr: Err("Invalid argument error: 9999900000 is too large to store in a {} of precision 7. Max is 9999999".to_string()) // max is 9.999999 + }, + // increase precision, decrease scale, always infallible + DecimalCastTestConfig { + input_prec: 5, + input_scale: 3, + input_repr: 99999, // 99.999 + output_prec: 10, + output_scale: 2, + expected_output_repr: Ok(10000), // 100.00 + }, + // increase precision, decrease scale, no rouding + DecimalCastTestConfig { + input_prec: 5, + input_scale: 3, + input_repr: 99994, // 99.994 + output_prec: 10, + output_scale: 2, + expected_output_repr: Ok(9999), // 99.99 + }, + // increase precision, don't change scale, always infallible + DecimalCastTestConfig { + input_prec: 5, + input_scale: 3, + input_repr: 99999, // 99.999 + output_prec: 10, + output_scale: 3, + expected_output_repr: Ok(99999), // 99.999 + }, + // decrease precision, increase scale, safe + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 999999, // 9.99999 + output_prec: 8, + output_scale: 7, + expected_output_repr: Ok(99999900), // 9.9999900 + }, + // decrease precision, increase scale, unsafe + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 9999999, // 99.99999 + output_prec: 8, + output_scale: 7, + expected_output_repr: Err("Invalid argument error: 999999900 is too large to store in a {} of precision 8. Max is 99999999".to_string()) // max is 9.9999999 + }, + // decrease precision, decrease scale, safe, infallible + DecimalCastTestConfig { + input_prec: 7, + input_scale: 4, + input_repr: 9999999, // 999.9999 + output_prec: 6, + output_scale: 2, + expected_output_repr: Ok(100000), + }, + // decrease precision, decrease scale, safe, fallible + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 12345678, // 123.45678 + output_prec: 8, + output_scale: 3, + expected_output_repr: Ok(123457), // 123.457 + }, + // decrease precision, decrease scale, unsafe + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 9999999, // 99.99999 + output_prec: 4, + output_scale: 3, + expected_output_repr: Err("Invalid argument error: 100000 is too large to store in a {} of precision 4. Max is 9999".to_string()) // max is 9.999 + }, + // decrease precision, same scale, safe + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 999999, // 9.99999 + output_prec: 6, + output_scale: 5, + expected_output_repr: Ok(999999), // 9.99999 + }, + // decrease precision, same scale, unsafe + DecimalCastTestConfig { + input_prec: 10, + input_scale: 5, + input_repr: 9999999, // 99.99999 + output_prec: 6, + output_scale: 5, + expected_output_repr: Err("Invalid argument error: 9999999 is too large to store in a {} of precision 6. Max is 999999".to_string()) // max is 9.99999 + }, + // same precision, increase scale, safe + DecimalCastTestConfig { + input_prec: 7, + input_scale: 4, + input_repr: 12345, // 1.2345 + output_prec: 7, + output_scale: 6, + expected_output_repr: Ok(1234500), // 1.234500 + }, + // same precision, increase scale, unsafe + DecimalCastTestConfig { + input_prec: 7, + input_scale: 4, + input_repr: 123456, // 12.3456 + output_prec: 7, + output_scale: 6, + expected_output_repr: Err("Invalid argument error: 12345600 is too large to store in a {} of precision 7. Max is 9999999".to_string()) // max is 9.99999 + }, + // same precision, decrease scale, infallible + DecimalCastTestConfig { + input_prec: 7, + input_scale: 5, + input_repr: 1234567, // 12.34567 + output_prec: 7, + output_scale: 4, + expected_output_repr: Ok(123457), // 12.3457 + }, + // same precision, same scale, infallible + DecimalCastTestConfig { + input_prec: 7, + input_scale: 5, + input_repr: 9999999, // 99.99999 + output_prec: 7, + output_scale: 5, + expected_output_repr: Ok(9999999), // 99.99999 + }, + // precision increase, input scale & output scale = 0, infallible + DecimalCastTestConfig { + input_prec: 7, + input_scale: 0, + input_repr: 1234567, // 1234567 + output_prec: 8, + output_scale: 0, + expected_output_repr: Ok(1234567), // 1234567 + }, + // precision decrease, input scale & output scale = 0, failure + DecimalCastTestConfig { + input_prec: 7, + input_scale: 0, + input_repr: 1234567, // 1234567 + output_prec: 6, + output_scale: 0, + expected_output_repr: Err("Invalid argument error: 1234567 is too large to store in a {} of precision 6. Max is 999999".to_string()) + }, + // precision decrease, input scale & output scale = 0, success + DecimalCastTestConfig { + input_prec: 7, + input_scale: 0, + input_repr: 123456, // 123456 + output_prec: 6, + output_scale: 0, + expected_output_repr: Ok(123456), // 123456 + }, + ]; + + for t in test_cases { + run_decimal_cast_test_case_between_multiple_types(t); + } + } + + #[test] + fn test_decimal_to_decimal_increase_scale_and_precision_unchecked() { + let test_cases = [ + DecimalCastTestConfig { + input_prec: 5, + input_scale: 0, + input_repr: 99999, + output_prec: 10, + output_scale: 5, + expected_output_repr: Ok(9999900000), + }, + DecimalCastTestConfig { + input_prec: 5, + input_scale: 0, + input_repr: -99999, + output_prec: 10, + output_scale: 5, + expected_output_repr: Ok(-9999900000), + }, + DecimalCastTestConfig { + input_prec: 5, + input_scale: 2, + input_repr: 99999, + output_prec: 10, + output_scale: 5, + expected_output_repr: Ok(99999000), + }, + DecimalCastTestConfig { + input_prec: 5, + input_scale: -2, + input_repr: -99999, + output_prec: 10, + output_scale: 3, + expected_output_repr: Ok(-9999900000), + }, + DecimalCastTestConfig { + input_prec: 5, + input_scale: 3, + input_repr: -12345, + output_prec: 6, + output_scale: 5, + expected_output_repr: Err("Invalid argument error: -1234500 is too small to store in a {} of precision 6. Min is -999999".to_string()) + }, + ]; + + for t in test_cases { + run_decimal_cast_test_case_between_multiple_types(t); + } + } + + #[test] + fn test_decimal_to_decimal_decrease_scale_and_precision_unchecked() { + let test_cases = [ + DecimalCastTestConfig { + input_prec: 5, + input_scale: 0, + input_repr: 99999, + output_scale: -3, + output_prec: 3, + expected_output_repr: Ok(100), + }, + DecimalCastTestConfig { + input_prec: 5, + input_scale: 0, + input_repr: -99999, + output_prec: 1, + output_scale: -5, + expected_output_repr: Ok(-1), + }, + DecimalCastTestConfig { + input_prec: 10, + input_scale: 2, + input_repr: 123456789, + output_prec: 5, + output_scale: -2, + expected_output_repr: Ok(12346), + }, + DecimalCastTestConfig { + input_prec: 10, + input_scale: 4, + input_repr: -9876543210, + output_prec: 7, + output_scale: 0, + expected_output_repr: Ok(-987654), + }, + DecimalCastTestConfig { + input_prec: 7, + input_scale: 4, + input_repr: 9999999, + output_prec: 6, + output_scale: 3, + expected_output_repr: + Err("Invalid argument error: 1000000 is too large to store in a {} of precision 6. Max is 999999".to_string()), + }, + ]; + for t in test_cases { + run_decimal_cast_test_case_between_multiple_types(t); + } + } + #[test] fn test_decimal_to_decimal_throw_error_on_precision_overflow_same_scale() { let array = vec![Some(123456789)]; - let array = create_decimal_array(array, 24, 2).unwrap(); - println!("{:?}", array); + let array = create_decimal128_array(array, 24, 2).unwrap(); let input_type = DataType::Decimal128(24, 2); let output_type = DataType::Decimal128(6, 2); assert!(can_cast_types(&input_type, &output_type)); @@ -9956,14 +10598,42 @@ mod tests { }; let result = cast_with_options(&array, &output_type, &options); assert_eq!(result.unwrap_err().to_string(), - "Invalid argument error: 123456790 is too large to store in a Decimal128 of precision 6. Max is 999999"); + "Invalid argument error: 123456789 is too large to store in a Decimal128 of precision 6. Max is 999999"); + } + + #[test] + fn test_decimal_to_decimal_same_scale() { + let array = vec![Some(520)]; + let array = create_decimal128_array(array, 4, 2).unwrap(); + let input_type = DataType::Decimal128(4, 2); + let output_type = DataType::Decimal128(3, 2); + assert!(can_cast_types(&input_type, &output_type)); + + let options = CastOptions { + safe: false, + ..Default::default() + }; + let result = cast_with_options(&array, &output_type, &options); + assert_eq!( + result.unwrap().as_primitive::().value(0), + 520 + ); + + // Cast 0 of decimal(3, 0) type to decimal(2, 0) + assert_eq!( + &cast( + &create_decimal128_array(vec![Some(0)], 3, 0).unwrap(), + &DataType::Decimal128(2, 0) + ) + .unwrap(), + &(Arc::new(create_decimal128_array(vec![Some(0)], 2, 0).unwrap()) as ArrayRef) + ); } #[test] fn test_decimal_to_decimal_throw_error_on_precision_overflow_lower_scale() { let array = vec![Some(123456789)]; - let array = create_decimal_array(array, 24, 2).unwrap(); - println!("{:?}", array); + let array = create_decimal128_array(array, 24, 4).unwrap(); let input_type = DataType::Decimal128(24, 4); let output_type = DataType::Decimal128(6, 2); assert!(can_cast_types(&input_type, &output_type)); @@ -9974,14 +10644,13 @@ mod tests { }; let result = cast_with_options(&array, &output_type, &options); assert_eq!(result.unwrap_err().to_string(), - "Invalid argument error: 123456790 is too large to store in a Decimal128 of precision 6. Max is 999999"); + "Invalid argument error: 1234568 is too large to store in a Decimal128 of precision 6. Max is 999999"); } #[test] fn test_decimal_to_decimal_throw_error_on_precision_overflow_greater_scale() { let array = vec![Some(123456789)]; - let array = create_decimal_array(array, 24, 2).unwrap(); - println!("{:?}", array); + let array = create_decimal128_array(array, 24, 2).unwrap(); let input_type = DataType::Decimal128(24, 2); let output_type = DataType::Decimal128(6, 3); assert!(can_cast_types(&input_type, &output_type)); @@ -9998,8 +10667,7 @@ mod tests { #[test] fn test_decimal_to_decimal_throw_error_on_precision_overflow_diff_type() { let array = vec![Some(123456789)]; - let array = create_decimal_array(array, 24, 2).unwrap(); - println!("{:?}", array); + let array = create_decimal128_array(array, 24, 2).unwrap(); let input_type = DataType::Decimal128(24, 2); let output_type = DataType::Decimal256(6, 2); assert!(can_cast_types(&input_type, &output_type)); @@ -10012,4 +10680,41 @@ mod tests { assert_eq!(result.unwrap_err().to_string(), "Invalid argument error: 123456789 is too large to store in a Decimal256 of precision 6. Max is 999999"); } + + #[test] + fn test_first_none() { + let array = Arc::new(ListArray::from_iter_primitive::(vec![ + None, + Some(vec![Some(1), Some(2)]), + ])) as ArrayRef; + let data_type = + DataType::FixedSizeList(FieldRef::new(Field::new("item", DataType::Int64, true)), 2); + let opt = CastOptions::default(); + let r = cast_with_options(&array, &data_type, &opt).unwrap(); + + let fixed_array = Arc::new(FixedSizeListArray::from_iter_primitive::( + vec![None, Some(vec![Some(1), Some(2)])], + 2, + )) as ArrayRef; + assert_eq!(*fixed_array, *r); + } + + #[test] + fn test_first_last_none() { + let array = Arc::new(ListArray::from_iter_primitive::(vec![ + None, + Some(vec![Some(1), Some(2)]), + None, + ])) as ArrayRef; + let data_type = + DataType::FixedSizeList(FieldRef::new(Field::new("item", DataType::Int64, true)), 2); + let opt = CastOptions::default(); + let r = cast_with_options(&array, &data_type, &opt).unwrap(); + + let fixed_array = Arc::new(FixedSizeListArray::from_iter_primitive::( + vec![None, Some(vec![Some(1), Some(2)]), None], + 2, + )) as ArrayRef; + assert_eq!(*fixed_array, *r); + } } diff --git a/arrow-cast/src/display.rs b/arrow-cast/src/display.rs index 669b8a664c2b..caa9804507d8 100644 --- a/arrow-cast/src/display.rs +++ b/arrow-cast/src/display.rs @@ -72,6 +72,8 @@ pub struct FormatOptions<'a> { time_format: TimeFormat<'a>, /// Duration format duration_format: DurationFormat, + /// Show types in visual representation batches + types_info: bool, } impl Default for FormatOptions<'_> { @@ -92,6 +94,7 @@ impl<'a> FormatOptions<'a> { timestamp_tz_format: None, time_format: None, duration_format: DurationFormat::ISO8601, + types_info: false, } } @@ -158,6 +161,18 @@ impl<'a> FormatOptions<'a> { ..self } } + + /// Overrides if types should be shown + /// + /// Defaults to [`false`] + pub const fn with_types_info(self, types_info: bool) -> Self { + Self { types_info, ..self } + } + + /// Returns true if type info should be included in visual representation of batches + pub const fn types_info(&self) -> bool { + self.types_info + } } /// Implements [`Display`] for a specific array value @@ -474,7 +489,7 @@ macro_rules! decimal_display { }; } -decimal_display!(Decimal128Type, Decimal256Type); +decimal_display!(Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type); fn write_timestamp( f: &mut dyn Write, @@ -575,6 +590,12 @@ temporal_display!(time32ms_to_time, time_format, Time32MillisecondType); temporal_display!(time64us_to_time, time_format, Time64MicrosecondType); temporal_display!(time64ns_to_time, time_format, Time64NanosecondType); +/// Derive [`DisplayIndexState`] for `PrimitiveArray<$t>` +/// +/// Arguments +/// * `$convert` - function to convert the value to an `Duration` +/// * `$t` - [`ArrowPrimitiveType`] of the array +/// * `$scale` - scale of the duration (passed to `duration_fmt`) macro_rules! duration_display { ($convert:ident, $t:ty, $scale:tt) => { impl<'a> DisplayIndexState<'a> for &'a PrimitiveArray<$t> { @@ -596,6 +617,34 @@ macro_rules! duration_display { }; } +/// Similar to [`duration_display`] but `$convert` returns an `Option` +macro_rules! duration_option_display { + ($convert:ident, $t:ty, $scale:tt) => { + impl<'a> DisplayIndexState<'a> for &'a PrimitiveArray<$t> { + type State = DurationFormat; + + fn prepare(&self, options: &FormatOptions<'a>) -> Result { + Ok(options.duration_format) + } + + fn write(&self, fmt: &Self::State, idx: usize, f: &mut dyn Write) -> FormatResult { + let v = self.value(idx); + match fmt { + DurationFormat::ISO8601 => match $convert(v) { + Some(td) => write!(f, "{}", td)?, + None => write!(f, "")?, + }, + DurationFormat::Pretty => match $convert(v) { + Some(_) => duration_fmt!(f, v, $scale)?, + None => write!(f, "")?, + }, + } + Ok(()) + } + } + }; +} + macro_rules! duration_fmt { ($f:ident, $v:expr, 0) => {{ let secs = $v; @@ -642,8 +691,8 @@ macro_rules! duration_fmt { }}; } -duration_display!(duration_s_to_duration, DurationSecondType, 0); -duration_display!(duration_ms_to_duration, DurationMillisecondType, 3); +duration_option_display!(try_duration_s_to_duration, DurationSecondType, 0); +duration_option_display!(try_duration_ms_to_duration, DurationMillisecondType, 3); duration_display!(duration_us_to_duration, DurationMicrosecondType, 6); duration_display!(duration_ns_to_duration, DurationNanosecondType, 9); @@ -727,12 +776,12 @@ impl Display for NanosecondsFormatter<'_> { let nanoseconds = self.nanoseconds % 1_000_000_000; if hours != 0 { - write!(f, "{prefix}{} hours", hours)?; + write!(f, "{prefix}{hours} hours")?; prefix = " "; } if mins != 0 { - write!(f, "{prefix}{} mins", mins)?; + write!(f, "{prefix}{mins} mins")?; prefix = " "; } @@ -770,12 +819,12 @@ impl Display for MillisecondsFormatter<'_> { let milliseconds = self.milliseconds % 1_000; if hours != 0 { - write!(f, "{prefix}{} hours", hours,)?; + write!(f, "{prefix}{hours} hours")?; prefix = " "; } if mins != 0 { - write!(f, "{prefix}{} mins", mins,)?; + write!(f, "{prefix}{mins} mins")?; prefix = " "; } @@ -1056,9 +1105,8 @@ pub fn lexical_to_string(n: N) -> String { #[cfg(test)] mod tests { - use arrow_array::builder::StringRunBuilder; - use super::*; + use arrow_array::builder::StringRunBuilder; /// Test to verify options can be constant. See #4580 const TEST_CONST_OPTIONS: FormatOptions<'static> = FormatOptions::new() diff --git a/arrow-cast/src/lib.rs b/arrow-cast/src/lib.rs index 6eac1be37c88..b042a7338519 100644 --- a/arrow-cast/src/lib.rs +++ b/arrow-cast/src/lib.rs @@ -17,6 +17,11 @@ //! Functions for converting from one data type to another in [Apache Arrow](https://docs.rs/arrow) +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] pub mod cast; pub use cast::*; diff --git a/arrow-cast/src/parse.rs b/arrow-cast/src/parse.rs index f4c4639c1c08..890719964d38 100644 --- a/arrow-cast/src/parse.rs +++ b/arrow-cast/src/parse.rs @@ -463,20 +463,11 @@ impl Parser for Float64Type { } } -/// This API is only stable since 1.70 so can't use it when current MSRV is lower -#[inline(always)] -fn is_some_and(opt: Option, f: impl FnOnce(T) -> bool) -> bool { - match opt { - None => false, - Some(x) => f(x), - } -} - macro_rules! parser_primitive { ($t:ty) => { impl Parser for $t { fn parse(string: &str) -> Option { - if !is_some_and(string.as_bytes().last(), |x| x.is_ascii_digit()) { + if !string.as_bytes().last().is_some_and(|x| x.is_ascii_digit()) { return None; } match atoi::FromRadix10SignedChecked::from_radix_10_signed_checked( @@ -595,6 +586,32 @@ const EPOCH_DAYS_FROM_CE: i32 = 719_163; const ERR_NANOSECONDS_NOT_SUPPORTED: &str = "The dates that can be represented as nanoseconds have to be between 1677-09-21T00:12:44.0 and 2262-04-11T23:47:16.854775804"; fn parse_date(string: &str) -> Option { + // If the date has an extended (signed) year such as "+10999-12-31" or "-0012-05-06" + // + // According to [ISO 8601], years have: + // Four digits or more for the year. Years in the range 0000 to 9999 will be pre-padded by + // zero to ensure four digits. Years outside that range will have a prefixed positive or negative symbol. + // + // [ISO 8601]: https://docs.oracle.com/en/java/javase/17/docs/api/java.base/java/time/format/DateTimeFormatter.html#ISO_LOCAL_DATE + if string.starts_with('+') || string.starts_with('-') { + // Skip the sign and look for the hyphen that terminates the year digits. + // According to ISO 8601 the unsigned part must be at least 4 digits. + let rest = &string[1..]; + let hyphen = rest.find('-')?; + if hyphen < 4 { + return None; + } + // The year substring is the sign and the digits (but not the separator) + // e.g. for "+10999-12-31", hyphen is 5 and s[..6] is "+10999" + let year: i32 = string[..hyphen + 1].parse().ok()?; + // The remainder should begin with a '-' which we strip off, leaving the month-day part. + let remainder = string[hyphen + 1..].strip_prefix('-')?; + let mut parts = remainder.splitn(2, '-'); + let month: u32 = parts.next()?.parse().ok()?; + let day: u32 = parts.next()?.parse().ok()?; + return NaiveDate::from_ymd_opt(year, month, day); + } + if string.len() > 10 { // Try to parse as datetime and return just the date part return string_to_datetime(&Utc, string) @@ -881,7 +898,7 @@ pub fn parse_decimal( for (_, b) in bs.by_ref() { if !b.is_ascii_digit() { if *b == b'e' || *b == b'E' { - result = match parse_e_notation::( + result = parse_e_notation::( s, digits as u16, fractionals as i16, @@ -889,10 +906,7 @@ pub fn parse_decimal( point_index, precision as u16, scale as i16, - ) { - Err(e) => return Err(e), - Ok(v) => v, - }; + )?; is_e_notation = true; @@ -926,7 +940,7 @@ pub fn parse_decimal( } } b'e' | b'E' => { - result = match parse_e_notation::( + result = parse_e_notation::( s, digits as u16, fractionals as i16, @@ -934,10 +948,7 @@ pub fn parse_decimal( index, precision as u16, scale as i16, - ) { - Err(e) => return Err(e), - Ok(v) => v, - }; + )?; is_e_notation = true; @@ -1224,8 +1235,7 @@ impl Interval { match (self.months, self.days, self.nanos) { (months, days, nanos) if days == 0 && nanos == 0 => Ok(months), _ => Err(ArrowError::InvalidArgumentError(format!( - "Unable to represent interval with days and nanos as year-months: {:?}", - self + "Unable to represent interval with days and nanos as year-months: {self:?}" ))), } } diff --git a/arrow-cast/src/pretty.rs b/arrow-cast/src/pretty.rs index ad3b952c327d..c3fc00e4b911 100644 --- a/arrow-cast/src/pretty.rs +++ b/arrow-cast/src/pretty.rs @@ -27,25 +27,118 @@ use std::fmt::Display; use comfy_table::{Cell, Table}; use arrow_array::{Array, ArrayRef, RecordBatch}; -use arrow_schema::ArrowError; +use arrow_schema::{ArrowError, SchemaRef}; use crate::display::{ArrayFormatter, FormatOptions}; -/// Create a visual representation of record batches +/// Create a visual representation of [`RecordBatch`]es +/// +/// Uses default values for display. See [`pretty_format_batches_with_options`] +/// for more control. +/// +/// # Example +/// ``` +/// # use std::sync::Arc; +/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +/// # use arrow_cast::pretty::pretty_format_batches; +/// # let batch = RecordBatch::try_from_iter(vec![ +/// # ("a", Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as ArrayRef), +/// # ("b", Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, Some("d"), Some("e")]))), +/// # ]).unwrap(); +/// // Note, returned object implements `Display` +/// let pretty_table = pretty_format_batches(&[batch]).unwrap(); +/// let table_str = format!("Batches:\n{pretty_table}"); +/// assert_eq!(table_str, +/// r#"Batches: +/// +---+---+ +/// | a | b | +/// +---+---+ +/// | 1 | a | +/// | 2 | b | +/// | 3 | | +/// | 4 | d | +/// | 5 | e | +/// +---+---+"#); +/// ``` pub fn pretty_format_batches(results: &[RecordBatch]) -> Result { let options = FormatOptions::default().with_display_error(true); pretty_format_batches_with_options(results, &options) } -/// Create a visual representation of record batches +/// Create a visual representation of [`RecordBatch`]es with a provided schema. +/// +/// Useful to display empty batches. +/// +/// # Example +/// ``` +/// # use std::sync::Arc; +/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +/// # use arrow_cast::pretty::pretty_format_batches_with_schema; +/// # use arrow_schema::{DataType, Field, Schema}; +/// let schema = Arc::new(Schema::new(vec![ +/// Field::new("a", DataType::Int32, false), +/// Field::new("b", DataType::Utf8, true), +/// ])); +/// // Note, returned object implements `Display` +/// let pretty_table = pretty_format_batches_with_schema(schema, &[]).unwrap(); +/// let table_str = format!("Batches:\n{pretty_table}"); +/// assert_eq!(table_str, +/// r#"Batches: +/// +---+---+ +/// | a | b | +/// +---+---+ +/// +---+---+"#); +/// ``` +pub fn pretty_format_batches_with_schema( + schema: SchemaRef, + results: &[RecordBatch], +) -> Result { + let options = FormatOptions::default().with_display_error(true); + create_table(Some(schema), results, &options) +} + +/// Create a visual representation of [`RecordBatch`]es with formatting options. +/// +/// # Arguments +/// * `results` - A slice of record batches to display +/// * `options` - [`FormatOptions`] that control the resulting display +/// +/// # Example +/// ``` +/// # use std::sync::Arc; +/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +/// # use arrow_cast::display::FormatOptions; +/// # use arrow_cast::pretty::{pretty_format_batches, pretty_format_batches_with_options}; +/// # let batch = RecordBatch::try_from_iter(vec![ +/// # ("a", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef), +/// # ("b", Arc::new(StringArray::from(vec![Some("a"), None]))), +/// # ]).unwrap(); +/// let options = FormatOptions::new() +/// .with_null(""); +/// // Note, returned object implements `Display` +/// let pretty_table = pretty_format_batches_with_options(&[batch], &options).unwrap(); +/// let table_str = format!("Batches:\n{pretty_table}"); +/// assert_eq!(table_str, +/// r#"Batches: +/// +---+--------+ +/// | a | b | +/// +---+--------+ +/// | 1 | a | +/// | 2 | | +/// +---+--------+"#); +/// ``` pub fn pretty_format_batches_with_options( results: &[RecordBatch], options: &FormatOptions, ) -> Result { - create_table(results, options) + create_table(None, results, options) } -/// Create a visual representation of columns +/// Create a visual representation of [`ArrayRef`] +/// +/// Uses default values for display. See [`pretty_format_columns_with_options`] +/// +/// See [`pretty_format_batches`] for an example pub fn pretty_format_columns( col_name: &str, results: &[ArrayRef], @@ -54,8 +147,10 @@ pub fn pretty_format_columns( pretty_format_columns_with_options(col_name, results, &options) } -/// Utility function to create a visual representation of columns with options -fn pretty_format_columns_with_options( +/// Create a visual representation of [`ArrayRef`] with formatting options. +/// +/// See [`pretty_format_batches_with_options`] for an example +pub fn pretty_format_columns_with_options( col_name: &str, results: &[ArrayRef], options: &FormatOptions, @@ -76,21 +171,41 @@ pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<(), ArrowEr } /// Convert a series of record batches into a table -fn create_table(results: &[RecordBatch], options: &FormatOptions) -> Result { +fn create_table( + schema_opt: Option, + results: &[RecordBatch], + options: &FormatOptions, +) -> Result { let mut table = Table::new(); table.load_preset("||--+-++| ++++++"); - if results.is_empty() { - return Ok(table); + let schema_opt = schema_opt.or_else(|| { + if results.is_empty() { + None + } else { + Some(results[0].schema()) + } + }); + + if let Some(schema) = schema_opt { + let mut header = Vec::new(); + for field in schema.fields() { + if options.types_info() { + header.push(Cell::new(format!( + "{}\n{}", + field.name(), + field.data_type() + ))) + } else { + header.push(Cell::new(field.name())); + } + } + table.set_header(header); } - let schema = results[0].schema(); - - let mut header = Vec::new(); - for field in schema.fields() { - header.push(Cell::new(field.name())); + if results.is_empty() { + return Ok(table); } - table.set_header(header); for batch in results { let formatters = batch @@ -150,7 +265,7 @@ mod tests { use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer}; use arrow_schema::*; - use crate::display::array_value_to_string; + use crate::display::{array_value_to_string, DurationFormat}; use super::*; @@ -1059,9 +1174,18 @@ mod tests { #[test] fn test_format_options() { - let options = FormatOptions::default().with_null("null"); - let array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]); - let batch = RecordBatch::try_from_iter([("my_column_name", Arc::new(array) as _)]).unwrap(); + let options = FormatOptions::default() + .with_null("null") + .with_types_info(true); + let int32_array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]); + let string_array = + StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz"), None]); + + let batch = RecordBatch::try_from_iter([ + ("my_int32_name", Arc::new(int32_array) as _), + ("my_string_name", Arc::new(string_array) as _), + ]) + .unwrap(); let column = pretty_format_columns_with_options( "my_column_name", @@ -1071,11 +1195,7 @@ mod tests { .unwrap() .to_string(); - let batch = pretty_format_batches_with_options(&[batch], &options) - .unwrap() - .to_string(); - - let expected = vec![ + let expected_column = vec![ "+----------------+", "| my_column_name |", "+----------------+", @@ -1088,9 +1208,78 @@ mod tests { ]; let actual: Vec<&str> = column.lines().collect(); - assert_eq!(expected, actual, "Actual result:\n{column}"); + assert_eq!(expected_column, actual, "Actual result:\n{column}"); + + let batch = pretty_format_batches_with_options(&[batch], &options) + .unwrap() + .to_string(); + + let expected_table = vec![ + "+---------------+----------------+", + "| my_int32_name | my_string_name |", + "| Int32 | Utf8 |", + "+---------------+----------------+", + "| 1 | foo |", + "| 2 | bar |", + "| null | null |", + "| 3 | baz |", + "| 4 | null |", + "+---------------+----------------+", + ]; let actual: Vec<&str> = batch.lines().collect(); - assert_eq!(expected, actual, "Actual result:\n{batch}"); + assert_eq!(expected_table, actual, "Actual result:\n{batch}"); + } + + #[test] + fn duration_pretty_and_iso_extremes() { + // Build [MIN, MAX, 3661, NULL] + let arr = DurationSecondArray::from(vec![Some(i64::MIN), Some(i64::MAX), Some(3661), None]); + let array: ArrayRef = Arc::new(arr); + + // Pretty formatting + let opts = FormatOptions::default().with_null("null"); + let opts = opts.with_duration_format(DurationFormat::Pretty); + let pretty = pretty_format_columns_with_options("pretty", &[array.clone()], &opts) + .unwrap() + .to_string(); + + // Expected output + let expected_pretty = vec![ + "+------------------------------+", + "| pretty |", + "+------------------------------+", + "| |", + "| |", + "| 0 days 1 hours 1 mins 1 secs |", + "| null |", + "+------------------------------+", + ]; + + let actual: Vec<&str> = pretty.lines().collect(); + assert_eq!(expected_pretty, actual, "Actual result:\n{pretty}"); + + // ISO8601 formatting + let opts_iso = FormatOptions::default() + .with_null("null") + .with_duration_format(DurationFormat::ISO8601); + let iso = pretty_format_columns_with_options("iso", &[array], &opts_iso) + .unwrap() + .to_string(); + + // Expected output + let expected_iso = vec![ + "+-----------+", + "| iso |", + "+-----------+", + "| |", + "| |", + "| PT3661S |", + "| null |", + "+-----------+", + ]; + + let actual: Vec<&str> = iso.lines().collect(); + assert_eq!(expected_iso, actual, "Actual result:\n{iso}"); } } diff --git a/arrow-csv/Cargo.toml b/arrow-csv/Cargo.toml index 8823924eb55b..c44ec01ce357 100644 --- a/arrow-csv/Cargo.toml +++ b/arrow-csv/Cargo.toml @@ -30,9 +30,11 @@ rust-version = { workspace = true } [lib] name = "arrow_csv" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow-array = { workspace = true } arrow-cast = { workspace = true } @@ -40,7 +42,6 @@ arrow-schema = { workspace = true } chrono = { workspace = true } csv = { version = "1.1", default-features = false } csv-core = { version = "0.1" } -lazy_static = { version = "1.4", default-features = false } regex = { version = "1.7.0", default-features = false, features = ["std", "unicode", "perf"] } [dev-dependencies] diff --git a/arrow-csv/src/lib.rs b/arrow-csv/src/lib.rs index 5ce1bc6c3396..a3552eda8a3e 100644 --- a/arrow-csv/src/lib.rs +++ b/arrow-csv/src/lib.rs @@ -17,6 +17,11 @@ //! Transfer data between the Arrow memory format and CSV (comma-separated values). +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] pub mod reader; @@ -46,8 +51,8 @@ fn map_csv_error(error: csv::Error) -> ArrowError { } => ArrowError::CsvError(format!( "Encountered unequal lengths between records on CSV file. Expected {} \ records, found {} records{}", - len, expected_len, + len, pos.as_ref() .map(|pos| format!(" at line {}", pos.line())) .unwrap_or_default(), diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs index d3d518316397..7b69df51b541 100644 --- a/arrow-csv/src/reader/mod.rs +++ b/arrow-csv/src/reader/mod.rs @@ -132,30 +132,30 @@ use arrow_cast::parse::{parse_decimal, string_to_datetime, Parser}; use arrow_schema::*; use chrono::{TimeZone, Utc}; use csv::StringRecord; -use lazy_static::lazy_static; use regex::{Regex, RegexSet}; use std::fmt::{self, Debug}; use std::fs::File; use std::io::{BufRead, BufReader as StdBufReader, Read}; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use crate::map_csv_error; use crate::reader::records::{RecordDecoder, StringRecords}; use arrow_array::timezone::Tz; -lazy_static! { - /// Order should match [`InferredDataType`] - static ref REGEX_SET: RegexSet = RegexSet::new([ +/// Order should match [`InferredDataType`] +static REGEX_SET: LazyLock = LazyLock::new(|| { + RegexSet::new([ r"(?i)^(true)$|^(false)$(?-i)", //BOOLEAN - r"^-?(\d+)$", //INTEGER + r"^-?(\d+)$", //INTEGER r"^-?((\d*\.\d+|\d+\.\d*)([eE][-+]?\d+)?|\d+([eE][-+]?\d+))$", //DECIMAL - r"^\d{4}-\d\d-\d\d$", //DATE32 + r"^\d{4}-\d\d-\d\d$", //DATE32 r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d(?:[^\d\.].*)?$", //Timestamp(Second) r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,3}(?:[^\d].*)?$", //Timestamp(Millisecond) r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,6}(?:[^\d].*)?$", //Timestamp(Microsecond) r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,9}(?:[^\d].*)?$", //Timestamp(Nanosecond) - ]).unwrap(); -} + ]) + .unwrap() +}); /// A wrapper over `Option` to check if the value is `NULL`. #[derive(Debug, Clone, Default)] @@ -221,6 +221,8 @@ impl InferredDataType { } else { 1 << m } + } else if string == "NaN" || string == "nan" || string == "inf" || string == "-inf" { + 1 << 2 // Float64 } else { 1 << 8 // Utf8 } @@ -652,6 +654,22 @@ fn parse( let field = &fields[i]; match field.data_type() { DataType::Boolean => build_boolean_array(line_number, rows, i, null_regex), + DataType::Decimal32(precision, scale) => build_decimal_array::( + line_number, + rows, + i, + *precision, + *scale, + null_regex, + ), + DataType::Decimal64(precision, scale) => build_decimal_array::( + line_number, + rows, + i, + *precision, + *scale, + null_regex, + ), DataType::Decimal128(precision, scale) => build_decimal_array::( line_number, rows, @@ -934,10 +952,12 @@ fn build_primitive_array( Some(e) => Ok(Some(e)), None => Err(ArrowError::ParseError(format!( // TODO: we should surface the underlying error here. - "Error while parsing value {} for column {} at line {}", + "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'", s, + T::DATA_TYPE, col_idx, - line_number + row_index + line_number + row_index, + row ))), } }) @@ -1020,10 +1040,12 @@ fn build_boolean_array( Some(e) => Ok(Some(e)), None => Err(ArrowError::ParseError(format!( // TODO: we should surface the underlying error here. - "Error while parsing value {} for column {} at line {}", + "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'", s, + "Boolean", col_idx, - line_number + row_index + line_number + row_index, + row ))), } }) @@ -1309,6 +1331,54 @@ mod tests { assert_eq!("0.290472", lng.value_as_string(9)); } + #[test] + fn test_csv_reader_with_decimal_3264() { + let schema = Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8, false), + Field::new("lat", DataType::Decimal32(9, 6), false), + Field::new("lng", DataType::Decimal64(16, 6), false), + ])); + + let file = File::open("test/data/decimal_test.csv").unwrap(); + + let mut csv = ReaderBuilder::new(schema).build(file).unwrap(); + let batch = csv.next().unwrap().unwrap(); + // access data from a primitive array + let lat = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!("57.653484", lat.value_as_string(0)); + assert_eq!("53.002666", lat.value_as_string(1)); + assert_eq!("52.412811", lat.value_as_string(2)); + assert_eq!("51.481583", lat.value_as_string(3)); + assert_eq!("12.123456", lat.value_as_string(4)); + assert_eq!("50.760000", lat.value_as_string(5)); + assert_eq!("0.123000", lat.value_as_string(6)); + assert_eq!("123.000000", lat.value_as_string(7)); + assert_eq!("123.000000", lat.value_as_string(8)); + assert_eq!("-50.760000", lat.value_as_string(9)); + + let lng = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!("-3.335724", lng.value_as_string(0)); + assert_eq!("-2.179404", lng.value_as_string(1)); + assert_eq!("-1.778197", lng.value_as_string(2)); + assert_eq!("-3.179090", lng.value_as_string(3)); + assert_eq!("-3.179090", lng.value_as_string(4)); + assert_eq!("0.290472", lng.value_as_string(5)); + assert_eq!("0.290472", lng.value_as_string(6)); + assert_eq!("0.290472", lng.value_as_string(7)); + assert_eq!("0.290472", lng.value_as_string(8)); + assert_eq!("0.290472", lng.value_as_string(9)); + } + #[test] fn test_csv_from_buf_reader() { let schema = Schema::new(vec![ @@ -1659,7 +1729,7 @@ mod tests { let mut csv = builder.build(file).unwrap(); let batch = csv.next().unwrap().unwrap(); - assert_eq!(7, batch.num_rows()); + assert_eq!(10, batch.num_rows()); assert_eq!(6, batch.num_columns()); let schema = batch.schema(); @@ -1758,10 +1828,8 @@ mod tests { assert_eq!(&DataType::Float64, schema.field(0).data_type()); } - #[test] - fn test_parse_invalid_csv() { - let file = File::open("test/data/various_types_invalid.csv").unwrap(); - + fn invalid_csv_helper(file_name: &str) -> String { + let file = File::open(file_name).unwrap(); let schema = Schema::new(vec![ Field::new("c_int", DataType::UInt64, false), Field::new("c_float", DataType::Float32, false), @@ -1776,16 +1844,32 @@ mod tests { .with_projection(vec![0, 1, 2, 3]); let mut csv = builder.build(file).unwrap(); - match csv.next() { - Some(e) => match e { - Err(e) => assert_eq!( - "ParseError(\"Error while parsing value 4.x4 for column 1 at line 4\")", - format!("{e:?}") - ), - Ok(_) => panic!("should have failed"), - }, - None => panic!("should have failed"), - } + + csv.next().unwrap().unwrap_err().to_string() + } + + #[test] + fn test_parse_invalid_csv_float() { + let file_name = "test/data/various_invalid_types/invalid_float.csv"; + + let error = invalid_csv_helper(file_name); + assert_eq!("Parser error: Error while parsing value '4.x4' as type 'Float32' for column 1 at line 4. Row data: '[4,4.x4,,false]'", error); + } + + #[test] + fn test_parse_invalid_csv_int() { + let file_name = "test/data/various_invalid_types/invalid_int.csv"; + + let error = invalid_csv_helper(file_name); + assert_eq!("Parser error: Error while parsing value '2.3' as type 'UInt64' for column 0 at line 2. Row data: '[2.3,2.2,2.22,false]'", error); + } + + #[test] + fn test_parse_invalid_csv_bool() { + let file_name = "test/data/various_invalid_types/invalid_bool.csv"; + + let error = invalid_csv_helper(file_name); + assert_eq!("Parser error: Error while parsing value 'none' as type 'Boolean' for column 3 at line 2. Row data: '[2,2.2,2.22,none]'", error); } /// Infer the data type of a record @@ -1803,6 +1887,10 @@ mod tests { assert_eq!(infer_field_schema("10.2"), DataType::Float64); assert_eq!(infer_field_schema(".2"), DataType::Float64); assert_eq!(infer_field_schema("2."), DataType::Float64); + assert_eq!(infer_field_schema("NaN"), DataType::Float64); + assert_eq!(infer_field_schema("nan"), DataType::Float64); + assert_eq!(infer_field_schema("inf"), DataType::Float64); + assert_eq!(infer_field_schema("-inf"), DataType::Float64); assert_eq!(infer_field_schema("true"), DataType::Boolean); assert_eq!(infer_field_schema("trUe"), DataType::Boolean); assert_eq!(infer_field_schema("false"), DataType::Boolean); @@ -2372,7 +2460,7 @@ mod tests { fn test_buffered() { let tests = [ ("test/data/uk_cities.csv", false, 37), - ("test/data/various_types.csv", true, 7), + ("test/data/various_types.csv", true, 10), ("test/data/decimal_test.csv", false, 10), ]; @@ -2609,7 +2697,7 @@ mod tests { .infer_schema(&mut read, None); assert!(result.is_err()); // Include line number in the error message to help locate and fix the issue - assert_eq!(result.err().unwrap().to_string(), "Csv error: Encountered unequal lengths between records on CSV file. Expected 2 records, found 3 records at line 3"); + assert_eq!(result.err().unwrap().to_string(), "Csv error: Encountered unequal lengths between records on CSV file. Expected 3 records, found 2 records at line 3"); } #[test] diff --git a/arrow-csv/src/reader/records.rs b/arrow-csv/src/reader/records.rs index a07fc9c94ffa..33927c93360a 100644 --- a/arrow-csv/src/reader/records.rs +++ b/arrow-csv/src/reader/records.rs @@ -290,6 +290,21 @@ impl<'a> StringRecord<'a> { } } +impl std::fmt::Display for StringRecord<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let num_fields = self.offsets.len() - 1; + write!(f, "[")?; + for i in 0..num_fields { + if i > 0 { + write!(f, ",")?; + } + write!(f, "{}", self.get(i))?; + } + write!(f, "]")?; + Ok(()) + } +} + #[cfg(test)] mod tests { use crate::reader::records::RecordDecoder; diff --git a/arrow-csv/src/writer.rs b/arrow-csv/src/writer.rs index c5a0a0b76d59..c2cb38a226b6 100644 --- a/arrow-csv/src/writer.rs +++ b/arrow-csv/src/writer.rs @@ -418,8 +418,8 @@ mod tests { use crate::ReaderBuilder; use arrow_array::builder::{ - BinaryBuilder, Decimal128Builder, Decimal256Builder, FixedSizeBinaryBuilder, - LargeBinaryBuilder, + BinaryBuilder, Decimal128Builder, Decimal256Builder, Decimal32Builder, Decimal64Builder, + FixedSizeBinaryBuilder, LargeBinaryBuilder, }; use arrow_array::types::*; use arrow_buffer::i256; @@ -496,25 +496,38 @@ sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo #[test] fn test_write_csv_decimal() { let schema = Schema::new(vec![ - Field::new("c1", DataType::Decimal128(38, 6), true), - Field::new("c2", DataType::Decimal256(76, 6), true), + Field::new("c1", DataType::Decimal32(9, 6), true), + Field::new("c2", DataType::Decimal64(17, 6), true), + Field::new("c3", DataType::Decimal128(38, 6), true), + Field::new("c4", DataType::Decimal256(76, 6), true), ]); - let mut c1_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(38, 6)); + let mut c1_builder = Decimal32Builder::new().with_data_type(DataType::Decimal32(9, 6)); c1_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]); let c1 = c1_builder.finish(); - let mut c2_builder = Decimal256Builder::new().with_data_type(DataType::Decimal256(76, 6)); - c2_builder.extend(vec![ + let mut c2_builder = Decimal64Builder::new().with_data_type(DataType::Decimal64(17, 6)); + c2_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]); + let c2 = c2_builder.finish(); + + let mut c3_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(38, 6)); + c3_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]); + let c3 = c3_builder.finish(); + + let mut c4_builder = Decimal256Builder::new().with_data_type(DataType::Decimal256(76, 6)); + c4_builder.extend(vec![ Some(i256::from_i128(-3335724)), Some(i256::from_i128(2179404)), None, Some(i256::from_i128(290472)), ]); - let c2 = c2_builder.finish(); + let c4 = c4_builder.finish(); - let batch = - RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap(); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)], + ) + .unwrap(); let mut file = tempfile::tempfile().unwrap(); @@ -530,15 +543,15 @@ 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(); - let expected = r#"c1,c2 --3.335724,-3.335724 -2.179404,2.179404 -, -0.290472,0.290472 --3.335724,-3.335724 -2.179404,2.179404 -, -0.290472,0.290472 + let expected = r#"c1,c2,c3,c4 +-3.335724,-3.335724,-3.335724,-3.335724 +2.179404,2.179404,2.179404,2.179404 +,,, +0.290472,0.290472,0.290472,0.290472 +-3.335724,-3.335724,-3.335724,-3.335724 +2.179404,2.179404,2.179404,2.179404 +,,, +0.290472,0.290472,0.290472,0.290472 "#; assert_eq!(expected, str::from_utf8(&buffer).unwrap()); } diff --git a/arrow-csv/test/data/various_invalid_types/invalid_bool.csv b/arrow-csv/test/data/various_invalid_types/invalid_bool.csv new file mode 100644 index 000000000000..81fd713df3de --- /dev/null +++ b/arrow-csv/test/data/various_invalid_types/invalid_bool.csv @@ -0,0 +1,6 @@ +c_int|c_float|c_string|c_bool +1|1.1|"1.11"|true +2|2.2|"2.22"|none +3|3.3|"3.33"|true +4|4.4|"4.4"|false +5|6.6|""|false diff --git a/arrow-csv/test/data/various_types_invalid.csv b/arrow-csv/test/data/various_invalid_types/invalid_float.csv similarity index 100% rename from arrow-csv/test/data/various_types_invalid.csv rename to arrow-csv/test/data/various_invalid_types/invalid_float.csv diff --git a/arrow-csv/test/data/various_invalid_types/invalid_int.csv b/arrow-csv/test/data/various_invalid_types/invalid_int.csv new file mode 100644 index 000000000000..b2046a89943d --- /dev/null +++ b/arrow-csv/test/data/various_invalid_types/invalid_int.csv @@ -0,0 +1,6 @@ +c_int|c_float|c_string|c_bool +1|1.1|"1.11"|true +2.3|2.2|"2.22"|false +3|3.3|"3.33"|true +4|4.4|"4.4"|false +5|6.6|""|false diff --git a/arrow-csv/test/data/various_invalid_types/null_in_non_nullable.csv b/arrow-csv/test/data/various_invalid_types/null_in_non_nullable.csv new file mode 100644 index 000000000000..2b4368d86068 --- /dev/null +++ b/arrow-csv/test/data/various_invalid_types/null_in_non_nullable.csv @@ -0,0 +1,6 @@ +c_int|c_float|c_string|c_bool +1|1.1|"1.11"|true +2|2.2|"2.22"|true +3|3.3|"3.33"|true +4|4.4||false +5|6.6|""|false diff --git a/arrow-csv/test/data/various_types.csv b/arrow-csv/test/data/various_types.csv index 570d07f5c221..66750090a93a 100644 --- a/arrow-csv/test/data/various_types.csv +++ b/arrow-csv/test/data/various_types.csv @@ -5,4 +5,7 @@ c_int|c_float|c_string|c_bool|c_date|c_datetime 4|4.4||false|| 5|6.6|""|false|1990-01-01|1990-01-01T03:00:00 4|4e6||false|| -4|4.0e-6||false|| \ No newline at end of file +4|4.0e-6||false|| +6|NaN||false|| +7|inf||false|| +8|-inf||false|| diff --git a/arrow-data/Cargo.toml b/arrow-data/Cargo.toml index c83f867523d5..fbed24fea1fa 100644 --- a/arrow-data/Cargo.toml +++ b/arrow-data/Cargo.toml @@ -30,7 +30,6 @@ rust-version = { workspace = true } [lib] name = "arrow_data" -path = "src/lib.rs" bench = false [features] @@ -42,7 +41,7 @@ force_validate = [] ffi = ["arrow-schema/ffi"] [package.metadata.docs.rs] -features = ["ffi"] +all-features = true [dependencies] diff --git a/arrow-data/src/byte_view.rs b/arrow-data/src/byte_view.rs index 3b3ec6246066..270f4f9948ac 100644 --- a/arrow-data/src/byte_view.rs +++ b/arrow-data/src/byte_view.rs @@ -18,6 +18,14 @@ use arrow_buffer::Buffer; use arrow_schema::ArrowError; +/// The maximum number of bytes that can be stored inline in a byte view. +/// +/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the +/// layout of the views. +/// +/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html +pub const MAX_INLINE_VIEW_LEN: u32 = 12; + /// Helper to access views of [`GenericByteViewArray`] (`StringViewArray` and /// `BinaryViewArray`) where the length is greater than 12 bytes. /// @@ -76,15 +84,15 @@ impl ByteView { /// See example on [`ByteView`] docs /// /// Notes: - /// * the length should always be greater than 12 (Data less than 12 - /// bytes is stored as an inline view) + /// * the length should always be greater than [`MAX_INLINE_VIEW_LEN`] + /// (Data less than 12 bytes is stored as an inline view) /// * buffer and offset are set to `0` /// /// # Panics /// If the prefix is not exactly 4 bytes #[inline] pub fn new(length: u32, prefix: &[u8]) -> Self { - debug_assert!(length > 12); + debug_assert!(length > MAX_INLINE_VIEW_LEN); Self { length, prefix: u32::from_le_bytes(prefix.try_into().unwrap()), @@ -159,8 +167,8 @@ where { for (idx, v) in views.iter().enumerate() { let len = *v as u32; - if len <= 12 { - if len < 12 && (v >> (32 + len * 8)) != 0 { + if len <= MAX_INLINE_VIEW_LEN { + if len < MAX_INLINE_VIEW_LEN && (v >> (32 + len * 8)) != 0 { return Err(ArrowError::InvalidArgumentError(format!( "View at index {idx} contained non-zero padding for string of length {len}", ))); diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs index a35b5e8629e9..fca19bc3aafe 100644 --- a/arrow-data/src/data.rs +++ b/arrow-data/src/data.rs @@ -83,6 +83,10 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff | DataType::Float16 | DataType::Float32 | DataType::Float64 + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) | DataType::Date32 | DataType::Time32(_) | DataType::Date64 @@ -139,10 +143,6 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff DataType::FixedSizeList(_, _) | DataType::Struct(_) | DataType::RunEndEncoded(_, _) => { [empty_buffer, MutableBuffer::new(0)] } - DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => [ - MutableBuffer::new(capacity * mem::size_of::()), - empty_buffer, - ], DataType::Union(_, mode) => { let type_ids = MutableBuffer::new(capacity * mem::size_of::()); match mode { @@ -203,26 +203,50 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff #[derive(Debug, Clone)] pub struct ArrayData { - /// The data type for this array data + /// The data type data_type: DataType, - /// The number of elements in this array data + /// The number of elements len: usize, - /// The offset into this array data, in number of items + /// The offset in number of items (not bytes). + /// + /// The offset applies to [`Self::child_data`] and [`Self::buffers`]. It + /// does NOT apply to [`Self::nulls`]. offset: usize, - /// The buffers for this array data. Note that depending on the array types, this - /// could hold different kinds of buffers (e.g., value buffer, value offset buffer) - /// at different positions. + /// The buffers that store the actual data for this array, as defined + /// in the [Arrow Spec]. + /// + /// Depending on the array types, [`Self::buffers`] can hold different + /// kinds of buffers (e.g., value buffer, value offset buffer) at different + /// positions. + /// + /// The buffer may be larger than needed. Some items at the beginning may be skipped if + /// there is an `offset`. Some items at the end may be skipped if the buffer is longer than + /// we need to satisfy `len`. + /// + /// [Arrow Spec](https://arrow.apache.org/docs/format/Columnar.html#physical-memory-layout) buffers: Vec, - /// The child(ren) of this array. Only non-empty for nested types, currently - /// `ListArray` and `StructArray`. + /// The child(ren) of this array. + /// + /// Only non-empty for nested types, such as `ListArray` and + /// `StructArray`. + /// + /// The first logical element in each child element begins at `offset`. + /// + /// If the child element also has an offset then these offsets are + /// cumulative. child_data: Vec, - /// The null bitmap. A `None` value for this indicates all values are non-null in - /// this array. + /// The null bitmap. + /// + /// `None` indicates all values are non-null in this array. + /// + /// [`Self::offset]` does not apply to the null bitmap. While the + /// BooleanBuffer may be sliced (have its own offset) internally, this + /// `NullBuffer` always represents exactly `len` elements. nulls: Option, } @@ -255,6 +279,10 @@ impl ArrayData { buffers: Vec, child_data: Vec, ) -> Self { + let mut skip_validation = UnsafeFlag::new(); + // SAFETY: caller responsible for ensuring data is valid + skip_validation.set(true); + ArrayDataBuilder { data_type, len, @@ -264,8 +292,11 @@ impl ArrayData { offset, buffers, child_data, + align_buffers: false, + skip_validation, } - .build_unchecked() + .build() + .unwrap() } /// Create a new ArrayData, validating that the provided buffers form a valid @@ -550,6 +581,7 @@ impl ArrayData { } /// Returns the `buffer` as a slice of type `T` starting at self.offset + /// /// # Panics /// This function panics if: /// * the buffer is not byte-aligned with type T, or @@ -606,7 +638,7 @@ impl ArrayData { ), DataType::Union(f, mode) => { let (id, _) = f.iter().next().unwrap(); - let ids = Buffer::from_iter(std::iter::repeat(id).take(len)); + let ids = Buffer::from_iter(std::iter::repeat_n(id, len)); let buffers = match mode { UnionMode::Sparse => vec![ids], UnionMode::Dense => { @@ -1011,7 +1043,7 @@ impl ArrayData { if values_data.len < expected_values_len { return Err(ArrowError::InvalidArgumentError(format!( "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}", - values_data.len, list_size, list_size, self.data_type + values_data.len, self.len, list_size, self.data_type ))); } @@ -1582,6 +1614,8 @@ pub fn layout(data_type: &DataType) -> DataTypeLayout { DataTypeLayout::new_fixed_width::() } DataType::Duration(_) => DataTypeLayout::new_fixed_width::(), + DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::(), + DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::(), DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::(), DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::(), DataType::FixedSizeBinary(size) => { @@ -1775,7 +1809,66 @@ impl PartialEq for ArrayData { } } -/// Builder for `ArrayData` type +/// A boolean flag that cannot be mutated outside of unsafe code. +/// +/// Defaults to a value of false. +/// +/// This structure is used to enforce safety in the [`ArrayDataBuilder`] +/// +/// [`ArrayDataBuilder`]: super::ArrayDataBuilder +/// +/// # Example +/// ```rust +/// use arrow_data::UnsafeFlag; +/// assert!(!UnsafeFlag::default().get()); // default is false +/// let mut flag = UnsafeFlag::new(); +/// assert!(!flag.get()); // defaults to false +/// // can only set it to true in unsafe code +/// unsafe { flag.set(true) }; +/// assert!(flag.get()); // now true +/// ``` +#[derive(Debug, Clone)] +#[doc(hidden)] +pub struct UnsafeFlag(bool); + +impl UnsafeFlag { + /// Creates a new `UnsafeFlag` with the value set to `false`. + /// + /// See examples on [`Self::new`] + #[inline] + pub const fn new() -> Self { + Self(false) + } + + /// Sets the value of the flag to the given value + /// + /// Note this can purposely only be done in `unsafe` code + /// + /// # Safety + /// + /// If set, the flag will be set to the given value. There is nothing + /// immediately unsafe about doing so, however, the flag can be used to + /// subsequently bypass safety checks in the [`ArrayDataBuilder`]. + #[inline] + pub unsafe fn set(&mut self, val: bool) { + self.0 = val; + } + + /// Returns the value of the flag + #[inline] + pub fn get(&self) -> bool { + self.0 + } +} + +// Manual impl to make it clear you can not construct unsafe with true +impl Default for UnsafeFlag { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`ArrayData`] type #[derive(Debug)] pub struct ArrayDataBuilder { data_type: DataType, @@ -1786,6 +1879,20 @@ pub struct ArrayDataBuilder { offset: usize, buffers: Vec, child_data: Vec, + /// Should buffers be realigned (copying if necessary)? + /// + /// Defaults to false. + align_buffers: bool, + /// Should data validation be skipped for this [`ArrayData`]? + /// + /// Defaults to false. + /// + /// # Safety + /// + /// This flag can only be set to true using `unsafe` APIs. However, once true + /// subsequent calls to `build()` may result in undefined behavior if the data + /// is not valid. + skip_validation: UnsafeFlag, } impl ArrayDataBuilder { @@ -1801,6 +1908,8 @@ impl ArrayDataBuilder { offset: 0, buffers: vec![], child_data: vec![], + align_buffers: false, + skip_validation: UnsafeFlag::new(), } } @@ -1877,51 +1986,85 @@ impl ArrayDataBuilder { /// Creates an array data, without any validation /// + /// Note: This is shorthand for + /// ```rust + /// # let mut builder = arrow_data::ArrayDataBuilder::new(arrow_schema::DataType::Null); + /// # let _ = unsafe { + /// builder.skip_validation(true).build().unwrap() + /// # }; + /// ``` + /// /// # Safety /// /// The same caveats as [`ArrayData::new_unchecked`] /// apply. - #[allow(clippy::let_and_return)] pub unsafe fn build_unchecked(self) -> ArrayData { - let data = self.build_impl(); - // Provide a force_validate mode - #[cfg(feature = "force_validate")] - data.validate_data().unwrap(); - data + self.skip_validation(true).build().unwrap() } - /// Same as [`Self::build_unchecked`] but ignoring `force_validate` feature flag - unsafe fn build_impl(self) -> ArrayData { - let nulls = self - .nulls + /// Creates an `ArrayData`, consuming `self` + /// + /// # Safety + /// + /// By default the underlying buffers are checked to ensure they are valid + /// Arrow data. However, if the [`Self::skip_validation`] flag has been set + /// to true (by the `unsafe` API) this validation is skipped. If the data is + /// not valid, undefined behavior will result. + pub fn build(self) -> Result { + let Self { + data_type, + len, + null_count, + null_bit_buffer, + nulls, + offset, + buffers, + child_data, + align_buffers, + skip_validation, + } = self; + + let nulls = nulls .or_else(|| { - let buffer = self.null_bit_buffer?; - let buffer = BooleanBuffer::new(buffer, self.offset, self.len); - Some(match self.null_count { - Some(n) => NullBuffer::new_unchecked(buffer, n), + let buffer = null_bit_buffer?; + let buffer = BooleanBuffer::new(buffer, offset, len); + Some(match null_count { + Some(n) => { + // SAFETY: call to `data.validate_data()` below validates the null buffer is valid + unsafe { NullBuffer::new_unchecked(buffer, n) } + } None => NullBuffer::new(buffer), }) }) .filter(|b| b.null_count() != 0); - ArrayData { - data_type: self.data_type, - len: self.len, - offset: self.offset, - buffers: self.buffers, - child_data: self.child_data, + let mut data = ArrayData { + data_type, + len, + offset, + buffers, + child_data, nulls, + }; + + if align_buffers { + data.align_buffers(); } - } - /// Creates an array data, validating all inputs - pub fn build(self) -> Result { - let data = unsafe { self.build_impl() }; - data.validate_data()?; + // SAFETY: `skip_validation` is only set to true using `unsafe` APIs + if !skip_validation.get() || cfg!(feature = "force_validate") { + data.validate_data()?; + } Ok(data) } /// Creates an array data, validating all inputs, and aligning any buffers + #[deprecated(since = "54.1.0", note = "Use ArrayData::align_buffers instead")] + pub fn build_aligned(self) -> Result { + self.align_buffers(true).build() + } + + /// Ensure that all buffers are aligned, copying data if necessary /// /// Rust requires that arrays are aligned to their corresponding primitive, /// see [`Layout::array`](std::alloc::Layout::array) and [`std::mem::align_of`]. @@ -1930,17 +2073,33 @@ impl ArrayDataBuilder { /// to allow for [slice](std::slice) based APIs. See [`BufferSpec::FixedWidth`]. /// /// As this alignment is architecture specific, and not guaranteed by all arrow implementations, - /// this method is provided to automatically copy buffers to a new correctly aligned allocation + /// this flag is provided to automatically copy buffers to a new correctly aligned allocation /// when necessary, making it useful when interacting with buffers produced by other systems, /// e.g. IPC or FFI. /// - /// This is unlike `[Self::build`] which will instead return an error on encountering + /// If this flag is not enabled, `[Self::build`] return an error on encountering /// insufficiently aligned buffers. - pub fn build_aligned(self) -> Result { - let mut data = unsafe { self.build_impl() }; - data.align_buffers(); - data.validate_data()?; - Ok(data) + pub fn align_buffers(mut self, align_buffers: bool) -> Self { + self.align_buffers = align_buffers; + self + } + + /// Skips validation of the data. + /// + /// If this flag is enabled, `[Self::build`] will skip validation of the + /// data + /// + /// If this flag is not enabled, `[Self::build`] will validate that all + /// buffers are valid and will return an error if any data is invalid. + /// Validation can be expensive. + /// + /// # Safety + /// + /// If validation is skipped, the buffers must form a valid Arrow array, + /// otherwise undefined behavior will result + pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self { + self.skip_validation.set(skip_validation); + self } } @@ -1955,6 +2114,8 @@ impl From for ArrayDataBuilder { nulls: d.nulls, null_bit_buffer: None, null_count: None, + align_buffers: false, + skip_validation: UnsafeFlag::new(), } } } diff --git a/arrow-data/src/decimal.rs b/arrow-data/src/decimal.rs index fe19db641236..35a7c08d8e47 100644 --- a/arrow-data/src/decimal.rs +++ b/arrow-data/src/decimal.rs @@ -15,24 +15,45 @@ // specific language governing permissions and limitations // under the License. -//! Defines maximum and minimum values for `decimal256` and `decimal128` types for varying precisions. +//! Maximum and minimum values for [`Decimal256`], [`Decimal128`], [`Decimal64`] and [`Decimal32`]. //! -//! Also provides functions to validate if a given decimal value is within the valid range of the decimal type. - +//! Also provides functions to validate if a given decimal value is within +//! the valid range of the decimal type. +//! +//! [`Decimal32`]: arrow_schema::DataType::Decimal32 +//! [`Decimal64`]: arrow_schema::DataType::Decimal64 +//! [`Decimal128`]: arrow_schema::DataType::Decimal128 +//! [`Decimal256`]: arrow_schema::DataType::Decimal256 use arrow_buffer::i256; use arrow_schema::ArrowError; pub use arrow_schema::{ DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, - DECIMAL_DEFAULT_SCALE, + DECIMAL32_DEFAULT_SCALE, DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, DECIMAL64_DEFAULT_SCALE, + DECIMAL64_MAX_PRECISION, DECIMAL64_MAX_SCALE, DECIMAL_DEFAULT_SCALE, }; -/// MAX decimal256 value of little-endian format for each precision. -/// Each element is the max value of signed 256-bit integer for the specified precision which -/// is encoded to the 32-byte width format of little-endian. +/// `MAX_DECIMAL256_FOR_EACH_PRECISION[p]` holds the maximum [`i256`] value that can +/// be stored in a [`Decimal256`] value of precision `p`. +/// +/// # Notes +/// +/// Each element is the max value of signed 256-bit integer for the specified +/// precision which is encoded to the 32-byte width format of little-endian. +/// /// The first element is unused and is inserted so that we can look up using /// precision as the index without the need to subtract 1 first. -pub(crate) const MAX_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION: [i256; 77] = [ +/// +/// # Example +/// ``` +/// # use arrow_buffer::i256; +/// # use arrow_data::decimal::MAX_DECIMAL256_FOR_EACH_PRECISION; +/// assert_eq!(MAX_DECIMAL256_FOR_EACH_PRECISION[3], i256::from(999)); +/// ``` +/// +/// [`Decimal256`]: arrow_schema::DataType::Decimal256 +/// [`i256`]: arrow_buffer::i256 +pub const MAX_DECIMAL256_FOR_EACH_PRECISION: [i256; 77] = [ i256::from_i128(0_i128), // unused first element i256::from_le_bytes([ 9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, @@ -340,12 +361,26 @@ pub(crate) const MAX_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION: [i256; 77] = [ ]), ]; -/// MIN decimal256 value of little-endian format for each precision. +/// `MIN_DECIMAL256_FOR_EACH_PRECISION[p]` holds the minimum [`i256`] value that can +/// be stored in a [`Decimal256`] value of precision `p`. +/// +/// # Notes +/// /// Each element is the min value of signed 256-bit integer for the specified precision which /// is encoded to the 76-byte width format of little-endian. +/// /// The first element is unused and is inserted so that we can look up using /// precision as the index without the need to subtract 1 first. -pub(crate) const MIN_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION: [i256; 77] = [ +/// # Example +/// ``` +/// # use arrow_buffer::i256; +/// # use arrow_data::decimal::MIN_DECIMAL256_FOR_EACH_PRECISION; +/// assert_eq!(MIN_DECIMAL256_FOR_EACH_PRECISION[3], i256::from(-999)); +/// ``` +/// +/// [`i256`]: arrow_buffer::i256 +/// [`Decimal256`]: arrow_schema::DataType::Decimal256 +pub const MIN_DECIMAL256_FOR_EACH_PRECISION: [i256; 77] = [ i256::from_i128(0_i128), // unused first element i256::from_le_bytes([ 247, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, @@ -654,7 +689,13 @@ pub(crate) const MIN_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION: [i256; 77] = [ ]; /// `MAX_DECIMAL_FOR_EACH_PRECISION[p-1]` holds the maximum `i128` value that can -/// be stored in [arrow_schema::DataType::Decimal128] value of precision `p` +/// be stored in a [`Decimal128`] value of precision `p` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 +#[deprecated( + since = "54.1.0", + note = "Use MAX_DECIMAL128_FOR_EACH_PRECISION (note indexes are different)" +)] #[allow(dead_code)] // no longer used but is part of our public API pub const MAX_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ 9, @@ -698,8 +739,14 @@ pub const MAX_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ ]; /// `MIN_DECIMAL_FOR_EACH_PRECISION[p-1]` holds the minimum `i128` value that can -/// be stored in a [arrow_schema::DataType::Decimal128] value of precision `p` +/// be stored in a [`Decimal128`] value of precision `p` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 #[allow(dead_code)] // no longer used but is part of our public API +#[deprecated( + since = "54.1.0", + note = "Use MIN_DECIMAL128_FOR_EACH_PRECISION (note indexes are different)" +)] pub const MIN_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ -9, -99, @@ -741,11 +788,22 @@ pub const MIN_DECIMAL_FOR_EACH_PRECISION: [i128; 38] = [ -99999999999999999999999999999999999999, ]; -/// `MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[p]` holds the maximum `i128` value that can -/// be stored in [arrow_schema::DataType::Decimal128] value of precision `p`. +/// `MAX_DECIMAL128_FOR_EACH_PRECISION[p]` holds the maximum `i128` value that can +/// be stored in [`Decimal128`] value of precision `p`. +/// +/// # Notes +/// /// The first element is unused and is inserted so that we can look up using /// precision as the index without the need to subtract 1 first. -pub(crate) const MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED: [i128; 39] = [ +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MAX_DECIMAL128_FOR_EACH_PRECISION; +/// assert_eq!(MAX_DECIMAL128_FOR_EACH_PRECISION[3], 999); +/// ``` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 +pub const MAX_DECIMAL128_FOR_EACH_PRECISION: [i128; 39] = [ 0, // unused first element 9, 99, @@ -788,10 +846,21 @@ pub(crate) const MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED: [i128; 39] = [ ]; /// `MIN_DECIMAL_FOR_EACH_PRECISION[p]` holds the minimum `i128` value that can -/// be stored in a [arrow_schema::DataType::Decimal128] value of precision `p`. +/// be stored in a [`Decimal128`] value of precision `p`. +/// +/// # Notes +/// /// The first element is unused and is inserted so that we can look up using /// precision as the index without the need to subtract 1 first. -pub(crate) const MIN_DECIMAL_FOR_EACH_PRECISION_ONE_BASED: [i128; 39] = [ +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MIN_DECIMAL128_FOR_EACH_PRECISION; +/// assert_eq!(MIN_DECIMAL128_FOR_EACH_PRECISION[3], -999); +/// ``` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 +pub const MIN_DECIMAL128_FOR_EACH_PRECISION: [i128; 39] = [ 0, // unused first element -9, -99, @@ -833,8 +902,198 @@ pub(crate) const MIN_DECIMAL_FOR_EACH_PRECISION_ONE_BASED: [i128; 39] = [ -99999999999999999999999999999999999999, ]; +/// `MAX_DECIMAL64_FOR_EACH_PRECISION[p]` holds the maximum `i64` value that can +/// be stored in [`Decimal64`] value of precision `p`. +/// +/// # Notes +/// +/// The first element is unused and is inserted so that we can look up using +/// precision as the index without the need to subtract 1 first. +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MAX_DECIMAL64_FOR_EACH_PRECISION; +/// assert_eq!(MAX_DECIMAL64_FOR_EACH_PRECISION[3], 999); +/// ``` +/// +/// [`Decimal64`]: arrow_schema::DataType::Decimal64 +pub const MAX_DECIMAL64_FOR_EACH_PRECISION: [i64; 19] = [ + 0, // unused first element + 9, + 99, + 999, + 9999, + 99999, + 999999, + 9999999, + 99999999, + 999999999, + 9999999999, + 99999999999, + 999999999999, + 9999999999999, + 99999999999999, + 999999999999999, + 9999999999999999, + 99999999999999999, + 999999999999999999, +]; + +/// `MIN_DECIMAL64_FOR_EACH_PRECISION[p]` holds the minimum `i64` value that can +/// be stored in a [`Decimal64`] value of precision `p`. +/// +/// # Notes +/// +/// The first element is unused and is inserted so that we can look up using +/// precision as the index without the need to subtract 1 first. +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MIN_DECIMAL64_FOR_EACH_PRECISION; +/// assert_eq!(MIN_DECIMAL64_FOR_EACH_PRECISION[3], -999); +/// ``` +/// +/// [`Decimal64`]: arrow_schema::DataType::Decimal64 +pub const MIN_DECIMAL64_FOR_EACH_PRECISION: [i64; 19] = [ + 0, // unused first element + -9, + -99, + -999, + -9999, + -99999, + -999999, + -9999999, + -99999999, + -999999999, + -9999999999, + -99999999999, + -999999999999, + -9999999999999, + -99999999999999, + -999999999999999, + -9999999999999999, + -99999999999999999, + -999999999999999999, +]; + +/// `MAX_DECIMAL32_FOR_EACH_PRECISION[p]` holds the maximum `i32` value that can +/// be stored in [`Decimal32`] value of precision `p`. +/// +/// # Notes +/// +/// The first element is unused and is inserted so that we can look up using +/// precision as the index without the need to subtract 1 first. +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MAX_DECIMAL32_FOR_EACH_PRECISION; +/// assert_eq!(MAX_DECIMAL32_FOR_EACH_PRECISION[3], 999); +/// ``` +/// +/// [`Decimal32`]: arrow_schema::DataType::Decimal32 +pub const MAX_DECIMAL32_FOR_EACH_PRECISION: [i32; 10] = [ + 0, // unused first element + 9, 99, 999, 9999, 99999, 999999, 9999999, 99999999, 999999999, +]; + +/// `MIN_DECIMAL32_FOR_EACH_PRECISION[p]` holds the minimum `ialue that can +/// be stored in a [`Decimal32`] value of precision `p`. +/// +/// # Notes +/// +/// The first element is unused and is inserted so that we can look up using +/// precision as the index without the need to subtract 1 first. +/// +/// # Example +/// ``` +/// # use arrow_data::decimal::MIN_DECIMAL32_FOR_EACH_PRECISION; +/// assert_eq!(MIN_DECIMAL32_FOR_EACH_PRECISION[3], -999); +/// ``` +/// +/// [`Decimal32`]: arrow_schema::DataType::Decimal32 +pub const MIN_DECIMAL32_FOR_EACH_PRECISION: [i32; 10] = [ + 0, // unused first element + -9, -99, -999, -9999, -99999, -999999, -9999999, -99999999, -999999999, +]; + +/// Validates that the specified `i32` value can be properly +/// interpreted as a [`Decimal32`] number with precision `precision` +/// +/// [`Decimal32`]: arrow_schema::DataType::Decimal32 +#[inline] +pub fn validate_decimal32_precision(value: i32, precision: u8) -> Result<(), ArrowError> { + if precision > DECIMAL32_MAX_PRECISION { + return Err(ArrowError::InvalidArgumentError(format!( + "Max precision of a Decimal32 is {DECIMAL32_MAX_PRECISION}, but got {precision}", + ))); + } + if value > MAX_DECIMAL32_FOR_EACH_PRECISION[precision as usize] { + Err(ArrowError::InvalidArgumentError(format!( + "{value} is too large to store in a Decimal32 of precision {precision}. Max is {}", + MAX_DECIMAL32_FOR_EACH_PRECISION[precision as usize] + ))) + } else if value < MIN_DECIMAL32_FOR_EACH_PRECISION[precision as usize] { + Err(ArrowError::InvalidArgumentError(format!( + "{value} is too small to store in a Decimal32 of precision {precision}. Min is {}", + MIN_DECIMAL32_FOR_EACH_PRECISION[precision as usize] + ))) + } else { + Ok(()) + } +} + +/// Returns true if the specified `i32` value can be properly +/// interpreted as a [`Decimal32`] number with precision `precision` +/// +/// [`Decimal32`]: arrow_schema::DataType::Decimal32 +#[inline] +pub fn is_validate_decimal32_precision(value: i32, precision: u8) -> bool { + precision <= DECIMAL32_MAX_PRECISION + && value >= MIN_DECIMAL32_FOR_EACH_PRECISION[precision as usize] + && value <= MAX_DECIMAL32_FOR_EACH_PRECISION[precision as usize] +} + +/// Validates that the specified `i64` value can be properly +/// interpreted as a [`Decimal64`] number with precision `precision` +/// +/// [`Decimal64`]: arrow_schema::DataType::Decimal64 +#[inline] +pub fn validate_decimal64_precision(value: i64, precision: u8) -> Result<(), ArrowError> { + if precision > DECIMAL64_MAX_PRECISION { + return Err(ArrowError::InvalidArgumentError(format!( + "Max precision of a Decimal64 is {DECIMAL64_MAX_PRECISION}, but got {precision}", + ))); + } + if value > MAX_DECIMAL64_FOR_EACH_PRECISION[precision as usize] { + Err(ArrowError::InvalidArgumentError(format!( + "{value} is too large to store in a Decimal64 of precision {precision}. Max is {}", + MAX_DECIMAL64_FOR_EACH_PRECISION[precision as usize] + ))) + } else if value < MIN_DECIMAL64_FOR_EACH_PRECISION[precision as usize] { + Err(ArrowError::InvalidArgumentError(format!( + "{value} is too small to store in a Decimal64 of precision {precision}. Min is {}", + MIN_DECIMAL64_FOR_EACH_PRECISION[precision as usize] + ))) + } else { + Ok(()) + } +} + +/// Returns true if the specified `i64` value can be properly +/// interpreted as a [`Decimal64`] number with precision `precision` +/// +/// [`Decimal64`]: arrow_schema::DataType::Decimal64 +#[inline] +pub fn is_validate_decimal64_precision(value: i64, precision: u8) -> bool { + precision <= DECIMAL64_MAX_PRECISION + && value >= MIN_DECIMAL64_FOR_EACH_PRECISION[precision as usize] + && value <= MAX_DECIMAL64_FOR_EACH_PRECISION[precision as usize] +} + /// Validates that the specified `i128` value can be properly -/// interpreted as a Decimal number with precision `precision` +/// interpreted as a [`Decimal128`] number with precision `precision` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 #[inline] pub fn validate_decimal_precision(value: i128, precision: u8) -> Result<(), ArrowError> { if precision > DECIMAL128_MAX_PRECISION { @@ -842,32 +1101,36 @@ pub fn validate_decimal_precision(value: i128, precision: u8) -> Result<(), Arro "Max precision of a Decimal128 is {DECIMAL128_MAX_PRECISION}, but got {precision}", ))); } - if value > MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] { + if value > MAX_DECIMAL128_FOR_EACH_PRECISION[precision as usize] { Err(ArrowError::InvalidArgumentError(format!( "{value} is too large to store in a Decimal128 of precision {precision}. Max is {}", - MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] + MAX_DECIMAL128_FOR_EACH_PRECISION[precision as usize] ))) - } else if value < MIN_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] { + } else if value < MIN_DECIMAL128_FOR_EACH_PRECISION[precision as usize] { Err(ArrowError::InvalidArgumentError(format!( "{value} is too small to store in a Decimal128 of precision {precision}. Min is {}", - MIN_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] + MIN_DECIMAL128_FOR_EACH_PRECISION[precision as usize] ))) } else { Ok(()) } } -/// Determines whether the specified `i128` value can be properly -/// interpreted as a Decimal number with precision `precision` +/// Returns true if the specified `i128` value can be properly +/// interpreted as a [`Decimal128`] number with precision `precision` +/// +/// [`Decimal128`]: arrow_schema::DataType::Decimal128 #[inline] pub fn is_validate_decimal_precision(value: i128, precision: u8) -> bool { precision <= DECIMAL128_MAX_PRECISION - && value >= MIN_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] - && value <= MAX_DECIMAL_FOR_EACH_PRECISION_ONE_BASED[precision as usize] + && value >= MIN_DECIMAL128_FOR_EACH_PRECISION[precision as usize] + && value <= MAX_DECIMAL128_FOR_EACH_PRECISION[precision as usize] } /// Validates that the specified `i256` of value can be properly -/// interpreted as a Decimal256 number with precision `precision` +/// interpreted as a [`Decimal256`] number with precision `precision` +/// +/// [`Decimal256`]: arrow_schema::DataType::Decimal256 #[inline] pub fn validate_decimal256_precision(value: i256, precision: u8) -> Result<(), ArrowError> { if precision > DECIMAL256_MAX_PRECISION { @@ -875,26 +1138,28 @@ pub fn validate_decimal256_precision(value: i256, precision: u8) -> Result<(), A "Max precision of a Decimal256 is {DECIMAL256_MAX_PRECISION}, but got {precision}", ))); } - if value > MAX_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] { + if value > MAX_DECIMAL256_FOR_EACH_PRECISION[precision as usize] { Err(ArrowError::InvalidArgumentError(format!( "{value:?} is too large to store in a Decimal256 of precision {precision}. Max is {:?}", - MAX_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] + MAX_DECIMAL256_FOR_EACH_PRECISION[precision as usize] ))) - } else if value < MIN_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] { + } else if value < MIN_DECIMAL256_FOR_EACH_PRECISION[precision as usize] { Err(ArrowError::InvalidArgumentError(format!( "{value:?} is too small to store in a Decimal256 of precision {precision}. Min is {:?}", - MIN_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] + MIN_DECIMAL256_FOR_EACH_PRECISION[precision as usize] ))) } else { Ok(()) } } -/// Determines whether the specified `i256` value can be properly -/// interpreted as a Decimal256 number with precision `precision` +/// Return true if the specified `i256` value can be properly +/// interpreted as a [`Decimal256`] number with precision `precision` +/// +/// [`Decimal256`]: arrow_schema::DataType::Decimal256 #[inline] pub fn is_validate_decimal256_precision(value: i256, precision: u8) -> bool { precision <= DECIMAL256_MAX_PRECISION - && value >= MIN_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] - && value <= MAX_DECIMAL_BYTES_FOR_LARGER_EACH_PRECISION[precision as usize] + && value >= MIN_DECIMAL256_FOR_EACH_PRECISION[precision as usize] + && value <= MAX_DECIMAL256_FOR_EACH_PRECISION[precision as usize] } diff --git a/arrow-data/src/equal/byte_view.rs b/arrow-data/src/equal/byte_view.rs index def395125366..7c6001e9a24e 100644 --- a/arrow-data/src/equal/byte_view.rs +++ b/arrow-data/src/equal/byte_view.rs @@ -69,6 +69,3 @@ pub(super) fn byte_view_equal( } true } - -#[cfg(test)] -mod tests {} diff --git a/arrow-data/src/equal/mod.rs b/arrow-data/src/equal/mod.rs index f24179b61700..1c16ee2f8a14 100644 --- a/arrow-data/src/equal/mod.rs +++ b/arrow-data/src/equal/mod.rs @@ -78,6 +78,8 @@ fn equal_values( DataType::Int64 => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), DataType::Float32 => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), DataType::Float64 => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), + DataType::Decimal32(_, _) => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), + DataType::Decimal64(_, _) => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), DataType::Decimal128(_, _) => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), DataType::Decimal256(_, _) => primitive_equal::(lhs, rhs, lhs_start, rhs_start, len), DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => { diff --git a/arrow-data/src/ffi.rs b/arrow-data/src/ffi.rs index 04599bc47a1c..3b446ef255fe 100644 --- a/arrow-data/src/ffi.rs +++ b/arrow-data/src/ffi.rs @@ -289,7 +289,7 @@ impl FFI_ArrowArray { /// Returns the buffer at the provided index /// /// # Panic - /// Panics if index exceeds the number of buffers or the buffer is not correctly aligned + /// Panics if index >= self.num_buffers() or the buffer is not correctly aligned #[inline] pub fn buffer(&self, index: usize) -> *const u8 { assert!(!self.buffers.is_null()); diff --git a/arrow-data/src/lib.rs b/arrow-data/src/lib.rs index a7feca6cd976..a023b1d98cb6 100644 --- a/arrow-data/src/lib.rs +++ b/arrow-data/src/lib.rs @@ -19,6 +19,11 @@ //! //! For a higher-level, strongly-typed interface see [arrow_array](https://docs.rs/arrow_array) +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] mod data; pub use data::*; diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs index 93b79e6a5eb8..5071bf8c4113 100644 --- a/arrow-data/src/transform/mod.rs +++ b/arrow-data/src/transform/mod.rs @@ -35,6 +35,7 @@ mod fixed_size_list; mod list; mod null; mod primitive; +mod run; mod structure; mod union; mod utils; @@ -256,6 +257,8 @@ fn build_extend(array: &ArrayData) -> Extend { | DataType::Duration(_) | DataType::Interval(IntervalUnit::DayTime) => primitive::build_extend::(array), DataType::Interval(IntervalUnit::MonthDayNano) => primitive::build_extend::(array), + DataType::Decimal32(_, _) => primitive::build_extend::(array), + DataType::Decimal64(_, _) => primitive::build_extend::(array), DataType::Decimal128(_, _) => primitive::build_extend::(array), DataType::Decimal256(_, _) => primitive::build_extend::(array), DataType::Utf8 | DataType::Binary => variable_size::build_extend::(array), @@ -275,7 +278,7 @@ fn build_extend(array: &ArrayData) -> Extend { UnionMode::Sparse => union::build_extend_sparse(array), UnionMode::Dense => union::build_extend_dense(array), }, - DataType::RunEndEncoded(_, _) => todo!(), + DataType::RunEndEncoded(_, _) => run::build_extend(array), } } @@ -302,6 +305,8 @@ fn build_extend_nulls(data_type: &DataType) -> ExtendNulls { | DataType::Duration(_) | DataType::Interval(IntervalUnit::DayTime) => primitive::extend_nulls::, DataType::Interval(IntervalUnit::MonthDayNano) => primitive::extend_nulls::, + DataType::Decimal32(_, _) => primitive::extend_nulls::, + DataType::Decimal64(_, _) => primitive::extend_nulls::, DataType::Decimal128(_, _) => primitive::extend_nulls::, DataType::Decimal256(_, _) => primitive::extend_nulls::, DataType::Utf8 | DataType::Binary => variable_size::extend_nulls::, @@ -331,7 +336,7 @@ fn build_extend_nulls(data_type: &DataType) -> ExtendNulls { UnionMode::Sparse => union::extend_nulls_sparse, UnionMode::Dense => union::extend_nulls_dense, }, - DataType::RunEndEncoded(_, _) => todo!(), + DataType::RunEndEncoded(_, _) => run::extend_nulls, }) } @@ -455,7 +460,9 @@ impl<'a> MutableArrayData<'a> { }; let child_data = match &data_type { - DataType::Decimal128(_, _) + DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) | DataType::Decimal256(_, _) | DataType::Null | DataType::Boolean @@ -767,7 +774,10 @@ impl<'a> MutableArrayData<'a> { let data = self.data; let buffers = match data.data_type { - DataType::Null | DataType::Struct(_) | DataType::FixedSizeList(_, _) => { + DataType::Null + | DataType::Struct(_) + | DataType::FixedSizeList(_, _) + | DataType::RunEndEncoded(_, _) => { vec![] } DataType::BinaryView | DataType::Utf8View => { @@ -793,13 +803,17 @@ impl<'a> MutableArrayData<'a> { _ => data.child_data.into_iter().map(|x| x.freeze()).collect(), }; - let nulls = data - .null_buffer - .map(|nulls| { - let bools = BooleanBuffer::new(nulls.into(), 0, data.len); - unsafe { NullBuffer::new_unchecked(bools, data.null_count) } - }) - .filter(|n| n.null_count() > 0); + let nulls = match data.data_type { + // RunEndEncoded and Null arrays cannot have top-level null bitmasks + DataType::RunEndEncoded(_, _) | DataType::Null => None, + _ => data + .null_buffer + .map(|nulls| { + let bools = BooleanBuffer::new(nulls.into(), 0, data.len); + unsafe { NullBuffer::new_unchecked(bools, data.null_count) } + }) + .filter(|n| n.null_count() > 0), + }; ArrayDataBuilder::new(data.data_type) .offset(0) diff --git a/arrow-data/src/transform/run.rs b/arrow-data/src/transform/run.rs new file mode 100644 index 000000000000..1ab6d0d31936 --- /dev/null +++ b/arrow-data/src/transform/run.rs @@ -0,0 +1,649 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::{ArrayData, Extend, _MutableArrayData}; +use arrow_buffer::{ArrowNativeType, Buffer, ToByteSlice}; +use arrow_schema::DataType; +use num::CheckedAdd; + +/// Generic helper to get the last run end value from a run ends array +fn get_last_run_end(run_ends_data: &super::MutableArrayData) -> T { + if run_ends_data.data.len == 0 { + T::default() + } else { + // Convert buffer to typed slice and get the last element + let buffer = Buffer::from(run_ends_data.data.buffer1.as_slice()); + let typed_slice: &[T] = buffer.typed_data(); + if typed_slice.len() >= run_ends_data.data.len { + typed_slice[run_ends_data.data.len - 1] + } else { + T::default() + } + } +} + +/// Extends the `MutableArrayData` with null values. +/// +/// For RunEndEncoded, this adds nulls by extending the run_ends array +/// and values array appropriately. +pub fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) { + if len == 0 { + return; + } + + // For REE, we always need to add a value entry when adding a new run + // The values array should have one entry per run, not per logical element + mutable.child_data[1].extend_nulls(1); + + // Determine the run end type from the data type + let run_end_type = if let DataType::RunEndEncoded(run_ends_field, _) = &mutable.data_type { + run_ends_field.data_type() + } else { + panic!("extend_nulls called on non-RunEndEncoded array"); + }; + + // Use a macro to handle all run end types generically + macro_rules! extend_nulls_impl { + ($run_end_type:ty) => {{ + let last_run_end = get_last_run_end::<$run_end_type>(&mutable.child_data[0]); + let new_value = last_run_end + .checked_add(<$run_end_type as ArrowNativeType>::usize_as(len)) + .expect("run end overflow"); + mutable.child_data[0] + .data + .buffer1 + .extend_from_slice(new_value.to_byte_slice()); + }}; + } + + // Apply the appropriate implementation based on run end type + match run_end_type { + DataType::Int16 => extend_nulls_impl!(i16), + DataType::Int32 => extend_nulls_impl!(i32), + DataType::Int64 => extend_nulls_impl!(i64), + _ => panic!("Invalid run end type for RunEndEncoded array: {run_end_type:?}"), + }; + + mutable.child_data[0].data.len += 1; +} + +/// Build run ends bytes and values range directly for batch processing +fn build_extend_arrays + CheckedAdd>( + buffer: &Buffer, + length: usize, + start: usize, + len: usize, + dest_last_run_end: T, +) -> (Vec, Option<(usize, usize)>) { + let mut run_ends_bytes = Vec::new(); + let mut values_range: Option<(usize, usize)> = None; + let end = start + len; + let mut prev_end = 0; + let mut current_run_end = dest_last_run_end; + + // Convert buffer to typed slice once + let typed_slice: &[T] = buffer.typed_data(); + + for i in 0..length { + if i < typed_slice.len() { + let run_end = typed_slice[i].to_usize().unwrap(); + + if prev_end <= start && run_end > start { + let start_offset = start - prev_end; + let end_offset = if run_end >= end { + end - prev_end + } else { + run_end - prev_end + }; + current_run_end = current_run_end + .checked_add(&T::usize_as(end_offset - start_offset)) + .expect("run end overflow"); + run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice()); + + // Start the range + values_range = Some((i, i + 1)); + } else if prev_end >= start && run_end <= end { + current_run_end = current_run_end + .checked_add(&T::usize_as(run_end - prev_end)) + .expect("run end overflow"); + run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice()); + + // Extend the range + values_range = Some((values_range.expect("Unreachable: values_range cannot be None when prev_end >= start && run_end <= end. \ + If prev_end >= start and run_end > prev_end (required for valid runs), then run_end > start, \ + which means the first condition (prev_end <= start && run_end > start) would have been true \ + and already set values_range to Some.").0, i + 1)); + } else if prev_end < end && run_end >= end { + current_run_end = current_run_end + .checked_add(&T::usize_as(end - prev_end)) + .expect("run end overflow"); + run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice()); + + // Extend the range and break + values_range = Some((values_range.expect("Unreachable: values_range cannot be None when prev_end < end && run_end >= end. \ + Due to sequential processing and monotonic prev_end advancement, if we reach a run \ + that spans beyond the slice end (run_end >= end), at least one previous condition \ + must have matched first to set values_range. Either the first condition matched when \ + the slice started (prev_end <= start && run_end > start), or the second condition \ + matched for runs within the slice (prev_end >= start && run_end <= end).").0, i + 1)); + break; + } + + prev_end = run_end; + if prev_end >= end { + break; + } + } else { + break; + } + } + (run_ends_bytes, values_range) +} + +/// Process extends using batch operations +fn process_extends_batch( + mutable: &mut _MutableArrayData, + source_array_idx: usize, + run_ends_bytes: Vec, + values_range: Option<(usize, usize)>, +) { + if run_ends_bytes.is_empty() { + return; + } + + // Batch extend the run_ends array with all bytes at once + mutable.child_data[0] + .data + .buffer1 + .extend_from_slice(&run_ends_bytes); + mutable.child_data[0].data.len += run_ends_bytes.len() / std::mem::size_of::(); + + // Batch extend the values array using the range + let (start_idx, end_idx) = + values_range.expect("values_range should be Some if run_ends_bytes is not empty"); + mutable.child_data[1].extend(source_array_idx, start_idx, end_idx); +} + +/// Returns a function that extends the run encoded array. +/// +/// It finds the physical indices in the source array that correspond to the logical range to copy, and adjusts the runs to the logical indices of the array to extend. The values are copied from the source array to the destination array verbatim. +pub fn build_extend(array: &ArrayData) -> Extend { + Box::new( + move |mutable: &mut _MutableArrayData, array_idx: usize, start: usize, len: usize| { + if len == 0 { + return; + } + + // We need to analyze the source array's run structure + let source_run_ends = &array.child_data()[0]; + let source_buffer = &source_run_ends.buffers()[0]; + + // Get the run end type from the mutable array + let dest_run_end_type = + if let DataType::RunEndEncoded(run_ends_field, _) = &mutable.data_type { + run_ends_field.data_type() + } else { + panic!("extend called on non-RunEndEncoded mutable array"); + }; + + // Build run ends and values indices directly for batch processing + macro_rules! build_and_process_impl { + ($run_end_type:ty) => {{ + let dest_last_run_end = + get_last_run_end::<$run_end_type>(&mutable.child_data[0]); + let (run_ends_bytes, values_range) = build_extend_arrays::<$run_end_type>( + source_buffer, + source_run_ends.len(), + start, + len, + dest_last_run_end, + ); + process_extends_batch::<$run_end_type>( + mutable, + array_idx, + run_ends_bytes, + values_range, + ); + }}; + } + + match dest_run_end_type { + DataType::Int16 => build_and_process_impl!(i16), + DataType::Int32 => build_and_process_impl!(i32), + DataType::Int64 => build_and_process_impl!(i64), + _ => panic!("Invalid run end type for RunEndEncoded array: {dest_run_end_type:?}",), + } + }, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::transform::MutableArrayData; + use crate::{ArrayData, ArrayDataBuilder}; + use arrow_buffer::Buffer; + use arrow_schema::{DataType, Field}; + use std::sync::Arc; + + fn create_run_array_data(run_ends: Vec, values: ArrayData) -> ArrayData { + let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int32, false)); + let values_field = Arc::new(Field::new("values", values.data_type().clone(), true)); + let data_type = DataType::RunEndEncoded(run_ends_field, values_field); + + let last_run_end = if run_ends.is_empty() { + 0 + } else { + run_ends[run_ends.len() - 1] as usize + }; + + let run_ends_buffer = Buffer::from_vec(run_ends); + let run_ends_data = ArrayDataBuilder::new(DataType::Int32) + .len(run_ends_buffer.len() / std::mem::size_of::()) + .add_buffer(run_ends_buffer) + .build() + .unwrap(); + + ArrayDataBuilder::new(data_type) + .len(last_run_end) + .add_child_data(run_ends_data) + .add_child_data(values) + .build() + .unwrap() + } + + fn create_run_array_data_int16(run_ends: Vec, values: ArrayData) -> ArrayData { + let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int16, false)); + let values_field = Arc::new(Field::new("values", values.data_type().clone(), true)); + let data_type = DataType::RunEndEncoded(run_ends_field, values_field); + + let last_run_end = if run_ends.is_empty() { + 0 + } else { + run_ends[run_ends.len() - 1] as usize + }; + + let run_ends_buffer = Buffer::from_vec(run_ends); + let run_ends_data = ArrayDataBuilder::new(DataType::Int16) + .len(run_ends_buffer.len() / std::mem::size_of::()) + .add_buffer(run_ends_buffer) + .build() + .unwrap(); + + ArrayDataBuilder::new(data_type) + .len(last_run_end) + .add_child_data(run_ends_data) + .add_child_data(values) + .build() + .unwrap() + } + + fn create_run_array_data_int64(run_ends: Vec, values: ArrayData) -> ArrayData { + let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int64, false)); + let values_field = Arc::new(Field::new("values", values.data_type().clone(), true)); + let data_type = DataType::RunEndEncoded(run_ends_field, values_field); + + let last_run_end = if run_ends.is_empty() { + 0 + } else { + run_ends[run_ends.len() - 1] as usize + }; + + let run_ends_buffer = Buffer::from_vec(run_ends); + let run_ends_data = ArrayDataBuilder::new(DataType::Int64) + .len(run_ends_buffer.len() / std::mem::size_of::()) + .add_buffer(run_ends_buffer) + .build() + .unwrap(); + + ArrayDataBuilder::new(data_type) + .len(last_run_end) + .add_child_data(run_ends_data) + .add_child_data(values) + .build() + .unwrap() + } + + fn create_int32_array_data(values: Vec) -> ArrayData { + let buffer = Buffer::from_vec(values); + ArrayDataBuilder::new(DataType::Int32) + .len(buffer.len() / std::mem::size_of::()) + .add_buffer(buffer) + .build() + .unwrap() + } + + fn create_string_dict_array_data(values: Vec<&str>, dict_values: Vec<&str>) -> ArrayData { + // Create dictionary values (strings) + let dict_offsets: Vec = dict_values + .iter() + .scan(0i32, |acc, s| { + let offset = *acc; + *acc += s.len() as i32; + Some(offset) + }) + .chain(std::iter::once( + dict_values.iter().map(|s| s.len()).sum::() as i32, + )) + .collect(); + + let dict_data: Vec = dict_values.iter().flat_map(|s| s.bytes()).collect(); + + let dict_array = ArrayDataBuilder::new(DataType::Utf8) + .len(dict_values.len()) + .add_buffer(Buffer::from_vec(dict_offsets)) + .add_buffer(Buffer::from_vec(dict_data)) + .build() + .unwrap(); + + // Create keys array + let keys: Vec = values + .iter() + .map(|v| dict_values.iter().position(|d| d == v).unwrap() as i32) + .collect(); + + // Create dictionary array + let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + + ArrayDataBuilder::new(dict_type) + .len(values.len()) + .add_buffer(Buffer::from_vec(keys)) + .add_child_data(dict_array) + .build() + .unwrap() + } + + #[test] + fn test_extend_nulls_int32() { + // Create values array with one value + let values = create_int32_array_data(vec![42]); + + // Create REE array with Int32 run ends + let ree_array = create_run_array_data(vec![5], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10); + + mutable.extend_nulls(3); + mutable.extend(0, 0, 5); + mutable.extend_nulls(3); + + // Verify the run ends were extended correctly + let result = mutable.freeze(); + let run_ends_buffer = &result.child_data()[0].buffers()[0]; + let run_ends_slice = run_ends_buffer.as_slice(); + + // Should have three run ends now + assert_eq!(result.child_data()[0].len(), 3); + let first_run_end = i32::from_ne_bytes(run_ends_slice[0..4].try_into().unwrap()); + let second_run_end = i32::from_ne_bytes(run_ends_slice[4..8].try_into().unwrap()); + let third_run_end = i32::from_ne_bytes(run_ends_slice[8..12].try_into().unwrap()); + assert_eq!(first_run_end, 3); + assert_eq!(second_run_end, 8); + assert_eq!(third_run_end, 11); + + // Verify the values array was extended correctly + assert_eq!(result.child_data()[1].len(), 3); // Should match run ends length + let values_buffer = &result.child_data()[1].buffers()[0]; + let values_slice = values_buffer.as_slice(); + + // Check the values in the buffer + let second_value = i32::from_ne_bytes(values_slice[4..8].try_into().unwrap()); + + // Second value should be the original value from the source array + assert_eq!(second_value, 42); + + // Verify the validity buffer shows the correct null pattern + let values_array = &result.child_data()[1]; + // First value should be null + assert!(values_array.is_null(0)); + // Second value should be valid + assert!(values_array.is_valid(1)); + // Third value should be null + assert!(values_array.is_null(2)); + } + + #[test] + fn test_extend_nulls_int16() { + // Create values array with one value + let values = create_int32_array_data(vec![42]); + + // Create REE array with Int16 run ends + let ree_array = create_run_array_data_int16(vec![5i16], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10); + + // First, we need to copy the existing data + mutable.extend(0, 0, 5); + + // Then add nulls + mutable.extend_nulls(3); + + // Verify the run ends were extended correctly + let result = mutable.freeze(); + let run_ends_buffer = &result.child_data()[0].buffers()[0]; + let run_ends_slice = run_ends_buffer.as_slice(); + + // Should have two run ends now: original 5 and new 8 (5 + 3) + assert_eq!(result.child_data()[0].len(), 2); + let first_run_end = i16::from_ne_bytes(run_ends_slice[0..2].try_into().unwrap()); + let second_run_end = i16::from_ne_bytes(run_ends_slice[2..4].try_into().unwrap()); + assert_eq!(first_run_end, 5); + assert_eq!(second_run_end, 8); + } + + #[test] + fn test_extend_nulls_int64() { + // Create values array with one value + let values = create_int32_array_data(vec![42]); + + // Create REE array with Int64 run ends + let ree_array = create_run_array_data_int64(vec![5i64], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10); + + // First, we need to copy the existing data + mutable.extend(0, 0, 5); + + // Then add nulls + mutable.extend_nulls(3); + + // Verify the run ends were extended correctly + let result = mutable.freeze(); + let run_ends_buffer = &result.child_data()[0].buffers()[0]; + let run_ends_slice = run_ends_buffer.as_slice(); + + // Should have two run ends now: original 5 and new 8 (5 + 3) + assert_eq!(result.child_data()[0].len(), 2); + let first_run_end = i64::from_ne_bytes(run_ends_slice[0..8].try_into().unwrap()); + let second_run_end = i64::from_ne_bytes(run_ends_slice[8..16].try_into().unwrap()); + assert_eq!(first_run_end, 5); + assert_eq!(second_run_end, 8); + } + + #[test] + fn test_extend_int32() { + // Create a simple REE array with Int32 run ends + let values = create_int32_array_data(vec![10, 20]); + + // Array: [10, 10, 20, 20, 20] (run_ends = [2, 5]) + let ree_array = create_run_array_data(vec![2, 5], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10); + + // Extend the entire array + mutable.extend(0, 0, 5); + + let result = mutable.freeze(); + + // Should have extended correctly + assert_eq!(result.len(), 5); // All 5 elements + + // Basic validation that we have the right structure + assert!(!result.child_data()[0].is_empty()); // Should have at least one run + assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length + } + + #[test] + fn test_extend_empty() { + let values = create_int32_array_data(vec![]); + let ree_array = create_run_array_data(vec![], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10); + mutable.extend(0, 0, 0); + + let result = mutable.freeze(); + assert_eq!(result.len(), 0); + assert_eq!(result.child_data()[0].len(), 0); + } + + #[test] + fn test_build_extend_arrays_int16() { + let buffer = Buffer::from_vec(vec![3i16, 5i16, 8i16]); + let (run_ends_bytes, values_range) = build_extend_arrays::(&buffer, 3, 2, 4, 0i16); + + // Logical array: [A, A, A, B, B, C, C, C] + // Requesting indices 2-6 should give us: + // - Part of first run (index 2) -> length 1 + // - All of second run -> length 2 + // - Part of third run -> length 1 + // Total length = 4, so run ends should be [1, 3, 4] + assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::()); + assert_eq!(values_range, Some((0, 3))); + + // Verify the bytes represent [1i16, 3i16, 4i16] + let expected_bytes = [1i16, 3i16, 4i16] + .iter() + .flat_map(|&val| val.to_ne_bytes()) + .collect::>(); + assert_eq!(run_ends_bytes, expected_bytes); + } + + #[test] + fn test_build_extend_arrays_int64() { + let buffer = Buffer::from_vec(vec![3i64, 5i64, 8i64]); + let (run_ends_bytes, values_range) = build_extend_arrays::(&buffer, 3, 2, 4, 0i64); + + // Same logic as above but with i64 + assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::()); + assert_eq!(values_range, Some((0, 3))); + + // Verify the bytes represent [1i64, 3i64, 4i64] + let expected_bytes = [1i64, 3i64, 4i64] + .iter() + .flat_map(|&val| val.to_ne_bytes()) + .collect::>(); + assert_eq!(run_ends_bytes, expected_bytes); + } + + #[test] + fn test_extend_string_dict() { + // Create a dictionary array with string values: ["hello", "world"] + let dict_values = vec!["hello", "world"]; + let values = create_string_dict_array_data(vec!["hello", "world"], dict_values); + + // Create REE array: [hello, hello, world, world, world] (run_ends = [2, 5]) + let ree_array = create_run_array_data(vec![2, 5], values); + + let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10); + + // Extend the entire array + mutable.extend(0, 0, 5); + + let result = mutable.freeze(); + + // Should have extended correctly + assert_eq!(result.len(), 5); // All 5 elements + + // Basic validation that we have the right structure + assert!(!result.child_data()[0].is_empty()); // Should have at least one run + assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length + + // Should have 2 runs since we have 2 different values + assert_eq!(result.child_data()[0].len(), 2); + assert_eq!(result.child_data()[1].len(), 2); + } + + #[test] + #[should_panic(expected = "run end overflow")] + fn test_extend_nulls_overflow_i16() { + let values = create_int32_array_data(vec![42]); + // Start with run end close to max to set up overflow condition + let ree_array = create_run_array_data_int16(vec![5], values); + let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10); + + // Extend the original data first to initialize state + mutable.extend(0, 0, 5_usize); + + // This should cause overflow: i16::MAX + 5 > i16::MAX + mutable.extend_nulls(i16::MAX as usize); + } + + #[test] + #[should_panic(expected = "run end overflow")] + fn test_extend_nulls_overflow_i32() { + let values = create_int32_array_data(vec![42]); + // Start with run end close to max to set up overflow condition + let ree_array = create_run_array_data(vec![10], values); + let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10); + + // Extend the original data first to initialize state + mutable.extend(0, 0, 10_usize); + + // This should cause overflow: (i32::MAX - 10) + 20 > i32::MAX + mutable.extend_nulls(i32::MAX as usize); + } + + #[test] + #[should_panic(expected = "run end overflow")] + fn test_build_extend_overflow_i16() { + // Create a source array with small run that will cause overflow when added + let values = create_int32_array_data(vec![10]); + let source_array = create_run_array_data_int16(vec![20], values); + + // Create a destination array with run end close to max + let dest_values = create_int32_array_data(vec![42]); + let dest_array = create_run_array_data_int16(vec![i16::MAX - 5], dest_values); + + let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10); + + // First extend the destination array to set up state + mutable.extend(1, 0, (i16::MAX - 5) as usize); + + // This should cause overflow: (i16::MAX - 5) + 20 > i16::MAX + mutable.extend(0, 0, 20); + } + + #[test] + #[should_panic(expected = "run end overflow")] + fn test_build_extend_overflow_i32() { + // Create a source array with small run that will cause overflow when added + let values = create_int32_array_data(vec![10]); + let source_array = create_run_array_data(vec![100], values); + + // Create a destination array with run end close to max + let dest_values = create_int32_array_data(vec![42]); + let dest_array = create_run_array_data(vec![i32::MAX - 50], dest_values); + + let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10); + + // First extend the destination array to set up state + mutable.extend(1, 0, (i32::MAX - 50) as usize); + + // This should cause overflow: (i32::MAX - 50) + 100 > i32::MAX + mutable.extend(0, 0, 100); + } +} diff --git a/arrow-flight/Cargo.toml b/arrow-flight/Cargo.toml index fbb295036a9b..ca0d1c5e4b3d 100644 --- a/arrow-flight/Cargo.toml +++ b/arrow-flight/Cargo.toml @@ -20,7 +20,7 @@ name = "arrow-flight" description = "Apache Arrow Flight" version = { workspace = true } edition = { workspace = true } -rust-version = "1.71.1" +rust-version = { workspace = true } authors = { workspace = true } homepage = { workspace = true } repository = { workspace = true } @@ -48,7 +48,7 @@ prost = { version = "0.13.1", default-features = false, features = ["prost-deriv # For Timestamp type prost-types = { version = "0.13.1", default-features = false } tokio = { version = "1.0", default-features = false, features = ["macros", "rt", "rt-multi-thread"], optional = true } -tonic = { version = "0.12.3", default-features = false, features = ["transport", "codegen", "prost"] } +tonic = { version = "0.13", default-features = false, features = ["transport", "codegen", "prost", "router"] } # CLI-related dependencies anyhow = { version = "1.0", optional = true } @@ -61,10 +61,16 @@ all-features = true [features] default = [] -flight-sql-experimental = ["dep:arrow-arith", "dep:arrow-data", "dep:arrow-ord", "dep:arrow-row", "dep:arrow-select", "dep:arrow-string", "dep:once_cell", "dep:paste"] -tls = ["tonic/tls"] +flight-sql = ["dep:arrow-arith", "dep:arrow-data", "dep:arrow-ord", "dep:arrow-row", "dep:arrow-select", "dep:arrow-string", "dep:once_cell", "dep:paste"] +# TODO: Remove in the next release +flight-sql-experimental = ["flight-sql"] +tls-aws-lc= ["tonic/tls-aws-lc"] +tls-native-roots = ["tonic/tls-native-roots"] +tls-ring = ["tonic/tls-ring"] +tls-webpki-roots = ["tonic/tls-webpki-roots"] + # Enable CLI tools -cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-roots", "dep:anyhow", "dep:clap", "dep:tracing-log", "dep:tracing-subscriber"] +cli = ["arrow-array/chrono-tz", "arrow-cast/prettyprint", "tonic/tls-webpki-roots", "dep:anyhow", "dep:clap", "dep:tracing-log", "dep:tracing-subscriber", "dep:tokio"] [dev-dependencies] arrow-cast = { workspace = true, features = ["prettyprint"] } @@ -83,18 +89,18 @@ uuid = { version = "1.10.0", features = ["v4"] } [[example]] name = "flight_sql_server" -required-features = ["flight-sql-experimental", "tls"] +required-features = ["flight-sql", "tls-ring"] [[bin]] name = "flight_sql_client" -required-features = ["cli", "flight-sql-experimental", "tls"] +required-features = ["cli", "flight-sql", "tls-ring"] [[test]] name = "flight_sql_client" path = "tests/flight_sql_client.rs" -required-features = ["flight-sql-experimental", "tls"] +required-features = ["flight-sql", "tls-ring"] [[test]] name = "flight_sql_client_cli" path = "tests/flight_sql_client_cli.rs" -required-features = ["cli", "flight-sql-experimental", "tls"] +required-features = ["cli", "flight-sql", "tls-ring"] diff --git a/arrow-flight/README.md b/arrow-flight/README.md index 3ffc8780c2f8..1cd8f5cfe21b 100644 --- a/arrow-flight/README.md +++ b/arrow-flight/README.md @@ -43,10 +43,16 @@ that demonstrate how to build a Flight server implemented with [tonic](https://d ## Feature Flags -- `flight-sql-experimental`: Enables experimental support for - [Apache Arrow FlightSQL], a protocol for interacting with SQL databases. +- `flight-sql`: Support for [Apache Arrow FlightSQL], a protocol for interacting with SQL databases. -- `tls`: Enables `tls` on `tonic` +You can enable TLS using the following features (not enabled by default) + +- `tls-aws-lc`: enables [tonic feature] `tls-aws-lc` +- `tls-native-roots`: enables [tonic feature] `tls-native-roots` +- `tls-ring`: enables [tonic feature] `tls-ring` +- `tls-webpki`: enables [tonic feature] `tls-webpki-roots` + +[tonic feature]: https://docs.rs/tonic/latest/tonic/#feature-flags ## CLI @@ -55,7 +61,7 @@ This crates offers a basic [Apache Arrow FlightSQL] command line interface. The client can be installed from the repository: ```console -$ cargo install --features=cli,flight-sql-experimental,tls --bin=flight_sql_client --path=. --locked +$ cargo install --features=cli,flight-sql,tls --bin=flight_sql_client --path=. --locked ``` The client comes with extensive help text: diff --git a/arrow-flight/examples/flight_sql_server.rs b/arrow-flight/examples/flight_sql_server.rs index 657298b4a8b3..f2837de7c788 100644 --- a/arrow-flight/examples/flight_sql_server.rs +++ b/arrow-flight/examples/flight_sql_server.rs @@ -112,6 +112,7 @@ static TABLES: Lazy> = Lazy::new(|| vec!["flight_sql.example.t pub struct FlightSqlServiceImpl {} impl FlightSqlServiceImpl { + #[allow(clippy::result_large_err)] fn check_token(&self, req: &Request) -> Result<(), Status> { let metadata = req.metadata(); let auth = metadata.get("authorization").ok_or_else(|| { @@ -188,7 +189,7 @@ impl FlightSqlService for FlightSqlServiceImpl { let result = Ok(result); let output = futures::stream::iter(vec![result]); - let token = format!("Bearer {}", FAKE_TOKEN); + let token = format!("Bearer {FAKE_TOKEN}"); let mut response: Response + Send>>> = Response::new(Box::pin(output)); response.metadata_mut().append( @@ -744,7 +745,7 @@ async fn main() -> Result<(), Box> { let addr_str = "0.0.0.0:50051"; let addr = addr_str.parse()?; - println!("Listening on {:?}", addr); + println!("Listening on {addr:?}"); if std::env::var("USE_TLS").ok().is_some() { let cert = std::fs::read_to_string("arrow-flight/examples/data/server.pem")?; @@ -813,7 +814,7 @@ mod tests { async fn bind_tcp() -> (TcpIncoming, SocketAddr) { let listener = TcpListener::bind("0.0.0.0:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let incoming = TcpIncoming::from_listener(listener, true, None).unwrap(); + let incoming = TcpIncoming::from(listener).with_nodelay(Some(true)); (incoming, addr) } diff --git a/arrow-flight/gen/Cargo.toml b/arrow-flight/gen/Cargo.toml index 6358227a8912..9e509e4fad43 100644 --- a/arrow-flight/gen/Cargo.toml +++ b/arrow-flight/gen/Cargo.toml @@ -32,5 +32,5 @@ publish = false [dependencies] # Pin specific version of the tonic-build dependencies to avoid auto-generated # (and checked in) arrow.flight.protocol.rs from changing -prost-build = { version = "=0.13.4", default-features = false } -tonic-build = { version = "=0.12.3", default-features = false, features = ["transport", "prost"] } +prost-build = { version = "=0.13.5", default-features = false } +tonic-build = { version = "=0.13.1", default-features = false, features = ["transport", "prost"] } diff --git a/arrow-flight/src/arrow.flight.protocol.rs b/arrow-flight/src/arrow.flight.protocol.rs index 0cd4f6948b77..a08ea01105e5 100644 --- a/arrow-flight/src/arrow.flight.protocol.rs +++ b/arrow-flight/src/arrow.flight.protocol.rs @@ -448,7 +448,7 @@ pub mod flight_service_client { } impl FlightServiceClient where - T: tonic::client::GrpcService, + T: tonic::client::GrpcService, T::Error: Into, T::ResponseBody: Body + std::marker::Send + 'static, ::Error: Into + std::marker::Send, @@ -469,13 +469,13 @@ pub mod flight_service_client { F: tonic::service::Interceptor, T::ResponseBody: Default, T: tonic::codegen::Service< - http::Request, + http::Request, Response = http::Response< - >::ResponseBody, + >::ResponseBody, >, >, , + http::Request, >>::Error: Into + std::marker::Send + std::marker::Sync, { FlightServiceClient::new(InterceptedService::new(inner, interceptor)) @@ -1098,7 +1098,7 @@ pub mod flight_service_server { B: Body + std::marker::Send + 'static, B::Error: Into + std::marker::Send + 'static, { - type Response = http::Response; + type Response = http::Response; type Error = std::convert::Infallible; type Future = BoxFuture; fn poll_ready( @@ -1571,7 +1571,9 @@ pub mod flight_service_server { } _ => { Box::pin(async move { - let mut response = http::Response::new(empty_body()); + let mut response = http::Response::new( + tonic::body::Body::default(), + ); let headers = response.headers_mut(); headers .insert( diff --git a/arrow-flight/src/client.rs b/arrow-flight/src/client.rs index 97d9899a9fb0..9b4c10e9a093 100644 --- a/arrow-flight/src/client.rs +++ b/arrow-flight/src/client.rs @@ -210,7 +210,7 @@ impl FlightClient { let (response_stream, trailers) = extract_lazy_trailers(response_stream); Ok(FlightRecordBatchStream::new_from_flight_data( - response_stream.map_err(FlightError::Tonic), + response_stream.map_err(|status| status.into()), ) .with_headers(md) .with_trailers(trailers)) @@ -470,7 +470,7 @@ impl FlightClient { .list_flights(request) .await? .into_inner() - .map_err(FlightError::Tonic); + .map_err(|status| status.into()); Ok(response.boxed()) } @@ -537,7 +537,7 @@ impl FlightClient { .list_actions(request) .await? .into_inner() - .map_err(FlightError::Tonic); + .map_err(|status| status.into()); Ok(action_stream.boxed()) } @@ -575,7 +575,7 @@ impl FlightClient { .do_action(request) .await? .into_inner() - .map_err(FlightError::Tonic) + .map_err(|status| status.into()) .map(|r| { r.map(|r| { // unwrap inner bytes diff --git a/arrow-flight/src/decode.rs b/arrow-flight/src/decode.rs index 7bafc384306b..70ce35a98952 100644 --- a/arrow-flight/src/decode.rs +++ b/arrow-flight/src/decode.rs @@ -138,12 +138,6 @@ impl FlightRecordBatchStream { self.trailers.as_ref().and_then(|trailers| trailers.get()) } - /// Has a message defining the schema been received yet? - #[deprecated = "use schema().is_some() instead"] - pub fn got_schema(&self) -> bool { - self.schema().is_some() - } - /// Return schema for the stream, if it has been received pub fn schema(&self) -> Option<&SchemaRef> { self.inner.schema() @@ -295,7 +289,7 @@ impl FlightDataDecoder { )); }; - let buffer = Buffer::from_bytes(data.data_body.into()); + let buffer = Buffer::from(data.data_body); let dictionary_batch = message.header_as_dictionary_batch().ok_or_else(|| { FlightError::protocol( "Could not get dictionary batch from DictionaryBatch message", diff --git a/arrow-flight/src/encode.rs b/arrow-flight/src/encode.rs index 315b7b3cb6e5..49910a3ee2b0 100644 --- a/arrow-flight/src/encode.rs +++ b/arrow-flight/src/encode.rs @@ -535,15 +535,13 @@ fn prepare_field_for_flight( ) .with_metadata(field.metadata().clone()) } else { - #[allow(deprecated)] - let dict_id = dictionary_tracker.set_dict_id(field.as_ref()); - + dictionary_tracker.next_dict_id(); #[allow(deprecated)] Field::new_dict( field.name(), field.data_type().clone(), field.is_nullable(), - dict_id, + 0, field.dict_is_ordered().unwrap_or_default(), ) .with_metadata(field.metadata().clone()) @@ -585,14 +583,13 @@ fn prepare_schema_for_flight( ) .with_metadata(field.metadata().clone()) } else { - #[allow(deprecated)] - let dict_id = dictionary_tracker.set_dict_id(field.as_ref()); + dictionary_tracker.next_dict_id(); #[allow(deprecated)] Field::new_dict( field.name(), field.data_type().clone(), field.is_nullable(), - dict_id, + 0, field.dict_is_ordered().unwrap_or_default(), ) .with_metadata(field.metadata().clone()) @@ -654,16 +651,10 @@ struct FlightIpcEncoder { impl FlightIpcEncoder { fn new(options: IpcWriteOptions, error_on_replacement: bool) -> Self { - #[allow(deprecated)] - let preserve_dict_id = options.preserve_dict_id(); Self { options, data_gen: IpcDataGenerator::default(), - #[allow(deprecated)] - dictionary_tracker: DictionaryTracker::new_with_preserve_dict_id( - error_on_replacement, - preserve_dict_id, - ), + dictionary_tracker: DictionaryTracker::new(error_on_replacement), } } @@ -1547,9 +1538,8 @@ mod tests { async fn verify_flight_round_trip(mut batches: Vec) { let expected_schema = batches.first().unwrap().schema(); - #[allow(deprecated)] let encoder = FlightDataEncoderBuilder::default() - .with_options(IpcWriteOptions::default().with_preserve_dict_id(false)) + .with_options(IpcWriteOptions::default()) .with_dictionary_handling(DictionaryHandling::Resend) .build(futures::stream::iter(batches.clone().into_iter().map(Ok))); @@ -1575,8 +1565,7 @@ mod tests { HashMap::from([("some_key".to_owned(), "some_value".to_owned())]), ); - #[allow(deprecated)] - let mut dictionary_tracker = DictionaryTracker::new_with_preserve_dict_id(false, true); + let mut dictionary_tracker = DictionaryTracker::new(false); let got = prepare_schema_for_flight(&schema, &mut dictionary_tracker, false); assert!(got.metadata().contains_key("some_key")); @@ -1606,9 +1595,7 @@ mod tests { options: &IpcWriteOptions, ) -> (Vec, FlightData) { let data_gen = IpcDataGenerator::default(); - #[allow(deprecated)] - let mut dictionary_tracker = - DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id()); + let mut dictionary_tracker = DictionaryTracker::new(false); let (encoded_dictionaries, encoded_batch) = data_gen .encoded_batch(batch, &mut dictionary_tracker, options) @@ -1695,9 +1682,9 @@ mod tests { #[tokio::test] async fn flight_data_size_even() { - let s1 = StringArray::from_iter_values(std::iter::repeat(".10 bytes.").take(1024)); + let s1 = StringArray::from_iter_values(std::iter::repeat_n(".10 bytes.", 1024)); let i1 = Int16Array::from_iter_values(0..1024); - let s2 = StringArray::from_iter_values(std::iter::repeat("6bytes").take(1024)); + let s2 = StringArray::from_iter_values(std::iter::repeat_n("6bytes", 1024)); let i2 = Int64Array::from_iter_values(0..1024); let batch = RecordBatch::try_from_iter(vec![ @@ -1708,7 +1695,7 @@ mod tests { ]) .unwrap(); - verify_encoded_split(batch, 112).await; + verify_encoded_split(batch, 120).await; } #[tokio::test] @@ -1719,7 +1706,7 @@ mod tests { // overage is much higher than ideal // https://github.com/apache/arrow-rs/issues/3478 - verify_encoded_split(batch, 4304).await; + verify_encoded_split(batch, 4312).await; } #[tokio::test] @@ -1755,7 +1742,7 @@ mod tests { // 5k over limit (which is 2x larger than limit of 5k) // overage is much higher than ideal // https://github.com/apache/arrow-rs/issues/3478 - verify_encoded_split(batch, 5800).await; + verify_encoded_split(batch, 5808).await; } #[tokio::test] @@ -1771,7 +1758,7 @@ mod tests { let batch = RecordBatch::try_from_iter(vec![("a1", Arc::new(array) as _)]).unwrap(); - verify_encoded_split(batch, 48).await; + verify_encoded_split(batch, 56).await; } #[tokio::test] @@ -1785,7 +1772,7 @@ mod tests { // overage is much higher than ideal // https://github.com/apache/arrow-rs/issues/3478 - verify_encoded_split(batch, 3328).await; + verify_encoded_split(batch, 3336).await; } #[tokio::test] @@ -1799,7 +1786,7 @@ mod tests { // overage is much higher than ideal // https://github.com/apache/arrow-rs/issues/3478 - verify_encoded_split(batch, 5280).await; + verify_encoded_split(batch, 5288).await; } #[tokio::test] @@ -1824,7 +1811,7 @@ mod tests { // overage is much higher than ideal // https://github.com/apache/arrow-rs/issues/3478 - verify_encoded_split(batch, 4128).await; + verify_encoded_split(batch, 4136).await; } /// Return size, in memory of flight data @@ -1833,7 +1820,7 @@ mod tests { .flight_descriptor .as_ref() .map(|descriptor| { - let path_len: usize = descriptor.path.iter().map(|p| p.as_bytes().len()).sum(); + let path_len: usize = descriptor.path.iter().map(|p| p.len()).sum(); std::mem::size_of_val(descriptor) + descriptor.cmd.len() + path_len }) diff --git a/arrow-flight/src/error.rs b/arrow-flight/src/error.rs index 499706e1ede7..d5ac568e9788 100644 --- a/arrow-flight/src/error.rs +++ b/arrow-flight/src/error.rs @@ -27,7 +27,7 @@ pub enum FlightError { /// Returned when functionality is not yet available. NotYetImplemented(String), /// Error from the underlying tonic library - Tonic(tonic::Status), + Tonic(Box), /// Some unexpected message was received ProtocolError(String), /// An error occurred during decoding @@ -51,12 +51,12 @@ impl FlightError { impl std::fmt::Display for FlightError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - FlightError::Arrow(source) => write!(f, "Arrow error: {}", source), - FlightError::NotYetImplemented(desc) => write!(f, "Not yet implemented: {}", desc), - FlightError::Tonic(source) => write!(f, "Tonic error: {}", source), - FlightError::ProtocolError(desc) => write!(f, "Protocol error: {}", desc), - FlightError::DecodeError(desc) => write!(f, "Decode error: {}", desc), - FlightError::ExternalError(source) => write!(f, "External error: {}", source), + FlightError::Arrow(source) => write!(f, "Arrow error: {source}"), + FlightError::NotYetImplemented(desc) => write!(f, "Not yet implemented: {desc}"), + FlightError::Tonic(source) => write!(f, "Tonic error: {source}"), + FlightError::ProtocolError(desc) => write!(f, "Protocol error: {desc}"), + FlightError::DecodeError(desc) => write!(f, "Decode error: {desc}"), + FlightError::ExternalError(source) => write!(f, "External error: {source}"), } } } @@ -74,7 +74,7 @@ impl Error for FlightError { impl From for FlightError { fn from(status: tonic::Status) -> Self { - Self::Tonic(status) + Self::Tonic(Box::new(status)) } } @@ -91,7 +91,7 @@ impl From for tonic::Status { match value { FlightError::Arrow(e) => tonic::Status::internal(e.to_string()), FlightError::NotYetImplemented(e) => tonic::Status::internal(e), - FlightError::Tonic(status) => status, + FlightError::Tonic(status) => *status, FlightError::ProtocolError(e) => tonic::Status::internal(e), FlightError::DecodeError(e) => tonic::Status::internal(e), FlightError::ExternalError(e) => tonic::Status::internal(e.to_string()), @@ -147,4 +147,10 @@ mod test { let source = root_error.downcast_ref::().unwrap(); assert!(matches!(source, FlightError::DecodeError(_))); } + + #[test] + fn test_error_size() { + // use Box in variants to keep this size down + assert_eq!(std::mem::size_of::(), 32); + } } diff --git a/arrow-flight/src/lib.rs b/arrow-flight/src/lib.rs index 1dd2700794f3..8043d5b4a72b 100644 --- a/arrow-flight/src/lib.rs +++ b/arrow-flight/src/lib.rs @@ -32,10 +32,16 @@ //! 2. Low level [tonic] generated [`flight_service_client`] and //! [`flight_service_server`]. //! -//! 3. Experimental support for [Flight SQL] in [`sql`]. Requires the -//! `flight-sql-experimental` feature of this crate to be activated. +//! 3. Support for [Flight SQL] in [`sql`]. Requires the +//! `flight-sql` feature of this crate to be activated. //! //! [Flight SQL]: https://arrow.apache.org/docs/format/FlightSql.html + +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![allow(rustdoc::invalid_html_tags)] #![warn(missing_docs)] // The unused_crate_dependencies lint does not work well for crates defining additional examples/bin targets @@ -122,7 +128,7 @@ mod trailers; pub mod utils; -#[cfg(feature = "flight-sql-experimental")] +#[cfg(feature = "flight-sql")] pub mod sql; mod streams; @@ -143,9 +149,7 @@ pub struct IpcMessage(pub Bytes); fn flight_schema_as_encoded_data(arrow_schema: &Schema, options: &IpcWriteOptions) -> EncodedData { let data_gen = writer::IpcDataGenerator::default(); - #[allow(deprecated)] - let mut dict_tracker = - writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id()); + let mut dict_tracker = writer::DictionaryTracker::new(false); data_gen.schema_to_bytes_with_dictionary_tracker(arrow_schema, &mut dict_tracker, options) } @@ -892,4 +896,26 @@ mod tests { let des_schema: Schema = (&result).try_into().unwrap(); assert_eq!(schema, des_schema); } + + #[test] + fn test_dict_schema() { + // Test for https://github.com/apache/arrow-rs/issues/7058 + let schema = Schema::new(vec![ + Field::new( + "a", + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + false, + ), + Field::new( + "b", + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + false, + ), + ]); + + let flight_info = FlightInfo::new().try_with_schema(&schema).unwrap(); + + let new_schema = Schema::try_from(flight_info).unwrap(); + assert_eq!(schema, new_schema); + } } diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs index a6e228737b3f..6791b68b757d 100644 --- a/arrow-flight/src/sql/client.rs +++ b/arrow-flight/src/sql/client.rs @@ -309,7 +309,7 @@ impl FlightSqlServiceClient { let (response_stream, trailers) = extract_lazy_trailers(response_stream); Ok(FlightRecordBatchStream::new_from_flight_data( - response_stream.map_err(FlightError::Tonic), + response_stream.map_err(|status| status.into()), ) .with_headers(md) .with_trailers(trailers)) @@ -721,7 +721,7 @@ pub fn arrow_data_from_flight_data( let dictionaries_by_field = HashMap::new(); let record_batch = read_record_batch( - &Buffer::from_bytes(flight_data.data_body.into()), + &Buffer::from(flight_data.data_body), ipc_record_batch, arrow_schema_ref.clone(), &dictionaries_by_field, diff --git a/arrow-flight/src/sql/metadata/db_schemas.rs b/arrow-flight/src/sql/metadata/db_schemas.rs index 303d11cd74ca..68e8b497336e 100644 --- a/arrow-flight/src/sql/metadata/db_schemas.rs +++ b/arrow-flight/src/sql/metadata/db_schemas.rs @@ -38,7 +38,7 @@ use crate::sql::CommandGetDbSchemas; /// Builds rows like this: /// /// * catalog_name: utf8, -/// * db_schema_name: utf8, +/// * db_schema_name: utf8 not null pub struct GetDbSchemasBuilder { // Specifies the Catalog to search for the tables. // - An empty string retrieves those without a catalog. @@ -177,7 +177,7 @@ fn get_db_schemas_schema() -> SchemaRef { /// The schema for GetDbSchemas static GET_DB_SCHEMAS_SCHEMA: Lazy = Lazy::new(|| { Arc::new(Schema::new(vec![ - Field::new("catalog_name", DataType::Utf8, false), + Field::new("catalog_name", DataType::Utf8, true), Field::new("db_schema_name", DataType::Utf8, false), ])) }); diff --git a/arrow-flight/src/sql/metadata/mod.rs b/arrow-flight/src/sql/metadata/mod.rs index fd71149a3180..66c12fce9af4 100644 --- a/arrow-flight/src/sql/metadata/mod.rs +++ b/arrow-flight/src/sql/metadata/mod.rs @@ -70,8 +70,7 @@ mod tests { let actual_lines: Vec<_> = formatted.trim().lines().collect(); assert_eq!( &actual_lines, expected_lines, - "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n", - expected_lines, actual_lines + "\n\nexpected:\n\n{expected_lines:#?}\nactual:\n\n{actual_lines:#?}\n\n", ); } } diff --git a/arrow-flight/src/sql/metadata/tables.rs b/arrow-flight/src/sql/metadata/tables.rs index 7ffb76fa1d5f..2cd16fdc231e 100644 --- a/arrow-flight/src/sql/metadata/tables.rs +++ b/arrow-flight/src/sql/metadata/tables.rs @@ -291,8 +291,8 @@ fn get_tables_schema(include_schema: bool) -> SchemaRef { /// The schema for GetTables without `table_schema` column static GET_TABLES_SCHEMA_WITHOUT_TABLE_SCHEMA: Lazy = Lazy::new(|| { Arc::new(Schema::new(vec![ - Field::new("catalog_name", DataType::Utf8, false), - Field::new("db_schema_name", DataType::Utf8, false), + Field::new("catalog_name", DataType::Utf8, true), + Field::new("db_schema_name", DataType::Utf8, true), Field::new("table_name", DataType::Utf8, false), Field::new("table_type", DataType::Utf8, false), ])) @@ -301,8 +301,8 @@ static GET_TABLES_SCHEMA_WITHOUT_TABLE_SCHEMA: Lazy = Lazy::new(|| { /// The schema for GetTables with `table_schema` column static GET_TABLES_SCHEMA_WITH_TABLE_SCHEMA: Lazy = Lazy::new(|| { Arc::new(Schema::new(vec![ - Field::new("catalog_name", DataType::Utf8, false), - Field::new("db_schema_name", DataType::Utf8, false), + Field::new("catalog_name", DataType::Utf8, true), + Field::new("db_schema_name", DataType::Utf8, true), Field::new("table_name", DataType::Utf8, false), Field::new("table_type", DataType::Utf8, false), Field::new("table_schema", DataType::Binary, false), diff --git a/arrow-flight/src/sql/mod.rs b/arrow-flight/src/sql/mod.rs index 94bb96a4f852..955f1904a6d6 100644 --- a/arrow-flight/src/sql/mod.rs +++ b/arrow-flight/src/sql/mod.rs @@ -45,7 +45,6 @@ use prost::Message; #[allow(clippy::all)] mod gen { - #![allow(rustdoc::unportable_markdown)] // Since this file is auto-generated, we suppress all warnings #![allow(missing_docs)] include!("arrow.flight.protocol.sql.rs"); @@ -114,6 +113,8 @@ pub mod client; pub mod metadata; pub mod server; +pub use crate::streams::FallibleRequestStream; + /// ProstMessageExt are useful utility methods for prost::Message types pub trait ProstMessageExt: prost::Message + Default { /// type_url for this Message diff --git a/arrow-flight/src/sql/server.rs b/arrow-flight/src/sql/server.rs index 8ab8a16dbb50..da5dc9945eee 100644 --- a/arrow-flight/src/sql/server.rs +++ b/arrow-flight/src/sql/server.rs @@ -17,12 +17,9 @@ //! Helper trait [`FlightSqlService`] for implementing a [`FlightService`] that implements FlightSQL. +use std::fmt::{Display, Formatter}; use std::pin::Pin; -use futures::{stream::Peekable, Stream, StreamExt}; -use prost::Message; -use tonic::{Request, Response, Status, Streaming}; - use super::{ ActionBeginSavepointRequest, ActionBeginSavepointResult, ActionBeginTransactionRequest, ActionBeginTransactionResult, ActionCancelQueryRequest, ActionCancelQueryResult, @@ -41,6 +38,9 @@ use crate::{ FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket, }; +use futures::{stream::Peekable, Stream, StreamExt}; +use prost::Message; +use tonic::{Request, Response, Status, Streaming}; pub(crate) static CREATE_PREPARED_STATEMENT: &str = "CreatePreparedStatement"; pub(crate) static CLOSE_PREPARED_STATEMENT: &str = "ClosePreparedStatement"; @@ -386,6 +386,15 @@ pub trait FlightSqlService: Sync + Send + Sized + 'static { ))) } + /// Implementors may override to handle do_put errors + async fn do_put_error_callback( + &self, + _request: Request, + error: DoPutError, + ) -> Result::DoPutStream>, Status> { + Err(Status::unimplemented(format!("Unhandled Error: {error}"))) + } + /// Execute an update SQL statement. async fn do_put_statement_update( &self, @@ -710,10 +719,21 @@ where // we wrap this stream in a `Peekable` one, which allows us to peek at // the first message without discarding it. let mut request = request.map(PeekableFlightDataStream::new); - let cmd = Pin::new(request.get_mut()).peek().await.unwrap().clone()?; + let mut stream = Pin::new(request.get_mut()); + + let peeked_item = stream.peek().await.cloned(); + let Some(cmd) = peeked_item else { + return self + .do_put_error_callback(request, DoPutError::MissingCommand) + .await; + }; - let message = - Any::decode(&*cmd.flight_descriptor.unwrap().cmd).map_err(decode_error_to_status)?; + let Some(flight_descriptor) = cmd?.flight_descriptor else { + return self + .do_put_error_callback(request, DoPutError::MissingFlightDescriptor) + .await; + }; + let message = Any::decode(flight_descriptor.cmd).map_err(decode_error_to_status)?; match Command::try_from(message).map_err(arrow_error_to_status)? { Command::CommandStatementUpdate(command) => { let record_count = self.do_put_statement_update(command, request).await?; @@ -968,6 +988,26 @@ where } } +/// Unrecoverable errors associated with `do_put` requests +pub enum DoPutError { + /// The first element in the request stream is missing the command + MissingCommand, + /// The first element in the request stream is missing the flight descriptor + MissingFlightDescriptor, +} +impl Display for DoPutError { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + DoPutError::MissingCommand => { + write!(f, "Command is missing.") + } + DoPutError::MissingFlightDescriptor => { + write!(f, "Flight descriptor is missing.") + } + } + } +} + fn decode_error_to_status(err: prost::DecodeError) -> Status { Status::invalid_argument(format!("{err:?}")) } diff --git a/arrow-flight/src/streams.rs b/arrow-flight/src/streams.rs index e532a80e1ebb..0cd3aa41a547 100644 --- a/arrow-flight/src/streams.rs +++ b/arrow-flight/src/streams.rs @@ -32,7 +32,7 @@ use std::task::{ready, Poll}; /// /// This can be used to accept a stream of `Result<_>` from a client API and send /// them to the remote server that wants only the successful results. -pub(crate) struct FallibleRequestStream { +pub struct FallibleRequestStream { /// sender to notify error sender: Option>, /// fallible stream @@ -40,7 +40,8 @@ pub(crate) struct FallibleRequestStream { } impl FallibleRequestStream { - pub(crate) fn new( + /// Create a FallibleRequestStream + pub fn new( sender: Sender, fallible_stream: Pin> + Send + 'static>>, ) -> Self { @@ -127,7 +128,7 @@ impl Stream for FallibleTonicResponseStream { match ready!(pinned.response_stream.poll_next_unpin(cx)) { Some(Ok(res)) => Poll::Ready(Some(Ok(res))), - Some(Err(status)) => Poll::Ready(Some(Err(FlightError::Tonic(status)))), + Some(Err(status)) => Poll::Ready(Some(Err(status.into()))), None => Poll::Ready(None), } } diff --git a/arrow-flight/src/utils.rs b/arrow-flight/src/utils.rs index 428dde73ca6c..a304aedcfaee 100644 --- a/arrow-flight/src/utils.rs +++ b/arrow-flight/src/utils.rs @@ -90,9 +90,7 @@ pub fn batches_to_flight_data( let mut flight_data = vec![]; let data_gen = writer::IpcDataGenerator::default(); - #[allow(deprecated)] - let mut dictionary_tracker = - writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id()); + let mut dictionary_tracker = writer::DictionaryTracker::new(false); for batch in batches.iter() { let (encoded_dictionaries, encoded_batch) = diff --git a/arrow-flight/tests/flight_sql_client.rs b/arrow-flight/tests/flight_sql_client.rs index 349da062a82d..f3b7114dbafa 100644 --- a/arrow-flight/tests/flight_sql_client.rs +++ b/arrow-flight/tests/flight_sql_client.rs @@ -22,21 +22,23 @@ use crate::common::utils::make_primitive_batch; use arrow_array::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::error::FlightError; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::sql::client::FlightSqlServiceClient; use arrow_flight::sql::server::{FlightSqlService, PeekableFlightDataStream}; use arrow_flight::sql::{ ActionBeginTransactionRequest, ActionBeginTransactionResult, ActionEndTransactionRequest, - CommandStatementIngest, EndTransaction, SqlInfo, TableDefinitionOptions, TableExistsOption, - TableNotExistOption, + CommandStatementIngest, EndTransaction, FallibleRequestStream, ProstMessageExt, SqlInfo, + TableDefinitionOptions, TableExistsOption, TableNotExistOption, }; -use arrow_flight::Action; +use arrow_flight::{Action, FlightData, FlightDescriptor}; use futures::{StreamExt, TryStreamExt}; +use prost::Message; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::Mutex; -use tonic::{Request, Status}; +use tonic::{IntoStreamingRequest, Request, Status}; use uuid::Uuid; #[tokio::test] @@ -116,6 +118,89 @@ pub async fn test_execute_ingest_error() { ); } +#[tokio::test] +pub async fn test_do_put_empty_stream() { + // Test for https://github.com/apache/arrow-rs/issues/7329 + + let test_server = FlightSqlServiceImpl::new(); + let fixture = TestFixture::new(test_server.service()).await; + let channel = fixture.channel().await; + let mut flight_sql_client = FlightSqlServiceClient::new(channel); + let cmd = make_ingest_command(); + + // Create an empty request stream + let input_data = futures::stream::iter(vec![]); + let flight_descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec()); + let flight_data_encoder = FlightDataEncoderBuilder::default() + .with_flight_descriptor(Some(flight_descriptor)) + .build(input_data); + let flight_data: Vec = Box::pin(flight_data_encoder).try_collect().await.unwrap(); + let request_stream = futures::stream::iter(flight_data); + + // Execute a `do_put` and verify that the server error contains the expected message + let err = flight_sql_client.do_put(request_stream).await.unwrap_err(); + assert!(err + .to_string() + .contains("Unhandled Error: Command is missing."),); +} + +#[tokio::test] +pub async fn test_do_put_first_element_err() { + // Test for https://github.com/apache/arrow-rs/issues/7329 + + let test_server = FlightSqlServiceImpl::new(); + let fixture = TestFixture::new(test_server.service()).await; + let channel = fixture.channel().await; + let mut flight_sql_client = FlightSqlServiceClient::new(channel); + let cmd = make_ingest_command(); + + let (sender, _receiver) = futures::channel::oneshot::channel(); + + // Create a fallible request stream such that the 1st element is a FlightError + let input_data = futures::stream::iter(vec![ + Err(FlightError::NotYetImplemented("random error".to_string())), + Ok(make_primitive_batch(5)), + ]); + let flight_descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec()); + let flight_data_encoder = FlightDataEncoderBuilder::default() + .with_flight_descriptor(Some(flight_descriptor)) + .build(input_data); + let flight_data: FallibleRequestStream = + FallibleRequestStream::new(sender, Box::pin(flight_data_encoder)); + let request_stream = flight_data.into_streaming_request(); + + // Execute a `do_put` and verify that the server error contains the expected message + let err = flight_sql_client.do_put(request_stream).await.unwrap_err(); + + assert!(err + .to_string() + .contains("Unhandled Error: Command is missing."),); +} + +#[tokio::test] +pub async fn test_do_put_missing_flight_descriptor() { + // Test for https://github.com/apache/arrow-rs/issues/7329 + + let test_server = FlightSqlServiceImpl::new(); + let fixture = TestFixture::new(test_server.service()).await; + let channel = fixture.channel().await; + let mut flight_sql_client = FlightSqlServiceClient::new(channel); + + // Create a request stream such that the flight descriptor is missing + let stream = futures::stream::iter(vec![Ok(make_primitive_batch(5))]); + let flight_data_encoder = FlightDataEncoderBuilder::default() + .with_flight_descriptor(None) + .build(stream); + let flight_data: Vec = Box::pin(flight_data_encoder).try_collect().await.unwrap(); + let request_stream = futures::stream::iter(flight_data); + + // Execute a `do_put` and verify that the server error contains the expected message + let err = flight_sql_client.do_put(request_stream).await.unwrap_err(); + assert!(err + .to_string() + .contains("Unhandled Error: Flight descriptor is missing."),); +} + fn make_ingest_command() -> CommandStatementIngest { CommandStatementIngest { table_definition_options: Some(TableDefinitionOptions { diff --git a/arrow-integration-test/Cargo.toml b/arrow-integration-test/Cargo.toml index 8afbfacff7c3..d560d4fd8363 100644 --- a/arrow-integration-test/Cargo.toml +++ b/arrow-integration-test/Cargo.toml @@ -30,9 +30,11 @@ rust-version = { workspace = true } [lib] name = "arrow_integration_test" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow = { workspace = true } arrow-buffer = { workspace = true } @@ -40,5 +42,3 @@ hex = { version = "0.4", default-features = false, features = ["std"] } serde = { version = "1.0", default-features = false, features = ["rc", "derive"] } serde_json = { version = "1.0", default-features = false, features = ["std"] } num = { version = "0.4", default-features = false, features = ["std"] } - -[build-dependencies] diff --git a/arrow-integration-test/src/datatype.rs b/arrow-integration-test/src/datatype.rs index e45e94c24e07..4c17fbe76be7 100644 --- a/arrow-integration-test/src/datatype.rs +++ b/arrow-integration-test/src/datatype.rs @@ -60,14 +60,14 @@ pub fn data_type_from_json(json: &serde_json::Value) -> Result { _ => 128, // Default bit width }; - if bit_width == 128 { - Ok(DataType::Decimal128(precision, scale)) - } else if bit_width == 256 { - Ok(DataType::Decimal256(precision, scale)) - } else { - Err(ArrowError::ParseError( + match bit_width { + 32 => Ok(DataType::Decimal32(precision, scale)), + 64 => Ok(DataType::Decimal64(precision, scale)), + 128 => Ok(DataType::Decimal128(precision, scale)), + 256 => Ok(DataType::Decimal256(precision, scale)), + _ => Err(ArrowError::ParseError( "Decimal bit_width invalid".to_string(), - )) + )), } } Some(s) if s == "floatingpoint" => match map.get("precision") { @@ -337,6 +337,12 @@ pub fn data_type_to_json(data_type: &DataType) -> serde_json::Value { TimeUnit::Nanosecond => "NANOSECOND", }}), DataType::Dictionary(_, _) => json!({ "name": "dictionary"}), + DataType::Decimal32(precision, scale) => { + json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 32}) + } + DataType::Decimal64(precision, scale) => { + json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 64}) + } DataType::Decimal128(precision, scale) => { json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 128}) } diff --git a/arrow-integration-test/src/lib.rs b/arrow-integration-test/src/lib.rs index f025009c22de..177a1c47f31f 100644 --- a/arrow-integration-test/src/lib.rs +++ b/arrow-integration-test/src/lib.rs @@ -21,6 +21,11 @@ //! //! This is not a canonical format, but provides a human-readable way of verifying language implementations +#![doc( + html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg", + html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![warn(missing_docs)] use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer}; use hex::decode; @@ -813,6 +818,42 @@ pub fn array_from_json( ))), } } + DataType::Decimal32(precision, scale) => { + let mut b = Decimal32Builder::with_capacity(json_col.count); + for (is_valid, value) in json_col + .validity + .as_ref() + .unwrap() + .iter() + .zip(json_col.data.unwrap()) + { + match is_valid { + 1 => b.append_value(value.as_str().unwrap().parse::().unwrap()), + _ => b.append_null(), + }; + } + Ok(Arc::new( + b.finish().with_precision_and_scale(*precision, *scale)?, + )) + } + DataType::Decimal64(precision, scale) => { + let mut b = Decimal64Builder::with_capacity(json_col.count); + for (is_valid, value) in json_col + .validity + .as_ref() + .unwrap() + .iter() + .zip(json_col.data.unwrap()) + { + match is_valid { + 1 => b.append_value(value.as_str().unwrap().parse::().unwrap()), + _ => b.append_null(), + }; + } + Ok(Arc::new( + b.finish().with_precision_and_scale(*precision, *scale)?, + )) + } DataType::Decimal128(precision, scale) => { let mut b = Decimal128Builder::with_capacity(json_col.count); for (is_valid, value) in json_col diff --git a/arrow-integration-testing/Cargo.toml b/arrow-integration-testing/Cargo.toml index 8654b4b92734..8e91fcbb3cb2 100644 --- a/arrow-integration-testing/Cargo.toml +++ b/arrow-integration-testing/Cargo.toml @@ -43,7 +43,7 @@ prost = { version = "0.13", default-features = false } serde = { version = "1.0", default-features = false, features = ["rc", "derive"] } serde_json = { version = "1.0", default-features = false, features = ["std"] } tokio = { version = "1.0", default-features = false, features = [ "rt-multi-thread"] } -tonic = { version = "0.12", default-features = false } +tonic = { version = "0.13", default-features = false } tracing-subscriber = { version = "0.3.1", default-features = false, features = ["fmt"], optional = true } flate2 = { version = "1", default-features = false, features = ["rust_backend"] } diff --git a/arrow-integration-testing/src/flight_client_scenarios/integration_test.rs b/arrow-integration-testing/src/flight_client_scenarios/integration_test.rs index 406419028d00..bd41ab602ee5 100644 --- a/arrow-integration-testing/src/flight_client_scenarios/integration_test.rs +++ b/arrow-integration-testing/src/flight_client_scenarios/integration_test.rs @@ -72,9 +72,7 @@ async fn upload_data( let (mut upload_tx, upload_rx) = mpsc::channel(10); let options = arrow::ipc::writer::IpcWriteOptions::default(); - #[allow(deprecated)] - let mut dict_tracker = - writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id()); + let mut dict_tracker = writer::DictionaryTracker::new(false); let data_gen = writer::IpcDataGenerator::default(); let data = IpcMessage( data_gen diff --git a/arrow-integration-testing/src/flight_client_scenarios/middleware.rs b/arrow-integration-testing/src/flight_client_scenarios/middleware.rs index b826ad456055..495825738aec 100644 --- a/arrow-integration-testing/src/flight_client_scenarios/middleware.rs +++ b/arrow-integration-testing/src/flight_client_scenarios/middleware.rs @@ -76,7 +76,7 @@ pub async fn run_scenario(host: &str, port: u16) -> Result { Ok(()) } -#[allow(clippy::unnecessary_wraps)] +#[allow(clippy::result_large_err)] fn middleware_interceptor(mut req: Request<()>) -> Result, Status> { let metadata = req.metadata_mut(); metadata.insert("x-middleware", "expected value".parse().unwrap()); diff --git a/arrow-integration-testing/src/flight_server_scenarios/integration_test.rs b/arrow-integration-testing/src/flight_server_scenarios/integration_test.rs index 92989a20393e..d608a4753723 100644 --- a/arrow-integration-testing/src/flight_server_scenarios/integration_test.rs +++ b/arrow-integration-testing/src/flight_server_scenarios/integration_test.rs @@ -119,9 +119,7 @@ impl FlightService for FlightServiceImpl { .ok_or_else(|| Status::not_found(format!("Could not find flight. {key}")))?; let options = arrow::ipc::writer::IpcWriteOptions::default(); - #[allow(deprecated)] - let mut dictionary_tracker = - writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id()); + let mut dictionary_tracker = writer::DictionaryTracker::new(false); let data_gen = writer::IpcDataGenerator::default(); let data = IpcMessage( data_gen diff --git a/arrow-integration-testing/src/lib.rs b/arrow-integration-testing/src/lib.rs index e669690ef4f5..10512a00eb9d 100644 --- a/arrow-integration-testing/src/lib.rs +++ b/arrow-integration-testing/src/lib.rs @@ -207,8 +207,7 @@ fn cdata_integration_import_schema_and_compare_to_json( // compare schemas if canonicalize_schema(&json_schema) != canonicalize_schema(&imported_schema) { return Err(ArrowError::ComputeError(format!( - "Schemas do not match.\n- JSON: {:?}\n- Imported: {:?}", - json_schema, imported_schema + "Schemas do not match.\n- JSON: {json_schema:?}\n- Imported: {imported_schema:?}", ))); } Ok(()) @@ -253,7 +252,7 @@ fn cdata_integration_import_batch_and_compare_to_json( fn result_to_c_error(result: &std::result::Result) -> *mut c_char { match result { Ok(_) => ptr::null_mut(), - Err(e) => CString::new(format!("{}", e)).unwrap().into_raw(), + Err(e) => CString::new(format!("{e}")).unwrap().into_raw(), } } diff --git a/arrow-ipc/Cargo.toml b/arrow-ipc/Cargo.toml index cf91b3a3415f..a1f826ef7d10 100644 --- a/arrow-ipc/Cargo.toml +++ b/arrow-ipc/Cargo.toml @@ -30,15 +30,17 @@ rust-version = { workspace = true } [lib] name = "arrow_ipc" -path = "src/lib.rs" bench = false +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } -flatbuffers = { version = "24.3.25", default-features = false } +flatbuffers = { version = "25.2.10", default-features = false } lz4_flex = { version = "0.11", default-features = false, features = ["std", "frame"], optional = true } zstd = { version = "0.13.0", default-features = false, optional = true } @@ -47,4 +49,17 @@ default = [] lz4 = ["lz4_flex"] [dev-dependencies] +criterion = "0.5.1" tempfile = "3.3" +tokio = "1.43.0" +# used in benches +memmap2 = "0.9.3" +bytes = "1.9" + +[[bench]] +name = "ipc_writer" +harness = false + +[[bench]] +name = "ipc_reader" +harness = false diff --git a/arrow-ipc/benches/ipc_reader.rs b/arrow-ipc/benches/ipc_reader.rs new file mode 100644 index 000000000000..ab77449eeb7d --- /dev/null +++ b/arrow-ipc/benches/ipc_reader.rs @@ -0,0 +1,307 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::builder::{Date32Builder, Decimal128Builder, Int32Builder}; +use arrow_array::{builder::StringBuilder, RecordBatch}; +use arrow_buffer::Buffer; +use arrow_ipc::convert::fb_to_schema; +use arrow_ipc::reader::{read_footer_length, FileDecoder, FileReader, StreamReader}; +use arrow_ipc::writer::{FileWriter, IpcWriteOptions, StreamWriter}; +use arrow_ipc::{root_as_footer, Block, CompressionType}; +use arrow_schema::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use std::io::{Cursor, Write}; +use std::sync::Arc; +use tempfile::tempdir; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("arrow_ipc_reader"); + + group.bench_function("StreamReader/read_10", |b| { + let buffer = ipc_stream(IpcWriteOptions::default()); + b.iter(move || { + let projection = None; + let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap(); + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + group.bench_function("StreamReader/no_validation/read_10", |b| { + let buffer = ipc_stream(IpcWriteOptions::default()); + b.iter(move || { + let projection = None; + let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap(); + unsafe { + // safety: we created a valid IPC file + reader = reader.with_skip_validation(true); + } + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + group.bench_function("StreamReader/read_10/zstd", |b| { + let buffer = ipc_stream( + IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + .unwrap(), + ); + b.iter(move || { + let projection = None; + let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap(); + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + group.bench_function("StreamReader/no_validation/read_10/zstd", |b| { + let buffer = ipc_stream( + IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + .unwrap(), + ); + b.iter(move || { + let projection = None; + let mut reader = StreamReader::try_new(buffer.as_slice(), projection).unwrap(); + unsafe { + // safety: we created a valid IPC file + reader = reader.with_skip_validation(true); + } + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + // --- Create IPC File --- + group.bench_function("FileReader/read_10", |b| { + let buffer = ipc_file(); + b.iter(move || { + let projection = None; + let cursor = Cursor::new(buffer.as_slice()); + let mut reader = FileReader::try_new(cursor, projection).unwrap(); + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + group.bench_function("FileReader/no_validation/read_10", |b| { + let buffer = ipc_file(); + b.iter(move || { + let projection = None; + let cursor = Cursor::new(buffer.as_slice()); + let mut reader = FileReader::try_new(cursor, projection).unwrap(); + unsafe { + // safety: we created a valid IPC file + reader = reader.with_skip_validation(true); + } + for _ in 0..10 { + reader.next().unwrap().unwrap(); + } + assert!(reader.next().is_none()); + }) + }); + + // write to an actual file + let dir = tempdir().unwrap(); + let path = dir.path().join("test.arrow"); + let mut file = std::fs::File::create(&path).unwrap(); + file.write_all(&ipc_file()).unwrap(); + drop(file); + + group.bench_function("FileReader/read_10/mmap", |b| { + let path = &path; + b.iter(move || { + let ipc_file = std::fs::File::open(path).expect("failed to open file"); + let mmap = unsafe { memmap2::Mmap::map(&ipc_file).expect("failed to mmap file") }; + + // Convert the mmap region to an Arrow `Buffer` to back the arrow arrays. + let bytes = bytes::Bytes::from_owner(mmap); + let buffer = Buffer::from(bytes); + let decoder = IPCBufferDecoder::new(buffer); + assert_eq!(decoder.num_batches(), 10); + + for i in 0..decoder.num_batches() { + decoder.get_batch(i); + } + }) + }); + + group.bench_function("FileReader/no_validation/read_10/mmap", |b| { + let path = &path; + b.iter(move || { + let ipc_file = std::fs::File::open(path).expect("failed to open file"); + let mmap = unsafe { memmap2::Mmap::map(&ipc_file).expect("failed to mmap file") }; + + // Convert the mmap region to an Arrow `Buffer` to back the arrow arrays. + let bytes = bytes::Bytes::from_owner(mmap); + let buffer = Buffer::from(bytes); + let decoder = IPCBufferDecoder::new(buffer); + let decoder = unsafe { decoder.with_skip_validation(true) }; + assert_eq!(decoder.num_batches(), 10); + + for i in 0..decoder.num_batches() { + decoder.get_batch(i); + } + }) + }); +} + +/// Return an IPC stream with 10 record batches +fn ipc_stream(options: IpcWriteOptions) -> Vec { + let batch = create_batch(8192, true); + let mut buffer = Vec::with_capacity(2 * 1024 * 1024); + let mut writer = + StreamWriter::try_new_with_options(&mut buffer, batch.schema().as_ref(), options).unwrap(); + for _ in 0..10 { + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + buffer +} + +/// Return an IPC file with 10 record batches +fn ipc_file() -> Vec { + let batch = create_batch(8192, true); + let mut buffer = Vec::with_capacity(2 * 1024 * 1024); + let mut writer = FileWriter::try_new(&mut buffer, batch.schema().as_ref()).unwrap(); + for _ in 0..10 { + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + buffer +} + +// copied from the zero_copy_ipc example. +// should we move this to an actual API? +/// Wrapper around the example in the `FileDecoder` which handles the +/// low level interaction with the Arrow IPC format. +struct IPCBufferDecoder { + /// Memory (or memory mapped) Buffer with the data + buffer: Buffer, + /// Decoder that reads Arrays that refers to the underlying buffers + decoder: FileDecoder, + /// Location of the batches within the buffer + batches: Vec, +} + +impl IPCBufferDecoder { + fn new(buffer: Buffer) -> Self { + let trailer_start = buffer.len() - 10; + let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap(); + let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap(); + + let schema = fb_to_schema(footer.schema().unwrap()); + + let mut decoder = FileDecoder::new(Arc::new(schema), footer.version()); + + // Read dictionaries + for block in footer.dictionaries().iter().flatten() { + let block_len = block.bodyLength() as usize + block.metaDataLength() as usize; + let data = buffer.slice_with_length(block.offset() as _, block_len); + decoder.read_dictionary(block, &data).unwrap(); + } + + // convert to Vec from the flatbuffers Vector to avoid having a direct dependency on flatbuffers + let batches = footer + .recordBatches() + .map(|b| b.iter().copied().collect()) + .unwrap_or_default(); + + Self { + buffer, + decoder, + batches, + } + } + + unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self { + self.decoder = self.decoder.with_skip_validation(skip_validation); + self + } + + fn num_batches(&self) -> usize { + self.batches.len() + } + + fn get_batch(&self, i: usize) -> RecordBatch { + let block = &self.batches[i]; + let block_len = block.bodyLength() as usize + block.metaDataLength() as usize; + let data = self + .buffer + .slice_with_length(block.offset() as _, block_len); + self.decoder + .read_record_batch(block, &data) + .unwrap() + .unwrap() + } +} + +fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Date32, true), + Field::new("c3", DataType::Decimal128(11, 2), true), + ])); + let mut a = Int32Builder::new(); + let mut b = StringBuilder::new(); + let mut c = Date32Builder::new(); + let mut d = Decimal128Builder::new() + .with_precision_and_scale(11, 2) + .unwrap(); + for i in 0..num_rows { + a.append_value(i as i32); + c.append_value(i as i32); + d.append_value((i * 1000000) as i128); + if allow_nulls && i % 10 == 0 { + b.append_null(); + } else { + b.append_value(format!("this is string number {i}")); + } + } + let a = a.finish(); + let b = b.finish(); + let c = c.finish(); + let d = d.finish(); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(a), Arc::new(b), Arc::new(c), Arc::new(d)], + ) + .unwrap() +} + +fn config() -> Criterion { + Criterion::default() +} + +criterion_group! { + name = benches; + config = config(); + targets = criterion_benchmark +} +criterion_main!(benches); diff --git a/arrow-ipc/benches/ipc_writer.rs b/arrow-ipc/benches/ipc_writer.rs new file mode 100644 index 000000000000..6b4d184b4556 --- /dev/null +++ b/arrow-ipc/benches/ipc_writer.rs @@ -0,0 +1,117 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::builder::{Date32Builder, Decimal128Builder, Int32Builder}; +use arrow_array::{builder::StringBuilder, RecordBatch}; +use arrow_ipc::writer::{FileWriter, IpcWriteOptions, StreamWriter}; +use arrow_ipc::CompressionType; +use arrow_schema::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("arrow_ipc_stream_writer"); + + group.bench_function("StreamWriter/write_10", |b| { + let batch = create_batch(8192, true); + let mut buffer = Vec::with_capacity(2 * 1024 * 1024); + b.iter(move || { + buffer.clear(); + let mut writer = StreamWriter::try_new(&mut buffer, batch.schema().as_ref()).unwrap(); + for _ in 0..10 { + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + }) + }); + + group.bench_function("StreamWriter/write_10/zstd", |b| { + let batch = create_batch(8192, true); + let mut buffer = Vec::with_capacity(2 * 1024 * 1024); + b.iter(move || { + buffer.clear(); + let options = IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + .unwrap(); + let mut writer = + StreamWriter::try_new_with_options(&mut buffer, batch.schema().as_ref(), options) + .unwrap(); + for _ in 0..10 { + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + }) + }); + + group.bench_function("FileWriter/write_10", |b| { + let batch = create_batch(8192, true); + let mut buffer = Vec::with_capacity(2 * 1024 * 1024); + b.iter(move || { + buffer.clear(); + let mut writer = FileWriter::try_new(&mut buffer, batch.schema().as_ref()).unwrap(); + for _ in 0..10 { + writer.write(&batch).unwrap(); + } + writer.finish().unwrap(); + }) + }); +} + +fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Date32, true), + Field::new("c3", DataType::Decimal128(11, 2), true), + ])); + let mut a = Int32Builder::new(); + let mut b = StringBuilder::new(); + let mut c = Date32Builder::new(); + let mut d = Decimal128Builder::new() + .with_precision_and_scale(11, 2) + .unwrap(); + for i in 0..num_rows { + a.append_value(i as i32); + c.append_value(i as i32); + d.append_value((i * 1000000) as i128); + if allow_nulls && i % 10 == 0 { + b.append_null(); + } else { + b.append_value(format!("this is string number {i}")); + } + } + let a = a.finish(); + let b = b.finish(); + let c = c.finish(); + let d = d.finish(); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(a), Arc::new(b), Arc::new(c), Arc::new(d)], + ) + .unwrap() +} + +fn config() -> Criterion { + Criterion::default() +} + +criterion_group! { + name = benches; + config = config(); + targets = criterion_benchmark +} +criterion_main!(benches); diff --git a/arrow-ipc/regen.sh b/arrow-ipc/regen.sh index 8d8862ccc7f4..b368bd1bc7cc 100755 --- a/arrow-ipc/regen.sh +++ b/arrow-ipc/regen.sh @@ -21,33 +21,36 @@ DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" # Change to the toplevel `arrow-rs` directory pushd $DIR/../ -echo "Build flatc from source ..." - -FB_URL="https://github.com/google/flatbuffers" -FB_DIR="arrow/.flatbuffers" -FLATC="$FB_DIR/bazel-bin/flatc" - -if [ -z $(which bazel) ]; then - echo "bazel is required to build flatc" - exit 1 -fi - -echo "Bazel version: $(bazel version | head -1 | awk -F':' '{print $2}')" - -if [ ! -e $FB_DIR ]; then - echo "git clone $FB_URL ..." - git clone -b master --no-tag --depth 1 $FB_URL $FB_DIR +if [ -z "$FLATC" ]; then + echo "Build flatc from source ..." + + FB_URL="https://github.com/google/flatbuffers" + FB_DIR="arrow/.flatbuffers" + FLATC="$FB_DIR/bazel-bin/flatc" + + if [ -z $(which bazel) ]; then + echo "bazel is required to build flatc" + exit 1 + fi + + echo "Bazel version: $(bazel version | head -1 | awk -F':' '{print $2}')" + + if [ ! -e $FB_DIR ]; then + echo "git clone $FB_URL ..." + git clone -b master --no-tag --depth 1 $FB_URL $FB_DIR + else + echo "git pull $FB_URL ..." + git -C $FB_DIR pull + fi + + pushd $FB_DIR + echo "run: bazel build :flatc ..." + bazel build :flatc + popd else - echo "git pull $FB_URL ..." - git -C $FB_DIR pull + echo "Using flatc $FLATC ..." fi -pushd $FB_DIR -echo "run: bazel build :flatc ..." -bazel build :flatc -popd - - # Execute the code generation: $FLATC --filename-suffix "" --rust -o arrow-ipc/src/gen/ format/*.fbs @@ -99,37 +102,38 @@ for f in `ls *.rs`; do fi echo "Modifying: $f" - sed -i '' '/extern crate flatbuffers;/d' $f - sed -i '' '/use self::flatbuffers::EndianScalar;/d' $f - sed -i '' '/\#\[allow(unused_imports, dead_code)\]/d' $f - sed -i '' '/pub mod org {/d' $f - sed -i '' '/pub mod apache {/d' $f - sed -i '' '/pub mod arrow {/d' $f - sed -i '' '/pub mod flatbuf {/d' $f - sed -i '' '/} \/\/ pub mod flatbuf/d' $f - sed -i '' '/} \/\/ pub mod arrow/d' $f - sed -i '' '/} \/\/ pub mod apache/d' $f - sed -i '' '/} \/\/ pub mod org/d' $f - sed -i '' '/use core::mem;/d' $f - sed -i '' '/use core::cmp::Ordering;/d' $f - sed -i '' '/use self::flatbuffers::{EndianScalar, Follow};/d' $f + sed --in-place='' '/extern crate flatbuffers;/d' $f + sed --in-place='' '/use self::flatbuffers::EndianScalar;/d' $f + sed --in-place='' '/\#\[allow(unused_imports, dead_code)\]/d' $f + sed --in-place='' '/pub mod org {/d' $f + sed --in-place='' '/pub mod apache {/d' $f + sed --in-place='' '/pub mod arrow {/d' $f + sed --in-place='' '/pub mod flatbuf {/d' $f + sed --in-place='' '/} \/\/ pub mod flatbuf/d' $f + sed --in-place='' '/} \/\/ pub mod arrow/d' $f + sed --in-place='' '/} \/\/ pub mod apache/d' $f + sed --in-place='' '/} \/\/ pub mod org/d' $f + sed --in-place='' '/use core::mem;/d' $f + sed --in-place='' '/use core::cmp::Ordering;/d' $f + sed --in-place='' '/use self::flatbuffers::{EndianScalar, Follow};/d' $f # required by flatc 1.12.0+ - sed -i '' "/\#\!\[allow(unused_imports, dead_code)\]/d" $f + sed --in-place='' "/\#\!\[allow(unused_imports, dead_code)\]/d" $f for name in ${names[@]}; do - sed -i '' "/use crate::${name}::\*;/d" $f - sed -i '' "s/use self::flatbuffers::Verifiable;/use flatbuffers::Verifiable;/g" $f + sed --in-place='' "/use crate::${name}::\*;/d" $f + sed --in-place='' "s/use self::flatbuffers::Verifiable;/use flatbuffers::Verifiable;/g" $f done # Replace all occurrences of "type__" with "type_", "TYPE__" with "TYPE_". - sed -i '' 's/type__/type_/g' $f - sed -i '' 's/TYPE__/TYPE_/g' $f + sed --in-place='' 's/type__/type_/g' $f + sed --in-place='' 's/TYPE__/TYPE_/g' $f # Some files need prefixes if [[ $f == "File.rs" ]]; then # Now prefix the file with the static contents echo -e "${PREFIX}" "${SCHEMA_IMPORT}" | cat - $f > temp && mv temp $f elif [[ $f == "Message.rs" ]]; then + sed --in-place='' 's/List/\`List\`/g' $f echo -e "${PREFIX}" "${SCHEMA_IMPORT}" "${SPARSE_TENSOR_IMPORT}" "${TENSOR_IMPORT}" | cat - $f > temp && mv temp $f elif [[ $f == "SparseTensor.rs" ]]; then echo -e "${PREFIX}" "${SCHEMA_IMPORT}" "${TENSOR_IMPORT}" | cat - $f > temp && mv temp $f diff --git a/arrow-ipc/src/convert.rs b/arrow-ipc/src/convert.rs index 37c5a19439c1..af0bdb1df3eb 100644 --- a/arrow-ipc/src/convert.rs +++ b/arrow-ipc/src/convert.rs @@ -19,6 +19,7 @@ use arrow_buffer::Buffer; use arrow_schema::*; +use core::panic; use flatbuffers::{ FlatBufferBuilder, ForwardsUOffset, UnionWIPOffset, Vector, Verifiable, Verifier, VerifierOptions, WIPOffset, @@ -28,7 +29,7 @@ use std::fmt::{Debug, Formatter}; use std::sync::Arc; use crate::writer::DictionaryTracker; -use crate::{size_prefixed_root_as_message, KeyValue, Message, CONTINUATION_MARKER}; +use crate::{KeyValue, Message, CONTINUATION_MARKER}; use DataType::*; /// Low level Arrow [Schema] to IPC bytes converter @@ -127,20 +128,17 @@ impl<'a> IpcSchemaEncoder<'a> { } } -/// Serialize a schema in IPC format -#[deprecated(since = "54.0.0", note = "Use `IpcSchemaConverter`.")] -pub fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder<'_> { - IpcSchemaEncoder::new().schema_to_fb(schema) -} - /// Push a key-value metadata into a FlatBufferBuilder and return [WIPOffset] pub fn metadata_to_fb<'a>( fbb: &mut FlatBufferBuilder<'a>, metadata: &HashMap, ) -> WIPOffset>>> { - let custom_metadata = metadata - .iter() - .map(|(k, v)| { + let mut ordered_keys = metadata.keys().collect::>(); + ordered_keys.sort(); + let custom_metadata = ordered_keys + .into_iter() + .map(|k| { + let v = metadata.get(k).unwrap(); let fb_key_name = fbb.create_string(k); let fb_val_name = fbb.create_string(v); @@ -255,32 +253,43 @@ pub fn try_schema_from_ipc_buffer(buffer: &[u8]) -> Result { // 4 bytes - an optional IPC_CONTINUATION_TOKEN prefix // 4 bytes - the byte length of the payload // a flatbuffer Message whose header is the Schema - if buffer.len() >= 4 { - // check continuation marker - let continuation_marker = &buffer[0..4]; - let begin_offset: usize = if continuation_marker.eq(&CONTINUATION_MARKER) { - // 4 bytes: CONTINUATION_MARKER - // 4 bytes: length - // buffer - 4 - } else { - // backward compatibility for buffer without the continuation marker - // 4 bytes: length - // buffer - 0 - }; - let msg = size_prefixed_root_as_message(&buffer[begin_offset..]).map_err(|err| { - ArrowError::ParseError(format!("Unable to convert flight info to a message: {err}")) - })?; - let ipc_schema = msg.header_as_schema().ok_or_else(|| { - ArrowError::ParseError("Unable to convert flight info to a schema".to_string()) - })?; - Ok(fb_to_schema(ipc_schema)) - } else { - Err(ArrowError::ParseError( + if buffer.len() < 4 { + return Err(ArrowError::ParseError( "The buffer length is less than 4 and missing the continuation marker or length of buffer".to_string() - )) + )); + } + + let (len, buffer) = if buffer[..4] == CONTINUATION_MARKER { + if buffer.len() < 8 { + return Err(ArrowError::ParseError( + "The buffer length is less than 8 and missing the length of buffer".to_string(), + )); + } + buffer[4..].split_at(4) + } else { + buffer.split_at(4) + }; + + let len = ::from_le_bytes(len.try_into().unwrap()); + if len < 0 { + return Err(ArrowError::ParseError(format!( + "The encapsulated message's reported length is negative ({len})" + ))); } + + if buffer.len() < len as usize { + let actual_len = buffer.len(); + return Err(ArrowError::ParseError( + format!("The buffer length ({actual_len}) is less than the encapsulated message's reported length ({len})") + )); + } + + let msg = crate::root_as_message(buffer) + .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))?; + let ipc_schema = msg.header_as_schema().ok_or_else(|| { + ArrowError::ParseError("Unable to convert flight info to a schema".to_string()) + })?; + Ok(fb_to_schema(ipc_schema)) } /// Get the Arrow data type from the flatbuffer Field table @@ -454,18 +463,14 @@ pub(crate) fn get_data_type(field: crate::Field, may_be_dictionary: bool) -> Dat crate::Type::Decimal => { let fsb = field.type_as_decimal().unwrap(); let bit_width = fsb.bitWidth(); - if bit_width == 128 { - DataType::Decimal128( - fsb.precision().try_into().unwrap(), - fsb.scale().try_into().unwrap(), - ) - } else if bit_width == 256 { - DataType::Decimal256( - fsb.precision().try_into().unwrap(), - fsb.scale().try_into().unwrap(), - ) - } else { - panic!("Unexpected decimal bit width {bit_width}") + let precision: u8 = fsb.precision().try_into().unwrap(); + let scale: i8 = fsb.scale().try_into().unwrap(); + match bit_width { + 32 => DataType::Decimal32(precision, scale), + 64 => DataType::Decimal64(precision, scale), + 128 => DataType::Decimal128(precision, scale), + 256 => DataType::Decimal256(precision, scale), + _ => panic!("Unexpected decimal bit width {bit_width}"), } } crate::Type::Union => { @@ -520,24 +525,13 @@ pub(crate) fn build_field<'a>( match dictionary_tracker { Some(tracker) => Some(get_fb_dictionary( index_type, - #[allow(deprecated)] - tracker.set_dict_id(field), - field - .dict_is_ordered() - .expect("All Dictionary types have `dict_is_ordered`"), - fbb, - )), - None => Some(get_fb_dictionary( - index_type, - #[allow(deprecated)] - field - .dict_id() - .expect("Dictionary type must have a dictionary id"), + tracker.next_dict_id(), field .dict_is_ordered() .expect("All Dictionary types have `dict_is_ordered`"), fbb, )), + None => panic!("IPC must no longer be used without dictionary tracker"), } } else { None @@ -833,6 +827,28 @@ pub(crate) fn get_fb_field_type<'a>( // type in the DictionaryEncoding metadata in the parent field get_fb_field_type(value_type, dictionary_tracker, fbb) } + Decimal32(precision, scale) => { + let mut builder = crate::DecimalBuilder::new(fbb); + builder.add_precision(*precision as i32); + builder.add_scale(*scale as i32); + builder.add_bitWidth(32); + FBFieldType { + type_type: crate::Type::Decimal, + type_: builder.finish().as_union_value(), + children: Some(fbb.create_vector(&empty_fields[..])), + } + } + Decimal64(precision, scale) => { + let mut builder = crate::DecimalBuilder::new(fbb); + builder.add_precision(*precision as i32); + builder.add_scale(*scale as i32); + builder.add_bitWidth(64); + FBFieldType { + type_type: crate::Type::Decimal, + type_: builder.finish().as_union_value(), + children: Some(fbb.create_vector(&empty_fields[..])), + } + } Decimal128(precision, scale) => { let mut builder = crate::DecimalBuilder::new(fbb); builder.add_precision(*precision as i32); diff --git a/arrow-ipc/src/gen/File.rs b/arrow-ipc/src/gen/File.rs index c0c2fb183237..427cf75de096 100644 --- a/arrow-ipc/src/gen/File.rs +++ b/arrow-ipc/src/gen/File.rs @@ -23,6 +23,8 @@ use flatbuffers::EndianScalar; use std::{cmp::Ordering, mem}; // automatically generated by the FlatBuffers compiler, do not modify +// @generated + // struct Block, aligned to 8 #[repr(transparent)] #[derive(Clone, Copy, PartialEq)] @@ -64,6 +66,10 @@ impl<'b> flatbuffers::Push for Block { let src = ::core::slice::from_raw_parts(self as *const Block as *const u8, Self::size()); dst.copy_from_slice(src); } + #[inline] + fn alignment() -> flatbuffers::PushAlignment { + flatbuffers::PushAlignment::new(8) + } } impl<'a> flatbuffers::Verifiable for Block { @@ -211,8 +217,8 @@ impl<'a> Footer<'a> { Footer { _tab: table } } #[allow(unused_mut)] - pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>( - _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>, + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>, args: &'args FooterArgs<'args>, ) -> flatbuffers::WIPOffset> { let mut builder = FooterBuilder::new(_fbb); @@ -344,11 +350,11 @@ impl<'a> Default for FooterArgs<'a> { } } -pub struct FooterBuilder<'a: 'b, 'b> { - fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>, +pub struct FooterBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, start_: flatbuffers::WIPOffset, } -impl<'a: 'b, 'b> FooterBuilder<'a, 'b> { +impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> FooterBuilder<'a, 'b, A> { #[inline] pub fn add_version(&mut self, version: MetadataVersion) { self.fbb_ @@ -388,7 +394,7 @@ impl<'a: 'b, 'b> FooterBuilder<'a, 'b> { ); } #[inline] - pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> FooterBuilder<'a, 'b> { + pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>) -> FooterBuilder<'a, 'b, A> { let start = _fbb.start_table(); FooterBuilder { fbb_: _fbb, @@ -474,16 +480,16 @@ pub unsafe fn size_prefixed_root_as_footer_unchecked(buf: &[u8]) -> Footer { flatbuffers::size_prefixed_root_unchecked::