diff --git a/.asf.yaml b/.asf.yaml index e96b43cf0..cb0520c17 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -29,6 +29,13 @@ github: rebase: false features: issues: true + protected_branches: + main: + required_status_checks: + # require branches to be up-to-date before merging + strict: true + # don't require any jobs to pass + contexts: [] staging: whoami: asf-staging diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index 91a099a61..000000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,12 +0,0 @@ -[target.x86_64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - -[target.aarch64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index acabad3ca..4c07b08bb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,57 +15,248 @@ # specific language governing permissions and limitations # under the License. -name: Python Release Build +# Reusable workflow for running building +# This ensures the same tests run for both debug (PRs) and release (main/tags) builds + +name: Build + on: - pull_request: - branches: ["main"] - push: - tags: ["*-rc*"] - branches: ["branch-*"] + workflow_call: + inputs: + build_mode: + description: 'Build mode: debug or release' + required: true + type: string + run_wheels: + description: 'Whether to build distribution wheels' + required: false + type: boolean + default: false + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 jobs: - build: + # ============================================ + # Linting Jobs + # ============================================ + lint-rust: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 + + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable + with: + toolchain: "nightly" + components: rustfmt + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + + - name: Check formatting + run: cargo +nightly fmt --all -- --check + + lint-python: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - name: Install Python uses: actions/setup-python@v5 with: python-version: "3.12" - - uses: astral-sh/setup-uv@v5 + - uses: astral-sh/setup-uv@v6 with: - enable-cache: true + enable-cache: true - # Use the --no-install-package to only install the dependencies - # but do not yet build the rust library - name: Install dependencies run: uv sync --dev --no-install-package datafusion - # Update output format to enable automatic inline annotations. - name: Run Ruff run: | uv run --no-project ruff check --output-format=github python/ uv run --no-project ruff format --check python/ + - name: Run codespell + run: | + uv run --no-project codespell --toml pyproject.toml + + lint-toml: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Install taplo + uses: taiki-e/install-action@v2 + with: + tool: taplo-cli + + # if you encounter an error, try running 'taplo format' to fix the formatting automatically. + - name: Check Cargo.toml formatting + run: taplo format --check + + check-crates-patch: + if: inputs.build_mode == 'release' && startsWith(github.ref, 'refs/tags/') + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Ensure [patch.crates-io] is empty + run: python3 dev/check_crates_patch.py + generate-license: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: astral-sh/setup-uv@v5 + - uses: actions/checkout@v6 + + - uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + + - name: Install cargo-license + uses: taiki-e/install-action@v2 with: - enable-cache: true + tool: cargo-license - name: Generate license file run: uv run --no-project python ./dev/create_license.py - - uses: actions/upload-artifact@v4 + + - uses: actions/upload-artifact@v6 with: name: python-wheel-license path: LICENSE.txt + # ============================================ + # Build - Linux x86_64 + # ============================================ + build-manylinux-x86_64: + needs: [generate-license, lint-rust, lint-python] + name: ManyLinux x86_64 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - run: rm LICENSE.txt + - name: Download LICENSE.txt + uses: actions/download-artifact@v7 + with: + name: python-wheel-license + path: . + + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }} + + - uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + + - name: Build (release mode) + uses: PyO3/maturin-action@v1 + if: inputs.build_mode == 'release' + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + args: --release --strip --features protoc,substrait --out dist + rustup-components: rust-std + + - name: Build (debug mode) + uses: PyO3/maturin-action@v1 + if: inputs.build_mode == 'debug' + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + args: --features protoc,substrait --out dist + rustup-components: rust-std + + - name: Build FFI test library + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/datafusion-ffi-example + args: --out dist + rustup-components: rust-std + + - name: Archive wheels + uses: actions/upload-artifact@v6 + with: + name: dist-manylinux-x86_64 + path: dist/* + + - name: Archive FFI test wheel + uses: actions/upload-artifact@v6 + with: + name: test-ffi-manylinux-x86_64 + path: examples/datafusion-ffi-example/dist/* + + # ============================================ + # Build - Linux ARM64 + # ============================================ + build-manylinux-aarch64: + needs: [generate-license, lint-rust, lint-python] + name: ManyLinux arm64 + runs-on: ubuntu-24.04-arm + steps: + - uses: actions/checkout@v6 + + - run: rm LICENSE.txt + - name: Download LICENSE.txt + uses: actions/download-artifact@v7 + with: + name: python-wheel-license + path: . + + - name: Setup Rust + uses: dtolnay/rust-toolchain@stable + + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }} + + - uses: astral-sh/setup-uv@v6 + with: + enable-cache: true + + - name: Build (release mode) + uses: PyO3/maturin-action@v1 + if: inputs.build_mode == 'release' + with: + target: aarch64-unknown-linux-gnu + manylinux: "2_28" + args: --release --strip --features protoc,substrait --out dist + rustup-components: rust-std + + - name: Build (debug mode) + uses: PyO3/maturin-action@v1 + if: inputs.build_mode == 'debug' + with: + target: aarch64-unknown-linux-gnu + manylinux: "2_28" + args: --features protoc,substrait --out dist + rustup-components: rust-std + + - name: Archive wheels + uses: actions/upload-artifact@v6 + if: inputs.build_mode == 'release' + with: + name: dist-manylinux-aarch64 + path: dist/* + + # ============================================ + # Build - macOS arm64 / Windows + # ============================================ build-python-mac-win: - needs: [generate-license] - name: Mac/Win + needs: [generate-license, lint-rust, lint-python] + name: macOS arm64 & Windows runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -73,35 +264,49 @@ jobs: python-version: ["3.10"] os: [macos-latest, windows-latest] steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} + - uses: actions/checkout@v6 - uses: dtolnay/rust-toolchain@stable - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: name: python-wheel-license path: . + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }} + + - uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + - name: Install Protoc uses: arduino/setup-protoc@v3 with: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@v5 - with: - enable-cache: true + - name: Install dependencies + run: uv sync --dev --no-install-package datafusion - - name: Build Python package - run: | - uv sync --dev --no-install-package datafusion - uv run --no-project maturin build --release --strip --features substrait + # Run clippy BEFORE maturin so we can avoid rebuilding. The features must match + # exactly the features used by maturin. Linux maturin builds need to happen in a + # container so only run this for our mac runner. + - name: Run Clippy + if: matrix.os != 'windows-latest' + run: cargo clippy --no-deps --all-targets --features substrait -- -D warnings + + - name: Build Python package (release mode) + if: inputs.build_mode == 'release' + run: uv run --no-project maturin build --release --strip --features substrait + + - name: Build Python package (debug mode) + if: inputs.build_mode != 'release' + run: uv run --no-project maturin build --features substrait - name: List Windows wheels if: matrix.os == 'windows-latest' @@ -115,127 +320,80 @@ jobs: run: find target/wheels/ - name: Archive wheels - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 + if: inputs.build_mode == 'release' with: name: dist-${{ matrix.os }} path: target/wheels/* + # ============================================ + # Build - macOS x86_64 (release only) + # ============================================ build-macos-x86_64: - needs: [generate-license] - name: Mac x86_64 - runs-on: macos-13 + if: inputs.build_mode == 'release' + needs: [generate-license, lint-rust, lint-python] + runs-on: macos-15-intel strategy: fail-fast: false matrix: python-version: ["3.10"] steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} + - uses: actions/checkout@v6 - uses: dtolnay/rust-toolchain@stable - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: name: python-wheel-license path: . + - name: Cache Cargo + uses: Swatinem/rust-cache@v2 + with: + key: ${{ inputs.build_mode }} + + - uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + - name: Install Protoc uses: arduino/setup-protoc@v3 with: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@v5 - with: - enable-cache: true + - name: Install dependencies + run: uv sync --dev --no-install-package datafusion - - name: Build Python package + - name: Build (release mode) run: | - uv sync --dev --no-install-package datafusion uv run --no-project maturin build --release --strip --features substrait - name: List Mac wheels run: find target/wheels/ - name: Archive wheels - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: dist-macos-aarch64 path: target/wheels/* - build-manylinux-x86_64: - needs: [generate-license] - name: Manylinux x86_64 - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - run: rm LICENSE.txt - - name: Download LICENSE.txt - uses: actions/download-artifact@v4 - with: - name: python-wheel-license - path: . - - run: cat LICENSE.txt - - name: Build wheels - uses: PyO3/maturin-action@v1 - env: - RUST_BACKTRACE: 1 - with: - rust-toolchain: nightly - target: x86_64 - manylinux: auto - rustup-components: rust-std rustfmt # Keep them in one line due to https://github.com/PyO3/maturin-action/issues/153 - args: --release --manylinux 2014 --features protoc,substrait - - name: Archive wheels - uses: actions/upload-artifact@v4 - with: - name: dist-manylinux-x86_64 - path: target/wheels/* - - build-manylinux-aarch64: - needs: [generate-license] - name: Manylinux arm64 - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - run: rm LICENSE.txt - - name: Download LICENSE.txt - uses: actions/download-artifact@v4 - with: - name: python-wheel-license - path: . - - run: cat LICENSE.txt - - name: Build wheels - uses: PyO3/maturin-action@v1 - env: - RUST_BACKTRACE: 1 - with: - rust-toolchain: nightly - target: aarch64 - # Use manylinux_2_28-cross because the manylinux2014-cross has GCC 4.8.5, which causes the build to fail - manylinux: 2_28 - rustup-components: rust-std rustfmt # Keep them in one line due to https://github.com/PyO3/maturin-action/issues/153 - args: --release --features protoc,substrait - - name: Archive wheels - uses: actions/upload-artifact@v4 - with: - name: dist-manylinux-aarch64 - path: target/wheels/* + # ============================================ + # Build - Source Distribution + # ============================================ build-sdist: needs: [generate-license] name: Source distribution + if: inputs.build_mode == 'release' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - run: rm LICENSE.txt - name: Download LICENSE.txt - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: name: python-wheel-license path: . @@ -249,16 +407,22 @@ jobs: args: --release --sdist --out dist --features protoc,substrait - name: Assert sdist build does not generate wheels run: | - if [ "$(ls -A target/wheels)" ]; then - echo "Error: Sdist build generated wheels" - exit 1 - else - echo "Directory is clean" - fi + if [ "$(ls -A target/wheels)" ]; then + echo "Error: Sdist build generated wheels" + exit 1 + else + echo "Directory is clean" + fi shell: bash - + + # ============================================ + # Build - Source Distribution + # ============================================ + merge-build-artifacts: runs-on: ubuntu-latest + name: Merge build artifacts + if: inputs.build_mode == 'release' needs: - build-python-mac-win - build-macos-x86_64 @@ -267,20 +431,104 @@ jobs: - build-sdist steps: - name: Merge Build Artifacts - uses: actions/upload-artifact/merge@v4 + uses: actions/upload-artifact/merge@v6 with: name: dist pattern: dist-* - - # NOTE: PyPI publish needs to be done manually for now after release passed the vote - # release: - # name: Publish in PyPI - # needs: [build-manylinux, build-python-mac-win] - # runs-on: ubuntu-latest - # steps: - # - uses: actions/download-artifact@v4 - # - name: Publish to PyPI - # uses: pypa/gh-action-pypi-publish@master - # with: - # user: __token__ - # password: ${{ secrets.pypi_password }} + + # ============================================ + # Build - Documentation + # ============================================ + # Documentation build job that runs after wheels are built + build-docs: + name: Build docs + runs-on: ubuntu-latest + needs: [build-manylinux-x86_64] # Only need the Linux wheel for docs + # Only run docs on main branch pushes, tags, or PRs + if: github.event_name == 'push' || github.event_name == 'pull_request' + steps: + - name: Set target branch + if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') + id: target-branch + run: | + set -x + if test '${{ github.ref }}' = 'refs/heads/main'; then + echo "value=asf-staging" >> "$GITHUB_OUTPUT" + elif test '${{ github.ref_type }}' = 'tag'; then + echo "value=asf-site" >> "$GITHUB_OUTPUT" + else + echo "Unsupported input: ${{ github.ref }} / ${{ github.ref_type }}" + exit 1 + fi + + - name: Checkout docs sources + uses: actions/checkout@v6 + + - name: Checkout docs target branch + if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') + uses: actions/checkout@v6 + with: + fetch-depth: 0 + ref: ${{ steps.target-branch.outputs.value }} + path: docs-target + + - name: Setup Python + uses: actions/setup-python@v6 + with: + python-version: "3.10" + + - name: Install dependencies + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + + # Download the Linux wheel built in the previous job + - name: Download pre-built Linux wheel + uses: actions/download-artifact@v7 + with: + name: dist-manylinux-x86_64 + path: wheels/ + + # Install from the pre-built wheels + - name: Install from pre-built wheels + run: | + set -x + uv venv + # Install documentation dependencies + uv sync --dev --no-install-package datafusion --group docs + # Install all pre-built wheels + WHEELS=$(find wheels/ -name "*.whl") + if [ -n "$WHEELS" ]; then + echo "Installing wheels:" + echo "$WHEELS" + uv pip install wheels/*.whl + else + echo "ERROR: No wheels found!" + exit 1 + fi + + - name: Build docs + run: | + set -x + cd docs + curl -O https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv + curl -O https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2021-01.parquet + uv run --no-project make html + + - name: Copy & push the generated HTML + if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') + run: | + set -x + cd docs-target + # delete anything but: 1) '.'; 2) '..'; 3) .git/ + find ./ | grep -vE "^./$|^../$|^./.git" | xargs rm -rf + cp ../.asf.yaml . + cp -r ../docs/build/html/* . + git status --porcelain + if [ "$(git status --porcelain)" != "" ]; then + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git add --all + git commit -m 'Publish built docs triggered by ${{ github.sha }}' + git push || git push --force + fi diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 000000000..ab284b522 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,41 @@ +# 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. + +# CI workflow for pull requests - runs tests in DEBUG mode for faster feedback + +name: CI + +on: + pull_request: + branches: ["main"] + +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + build: + uses: ./.github/workflows/build.yml + with: + build_mode: debug + run_wheels: false + secrets: inherit + + test: + needs: build + uses: ./.github/workflows/test.yml + secrets: inherit diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 44481818e..2c8ecbc5e 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -25,10 +25,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Setup Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.14" - name: Audit licenses run: ./dev/release/run-rat.sh . diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml deleted file mode 100644 index 9037e0a5c..000000000 --- a/.github/workflows/docs.yaml +++ /dev/null @@ -1,95 +0,0 @@ -on: - push: - branches: - - main - tags-ignore: - - "**-rc**" - pull_request: - branches: - - main - -name: Deploy DataFusion Python site - -jobs: - debug-github-context: - name: Print github context - runs-on: ubuntu-latest - steps: - - name: Dump GitHub context - env: - GITHUB_CONTEXT: ${{ toJson(github) }} - run: | - echo "$GITHUB_CONTEXT" - build-docs: - name: Build docs - runs-on: ubuntu-latest - steps: - - name: Set target branch - if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') - id: target-branch - run: | - set -x - if test '${{ github.ref }}' = 'refs/heads/main'; then - echo "value=asf-staging" >> "$GITHUB_OUTPUT" - elif test '${{ github.ref_type }}' = 'tag'; then - echo "value=asf-site" >> "$GITHUB_OUTPUT" - else - echo "Unsupported input: ${{ github.ref }} / ${{ github.ref_type }}" - exit 1 - fi - - name: Checkout docs sources - uses: actions/checkout@v4 - - name: Checkout docs target branch - if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') - uses: actions/checkout@v4 - with: - fetch-depth: 0 - ref: ${{ steps.target-branch.outputs.value }} - path: docs-target - - name: Setup Python - uses: actions/setup-python@v5 - with: - python-version: "3.11" - - - name: Install Protoc - uses: arduino/setup-protoc@v3 - with: - version: '27.4' - repo-token: ${{ secrets.GITHUB_TOKEN }} - - - name: Install dependencies and build - uses: astral-sh/setup-uv@v5 - with: - enable-cache: true - - - name: Build repo - run: | - uv venv - uv sync --dev --no-install-package datafusion --group docs - uv run --no-project maturin develop --uv - - - name: Build docs - run: | - set -x - cd docs - curl -O https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv - curl -O https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2021-01.parquet - uv run --no-project make html - - - name: Copy & push the generated HTML - if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || github.ref_type == 'tag') - run: | - set -x - cd docs-target - # delete anything but: 1) '.'; 2) '..'; 3) .git/ - find ./ | grep -vE "^./$|^../$|^./.git" | xargs rm -rf - cp ../.asf.yaml . - cp -r ../docs/build/html/* . - git status --porcelain - if [ "$(git status --porcelain)" != "" ]; then - git config user.name "github-actions[bot]" - git config user.email "github-actions[bot]@users.noreply.github.com" - git add --all - git commit -m 'Publish built docs triggered by ${{ github.sha }}' - git push || git push --force - fi diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 000000000..bddc89eac --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,49 @@ +# 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. + +# Release workflow - runs tests in RELEASE mode and builds distribution wheels +# Triggered on: +# - Merges to main +# - Release candidate tags (*-rc*) +# - Release tags (e.g., 45.0.0) + +name: Release Build + +on: + push: + branches: + - "main" + tags: + - "*-rc*" # Release candidates (e.g., 45.0.0-rc1) + - "[0-9]+.*" # Release tags (e.g., 45.0.0) + +concurrency: + group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + build: + uses: ./.github/workflows/build.yml + with: + build_mode: release + run_wheels: true + secrets: inherit + + test: + needs: build + uses: ./.github/workflows/test.yml + secrets: inherit diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yml similarity index 56% rename from .github/workflows/test.yaml rename to .github/workflows/test.yml index da3582766..a2f304aa5 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yml @@ -15,16 +15,13 @@ # specific language governing permissions and limitations # under the License. -name: Python test -on: - push: - branches: [main] - pull_request: - branches: [main] +# Reusable workflow for running tests +# This ensures the same tests run for both debug (PRs) and release (main/tags) builds + +name: Test -concurrency: - group: ${{ github.repository }}-${{ github.head_ref || github.sha }}-${{ github.workflow }} - cancel-in-progress: true +on: + workflow_call: jobs: test-matrix: @@ -33,71 +30,80 @@ jobs: fail-fast: false matrix: python-version: - - "3.9" - "3.10" - "3.11" - "3.12" - "3.13" + - "3.14" toolchain: - "stable" steps: - - uses: actions/checkout@v4 - - - name: Setup Rust Toolchain - uses: dtolnay/rust-toolchain@stable - id: rust-toolchain - with: - components: clippy,rustfmt - - - name: Install Protoc - uses: arduino/setup-protoc@v3 - with: - version: '27.4' - repo-token: ${{ secrets.GITHUB_TOKEN }} + - uses: actions/checkout@v6 - name: Setup Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - name: Cache Cargo - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: ~/.cargo - key: cargo-cache-${{ steps.rust-toolchain.outputs.cachekey }}-${{ hashFiles('Cargo.lock') }} + key: cargo-cache-${{ matrix.toolchain }}-${{ hashFiles('Cargo.lock') }} - - name: Check Formatting - if: ${{ matrix.python-version == '3.10' && matrix.toolchain == 'stable' }} - run: cargo fmt -- --check + - name: Install dependencies + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true - - name: Run Clippy - if: ${{ matrix.python-version == '3.10' && matrix.toolchain == 'stable' }} - run: cargo clippy --all-targets --all-features -- -D clippy::all -D warnings -A clippy::redundant_closure + # Download the Linux wheel built in the build workflow + - name: Download pre-built Linux wheel + uses: actions/download-artifact@v7 + with: + name: dist-manylinux-x86_64 + path: wheels/ - - name: Install dependencies and build - uses: astral-sh/setup-uv@v5 + # Download the FFI test wheel + - name: Download pre-built FFI test wheel + uses: actions/download-artifact@v7 with: - enable-cache: true + name: test-ffi-manylinux-x86_64 + path: wheels/ + + # Install from the pre-built wheels + - name: Install from pre-built wheels + run: | + set -x + uv venv + # Install development dependencies + uv sync --dev --no-install-package datafusion + # Install all pre-built wheels + WHEELS=$(find wheels/ -name "*.whl") + if [ -n "$WHEELS" ]; then + echo "Installing wheels:" + echo "$WHEELS" + uv pip install wheels/*.whl + else + echo "ERROR: No wheels found!" + exit 1 + fi - name: Run tests env: RUST_BACKTRACE: 1 run: | git submodule update --init - uv sync --dev --no-install-package datafusion - uv run --no-project maturin develop --uv - uv run --no-project pytest -v . + uv run --no-project pytest -v --import-mode=importlib - name: FFI unit tests run: | - cd examples/ffi-table-provider - uv run --no-project maturin develop --uv - uv run --no-project pytest python/tests/_test_table_provider.py + cd examples/datafusion-ffi-example + uv run --no-project pytest python/tests/_test*.py - name: Cache the generated dataset id: cache-tpch-dataset - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: benchmarks/tpch/data key: tpch-data-2.18.0 diff --git a/.github/workflows/verify-release-candidate.yml b/.github/workflows/verify-release-candidate.yml new file mode 100644 index 000000000..a10a4faa9 --- /dev/null +++ b/.github/workflows/verify-release-candidate.yml @@ -0,0 +1,78 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +name: Verify Release Candidate + +# NOTE: This workflow is intended to be run manually via workflow_dispatch. + +on: + workflow_dispatch: + inputs: + version: + description: Version number (e.g., 52.0.0) + required: true + type: string + rc_number: + description: Release candidate number (e.g., 0) + required: true + type: string + +concurrency: + group: ${{ github.repository }}-${{ github.ref }}-${{ github.workflow }} + cancel-in-progress: true + +jobs: + verify: + name: Verify RC (${{ matrix.os }}-${{ matrix.arch }}) + strategy: + fail-fast: false + matrix: + include: + # Linux + - os: linux + arch: x64 + runner: ubuntu-latest + - os: linux + arch: arm64 + runner: ubuntu-24.04-arm + + # macOS + - os: macos + arch: arm64 + runner: macos-latest + - os: macos + arch: x64 + runner: macos-15-intel + + # Windows + - os: windows + arch: x64 + runner: windows-latest + runs-on: ${{ matrix.runner }} + steps: + - name: Checkout repository + uses: actions/checkout@v6 + + - name: Set up protoc + uses: arduino/setup-protoc@v3 + with: + version: "27.4" + repo-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Run release candidate verification + shell: bash + run: ./dev/release/verify-release-candidate.sh "${{ inputs.version }}" "${{ inputs.rc_number }}" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index abcfcf321..8ae6a4e32 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,7 +22,7 @@ repos: - id: actionlint-docker - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.9.10 + rev: v0.15.1 hooks: # Run the linter. - id: ruff @@ -33,7 +33,7 @@ repos: - id: rust-fmt name: Rust fmt description: Run cargo fmt on files included in the commit. rustfmt should be installed before-hand. - entry: cargo fmt --all -- + entry: cargo +nightly fmt --all -- pass_filenames: true types: [file, rust] language: system @@ -45,5 +45,13 @@ repos: types: [file, rust] language: system + - repo: https://github.com/codespell-project/codespell + rev: v2.4.1 + hooks: + - id: codespell + args: [ --toml, "pyproject.toml"] + additional_dependencies: + - tomli + default_language_version: python: python3 diff --git a/Cargo.lock b/Cargo.lock index 5c7f2bf3c..e44c84b97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "abi_stable" @@ -50,36 +50,21 @@ dependencies = [ "core_extensions", ] -[[package]] -name = "addr2line" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1" -dependencies = [ - "gimli", -] - [[package]] name = "adler2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" - -[[package]] -name = "adler32" -version = "1.2.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "ahash" -version = "0.8.11" +version = "0.8.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", "const-random", - "getrandom 0.2.15", + "getrandom 0.3.4", "once_cell", "version_check", "zerocopy", @@ -87,9 +72,9 @@ dependencies = [ [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" dependencies = [ "memchr", ] @@ -115,12 +100,6 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" -[[package]] -name = "android-tzdata" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" - [[package]] name = "android_system_properties" version = "0.1.5" @@ -132,22 +111,24 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.95" +version = "1.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" +checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" [[package]] name = "apache-avro" -version = "0.17.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aef82843a0ec9f8b19567445ad2421ceeb1d711514384bdd3d49fe37102ee13" +checksum = "36fa98bc79671c7981272d91a8753a928ff6a1cd8e4f20a44c45bd5d313840bf" dependencies = [ "bigdecimal", - "bzip2 0.4.4", + "bon", + "bzip2", "crc32fast", "digest", - "libflate", + "liblzma", "log", + "miniz_oxide", "num-bigint", "quad-rand", "rand", @@ -158,13 +139,29 @@ dependencies = [ "snap", "strum", "strum_macros", - "thiserror 1.0.69", - "typed-builder", + "thiserror", "uuid", - "xz2", "zstd", ] +[[package]] +name = "ar_archive_writer" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eb93bbb63b9c227414f6eb3a0adfddca591a8ce1e9b60661bb08969b87e340b" +dependencies = [ + "object", +] + +[[package]] +name = "arc-swap" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f3647c145568cec02c42054e07bdf9a5a698e15b466fb2341bfc393cd24aa5" +dependencies = [ + "rustversion", +] + [[package]] name = "arrayref" version = "0.3.9" @@ -179,9 +176,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "755b6da235ac356a869393c23668c663720b8749dd6f15e52b6c214b4b964cc7" +checksum = "602268ce9f569f282cedb9a9f6bac569b680af47b9b077d515900c03c5d190da" dependencies = [ "arrow-arith", "arrow-array", @@ -192,32 +189,32 @@ dependencies = [ "arrow-ipc", "arrow-json", "arrow-ord", + "arrow-pyarrow", "arrow-row", "arrow-schema", "arrow-select", "arrow-string", - "pyo3", ] [[package]] name = "arrow-arith" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64656a1e0b13ca766f8440752e9a93e11014eec7b67909986f83ed0ab1fe37b8" +checksum = "cd53c6bf277dea91f136ae8e3a5d7041b44b5e489e244e637d00ae302051f56f" dependencies = [ "arrow-array", "arrow-buffer", "arrow-data", "arrow-schema", "chrono", - "num", + "num-traits", ] [[package]] name = "arrow-array" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57a4a6d2896083cfbdf84a71a863b22460d0708f8206a8373c52e326cc72ea1a" +checksum = "e53796e07a6525edaf7dc28b540d477a934aff14af97967ad1d5550878969b9e" dependencies = [ "ahash", "arrow-buffer", @@ -226,47 +223,51 @@ dependencies = [ "chrono", "chrono-tz", "half", - "hashbrown 0.15.2", - "num", + "hashbrown 0.16.1", + "num-complex", + "num-integer", + "num-traits", ] [[package]] name = "arrow-buffer" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cef870583ce5e4f3b123c181706f2002fb134960f9a911900f64ba4830c7a43a" +checksum = "f2c1a85bb2e94ee10b76531d8bc3ce9b7b4c0d508cabfb17d477f63f2617bd20" dependencies = [ "bytes", "half", - "num", + "num-bigint", + "num-traits", ] [[package]] name = "arrow-cast" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ac7eba5a987f8b4a7d9629206ba48e19a1991762795bbe5d08497b7736017ee" +checksum = "89fb245db6b0e234ed8e15b644edb8664673fefe630575e94e62cd9d489a8a26" dependencies = [ "arrow-array", "arrow-buffer", "arrow-data", + "arrow-ord", "arrow-schema", "arrow-select", "atoi", - "base64 0.22.1", + "base64", "chrono", "comfy-table", "half", "lexical-core", - "num", + "num-traits", "ryu", ] [[package]] name = "arrow-csv" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90f12542b8164398fc9ec595ff783c4cf6044daa89622c5a7201be920e4c0d4c" +checksum = "d374882fb465a194462527c0c15a93aa19a554cf690a6b77a26b2a02539937a7" dependencies = [ "arrow-array", "arrow-cast", @@ -274,41 +275,43 @@ dependencies = [ "chrono", "csv", "csv-core", - "lazy_static", "regex", ] [[package]] name = "arrow-data" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b095e8a4f3c309544935d53e04c3bfe4eea4e71c3de6fe0416d1f08bb4441a83" +checksum = "189d210bc4244c715fa3ed9e6e22864673cccb73d5da28c2723fb2e527329b33" dependencies = [ "arrow-buffer", "arrow-schema", "half", - "num", + "num-integer", + "num-traits", ] [[package]] name = "arrow-ipc" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65c63da4afedde2b25ef69825cd4663ca76f78f79ffe2d057695742099130ff6" +checksum = "7968c2e5210c41f4909b2ef76f6e05e172b99021c2def5edf3cc48fdd39d1d6c" dependencies = [ "arrow-array", "arrow-buffer", "arrow-data", "arrow-schema", + "arrow-select", "flatbuffers", "lz4_flex", + "zstd", ] [[package]] name = "arrow-json" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9551d9400532f23a370cabbea1dc5a53c49230397d41f96c4c8eedf306199305" +checksum = "92111dba5bf900f443488e01f00d8c4ddc2f47f5c50039d18120287b580baa22" dependencies = [ "arrow-array", "arrow-buffer", @@ -318,17 +321,21 @@ dependencies = [ "chrono", "half", "indexmap", + "itoa", "lexical-core", - "num", - "serde", + "memchr", + "num-traits", + "ryu", + "serde_core", "serde_json", + "simdutf8", ] [[package]] name = "arrow-ord" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c07223476f8219d1ace8cd8d85fa18c4ebd8d945013f25ef5c72e85085ca4ee" +checksum = "211136cb253577ee1a6665f741a13136d4e563f64f5093ffd6fb837af90b9495" dependencies = [ "arrow-array", "arrow-buffer", @@ -337,11 +344,23 @@ dependencies = [ "arrow-select", ] +[[package]] +name = "arrow-pyarrow" +version = "58.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "205437da4c0877c756c81bfe847a621d0a740cd00a155109d65510a1a62ebcd9" +dependencies = [ + "arrow-array", + "arrow-data", + "arrow-schema", + "pyo3", +] + [[package]] name = "arrow-row" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91b194b38bfd89feabc23e798238989c6648b2506ad639be42ec8eb1658d82c4" +checksum = "8e0f20145f9f5ea3fe383e2ba7a7487bf19be36aa9dbf5dd6a1f92f657179663" dependencies = [ "arrow-array", "arrow-buffer", @@ -352,32 +371,34 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f40f6be8f78af1ab610db7d9b236e21d587b7168e368a36275d2e5670096735" +checksum = "1b47e0ca91cc438d2c7879fe95e0bca5329fff28649e30a88c6f760b1faeddcb" dependencies = [ - "bitflags 2.8.0", + "bitflags", + "serde_core", + "serde_json", ] [[package]] name = "arrow-select" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac265273864a820c4a179fc67182ccc41ea9151b97024e1be956f0f2369c2539" +checksum = "750a7d1dda177735f5e82a314485b6915c7cccdbb278262ac44090f4aba4a325" dependencies = [ "ahash", "arrow-array", "arrow-buffer", "arrow-data", "arrow-schema", - "num", + "num-traits", ] [[package]] name = "arrow-string" -version = "54.2.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c8eed43be4ead49128370f7131f054839d3d6003e52aebf64322470b8fbd0" +checksum = "e1eab1208bc4fe55d768cdc9b9f3d9df5a794cdb3ee2586bf89f9b30dc31ad8c" dependencies = [ "arrow-array", "arrow-buffer", @@ -385,7 +406,7 @@ dependencies = [ "arrow-schema", "arrow-select", "memchr", - "num", + "num-traits", "regex", "regex-syntax", ] @@ -404,19 +425,14 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.18" +version = "0.4.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df895a515f70646414f4b45c0b79082783b80552b373a68283012928df56f522" +checksum = "7d67d43201f4d20c78bcda740c142ca52482d81da80681533d33bf3f0596c8e2" dependencies = [ - "bzip2 0.4.4", - "flate2", - "futures-core", - "memchr", + "compression-codecs", + "compression-core", "pin-project-lite", "tokio", - "xz2", - "zstd", - "zstd-safe", ] [[package]] @@ -436,18 +452,18 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "async-trait" -version = "0.1.86" +version = "0.1.89" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "644dd749086bf3771a2fbc5f256fdb982d53f011c7d5d560304eafeecebce79d" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -467,30 +483,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" - -[[package]] -name = "backtrace" -version = "0.3.74" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a" -dependencies = [ - "addr2line", - "cfg-if", - "libc", - "miniz_oxide", - "object", - "rustc-demangle", - "windows-targets", -] - -[[package]] -name = "base64" -version = "0.21.7" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "base64" @@ -500,9 +495,9 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "bigdecimal" -version = "0.4.7" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f31f3af01c5c65a07985c804d3366560e6fa7883d640a122819b14ec327482c" +checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", @@ -514,15 +509,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" - -[[package]] -name = "bitflags" -version = "2.8.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f68f53c83ab957f72c32642f3868eec03eb974d1fb82e453128456482613d36" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" [[package]] name = "blake2" @@ -535,15 +524,16 @@ dependencies = [ [[package]] name = "blake3" -version = "1.5.5" +version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8ee0c1824c4dea5b5f81736aff91bae041d2c07ee1192bec91054e10e3e601e" +checksum = "2468ef7d57b3fb7e16b576e8377cdbde2320c60e1491e961d11da40fc4f02a2d" dependencies = [ "arrayref", "arrayvec", "cc", "cfg-if", "constant_time_eq", + "cpufeatures", ] [[package]] @@ -555,11 +545,36 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bon" +version = "3.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d13a61f2963b88eef9c1be03df65d42f6996dfeac1054870d950fcf66686f83" +dependencies = [ + "bon-macros", + "rustversion", +] + +[[package]] +name = "bon-macros" +version = "3.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d314cc62af2b6b0c65780555abb4d02a03dd3b799cd42419044f0c38d99738c0" +dependencies = [ + "darling", + "ident_case", + "prettyplease", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.117", +] + [[package]] name = "brotli" -version = "7.0.0" +version = "8.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc97b8f16f944bba54f0433f07e30be199b6dc2bd25937444bbad560bcea29bd" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -568,9 +583,9 @@ dependencies = [ [[package]] name = "brotli-decompressor" -version = "4.0.2" +version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74fa05ad7d803d413eb8380983b092cbbaf9a85f151b871360e7b00cd7060b37" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -578,9 +593,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.17.0" +version = "3.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf" +checksum = "5c6f81257d10a0f602a294ae4182251151ff97dbb504ef9afcdda4a64b24d9b4" [[package]] name = "byteorder" @@ -590,46 +605,26 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9" - -[[package]] -name = "bzip2" -version = "0.4.4" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" -dependencies = [ - "bzip2-sys", - "libc", -] +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" [[package]] name = "bzip2" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75b89e7c29231c673a61a46e722602bcd138298f6b9e81e71119693534585f5c" -dependencies = [ - "bzip2-sys", -] - -[[package]] -name = "bzip2-sys" -version = "0.1.12+1.0.8" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72ebc2f1a417f01e1da30ef264ee86ae31d2dcd2d603ea283d3c244a883ca2a9" +checksum = "f3a53fac24f34a81bc9954b5d6cfce0c21e18ec6959f44f56e8e90e4bb7c346c" dependencies = [ - "cc", - "libc", - "pkg-config", + "libbz2-rs-sys", ] [[package]] name = "cc" -version = "1.2.14" +version = "1.2.56" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c3d1b2e905a3a7b00a6141adb0e4c0bb941d11caf55349d863942a1cc44e3c9" +checksum = "aebf35691d1bfb0ac386a69bac2fde4dd276fb618cf8bf4f5318fe285e821bb2" dependencies = [ + "find-msvc-tools", "jobserver", "libc", "shlex", @@ -637,9 +632,9 @@ dependencies = [ [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" @@ -649,57 +644,66 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chrono" -version = "0.4.39" +version = "0.4.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" dependencies = [ - "android-tzdata", "iana-time-zone", "num-traits", "serde", - "windows-targets", + "windows-link", ] [[package]] name = "chrono-tz" -version = "0.10.1" +version = "0.10.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c6ac4f2c0bf0f44e9161aec9675e1050aa4a530663c4a9e37e108fa948bca9f" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "chrono-tz-build", "phf", ] -[[package]] -name = "chrono-tz-build" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94fea34d77a245229e7746bd2beb786cd2a896f306ff491fb8cecb3074b10a7" -dependencies = [ - "parse-zoneinfo", - "phf_codegen", -] - [[package]] name = "cmake" -version = "0.1.54" +version = "0.1.57" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7caa3f9de89ddbe2c607f4101924c5abec803763ae9534e4f4d7d8f84aa81f0" +checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" dependencies = [ "cc", ] [[package]] name = "comfy-table" -version = "7.1.4" +version = "7.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" +checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47" dependencies = [ "unicode-segmentation", "unicode-width", ] +[[package]] +name = "compression-codecs" +version = "0.4.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb7b51a7d9c967fc26773061ba86150f19c50c0d65c887cb1fbe295fd16619b7" +dependencies = [ + "bzip2", + "compression-core", + "flate2", + "liblzma", + "memchr", + "zstd", + "zstd-safe", +] + +[[package]] +name = "compression-core" +version = "0.4.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75984efb6ed102a0d42db99afb6c1948f0380d1d91808d5529916e6c08b49d8d" + [[package]] name = "const-random" version = "0.1.18" @@ -715,28 +719,31 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" dependencies = [ - "getrandom 0.2.15", + "getrandom 0.2.17", "once_cell", "tiny-keccak", ] [[package]] name = "const_panic" -version = "0.2.12" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2459fc9262a1aa204eb4b5764ad4f189caec88aea9634389c0a25f8be7f6265e" +checksum = "e262cdaac42494e3ae34c43969f9cdeb7da178bdb4b66fa6a1ea2edb4c8ae652" +dependencies = [ + "typewit", +] [[package]] name = "constant_time_eq" -version = "0.3.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" [[package]] name = "core-foundation" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b55271e5c8c478ad3f38ad24ef34923091e0548492a266d19b3c0b4d82574c63" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" dependencies = [ "core-foundation-sys", "libc", @@ -748,29 +755,20 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core2" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b49ba7ef1ad6107f8824dbe97de947cbaac53c44e7f9756a1fba0d37c1eec505" -dependencies = [ - "memchr", -] - [[package]] name = "core_extensions" -version = "1.5.3" +version = "1.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92c71dc07c9721607e7a16108336048ee978c3a8b129294534272e8bac96c0ee" +checksum = "42bb5e5d0269fd4f739ea6cedaf29c16d81c27a7ce7582008e90eb50dcd57003" dependencies = [ "core_extensions_proc_macros", ] [[package]] name = "core_extensions_proc_macros" -version = "1.5.3" +version = "1.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69f3b219d28b6e3b4ac87bc1fc522e0803ab22e055da177bff0068c4150c61a6" +checksum = "533d38ecd2709b7608fb8e18e4504deb99e9a72879e6aa66373a76d8dc4259ea" [[package]] name = "cpufeatures" @@ -783,18 +781,18 @@ dependencies = [ [[package]] name = "crc32fast" -version = "1.4.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" dependencies = [ "cfg-if", ] [[package]] name = "crossbeam-channel" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06ba6d68e24814cb8de6bb986db8222d3a027d15872cabc0d18817bc3c0e4471" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" dependencies = [ "crossbeam-utils", ] @@ -807,46 +805,84 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crunchy" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43da5946c66ffcc7745f48db692ffbb10a83bfe0afd96235c5c2a4fb23994929" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "crypto-common" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", "typenum", ] +[[package]] +name = "cstr" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68523903c8ae5aacfa32a0d9ae60cadeb764e1da14ee0d26b1f3089f13a54636" +dependencies = [ + "proc-macro2", + "quote", +] + [[package]] name = "csv" -version = "1.3.1" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acdc4883a9c96732e4733212c01447ebd805833b7275a73ca3ee080fd77afdaf" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" dependencies = [ "csv-core", "itoa", "ryu", - "serde", + "serde_core", ] [[package]] name = "csv-core" -version = "0.1.12" +version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d02f3b0da4c6504f86e9cd789d8dbafab48c2321be74e9987593de5a894d93d" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" dependencies = [ "memchr", ] [[package]] -name = "dary_heap" -version = "0.3.7" +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.117", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.117", +] [[package]] name = "dashmap" @@ -864,25 +900,28 @@ dependencies = [ [[package]] name = "datafusion" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eae420e7a5b0b7f1c39364cc76cbcd0f5fdc416b2514ae3847c2676bbd60702a" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ - "apache-avro", "arrow", - "arrow-array", - "arrow-ipc", "arrow-schema", - "async-compression", "async-trait", "bytes", - "bzip2 0.5.1", + "bzip2", "chrono", "datafusion-catalog", + "datafusion-catalog-listing", "datafusion-common", "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-datasource-arrow", + "datafusion-datasource-avro", + "datafusion-datasource-csv", + "datafusion-datasource-json", + "datafusion-datasource-parquet", "datafusion-execution", "datafusion-expr", + "datafusion-expr-common", "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-nested", @@ -890,16 +929,17 @@ dependencies = [ "datafusion-functions-window", "datafusion-optimizer", "datafusion-physical-expr", + "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-optimizer", "datafusion-physical-plan", + "datafusion-session", "datafusion-sql", "flate2", "futures", - "glob", - "itertools 0.14.0", + "itertools", + "liblzma", "log", - "num-traits", "object_store", "parking_lot", "parquet", @@ -908,51 +948,71 @@ dependencies = [ "sqlparser", "tempfile", "tokio", - "tokio-util", "url", "uuid", - "xz2", "zstd", ] [[package]] name = "datafusion-catalog" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f27987bc22b810939e8dfecc55571e9d50355d6ea8ec1c47af8383a76a6d0e1" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "async-trait", "dashmap", "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", - "datafusion-sql", + "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", + "object_store", "parking_lot", - "sqlparser", + "tokio", +] + +[[package]] +name = "datafusion-catalog-listing" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "async-trait", + "datafusion-catalog", + "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-adapter", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "futures", + "itertools", + "log", + "object_store", ] [[package]] name = "datafusion-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3f6d5b8c9408cc692f7c194b8aa0c0f9b253e065a8d960ad9cdc2a13e697602" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "apache-avro", "arrow", - "arrow-array", - "arrow-buffer", "arrow-ipc", - "arrow-schema", - "base64 0.22.1", + "chrono", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", + "itertools", "libc", "log", "object_store", @@ -966,47 +1026,199 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d4603c8e8a4baf77660ab7074cc66fc15cc8a18f2ce9dfadb755fc6ee294e48" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ + "futures", "log", "tokio", ] [[package]] -name = "datafusion-doc" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5bf4bc68623a5cf231eed601ed6eb41f46a37c4d15d11a0bff24cbc8396cd66" - -[[package]] -name = "datafusion-execution" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88b491c012cdf8e051053426013429a76f74ee3c2db68496c79c323ca1084d27" +name = "datafusion-datasource" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", - "dashmap", + "async-compression", + "async-trait", + "bytes", + "bzip2", + "chrono", "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-adapter", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "flate2", "futures", + "glob", + "itertools", + "liblzma", "log", "object_store", - "parking_lot", "rand", - "tempfile", + "tokio", + "tokio-util", "url", + "zstd", ] [[package]] -name = "datafusion-expr" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5a181408d4fc5dc22f9252781a8f39f2d0e5d1b33ec9bde242844980a2689c1" +name = "datafusion-datasource-arrow" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", - "chrono", + "arrow-ipc", + "async-trait", + "bytes", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "itertools", + "object_store", + "tokio", +] + +[[package]] +name = "datafusion-datasource-avro" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "apache-avro", + "arrow", + "async-trait", + "bytes", + "datafusion-common", + "datafusion-datasource", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "num-traits", + "object_store", +] + +[[package]] +name = "datafusion-datasource-csv" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "object_store", + "regex", + "tokio", +] + +[[package]] +name = "datafusion-datasource-json" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "object_store", + "serde_json", + "tokio", + "tokio-stream", +] + +[[package]] +name = "datafusion-datasource-parquet" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate-common", + "datafusion-physical-expr", + "datafusion-physical-expr-adapter", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-pruning", + "datafusion-session", + "futures", + "itertools", + "log", + "object_store", + "parking_lot", + "parquet", + "tokio", +] + +[[package]] +name = "datafusion-doc" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" + +[[package]] +name = "datafusion-execution" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "arrow-buffer", + "async-trait", + "chrono", + "dashmap", + "datafusion-common", + "datafusion-expr", + "datafusion-physical-expr-common", + "futures", + "log", + "object_store", + "parking_lot", + "rand", + "tempfile", + "url", +] + +[[package]] +name = "datafusion-expr" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "async-trait", + "chrono", "datafusion-common", "datafusion-doc", "datafusion-expr-common", @@ -1014,6 +1226,7 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr-common", "indexmap", + "itertools", "paste", "recursive", "serde_json", @@ -1022,30 +1235,38 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1129b48e8534d8c03c6543bcdccef0b55c8ac0c1272a15a56c67068b6eb1885" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "datafusion-common", - "itertools 0.14.0", + "indexmap", + "itertools", "paste", ] [[package]] name = "datafusion-ffi" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff47a79d442207c168c6e3e1d970c248589c148e4800e5b285ac1b2cb1a230f8" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "abi_stable", "arrow", - "arrow-array", "arrow-schema", "async-ffi", "async-trait", - "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", "datafusion-proto", + "datafusion-proto-common", + "datafusion-session", "futures", "log", "prost", @@ -1053,29 +1274,50 @@ dependencies = [ "tokio", ] +[[package]] +name = "datafusion-ffi-example" +version = "52.0.0" +dependencies = [ + "arrow", + "arrow-array", + "arrow-schema", + "async-trait", + "datafusion-catalog", + "datafusion-common", + "datafusion-expr", + "datafusion-ffi", + "datafusion-functions-aggregate", + "datafusion-functions-window", + "datafusion-python-util", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + [[package]] name = "datafusion-functions" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6125874e4856dfb09b59886784fcb74cde5cfc5930b3a80a1a728ef7a010df6b" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "arrow-buffer", - "base64 0.22.1", + "base64", "blake2", "blake3", "chrono", + "chrono-tz", "datafusion-common", "datafusion-doc", "datafusion-execution", "datafusion-expr", "datafusion-expr-common", "datafusion-macros", - "hashbrown 0.14.5", "hex", - "itertools 0.14.0", + "itertools", "log", "md-5", + "memchr", + "num-traits", "rand", "regex", "sha2", @@ -1085,14 +1327,11 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3add7b1d3888e05e7c95f2b281af900ca69ebdcb21069ba679b33bde8b3b9d6" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "arrow", - "arrow-buffer", - "arrow-schema", "datafusion-common", "datafusion-doc", "datafusion-execution", @@ -1103,14 +1342,14 @@ dependencies = [ "datafusion-physical-expr-common", "half", "log", + "num-traits", "paste", ] [[package]] name = "datafusion-functions-aggregate-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e18baa4cfc3d2f144f74148ed68a1f92337f5072b6dde204a0dbbdf3324989c" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "arrow", @@ -1121,33 +1360,32 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ec5ee8cecb0dc370291279673097ddabec03a011f73f30d7f1096457127e03e" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", - "arrow-array", - "arrow-buffer", "arrow-ord", - "arrow-schema", "datafusion-common", "datafusion-doc", "datafusion-execution", "datafusion-expr", + "datafusion-expr-common", "datafusion-functions", "datafusion-functions-aggregate", + "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", - "itertools 0.14.0", + "hashbrown 0.16.1", + "itertools", + "itoa", "log", "paste", ] [[package]] name = "datafusion-functions-table" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c403ddd473bbb0952ba880008428b3c7febf0ed3ce1eec35a205db20efb2a36" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "async-trait", @@ -1161,10 +1399,10 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ab18c2fb835614d06a75f24a9e09136d3a8c12a92d97c95a6af316a1787a9c5" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ + "arrow", "datafusion-common", "datafusion-doc", "datafusion-expr", @@ -1178,9 +1416,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a77b73bc15e7d1967121fdc7a55d819bfb9d6c03766a6c322247dce9094a53a4" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1188,28 +1425,27 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09369b8d962291e808977cf94d495fd8b5b38647232d7ef562c27ac0f495b0af" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ - "datafusion-expr", + "datafusion-doc", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "datafusion-optimizer" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2403a7e4a84637f3de7d8d4d7a9ccc0cc4be92d89b0161ba3ee5be82f0531c54" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "chrono", "datafusion-common", "datafusion-expr", + "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools 0.14.0", + "itertools", "log", "recursive", "regex", @@ -1218,52 +1454,63 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86ff72ac702b62dbf2650c4e1d715ebd3e4aab14e3885e72e8549e250307347c" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "arrow", - "arrow-array", - "arrow-buffer", - "arrow-schema", "datafusion-common", "datafusion-expr", "datafusion-expr-common", "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", - "itertools 0.14.0", - "log", + "itertools", + "parking_lot", "paste", "petgraph", + "recursive", + "tokio", +] + +[[package]] +name = "datafusion-physical-expr-adapter" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "datafusion-common", + "datafusion-expr", + "datafusion-functions", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "itertools", ] [[package]] name = "datafusion-physical-expr-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60982b7d684e25579ee29754b4333057ed62e2cc925383c5f0bd8cab7962f435" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "arrow", - "arrow-buffer", + "chrono", "datafusion-common", "datafusion-expr-common", - "hashbrown 0.14.5", - "itertools 0.14.0", + "hashbrown 0.16.1", + "indexmap", + "itertools", + "parking_lot", ] [[package]] name = "datafusion-physical-optimizer" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac5e85c189d5238a5cf181a624e450c4cd4c66ac77ca551d6f3ff9080bac90bb" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", - "arrow-schema", "datafusion-common", "datafusion-execution", "datafusion-expr", @@ -1271,40 +1518,37 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "futures", - "itertools 0.14.0", - "log", + "datafusion-pruning", + "itertools", "recursive", - "url", ] [[package]] name = "datafusion-physical-plan" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c36bf163956d7e2542657c78b3383fdc78f791317ef358a359feffcdb968106f" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "ahash", "arrow", - "arrow-array", - "arrow-buffer", "arrow-ord", "arrow-schema", "async-trait", - "chrono", "datafusion-common", "datafusion-common-runtime", "datafusion-execution", "datafusion-expr", + "datafusion-functions", + "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr", "datafusion-physical-expr-common", "futures", "half", - "hashbrown 0.14.5", + "hashbrown 0.16.1", "indexmap", - "itertools 0.14.0", + "itertools", "log", + "num-traits", "parking_lot", "pin-project-lite", "tokio", @@ -1312,66 +1556,123 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2db5d79f0c974041787b899d24dc91bdab2ff112d1942dd71356a4ce3b407e6c" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "chrono", - "datafusion", + "datafusion-catalog", + "datafusion-catalog-listing", "datafusion-common", + "datafusion-datasource", + "datafusion-datasource-arrow", + "datafusion-datasource-csv", + "datafusion-datasource-json", + "datafusion-datasource-parquet", + "datafusion-execution", "datafusion-expr", + "datafusion-functions-table", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", "datafusion-proto-common", "object_store", "prost", + "rand", ] [[package]] name = "datafusion-proto-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de21bde1603aac0ff32cf478e47081be6e3583c6861fe8f57034da911efe7578" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", "datafusion-common", "prost", ] +[[package]] +name = "datafusion-pruning" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "arrow", + "datafusion-common", + "datafusion-datasource", + "datafusion-expr-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "itertools", + "log", +] + [[package]] name = "datafusion-python" -version = "45.2.0" +version = "52.0.0" dependencies = [ "arrow", + "arrow-select", "async-trait", + "cstr", "datafusion", "datafusion-ffi", "datafusion-proto", + "datafusion-python-util", "datafusion-substrait", "futures", + "log", "mimalloc", "object_store", + "parking_lot", "prost", "prost-types", "pyo3", "pyo3-async-runtimes", "pyo3-build-config", + "pyo3-log", + "serde_json", "tokio", "url", "uuid", ] +[[package]] +name = "datafusion-python-util" +version = "52.0.0" +dependencies = [ + "arrow", + "datafusion", + "datafusion-ffi", + "prost", + "pyo3", + "tokio", +] + +[[package]] +name = "datafusion-session" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" +dependencies = [ + "async-trait", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-plan", + "parking_lot", +] + [[package]] name = "datafusion-sql" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13caa4daede211ecec53c78b13c503b592794d125f9a3cc3afe992edf9e7f43" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ "arrow", - "arrow-array", - "arrow-schema", "bigdecimal", + "chrono", "datafusion-common", "datafusion-expr", + "datafusion-functions-nested", "indexmap", "log", "recursive", @@ -1381,20 +1682,20 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1634405abd8bd3c64c352f2da2f2aec6d80a815930257e0db0ce4ff5daf00944" +version = "53.0.0" +source = "git+https://github.com/apache/datafusion.git?rev=35749607f585b3bf25b66b7d2289c56c18d03e4f#35749607f585b3bf25b66b7d2289c56c18d03e4f" dependencies = [ - "arrow-buffer", "async-recursion", "async-trait", "chrono", "datafusion", - "itertools 0.14.0", + "half", + "itertools", "object_store", "pbjson-types", "prost", "substrait", + "tokio", "url", ] @@ -1417,20 +1718,20 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "dyn-clone" -version = "1.0.18" +version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "feeef44e73baff3a26d371801df019877a9866a8c493d315ab00177843314f35" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" [[package]] name = "either" -version = "1.13.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" [[package]] name = "equivalent" @@ -1440,12 +1741,12 @@ checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" -version = "0.3.10" +version = "0.3.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -1454,6 +1755,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + [[package]] name = "fixedbitset" version = "0.5.7" @@ -1462,22 +1769,23 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flatbuffers" -version = "24.12.23" +version = "25.12.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" dependencies = [ - "bitflags 1.3.2", + "bitflags", "rustc_version", ] [[package]] name = "flate2" -version = "1.0.35" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c936bfdafb507ebbf50b8074c54fa31c5be9a1e7e5f467dd659697041407d07c" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", + "zlib-rs", ] [[package]] @@ -1488,24 +1796,30 @@ checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" [[package]] name = "foldhash" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0d2fde1f7b3d48b8395d5f2de76c18a528bd6a9cdde438df747bfcba3e05d6f" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" [[package]] name = "form_urlencoded" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" dependencies = [ "percent-encoding", ] [[package]] name = "futures" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" dependencies = [ "futures-channel", "futures-core", @@ -1518,9 +1832,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", "futures-sink", @@ -1528,15 +1842,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" dependencies = [ "futures-core", "futures-task", @@ -1545,38 +1859,38 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" [[package]] name = "futures-task" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-channel", "futures-core", @@ -1586,7 +1900,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -1611,46 +1924,55 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.15" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", "js-sys", "libc", - "wasi 0.11.0+wasi-snapshot-preview1", + "wasi", "wasm-bindgen", ] [[package]] name = "getrandom" -version = "0.3.1" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", - "wasi 0.13.3+wasi-0.2.2", - "windows-targets", + "r-efi", + "wasip2", + "wasm-bindgen", ] [[package]] -name = "gimli" -version = "0.31.1" +name = "getrandom" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", + "wasip3", +] [[package]] name = "glob" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" [[package]] name = "h2" -version = "0.4.7" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccae279728d634d083c00f6099cb58f01cc99c145b84b8be2f6c74618d79922e" +checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" dependencies = [ "atomic-waker", "bytes", @@ -1667,13 +1989,14 @@ dependencies = [ [[package]] name = "half" -version = "2.4.1" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" dependencies = [ "cfg-if", "crunchy", "num-traits", + "zerocopy", ] [[package]] @@ -1681,20 +2004,25 @@ name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "ahash", - "allocator-api2", + "foldhash 0.1.5", ] [[package]] name = "hashbrown" -version = "0.15.2" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" dependencies = [ "allocator-api2", "equivalent", - "foldhash", + "foldhash 0.2.0", ] [[package]] @@ -1711,12 +2039,11 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "http" -version = "1.2.0" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f16ca2af56261c99fba8bac40a10251ce8188205a4c448fbb745a2e4daa76fea" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" dependencies = [ "bytes", - "fnv", "itoa", ] @@ -1732,12 +2059,12 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "793429d76616a256bcb62c2a2ec2bed781c8307e797e2598c50010f2bee2544f" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", - "futures-util", + "futures-core", "http", "http-body", "pin-project-lite", @@ -1745,31 +2072,33 @@ dependencies = [ [[package]] name = "httparse" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2d708df4e7140240a16cd6ab0ab65c972d7433ab77819ea693fde9c43811e2a" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" [[package]] name = "humantime" -version = "2.1.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" +checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" [[package]] name = "hyper" -version = "1.6.0" +version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc2b571658e38e0c01b1fdca3bbbe93c00d3d71693ff2770043f8c29bc7d6f80" +checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" dependencies = [ + "atomic-waker", "bytes", "futures-channel", - "futures-util", + "futures-core", "h2", "http", "http-body", "httparse", "itoa", "pin-project-lite", + "pin-utils", "smallvec", "tokio", "want", @@ -1777,11 +2106,10 @@ dependencies = [ [[package]] name = "hyper-rustls" -version = "0.27.5" +version = "0.27.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d191583f3da1305256f22463b9bb0471acad48a4e534a5218b9963e9c1f59b2" +checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ - "futures-util", "http", "hyper", "hyper-util", @@ -1795,16 +2123,20 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.10" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df2dcfbe0677734ab2f3ffa7fa7bfd4706bfdc1ef393f2ee30184aed67e631b4" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ + "base64", "bytes", "futures-channel", "futures-util", "http", "http-body", "hyper", + "ipnet", + "libc", + "percent-encoding", "pin-project-lite", "socket2", "tokio", @@ -1814,14 +2146,15 @@ dependencies = [ [[package]] name = "iana-time-zone" -version = "0.1.61" +version = "0.1.65" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" dependencies = [ "android_system_properties", "core-foundation-sys", "iana-time-zone-haiku", "js-sys", + "log", "wasm-bindgen", "windows-core", ] @@ -1837,21 +2170,22 @@ dependencies = [ [[package]] name = "icu_collections" -version = "1.5.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" dependencies = [ "displaydoc", + "potential_utf", "yoke", "zerofrom", "zerovec", ] [[package]] -name = "icu_locid" -version = "1.5.0" +name = "icu_locale_core" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" dependencies = [ "displaydoc", "litemap", @@ -1860,104 +2194,78 @@ dependencies = [ "zerovec", ] -[[package]] -name = "icu_locid_transform" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_locid_transform_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_locid_transform_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" - [[package]] name = "icu_normalizer" -version = "1.5.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" dependencies = [ - "displaydoc", "icu_collections", "icu_normalizer_data", "icu_properties", "icu_provider", "smallvec", - "utf16_iter", - "utf8_iter", - "write16", "zerovec", ] [[package]] name = "icu_normalizer_data" -version = "1.5.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" [[package]] name = "icu_properties" -version = "1.5.1" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" dependencies = [ - "displaydoc", "icu_collections", - "icu_locid_transform", + "icu_locale_core", "icu_properties_data", "icu_provider", - "tinystr", + "zerotrie", "zerovec", ] [[package]] name = "icu_properties_data" -version = "1.5.0" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" [[package]] name = "icu_provider" -version = "1.5.0" +version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" dependencies = [ "displaydoc", - "icu_locid", - "icu_provider_macros", - "stable_deref_trait", - "tinystr", + "icu_locale_core", "writeable", "yoke", "zerofrom", + "zerotrie", "zerovec", ] [[package]] -name = "icu_provider_macros" -version = "1.5.0" +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "ident_case" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.98", -] +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" [[package]] name = "idna" -version = "1.0.3" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" dependencies = [ "idna_adapter", "smallvec", @@ -1966,9 +2274,9 @@ dependencies = [ [[package]] name = "idna_adapter" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" dependencies = [ "icu_normalizer", "icu_properties", @@ -1976,20 +2284,16 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.7.1" +version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" dependencies = [ "equivalent", - "hashbrown 0.15.2", + "hashbrown 0.16.1", + "serde", + "serde_core", ] -[[package]] -name = "indoc" -version = "2.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" - [[package]] name = "integer-encoding" version = "3.0.4" @@ -2003,12 +2307,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "469fb0b9cefa57e3ef31275ee7cacb78f2fdca44e4765491884a2b119d4eb130" [[package]] -name = "itertools" -version = "0.13.0" +name = "iri-string" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +checksum = "c91338f0783edbd6195decb37bae672fd3b165faffb89bf7b9e6942f8b1a731a" dependencies = [ - "either", + "memchr", + "serde", ] [[package]] @@ -2022,40 +2327,41 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.14" +version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" [[package]] name = "jobserver" -version = "0.1.32" +version = "0.1.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" dependencies = [ + "getrandom 0.3.4", "libc", ] [[package]] name = "js-sys" -version = "0.3.77" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" +checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" dependencies = [ "once_cell", "wasm-bindgen", ] [[package]] -name = "lazy_static" -version = "1.5.0" +name = "leb128fmt" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "lexical-core" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b765c31809609075565a70b4b71402281283aeda7ecaf4818ac14a7b2ade8958" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" dependencies = [ "lexical-parse-float", "lexical-parse-integer", @@ -2066,106 +2372,101 @@ dependencies = [ [[package]] name = "lexical-parse-float" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de6f9cb01fb0b08060209a057c048fcbab8717b4c1ecd2eac66ebfe39a65b0f2" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" dependencies = [ "lexical-parse-integer", "lexical-util", - "static_assertions", ] [[package]] name = "lexical-parse-integer" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72207aae22fc0a121ba7b6d479e42cbfea549af1479c3f3a4f12c70dd66df12e" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" dependencies = [ "lexical-util", - "static_assertions", ] [[package]] name = "lexical-util" -version = "1.0.6" +version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a82e24bf537fd24c177ffbbdc6ebcc8d54732c35b50a3f28cc3f4e4c949a0b3" -dependencies = [ - "static_assertions", -] +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" [[package]] name = "lexical-write-float" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5afc668a27f460fb45a81a757b6bf2f43c2d7e30cb5a2dcd3abf294c78d62bd" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" dependencies = [ "lexical-util", "lexical-write-integer", - "static_assertions", ] [[package]] name = "lexical-write-integer" -version = "1.0.5" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "629ddff1a914a836fb245616a7888b62903aae58fa771e1d83943035efa0f978" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" dependencies = [ "lexical-util", - "static_assertions", ] +[[package]] +name = "libbz2-rs-sys" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4a545a15244c7d945065b5d392b2d2d7f21526fba56ce51467b06ed445e8f7" + [[package]] name = "libc" -version = "0.2.169" +version = "0.2.182" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a" +checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" [[package]] -name = "libflate" -version = "2.1.0" +name = "libloading" +version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45d9dfdc14ea4ef0900c1cddbc8dcd553fbaacd8a4a282cf4018ae9dd04fb21e" +checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" dependencies = [ - "adler32", - "core2", - "crc32fast", - "dary_heap", - "libflate_lz77", + "cfg-if", + "winapi", ] [[package]] -name = "libflate_lz77" -version = "2.1.0" +name = "liblzma" +version = "0.4.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6e0d73b369f386f1c44abd9c570d5318f55ccde816ff4b562fa452e5182863d" +checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" dependencies = [ - "core2", - "hashbrown 0.14.5", - "rle-decode-fast", + "liblzma-sys", ] [[package]] -name = "libloading" -version = "0.7.4" +name = "liblzma-sys" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +checksum = "9f2db66f3268487b5033077f266da6777d057949b8f93c8ad82e441df25e6186" dependencies = [ - "cfg-if", - "winapi", + "cc", + "libc", + "pkg-config", ] [[package]] name = "libm" -version = "0.2.11" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.39" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23aa6811d3bd4deb8a84dde645f943476d13b248d818edcf8ce0b2f37f036b44" +checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" dependencies = [ "cc", "libc", @@ -2173,50 +2474,44 @@ dependencies = [ [[package]] name = "linux-raw-sys" -version = "0.4.15" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" +checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039" [[package]] name = "litemap" -version = "0.7.4" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104" +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" [[package]] name = "lock_api" -version = "0.4.12" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" dependencies = [ - "autocfg", "scopeguard", ] [[package]] name = "log" -version = "0.4.25" +version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" [[package]] -name = "lz4_flex" -version = "0.11.3" +name = "lru-slab" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75761162ae2b0e580d7e7c390558127e5f01b4194debd6221fd8c207fc80e3f5" -dependencies = [ - "twox-hash", -] +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] -name = "lzma-sys" -version = "0.1.20" +name = "lz4_flex" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +checksum = "ab6473172471198271ff72e9379150e9dfd70d8e533e0752a27e515b48dd375e" dependencies = [ - "cc", - "libc", - "pkg-config", + "twox-hash", ] [[package]] @@ -2231,73 +2526,45 @@ dependencies = [ [[package]] name = "memchr" -version = "2.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" - -[[package]] -name = "memoffset" -version = "0.9.1" +version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "mimalloc" -version = "0.1.43" +version = "0.1.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68914350ae34959d83f732418d51e2427a794055d0b9529f48259ac07af65633" +checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" dependencies = [ "libmimalloc-sys", ] -[[package]] -name = "mime" -version = "0.3.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" - [[package]] name = "miniz_oxide" -version = "0.8.4" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3b1c9bd4fe1f0f8b387f6eb9eb3b4a1aa26185e5750efb9140301703f62cd1b" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" dependencies = [ "adler2", + "simd-adler32", ] [[package]] name = "mio" -version = "1.0.3" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" +checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", - "wasi 0.11.0+wasi-snapshot-preview1", - "windows-sys 0.52.0", + "wasi", + "windows-sys 0.61.2", ] [[package]] name = "multimap" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "defc4c55412d89136f966bbb339008b474350e5e6e78d2714439c386b3137a03" - -[[package]] -name = "num" -version = "0.4.3" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" -dependencies = [ - "num-bigint", - "num-complex", - "num-integer", - "num-iter", - "num-rational", - "num-traits", -] +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" [[package]] name = "num-bigint" @@ -2328,28 +2595,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "num-iter" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" -dependencies = [ - "autocfg", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.19" @@ -2362,28 +2607,31 @@ dependencies = [ [[package]] name = "object" -version = "0.36.7" +version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" dependencies = [ "memchr", ] [[package]] name = "object_store" -version = "0.11.2" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cfccb68961a56facde1163f9319e0d15743352344e7808a11795fb99698dcaf" +checksum = "c2858065e55c148d294a9f3aae3b0fa9458edadb41a108397094566f4e3c0dfb" dependencies = [ "async-trait", - "base64 0.22.1", + "base64", "bytes", "chrono", + "form_urlencoded", "futures", + "http", + "http-body-util", "httparse", "humantime", "hyper", - "itertools 0.13.0", + "itertools", "md-5", "parking_lot", "percent-encoding", @@ -2391,27 +2639,30 @@ dependencies = [ "rand", "reqwest", "ring", - "rustls-pemfile", + "rustls-pki-types", "serde", "serde_json", - "snafu", + "serde_urlencoded", + "thiserror", "tokio", "tracing", "url", "walkdir", + "wasm-bindgen-futures", + "web-time", ] [[package]] name = "once_cell" -version = "1.20.3" +version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "openssl-probe" -version = "0.1.6" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" [[package]] name = "ordered-float" @@ -2424,9 +2675,9 @@ dependencies = [ [[package]] name = "parking_lot" -version = "0.12.3" +version = "0.12.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" dependencies = [ "lock_api", "parking_lot_core", @@ -2434,42 +2685,42 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.10" +version = "0.9.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-targets", + "windows-link", ] [[package]] name = "parquet" -version = "54.1.0" +version = "58.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a01a0efa30bbd601ae85b375c728efdb211ade54390281628a7b16708beb235" +checksum = "3f491d0ef1b510194426ee67ddc18a9b747ef3c42050c19322a2cd2e1666c29b" dependencies = [ "ahash", "arrow-array", "arrow-buffer", - "arrow-cast", "arrow-data", "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64", "brotli", "bytes", "chrono", "flate2", "futures", "half", - "hashbrown 0.15.2", + "hashbrown 0.16.1", "lz4_flex", - "num", "num-bigint", + "num-integer", + "num-traits", "object_store", "paste", "seq-macro", @@ -2479,16 +2730,6 @@ dependencies = [ "tokio", "twox-hash", "zstd", - "zstd-sys", -] - -[[package]] -name = "parse-zoneinfo" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f2a05b18d44e2957b88f96ba460715e295bc1d7510468a2f3d3b44535d26c24" -dependencies = [ - "regex", ] [[package]] @@ -2499,31 +2740,31 @@ checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" [[package]] name = "pbjson" -version = "0.7.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7e6349fa080353f4a597daffd05cb81572a9c031a6d4fff7e504947496fcc68" +checksum = "898bac3fa00d0ba57a4e8289837e965baa2dee8c3749f3b11d45a64b4223d9c3" dependencies = [ - "base64 0.21.7", + "base64", "serde", ] [[package]] name = "pbjson-build" -version = "0.7.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eea3058763d6e656105d1403cb04e0a41b7bbac6362d413e7c33be0c32279c9" +checksum = "af22d08a625a2213a78dbb0ffa253318c5c79ce3133d32d296655a7bdfb02095" dependencies = [ "heck", - "itertools 0.13.0", + "itertools", "prost", "prost-types", ] [[package]] name = "pbjson-types" -version = "0.7.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e54e5e7bfb1652f95bc361d76f3c780d8e526b134b85417e774166ee941f0887" +checksum = "8e748e28374f10a330ee3bb9f29b828c0ac79831a32bab65015ad9b661ead526" dependencies = [ "bytes", "chrono", @@ -2536,54 +2777,36 @@ dependencies = [ [[package]] name = "percent-encoding" -version = "2.3.1" +version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" [[package]] name = "petgraph" -version = "0.7.1" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", + "hashbrown 0.15.5", "indexmap", + "serde", ] [[package]] name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_codegen" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" -dependencies = [ - "phf_generator", - "phf_shared", -] - -[[package]] -name = "phf_generator" -version = "0.11.3" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" dependencies = [ "phf_shared", - "rand", ] [[package]] name = "phf_shared" -version = "0.11.3" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" dependencies = [ "siphasher", ] @@ -2602,49 +2825,58 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "pkg-config" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" [[package]] name = "portable-atomic" -version = "1.10.0" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "potential_utf" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6" +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +dependencies = [ + "zerovec", +] [[package]] name = "ppv-lite86" -version = "0.2.20" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ "zerocopy", ] [[package]] name = "prettyplease" -version = "0.2.29" +version = "0.2.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6924ced06e1f7dfe3fa48d57b9f74f55d8915f5036121bef647ef4b204895fac" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "proc-macro2" -version = "1.0.93" +version = "1.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" dependencies = [ "unicode-ident", ] [[package]] name = "prost" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" dependencies = [ "bytes", "prost-derive", @@ -2652,42 +2884,41 @@ dependencies = [ [[package]] name = "prost-build" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" +checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ "heck", - "itertools 0.14.0", + "itertools", "log", "multimap", - "once_cell", "petgraph", "prettyplease", "prost", "prost-types", "regex", - "syn 2.0.98", + "syn 2.0.117", "tempfile", ] [[package]] name = "prost-derive" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "prost-types" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" dependencies = [ "prost", ] @@ -2703,38 +2934,36 @@ dependencies = [ [[package]] name = "psm" -version = "0.1.25" +version = "0.1.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f58e5423e24c18cc840e1c98370b3993c6649cd1678b4d24318bcf0a083cbe88" +checksum = "3852766467df634d74f0b2d7819bf8dc483a0eb2e3b0f50f756f9cfe8b0d18d8" dependencies = [ + "ar_archive_writer", "cc", ] [[package]] name = "pyo3" -version = "0.23.4" +version = "0.28.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57fe09249128b3173d092de9523eaa75136bf7ba85e0d69eca241c7939c933cc" +checksum = "cf85e27e86080aafd5a22eae58a162e133a589551542b3e5cee4beb27e54f8e1" dependencies = [ - "cfg-if", - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-async-runtimes" -version = "0.23.0" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e" +checksum = "9e7364a95bf00e8377bbf9b0f09d7ff9715a29d8fcf93b47d1a967363b973178" dependencies = [ - "futures", + "futures-channel", + "futures-util", "once_cell", "pin-project-lite", "pyo3", @@ -2743,47 +2972,57 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.23.4" +version = "0.28.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cd3927b5a78757a0d71aa9dff669f903b1eb64b54142a9bd9f757f8fde65fd7" +checksum = "8bf94ee265674bf76c09fa430b0e99c26e319c945d96ca0d5a8215f31bf81cf7" dependencies = [ - "once_cell", "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.23.4" +version = "0.28.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dab6bb2102bd8f991e7749f130a70d05dd557613e39ed2deeee8e9ca0c4d548d" +checksum = "491aa5fc66d8059dd44a75f4580a2962c1862a1c2945359db36f6c2818b748dc" dependencies = [ "libc", "pyo3-build-config", ] +[[package]] +name = "pyo3-log" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c2ec80932c5c3b2d4fbc578c9b56b2d4502098587edb8bef5b6bfcad43682e" +dependencies = [ + "arc-swap", + "log", + "pyo3", +] + [[package]] name = "pyo3-macros" -version = "0.23.4" +version = "0.28.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91871864b353fd5ffcb3f91f2f703a22a9797c91b9ab497b1acac7b07ae509c7" +checksum = "f5d671734e9d7a43449f8480f8b38115df67bef8d21f76837fa75ee7aaa5e52e" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "pyo3-macros-backend" -version = "0.23.4" +version = "0.28.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43abc3b80bc20f3facd86cd3c60beed58c3e2aa26213f3cda368de39c60a27e4" +checksum = "22faaa1ce6c430a1f71658760497291065e6450d7b5dc2bcf254d49f66ee700a" dependencies = [ "heck", "proc-macro2", "pyo3-build-config", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -2794,9 +3033,9 @@ checksum = "5a651516ddc9168ebd67b24afd085a718be02f8858fe406591b013d101ce2f40" [[package]] name = "quick-xml" -version = "0.37.2" +version = "0.38.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "165859e9e55f79d67b96c5d96f4e88b6f2695a1972849c15a6a3f5c59fc2c003" +checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c" dependencies = [ "memchr", "serde", @@ -2804,37 +3043,40 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.6" +version = "0.11.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62e96808277ec6f97351a2380e6c25114bc9e67037775464979f3037c92d05ef" +checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" dependencies = [ "bytes", + "cfg_aliases", "pin-project-lite", "quinn-proto", "quinn-udp", "rustc-hash", "rustls", "socket2", - "thiserror 2.0.11", + "thiserror", "tokio", "tracing", + "web-time", ] [[package]] name = "quinn-proto" -version = "0.11.9" +version = "0.11.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2fe5ef3495d7d2e377ff17b1a8ce2ee2ec2a18cde8b6ad6619d65d0701c135d" +checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31" dependencies = [ "bytes", - "getrandom 0.2.15", + "getrandom 0.3.4", + "lru-slab", "rand", "ring", "rustc-hash", "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.11", + "thiserror", "tinyvec", "tracing", "web-time", @@ -2842,43 +3084,48 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.10" +version = "0.5.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e46f3055866785f6b92bc6164b76be02ca8f2eb4b002c0354b28cf4c119e5944" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" dependencies = [ "cfg_aliases", "libc", "once_cell", "socket2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] name = "quote" -version = "1.0.38" +version = "1.0.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc" +checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "rand" -version = "0.8.5" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ - "libc", "rand_chacha", "rand_core", ] [[package]] name = "rand_chacha" -version = "0.3.1" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", "rand_core", @@ -2886,11 +3133,11 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.6.4" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" dependencies = [ - "getrandom 0.2.15", + "getrandom 0.3.4", ] [[package]] @@ -2910,23 +3157,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "redox_syscall" -version = "0.5.8" +version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03a862b389f93e68874fbf580b9de08dd02facb9a788ebadaf4a3fd33cf58834" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags 2.8.0", + "bitflags", ] [[package]] name = "regex" -version = "1.11.1" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", @@ -2936,9 +3183,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.9" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" dependencies = [ "aho-corasick", "memchr", @@ -2947,23 +3194,23 @@ dependencies = [ [[package]] name = "regex-lite" -version = "0.1.6" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.5" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" [[package]] name = "regress" -version = "0.10.3" +version = "0.10.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ef7fa9ed0256d64a688a3747d0fef7a88851c18a5e1d57f115f38ec2e09366" +checksum = "2057b2325e68a893284d1538021ab90279adac1139957ca2a74426c6f118fb48" dependencies = [ - "hashbrown 0.15.2", + "hashbrown 0.16.1", "memchr", ] @@ -2978,11 +3225,11 @@ dependencies = [ [[package]] name = "reqwest" -version = "0.12.12" +version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43e734407157c3c2034e0258f5e4473ddb361b1e85f95a66690d67264d7cd1da" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "futures-core", "futures-util", @@ -2993,17 +3240,13 @@ dependencies = [ "hyper", "hyper-rustls", "hyper-util", - "ipnet", "js-sys", "log", - "mime", - "once_cell", "percent-encoding", "pin-project-lite", "quinn", "rustls", "rustls-native-certs", - "rustls-pemfile", "rustls-pki-types", "serde", "serde_json", @@ -3013,41 +3256,29 @@ dependencies = [ "tokio-rustls", "tokio-util", "tower", + "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", "wasm-streams", "web-sys", - "windows-registry", ] [[package]] name = "ring" -version = "0.17.9" +version = "0.17.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e75ec5e92c4d8aede845126adc388046234541629e76029599ed35a003c7ed24" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" dependencies = [ "cc", "cfg-if", - "getrandom 0.2.15", + "getrandom 0.2.17", "libc", "untrusted", "windows-sys 0.52.0", ] -[[package]] -name = "rle-decode-fast" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3582f63211428f83597b51b2ddb88e2a91a9d52d12831f9d08f5e624e8977422" - -[[package]] -name = "rustc-demangle" -version = "0.1.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" - [[package]] name = "rustc-hash" version = "2.1.1" @@ -3065,22 +3296,22 @@ dependencies = [ [[package]] name = "rustix" -version = "0.38.44" +version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +checksum = "146c9e247ccc180c1f61615433868c99f3de3ae256a30a43b49f67c2d9171f34" dependencies = [ - "bitflags 2.8.0", + "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] name = "rustls" -version = "0.23.23" +version = "0.23.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47796c98c480fce5406ef69d1c76378375492c3b0a0de587be0c1d9feb12f395" +checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b" dependencies = [ "once_cell", "ring", @@ -3092,9 +3323,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.1" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fcff2dd52b58a8d98a70243663a0d234c4e2b79235637849d15913394a247d3" +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -3102,29 +3333,21 @@ dependencies = [ "security-framework", ] -[[package]] -name = "rustls-pemfile" -version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "rustls-pki-types" -version = "1.11.0" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "917ce264624a4b4db1c364dcc35bfca9ded014d0a958cd47ad3e960e988ea51c" +checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" dependencies = [ "web-time", + "zeroize", ] [[package]] name = "rustls-webpki" -version = "0.102.8" +version = "0.103.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53" dependencies = [ "ring", "rustls-pki-types", @@ -3133,15 +3356,15 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.19" +version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7c45b9784283f1b2e7fb61b42047c2fd678ef0960d4f6f1eba131594cc369d4" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "ryu" -version = "1.0.19" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "same-file" @@ -3154,18 +3377,18 @@ dependencies = [ [[package]] name = "schannel" -version = "0.1.27" +version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f29ebaa345f945cec9fbbc532eb307f0fdad8161f281b6369539c8d84876b3d" +checksum = "891d81b926048e76efe18581bf793546b4c0eaf8448d72be8de2bbee5fd166e1" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] name = "schemars" -version = "0.8.21" +version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09c024468a378b7e36765cd36702b7a90cc3cba11654f6685c8f233408e89e92" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" dependencies = [ "dyn-clone", "schemars_derive", @@ -3175,14 +3398,14 @@ dependencies = [ [[package]] name = "schemars_derive" -version = "0.8.21" +version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1eee588578aff73f856ab961cd2f79e36bc45d7ded33a7562adba4667aecc0e" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" dependencies = [ "proc-macro2", "quote", "serde_derive_internals", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -3193,11 +3416,11 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "security-framework" -version = "3.2.0" +version = "3.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "271720403f46ca04f7ba6f55d438f8bd878d6b8ca0a1046e8228c4145bcbb316" +checksum = "d17b898a6d6948c3a8ee4372c17cb384f90d2e6e912ef00895b14fd7ab54ec38" dependencies = [ - "bitflags 2.8.0", + "bitflags", "core-foundation", "core-foundation-sys", "libc", @@ -3206,9 +3429,9 @@ dependencies = [ [[package]] name = "security-framework-sys" -version = "2.14.0" +version = "2.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49db231d56a190491cb4aeda9527f1ad45345af50b0851622a7adb8c03b01c32" +checksum = "321c8673b092a9a42605034a9879d73cb79101ed5fd117bc9a597b89b4e9e61a" dependencies = [ "core-foundation-sys", "libc", @@ -3216,46 +3439,58 @@ dependencies = [ [[package]] name = "semver" -version = "1.0.25" +version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f79dfe2d285b0488816f30e700a7438c5a73d816b5b7d3ac72fbc48b0d185e03" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" dependencies = [ "serde", + "serde_core", ] [[package]] name = "seq-macro" -version = "0.3.5" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.217" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ + "serde_core", "serde_derive", ] [[package]] name = "serde_bytes" -version = "0.11.15" +version = "0.11.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "387cc504cb06bb40a96c8e04e951fe01854cf6bc921053c954e4a606d9675c6a" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" dependencies = [ "serde", + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.217" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -3266,19 +3501,20 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "serde_json" -version = "1.0.138" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d434192e7da787e94a6ea7e9670b26a036d0ca41e0b7efb2676dd32bae872949" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ "itoa", "memchr", - "ryu", "serde", + "serde_core", + "zmij", ] [[package]] @@ -3290,7 +3526,7 @@ dependencies = [ "proc-macro2", "quote", "serde", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -3320,9 +3556,9 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.8" +version = "0.10.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures", @@ -3335,6 +3571,12 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + [[package]] name = "simdutf8" version = "0.1.5" @@ -3343,45 +3585,21 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] name = "siphasher" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56199f7ddabf13fe5074ce809e7d3f42b42ae711800501b5b16ea82ad029c39d" +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" [[package]] name = "slab" -version = "0.4.9" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" -dependencies = [ - "autocfg", -] +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "smallvec" -version = "1.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" - -[[package]] -name = "snafu" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "223891c85e2a29c3fe8fb900c1fae5e69c2e42415e3177752e8718475efa5019" -dependencies = [ - "snafu-derive", -] - -[[package]] -name = "snafu-derive" -version = "0.8.5" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03c3c6b7927ffe7ecaa769ee0e3994da3b8cafc8f444578982c83ecb161af917" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.98", -] +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "snap" @@ -3391,46 +3609,47 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" [[package]] name = "socket2" -version = "0.5.8" +version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c970269d99b64e60ec3bd6ad27270092a5394c4e309314b18ae3fe575695fbe8" +checksum = "86f4aa3ad99f2088c990dfa82d367e19cb29268ed67c574d10d0a4bfe71f07e0" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] name = "sqlparser" -version = "0.53.0" +version = "0.61.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05a528114c392209b3264855ad491fcce534b94a38771b0a0b97a79379275ce8" +checksum = "dbf5ea8d4d7c808e1af1cbabebca9a2abe603bcefc22294c5b95018d53200cb7" dependencies = [ "log", + "recursive", "sqlparser_derive", ] [[package]] name = "sqlparser_derive" -version = "0.3.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da5fc6819faabb412da764b99d3b713bb55083c11e7e0c00144d386cd6a1939c" +checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "stable_deref_trait" -version = "1.2.0" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stacker" -version = "0.1.18" +version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d08feb8f695b465baed819b03c128dc23f57a694510ab1f06c77f763975685e" +checksum = "08d74a23609d509411d10e2176dc2a4346e3b4aea2e7b1869f19fdedbc71c013" dependencies = [ "cc", "cfg-if", @@ -3440,35 +3659,34 @@ dependencies = [ ] [[package]] -name = "static_assertions" -version = "1.1.0" +name = "strsim" +version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "strum" -version = "0.26.3" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" [[package]] name = "strum_macros" -version = "0.26.4" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" dependencies = [ "heck", "proc-macro2", "quote", - "rustversion", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "substrait" -version = "0.52.3" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5db15789cecbfdf6b1fcf2db807e767c92273bdc407ac057c2194b070c597756" +checksum = "62fc4b483a129b9772ccb9c3f7945a472112fdd9140da87f8a4e7f1d44e045d0" dependencies = [ "heck", "pbjson", @@ -3485,7 +3703,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "syn 2.0.98", + "syn 2.0.117", "typify", "walkdir", ] @@ -3509,9 +3727,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.98" +version = "2.0.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" dependencies = [ "proc-macro2", "quote", @@ -3529,73 +3747,52 @@ dependencies = [ [[package]] name = "synstructure" -version = "0.13.1" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "tempfile" -version = "3.16.0" +version = "3.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38c246215d7d24f48ae091a2902398798e05d978b24315d6efbc00ede9a8bb91" +checksum = "0136791f7c95b1f6dd99f9cc786b91bb81c3800b639b3478e561ddb7be95e5f1" dependencies = [ - "cfg-if", "fastrand", - "getrandom 0.3.1", + "getrandom 0.4.1", "once_cell", "rustix", - "windows-sys 0.59.0", -] - -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", + "windows-sys 0.61.2", ] [[package]] name = "thiserror" -version = "2.0.11" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl 2.0.11", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.98", + "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.11" +version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] @@ -3620,9 +3817,9 @@ dependencies = [ [[package]] name = "tinystr" -version = "0.7.6" +version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" dependencies = [ "displaydoc", "zerovec", @@ -3630,9 +3827,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.8.1" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "022db8904dfa342efe721985167e9fcd16c29b226db4397ed752a761cfce81e8" +checksum = "bfa5fdc3bce6191a1dbc8c02d5c8bffcf557bafa17c124c5264a458f1b0613fa" dependencies = [ "tinyvec_macros", ] @@ -3645,46 +3842,57 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.43.0" +version = "1.49.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d61fa4ffa3de412bfea335c6ecff681de2b609ba3c77ef3e00e521813a9ed9e" +checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" dependencies = [ - "backtrace", "bytes", "libc", "mio", "pin-project-lite", "socket2", "tokio-macros", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] name = "tokio-macros" -version = "2.5.0" +version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" +checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "tokio-rustls" -version = "0.26.1" +version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f6d0975eaace0cf0fcadee4e4aaa5da15b5c079146f2cffb67c113be122bf37" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ "rustls", "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-util" -version = "0.7.13" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7fcaa8d55a2bdd6b83ace262b016eca0d79ee02818c5c1bcdf0305114081078" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", @@ -3695,9 +3903,9 @@ dependencies = [ [[package]] name = "tower" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", @@ -3708,6 +3916,24 @@ dependencies = [ "tower-service", ] +[[package]] +name = "tower-http" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "iri-string", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", +] + [[package]] name = "tower-layer" version = "0.3.3" @@ -3722,9 +3948,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracing" -version = "0.1.41" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", "tracing-attributes", @@ -3733,20 +3959,20 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.28" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "tracing-core" -version = "0.1.33" +version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", ] @@ -3774,13 +4000,9 @@ checksum = "e78122066b0cb818b8afd08f7ed22f7fdbc3e90815035726f0840d0d26c0747a" [[package]] name = "twox-hash" -version = "1.6.3" +version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" -dependencies = [ - "cfg-if", - "static_assertions", -] +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" [[package]] name = "typed-arena" @@ -3789,36 +4011,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" [[package]] -name = "typed-builder" -version = "0.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06fbd5b8de54c5f7c91f6fe4cebb949be2125d7758e630bb58b1d831dbce600" -dependencies = [ - "typed-builder-macro", -] - -[[package]] -name = "typed-builder-macro" -version = "0.19.1" +name = "typenum" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9534daa9fd3ed0bd911d462a37f172228077e7abf18c18a5f67199d959205f8" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.98", -] +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" [[package]] -name = "typenum" -version = "1.17.0" +name = "typewit" +version = "1.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +checksum = "f8c1ae7cc0fdb8b842d65d127cb981574b0d2b249b74d1c7a2986863dc134f71" [[package]] name = "typify" -version = "0.3.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e03ba3643450cfd95a1aca2e1938fef63c1c1994489337998aff4ad771f21ef8" +checksum = "e6d5bcc6f62eb1fa8aa4098f39b29f93dcb914e17158b76c50360911257aa629" dependencies = [ "typify-impl", "typify-macro", @@ -3826,9 +4034,9 @@ dependencies = [ [[package]] name = "typify-impl" -version = "0.3.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bce48219a2f3154aaa2c56cbf027728b24a3c8fe0a47ed6399781de2b3f3eeaf" +checksum = "a1eb359f7ffa4f9ebe947fa11a1b2da054564502968db5f317b7e37693cb2240" dependencies = [ "heck", "log", @@ -3839,16 +4047,16 @@ dependencies = [ "semver", "serde", "serde_json", - "syn 2.0.98", - "thiserror 2.0.11", + "syn 2.0.117", + "thiserror", "unicode-ident", ] [[package]] name = "typify-macro" -version = "0.3.0" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68b5780d745920ed73c5b7447496a9b5c42ed2681a9b70859377aec423ecf02b" +checksum = "911c32f3c8514b048c1b228361bebb5e6d73aeec01696e8cc0e82e2ffef8ab7a" dependencies = [ "proc-macro2", "quote", @@ -3857,15 +4065,15 @@ dependencies = [ "serde", "serde_json", "serde_tokenstream", - "syn 2.0.98", + "syn 2.0.117", "typify-impl", ] [[package]] name = "unicode-ident" -version = "1.0.16" +version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a210d160f08b701c8721ba1c726c11662f877ea6b7094007e1ca9a1041945034" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" [[package]] name = "unicode-segmentation" @@ -3875,15 +4083,15 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" [[package]] name = "unicode-width" -version = "0.2.0" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" [[package]] -name = "unindent" -version = "0.2.3" +name = "unicode-xid" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" [[package]] name = "unsafe-libyaml" @@ -3899,21 +4107,16 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" [[package]] name = "url" -version = "2.5.4" +version = "2.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" dependencies = [ "form_urlencoded", "idna", "percent-encoding", + "serde", ] -[[package]] -name = "utf16_iter" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" - [[package]] name = "utf8_iter" version = "1.0.4" @@ -3922,12 +4125,14 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.13.1" +version = "1.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced87ca4be083373936a67f8de945faa23b6b42384bd5b64434850802c6dccd0" +checksum = "b672338555252d43fd2240c714dc444b8c6fb0a5c5335e65a07bba7742735ddb" dependencies = [ - "getrandom 0.3.1", - "serde", + "getrandom 0.4.1", + "js-sys", + "serde_core", + "wasm-bindgen", ] [[package]] @@ -3957,52 +4162,49 @@ dependencies = [ [[package]] name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] -name = "wasi" -version = "0.13.3+wasi-0.2.2" +name = "wasip2" +version = "1.0.2+wasi-0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" dependencies = [ - "wit-bindgen-rt", + "wit-bindgen", ] [[package]] -name = "wasm-bindgen" -version = "0.2.100" +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" dependencies = [ - "cfg-if", - "once_cell", - "rustversion", - "wasm-bindgen-macro", + "wit-bindgen", ] [[package]] -name = "wasm-bindgen-backend" -version = "0.2.100" +name = "wasm-bindgen" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" +checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" dependencies = [ - "bumpalo", - "log", - "proc-macro2", - "quote", - "syn 2.0.98", + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-futures" -version = "0.4.50" +version = "0.4.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "555d470ec0bc3bb57890405e5d4322cc9ea83cebb085523ced7be4144dac1e61" +checksum = "70a6e77fd0ae8029c9ea0063f87c46fde723e7d887703d74ad2616d792e51e6f" dependencies = [ "cfg-if", + "futures-util", "js-sys", "once_cell", "wasm-bindgen", @@ -4011,9 +4213,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.100" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" +checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4021,26 +4223,48 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.100" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" +checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" dependencies = [ + "bumpalo", "proc-macro2", "quote", - "syn 2.0.98", - "wasm-bindgen-backend", + "syn 2.0.117", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.100" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + [[package]] name = "wasm-streams" version = "0.4.2" @@ -4054,11 +4278,23 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + [[package]] name = "web-sys" -version = "0.3.77" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33b6dd2ef9186f1f2072e409e99cd22a975331a6b3591b12c764e0e55c60d5d2" +checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" dependencies = [ "js-sys", "wasm-bindgen", @@ -4092,11 +4328,11 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.9" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4107,41 +4343,61 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "windows-core" -version = "0.52.0" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ - "windows-targets", + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", ] [[package]] -name = "windows-registry" -version = "0.2.0" +name = "windows-implement" +version = "0.60.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e400001bb720a623c1c69032f8e3e4cf09984deec740f007dd2b03ec864804b0" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ - "windows-result", - "windows-strings", - "windows-targets", + "proc-macro2", + "quote", + "syn 2.0.117", ] +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-result" -version = "0.2.0" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d1043d8214f791817bab27572aaa8af63732e11bf84aa21a45a78d6c317ae0e" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" dependencies = [ - "windows-targets", + "windows-link", ] [[package]] name = "windows-strings" -version = "0.1.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cd9b125c486025df0eabcb585e62173c6c9eddcec5d117d3b6e8c30e2ee4d10" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" dependencies = [ - "windows-result", - "windows-targets", + "windows-link", ] [[package]] @@ -4150,7 +4406,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -4159,7 +4415,25 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", ] [[package]] @@ -4168,14 +4442,31 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", ] [[package]] @@ -4184,42 +4475,84 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" @@ -4227,42 +4560,111 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] -name = "wit-bindgen-rt" -version = "0.33.0" +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "wit-bindgen" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" dependencies = [ - "bitflags 2.8.0", + "wit-bindgen-rust-macro", ] [[package]] -name = "write16" -version = "1.0.0" +name = "wit-bindgen-core" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] [[package]] -name = "writeable" -version = "0.5.5" +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn 2.0.117", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.117", + "wit-bindgen-core", + "wit-bindgen-rust", +] [[package]] -name = "xz2" -version = "0.1.7" +name = "wit-component" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ - "lzma-sys", + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", ] [[package]] -name = "yoke" -version = "0.7.5" +name = "wit-parser" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "writeable" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" + +[[package]] +name = "yoke" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +dependencies = [ "stable_deref_trait", "yoke-derive", "zerofrom", @@ -4270,69 +4672,79 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.7.5" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", "synstructure", ] [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "db6d35d663eadb6c932438e763b262fe1a70987f9ae936e60158176d710cae4a" dependencies = [ - "byteorder", "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.35" +version = "0.8.39" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +checksum = "4122cd3169e94605190e77839c9a40d40ed048d305bfdc146e7df40ab0f3e517" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] [[package]] name = "zerofrom" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" dependencies = [ "zerofrom-derive", ] [[package]] name = "zerofrom-derive" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", "synstructure", ] [[package]] name = "zeroize" -version = "1.8.1" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + +[[package]] +name = "zerotrie" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde" +checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] [[package]] name = "zerovec" -version = "0.10.4" +version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" +checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" dependencies = [ "yoke", "zerofrom", @@ -4341,38 +4753,50 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.10.3" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" +checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.117", ] +[[package]] +name = "zlib-rs" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c745c48e1007337ed136dc99df34128b9faa6ed542d80a1c673cf55a6d7236c8" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + [[package]] name = "zstd" -version = "0.13.2" +version = "0.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fcf2b778a664581e31e389454a7072dab1647606d44f7feea22cd5abb9c9f3f9" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" dependencies = [ "zstd-safe", ] [[package]] name = "zstd-safe" -version = "7.2.1" +version = "7.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54a3ab4db68cea366acc5c897c7b4d4d1b8994a9cd6e6f841f8964566a419059" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" dependencies = [ "zstd-sys", ] [[package]] name = "zstd-sys" -version = "2.0.13+zstd.1.5.6" +version = "2.0.16+zstd.1.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38ff0f21cfee8f97d94cef41359e0c89aa6113028ab0291aa8ca0038995a95aa" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" dependencies = [ "cc", "pkg-config", diff --git a/Cargo.toml b/Cargo.toml index 50967a219..19b79daf8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,49 +15,67 @@ # specific language governing permissions and limitations # under the License. -[package] -name = "datafusion-python" -version = "45.2.0" +[workspace.package] +version = "52.0.0" homepage = "https://datafusion.apache.org/python" repository = "https://github.com/apache/datafusion-python" authors = ["Apache DataFusion "] description = "Apache DataFusion DataFrame and SQL Query Engine" readme = "README.md" license = "Apache-2.0" -edition = "2021" -rust-version = "1.78" -include = ["/src", "/datafusion", "/LICENSE.txt", "pyproject.toml", "Cargo.toml", "Cargo.lock"] +edition = "2024" +rust-version = "1.88" -[features] -default = ["mimalloc"] -protoc = [ "datafusion-substrait/protoc" ] -substrait = ["dep:datafusion-substrait"] +[workspace] +members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"] +resolver = "3" -[dependencies] -tokio = { version = "1.42", features = ["macros", "rt", "rt-multi-thread", "sync"] } -pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"] } -pyo3-async-runtimes = { version = "0.23", features = ["tokio-runtime"]} -arrow = { version = "54", features = ["pyarrow"] } -datafusion = { version = "45.0.0", features = ["avro", "unicode_expressions"] } -datafusion-substrait = { version = "45.0.0", optional = true } -datafusion-proto = { version = "45.0.0" } -datafusion-ffi = { version = "45.0.0" } -prost = "0.13" # keep in line with `datafusion-substrait` -uuid = { version = "1.12", features = ["v4"] } -mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] } -async-trait = "0.1" +[workspace.dependencies] +tokio = { version = "1.49" } +pyo3 = { version = "0.28" } +pyo3-async-runtimes = { version = "0.28" } +pyo3-log = "0.13.3" +arrow = { version = "58" } +arrow-array = { version = "58" } +arrow-schema = { version = "58" } +arrow-select = { version = "58" } +datafusion = { version = "53" } +datafusion-substrait = { version = "53" } +datafusion-proto = { version = "53" } +datafusion-ffi = { version = "53" } +datafusion-catalog = { version = "53", default-features = false } +datafusion-common = { version = "53", default-features = false } +datafusion-functions-aggregate = { version = "53" } +datafusion-functions-window = { version = "53" } +datafusion-expr = { version = "53" } +prost = "0.14.3" +serde_json = "1" +uuid = { version = "1.21" } +mimalloc = { version = "0.1", default-features = false } +async-trait = "0.1.89" futures = "0.3" -object_store = { version = "0.11.0", features = ["aws", "gcp", "azure", "http"] } +cstr = "0.2" +object_store = { version = "0.13.1" } url = "2" - -[build-dependencies] -prost-types = "0.13" # keep in line with `datafusion-substrait` -pyo3-build-config = "0.23" - -[lib] -name = "datafusion_python" -crate-type = ["cdylib", "rlib"] +log = "0.4.29" +parking_lot = "0.12" +prost-types = "0.14.3" # keep in line with `datafusion-substrait` +pyo3-build-config = "0.28" +datafusion-python-util = { path = "crates/util" } [profile.release] lto = true codegen-units = 1 + +# We cannot publish to crates.io with any patches in the below section. Developers +# must remove any entries in this section before creating a release candidate. +[patch.crates-io] +datafusion = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-substrait = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-proto = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-ffi = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-catalog = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-common = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-functions-window = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } +datafusion-expr = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f", submodules = false } diff --git a/README.md b/README.md index 4f80dbe18..c24257876 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,10 @@ DataFusion's Python bindings can be used as a foundation for building new data s - Serialize and deserialize query plans in Substrait format. - Experimental support for transpiling SQL queries to DataFrame calls with Polars, Pandas, and cuDF. +For tips on tuning parallelism, see +[Maximizing CPU Usage](docs/source/user-guide/configuration.rst#maximizing-cpu-usage) +in the configuration guide. + ## Example Usage The following example demonstrates running a SQL query against a Parquet file using DataFusion, storing the results @@ -216,6 +220,8 @@ You can verify the installation by running: This assumes that you have rust and cargo installed. We use the workflow recommended by [pyo3](https://github.com/PyO3/pyo3) and [maturin](https://github.com/PyO3/maturin). The Maturin tools used in this workflow can be installed either via `uv` or `pip`. Both approaches should offer the same experience. It is recommended to use `uv` since it has significant performance improvements over `pip`. +Currently for protobuf support either [protobuf](https://protobuf.dev/installation/) or cmake must be installed. + Bootstrap (`uv`): By default `uv` will attempt to build the datafusion python package. For our development we prefer to build manually. This means @@ -225,7 +231,9 @@ and for `uv run` commands the additional parameter `--no-project` ```bash # fetch this repo git clone git@github.com:apache/datafusion-python.git -# create the virtual enviornment +# cd to the repo root +cd datafusion-python/ +# create the virtual environment uv sync --dev --no-install-package datafusion # activate the environment source .venv/bin/activate @@ -236,6 +244,8 @@ Bootstrap (`pip`): ```bash # fetch this repo git clone git@github.com:apache/datafusion-python.git +# cd to the repo root +cd datafusion-python/ # prepare development environment (used to build wheel / install in development) python3 -m venv .venv # activate the venv @@ -265,7 +275,16 @@ needing to activate the virtual environment: ```bash uv run --no-project maturin develop --uv -uv --no-project pytest . +uv run --no-project pytest . +``` + +To run the FFI tests within the examples folder, after you have built +`datafusion-python` with the previous commands: + +```bash +cd examples/datafusion-ffi-example +uv run --no-project maturin develop --uv +uv run --no-project pytest python/tests/_test_*py ``` ### Running & Installing pre-commit hooks @@ -278,7 +297,9 @@ Our pre-commit hooks can be installed by running `pre-commit install`, which wil your DATAFUSION_PYTHON_ROOT/.github directory and run each time you perform a commit, failing to complete the commit if an offending lint is found allowing you to make changes locally before pushing. -The pre-commit hooks can also be run adhoc without installing them by simply running `pre-commit run --all-files` +The pre-commit hooks can also be run adhoc without installing them by simply running `pre-commit run --all-files`. + +NOTE: the current `pre-commit` hooks require docker, and cmake. See note on protobuf above. ## Running linters without using pre-commit diff --git a/benchmarks/db-benchmark/groupby-datafusion.py b/benchmarks/db-benchmark/groupby-datafusion.py index f9e8d638b..533166695 100644 --- a/benchmarks/db-benchmark/groupby-datafusion.py +++ b/benchmarks/db-benchmark/groupby-datafusion.py @@ -18,6 +18,7 @@ import gc import os import timeit +from pathlib import Path import datafusion as df import pyarrow as pa @@ -34,7 +35,7 @@ print("# groupby-datafusion.py", flush=True) -exec(open("./_helpers/helpers.py").read()) +exec(Path.open("./_helpers/helpers.py").read()) def ans_shape(batches) -> tuple[int, int]: @@ -65,7 +66,7 @@ def execute(df) -> list: sql = True data_name = os.environ["SRC_DATANAME"] -src_grp = os.path.join("data", data_name + ".csv") +src_grp = "data" / data_name / ".csv" print("loading dataset %s" % src_grp, flush=True) schema = pa.schema( diff --git a/benchmarks/db-benchmark/join-datafusion.py b/benchmarks/db-benchmark/join-datafusion.py index 039868031..3be296c81 100755 --- a/benchmarks/db-benchmark/join-datafusion.py +++ b/benchmarks/db-benchmark/join-datafusion.py @@ -18,6 +18,7 @@ import gc import os import timeit +from pathlib import Path import datafusion as df from datafusion import col @@ -26,7 +27,7 @@ print("# join-datafusion.py", flush=True) -exec(open("./_helpers/helpers.py").read()) +exec(Path.open("./_helpers/helpers.py").read()) def ans_shape(batches) -> tuple[int, int]: @@ -49,12 +50,12 @@ def ans_shape(batches) -> tuple[int, int]: on_disk = "FALSE" data_name = os.environ["SRC_DATANAME"] -src_jn_x = os.path.join("data", data_name + ".csv") +src_jn_x = "data" / data_name / ".csv" y_data_name = join_to_tbls(data_name) src_jn_y = [ - os.path.join("data", y_data_name[0] + ".csv"), - os.path.join("data", y_data_name[1] + ".csv"), - os.path.join("data", y_data_name[2] + ".csv"), + "data" / y_data_name[0] / ".csv", + "data" / y_data_name[1] / ".csv", + "data" / y_data_name[2] / ".csv", ] if len(src_jn_y) != 3: error_msg = "Something went wrong in preparing files used for join" diff --git a/benchmarks/max_cpu_usage.py b/benchmarks/max_cpu_usage.py new file mode 100644 index 000000000..ae73baad6 --- /dev/null +++ b/benchmarks/max_cpu_usage.py @@ -0,0 +1,107 @@ +# 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. +"""Benchmark script showing how to maximize CPU usage. + +This script demonstrates one example of tuning DataFusion for improved parallelism +and CPU utilization. It uses synthetic in-memory data and performs simple aggregation +operations to showcase the impact of partitioning configuration. + +IMPORTANT: This is a simplified example designed to illustrate partitioning concepts. +Actual performance in your applications may vary significantly based on many factors: + +- Type of table providers (Parquet files, CSV, databases, etc.) +- I/O operations and storage characteristics (local disk, network, cloud storage) +- Query complexity and operation types (joins, window functions, complex expressions) +- Data distribution and size characteristics +- Memory available and hardware specifications +- Network latency for distributed data sources + +It is strongly recommended that you create similar benchmarks tailored to your specific: +- Hardware configuration +- Data sources and formats +- Typical query patterns and workloads +- Performance requirements + +This will give you more accurate insights into how DataFusion configuration options +will affect your particular use case. +""" + +from __future__ import annotations + +import argparse +import multiprocessing +import time + +import pyarrow as pa +from datafusion import SessionConfig, SessionContext, col +from datafusion import functions as f + + +def main(num_rows: int, partitions: int) -> None: + """Run a simple aggregation after repartitioning. + + This function demonstrates basic partitioning concepts using synthetic data. + Real-world performance will depend on your specific data sources, query types, + and system configuration. + """ + # Create some example data (synthetic in-memory data for demonstration) + # Note: Real applications typically work with files, databases, or other + # data sources that have different I/O and distribution characteristics + array = pa.array(range(num_rows)) + batch = pa.record_batch([array], names=["a"]) + + # Configure the session to use a higher target partition count and + # enable automatic repartitioning. + config = ( + SessionConfig() + .with_target_partitions(partitions) + .with_repartition_joins(enabled=True) + .with_repartition_aggregations(enabled=True) + .with_repartition_windows(enabled=True) + ) + ctx = SessionContext(config) + + # Register the input data and repartition manually to ensure that all + # partitions are used. + df = ctx.create_dataframe([[batch]]).repartition(partitions) + + start = time.time() + df = df.aggregate([], [f.sum(col("a"))]) + df.collect() + end = time.time() + + print( + f"Processed {num_rows} rows using {partitions} partitions in {end - start:.3f}s" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--rows", + type=int, + default=1_000_000, + help="Number of rows in the generated dataset", + ) + parser.add_argument( + "--partitions", + type=int, + default=multiprocessing.cpu_count(), + help="Target number of partitions to use", + ) + args = parser.parse_args() + main(args.rows, args.partitions) diff --git a/benchmarks/tpch/tpch.py b/benchmarks/tpch/tpch.py index 2d1bbae5b..ffee5554c 100644 --- a/benchmarks/tpch/tpch.py +++ b/benchmarks/tpch/tpch.py @@ -17,12 +17,13 @@ import argparse import time +from pathlib import Path from datafusion import SessionContext def bench(data_path, query_path) -> None: - with open("results.csv", "w") as results: + with Path("results.csv").open("w") as results: # register tables start = time.time() total_time_millis = 0 @@ -45,7 +46,7 @@ def bench(data_path, query_path) -> None: print("Configuration:\n", ctx) # register tables - with open("create_tables.sql") as f: + with Path("create_tables.sql").open() as f: sql = "" for line in f.readlines(): if line.startswith("--"): @@ -65,7 +66,7 @@ def bench(data_path, query_path) -> None: # run queries for query in range(1, 23): - with open(f"{query_path}/q{query}.sql") as f: + with Path(f"{query_path}/q{query}.sql").open() as f: text = f.read() tmp = text.split(";") queries = [s.strip() for s in tmp if len(s.strip()) > 0] diff --git a/ci/scripts/rust_fmt.sh b/ci/scripts/rust_fmt.sh index 9d8325877..05cb6b208 100755 --- a/ci/scripts/rust_fmt.sh +++ b/ci/scripts/rust_fmt.sh @@ -18,4 +18,4 @@ # under the License. set -ex -cargo fmt --all -- --check +cargo +nightly fmt --all -- --check diff --git a/conftest.py b/conftest.py new file mode 100644 index 000000000..1c89f92bc --- /dev/null +++ b/conftest.py @@ -0,0 +1,29 @@ +# 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. + +"""Pytest configuration for doctest namespace injection.""" + +import datafusion as dfn +import numpy as np +import pytest + + +@pytest.fixture(autouse=True) +def _doctest_namespace(doctest_namespace: dict) -> None: + """Add common imports to the doctest namespace.""" + doctest_namespace["dfn"] = dfn + doctest_namespace["np"] = np diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml new file mode 100644 index 000000000..3e2b01c8e --- /dev/null +++ b/crates/core/Cargo.toml @@ -0,0 +1,82 @@ +# 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. + +[package] +name = "datafusion-python" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true +include = [ + "src", + "../LICENSE.txt", + "build.rs", + "../pyproject.toml", + "Cargo.toml", + "../Cargo.lock", +] + +[dependencies] +tokio = { workspace = true, features = [ + "macros", + "rt", + "rt-multi-thread", + "sync", +] } +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] } +pyo3-log = { workspace = true } +arrow = { workspace = true, features = ["pyarrow"] } +arrow-select = { workspace = true } +datafusion = { workspace = true, features = ["avro", "unicode_expressions"] } +datafusion-substrait = { workspace = true, optional = true } +datafusion-proto = { workspace = true } +datafusion-ffi = { workspace = true } +prost = { workspace = true } # keep in line with `datafusion-substrait` +serde_json = { workspace = true } +uuid = { workspace = true, features = ["v4"] } +mimalloc = { workspace = true, optional = true, features = [ + "local_dynamic_tls", +] } +async-trait = { workspace = true } +futures = { workspace = true } +cstr = { workspace = true } +object_store = { workspace = true, features = ["aws", "gcp", "azure", "http"] } +url = { workspace = true } +log = { workspace = true } +parking_lot = { workspace = true } +datafusion-python-util = { workspace = true } + +[build-dependencies] +prost-types = { workspace = true } +pyo3-build-config = { workspace = true } + +[features] +default = ["mimalloc"] +protoc = ["datafusion-substrait/protoc"] +substrait = ["dep:datafusion-substrait"] + +[lib] +name = "datafusion_python" +crate-type = ["cdylib", "rlib"] diff --git a/build.rs b/crates/core/build.rs similarity index 100% rename from build.rs rename to crates/core/build.rs diff --git a/crates/core/src/array.rs b/crates/core/src/array.rs new file mode 100644 index 000000000..99e63ef50 --- /dev/null +++ b/crates/core/src/array.rs @@ -0,0 +1,92 @@ +// 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 std::ptr::NonNull; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef}; +use arrow::datatypes::{Field, FieldRef}; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::pyarrow::ToPyArrow; +use datafusion_python_util::validate_pycapsule; +use pyo3::ffi::c_str; +use pyo3::prelude::{PyAnyMethods, PyCapsuleMethods}; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; + +use crate::errors::PyDataFusionResult; + +/// A Python object which implements the Arrow PyCapsule for importing +/// into other libraries. +#[pyclass( + from_py_object, + name = "ArrowArrayExportable", + module = "datafusion", + frozen +)] +#[derive(Clone)] +pub struct PyArrowArrayExportable { + array: ArrayRef, + field: FieldRef, +} + +#[pymethods] +impl PyArrowArrayExportable { + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_array__<'py>( + &'py self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyDataFusionResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { + let field = if let Some(schema_capsule) = requested_schema { + validate_pycapsule(&schema_capsule, "arrow_schema")?; + + let data: NonNull = schema_capsule + .pointer_checked(Some(c_str!("arrow_schema")))? + .cast(); + let schema_ptr = unsafe { data.as_ref() }; + let desired_field = Field::try_from(schema_ptr)?; + + Arc::new(desired_field) + } else { + Arc::clone(&self.field) + }; + + let ffi_schema = FFI_ArrowSchema::try_from(&field)?; + let schema_capsule = PyCapsule::new(py, ffi_schema, Some(cr"arrow_schema".into()))?; + + let ffi_array = FFI_ArrowArray::new(&self.array.to_data()); + let array_capsule = PyCapsule::new(py, ffi_array, Some(cr"arrow_array".into()))?; + + Ok((schema_capsule, array_capsule)) + } +} + +impl ToPyArrow for PyArrowArrayExportable { + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { + let module = py.import("pyarrow")?; + let method = module.getattr("array")?; + let array = method.call((self.clone(),), None)?; + Ok(array) + } +} + +impl PyArrowArrayExportable { + pub fn new(array: ArrayRef, field: FieldRef) -> Self { + Self { array, field } + } +} diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs new file mode 100644 index 000000000..f707e7e5c --- /dev/null +++ b/crates/core/src/catalog.rs @@ -0,0 +1,731 @@ +// 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 std::any::Any; +use std::collections::HashSet; +use std::ptr::NonNull; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::{ + CatalogProvider, CatalogProviderList, MemoryCatalogProvider, MemoryCatalogProviderList, + MemorySchemaProvider, SchemaProvider, +}; +use datafusion::common::DataFusionError; +use datafusion::datasource::TableProvider; +use datafusion_ffi::catalog_provider::FFI_CatalogProvider; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::schema_provider::FFI_SchemaProvider; +use datafusion_python_util::{ + create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, validate_pycapsule, + wait_for_future, +}; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyKeyError; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::context::PySessionContext; +use crate::dataset::Dataset; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err}; +use crate::table::PyTable; + +#[pyclass( + from_py_object, + frozen, + name = "RawCatalogList", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PyCatalogList { + pub catalog_list: Arc, + codec: Arc, +} + +#[pyclass( + from_py_object, + frozen, + name = "RawCatalog", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PyCatalog { + pub catalog: Arc, + codec: Arc, +} + +#[pyclass( + from_py_object, + frozen, + name = "RawSchema", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PySchema { + pub schema: Arc, + codec: Arc, +} + +impl PyCatalog { + pub(crate) fn new_from_parts( + catalog: Arc, + codec: Arc, + ) -> Self { + Self { catalog, codec } + } +} + +impl PySchema { + pub(crate) fn new_from_parts( + schema: Arc, + codec: Arc, + ) -> Self { + Self { schema, codec } + } +} + +#[pymethods] +impl PyCatalogList { + #[new] + pub fn new( + py: Python, + catalog_list: Py, + session: Option>, + ) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog_list = Arc::new(RustWrappedPyCatalogProviderList::new( + catalog_list, + codec.clone(), + )) as Arc; + Ok(Self { + catalog_list, + codec, + }) + } + + #[staticmethod] + pub fn memory_catalog_list(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog_list = + Arc::new(MemoryCatalogProviderList::default()) as Arc; + Ok(Self { + catalog_list, + codec, + }) + } + + pub fn catalog_names(&self) -> HashSet { + self.catalog_list.catalog_names().into_iter().collect() + } + + #[pyo3(signature = (name="public"))] + pub fn catalog(&self, name: &str) -> PyResult> { + let catalog = self + .catalog_list + .catalog(name) + .ok_or(PyKeyError::new_err(format!( + "Schema with name {name} doesn't exist." + )))?; + + Python::attach(|py| { + match catalog + .as_any() + .downcast_ref::() + { + Some(wrapped_catalog) => Ok(wrapped_catalog.catalog_provider.clone_ref(py)), + None => PyCatalog::new_from_parts(catalog, self.codec.clone()).into_py_any(py), + } + }) + } + + pub fn register_catalog(&self, name: &str, catalog_provider: Bound<'_, PyAny>) -> PyResult<()> { + let provider = extract_catalog_provider_from_pyobj(catalog_provider, self.codec.as_ref())?; + + let _ = self + .catalog_list + .register_catalog(name.to_owned(), provider); + + Ok(()) + } + + pub fn __repr__(&self) -> PyResult { + let mut names: Vec = self.catalog_names().into_iter().collect(); + names.sort(); + Ok(format!("CatalogList(catalog_names=[{}])", names.join(", "))) + } +} + +#[pymethods] +impl PyCatalog { + #[new] + pub fn new(py: Python, catalog: Py, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog = Arc::new(RustWrappedPyCatalogProvider::new(catalog, codec.clone())) + as Arc; + Ok(Self { catalog, codec }) + } + + #[staticmethod] + pub fn memory_catalog(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let catalog = Arc::new(MemoryCatalogProvider::default()) as Arc; + Ok(Self { catalog, codec }) + } + + pub fn schema_names(&self) -> HashSet { + self.catalog.schema_names().into_iter().collect() + } + + #[pyo3(signature = (name="public"))] + pub fn schema(&self, name: &str) -> PyResult> { + let schema = self + .catalog + .schema(name) + .ok_or(PyKeyError::new_err(format!( + "Schema with name {name} doesn't exist." + )))?; + + Python::attach(|py| { + match schema + .as_any() + .downcast_ref::() + { + Some(wrapped_schema) => Ok(wrapped_schema.schema_provider.clone_ref(py)), + None => PySchema::new_from_parts(schema, self.codec.clone()).into_py_any(py), + } + }) + } + + pub fn register_schema(&self, name: &str, schema_provider: Bound<'_, PyAny>) -> PyResult<()> { + let provider = extract_schema_provider_from_pyobj(schema_provider, self.codec.as_ref())?; + + let _ = self + .catalog + .register_schema(name, provider) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + pub fn deregister_schema(&self, name: &str, cascade: bool) -> PyResult<()> { + let _ = self + .catalog + .deregister_schema(name, cascade) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + pub fn __repr__(&self) -> PyResult { + let mut names: Vec = self.schema_names().into_iter().collect(); + names.sort(); + Ok(format!("Catalog(schema_names=[{}])", names.join(", "))) + } +} + +#[pymethods] +impl PySchema { + #[new] + pub fn new( + py: Python, + schema_provider: Py, + session: Option>, + ) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let schema = + Arc::new(RustWrappedPySchemaProvider::new(schema_provider)) as Arc; + Ok(Self { schema, codec }) + } + + #[staticmethod] + fn memory_schema(py: Python, session: Option>) -> PyResult { + let codec = extract_logical_extension_codec(py, session)?; + let schema = Arc::new(MemorySchemaProvider::default()) as Arc; + Ok(Self { schema, codec }) + } + + #[getter] + fn table_names(&self) -> HashSet { + self.schema.table_names().into_iter().collect() + } + + fn table(&self, name: &str, py: Python) -> PyDataFusionResult { + if let Some(table) = wait_for_future(py, self.schema.table(name))?? { + Ok(PyTable::from(table)) + } else { + Err(PyDataFusionError::Common(format!( + "Table not found: {name}" + ))) + } + } + + fn __repr__(&self) -> PyResult { + let mut names: Vec = self.table_names().into_iter().collect(); + names.sort(); + Ok(format!("Schema(table_names=[{}])", names.join(";"))) + } + + fn register_table(&self, name: &str, table_provider: Bound<'_, PyAny>) -> PyResult<()> { + let py = table_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, self.codec.as_ref())? + .as_any() + .clone(); + + let table = PyTable::new(table_provider, Some(codec_capsule))?; + + let _ = self + .schema + .register_table(name.to_string(), table.table) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + fn deregister_table(&self, name: &str) -> PyResult<()> { + let _ = self + .schema + .deregister_table(name) + .map_err(py_datafusion_err)?; + + Ok(()) + } + + fn table_exist(&self, name: &str) -> bool { + self.schema.table_exist(name) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPySchemaProvider { + schema_provider: Py, + owner_name: Option, +} + +impl RustWrappedPySchemaProvider { + pub fn new(schema_provider: Py) -> Self { + let owner_name = Python::attach(|py| { + schema_provider + .bind(py) + .getattr("owner_name") + .ok() + .map(|name| name.to_string()) + }); + + Self { + schema_provider, + owner_name, + } + } + + fn table_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let py_table_method = provider.getattr("table")?; + + let py_table = py_table_method.call((name,), None)?; + if py_table.is_none() { + return Ok(None); + } + + let table = PyTable::new(py_table, None)?; + + Ok(Some(table.table)) + }) + } +} + +#[async_trait] +impl SchemaProvider for RustWrappedPySchemaProvider { + fn owner_name(&self) -> Option<&str> { + self.owner_name.as_deref() + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn table_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + + provider + .getattr("table_names") + .and_then(|names| names.extract::>()) + .unwrap_or_else(|err| { + log::error!("Unable to get table_names: {err}"); + Vec::default() + }) + }) + } + + async fn table( + &self, + name: &str, + ) -> datafusion::common::Result>, DataFusionError> { + self.table_inner(name) + .map_err(|e| DataFusionError::External(Box::new(e))) + } + + fn register_table( + &self, + name: String, + table: Arc, + ) -> datafusion::common::Result>> { + let py_table = PyTable::from(table); + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let _ = provider + .call_method1("register_table", (name, py_table)) + .map_err(to_datafusion_err)?; + // Since the definition of `register_table` says that an error + // will be returned if the table already exists, there is no + // case where we want to return a table provider as output. + Ok(None) + }) + } + + fn deregister_table( + &self, + name: &str, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + let table = provider + .call_method1("deregister_table", (name,)) + .map_err(to_datafusion_err)?; + if table.is_none() { + return Ok(None); + } + + // If we can turn this table provider into a `Dataset`, return it. + // Otherwise, return None. + let dataset = match Dataset::new(&table, py) { + Ok(dataset) => Some(Arc::new(dataset) as Arc), + Err(_) => None, + }; + + Ok(dataset) + }) + } + + fn table_exist(&self, name: &str) -> bool { + Python::attach(|py| { + let provider = self.schema_provider.bind(py); + provider + .call_method1("table_exist", (name,)) + .and_then(|pyobj| pyobj.extract()) + .unwrap_or(false) + }) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPyCatalogProvider { + pub(crate) catalog_provider: Py, + codec: Arc, +} + +impl RustWrappedPyCatalogProvider { + pub fn new(catalog_provider: Py, codec: Arc) -> Self { + Self { + catalog_provider, + codec, + } + } + + fn schema_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + + let py_schema = provider.call_method1("schema", (name,))?; + if py_schema.is_none() { + return Ok(None); + } + + extract_schema_provider_from_pyobj(py_schema, self.codec.as_ref()).map(Some) + }) + } +} + +#[async_trait] +impl CatalogProvider for RustWrappedPyCatalogProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + provider + .call_method0("schema_names") + .and_then(|names| names.extract::>()) + .map(|names| names.into_iter().collect()) + .unwrap_or_else(|err| { + log::error!("Unable to get schema_names: {err}"); + Vec::default() + }) + }) + } + + fn schema(&self, name: &str) -> Option> { + self.schema_inner(name).unwrap_or_else(|err| { + log::error!("CatalogProvider schema returned error: {err}"); + None + }) + } + + fn register_schema( + &self, + name: &str, + schema: Arc, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + let py_schema = match schema + .as_any() + .downcast_ref::() + { + Some(wrapped_schema) => wrapped_schema.schema_provider.as_any(), + None => &PySchema::new_from_parts(schema, self.codec.clone()) + .into_py_any(py) + .map_err(to_datafusion_err)?, + }; + + let provider = self.catalog_provider.bind(py); + let schema = provider + .call_method1("register_schema", (name, py_schema)) + .map_err(to_datafusion_err)?; + if schema.is_none() { + return Ok(None); + } + + let schema = Arc::new(RustWrappedPySchemaProvider::new(schema.into())) + as Arc; + + Ok(Some(schema)) + }) + } + + fn deregister_schema( + &self, + name: &str, + cascade: bool, + ) -> datafusion::common::Result>> { + Python::attach(|py| { + let provider = self.catalog_provider.bind(py); + let schema = provider + .call_method1("deregister_schema", (name, cascade)) + .map_err(to_datafusion_err)?; + if schema.is_none() { + return Ok(None); + } + + let schema = Arc::new(RustWrappedPySchemaProvider::new(schema.into())) + as Arc; + + Ok(Some(schema)) + }) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPyCatalogProviderList { + pub(crate) catalog_provider_list: Py, + codec: Arc, +} + +impl RustWrappedPyCatalogProviderList { + pub fn new(catalog_provider_list: Py, codec: Arc) -> Self { + Self { + catalog_provider_list, + codec, + } + } + + fn catalog_inner(&self, name: &str) -> PyResult>> { + Python::attach(|py| { + let provider = self.catalog_provider_list.bind(py); + + let py_schema = provider.call_method1("catalog", (name,))?; + if py_schema.is_none() { + return Ok(None); + } + + extract_catalog_provider_from_pyobj(py_schema, self.codec.as_ref()).map(Some) + }) + } +} + +#[async_trait] +impl CatalogProviderList for RustWrappedPyCatalogProviderList { + fn as_any(&self) -> &dyn Any { + self + } + + fn catalog_names(&self) -> Vec { + Python::attach(|py| { + let provider = self.catalog_provider_list.bind(py); + provider + .call_method0("catalog_names") + .and_then(|names| names.extract::>()) + .map(|names| names.into_iter().collect()) + .unwrap_or_else(|err| { + log::error!("Unable to get catalog_names: {err}"); + Vec::default() + }) + }) + } + + fn catalog(&self, name: &str) -> Option> { + self.catalog_inner(name).unwrap_or_else(|err| { + log::error!("CatalogProvider catalog returned error: {err}"); + None + }) + } + + fn register_catalog( + &self, + name: String, + catalog: Arc, + ) -> Option> { + Python::attach(|py| { + let py_catalog = match catalog + .as_any() + .downcast_ref::() + { + Some(wrapped_schema) => wrapped_schema.catalog_provider.as_any().clone_ref(py), + None => { + match PyCatalog::new_from_parts(catalog, self.codec.clone()).into_py_any(py) { + Ok(c) => c, + Err(err) => { + log::error!( + "register_catalog returned error during conversion to PyAny: {err}" + ); + return None; + } + } + } + }; + + let provider = self.catalog_provider_list.bind(py); + let catalog = match provider.call_method1("register_catalog", (name, py_catalog)) { + Ok(c) => c, + Err(err) => { + log::error!("register_catalog returned error: {err}"); + return None; + } + }; + if catalog.is_none() { + return None; + } + + let catalog = Arc::new(RustWrappedPyCatalogProvider::new( + catalog.into(), + self.codec.clone(), + )) as Arc; + + Some(catalog) + }) + } +} + +fn extract_catalog_provider_from_pyobj( + mut catalog_provider: Bound, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + if catalog_provider.hasattr("__datafusion_catalog_provider__")? { + let py = catalog_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, codec)?; + catalog_provider = catalog_provider + .getattr("__datafusion_catalog_provider__")? + .call1((codec_capsule,))?; + } + + let provider = if let Ok(capsule) = catalog_provider.cast::() { + validate_pycapsule(capsule, "datafusion_catalog_provider")?; + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_catalog_provider")))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider as Arc + } else { + match catalog_provider.extract::() { + Ok(py_catalog) => py_catalog.catalog, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + catalog_provider.into(), + Arc::new(codec.clone()), + )) as Arc, + } + }; + + Ok(provider) +} + +fn extract_schema_provider_from_pyobj( + mut schema_provider: Bound, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + if schema_provider.hasattr("__datafusion_schema_provider__")? { + let py = schema_provider.py(); + let codec_capsule = create_logical_extension_capsule(py, codec)?; + schema_provider = schema_provider + .getattr("__datafusion_schema_provider__")? + .call1((codec_capsule,))?; + } + + let provider = if let Ok(capsule) = schema_provider.cast::() { + validate_pycapsule(capsule, "datafusion_schema_provider")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_schema_provider")))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider as Arc + } else { + match schema_provider.extract::() { + Ok(py_schema) => py_schema.schema, + Err(_) => Arc::new(RustWrappedPySchemaProvider::new(schema_provider.into())) + as Arc, + } + }; + + Ok(provider) +} + +fn extract_logical_extension_codec( + py: Python, + obj: Option>, +) -> PyResult> { + let obj = match obj { + Some(obj) => obj, + None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, + }; + ffi_logical_codec_from_pycapsule(obj).map(Arc::new) +} + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + Ok(()) +} diff --git a/src/common.rs b/crates/core/src/common.rs similarity index 92% rename from src/common.rs rename to crates/core/src/common.rs index 453bf67a4..88d2fdd5f 100644 --- a/src/common.rs +++ b/crates/core/src/common.rs @@ -36,5 +36,8 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/src/common/data_type.rs b/crates/core/src/common/data_type.rs similarity index 72% rename from src/common/data_type.rs rename to crates/core/src/common/data_type.rs index f5f8a6b06..af4179806 100644 --- a/src/common/data_type.rs +++ b/crates/core/src/common/data_type.rs @@ -15,14 +15,18 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use datafusion::arrow::array::Array; use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; -use datafusion::common::{DataFusionError, ScalarValue}; -use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment; -use pyo3::{exceptions::PyValueError, prelude::*}; - -use crate::errors::py_datafusion_err; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::expr::NullTreatment as DFNullTreatment; +use pyo3::exceptions::{PyNotImplementedError, PyValueError}; +use pyo3::prelude::*; +/// A [`ScalarValue`] wrapped in a Python object. This struct allows for conversion +/// from a variety of Python objects into a [`ScalarValue`]. See +/// ``FromPyArrow::from_pyarrow_bound`` conversion details. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] pub struct PyScalarValue(pub ScalarValue); @@ -38,7 +42,14 @@ impl From for ScalarValue { } #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(eq, eq_int, name = "RexType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "RexType", + module = "datafusion.common" +)] pub enum RexType { Alias, Literal, @@ -57,8 +68,14 @@ pub enum RexType { /// and manageable location. Therefore this structure exists /// to map those types and provide a simple place for developers /// to map types from one system to another. +// TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now #[derive(Debug, Clone)] -#[pyclass(name = "DataTypeMap", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "DataTypeMap", + module = "datafusion.common", + subclass +)] pub struct DataTypeMap { #[pyo3(get, set)] pub arrow_type: PyDataType, @@ -171,9 +188,7 @@ impl DataTypeMap { PythonType::Datetime, SqlType::DATE, )), - DataType::Duration(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), + DataType::Duration(_) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), DataType::Interval(interval_unit) => Ok(DataTypeMap::new( DataType::Interval(*interval_unit), PythonType::Datetime, @@ -188,9 +203,9 @@ impl DataTypeMap { PythonType::Bytes, SqlType::BINARY, )), - DataType::FixedSizeBinary(_) => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", arrow_type)), - )), + DataType::FixedSizeBinary(_) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } DataType::LargeBinary => Ok(DataTypeMap::new( DataType::LargeBinary, PythonType::Bytes, @@ -206,25 +221,28 @@ impl DataTypeMap { PythonType::Str, SqlType::VARCHAR, )), - DataType::List(_) => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - arrow_type - )))), - DataType::FixedSizeList(_, _) => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", arrow_type)), + DataType::List(_) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::FixedSizeList(_, _) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } + DataType::LargeList(_) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } + DataType::Struct(_) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::Union(_, _) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::Dictionary(_, _) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } + DataType::Decimal32(precision, scale) => Ok(DataTypeMap::new( + DataType::Decimal32(*precision, *scale), + PythonType::Float, + SqlType::DECIMAL, + )), + DataType::Decimal64(precision, scale) => Ok(DataTypeMap::new( + DataType::Decimal64(*precision, *scale), + PythonType::Float, + SqlType::DECIMAL, )), - DataType::LargeList(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::Struct(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::Union(_, _) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::Dictionary(_, _) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), DataType::Decimal128(precision, scale) => Ok(DataTypeMap::new( DataType::Decimal128(*precision, *scale), PythonType::Float, @@ -235,25 +253,16 @@ impl DataTypeMap { PythonType::Float, SqlType::DECIMAL, )), - DataType::Map(_, _) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::RunEndEncoded(_, _) => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", arrow_type)), - )), - DataType::BinaryView => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::Utf8View => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - arrow_type - )))), - DataType::ListView(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), - DataType::LargeListView(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", arrow_type), - ))), + DataType::Map(_, _) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::RunEndEncoded(_, _) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } + DataType::BinaryView => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::Utf8View => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::ListView(_) => Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))), + DataType::LargeListView(_) => { + Err(PyNotImplementedError::new_err(format!("{arrow_type:?}"))) + } } } @@ -269,6 +278,12 @@ impl DataTypeMap { ScalarValue::Float16(_) => Ok(DataType::Float16), ScalarValue::Float32(_) => Ok(DataType::Float32), ScalarValue::Float64(_) => Ok(DataType::Float64), + ScalarValue::Decimal32(_, precision, scale) => { + Ok(DataType::Decimal32(*precision, *scale)) + } + ScalarValue::Decimal64(_, precision, scale) => { + Ok(DataType::Decimal64(*precision, *scale)) + } ScalarValue::Decimal128(_, precision, scale) => { Ok(DataType::Decimal128(*precision, *scale)) } @@ -319,33 +334,37 @@ impl DataTypeMap { Ok(DataType::Interval(IntervalUnit::MonthDayNano)) } ScalarValue::List(arr) => Ok(arr.data_type().to_owned()), - ScalarValue::Struct(_fields) => Err(py_datafusion_err( - DataFusionError::NotImplemented("ScalarValue::Struct".to_string()), + ScalarValue::Struct(_fields) => Err(PyNotImplementedError::new_err( + "ScalarValue::Struct".to_string(), )), ScalarValue::FixedSizeBinary(size, _) => Ok(DataType::FixedSizeBinary(*size)), ScalarValue::FixedSizeList(_array_ref) => { // The FieldRef was removed from ScalarValue::FixedSizeList in // https://github.com/apache/arrow-datafusion/pull/8221, so we can no // longer convert back to a DataType here - Err(py_datafusion_err(DataFusionError::NotImplemented( + Err(PyNotImplementedError::new_err( "ScalarValue::FixedSizeList".to_string(), - ))) + )) } - ScalarValue::LargeList(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( + ScalarValue::LargeList(_) => Err(PyNotImplementedError::new_err( "ScalarValue::LargeList".to_string(), - ))), + )), ScalarValue::DurationSecond(_) => Ok(DataType::Duration(TimeUnit::Second)), ScalarValue::DurationMillisecond(_) => Ok(DataType::Duration(TimeUnit::Millisecond)), ScalarValue::DurationMicrosecond(_) => Ok(DataType::Duration(TimeUnit::Microsecond)), ScalarValue::DurationNanosecond(_) => Ok(DataType::Duration(TimeUnit::Nanosecond)), - ScalarValue::Union(_, _, _) => Err(py_datafusion_err(DataFusionError::NotImplemented( + ScalarValue::Union(_, _, _) => Err(PyNotImplementedError::new_err( "ScalarValue::LargeList".to_string(), - ))), + )), ScalarValue::Utf8View(_) => Ok(DataType::Utf8View), ScalarValue::BinaryView(_) => Ok(DataType::BinaryView), - ScalarValue::Map(_) => Err(py_datafusion_err(DataFusionError::NotImplemented( + ScalarValue::Map(_) => Err(PyNotImplementedError::new_err( "ScalarValue::Map".to_string(), - ))), + )), + ScalarValue::RunEndEncoded(field1, field2, _) => Ok(DataType::RunEndEncoded( + Arc::clone(field1), + Arc::clone(field2), + )), } } } @@ -379,8 +398,7 @@ impl DataTypeMap { "double" => Ok(DataType::Float64), "byte_array" => Ok(DataType::Utf8), _ => Err(PyValueError::new_err(format!( - "Unable to determine Arrow Data Type from Parquet String type: {:?}", - parquet_str_type + "Unable to determine Arrow Data Type from Parquet String type: {parquet_str_type:?}" ))), }; DataTypeMap::map_from_arrow_type(&arrow_dtype?) @@ -403,14 +421,8 @@ impl DataTypeMap { #[pyo3(name = "sql")] pub fn py_map_from_sql_type(sql_type: &SqlType) -> PyResult { match sql_type { - SqlType::ANY => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::ARRAY => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::ANY => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::ARRAY => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::BIGINT => Ok(DataTypeMap::new( DataType::Int64, PythonType::Int, @@ -431,13 +443,8 @@ impl DataTypeMap { PythonType::Int, SqlType::CHAR, )), - SqlType::COLUMN_LIST => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::CURSOR => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::COLUMN_LIST => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::CURSOR => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::DATE => Ok(DataTypeMap::new( DataType::Date64, PythonType::Datetime, @@ -448,139 +455,88 @@ impl DataTypeMap { PythonType::Float, SqlType::DECIMAL, )), - SqlType::DISTINCT => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::DISTINCT => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::DOUBLE => Ok(DataTypeMap::new( DataType::Decimal256(1, 1), PythonType::Float, SqlType::DOUBLE, )), - SqlType::DYNAMIC_STAR => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), + SqlType::DYNAMIC_STAR => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::FLOAT => Ok(DataTypeMap::new( DataType::Decimal128(1, 1), PythonType::Float, SqlType::FLOAT, )), - SqlType::GEOMETRY => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::GEOMETRY => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::INTEGER => Ok(DataTypeMap::new( DataType::Int8, PythonType::Int, SqlType::INTEGER, )), - SqlType::INTERVAL => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::INTERVAL_DAY => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_DAY_HOUR => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_DAY_MINUTE => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::INTERVAL_DAY_SECOND => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::INTERVAL_HOUR => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_HOUR_MINUTE => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::INTERVAL_HOUR_SECOND => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::INTERVAL_MINUTE => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_MINUTE_SECOND => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::INTERVAL_MONTH => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_SECOND => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_YEAR => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::INTERVAL_YEAR_MONTH => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::MAP => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::MULTISET => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::INTERVAL => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::INTERVAL_DAY => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::INTERVAL_DAY_HOUR => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_DAY_MINUTE => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_DAY_SECOND => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_HOUR => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::INTERVAL_HOUR_MINUTE => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_HOUR_SECOND => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_MINUTE => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_MINUTE_SECOND => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_MONTH => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::INTERVAL_SECOND => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::INTERVAL_YEAR => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::INTERVAL_YEAR_MONTH => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::MAP => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::MULTISET => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::NULL => Ok(DataTypeMap::new( DataType::Null, PythonType::None, SqlType::NULL, )), - SqlType::OTHER => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::REAL => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::ROW => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::SARG => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::OTHER => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::REAL => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::ROW => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::SARG => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::SMALLINT => Ok(DataTypeMap::new( DataType::Int16, PythonType::Int, SqlType::SMALLINT, )), - SqlType::STRUCTURED => Err(py_datafusion_err(DataFusionError::NotImplemented( - format!("{:?}", sql_type), - ))), - SqlType::SYMBOL => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::TIME => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::TIME_WITH_LOCAL_TIME_ZONE => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), - SqlType::TIMESTAMP => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), - SqlType::TIMESTAMP_WITH_LOCAL_TIME_ZONE => Err(py_datafusion_err( - DataFusionError::NotImplemented(format!("{:?}", sql_type)), - )), + SqlType::STRUCTURED => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::SYMBOL => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::TIME => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::TIME_WITH_LOCAL_TIME_ZONE => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } + SqlType::TIMESTAMP => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), + SqlType::TIMESTAMP_WITH_LOCAL_TIME_ZONE => { + Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))) + } SqlType::TINYINT => Ok(DataTypeMap::new( DataType::Int8, PythonType::Int, SqlType::TINYINT, )), - SqlType::UNKNOWN => Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - sql_type - )))), + SqlType::UNKNOWN => Err(PyNotImplementedError::new_err(format!("{sql_type:?}"))), SqlType::VARBINARY => Ok(DataTypeMap::new( DataType::LargeBinary, PythonType::Bytes, @@ -595,7 +551,7 @@ impl DataTypeMap { } /// Unfortunately PyO3 does not allow for us to expose the DataType as an enum since - /// we cannot directly annotae the Enum instance of dependency code. Therefore, here + /// we cannot directly annotate the Enum instance of dependency code. Therefore, here /// we provide an enum to mimic it. #[pyo3(name = "friendly_arrow_type_name")] pub fn friendly_arrow_type_name(&self) -> PyResult<&str> { @@ -631,6 +587,8 @@ impl DataTypeMap { DataType::Struct(_) => "Struct", DataType::Union(_, _) => "Union", DataType::Dictionary(_, _) => "Dictionary", + DataType::Decimal32(_, _) => "Decimal32", + DataType::Decimal64(_, _) => "Decimal64", DataType::Decimal128(_, _) => "Decimal128", DataType::Decimal256(_, _) => "Decimal256", DataType::Map(_, _) => "Map", @@ -647,7 +605,12 @@ impl DataTypeMap { /// Since `DataType` exists in another package we cannot make that happen here so we wrap /// `DataType` as `PyDataType` This exists solely to satisfy those constraints. #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(name = "DataType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + name = "DataType", + module = "datafusion.common" +)] pub struct PyDataType { pub data_type: DataType, } @@ -682,8 +645,7 @@ impl PyDataType { "datetime64" => Ok(DataType::Date64), "object" => Ok(DataType::Utf8), _ => Err(PyValueError::new_err(format!( - "Unable to determine Arrow Data Type from Arrow String type: {:?}", - arrow_str_type + "Unable to determine Arrow Data Type from Arrow String type: {arrow_str_type:?}" ))), }; Ok(PyDataType { @@ -706,7 +668,14 @@ impl From for PyDataType { /// Represents the possible Python types that can be mapped to the SQL types #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(eq, eq_int, name = "PythonType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "PythonType", + module = "datafusion.common" +)] pub enum PythonType { Array, Bool, @@ -726,7 +695,14 @@ pub enum PythonType { #[allow(non_camel_case_types)] #[allow(clippy::upper_case_acronyms)] #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(eq, eq_int, name = "SqlType", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "SqlType", + module = "datafusion.common" +)] pub enum SqlType { ANY, ARRAY, @@ -784,7 +760,14 @@ pub enum SqlType { #[allow(non_camel_case_types)] #[allow(clippy::upper_case_acronyms)] #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[pyclass(eq, eq_int, name = "NullTreatment", module = "datafusion.common")] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "NullTreatment", + module = "datafusion.common" +)] pub enum NullTreatment { IGNORE_NULLS, RESPECT_NULLS, diff --git a/src/common/df_schema.rs b/crates/core/src/common/df_schema.rs similarity index 93% rename from src/common/df_schema.rs rename to crates/core/src/common/df_schema.rs index 4e1d84060..9167e772e 100644 --- a/src/common/df_schema.rs +++ b/crates/core/src/common/df_schema.rs @@ -21,7 +21,13 @@ use datafusion::common::DFSchema; use pyo3::prelude::*; #[derive(Debug, Clone)] -#[pyclass(name = "DFSchema", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DFSchema", + module = "datafusion.common", + subclass +)] pub struct PyDFSchema { schema: Arc, } diff --git a/src/common/function.rs b/crates/core/src/common/function.rs similarity index 93% rename from src/common/function.rs rename to crates/core/src/common/function.rs index a8d752f16..41cab515f 100644 --- a/src/common/function.rs +++ b/crates/core/src/common/function.rs @@ -22,7 +22,13 @@ use pyo3::prelude::*; use super::data_type::PyDataType; -#[pyclass(name = "SqlFunction", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SqlFunction", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlFunction { pub name: String, diff --git a/src/common/schema.rs b/crates/core/src/common/schema.rs similarity index 50% rename from src/common/schema.rs rename to crates/core/src/common/schema.rs index 66ce925ae..29a27b204 100644 --- a/src/common/schema.rs +++ b/crates/core/src/common/schema.rs @@ -15,30 +15,46 @@ // specific language governing permissions and limitations // under the License. -use std::{any::Any, borrow::Cow}; +use std::any::Any; +use std::borrow::Cow; +use std::fmt::{self, Display, Formatter}; +use std::sync::Arc; +use arrow::datatypes::Schema; +use arrow::pyarrow::PyArrowType; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::Constraints; +use datafusion::datasource::TableType; +use datafusion::logical_expr::utils::split_conjunction; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableSource}; +use parking_lot::RwLock; use pyo3::prelude::*; -use datafusion::logical_expr::utils::split_conjunction; - -use super::{data_type::DataTypeMap, function::SqlFunction}; +use super::data_type::DataTypeMap; +use super::function::SqlFunction; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "SqlSchema", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "SqlSchema", + module = "datafusion.common", + subclass, + frozen +)] #[derive(Debug, Clone)] pub struct SqlSchema { - #[pyo3(get, set)] - pub name: String, - #[pyo3(get, set)] - pub tables: Vec, - #[pyo3(get, set)] - pub views: Vec, - #[pyo3(get, set)] - pub functions: Vec, + name: Arc>, + tables: Arc>>, + views: Arc>>, + functions: Arc>>, } -#[pyclass(name = "SqlTable", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "SqlTable", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlTable { #[pyo3(get, set)] @@ -82,7 +98,12 @@ impl SqlTable { } } -#[pyclass(name = "SqlView", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + name = "SqlView", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlView { #[pyo3(get, set)] @@ -96,28 +117,70 @@ impl SqlSchema { #[new] pub fn new(schema_name: &str) -> Self { Self { - name: schema_name.to_owned(), - tables: Vec::new(), - views: Vec::new(), - functions: Vec::new(), + name: Arc::new(RwLock::new(schema_name.to_owned())), + tables: Arc::new(RwLock::new(Vec::new())), + views: Arc::new(RwLock::new(Vec::new())), + functions: Arc::new(RwLock::new(Vec::new())), } } + #[getter] + fn name(&self) -> PyResult { + Ok(self.name.read().clone()) + } + + #[setter] + fn set_name(&self, value: String) -> PyResult<()> { + *self.name.write() = value; + Ok(()) + } + + #[getter] + fn tables(&self) -> PyResult> { + Ok(self.tables.read().clone()) + } + + #[setter] + fn set_tables(&self, tables: Vec) -> PyResult<()> { + *self.tables.write() = tables; + Ok(()) + } + + #[getter] + fn views(&self) -> PyResult> { + Ok(self.views.read().clone()) + } + + #[setter] + fn set_views(&self, views: Vec) -> PyResult<()> { + *self.views.write() = views; + Ok(()) + } + + #[getter] + fn functions(&self) -> PyResult> { + Ok(self.functions.read().clone()) + } + + #[setter] + fn set_functions(&self, functions: Vec) -> PyResult<()> { + *self.functions.write() = functions; + Ok(()) + } + pub fn table_by_name(&self, table_name: &str) -> Option { - for tbl in &self.tables { - if tbl.name.eq(table_name) { - return Some(tbl.clone()); - } - } - None + let tables = self.tables.read(); + tables.iter().find(|tbl| tbl.name.eq(table_name)).cloned() } - pub fn add_table(&mut self, table: SqlTable) { - self.tables.push(table); + pub fn add_table(&self, table: SqlTable) { + let mut tables = self.tables.write(); + tables.push(table); } - pub fn drop_table(&mut self, table_name: String) { - self.tables.retain(|x| !x.name.eq(&table_name)); + pub fn drop_table(&self, table_name: String) { + let mut tables = self.tables.write(); + tables.retain(|x| !x.name.eq(&table_name)); } } @@ -190,7 +253,7 @@ impl TableSource for SqlTableSource { .collect() } - fn get_logical_plan(&self) -> Option> { + fn get_logical_plan(&self) -> Option> { None } } @@ -200,7 +263,13 @@ fn is_supported_push_down_expr(_expr: &Expr) -> bool { true } -#[pyclass(name = "SqlStatistics", module = "datafusion.common", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SqlStatistics", + module = "datafusion.common", + subclass +)] #[derive(Debug, Clone)] pub struct SqlStatistics { row_count: f64, @@ -218,3 +287,103 @@ impl SqlStatistics { self.row_count } } + +#[pyclass( + from_py_object, + frozen, + name = "Constraints", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyConstraints { + pub constraints: Constraints, +} + +impl From for Constraints { + fn from(constraints: PyConstraints) -> Self { + constraints.constraints + } +} + +impl From for PyConstraints { + fn from(constraints: Constraints) -> Self { + PyConstraints { constraints } + } +} + +impl Display for PyConstraints { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "Constraints: {:?}", self.constraints) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "TableType", + module = "datafusion.common" +)] +pub enum PyTableType { + Base, + View, + Temporary, +} + +impl From for datafusion::logical_expr::TableType { + fn from(table_type: PyTableType) -> Self { + match table_type { + PyTableType::Base => datafusion::logical_expr::TableType::Base, + PyTableType::View => datafusion::logical_expr::TableType::View, + PyTableType::Temporary => datafusion::logical_expr::TableType::Temporary, + } + } +} + +impl From for PyTableType { + fn from(table_type: TableType) -> Self { + match table_type { + datafusion::logical_expr::TableType::Base => PyTableType::Base, + datafusion::logical_expr::TableType::View => PyTableType::View, + datafusion::logical_expr::TableType::Temporary => PyTableType::Temporary, + } + } +} + +#[pyclass( + from_py_object, + frozen, + name = "TableSource", + module = "datafusion.common", + subclass +)] +#[derive(Clone)] +pub struct PyTableSource { + pub table_source: Arc, +} + +#[pymethods] +impl PyTableSource { + pub fn schema(&self) -> PyArrowType { + (*self.table_source.schema()).clone().into() + } + + pub fn constraints(&self) -> Option { + self.table_source.constraints().map(|c| PyConstraints { + constraints: c.clone(), + }) + } + + pub fn table_type(&self) -> PyTableType { + self.table_source.table_type().into() + } + + pub fn get_logical_plan(&self) -> Option { + self.table_source + .get_logical_plan() + .map(|plan| PyLogicalPlan::new(plan.into_owned())) + } +} diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs new file mode 100644 index 000000000..fdb693a12 --- /dev/null +++ b/crates/core/src/config.rs @@ -0,0 +1,104 @@ +// 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 std::sync::Arc; + +use datafusion::config::ConfigOptions; +use parking_lot::RwLock; +use pyo3::prelude::*; +use pyo3::types::*; + +use crate::common::data_type::PyScalarValue; +use crate::errors::PyDataFusionResult; +#[pyclass( + from_py_object, + name = "Config", + module = "datafusion", + subclass, + frozen +)] +#[derive(Clone)] +pub(crate) struct PyConfig { + config: Arc>, +} + +#[pymethods] +impl PyConfig { + #[new] + fn py_new() -> Self { + Self { + config: Arc::new(RwLock::new(ConfigOptions::new())), + } + } + + /// Get configurations from environment variables + #[staticmethod] + pub fn from_env() -> PyDataFusionResult { + Ok(Self { + config: Arc::new(RwLock::new(ConfigOptions::from_env()?)), + }) + } + + /// Get a configuration option + pub fn get<'py>(&self, key: &str, py: Python<'py>) -> PyResult> { + let value: Option> = { + let options = self.config.read(); + options + .entries() + .into_iter() + .find_map(|entry| (entry.key == key).then_some(entry.value.clone())) + }; + + match value { + Some(value) => Ok(value.into_pyobject(py)?), + None => Ok(None::.into_pyobject(py)?), + } + } + + /// Set a configuration option + pub fn set(&self, key: &str, value: Py, py: Python) -> PyDataFusionResult<()> { + let scalar_value: PyScalarValue = value.extract(py)?; + let mut options = self.config.write(); + options.set(key, scalar_value.0.to_string().as_str())?; + Ok(()) + } + + /// Get all configuration options + pub fn get_all(&self, py: Python) -> PyResult> { + let entries: Vec<(String, Option)> = { + let options = self.config.read(); + options + .entries() + .into_iter() + .map(|entry| (entry.key.clone(), entry.value.clone())) + .collect() + }; + + let dict = PyDict::new(py); + for (key, value) in entries { + dict.set_item(key, value.into_pyobject(py)?)?; + } + Ok(dict.into()) + } + + fn __repr__(&self, py: Python) -> PyResult { + match self.get_all(py) { + Ok(result) => Ok(format!("Config({result})")), + Err(err) => Ok(format!("Error: {:?}", err.to_string())), + } + } +} diff --git a/src/context.rs b/crates/core/src/context.rs similarity index 63% rename from src/context.rs rename to crates/core/src/context.rs index 0db0f4d7e..200b6470b 100644 --- a/src/context.rs +++ b/crates/core/src/context.rs @@ -17,63 +17,85 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; +use std::ptr::NonNull; use std::str::FromStr; use std::sync::Arc; use arrow::array::RecordBatchReader; use arrow::ffi_stream::ArrowArrayStreamReader; use arrow::pyarrow::FromPyArrow; -use datafusion::execution::session_state::SessionStateBuilder; -use object_store::ObjectStore; -use url::Url; -use uuid::Uuid; - -use pyo3::exceptions::{PyKeyError, PyValueError}; -use pyo3::prelude::*; - -use crate::catalog::{PyCatalog, PyTable}; -use crate::dataframe::PyDataFrame; -use crate::dataset::Dataset; -use crate::errors::{py_datafusion_err, PyDataFusionResult}; -use crate::expr::sort_expr::PySortExpr; -use crate::physical_plan::PyExecutionPlan; -use crate::record_batch::PyRecordBatchStream; -use crate::sql::exceptions::py_value_err; -use crate::sql::logical::PyLogicalPlan; -use crate::store::StorageContexts; -use crate::udaf::PyAggregateUDF; -use crate::udf::PyScalarUDF; -use crate::udwf::PyWindowUDF; -use crate::utils::{get_global_ctx, get_tokio_runtime, validate_pycapsule, wait_for_future}; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::common::TableReference; -use datafusion::common::{exec_err, ScalarValue}; +use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory}; +use datafusion::common::{ScalarValue, TableReference, exec_err}; use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; -use datafusion::datasource::MemTable; -use datafusion::datasource::TableProvider; +use datafusion::datasource::{MemTable, TableProvider}; +use datafusion::execution::TaskContextProvider; use datafusion::execution::context::{ DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext, }; -use datafusion::execution::disk_manager::DiskManagerConfig; +use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool}; use datafusion::execution::options::ReadOptions; use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::{ - AvroReadOptions, CsvReadOptions, DataFrame, NdJsonReadOptions, ParquetReadOptions, + AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, +}; +use datafusion_ffi::catalog_provider::FFI_CatalogProvider; +use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList; +use datafusion_ffi::execution::FFI_TaskContextProvider; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; +use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_python_util::{ + create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx, + get_tokio_runtime, spawn_future, validate_pycapsule, wait_for_future, }; -use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider}; -use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType}; -use tokio::task::JoinHandle; +use object_store::ObjectStore; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyKeyError, PyValueError}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; +use url::Url; +use uuid::Uuid; + +use crate::catalog::{ + PyCatalog, PyCatalogList, RustWrappedPyCatalogProvider, RustWrappedPyCatalogProviderList, +}; +use crate::common::data_type::PyScalarValue; +use crate::dataframe::PyDataFrame; +use crate::dataset::Dataset; +use crate::errors::{ + PyDataFusionError, PyDataFusionResult, from_datafusion_error, py_datafusion_err, +}; +use crate::expr::sort_expr::PySortExpr; +use crate::options::PyCsvReadOptions; +use crate::physical_plan::PyExecutionPlan; +use crate::record_batch::PyRecordBatchStream; +use crate::sql::logical::PyLogicalPlan; +use crate::sql::util::replace_placeholders_with_strings; +use crate::store::StorageContexts; +use crate::table::{PyTable, RustWrappedPyTableProviderFactory}; +use crate::udaf::PyAggregateUDF; +use crate::udf::PyScalarUDF; +use crate::udtf::PyTableFunction; +use crate::udwf::PyWindowUDF; /// Configuration options for a SessionContext -#[pyclass(name = "SessionConfig", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SessionConfig", + module = "datafusion", + subclass +)] #[derive(Clone, Default)] pub struct PySessionConfig { pub config: SessionConfig, @@ -166,7 +188,13 @@ impl PySessionConfig { } /// Runtime options for a SessionContext -#[pyclass(name = "RuntimeEnvBuilder", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RuntimeEnvBuilder", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PyRuntimeEnvBuilder { pub builder: RuntimeEnvBuilder, @@ -182,22 +210,49 @@ impl PyRuntimeEnvBuilder { } fn with_disk_manager_disabled(&self) -> Self { - let mut builder = self.builder.clone(); - builder = builder.with_disk_manager(DiskManagerConfig::Disabled); - Self { builder } + let mut runtime_builder = self.builder.clone(); + + let mut disk_mgr_builder = runtime_builder + .disk_manager_builder + .clone() + .unwrap_or_default(); + disk_mgr_builder.set_mode(DiskManagerMode::Disabled); + + runtime_builder = runtime_builder.with_disk_manager_builder(disk_mgr_builder); + Self { + builder: runtime_builder, + } } fn with_disk_manager_os(&self) -> Self { - let builder = self.builder.clone(); - let builder = builder.with_disk_manager(DiskManagerConfig::NewOs); - Self { builder } + let mut runtime_builder = self.builder.clone(); + + let mut disk_mgr_builder = runtime_builder + .disk_manager_builder + .clone() + .unwrap_or_default(); + disk_mgr_builder.set_mode(DiskManagerMode::OsTmpDirectory); + + runtime_builder = runtime_builder.with_disk_manager_builder(disk_mgr_builder); + Self { + builder: runtime_builder, + } } fn with_disk_manager_specified(&self, paths: Vec) -> Self { - let builder = self.builder.clone(); let paths = paths.iter().map(|s| s.into()).collect(); - let builder = builder.with_disk_manager(DiskManagerConfig::NewSpecified(paths)); - Self { builder } + let mut runtime_builder = self.builder.clone(); + + let mut disk_mgr_builder = runtime_builder + .disk_manager_builder + .clone() + .unwrap_or_default(); + disk_mgr_builder.set_mode(DiskManagerMode::Directories(paths)); + + runtime_builder = runtime_builder.with_disk_manager_builder(disk_mgr_builder); + Self { + builder: runtime_builder, + } } fn with_unbounded_memory_pool(&self) -> Self { @@ -226,7 +281,13 @@ impl PyRuntimeEnvBuilder { } /// `PySQLOptions` allows you to specify options to the sql execution. -#[pyclass(name = "SQLOptions", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SQLOptions", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PySQLOptions { pub options: SQLOptions, @@ -265,10 +326,17 @@ impl PySQLOptions { /// `PySessionContext` is able to plan and execute DataFusion plans. /// It has a powerful optimizer, a physical planner for local execution, and a /// multi-threaded execution engine to perform the execution. -#[pyclass(name = "SessionContext", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SessionContext", + module = "datafusion", + subclass +)] #[derive(Clone)] pub struct PySessionContext { - pub ctx: SessionContext, + pub ctx: Arc, + logical_codec: Arc, } #[pymethods] @@ -295,29 +363,30 @@ impl PySessionContext { .with_runtime_env(runtime) .with_default_features() .build(); - Ok(PySessionContext { - ctx: SessionContext::new_with_state(session_state), - }) + let ctx = Arc::new(SessionContext::new_with_state(session_state)); + let logical_codec = Self::default_logical_codec(&ctx); + Ok(PySessionContext { ctx, logical_codec }) } pub fn enable_url_table(&self) -> PyResult { Ok(PySessionContext { - ctx: self.ctx.clone().enable_url_table(), + ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), + logical_codec: Arc::clone(&self.logical_codec), }) } - #[classmethod] + #[staticmethod] #[pyo3(signature = ())] - fn global_ctx(_cls: &Bound<'_, PyType>) -> PyResult { - Ok(Self { - ctx: get_global_ctx().clone(), - }) + pub fn global_ctx() -> PyResult { + let ctx = get_global_ctx().clone(); + let logical_codec = Self::default_logical_codec(&ctx); + Ok(Self { ctx, logical_codec }) } /// Register an object store with the given name #[pyo3(signature = (scheme, store, host=None))] pub fn register_object_store( - &mut self, + &self, scheme: &str, store: StorageContexts, host: Option<&str>, @@ -337,7 +406,7 @@ impl PySessionContext { } else { &upstream_host }; - let url_string = format!("{}{}", scheme, derived_host); + let url_string = format!("{scheme}{derived_host}"); let url = Url::parse(&url_string).unwrap(); self.ctx.runtime_env().register_object_store(&url, store); Ok(()) @@ -349,10 +418,10 @@ impl PySessionContext { schema=None, file_sort_order=None))] pub fn register_listing_table( - &mut self, + &self, name: &str, path: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, file_extension: &str, schema: Option>, file_sort_order: Option>>, @@ -360,7 +429,12 @@ impl PySessionContext { ) -> PyDataFusionResult<()> { let options = ListingOptions::new(Arc::new(ParquetFormat::new())) .with_file_extension(file_extension) - .with_table_partition_cols(convert_table_partition_cols(table_partition_cols)?) + .with_table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ) .with_file_sort_order( file_sort_order .unwrap_or_default() @@ -374,49 +448,65 @@ impl PySessionContext { None => { let state = self.ctx.state(); let schema = options.infer_schema(&state, &table_path); - wait_for_future(py, schema)? + wait_for_future(py, schema)?? } }; let config = ListingTableConfig::new(table_path) .with_listing_options(options) .with_schema(resolved_schema); let table = ListingTable::try_new(config)?; - self.register_table( - name, - &PyTable { - table: Arc::new(table), - }, - )?; + self.ctx.register_table(name, Arc::new(table))?; Ok(()) } - /// Returns a PyDataFrame whose plan corresponds to the SQL statement. - pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult { - let result = self.ctx.sql(query); - let df = wait_for_future(py, result)?; - Ok(PyDataFrame::new(df)) + pub fn register_udtf(&self, func: PyTableFunction) { + let name = func.name.clone(); + let func = Arc::new(func); + self.ctx.register_udtf(&name, func); } - #[pyo3(signature = (query, options=None))] + #[pyo3(signature = (query, options=None, param_values=HashMap::default(), param_strings=HashMap::default()))] pub fn sql_with_options( - &mut self, - query: &str, - options: Option, + &self, py: Python, + mut query: String, + options: Option, + param_values: HashMap, + param_strings: HashMap, ) -> PyDataFusionResult { let options = if let Some(options) = options { options.options } else { SQLOptions::new() }; - let result = self.ctx.sql_with_options(query, options); - let df = wait_for_future(py, result)?; + + let param_values = param_values + .into_iter() + .map(|(name, value)| (name, ScalarValue::from(value))) + .collect::>(); + + let state = self.ctx.state(); + let dialect = state.config().options().sql_parser.dialect.as_ref(); + + if !param_strings.is_empty() { + query = replace_placeholders_with_strings(&query, dialect, param_strings)?; + } + + let mut df = wait_for_future(py, async { + self.ctx.sql_with_options(&query, options).await + })? + .map_err(from_datafusion_error)?; + + if !param_values.is_empty() { + df = df.with_param_values(param_values)?; + } + Ok(PyDataFrame::new(df)) } #[pyo3(signature = (partitions, name=None, schema=None))] pub fn create_dataframe( - &mut self, + &self, partitions: PyArrowType>>, name: Option<&str>, schema: Option>, @@ -444,21 +534,21 @@ impl PySessionContext { self.ctx.register_table(&*table_name, Arc::new(table))?; - let table = wait_for_future(py, self._table(&table_name))?; + let table = wait_for_future(py, self._table(&table_name))??; let df = PyDataFrame::new(table); Ok(df) } /// Create a DataFrame from an existing logical plan - pub fn create_dataframe_from_logical_plan(&mut self, plan: PyLogicalPlan) -> PyDataFrame { + pub fn create_dataframe_from_logical_plan(&self, plan: PyLogicalPlan) -> PyDataFrame { PyDataFrame::new(DataFrame::new(self.ctx.state(), plan.plan.as_ref().clone())) } /// Construct datafusion dataframe from Python list #[pyo3(signature = (data, name=None))] pub fn from_pylist( - &mut self, + &self, data: Bound<'_, PyList>, name: Option<&str>, ) -> PyResult { @@ -478,7 +568,7 @@ impl PySessionContext { /// Construct datafusion dataframe from Python dictionary #[pyo3(signature = (data, name=None))] pub fn from_pydict( - &mut self, + &self, data: Bound<'_, PyDict>, name: Option<&str>, ) -> PyResult { @@ -498,7 +588,7 @@ impl PySessionContext { /// Construct datafusion dataframe from Arrow Table #[pyo3(signature = (data, name=None))] pub fn from_arrow( - &mut self, + &self, data: Bound<'_, PyAny>, name: Option<&str>, py: Python, @@ -518,7 +608,7 @@ impl PySessionContext { (array.schema().as_ref().to_owned(), vec![array]) } else { - return Err(crate::errors::PyDataFusionError::Common( + return Err(PyDataFusionError::Common( "Expected either a Arrow Array or Arrow Stream in from_arrow().".to_string(), )); }; @@ -532,11 +622,7 @@ impl PySessionContext { /// Construct datafusion dataframe from pandas #[allow(clippy::wrong_self_convention)] #[pyo3(signature = (data, name=None))] - pub fn from_pandas( - &mut self, - data: Bound<'_, PyAny>, - name: Option<&str>, - ) -> PyResult { + pub fn from_pandas(&self, data: Bound<'_, PyAny>, name: Option<&str>) -> PyResult { // Obtain GIL token let py = data.py(); @@ -552,11 +638,7 @@ impl PySessionContext { /// Construct datafusion dataframe from polars #[pyo3(signature = (data, name=None))] - pub fn from_polars( - &mut self, - data: Bound<'_, PyAny>, - name: Option<&str>, - ) -> PyResult { + pub fn from_polars(&self, data: Bound<'_, PyAny>, name: Option<&str>) -> PyResult { // Convert Polars dataframe to Arrow Table let table = data.call_method0("to_arrow")?; @@ -565,43 +647,143 @@ impl PySessionContext { Ok(df) } - pub fn register_table(&mut self, name: &str, table: &PyTable) -> PyDataFusionResult<()> { - self.ctx.register_table(name, table.table())?; + pub fn register_table(&self, name: &str, table: Bound<'_, PyAny>) -> PyDataFusionResult<()> { + let session = self.clone().into_bound_py_any(table.py())?; + let table = PyTable::new(table, Some(session))?; + + self.ctx.register_table(name, table.table)?; Ok(()) } - pub fn deregister_table(&mut self, name: &str) -> PyDataFusionResult<()> { + pub fn deregister_table(&self, name: &str) -> PyDataFusionResult<()> { self.ctx.deregister_table(name)?; Ok(()) } - /// Construct datafusion dataframe from Arrow Table - pub fn register_table_provider( - &mut self, - name: &str, - provider: Bound<'_, PyAny>, + pub fn register_table_factory( + &self, + format: &str, + mut factory: Bound<'_, PyAny>, ) -> PyDataFusionResult<()> { - if provider.hasattr("__datafusion_table_provider__")? { - let capsule = provider.getattr("__datafusion_table_provider__")?.call0()?; - let capsule = capsule.downcast::().map_err(py_datafusion_err)?; - validate_pycapsule(capsule, "datafusion_table_provider")?; + if factory.hasattr("__datafusion_table_provider_factory__")? { + let py = factory.py(); + let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + factory = factory + .getattr("__datafusion_table_provider_factory__")? + .call1((codec_capsule,))?; + } - let provider = unsafe { capsule.reference::() }; - let provider: ForeignTableProvider = provider.into(); + let factory: Arc = + if let Ok(capsule) = factory.cast::().map_err(py_datafusion_err) { + validate_pycapsule(capsule, "datafusion_table_provider_factory")?; - let _ = self.ctx.register_table(name, Arc::new(provider))?; + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_table_provider_factory")))? + .cast(); + let factory = unsafe { data.as_ref() }; + factory.into() + } else { + Arc::new(RustWrappedPyTableProviderFactory::new( + factory.into(), + self.logical_codec.clone(), + )) + }; - Ok(()) + let st = self.ctx.state_ref(); + let mut lock = st.write(); + lock.table_factories_mut() + .insert(format.to_owned(), factory); + + Ok(()) + } + + pub fn register_catalog_provider_list( + &self, + mut provider: Bound, + ) -> PyDataFusionResult<()> { + if provider.hasattr("__datafusion_catalog_provider_list__")? { + let py = provider.py(); + let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + provider = provider + .getattr("__datafusion_catalog_provider_list__")? + .call1((codec_capsule,))?; + } + + let provider = if let Ok(capsule) = provider.cast::().map_err(py_datafusion_err) + { + validate_pycapsule(capsule, "datafusion_catalog_provider_list")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_catalog_provider_list")))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider as Arc } else { - Err(crate::errors::PyDataFusionError::Common( - "__datafusion_table_provider__ does not exist on Table Provider object." - .to_string(), - )) + match provider.extract::() { + Ok(py_catalog_list) => py_catalog_list.catalog_list, + Err(_) => Arc::new(RustWrappedPyCatalogProviderList::new( + provider.into(), + Arc::clone(&self.logical_codec), + )) as Arc, + } + }; + + self.ctx.register_catalog_list(provider); + + Ok(()) + } + + pub fn register_catalog_provider( + &self, + name: &str, + mut provider: Bound<'_, PyAny>, + ) -> PyDataFusionResult<()> { + if provider.hasattr("__datafusion_catalog_provider__")? { + let py = provider.py(); + let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?; + provider = provider + .getattr("__datafusion_catalog_provider__")? + .call1((codec_capsule,))?; } + + let provider = if let Ok(capsule) = provider.cast::().map_err(py_datafusion_err) + { + validate_pycapsule(capsule, "datafusion_catalog_provider")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_catalog_provider")))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + provider as Arc + } else { + match provider.extract::() { + Ok(py_catalog) => py_catalog.catalog, + Err(_) => Arc::new(RustWrappedPyCatalogProvider::new( + provider.into(), + Arc::clone(&self.logical_codec), + )) as Arc, + } + }; + + let _ = self.ctx.register_catalog(name, provider); + + Ok(()) + } + + /// Construct datafusion dataframe from Arrow Table + pub fn register_table_provider( + &self, + name: &str, + provider: Bound<'_, PyAny>, + ) -> PyDataFusionResult<()> { + // Deprecated: use `register_table` instead + self.register_table(name, provider) } pub fn register_record_batches( - &mut self, + &self, name: &str, partitions: PyArrowType>>, ) -> PyDataFusionResult<()> { @@ -619,10 +801,10 @@ impl PySessionContext { schema=None, file_sort_order=None))] pub fn register_parquet( - &mut self, + &self, name: &str, path: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, parquet_pruning: bool, file_extension: &str, skip_metadata: bool, @@ -631,7 +813,12 @@ impl PySessionContext { py: Python, ) -> PyDataFusionResult<()> { let mut options = ParquetReadOptions::default() - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?) + .table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ) .parquet_pruning(parquet_pruning) .skip_metadata(skip_metadata); options.file_extension = file_extension; @@ -643,54 +830,33 @@ impl PySessionContext { .collect(); let result = self.ctx.register_parquet(name, path, options); - wait_for_future(py, result)?; + wait_for_future(py, result)??; Ok(()) } - #[allow(clippy::too_many_arguments)] #[pyo3(signature = (name, path, - schema=None, - has_header=true, - delimiter=",", - schema_infer_max_records=1000, - file_extension=".csv", - file_compression_type=None))] + options=None))] pub fn register_csv( - &mut self, + &self, name: &str, path: &Bound<'_, PyAny>, - schema: Option>, - has_header: bool, - delimiter: &str, - schema_infer_max_records: usize, - file_extension: &str, - file_compression_type: Option, + options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult<()> { - let delimiter = delimiter.as_bytes(); - if delimiter.len() != 1 { - return Err(crate::errors::PyDataFusionError::PythonError(py_value_err( - "Delimiter must be a single character", - ))); - } - - let mut options = CsvReadOptions::new() - .has_header(has_header) - .delimiter(delimiter[0]) - .schema_infer_max_records(schema_infer_max_records) - .file_extension(file_extension) - .file_compression_type(parse_file_compression_type(file_compression_type)?); - options.schema = schema.as_ref().map(|x| &x.0); + let options = options + .map(|opts| opts.try_into()) + .transpose()? + .unwrap_or_default(); if path.is_instance_of::() { let paths = path.extract::>()?; let result = self.register_csv_from_multiple_paths(name, paths, options); - wait_for_future(py, result)?; + wait_for_future(py, result)??; } else { let path = path.extract::()?; let result = self.ctx.register_csv(name, &path, options); - wait_for_future(py, result)?; + wait_for_future(py, result)??; } Ok(()) @@ -705,13 +871,13 @@ impl PySessionContext { table_partition_cols=vec![], file_compression_type=None))] pub fn register_json( - &mut self, + &self, name: &str, path: PathBuf, schema: Option>, schema_infer_max_records: usize, file_extension: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, file_compression_type: Option, py: Python, ) -> PyDataFusionResult<()> { @@ -719,15 +885,20 @@ impl PySessionContext { .to_str() .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - let mut options = NdJsonReadOptions::default() + let mut options = JsonReadOptions::default() .file_compression_type(parse_file_compression_type(file_compression_type)?) - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?); + .table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ); options.schema_infer_max_records = schema_infer_max_records; options.file_extension = file_extension; options.schema = schema.as_ref().map(|x| &x.0); let result = self.ctx.register_json(name, path, options); - wait_for_future(py, result)?; + wait_for_future(py, result)??; Ok(()) } @@ -739,25 +910,29 @@ impl PySessionContext { file_extension=".avro", table_partition_cols=vec![]))] pub fn register_avro( - &mut self, + &self, name: &str, path: PathBuf, schema: Option>, file_extension: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, py: Python, ) -> PyDataFusionResult<()> { let path = path .to_str() .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - let mut options = AvroReadOptions::default() - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?); + let mut options = AvroReadOptions::default().table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ); options.file_extension = file_extension; options.schema = schema.as_ref().map(|x| &x.0); let result = self.ctx.register_avro(name, path, options); - wait_for_future(py, result)?; + wait_for_future(py, result)??; Ok(()) } @@ -776,32 +951,43 @@ impl PySessionContext { Ok(()) } - pub fn register_udf(&mut self, udf: PyScalarUDF) -> PyResult<()> { + pub fn register_udf(&self, udf: PyScalarUDF) -> PyResult<()> { self.ctx.register_udf(udf.function); Ok(()) } - pub fn register_udaf(&mut self, udaf: PyAggregateUDF) -> PyResult<()> { + pub fn register_udaf(&self, udaf: PyAggregateUDF) -> PyResult<()> { self.ctx.register_udaf(udaf.function); Ok(()) } - pub fn register_udwf(&mut self, udwf: PyWindowUDF) -> PyResult<()> { + pub fn register_udwf(&self, udwf: PyWindowUDF) -> PyResult<()> { self.ctx.register_udwf(udwf.function); Ok(()) } #[pyo3(signature = (name="datafusion"))] - pub fn catalog(&self, name: &str) -> PyResult { - match self.ctx.catalog(name) { - Some(catalog) => Ok(PyCatalog::new(catalog)), - None => Err(PyKeyError::new_err(format!( - "Catalog with name {} doesn't exist.", - &name, - ))), + pub fn catalog(&self, py: Python, name: &str) -> PyResult> { + let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!( + "Catalog with name {name} doesn't exist." + )))?; + + match catalog + .as_any() + .downcast_ref::() + { + Some(wrapped_schema) => Ok(wrapped_schema.catalog_provider.clone_ref(py)), + None => Ok( + PyCatalog::new_from_parts(catalog, Arc::clone(&self.logical_codec)) + .into_py_any(py)?, + ), } } + pub fn catalog_names(&self) -> HashSet { + self.ctx.catalog_names().into_iter().collect() + } + pub fn tables(&self) -> HashSet { self.ctx .catalog_names() @@ -818,9 +1004,19 @@ impl PySessionContext { } pub fn table(&self, name: &str, py: Python) -> PyResult { - let x = wait_for_future(py, self.ctx.table(name)) + let res = wait_for_future(py, self.ctx.table(name)) .map_err(|e| PyKeyError::new_err(e.to_string()))?; - Ok(PyDataFrame::new(x)) + match res { + Ok(df) => Ok(PyDataFrame::new(df)), + Err(e) => { + if let datafusion::error::DataFusionError::Plan(msg) = &e + && msg.contains("No table named") + { + return Err(PyKeyError::new_err(msg.to_string())); + } + Err(py_datafusion_err(e)) + } + } } pub fn table_exist(&self, name: &str) -> PyDataFusionResult { @@ -838,82 +1034,63 @@ impl PySessionContext { #[allow(clippy::too_many_arguments)] #[pyo3(signature = (path, schema=None, schema_infer_max_records=1000, file_extension=".json", table_partition_cols=vec![], file_compression_type=None))] pub fn read_json( - &mut self, + &self, path: PathBuf, schema: Option>, schema_infer_max_records: usize, file_extension: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, file_compression_type: Option, py: Python, ) -> PyDataFusionResult { let path = path .to_str() .ok_or_else(|| PyValueError::new_err("Unable to convert path to a string"))?; - let mut options = NdJsonReadOptions::default() - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?) + let mut options = JsonReadOptions::default() + .table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ) .file_compression_type(parse_file_compression_type(file_compression_type)?); options.schema_infer_max_records = schema_infer_max_records; options.file_extension = file_extension; let df = if let Some(schema) = schema { options.schema = Some(&schema.0); let result = self.ctx.read_json(path, options); - wait_for_future(py, result)? + wait_for_future(py, result)?? } else { let result = self.ctx.read_json(path, options); - wait_for_future(py, result)? + wait_for_future(py, result)?? }; Ok(PyDataFrame::new(df)) } - #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( path, - schema=None, - has_header=true, - delimiter=",", - schema_infer_max_records=1000, - file_extension=".csv", - table_partition_cols=vec![], - file_compression_type=None))] + options=None))] pub fn read_csv( &self, path: &Bound<'_, PyAny>, - schema: Option>, - has_header: bool, - delimiter: &str, - schema_infer_max_records: usize, - file_extension: &str, - table_partition_cols: Vec<(String, String)>, - file_compression_type: Option, + options: Option<&PyCsvReadOptions>, py: Python, ) -> PyDataFusionResult { - let delimiter = delimiter.as_bytes(); - if delimiter.len() != 1 { - return Err(crate::errors::PyDataFusionError::PythonError(py_value_err( - "Delimiter must be a single character", - ))); - }; - - let mut options = CsvReadOptions::new() - .has_header(has_header) - .delimiter(delimiter[0]) - .schema_infer_max_records(schema_infer_max_records) - .file_extension(file_extension) - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?) - .file_compression_type(parse_file_compression_type(file_compression_type)?); - options.schema = schema.as_ref().map(|x| &x.0); + let options = options + .map(|opts| opts.try_into()) + .transpose()? + .unwrap_or_default(); if path.is_instance_of::() { let paths = path.extract::>()?; let paths = paths.iter().map(|p| p as &str).collect::>(); let result = self.ctx.read_csv(paths, options); - let df = PyDataFrame::new(wait_for_future(py, result)?); + let df = PyDataFrame::new(wait_for_future(py, result)??); Ok(df) } else { let path = path.extract::()?; let result = self.ctx.read_csv(path, options); - let df = PyDataFrame::new(wait_for_future(py, result)?); + let df = PyDataFrame::new(wait_for_future(py, result)??); Ok(df) } } @@ -930,7 +1107,7 @@ impl PySessionContext { pub fn read_parquet( &self, path: &str, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, parquet_pruning: bool, file_extension: &str, skip_metadata: bool, @@ -939,7 +1116,12 @@ impl PySessionContext { py: Python, ) -> PyDataFusionResult { let mut options = ParquetReadOptions::default() - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?) + .table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ) .parquet_pruning(parquet_pruning) .skip_metadata(skip_metadata); options.file_extension = file_extension; @@ -951,7 +1133,7 @@ impl PySessionContext { .collect(); let result = self.ctx.read_parquet(path, options); - let df = PyDataFrame::new(wait_for_future(py, result)?); + let df = PyDataFrame::new(wait_for_future(py, result)??); Ok(df) } @@ -961,25 +1143,31 @@ impl PySessionContext { &self, path: &str, schema: Option>, - table_partition_cols: Vec<(String, String)>, + table_partition_cols: Vec<(String, PyArrowType)>, file_extension: &str, py: Python, ) -> PyDataFusionResult { - let mut options = AvroReadOptions::default() - .table_partition_cols(convert_table_partition_cols(table_partition_cols)?); + let mut options = AvroReadOptions::default().table_partition_cols( + table_partition_cols + .into_iter() + .map(|(name, ty)| (name, ty.0)) + .collect::>(), + ); options.file_extension = file_extension; let df = if let Some(schema) = schema { options.schema = Some(&schema.0); let read_future = self.ctx.read_avro(path, options); - wait_for_future(py, read_future)? + wait_for_future(py, read_future)?? } else { let read_future = self.ctx.read_avro(path, options); - wait_for_future(py, read_future)? + wait_for_future(py, read_future)?? }; Ok(PyDataFrame::new(df)) } - pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult { + pub fn read_table(&self, table: Bound<'_, PyAny>) -> PyDataFusionResult { + let session = self.clone().into_bound_py_any(table.py())?; + let table = PyTable::new(table, Some(session))?; let df = self.ctx.read_table(table.table())?; Ok(PyDataFrame::new(df)) } @@ -1009,13 +1197,42 @@ impl PySessionContext { py: Python, ) -> PyDataFusionResult { let ctx: TaskContext = TaskContext::from(&self.ctx.state()); - // create a Tokio runtime to run the async code - let rt = &get_tokio_runtime().0; let plan = plan.plan.clone(); - let fut: JoinHandle> = - rt.spawn(async move { plan.execute(part, Arc::new(ctx)) }); - let stream = wait_for_future(py, fut).map_err(py_datafusion_err)?; - Ok(PyRecordBatchStream::new(stream?)) + let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?; + Ok(PyRecordBatchStream::new(stream)) + } + + pub fn __datafusion_task_context_provider__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let name = cr"datafusion_task_context_provider".into(); + + let ctx_provider = Arc::clone(&self.ctx) as Arc; + let ffi_ctx_provider = FFI_TaskContextProvider::from(&ctx_provider); + + PyCapsule::new(py, ffi_ctx_provider, Some(name)) + } + + pub fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + create_logical_extension_capsule(py, self.logical_codec.as_ref()) + } + + pub fn with_logical_extension_codec<'py>( + &self, + codec: Bound<'py, PyAny>, + ) -> PyDataFusionResult { + let logical_codec = Arc::new(ffi_logical_codec_from_pycapsule(codec)?); + + Ok({ + Self { + ctx: Arc::clone(&self.ctx), + logical_codec, + } + }) } } @@ -1063,21 +1280,17 @@ impl PySessionContext { .register_table(TableReference::Bare { table: name.into() }, Arc::new(table))?; Ok(()) } -} -pub fn convert_table_partition_cols( - table_partition_cols: Vec<(String, String)>, -) -> PyDataFusionResult> { - table_partition_cols - .into_iter() - .map(|(name, ty)| match ty.as_str() { - "string" => Ok((name, DataType::Utf8)), - "int" => Ok((name, DataType::Int32)), - _ => Err(crate::errors::PyDataFusionError::Common(format!( - "Unsupported data type '{ty}' for partition column. Supported types are 'string' and 'int'" - ))), - }) - .collect::, _>>() + fn default_logical_codec(ctx: &Arc) -> Arc { + let codec = Arc::new(DefaultLogicalExtensionCodec {}); + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::clone(ctx) as Arc; + Arc::new(FFI_LogicalExtensionCodec::new( + codec, + Some(runtime), + &ctx_provider, + )) + } } pub fn parse_file_compression_type( @@ -1091,12 +1304,15 @@ pub fn parse_file_compression_type( impl From for SessionContext { fn from(ctx: PySessionContext) -> SessionContext { - ctx.ctx + ctx.ctx.as_ref().clone() } } impl From for PySessionContext { fn from(ctx: SessionContext) -> PySessionContext { - PySessionContext { ctx } + let ctx = Arc::new(ctx); + let logical_codec = Self::default_logical_codec(&ctx); + + PySessionContext { ctx, logical_codec } } } diff --git a/crates/core/src/dataframe.rs b/crates/core/src/dataframe.rs new file mode 100644 index 000000000..29fc05ed3 --- /dev/null +++ b/crates/core/src/dataframe.rs @@ -0,0 +1,1472 @@ +// 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 std::collections::HashMap; +use std::ffi::{CStr, CString}; +use std::ptr::NonNull; +use std::str::FromStr; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, RecordBatchReader, new_null_array}; +use arrow::compute::can_cast_types; +use arrow::error::ArrowError; +use arrow::ffi::FFI_ArrowSchema; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow::pyarrow::FromPyArrow; +use cstr::cstr; +use datafusion::arrow::datatypes::{Schema, SchemaRef}; +use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow}; +use datafusion::arrow::util::pretty; +use datafusion::catalog::TableProvider; +use datafusion::common::UnnestOptions; +use datafusion::config::{CsvOptions, ParquetColumnOptions, ParquetOptions, TableParquetOptions}; +use datafusion::dataframe::{DataFrame, DataFrameWriteOptions}; +use datafusion::error::DataFusionError; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::logical_expr::SortExpr; +use datafusion::logical_expr::dml::InsertOp; +use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; +use datafusion::prelude::*; +use datafusion_python_util::{is_ipython_env, spawn_future, validate_pycapsule, wait_for_future}; +use futures::{StreamExt, TryStreamExt}; +use parking_lot::Mutex; +use pyo3::PyErr; +use pyo3::exceptions::PyValueError; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods}; + +use crate::common::data_type::PyScalarValue; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err}; +use crate::expr::PyExpr; +use crate::expr::sort_expr::{PySortExpr, to_sort_expressions}; +use crate::physical_plan::PyExecutionPlan; +use crate::record_batch::{PyRecordBatchStream, poll_next_batch}; +use crate::sql::logical::PyLogicalPlan; +use crate::table::{PyTable, TempViewTable}; + +/// File-level static CStr for the Arrow array stream capsule name. +static ARROW_ARRAY_STREAM_NAME: &CStr = cstr!("arrow_array_stream"); + +// Type aliases to simplify very complex types used in this file and +// avoid compiler complaints about deeply nested types in struct fields. +type CachedBatches = Option<(Vec, bool)>; +type SharedCachedBatches = Arc>; + +/// Configuration for DataFrame display formatting +#[derive(Debug, Clone)] +pub struct FormatterConfig { + /// Maximum memory in bytes to use for display (default: 2MB) + pub max_bytes: usize, + /// Minimum number of rows to display (default: 10) + pub min_rows: usize, + /// Maximum number of rows to include in __repr__ output (default: 10) + pub max_rows: usize, +} + +impl Default for FormatterConfig { + fn default() -> Self { + Self { + max_bytes: 2 * 1024 * 1024, // 2MB + min_rows: 10, + max_rows: 10, + } + } +} + +impl FormatterConfig { + /// Validates that all configuration values are positive integers. + /// + /// # Returns + /// + /// `Ok(())` if all values are valid, or an `Err` with a descriptive error message. + pub fn validate(&self) -> Result<(), String> { + if self.max_bytes == 0 { + return Err("max_bytes must be a positive integer".to_string()); + } + + if self.min_rows == 0 { + return Err("min_rows must be a positive integer".to_string()); + } + + if self.max_rows == 0 { + return Err("max_rows must be a positive integer".to_string()); + } + + if self.min_rows > self.max_rows { + return Err("min_rows must be less than or equal to max_rows".to_string()); + } + + Ok(()) + } +} + +/// Holds the Python formatter and its configuration +struct PythonFormatter<'py> { + /// The Python formatter object + formatter: Bound<'py, PyAny>, + /// The formatter configuration + config: FormatterConfig, +} + +/// Get the Python formatter and its configuration +fn get_python_formatter_with_config(py: Python) -> PyResult { + let formatter = import_python_formatter(py)?; + let config = build_formatter_config_from_python(&formatter)?; + Ok(PythonFormatter { formatter, config }) +} + +/// Get the Python formatter from the datafusion.dataframe_formatter module +fn import_python_formatter(py: Python<'_>) -> PyResult> { + let formatter_module = py.import("datafusion.dataframe_formatter")?; + let get_formatter = formatter_module.getattr("get_formatter")?; + get_formatter.call0() +} + +// Helper function to extract attributes with fallback to default +fn get_attr<'a, T>(py_object: &'a Bound<'a, PyAny>, attr_name: &str, default_value: T) -> T +where + T: for<'py> pyo3::FromPyObject<'py, 'py> + Clone, +{ + py_object + .getattr(attr_name) + .and_then(|v| v.extract::().map_err(Into::::into)) + .unwrap_or_else(|_| default_value.clone()) +} + +/// Helper function to create a FormatterConfig from a Python formatter object +fn build_formatter_config_from_python(formatter: &Bound<'_, PyAny>) -> PyResult { + let default_config = FormatterConfig::default(); + let max_bytes = get_attr(formatter, "max_memory_bytes", default_config.max_bytes); + let min_rows = get_attr(formatter, "min_rows", default_config.min_rows); + + // Backward compatibility: Try max_rows first (new name), fall back to repr_rows (deprecated), + // then use default. This ensures backward compatibility with custom formatter implementations + // during the deprecation period. + let max_rows = get_attr(formatter, "max_rows", 0usize); + let max_rows = if max_rows > 0 { + // max_rows attribute exists and has a value + max_rows + } else { + // Try the deprecated repr_rows attribute + let repr_rows = get_attr(formatter, "repr_rows", 0usize); + if repr_rows > 0 { + repr_rows + } else { + // Use default + default_config.max_rows + } + }; + + let config = FormatterConfig { + max_bytes, + min_rows, + max_rows, + }; + + // Return the validated config, converting String error to PyErr + config.validate().map_err(PyValueError::new_err)?; + Ok(config) +} + +/// Python mapping of `ParquetOptions` (includes just the writer-related options). +#[pyclass( + from_py_object, + frozen, + name = "ParquetWriterOptions", + module = "datafusion", + subclass +)] +#[derive(Clone, Default)] +pub struct PyParquetWriterOptions { + options: ParquetOptions, +} + +#[pymethods] +impl PyParquetWriterOptions { + #[new] + #[allow(clippy::too_many_arguments)] + pub fn new( + data_pagesize_limit: usize, + write_batch_size: usize, + writer_version: &str, + skip_arrow_metadata: bool, + compression: Option, + dictionary_enabled: Option, + dictionary_page_size_limit: usize, + statistics_enabled: Option, + max_row_group_size: usize, + created_by: String, + column_index_truncate_length: Option, + statistics_truncate_length: Option, + data_page_row_count_limit: usize, + encoding: Option, + bloom_filter_on_write: bool, + bloom_filter_fpp: Option, + bloom_filter_ndv: Option, + allow_single_file_parallelism: bool, + maximum_parallel_row_group_writers: usize, + maximum_buffered_record_batches_per_stream: usize, + ) -> PyResult { + let writer_version = + datafusion::common::parquet_config::DFParquetWriterVersion::from_str(writer_version) + .map_err(py_datafusion_err)?; + Ok(Self { + options: ParquetOptions { + data_pagesize_limit, + write_batch_size, + writer_version, + skip_arrow_metadata, + compression, + dictionary_enabled, + dictionary_page_size_limit, + statistics_enabled, + max_row_group_size, + created_by, + column_index_truncate_length, + statistics_truncate_length, + data_page_row_count_limit, + encoding, + bloom_filter_on_write, + bloom_filter_fpp, + bloom_filter_ndv, + allow_single_file_parallelism, + maximum_parallel_row_group_writers, + maximum_buffered_record_batches_per_stream, + ..Default::default() + }, + }) + } +} + +/// Python mapping of `ParquetColumnOptions`. +#[pyclass( + from_py_object, + frozen, + name = "ParquetColumnOptions", + module = "datafusion", + subclass +)] +#[derive(Clone, Default)] +pub struct PyParquetColumnOptions { + options: ParquetColumnOptions, +} + +#[pymethods] +impl PyParquetColumnOptions { + #[new] + pub fn new( + bloom_filter_enabled: Option, + encoding: Option, + dictionary_enabled: Option, + compression: Option, + statistics_enabled: Option, + bloom_filter_fpp: Option, + bloom_filter_ndv: Option, + ) -> Self { + Self { + options: ParquetColumnOptions { + bloom_filter_enabled, + encoding, + dictionary_enabled, + compression, + statistics_enabled, + bloom_filter_fpp, + bloom_filter_ndv, + }, + } + } +} + +/// A PyDataFrame is a representation of a logical plan and an API to compose statements. +/// Use it to build a plan and `.collect()` to execute the plan and collect the result. +/// The actual execution of a plan runs natively on Rust and Arrow on a multi-threaded environment. +#[pyclass( + from_py_object, + name = "DataFrame", + module = "datafusion", + subclass, + frozen +)] +#[derive(Clone)] +pub struct PyDataFrame { + df: Arc, + + // In IPython environment cache batches between __repr__ and _repr_html_ calls. + batches: SharedCachedBatches, +} + +impl PyDataFrame { + /// creates a new PyDataFrame + pub fn new(df: DataFrame) -> Self { + Self { + df: Arc::new(df), + batches: Arc::new(Mutex::new(None)), + } + } + + /// Return a clone of the inner Arc for crate-local callers. + pub(crate) fn inner_df(&self) -> Arc { + Arc::clone(&self.df) + } + + fn prepare_repr_string<'py>( + &self, + py: Python<'py>, + as_html: bool, + ) -> PyDataFusionResult { + // Get the Python formatter and config + let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?; + + let is_ipython = *is_ipython_env(py); + + let (cached_batches, should_cache) = { + let mut cache = self.batches.lock(); + let should_cache = is_ipython && cache.is_none(); + let batches = cache.take(); + (batches, should_cache) + }; + + let (batches, has_more) = match cached_batches { + Some(b) => b, + None => wait_for_future( + py, + collect_record_batches_to_display(self.df.as_ref().clone(), config), + )??, + }; + + if batches.is_empty() { + // This should not be reached, but do it for safety since we index into the vector below + return Ok("No data to display".to_string()); + } + + let table_uuid = uuid::Uuid::new_v4().to_string(); + + // Convert record batches to Py list + let py_batches = batches + .iter() + .map(|rb| rb.to_pyarrow(py)) + .collect::>>>()?; + + let py_schema = self.schema().into_pyobject(py)?; + + let kwargs = pyo3::types::PyDict::new(py); + let py_batches_list = PyList::new(py, py_batches.as_slice())?; + kwargs.set_item("batches", py_batches_list)?; + kwargs.set_item("schema", py_schema)?; + kwargs.set_item("has_more", has_more)?; + kwargs.set_item("table_uuid", table_uuid)?; + + let method_name = match as_html { + true => "format_html", + false => "format_str", + }; + + let html_result = formatter.call_method(method_name, (), Some(&kwargs))?; + let html_str: String = html_result.extract()?; + + if should_cache { + let mut cache = self.batches.lock(); + *cache = Some((batches.clone(), has_more)); + } + + Ok(html_str) + } + + async fn collect_column_inner(&self, column: &str) -> Result { + let batches = self + .df + .as_ref() + .clone() + .select_columns(&[column])? + .collect() + .await?; + + let arrays = batches + .iter() + .map(|b| b.column(0).as_ref()) + .collect::>(); + + arrow_select::concat::concat(&arrays).map_err(Into::into) + } +} + +/// Synchronous wrapper around partitioned [`SendableRecordBatchStream`]s used +/// for the `__arrow_c_stream__` implementation. +/// +/// It drains each partition's stream sequentially, yielding record batches in +/// their original partition order. When a `projection` is set, each batch is +/// converted via `record_batch_into_schema` to apply schema changes per batch. +struct PartitionedDataFrameStreamReader { + streams: Vec, + schema: SchemaRef, + projection: Option, + current: usize, +} + +impl Iterator for PartitionedDataFrameStreamReader { + type Item = Result; + + fn next(&mut self) -> Option { + while self.current < self.streams.len() { + let stream = &mut self.streams[self.current]; + let fut = poll_next_batch(stream); + let result = Python::attach(|py| wait_for_future(py, fut)); + + match result { + Ok(Ok(Some(batch))) => { + let batch = if let Some(ref schema) = self.projection { + match record_batch_into_schema(batch, schema.as_ref()) { + Ok(b) => b, + Err(e) => return Some(Err(e)), + } + } else { + batch + }; + return Some(Ok(batch)); + } + Ok(Ok(None)) => { + self.current += 1; + continue; + } + Ok(Err(e)) => { + return Some(Err(ArrowError::ExternalError(Box::new(e)))); + } + Err(e) => { + return Some(Err(ArrowError::ExternalError(Box::new(e)))); + } + } + } + + None + } +} + +impl RecordBatchReader for PartitionedDataFrameStreamReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +#[pymethods] +impl PyDataFrame { + /// Enable selection for `df[col]`, `df[col1, col2, col3]`, and `df[[col1, col2, col3]]` + fn __getitem__(&self, key: Bound<'_, PyAny>) -> PyDataFusionResult { + if let Ok(key) = key.extract::() { + // df[col] + self.select_columns(vec![key]) + } else if let Ok(tuple) = key.cast::() { + // df[col1, col2, col3] + let keys = tuple + .iter() + .map(|item| item.extract::()) + .collect::>>()?; + self.select_columns(keys) + } else if let Ok(keys) = key.extract::>() { + // df[[col1, col2, col3]] + self.select_columns(keys) + } else { + let message = "DataFrame can only be indexed by string index or indices".to_string(); + Err(PyDataFusionError::Common(message)) + } + } + + fn __repr__(&self, py: Python) -> PyDataFusionResult { + self.prepare_repr_string(py, false) + } + + #[staticmethod] + #[expect(unused_variables)] + fn default_str_repr<'py>( + batches: Vec>, + schema: &Bound<'py, PyAny>, + has_more: bool, + table_uuid: &str, + ) -> PyResult { + let batches = batches + .into_iter() + .map(|batch| RecordBatch::from_pyarrow_bound(&batch)) + .collect::>>()? + .into_iter() + .filter(|batch| batch.num_rows() > 0) + .collect::>(); + + if batches.is_empty() { + return Ok("No data to display".to_owned()); + } + + let batches_as_displ = + pretty::pretty_format_batches(&batches).map_err(py_datafusion_err)?; + + let additional_str = match has_more { + true => "\nData truncated.", + false => "", + }; + + Ok(format!("DataFrame()\n{batches_as_displ}{additional_str}")) + } + + fn _repr_html_(&self, py: Python) -> PyDataFusionResult { + self.prepare_repr_string(py, true) + } + + /// Calculate summary statistics for a DataFrame + fn describe(&self, py: Python) -> PyDataFusionResult { + let df = self.df.as_ref().clone(); + let stat_df = wait_for_future(py, df.describe())??; + Ok(Self::new(stat_df)) + } + + /// Returns the schema from the logical plan + fn schema(&self) -> PyArrowType { + PyArrowType(self.df.schema().as_arrow().clone()) + } + + /// Convert this DataFrame into a Table Provider that can be used in register_table + /// By convention, into_... methods consume self and return the new object. + /// Disabling the clippy lint, so we can use &self + /// because we're working with Python bindings + /// where objects are shared + #[allow(clippy::wrong_self_convention)] + pub fn into_view(&self, temporary: bool) -> PyDataFusionResult { + let table_provider = if temporary { + Arc::new(TempViewTable::new(Arc::clone(&self.df))) as Arc + } else { + // Call the underlying Rust DataFrame::into_view method. + // Note that the Rust method consumes self; here we clone the inner Arc + // so that we don't invalidate this PyDataFrame. + self.df.as_ref().clone().into_view() + }; + Ok(PyTable::from(table_provider)) + } + + #[pyo3(signature = (*args))] + fn select_columns(&self, args: Vec) -> PyDataFusionResult { + let args = args.iter().map(|s| s.as_ref()).collect::>(); + let df = self.df.as_ref().clone().select_columns(&args)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (*args))] + fn select_exprs(&self, args: Vec) -> PyDataFusionResult { + let args = args.iter().map(|s| s.as_ref()).collect::>(); + let df = self.df.as_ref().clone().select_exprs(&args)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (*args))] + fn select(&self, args: Vec) -> PyDataFusionResult { + let expr: Vec = args.into_iter().map(|e| e.into()).collect(); + let df = self.df.as_ref().clone().select(expr)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (*args))] + fn drop(&self, args: Vec) -> PyDataFusionResult { + let cols = args.iter().map(|s| s.as_ref()).collect::>(); + let df = self.df.as_ref().clone().drop_columns(&cols)?; + Ok(Self::new(df)) + } + + fn filter(&self, predicate: PyExpr) -> PyDataFusionResult { + let df = self.df.as_ref().clone().filter(predicate.into())?; + Ok(Self::new(df)) + } + + fn parse_sql_expr(&self, expr: PyBackedStr) -> PyDataFusionResult { + self.df + .as_ref() + .parse_sql_expr(&expr) + .map(PyExpr::from) + .map_err(PyDataFusionError::from) + } + + fn with_column(&self, name: &str, expr: PyExpr) -> PyDataFusionResult { + let df = self.df.as_ref().clone().with_column(name, expr.into())?; + Ok(Self::new(df)) + } + + fn with_columns(&self, exprs: Vec) -> PyDataFusionResult { + let mut df = self.df.as_ref().clone(); + for expr in exprs { + let expr: Expr = expr.into(); + let name = format!("{}", expr.schema_name()); + df = df.with_column(name.as_str(), expr)? + } + Ok(Self::new(df)) + } + + /// Rename one column by applying a new projection. This is a no-op if the column to be + /// renamed does not exist. + fn with_column_renamed(&self, old_name: &str, new_name: &str) -> PyDataFusionResult { + let df = self + .df + .as_ref() + .clone() + .with_column_renamed(old_name, new_name)?; + Ok(Self::new(df)) + } + + fn aggregate(&self, group_by: Vec, aggs: Vec) -> PyDataFusionResult { + let group_by = group_by.into_iter().map(|e| e.into()).collect(); + let aggs = aggs.into_iter().map(|e| e.into()).collect(); + let df = self.df.as_ref().clone().aggregate(group_by, aggs)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (*exprs))] + fn sort(&self, exprs: Vec) -> PyDataFusionResult { + let exprs = to_sort_expressions(exprs); + let df = self.df.as_ref().clone().sort(exprs)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (count, offset=0))] + fn limit(&self, count: usize, offset: usize) -> PyDataFusionResult { + let df = self.df.as_ref().clone().limit(offset, Some(count))?; + Ok(Self::new(df)) + } + + /// Executes the plan, returning a list of `RecordBatch`es. + /// Unless some order is specified in the plan, there is no + /// guarantee of the order of the result. + fn collect<'py>(&self, py: Python<'py>) -> PyResult>> { + let batches = wait_for_future(py, self.df.as_ref().clone().collect())? + .map_err(PyDataFusionError::from)?; + // cannot use PyResult> return type due to + // https://github.com/PyO3/pyo3/issues/1813 + batches.into_iter().map(|rb| rb.to_pyarrow(py)).collect() + } + + /// Cache DataFrame. + fn cache(&self, py: Python) -> PyDataFusionResult { + let df = wait_for_future(py, self.df.as_ref().clone().cache())??; + Ok(Self::new(df)) + } + + /// Executes this DataFrame and collects all results into a vector of vector of RecordBatch + /// maintaining the input partitioning. + fn collect_partitioned<'py>(&self, py: Python<'py>) -> PyResult>>> { + let batches = wait_for_future(py, self.df.as_ref().clone().collect_partitioned())? + .map_err(PyDataFusionError::from)?; + + batches + .into_iter() + .map(|rbs| rbs.into_iter().map(|rb| rb.to_pyarrow(py)).collect()) + .collect() + } + + fn collect_column<'py>(&self, py: Python<'py>, column: &str) -> PyResult> { + wait_for_future(py, self.collect_column_inner(column))? + .map_err(PyDataFusionError::from)? + .to_data() + .to_pyarrow(py) + } + + /// Print the result, 20 lines by default + #[pyo3(signature = (num=20))] + fn show(&self, py: Python, num: usize) -> PyDataFusionResult<()> { + let df = self.df.as_ref().clone().limit(0, Some(num))?; + print_dataframe(py, df) + } + + /// Filter out duplicate rows + fn distinct(&self) -> PyDataFusionResult { + let df = self.df.as_ref().clone().distinct()?; + Ok(Self::new(df)) + } + + fn join( + &self, + right: PyDataFrame, + how: &str, + left_on: Vec, + right_on: Vec, + coalesce_keys: bool, + ) -> PyDataFusionResult { + let join_type = match how { + "inner" => JoinType::Inner, + "left" => JoinType::Left, + "right" => JoinType::Right, + "full" => JoinType::Full, + "semi" => JoinType::LeftSemi, + "anti" => JoinType::LeftAnti, + how => { + return Err(PyDataFusionError::Common(format!( + "The join type {how} does not exist or is not implemented" + ))); + } + }; + + let left_keys = left_on.iter().map(|s| s.as_ref()).collect::>(); + let right_keys = right_on.iter().map(|s| s.as_ref()).collect::>(); + + let mut df = self.df.as_ref().clone().join( + right.df.as_ref().clone(), + join_type, + &left_keys, + &right_keys, + None, + )?; + + if coalesce_keys { + let mutual_keys = left_keys + .iter() + .zip(right_keys.iter()) + .filter(|(l, r)| l == r) + .map(|(key, _)| *key) + .collect::>(); + + let fields_to_coalesce = mutual_keys + .iter() + .map(|name| { + let qualified_fields = df + .logical_plan() + .schema() + .qualified_fields_with_unqualified_name(name); + (*name, qualified_fields) + }) + .filter(|(_, fields)| fields.len() == 2) + .collect::>(); + + let expr: Vec = df + .logical_plan() + .schema() + .fields() + .into_iter() + .enumerate() + .map(|(idx, _)| df.logical_plan().schema().qualified_field(idx)) + .filter_map(|(qualifier, field)| { + if let Some((key_name, qualified_fields)) = fields_to_coalesce + .iter() + .find(|(_, qf)| qf.contains(&(qualifier, field))) + { + // Only add the coalesce expression once (when we encounter the first field) + // Skip the second field (it's already included in to coalesce) + if (qualifier, field) == qualified_fields[0] { + let left_col = Expr::Column(Column::from(qualified_fields[0])); + let right_col = Expr::Column(Column::from(qualified_fields[1])); + return Some(coalesce(vec![left_col, right_col]).alias(*key_name)); + } + None + } else { + Some(Expr::Column(Column::from((qualifier, field)))) + } + }) + .collect(); + df = df.select(expr)?; + } + + Ok(Self::new(df)) + } + + fn join_on( + &self, + right: PyDataFrame, + on_exprs: Vec, + how: &str, + ) -> PyDataFusionResult { + let join_type = match how { + "inner" => JoinType::Inner, + "left" => JoinType::Left, + "right" => JoinType::Right, + "full" => JoinType::Full, + "semi" => JoinType::LeftSemi, + "anti" => JoinType::LeftAnti, + how => { + return Err(PyDataFusionError::Common(format!( + "The join type {how} does not exist or is not implemented" + ))); + } + }; + let exprs: Vec = on_exprs.into_iter().map(|e| e.into()).collect(); + + let df = self + .df + .as_ref() + .clone() + .join_on(right.df.as_ref().clone(), join_type, exprs)?; + Ok(Self::new(df)) + } + + /// Print the query plan + #[pyo3(signature = (verbose=false, analyze=false))] + fn explain(&self, py: Python, verbose: bool, analyze: bool) -> PyDataFusionResult<()> { + let df = self.df.as_ref().clone().explain(verbose, analyze)?; + print_dataframe(py, df) + } + + /// Get the logical plan for this `DataFrame` + fn logical_plan(&self) -> PyResult { + Ok(self.df.as_ref().clone().logical_plan().clone().into()) + } + + /// Get the optimized logical plan for this `DataFrame` + fn optimized_logical_plan(&self) -> PyDataFusionResult { + Ok(self.df.as_ref().clone().into_optimized_plan()?.into()) + } + + /// Get the execution plan for this `DataFrame` + fn execution_plan(&self, py: Python) -> PyDataFusionResult { + let plan = wait_for_future(py, self.df.as_ref().clone().create_physical_plan())??; + Ok(plan.into()) + } + + /// Repartition a `DataFrame` based on a logical partitioning scheme. + fn repartition(&self, num: usize) -> PyDataFusionResult { + let new_df = self + .df + .as_ref() + .clone() + .repartition(Partitioning::RoundRobinBatch(num))?; + Ok(Self::new(new_df)) + } + + /// Repartition a `DataFrame` based on a logical partitioning scheme. + #[pyo3(signature = (*args, num))] + fn repartition_by_hash(&self, args: Vec, num: usize) -> PyDataFusionResult { + let expr = args.into_iter().map(|py_expr| py_expr.into()).collect(); + let new_df = self + .df + .as_ref() + .clone() + .repartition(Partitioning::Hash(expr, num))?; + Ok(Self::new(new_df)) + } + + /// Calculate the union of two `DataFrame`s, preserving duplicate rows.The + /// two `DataFrame`s must have exactly the same schema + #[pyo3(signature = (py_df, distinct=false))] + fn union(&self, py_df: PyDataFrame, distinct: bool) -> PyDataFusionResult { + let new_df = if distinct { + self.df + .as_ref() + .clone() + .union_distinct(py_df.df.as_ref().clone())? + } else { + self.df.as_ref().clone().union(py_df.df.as_ref().clone())? + }; + + Ok(Self::new(new_df)) + } + + /// Calculate the distinct union of two `DataFrame`s. The + /// two `DataFrame`s must have exactly the same schema + fn union_distinct(&self, py_df: PyDataFrame) -> PyDataFusionResult { + let new_df = self + .df + .as_ref() + .clone() + .union_distinct(py_df.df.as_ref().clone())?; + Ok(Self::new(new_df)) + } + + #[pyo3(signature = (column, preserve_nulls=true))] + fn unnest_column(&self, column: &str, preserve_nulls: bool) -> PyDataFusionResult { + // TODO: expose RecursionUnnestOptions + // REF: https://github.com/apache/datafusion/pull/11577 + let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls); + let df = self + .df + .as_ref() + .clone() + .unnest_columns_with_options(&[column], unnest_options)?; + Ok(Self::new(df)) + } + + #[pyo3(signature = (columns, preserve_nulls=true))] + fn unnest_columns( + &self, + columns: Vec, + preserve_nulls: bool, + ) -> PyDataFusionResult { + // TODO: expose RecursionUnnestOptions + // REF: https://github.com/apache/datafusion/pull/11577 + let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls); + let cols = columns.iter().map(|s| s.as_ref()).collect::>(); + let df = self + .df + .as_ref() + .clone() + .unnest_columns_with_options(&cols, unnest_options)?; + Ok(Self::new(df)) + } + + /// Calculate the intersection of two `DataFrame`s. The two `DataFrame`s must have exactly the same schema + fn intersect(&self, py_df: PyDataFrame) -> PyDataFusionResult { + let new_df = self + .df + .as_ref() + .clone() + .intersect(py_df.df.as_ref().clone())?; + Ok(Self::new(new_df)) + } + + /// Calculate the exception of two `DataFrame`s. The two `DataFrame`s must have exactly the same schema + fn except_all(&self, py_df: PyDataFrame) -> PyDataFusionResult { + let new_df = self.df.as_ref().clone().except(py_df.df.as_ref().clone())?; + Ok(Self::new(new_df)) + } + + /// Write a `DataFrame` to a CSV file. + fn write_csv( + &self, + py: Python, + path: &str, + with_header: bool, + write_options: Option, + ) -> PyDataFusionResult<()> { + let csv_options = CsvOptions { + has_header: Some(with_header), + ..Default::default() + }; + let write_options = write_options + .map(DataFrameWriteOptions::from) + .unwrap_or_default(); + + wait_for_future( + py, + self.df + .as_ref() + .clone() + .write_csv(path, write_options, Some(csv_options)), + )??; + Ok(()) + } + + /// Write a `DataFrame` to a Parquet file. + #[pyo3(signature = ( + path, + compression="zstd", + compression_level=None, + write_options=None, + ))] + fn write_parquet( + &self, + path: &str, + compression: &str, + compression_level: Option, + write_options: Option, + py: Python, + ) -> PyDataFusionResult<()> { + fn verify_compression_level(cl: Option) -> Result { + cl.ok_or(PyValueError::new_err("compression_level is not defined")) + } + + let _validated = match compression.to_lowercase().as_str() { + "snappy" => Compression::SNAPPY, + "gzip" => Compression::GZIP( + GzipLevel::try_new(compression_level.unwrap_or(6)) + .map_err(|e| PyValueError::new_err(format!("{e}")))?, + ), + "brotli" => Compression::BROTLI( + BrotliLevel::try_new(verify_compression_level(compression_level)?) + .map_err(|e| PyValueError::new_err(format!("{e}")))?, + ), + "zstd" => Compression::ZSTD( + ZstdLevel::try_new(verify_compression_level(compression_level)? as i32) + .map_err(|e| PyValueError::new_err(format!("{e}")))?, + ), + "lzo" => Compression::LZO, + "lz4" => Compression::LZ4, + "lz4_raw" => Compression::LZ4_RAW, + "uncompressed" => Compression::UNCOMPRESSED, + _ => { + return Err(PyDataFusionError::Common(format!( + "Unrecognized compression type {compression}" + ))); + } + }; + + let mut compression_string = compression.to_string(); + if let Some(level) = compression_level { + compression_string.push_str(&format!("({level})")); + } + + let mut options = TableParquetOptions::default(); + options.global.compression = Some(compression_string); + let write_options = write_options + .map(DataFrameWriteOptions::from) + .unwrap_or_default(); + + wait_for_future( + py, + self.df + .as_ref() + .clone() + .write_parquet(path, write_options, Option::from(options)), + )??; + Ok(()) + } + + /// Write a `DataFrame` to a Parquet file, using advanced options. + fn write_parquet_with_options( + &self, + path: &str, + options: PyParquetWriterOptions, + column_specific_options: HashMap, + write_options: Option, + py: Python, + ) -> PyDataFusionResult<()> { + let table_options = TableParquetOptions { + global: options.options, + column_specific_options: column_specific_options + .into_iter() + .map(|(k, v)| (k, v.options)) + .collect(), + ..Default::default() + }; + let write_options = write_options + .map(DataFrameWriteOptions::from) + .unwrap_or_default(); + wait_for_future( + py, + self.df.as_ref().clone().write_parquet( + path, + write_options, + Option::from(table_options), + ), + )??; + Ok(()) + } + + /// Executes a query and writes the results to a partitioned JSON file. + fn write_json( + &self, + path: &str, + py: Python, + write_options: Option, + ) -> PyDataFusionResult<()> { + let write_options = write_options + .map(DataFrameWriteOptions::from) + .unwrap_or_default(); + wait_for_future( + py, + self.df + .as_ref() + .clone() + .write_json(path, write_options, None), + )??; + Ok(()) + } + + fn write_table( + &self, + py: Python, + table_name: &str, + write_options: Option, + ) -> PyDataFusionResult<()> { + let write_options = write_options + .map(DataFrameWriteOptions::from) + .unwrap_or_default(); + wait_for_future( + py, + self.df + .as_ref() + .clone() + .write_table(table_name, write_options), + )??; + Ok(()) + } + + /// Convert to Arrow Table + /// Collect the batches and pass to Arrow Table + fn to_arrow_table(&self, py: Python<'_>) -> PyResult> { + let batches = self.collect(py)?.into_pyobject(py)?; + + // only use the DataFrame's schema if there are no batches, otherwise let the schema be + // determined from the batches (avoids some inconsistencies with nullable columns) + let args = if batches.len()? == 0 { + let schema = self.schema().into_pyobject(py)?; + PyTuple::new(py, &[batches, schema])? + } else { + PyTuple::new(py, &[batches])? + }; + + // Instantiate pyarrow Table object and use its from_batches method + let table_class = py.import("pyarrow")?.getattr("Table")?; + let table: Py = table_class.call_method1("from_batches", args)?.into(); + Ok(table) + } + + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &'py self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyDataFusionResult> { + let df = self.df.as_ref().clone(); + let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?; + + let mut schema: Schema = self.df.schema().to_owned().as_arrow().clone(); + let mut projection: Option = None; + + if let Some(schema_capsule) = requested_schema { + validate_pycapsule(&schema_capsule, "arrow_schema")?; + + let data: NonNull = schema_capsule + .pointer_checked(Some(c_str!("arrow_schema")))? + .cast(); + let schema_ptr = unsafe { data.as_ref() }; + let desired_schema = Schema::try_from(schema_ptr)?; + + schema = project_schema(schema, desired_schema)?; + projection = Some(Arc::new(schema.clone())); + } + + let schema_ref = Arc::new(schema.clone()); + + let reader = PartitionedDataFrameStreamReader { + streams, + schema: schema_ref, + projection, + current: 0, + }; + let reader: Box = Box::new(reader); + + // Create the Arrow stream and wrap it in a PyCapsule. The default + // destructor provided by PyO3 will drop the stream unless ownership is + // transferred to PyArrow during import. + let stream = FFI_ArrowArrayStream::new(reader); + let name = CString::new(ARROW_ARRAY_STREAM_NAME.to_bytes()).unwrap(); + let capsule = PyCapsule::new(py, stream, Some(name))?; + Ok(capsule) + } + + fn execute_stream(&self, py: Python) -> PyDataFusionResult { + let df = self.df.as_ref().clone(); + let stream = spawn_future(py, async move { df.execute_stream().await })?; + Ok(PyRecordBatchStream::new(stream)) + } + + fn execute_stream_partitioned(&self, py: Python) -> PyResult> { + let df = self.df.as_ref().clone(); + let streams = spawn_future(py, async move { df.execute_stream_partitioned().await })?; + Ok(streams.into_iter().map(PyRecordBatchStream::new).collect()) + } + + /// Convert to pandas dataframe with pyarrow + /// Collect the batches, pass to Arrow Table & then convert to Pandas DataFrame + fn to_pandas(&self, py: Python<'_>) -> PyResult> { + let table = self.to_arrow_table(py)?; + + // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pandas + let result = table.call_method0(py, "to_pandas")?; + Ok(result) + } + + /// Convert to Python list using pyarrow + /// Each list item represents one row encoded as dictionary + fn to_pylist(&self, py: Python<'_>) -> PyResult> { + let table = self.to_arrow_table(py)?; + + // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pylist + let result = table.call_method0(py, "to_pylist")?; + Ok(result) + } + + /// Convert to Python dictionary using pyarrow + /// Each dictionary key is a column and the dictionary value represents the column values + fn to_pydict(&self, py: Python) -> PyResult> { + let table = self.to_arrow_table(py)?; + + // See also: https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.to_pydict + let result = table.call_method0(py, "to_pydict")?; + Ok(result) + } + + /// Convert to polars dataframe with pyarrow + /// Collect the batches, pass to Arrow Table & then convert to polars DataFrame + fn to_polars(&self, py: Python<'_>) -> PyResult> { + let table = self.to_arrow_table(py)?; + let dataframe = py.import("polars")?.getattr("DataFrame")?; + let args = PyTuple::new(py, &[table])?; + let result: Py = dataframe.call1(args)?.into(); + Ok(result) + } + + // Executes this DataFrame to get the total number of rows. + fn count(&self, py: Python) -> PyDataFusionResult { + Ok(wait_for_future(py, self.df.as_ref().clone().count())??) + } + + /// Fill null values with a specified value for specific columns + #[pyo3(signature = (value, columns=None))] + fn fill_null( + &self, + value: Py, + columns: Option>, + py: Python, + ) -> PyDataFusionResult { + let scalar_value: PyScalarValue = value.extract(py)?; + + let cols = match columns { + Some(col_names) => col_names.iter().map(|c| c.to_string()).collect(), + None => Vec::new(), // Empty vector means fill null for all columns + }; + + let df = self.df.as_ref().clone().fill_null(scalar_value.0, cols)?; + Ok(Self::new(df)) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "InsertOp", + module = "datafusion" +)] +pub enum PyInsertOp { + APPEND, + REPLACE, + OVERWRITE, +} + +impl From for InsertOp { + fn from(value: PyInsertOp) -> Self { + match value { + PyInsertOp::APPEND => InsertOp::Append, + PyInsertOp::REPLACE => InsertOp::Replace, + PyInsertOp::OVERWRITE => InsertOp::Overwrite, + } + } +} + +#[derive(Debug, Clone)] +#[pyclass( + from_py_object, + frozen, + name = "DataFrameWriteOptions", + module = "datafusion" +)] +pub struct PyDataFrameWriteOptions { + insert_operation: InsertOp, + single_file_output: bool, + partition_by: Vec, + sort_by: Vec, +} + +impl From for DataFrameWriteOptions { + fn from(value: PyDataFrameWriteOptions) -> Self { + DataFrameWriteOptions::new() + .with_insert_operation(value.insert_operation) + .with_single_file_output(value.single_file_output) + .with_partition_by(value.partition_by) + .with_sort_by(value.sort_by) + } +} + +#[pymethods] +impl PyDataFrameWriteOptions { + #[new] + fn new( + insert_operation: Option, + single_file_output: bool, + partition_by: Option>, + sort_by: Option>, + ) -> Self { + let insert_operation = insert_operation.map(Into::into).unwrap_or(InsertOp::Append); + let sort_by = sort_by + .unwrap_or_default() + .into_iter() + .map(Into::into) + .collect(); + Self { + insert_operation, + single_file_output, + partition_by: partition_by.unwrap_or_default(), + sort_by, + } + } +} + +/// Print DataFrame +fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> { + // Get string representation of record batches + let batches = wait_for_future(py, df.collect())??; + let result = if batches.is_empty() { + "DataFrame has no rows".to_string() + } else { + match pretty::pretty_format_batches(&batches) { + Ok(batch) => format!("DataFrame()\n{batch}"), + Err(err) => format!("Error: {:?}", err.to_string()), + } + }; + + // Import the Python 'builtins' module to access the print function + // Note that println! does not print to the Python debug console and is not visible in notebooks for instance + let print = py.import("builtins")?.getattr("print")?; + print.call1((result,))?; + Ok(()) +} + +fn project_schema(from_schema: Schema, to_schema: Schema) -> Result { + let merged_schema = Schema::try_merge(vec![from_schema, to_schema.clone()])?; + + let project_indices: Vec = to_schema + .fields + .iter() + .map(|field| field.name()) + .filter_map(|field_name| merged_schema.index_of(field_name).ok()) + .collect(); + + merged_schema.project(&project_indices) +} +// NOTE: `arrow::compute::cast` in combination with `RecordBatch::try_select` or +// DataFusion's `schema::cast_record_batch` do not fully cover the required +// transformations here. They will not create missing columns and may insert +// nulls for non-nullable fields without erroring. To maintain current behavior +// we perform the casting and null checks manually. +fn record_batch_into_schema( + record_batch: RecordBatch, + schema: &Schema, +) -> Result { + let schema = Arc::new(schema.clone()); + let base_schema = record_batch.schema(); + if base_schema.fields().is_empty() { + // Nothing to project + return Ok(RecordBatch::new_empty(schema)); + } + + let array_size = record_batch.column(0).len(); + let mut data_arrays = Vec::with_capacity(schema.fields().len()); + + for field in schema.fields() { + let desired_data_type = field.data_type(); + if let Some(original_data) = record_batch.column_by_name(field.name()) { + let original_data_type = original_data.data_type(); + + if can_cast_types(original_data_type, desired_data_type) { + data_arrays.push(arrow::compute::kernels::cast( + original_data, + desired_data_type, + )?); + } else if field.is_nullable() { + data_arrays.push(new_null_array(desired_data_type, array_size)); + } else { + return Err(ArrowError::CastError(format!( + "Attempting to cast to non-nullable and non-castable field {} during schema projection.", + field.name() + ))); + } + } else { + if !field.is_nullable() { + return Err(ArrowError::CastError(format!( + "Attempting to set null to non-nullable field {} during schema projection.", + field.name() + ))); + } + data_arrays.push(new_null_array(desired_data_type, array_size)); + } + } + + RecordBatch::try_new(schema, data_arrays) +} + +/// This is a helper function to return the first non-empty record batch from executing a DataFrame. +/// It additionally returns a bool, which indicates if there are more record batches available. +/// We do this so we can determine if we should indicate to the user that the data has been +/// truncated. This collects until we have archived both of these two conditions +/// +/// - We have collected our minimum number of rows +/// - We have reached our limit, either data size or maximum number of rows +/// +/// Otherwise it will return when the stream has exhausted. If you want a specific number of +/// rows, set min_rows == max_rows. +async fn collect_record_batches_to_display( + df: DataFrame, + config: FormatterConfig, +) -> Result<(Vec, bool), DataFusionError> { + let FormatterConfig { + max_bytes, + min_rows, + max_rows, + } = config; + + let partitioned_stream = df.execute_stream_partitioned().await?; + let mut stream = futures::stream::iter(partitioned_stream).flatten(); + let mut size_estimate_so_far = 0; + let mut rows_so_far = 0; + let mut record_batches = Vec::default(); + let mut has_more = false; + + // Collect rows until we hit a limit (memory or max_rows) OR reach the guaranteed minimum. + // The minimum rows constraint overrides both memory and row limits to ensure a baseline + // of data is always displayed, even if it temporarily exceeds those limits. + // This provides better UX by guaranteeing users see at least min_rows rows. + while (size_estimate_so_far < max_bytes && rows_so_far < max_rows) || rows_so_far < min_rows { + let mut rb = match stream.next().await { + None => { + break; + } + Some(Ok(r)) => r, + Some(Err(e)) => return Err(e), + }; + + let mut rows_in_rb = rb.num_rows(); + if rows_in_rb > 0 { + size_estimate_so_far += rb.get_array_memory_size(); + + // When memory limit is exceeded, scale back row count proportionally to stay within budget + if size_estimate_so_far > max_bytes { + let ratio = max_bytes as f32 / size_estimate_so_far as f32; + let total_rows = rows_in_rb + rows_so_far; + + // Calculate reduced rows maintaining the memory/data proportion + let mut reduced_row_num = (total_rows as f32 * ratio).round() as usize; + // Ensure we always respect the minimum rows guarantee + if reduced_row_num < min_rows { + reduced_row_num = min_rows.min(total_rows); + } + + let limited_rows_this_rb = reduced_row_num - rows_so_far; + if limited_rows_this_rb < rows_in_rb { + rows_in_rb = limited_rows_this_rb; + rb = rb.slice(0, limited_rows_this_rb); + has_more = true; + } + } + + if rows_in_rb + rows_so_far > max_rows { + rb = rb.slice(0, max_rows - rows_so_far); + has_more = true; + } + + rows_so_far += rb.num_rows(); + record_batches.push(rb); + } + } + + if record_batches.is_empty() { + return Ok((Vec::default(), false)); + } + + if !has_more { + // Data was not already truncated, so check to see if more record batches remain + has_more = match stream.try_next().await { + Ok(None) => false, // reached end + Ok(Some(_)) => true, + Err(_) => false, // Stream disconnected + }; + } + + Ok((record_batches, has_more)) +} diff --git a/src/dataset.rs b/crates/core/src/dataset.rs similarity index 95% rename from src/dataset.rs rename to crates/core/src/dataset.rs index 0baf4da2a..dbeafcd9f 100644 --- a/src/dataset.rs +++ b/crates/core/src/dataset.rs @@ -15,25 +15,22 @@ // specific language governing permissions and limitations // under the License. -use datafusion::catalog::Session; -use pyo3::exceptions::PyValueError; -/// Implements a Datafusion TableProvider that delegates to a PyArrow Dataset -/// This allows us to use PyArrow Datasets as Datafusion tables while pushing down projections and filters -use pyo3::prelude::*; -use pyo3::types::PyType; - use std::any::Any; use std::sync::Arc; use async_trait::async_trait; - use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::pyarrow::PyArrowType; +use datafusion::catalog::Session; use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::{DataFusionError, Result as DFResult}; -use datafusion::logical_expr::Expr; -use datafusion::logical_expr::TableProviderFilterPushDown; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; +use pyo3::exceptions::PyValueError; +/// Implements a Datafusion TableProvider that delegates to a PyArrow Dataset +/// This allows us to use PyArrow Datasets as Datafusion tables while pushing down projections and filters +use pyo3::prelude::*; +use pyo3::types::PyType; use crate::dataset_exec::DatasetExec; use crate::pyarrow_filter_expression::PyArrowFilterExpression; @@ -41,7 +38,7 @@ use crate::pyarrow_filter_expression::PyArrowFilterExpression; // Wraps a pyarrow.dataset.Dataset class and implements a Datafusion TableProvider around it #[derive(Debug)] pub(crate) struct Dataset { - dataset: PyObject, + dataset: Py, } impl Dataset { @@ -50,7 +47,7 @@ impl Dataset { // Ensure that we were passed an instance of pyarrow.dataset.Dataset let ds = PyModule::import(py, "pyarrow.dataset")?; let ds_attr = ds.getattr("Dataset")?; - let ds_type = ds_attr.downcast::()?; + let ds_type = ds_attr.cast::()?; if dataset.is_instance(ds_type)? { Ok(Dataset { dataset: dataset.clone().unbind(), @@ -73,7 +70,7 @@ impl TableProvider for Dataset { /// Get a reference to the schema for this table fn schema(&self) -> SchemaRef { - Python::with_gil(|py| { + Python::attach(|py| { let dataset = self.dataset.bind(py); // This can panic but since we checked that self.dataset is a pyarrow.dataset.Dataset it should never Arc::new( @@ -107,7 +104,7 @@ impl TableProvider for Dataset { // The datasource should return *at least* this number of rows if available. _limit: Option, ) -> DFResult> { - Python::with_gil(|py| { + Python::attach(|py| { let plan: Arc = Arc::new( DatasetExec::new(py, self.dataset.bind(py), projection.cloned(), filters) .map_err(|err| DataFusionError::External(Box::new(err)))?, diff --git a/src/dataset_exec.rs b/crates/core/src/dataset_exec.rs similarity index 92% rename from src/dataset_exec.rs rename to crates/core/src/dataset_exec.rs index 445e4fe74..e3c058c07 100644 --- a/src/dataset_exec.rs +++ b/crates/core/src/dataset_exec.rs @@ -15,32 +15,29 @@ // specific language governing permissions and limitations // under the License. -use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; -/// Implements a Datafusion physical ExecutionPlan that delegates to a PyArrow Dataset -/// This actually performs the projection, filtering and scanning of a Dataset -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyIterator, PyList}; - use std::any::Any; use std::sync::Arc; -use futures::{stream, TryStreamExt}; - use datafusion::arrow::datatypes::SchemaRef; -use datafusion::arrow::error::ArrowError; -use datafusion::arrow::error::Result as ArrowResult; +use datafusion::arrow::error::{ArrowError, Result as ArrowResult}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; use datafusion::error::{DataFusionError as InnerDataFusionError, Result as DFResult}; use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::Expr; +use datafusion::logical_expr::utils::conjunction; use datafusion::physical_expr::{EquivalenceProperties, LexOrdering}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - SendableRecordBatchStream, Statistics, + PlanProperties, SendableRecordBatchStream, Statistics, }; +use futures::{TryStreamExt, stream}; +/// Implements a Datafusion physical ExecutionPlan that delegates to a PyArrow Dataset +/// This actually performs the projection, filtering and scanning of a Dataset +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyIterator, PyList}; use crate::errors::PyDataFusionResult; use crate::pyarrow_filter_expression::PyArrowFilterExpression; @@ -53,7 +50,7 @@ impl Iterator for PyArrowBatchesAdapter { type Item = ArrowResult; fn next(&mut self) -> Option { - Python::with_gil(|py| { + Python::attach(|py| { let mut batches = self.batches.clone_ref(py).into_bound(py); Some( batches @@ -68,13 +65,13 @@ impl Iterator for PyArrowBatchesAdapter { // Wraps a pyarrow.dataset.Dataset class and implements a Datafusion ExecutionPlan around it #[derive(Debug)] pub(crate) struct DatasetExec { - dataset: PyObject, + dataset: Py, schema: SchemaRef, fragments: Py, columns: Option>, - filter_expr: Option, + filter_expr: Option>, projected_statistics: Statistics, - plan_properties: datafusion::physical_plan::PlanProperties, + plan_properties: Arc, } impl DatasetExec { @@ -97,7 +94,7 @@ impl DatasetExec { .collect() }); let columns: Option> = columns.transpose()?; - let filter_expr: Option = conjunction(filters.to_owned()) + let filter_expr: Option> = conjunction(filters.to_owned()) .map(|filters| { PyArrowFilterExpression::try_from(&filters) .map(|filter_expr| filter_expr.inner().clone_ref(py)) @@ -131,15 +128,15 @@ impl DatasetExec { )?; let fragments_iter = pylist.call1((fragments_iterator,))?; - let fragments = fragments_iter.downcast::().map_err(PyErr::from)?; + let fragments = fragments_iter.cast::().map_err(PyErr::from)?; let projected_statistics = Statistics::new_unknown(&schema); - let plan_properties = datafusion::physical_plan::PlanProperties::new( + let plan_properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(schema.clone()), Partitioning::UnknownPartitioning(fragments.len()), EmissionType::Final, Boundedness::Bounded, - ); + )); Ok(DatasetExec { dataset: dataset.clone().unbind(), @@ -187,7 +184,7 @@ impl ExecutionPlan for DatasetExec { context: Arc, ) -> DFResult { let batch_size = context.session_config().batch_size(); - Python::with_gil(|py| { + Python::attach(|py| { let dataset = self.dataset.bind(py); let fragments = self.fragments.bind(py); let fragment = fragments @@ -238,11 +235,11 @@ impl ExecutionPlan for DatasetExec { }) } - fn statistics(&self) -> DFResult { + fn partition_statistics(&self, _partition: Option) -> DFResult { Ok(self.projected_statistics.clone()) } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } } @@ -272,10 +269,12 @@ impl ExecutionPlanProperties for DatasetExec { impl DisplayAs for DatasetExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { - Python::with_gil(|py| { + Python::attach(|py| { let number_of_fragments = self.fragments.bind(py).len(); match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => { let projected_columns: Vec = self .schema .fields() diff --git a/crates/core/src/errors.rs b/crates/core/src/errors.rs new file mode 100644 index 000000000..8babc5a56 --- /dev/null +++ b/crates/core/src/errors.rs @@ -0,0 +1,18 @@ +// 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. + +pub use datafusion_python_util::errors::*; diff --git a/src/expr.rs b/crates/core/src/expr.rs similarity index 79% rename from src/expr.rs rename to crates/core/src/expr.rs index d3c528eb4..c4f2a12da 100644 --- a/src/expr.rs +++ b/crates/core/src/expr.rs @@ -15,29 +15,36 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::utils::exprlist_to_fields; -use datafusion::logical_expr::{ - ExprFuncBuilder, ExprFunctionExt, LogicalPlan, WindowFunctionDefinition, -}; -use pyo3::IntoPyObjectExt; -use pyo3::{basic::CompareOp, prelude::*}; +use std::collections::HashMap; use std::convert::{From, Into}; use std::sync::Arc; -use window::PyWindowFrame; use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::functions::core::expr_ext::FieldAccessor; +use datafusion::logical_expr::expr::{ + AggregateFunction, AggregateFunctionParams, FieldMetadata, InList, InSubquery, ScalarFunction, + SetComparison, WindowFunction, +}; +use datafusion::logical_expr::utils::exprlist_to_fields; use datafusion::logical_expr::{ - col, - expr::{AggregateFunction, InList, InSubquery, ScalarFunction, WindowFunction}, - lit, Between, BinaryExpr, Case, Cast, Expr, Like, Operator, TryCast, + Between, BinaryExpr, Case, Cast, Expr, ExprFuncBuilder, ExprFunctionExt, Like, LogicalPlan, + Operator, TryCast, WindowFunctionDefinition, col, lit, lit_with_metadata, }; +use pyo3::IntoPyObjectExt; +use pyo3::basic::CompareOp; +use pyo3::prelude::*; +use window::PyWindowFrame; -use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType}; -use crate::errors::{ - py_runtime_err, py_type_err, py_unsupported_variant_err, PyDataFusionError, PyDataFusionResult, +use self::alias::PyAlias; +use self::bool_expr::{ + PyIsFalse, PyIsNotFalse, PyIsNotNull, PyIsNotTrue, PyIsNotUnknown, PyIsNull, PyIsTrue, + PyIsUnknown, PyNegative, PyNot, }; +use self::like::{PyILike, PyLike, PySimilarTo}; +use self::scalar_variable::PyScalarVariable; +use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType}; +use crate::errors::{PyDataFusionResult, py_runtime_err, py_type_err, py_unsupported_variant_err}; use crate::expr::aggregate_expr::PyAggregateFunction; use crate::expr::binary_expr::PyBinaryExpr; use crate::expr::column::PyColumn; @@ -46,14 +53,6 @@ use crate::functions::add_builder_fns_to_window; use crate::pyarrow_util::scalar_to_pyarrow; use crate::sql::logical::PyLogicalPlan; -use self::alias::PyAlias; -use self::bool_expr::{ - PyIsFalse, PyIsNotFalse, PyIsNotNull, PyIsNotTrue, PyIsNotUnknown, PyIsNull, PyIsTrue, - PyIsUnknown, PyNegative, PyNot, -}; -use self::like::{PyILike, PyLike, PySimilarTo}; -use self::scalar_variable::PyScalarVariable; - pub mod aggregate; pub mod aggregate_expr; pub mod alias; @@ -65,10 +64,21 @@ pub mod case; pub mod cast; pub mod column; pub mod conditional_expr; +pub mod copy_to; +pub mod create_catalog; +pub mod create_catalog_schema; +pub mod create_external_table; +pub mod create_function; +pub mod create_index; pub mod create_memory_table; pub mod create_view; +pub mod describe_table; pub mod distinct; +pub mod dml; +pub mod drop_catalog_schema; +pub mod drop_function; pub mod drop_table; +pub mod drop_view; pub mod empty_relation; pub mod exists; pub mod explain; @@ -84,24 +94,34 @@ pub mod literal; pub mod logical_node; pub mod placeholder; pub mod projection; +pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; pub mod scalar_variable; +pub mod set_comparison; pub mod signature; pub mod sort; pub mod sort_expr; +pub mod statement; pub mod subquery; pub mod subquery_alias; pub mod table_scan; pub mod union; pub mod unnest; pub mod unnest_expr; +pub mod values; pub mod window; -use sort_expr::{to_sort_expressions, PySortExpr}; +use sort_expr::{PySortExpr, to_sort_expressions}; /// A PyExpr that can be used on a DataFrame -#[pyclass(name = "RawExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "RawExpr", + module = "datafusion.expr", + subclass +)] #[derive(Debug, Clone)] pub struct PyExpr { pub expr: Expr, @@ -128,15 +148,18 @@ pub fn py_expr_list(expr: &[Expr]) -> PyResult> { impl PyExpr { /// Return the specific expression fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { - Python::with_gil(|_| { - match &self.expr { + Python::attach(|_| match &self.expr { Expr::Alias(alias) => Ok(PyAlias::from(alias.clone()).into_bound_py_any(py)?), Expr::Column(col) => Ok(PyColumn::from(col.clone()).into_bound_py_any(py)?), - Expr::ScalarVariable(data_type, variables) => { - Ok(PyScalarVariable::new(data_type, variables).into_bound_py_any(py)?) + Expr::ScalarVariable(field, variables) => { + Ok(PyScalarVariable::new(field, variables).into_bound_py_any(py)?) } Expr::Like(value) => Ok(PyLike::from(value.clone()).into_bound_py_any(py)?), - Expr::Literal(value) => Ok(PyLiteral::from(value.clone()).into_bound_py_any(py)?), + Expr::Literal(value, metadata) => Ok(PyLiteral::new_with_metadata( + value.clone(), + metadata.clone(), + ) + .into_bound_py_any(py)?), Expr::BinaryExpr(expr) => Ok(PyBinaryExpr::from(expr.clone()).into_bound_py_any(py)?), Expr::Not(expr) => Ok(PyNot::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotNull(expr) => Ok(PyIsNotNull::new(*expr.clone()).into_bound_py_any(py)?), @@ -146,25 +169,29 @@ impl PyExpr { Expr::IsUnknown(expr) => Ok(PyIsUnknown::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotTrue(expr) => Ok(PyIsNotTrue::new(*expr.clone()).into_bound_py_any(py)?), Expr::IsNotFalse(expr) => Ok(PyIsNotFalse::new(*expr.clone()).into_bound_py_any(py)?), - Expr::IsNotUnknown(expr) => Ok(PyIsNotUnknown::new(*expr.clone()).into_bound_py_any(py)?), + Expr::IsNotUnknown(expr) => { + Ok(PyIsNotUnknown::new(*expr.clone()).into_bound_py_any(py)?) + } Expr::Negative(expr) => Ok(PyNegative::new(*expr.clone()).into_bound_py_any(py)?), Expr::AggregateFunction(expr) => { Ok(PyAggregateFunction::from(expr.clone()).into_bound_py_any(py)?) } Expr::SimilarTo(value) => Ok(PySimilarTo::from(value.clone()).into_bound_py_any(py)?), - Expr::Between(value) => Ok(between::PyBetween::from(value.clone()).into_bound_py_any(py)?), + Expr::Between(value) => { + Ok(between::PyBetween::from(value.clone()).into_bound_py_any(py)?) + } Expr::Case(value) => Ok(case::PyCase::from(value.clone()).into_bound_py_any(py)?), Expr::Cast(value) => Ok(cast::PyCast::from(value.clone()).into_bound_py_any(py)?), Expr::TryCast(value) => Ok(cast::PyTryCast::from(value.clone()).into_bound_py_any(py)?), Expr::ScalarFunction(value) => Err(py_unsupported_variant_err(format!( - "Converting Expr::ScalarFunction to a Python object is not implemented: {:?}", - value + "Converting Expr::ScalarFunction to a Python object is not implemented: {value:?}" ))), Expr::WindowFunction(value) => Err(py_unsupported_variant_err(format!( - "Converting Expr::WindowFunction to a Python object is not implemented: {:?}", - value + "Converting Expr::WindowFunction to a Python object is not implemented: {value:?}" ))), - Expr::InList(value) => Ok(in_list::PyInList::from(value.clone()).into_bound_py_any(py)?), + Expr::InList(value) => { + Ok(in_list::PyInList::from(value.clone()).into_bound_py_any(py)?) + } Expr::Exists(value) => Ok(exists::PyExists::from(value.clone()).into_bound_py_any(py)?), Expr::InSubquery(value) => { Ok(in_subquery::PyInSubquery::from(value.clone()).into_bound_py_any(py)?) @@ -172,9 +199,9 @@ impl PyExpr { Expr::ScalarSubquery(value) => { Ok(scalar_subquery::PyScalarSubquery::from(value.clone()).into_bound_py_any(py)?) } + #[allow(deprecated)] Expr::Wildcard { qualifier, options } => Err(py_unsupported_variant_err(format!( - "Converting Expr::Wildcard to a Python object is not implemented : {:?} {:?}", - qualifier, options + "Converting Expr::Wildcard to a Python object is not implemented : {qualifier:?} {options:?}" ))), Expr::GroupingSet(value) => { Ok(grouping_set::PyGroupingSet::from(value.clone()).into_bound_py_any(py)?) @@ -182,12 +209,17 @@ impl PyExpr { Expr::Placeholder(value) => { Ok(placeholder::PyPlaceholder::from(value.clone()).into_bound_py_any(py)?) } - Expr::OuterReferenceColumn(data_type, column) => Err(py_unsupported_variant_err(format!( - "Converting Expr::OuterReferenceColumn to a Python object is not implemented: {:?} - {:?}", - data_type, column - ))), - Expr::Unnest(value) => Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_bound_py_any(py)?), - } + Expr::OuterReferenceColumn(data_type, column) => { + Err(py_unsupported_variant_err(format!( + "Converting Expr::OuterReferenceColumn to a Python object is not implemented: {data_type:?} - {column:?}" + ))) + } + Expr::Unnest(value) => { + Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_bound_py_any(py)?) + } + Expr::SetComparison(value) => { + Ok(set_comparison::PySetComparison::from(value.clone()).into_bound_py_any(py)?) + } }) } @@ -267,14 +299,25 @@ impl PyExpr { lit(value.0).into() } + #[staticmethod] + pub fn literal_with_metadata( + value: PyScalarValue, + metadata: HashMap, + ) -> PyExpr { + let metadata = FieldMetadata::new(metadata.into_iter().collect()); + lit_with_metadata(value.0, Some(metadata)).into() + } + #[staticmethod] pub fn column(value: &str) -> PyExpr { col(value).into() } /// assign a name to the PyExpr - pub fn alias(&self, name: &str) -> PyExpr { - self.expr.clone().alias(name).into() + #[pyo3(signature = (name, metadata=None))] + pub fn alias(&self, name: &str, metadata: Option>) -> PyExpr { + let metadata = metadata.map(|m| FieldMetadata::new(m.into_iter().collect())); + self.expr.clone().alias_with_metadata(name, metadata).into() } /// Create a sort PyExpr from an existing PyExpr. @@ -332,7 +375,6 @@ impl PyExpr { | Expr::AggregateFunction { .. } | Expr::WindowFunction { .. } | Expr::InList { .. } - | Expr::Wildcard { .. } | Expr::Exists { .. } | Expr::InSubquery { .. } | Expr::GroupingSet(..) @@ -344,8 +386,13 @@ impl PyExpr { | Expr::Placeholder { .. } | Expr::OuterReferenceColumn(_, _) | Expr::Unnest(_) - | Expr::IsNotUnknown(_) => RexType::Call, + | Expr::IsNotUnknown(_) + | Expr::SetComparison(_) => RexType::Call, Expr::ScalarSubquery(..) => RexType::ScalarSubquery, + #[allow(deprecated)] + Expr::Wildcard { .. } => { + return Err(py_unsupported_variant_err("Expr::Wildcard is unsupported")); + } }) } @@ -355,10 +402,10 @@ impl PyExpr { Self::_types(&self.expr) } - /// Extracts the Expr value into a PyObject that can be shared with Python - pub fn python_value(&self, py: Python) -> PyResult { + /// Extracts the Expr value into a Py that can be shared with Python + pub fn python_value<'py>(&self, py: Python<'py>) -> PyResult> { match &self.expr { - Expr::Literal(scalar_value) => scalar_to_pyarrow(scalar_value, py), + Expr::Literal(scalar_value, _) => scalar_to_pyarrow(scalar_value, py), _ => Err(py_type_err(format!( "Non Expr::Literal encountered in types: {:?}", &self.expr @@ -391,12 +438,21 @@ impl PyExpr { | Expr::Negative(expr) | Expr::Cast(Cast { expr, .. }) | Expr::TryCast(TryCast { expr, .. }) - | Expr::InSubquery(InSubquery { expr, .. }) => Ok(vec![PyExpr::from(*expr.clone())]), + | Expr::InSubquery(InSubquery { expr, .. }) + | Expr::SetComparison(SetComparison { expr, .. }) => { + Ok(vec![PyExpr::from(*expr.clone())]) + } // Expr variants containing a collection of Expr(s) for operands - Expr::AggregateFunction(AggregateFunction { args, .. }) - | Expr::ScalarFunction(ScalarFunction { args, .. }) - | Expr::WindowFunction(WindowFunction { args, .. }) => { + Expr::AggregateFunction(AggregateFunction { + params: AggregateFunctionParams { args, .. }, + .. + }) + | Expr::ScalarFunction(ScalarFunction { args, .. }) => { + Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect()) + } + Expr::WindowFunction(boxed_window_fn) => { + let args = &boxed_window_fn.params.args; Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect()) } @@ -465,13 +521,17 @@ impl PyExpr { Expr::GroupingSet(..) | Expr::Unnest(_) | Expr::OuterReferenceColumn(_, _) - | Expr::Wildcard { .. } | Expr::ScalarSubquery(..) | Expr::Placeholder { .. } | Expr::Exists { .. } => Err(py_runtime_err(format!( "Unimplemented Expr type: {}", self.expr ))), + + #[allow(deprecated)] + Expr::Wildcard { .. } => { + Err(py_unsupported_variant_err("Expr::Wildcard is unsupported")) + } } } @@ -521,7 +581,7 @@ impl PyExpr { return Err(py_type_err(format!( "Catch all triggered in get_operator_name: {:?}", &self.expr - ))) + ))); } }) } @@ -573,10 +633,10 @@ impl PyExpr { ) -> PyDataFusionResult { match &self.expr { Expr::AggregateFunction(agg_fn) => { - let window_fn = Expr::WindowFunction(WindowFunction::new( + let window_fn = Expr::WindowFunction(Box::new(WindowFunction::new( WindowFunctionDefinition::AggregateUDF(agg_fn.func.clone()), - agg_fn.args.clone(), - )); + agg_fn.params.args.clone(), + ))); add_builder_fns_to_window( window_fn, @@ -593,16 +653,22 @@ impl PyExpr { order_by, null_treatment, ), - _ => Err( - PyDataFusionError::ExecutionError(datafusion::error::DataFusionError::Plan( - format!("Using {} with `over` is not allowed. Must use an aggregate or window function.", self.expr.variant_name()), - )) - ), + _ => Err(datafusion::error::DataFusionError::Plan(format!( + "Using {} with `over` is not allowed. Must use an aggregate or window function.", + self.expr.variant_name() + )) + .into()), } } } -#[pyclass(name = "ExprFuncBuilder", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ExprFuncBuilder", + module = "datafusion.expr", + subclass +)] #[derive(Debug, Clone)] pub struct PyExprFuncBuilder { pub builder: ExprFuncBuilder, @@ -663,16 +729,8 @@ impl PyExpr { /// Create a [Field] representing an [Expr], given an input [LogicalPlan] to resolve against pub fn expr_to_field(expr: &Expr, input_plan: &LogicalPlan) -> PyDataFusionResult> { - match expr { - Expr::Wildcard { .. } => { - // Since * could be any of the valid column names just return the first one - Ok(Arc::new(input_plan.schema().field(0).clone())) - } - _ => { - let fields = exprlist_to_fields(&[expr.clone()], input_plan)?; - Ok(fields[0].1.clone()) - } - } + let fields = exprlist_to_fields(std::slice::from_ref(expr), input_plan)?; + Ok(fields[0].1.clone()) } fn _types(expr: &Expr) -> PyResult { match expr { @@ -709,15 +767,25 @@ impl PyExpr { | Operator::BitwiseXor | Operator::BitwiseAnd | Operator::BitwiseOr => DataTypeMap::map_from_arrow_type(&DataType::Binary), - Operator::AtArrow | Operator::ArrowAt => { - Err(py_type_err(format!("Unsupported expr: ${op}"))) - } + Operator::AtArrow + | Operator::ArrowAt + | Operator::Arrow + | Operator::LongArrow + | Operator::HashArrow + | Operator::HashLongArrow + | Operator::AtAt + | Operator::IntegerDivide + | Operator::HashMinus + | Operator::AtQuestion + | Operator::Question + | Operator::QuestionAnd + | Operator::QuestionPipe + | Operator::Colon => Err(py_type_err(format!("Unsupported expr: ${op}"))), }, Expr::Cast(Cast { expr: _, data_type }) => DataTypeMap::map_from_arrow_type(data_type), - Expr::Literal(scalar_value) => DataTypeMap::map_from_scalar_value(scalar_value), + Expr::Literal(scalar_value, _) => DataTypeMap::map_from_scalar_value(scalar_value), _ => Err(py_type_err(format!( - "Non Expr::Literal encountered in types: {:?}", - expr + "Non Expr::Literal encountered in types: {expr:?}" ))), } } @@ -785,5 +853,32 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) } diff --git a/src/expr/aggregate.rs b/crates/core/src/expr/aggregate.rs similarity index 89% rename from src/expr/aggregate.rs rename to crates/core/src/expr/aggregate.rs index 8fc9da5b0..5a6a771a7 100644 --- a/src/expr/aggregate.rs +++ b/crates/core/src/expr/aggregate.rs @@ -15,12 +15,14 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::common::DataFusionError; -use datafusion::logical_expr::expr::{AggregateFunction, Alias}; -use datafusion::logical_expr::logical_plan::Aggregate; use datafusion::logical_expr::Expr; -use pyo3::{prelude::*, IntoPyObjectExt}; -use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::expr::{AggregateFunction, AggregateFunctionParams, Alias}; +use datafusion::logical_expr::logical_plan::Aggregate; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; @@ -28,7 +30,13 @@ use crate::errors::py_type_err; use crate::expr::PyExpr; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Aggregate", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Aggregate", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAggregate { aggregate: Aggregate, @@ -116,7 +124,7 @@ impl PyAggregate { } fn __repr__(&self) -> PyResult { - Ok(format!("Aggregate({})", self)) + Ok(format!("Aggregate({self})")) } } @@ -126,9 +134,11 @@ impl PyAggregate { match expr { // TODO: This Alias logic seems to be returning some strange results that we should investigate Expr::Alias(Alias { expr, .. }) => self._aggregation_arguments(expr.as_ref()), - Expr::AggregateFunction(AggregateFunction { func: _, args, .. }) => { - Ok(args.iter().map(|e| PyExpr::from(e.clone())).collect()) - } + Expr::AggregateFunction(AggregateFunction { + func: _, + params: AggregateFunctionParams { args, .. }, + .. + }) => Ok(args.iter().map(|e| PyExpr::from(e.clone())).collect()), _ => Err(py_type_err( "Encountered a non Aggregate type in aggregation_arguments", )), diff --git a/src/expr/aggregate_expr.rs b/crates/core/src/expr/aggregate_expr.rs similarity index 84% rename from src/expr/aggregate_expr.rs rename to crates/core/src/expr/aggregate_expr.rs index 09471097f..88e47999f 100644 --- a/src/expr/aggregate_expr.rs +++ b/crates/core/src/expr/aggregate_expr.rs @@ -15,12 +15,20 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; +use std::fmt::{Display, Formatter}; + use datafusion::logical_expr::expr::AggregateFunction; use pyo3::prelude::*; -use std::fmt::{Display, Formatter}; -#[pyclass(name = "AggregateFunction", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "AggregateFunction", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAggregateFunction { aggr: AggregateFunction, @@ -40,7 +48,13 @@ impl From for PyAggregateFunction { impl Display for PyAggregateFunction { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { - let args: Vec = self.aggr.args.iter().map(|expr| expr.to_string()).collect(); + let args: Vec = self + .aggr + .params + .args + .iter() + .map(|expr| expr.to_string()) + .collect(); write!(f, "{}({})", self.aggr.func.name(), args.join(", ")) } } @@ -54,12 +68,13 @@ impl PyAggregateFunction { /// is this a distinct aggregate such as `COUNT(DISTINCT expr)` fn is_distinct(&self) -> bool { - self.aggr.distinct + self.aggr.params.distinct } /// Get the arguments to the aggregate function fn args(&self) -> Vec { self.aggr + .params .args .iter() .map(|expr| PyExpr::from(expr.clone())) @@ -68,6 +83,6 @@ impl PyAggregateFunction { /// Get a String representation of this column fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } } diff --git a/src/expr/alias.rs b/crates/core/src/expr/alias.rs similarity index 92% rename from src/expr/alias.rs rename to crates/core/src/expr/alias.rs index e8e03cfad..b76e82e22 100644 --- a/src/expr/alias.rs +++ b/crates/core/src/expr/alias.rs @@ -15,13 +15,20 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; -use pyo3::prelude::*; use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::expr::Alias; +use pyo3::prelude::*; + +use crate::expr::PyExpr; -#[pyclass(name = "Alias", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Alias", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAlias { alias: Alias, @@ -64,6 +71,6 @@ impl PyAlias { /// Get a String representation of this column fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } } diff --git a/src/expr/analyze.rs b/crates/core/src/expr/analyze.rs similarity index 91% rename from src/expr/analyze.rs rename to crates/core/src/expr/analyze.rs index 62f93cd26..137765fe1 100644 --- a/src/expr/analyze.rs +++ b/crates/core/src/expr/analyze.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Analyze; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::logical_plan::Analyze; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use super::logical_node::LogicalNode; use crate::common::df_schema::PyDFSchema; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Analyze", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Analyze", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyAnalyze { analyze: Analyze, @@ -69,7 +77,7 @@ impl PyAnalyze { } fn __repr__(&self) -> PyResult { - Ok(format!("Analyze({})", self)) + Ok(format!("Analyze({self})")) } } diff --git a/src/expr/between.rs b/crates/core/src/expr/between.rs similarity index 93% rename from src/expr/between.rs rename to crates/core/src/expr/between.rs index a2cac1442..6943b6c3b 100644 --- a/src/expr/between.rs +++ b/crates/core/src/expr/between.rs @@ -15,12 +15,20 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; +use std::fmt::{self, Display, Formatter}; + use datafusion::logical_expr::expr::Between; use pyo3::prelude::*; -use std::fmt::{self, Display, Formatter}; -#[pyclass(name = "Between", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "Between", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyBetween { between: Between, @@ -71,6 +79,6 @@ impl PyBetween { } fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } } diff --git a/src/expr/binary_expr.rs b/crates/core/src/expr/binary_expr.rs similarity index 93% rename from src/expr/binary_expr.rs rename to crates/core/src/expr/binary_expr.rs index 740299211..2326ba705 100644 --- a/src/expr/binary_expr.rs +++ b/crates/core/src/expr/binary_expr.rs @@ -15,11 +15,18 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; use datafusion::logical_expr::BinaryExpr; use pyo3::prelude::*; -#[pyclass(name = "BinaryExpr", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "BinaryExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyBinaryExpr { expr: BinaryExpr, diff --git a/src/expr/bool_expr.rs b/crates/core/src/expr/bool_expr.rs similarity index 83% rename from src/expr/bool_expr.rs rename to crates/core/src/expr/bool_expr.rs index e67e25d74..9e374c7e2 100644 --- a/src/expr/bool_expr.rs +++ b/crates/core/src/expr/bool_expr.rs @@ -15,13 +15,20 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::logical_expr::Expr; use pyo3::prelude::*; -use std::fmt::{self, Display, Formatter}; use super::PyExpr; -#[pyclass(name = "Not", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Not", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyNot { expr: Expr, @@ -51,7 +58,13 @@ impl PyNot { } } -#[pyclass(name = "IsNotNull", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotNull", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotNull { expr: Expr, @@ -81,7 +94,13 @@ impl PyIsNotNull { } } -#[pyclass(name = "IsNull", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNull", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNull { expr: Expr, @@ -111,7 +130,13 @@ impl PyIsNull { } } -#[pyclass(name = "IsTrue", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsTrue", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsTrue { expr: Expr, @@ -141,7 +166,13 @@ impl PyIsTrue { } } -#[pyclass(name = "IsFalse", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsFalse", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsFalse { expr: Expr, @@ -171,7 +202,13 @@ impl PyIsFalse { } } -#[pyclass(name = "IsUnknown", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsUnknown", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsUnknown { expr: Expr, @@ -201,7 +238,13 @@ impl PyIsUnknown { } } -#[pyclass(name = "IsNotTrue", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotTrue", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotTrue { expr: Expr, @@ -231,7 +274,13 @@ impl PyIsNotTrue { } } -#[pyclass(name = "IsNotFalse", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotFalse", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotFalse { expr: Expr, @@ -261,7 +310,13 @@ impl PyIsNotFalse { } } -#[pyclass(name = "IsNotUnknown", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "IsNotUnknown", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyIsNotUnknown { expr: Expr, @@ -291,7 +346,13 @@ impl PyIsNotUnknown { } } -#[pyclass(name = "Negative", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Negative", + module = "datafusion.expr", + subclass +)] #[derive(Clone, Debug)] pub struct PyNegative { expr: Expr, diff --git a/src/expr/case.rs b/crates/core/src/expr/case.rs similarity index 93% rename from src/expr/case.rs rename to crates/core/src/expr/case.rs index 92e28ba56..4f00449d8 100644 --- a/src/expr/case.rs +++ b/crates/core/src/expr/case.rs @@ -15,11 +15,18 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; use datafusion::logical_expr::Case; use pyo3::prelude::*; -#[pyclass(name = "Case", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "Case", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCase { case: Case, diff --git a/src/expr/cast.rs b/crates/core/src/expr/cast.rs similarity index 87% rename from src/expr/cast.rs rename to crates/core/src/expr/cast.rs index b8faea634..37d603538 100644 --- a/src/expr/cast.rs +++ b/crates/core/src/expr/cast.rs @@ -15,11 +15,19 @@ // specific language governing permissions and limitations // under the License. -use crate::{common::data_type::PyDataType, expr::PyExpr}; use datafusion::logical_expr::{Cast, TryCast}; use pyo3::prelude::*; -#[pyclass(name = "Cast", module = "datafusion.expr", subclass)] +use crate::common::data_type::PyDataType; +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "Cast", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCast { cast: Cast, @@ -48,7 +56,7 @@ impl PyCast { } } -#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)] +#[pyclass(from_py_object, name = "TryCast", module = "datafusion.expr", subclass)] #[derive(Clone)] pub struct PyTryCast { try_cast: TryCast, diff --git a/src/expr/column.rs b/crates/core/src/expr/column.rs similarity index 90% rename from src/expr/column.rs rename to crates/core/src/expr/column.rs index 365dbc0d2..c1238f98a 100644 --- a/src/expr/column.rs +++ b/crates/core/src/expr/column.rs @@ -18,7 +18,13 @@ use datafusion::common::Column; use pyo3::prelude::*; -#[pyclass(name = "Column", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Column", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyColumn { pub col: Column, @@ -45,7 +51,7 @@ impl PyColumn { /// Get the column relation fn relation(&self) -> Option { - self.col.relation.as_ref().map(|r| format!("{}", r)) + self.col.relation.as_ref().map(|r| format!("{r}")) } /// Get the fully-qualified column name diff --git a/crates/core/src/expr/conditional_expr.rs b/crates/core/src/expr/conditional_expr.rs new file mode 100644 index 000000000..ea21fdb20 --- /dev/null +++ b/crates/core/src/expr/conditional_expr.rs @@ -0,0 +1,84 @@ +// 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 datafusion::logical_expr::conditional_expressions::CaseBuilder; +use datafusion::prelude::Expr; +use pyo3::prelude::*; + +use crate::errors::PyDataFusionResult; +use crate::expr::PyExpr; + +// TODO(tsaucer) replace this all with CaseBuilder after it implements Clone +#[derive(Clone, Debug)] +#[pyclass( + from_py_object, + name = "CaseBuilder", + module = "datafusion.expr", + subclass, + frozen +)] +pub struct PyCaseBuilder { + expr: Option, + when: Vec, + then: Vec, +} + +#[pymethods] +impl PyCaseBuilder { + #[new] + pub fn new(expr: Option) -> Self { + Self { + expr: expr.map(Into::into), + when: vec![], + then: vec![], + } + } + + pub fn when(&self, when: PyExpr, then: PyExpr) -> PyCaseBuilder { + let mut case_builder = self.clone(); + case_builder.when.push(when.into()); + case_builder.then.push(then.into()); + + case_builder + } + + fn otherwise(&self, else_expr: PyExpr) -> PyDataFusionResult { + let case_builder = CaseBuilder::new( + self.expr.clone().map(Box::new), + self.when.clone(), + self.then.clone(), + Some(Box::new(else_expr.into())), + ); + + let expr = case_builder.end()?; + + Ok(expr.into()) + } + + fn end(&self) -> PyDataFusionResult { + let case_builder = CaseBuilder::new( + self.expr.clone().map(Box::new), + self.when.clone(), + self.then.clone(), + None, + ); + + let expr = case_builder.end()?; + + Ok(expr.into()) + } +} diff --git a/crates/core/src/expr/copy_to.rs b/crates/core/src/expr/copy_to.rs new file mode 100644 index 000000000..78e53cdff --- /dev/null +++ b/crates/core/src/expr/copy_to.rs @@ -0,0 +1,149 @@ +// 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 std::collections::HashMap; +use std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::common::file_options::file_type::FileType; +use datafusion::logical_expr::dml::CopyTo; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CopyTo", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCopyTo { + copy: CopyTo, +} + +impl From for CopyTo { + fn from(copy: PyCopyTo) -> Self { + copy.copy + } +} + +impl From for PyCopyTo { + fn from(copy: CopyTo) -> PyCopyTo { + PyCopyTo { copy } + } +} + +impl Display for PyCopyTo { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "CopyTo: {:?}", self.copy.output_url) + } +} + +impl LogicalNode for PyCopyTo { + fn inputs(&self) -> Vec { + vec![PyLogicalPlan::from((*self.copy.input).clone())] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyCopyTo { + #[new] + pub fn new( + input: PyLogicalPlan, + output_url: String, + partition_by: Vec, + file_type: PyFileType, + options: HashMap, + ) -> Self { + PyCopyTo { + copy: CopyTo::new( + input.plan(), + output_url, + partition_by, + file_type.file_type, + options, + ), + } + } + + fn input(&self) -> PyLogicalPlan { + PyLogicalPlan::from((*self.copy.input).clone()) + } + + fn output_url(&self) -> String { + self.copy.output_url.clone() + } + + fn partition_by(&self) -> Vec { + self.copy.partition_by.clone() + } + + fn file_type(&self) -> PyFileType { + PyFileType { + file_type: self.copy.file_type.clone(), + } + } + + fn options(&self) -> HashMap { + self.copy.options.clone() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CopyTo({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CopyTo".to_string()) + } +} + +#[pyclass( + from_py_object, + frozen, + name = "FileType", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyFileType { + file_type: Arc, +} + +impl Display for PyFileType { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "FileType: {}", self.file_type) + } +} + +#[pymethods] +impl PyFileType { + fn __repr__(&self) -> PyResult { + Ok(format!("FileType({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("FileType".to_string()) + } +} diff --git a/crates/core/src/expr/create_catalog.rs b/crates/core/src/expr/create_catalog.rs new file mode 100644 index 000000000..fa95980c0 --- /dev/null +++ b/crates/core/src/expr/create_catalog.rs @@ -0,0 +1,105 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::CreateCatalog; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateCatalog", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateCatalog { + create: CreateCatalog, +} + +impl From for CreateCatalog { + fn from(create: PyCreateCatalog) -> Self { + create.create + } +} + +impl From for PyCreateCatalog { + fn from(create: CreateCatalog) -> PyCreateCatalog { + PyCreateCatalog { create } + } +} + +impl Display for PyCreateCatalog { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "CreateCatalog: {:?}", self.create.catalog_name) + } +} + +#[pymethods] +impl PyCreateCatalog { + #[new] + pub fn new( + catalog_name: String, + if_not_exists: bool, + schema: PyDFSchema, + ) -> PyResult { + Ok(PyCreateCatalog { + create: CreateCatalog { + catalog_name, + if_not_exists, + schema: Arc::new(schema.into()), + }, + }) + } + + pub fn catalog_name(&self) -> String { + self.create.catalog_name.clone() + } + + pub fn if_not_exists(&self) -> bool { + self.create.if_not_exists + } + + pub fn schema(&self) -> PyDFSchema { + (*self.create.schema).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateCatalog({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CreateCatalog".to_string()) + } +} + +impl LogicalNode for PyCreateCatalog { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/crates/core/src/expr/create_catalog_schema.rs b/crates/core/src/expr/create_catalog_schema.rs new file mode 100644 index 000000000..d836284a0 --- /dev/null +++ b/crates/core/src/expr/create_catalog_schema.rs @@ -0,0 +1,105 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::CreateCatalogSchema; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateCatalogSchema", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateCatalogSchema { + create: CreateCatalogSchema, +} + +impl From for CreateCatalogSchema { + fn from(create: PyCreateCatalogSchema) -> Self { + create.create + } +} + +impl From for PyCreateCatalogSchema { + fn from(create: CreateCatalogSchema) -> PyCreateCatalogSchema { + PyCreateCatalogSchema { create } + } +} + +impl Display for PyCreateCatalogSchema { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "CreateCatalogSchema: {:?}", self.create.schema_name) + } +} + +#[pymethods] +impl PyCreateCatalogSchema { + #[new] + pub fn new( + schema_name: String, + if_not_exists: bool, + schema: PyDFSchema, + ) -> PyResult { + Ok(PyCreateCatalogSchema { + create: CreateCatalogSchema { + schema_name, + if_not_exists, + schema: Arc::new(schema.into()), + }, + }) + } + + pub fn schema_name(&self) -> String { + self.create.schema_name.clone() + } + + pub fn if_not_exists(&self) -> bool { + self.create.if_not_exists + } + + pub fn schema(&self) -> PyDFSchema { + (*self.create.schema).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateCatalogSchema({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CreateCatalogSchema".to_string()) + } +} + +impl LogicalNode for PyCreateCatalogSchema { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/crates/core/src/expr/create_external_table.rs b/crates/core/src/expr/create_external_table.rs new file mode 100644 index 000000000..980eea131 --- /dev/null +++ b/crates/core/src/expr/create_external_table.rs @@ -0,0 +1,192 @@ +// 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 std::collections::HashMap; +use std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::CreateExternalTable; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use super::sort_expr::PySortExpr; +use crate::common::df_schema::PyDFSchema; +use crate::common::schema::PyConstraints; +use crate::expr::PyExpr; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateExternalTable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateExternalTable { + create: CreateExternalTable, +} + +impl From for CreateExternalTable { + fn from(create: PyCreateExternalTable) -> Self { + create.create + } +} + +impl From for PyCreateExternalTable { + fn from(create: CreateExternalTable) -> PyCreateExternalTable { + PyCreateExternalTable { create } + } +} + +impl Display for PyCreateExternalTable { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "CreateExternalTable: {:?}{}", + self.create.name, self.create.constraints + ) + } +} + +#[pymethods] +impl PyCreateExternalTable { + #[allow(clippy::too_many_arguments)] + #[new] + #[pyo3(signature = (schema, name, location, file_type, table_partition_cols, if_not_exists, or_replace, temporary, order_exprs, unbounded, options, constraints, column_defaults, definition=None))] + pub fn new( + schema: PyDFSchema, + name: String, + location: String, + file_type: String, + table_partition_cols: Vec, + if_not_exists: bool, + or_replace: bool, + temporary: bool, + order_exprs: Vec>, + unbounded: bool, + options: HashMap, + constraints: PyConstraints, + column_defaults: HashMap, + definition: Option, + ) -> Self { + let create = CreateExternalTable { + schema: Arc::new(schema.into()), + name: name.into(), + location, + file_type, + table_partition_cols, + if_not_exists, + or_replace, + temporary, + definition, + order_exprs: order_exprs + .into_iter() + .map(|vec| vec.into_iter().map(|s| s.into()).collect::>()) + .collect::>(), + unbounded, + options, + constraints: constraints.constraints, + column_defaults: column_defaults + .into_iter() + .map(|(k, v)| (k, v.into())) + .collect(), + }; + PyCreateExternalTable { create } + } + + pub fn schema(&self) -> PyDFSchema { + (*self.create.schema).clone().into() + } + + pub fn name(&self) -> PyResult { + Ok(self.create.name.to_string()) + } + + pub fn location(&self) -> String { + self.create.location.clone() + } + + pub fn file_type(&self) -> String { + self.create.file_type.clone() + } + + pub fn table_partition_cols(&self) -> Vec { + self.create.table_partition_cols.clone() + } + + pub fn if_not_exists(&self) -> bool { + self.create.if_not_exists + } + + pub fn temporary(&self) -> bool { + self.create.temporary + } + + pub fn definition(&self) -> Option { + self.create.definition.clone() + } + + pub fn order_exprs(&self) -> Vec> { + self.create + .order_exprs + .iter() + .map(|vec| vec.iter().map(|s| s.clone().into()).collect()) + .collect() + } + + pub fn unbounded(&self) -> bool { + self.create.unbounded + } + + pub fn options(&self) -> HashMap { + self.create.options.clone() + } + + pub fn constraints(&self) -> PyConstraints { + PyConstraints { + constraints: self.create.constraints.clone(), + } + } + + pub fn column_defaults(&self) -> HashMap { + self.create + .column_defaults + .iter() + .map(|(k, v)| (k.clone(), v.clone().into())) + .collect() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateExternalTable({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CreateExternalTable".to_string()) + } +} + +impl LogicalNode for PyCreateExternalTable { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/crates/core/src/expr/create_function.rs b/crates/core/src/expr/create_function.rs new file mode 100644 index 000000000..622858913 --- /dev/null +++ b/crates/core/src/expr/create_function.rs @@ -0,0 +1,207 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::{ + CreateFunction, CreateFunctionBody, OperateFunctionArg, Volatility, +}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::PyExpr; +use super::logical_node::LogicalNode; +use crate::common::data_type::PyDataType; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateFunction", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateFunction { + create: CreateFunction, +} + +impl From for CreateFunction { + fn from(create: PyCreateFunction) -> Self { + create.create + } +} + +impl From for PyCreateFunction { + fn from(create: CreateFunction) -> PyCreateFunction { + PyCreateFunction { create } + } +} + +impl Display for PyCreateFunction { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "CreateFunction: name {:?}", self.create.name) + } +} + +#[pyclass( + from_py_object, + frozen, + name = "OperateFunctionArg", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyOperateFunctionArg { + arg: OperateFunctionArg, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "Volatility", + module = "datafusion.expr" +)] +pub enum PyVolatility { + Immutable, + Stable, + Volatile, +} + +#[pyclass( + from_py_object, + frozen, + name = "CreateFunctionBody", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateFunctionBody { + body: CreateFunctionBody, +} + +#[pymethods] +impl PyCreateFunctionBody { + pub fn language(&self) -> Option { + self.body + .language + .as_ref() + .map(|language| language.to_string()) + } + + pub fn behavior(&self) -> Option { + self.body.behavior.as_ref().map(|behavior| match behavior { + Volatility::Immutable => PyVolatility::Immutable, + Volatility::Stable => PyVolatility::Stable, + Volatility::Volatile => PyVolatility::Volatile, + }) + } + + pub fn function_body(&self) -> Option { + self.body + .function_body + .as_ref() + .map(|function_body| function_body.clone().into()) + } +} + +#[pymethods] +impl PyCreateFunction { + #[new] + #[pyo3(signature = (or_replace, temporary, name, params, schema, return_type=None, args=None))] + pub fn new( + or_replace: bool, + temporary: bool, + name: String, + params: PyCreateFunctionBody, + schema: PyDFSchema, + return_type: Option, + args: Option>, + ) -> Self { + PyCreateFunction { + create: CreateFunction { + or_replace, + temporary, + name, + args: args.map(|args| args.into_iter().map(|arg| arg.arg).collect()), + return_type: return_type.map(|return_type| return_type.data_type), + params: params.body, + schema: Arc::new(schema.into()), + }, + } + } + + pub fn or_replace(&self) -> bool { + self.create.or_replace + } + + pub fn temporary(&self) -> bool { + self.create.temporary + } + + pub fn name(&self) -> String { + self.create.name.clone() + } + + pub fn params(&self) -> PyCreateFunctionBody { + PyCreateFunctionBody { + body: self.create.params.clone(), + } + } + + pub fn schema(&self) -> PyDFSchema { + (*self.create.schema).clone().into() + } + + pub fn return_type(&self) -> Option { + self.create + .return_type + .as_ref() + .map(|return_type| return_type.clone().into()) + } + + pub fn args(&self) -> Option> { + self.create.args.as_ref().map(|args| { + args.iter() + .map(|arg| PyOperateFunctionArg { arg: arg.clone() }) + .collect() + }) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateFunction({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CreateFunction".to_string()) + } +} + +impl LogicalNode for PyCreateFunction { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/crates/core/src/expr/create_index.rs b/crates/core/src/expr/create_index.rs new file mode 100644 index 000000000..5f9bd11e8 --- /dev/null +++ b/crates/core/src/expr/create_index.rs @@ -0,0 +1,135 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::CreateIndex; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use super::sort_expr::PySortExpr; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateIndex", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyCreateIndex { + create: CreateIndex, +} + +impl From for CreateIndex { + fn from(create: PyCreateIndex) -> Self { + create.create + } +} + +impl From for PyCreateIndex { + fn from(create: CreateIndex) -> PyCreateIndex { + PyCreateIndex { create } + } +} + +impl Display for PyCreateIndex { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "CreateIndex: {:?}", self.create.name) + } +} + +#[pymethods] +impl PyCreateIndex { + #[new] + #[pyo3(signature = (table, columns, unique, if_not_exists, schema, name=None, using=None))] + pub fn new( + table: String, + columns: Vec, + unique: bool, + if_not_exists: bool, + schema: PyDFSchema, + name: Option, + using: Option, + ) -> PyResult { + Ok(PyCreateIndex { + create: CreateIndex { + name, + table: table.into(), + using, + columns: columns.iter().map(|c| c.clone().into()).collect(), + unique, + if_not_exists, + schema: Arc::new(schema.into()), + }, + }) + } + + pub fn name(&self) -> Option { + self.create.name.clone() + } + + pub fn table(&self) -> PyResult { + Ok(self.create.table.to_string()) + } + + pub fn using(&self) -> Option { + self.create.using.clone() + } + + pub fn columns(&self) -> Vec { + self.create + .columns + .iter() + .map(|c| c.clone().into()) + .collect() + } + + pub fn unique(&self) -> bool { + self.create.unique + } + + pub fn if_not_exists(&self) -> bool { + self.create.if_not_exists + } + + pub fn schema(&self) -> PyDFSchema { + (*self.create.schema).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("CreateIndex({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("CreateIndex".to_string()) + } +} + +impl LogicalNode for PyCreateIndex { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/src/expr/create_memory_table.rs b/crates/core/src/expr/create_memory_table.rs similarity index 92% rename from src/expr/create_memory_table.rs rename to crates/core/src/expr/create_memory_table.rs index 8872b2d47..3214dab0e 100644 --- a/src/expr/create_memory_table.rs +++ b/crates/core/src/expr/create_memory_table.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::CreateMemoryTable; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::sql::logical::PyLogicalPlan; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "CreateMemoryTable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "CreateMemoryTable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateMemoryTable { create: CreateMemoryTable, @@ -78,7 +84,7 @@ impl PyCreateMemoryTable { } fn __repr__(&self) -> PyResult { - Ok(format!("CreateMemoryTable({})", self)) + Ok(format!("CreateMemoryTable({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/create_view.rs b/crates/core/src/expr/create_view.rs similarity index 90% rename from src/expr/create_view.rs rename to crates/core/src/expr/create_view.rs index 87bb76876..6941ef769 100644 --- a/src/expr/create_view.rs +++ b/crates/core/src/expr/create_view.rs @@ -18,13 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::{CreateView, DdlStatement, LogicalPlan}; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::{errors::py_type_err, sql::logical::PyLogicalPlan}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; - -#[pyclass(name = "CreateView", module = "datafusion.expr", subclass)] +use crate::errors::py_type_err; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "CreateView", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyCreateView { create: CreateView, @@ -75,7 +82,7 @@ impl PyCreateView { } fn __repr__(&self) -> PyResult { - Ok(format!("CreateView({})", self)) + Ok(format!("CreateView({self})")) } fn __name__(&self) -> PyResult { diff --git a/crates/core/src/expr/describe_table.rs b/crates/core/src/expr/describe_table.rs new file mode 100644 index 000000000..73955bb34 --- /dev/null +++ b/crates/core/src/expr/describe_table.rs @@ -0,0 +1,98 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use arrow::datatypes::Schema; +use arrow::pyarrow::PyArrowType; +use datafusion::logical_expr::DescribeTable; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "DescribeTable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDescribeTable { + describe: DescribeTable, +} + +impl Display for PyDescribeTable { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "DescribeTable") + } +} + +#[pymethods] +impl PyDescribeTable { + #[new] + fn new(schema: PyArrowType, output_schema: PyDFSchema) -> Self { + Self { + describe: DescribeTable { + schema: Arc::new(schema.0), + output_schema: Arc::new(output_schema.into()), + }, + } + } + + pub fn schema(&self) -> PyArrowType { + (*self.describe.schema).clone().into() + } + + pub fn output_schema(&self) -> PyDFSchema { + (*self.describe.output_schema).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok(format!("DescribeTable({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("DescribeTable".to_string()) + } +} + +impl From for DescribeTable { + fn from(describe: PyDescribeTable) -> Self { + describe.describe + } +} + +impl From for PyDescribeTable { + fn from(describe: DescribeTable) -> PyDescribeTable { + PyDescribeTable { describe } + } +} + +impl LogicalNode for PyDescribeTable { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/src/expr/distinct.rs b/crates/core/src/expr/distinct.rs similarity index 91% rename from src/expr/distinct.rs rename to crates/core/src/expr/distinct.rs index b62b776f8..68c2a17fe 100644 --- a/src/expr/distinct.rs +++ b/crates/core/src/expr/distinct.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::Distinct; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::sql::logical::PyLogicalPlan; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Distinct", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Distinct", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDistinct { distinct: Distinct, @@ -48,8 +54,7 @@ impl Display for PyDistinct { Distinct::All(input) => write!( f, "Distinct ALL - \nInput: {:?}", - input, + \nInput: {input:?}", ), Distinct::On(distinct_on) => { write!( @@ -71,7 +76,7 @@ impl PyDistinct { } fn __repr__(&self) -> PyResult { - Ok(format!("Distinct({})", self)) + Ok(format!("Distinct({self})")) } fn __name__(&self) -> PyResult { diff --git a/crates/core/src/expr/dml.rs b/crates/core/src/expr/dml.rs new file mode 100644 index 000000000..26f975820 --- /dev/null +++ b/crates/core/src/expr/dml.rs @@ -0,0 +1,149 @@ +// 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 datafusion::logical_expr::dml::InsertOp; +use datafusion::logical_expr::{DmlStatement, WriteOp}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::common::schema::PyTableSource; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "DmlStatement", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDmlStatement { + dml: DmlStatement, +} + +impl From for DmlStatement { + fn from(dml: PyDmlStatement) -> Self { + dml.dml + } +} + +impl From for PyDmlStatement { + fn from(dml: DmlStatement) -> PyDmlStatement { + PyDmlStatement { dml } + } +} + +impl LogicalNode for PyDmlStatement { + fn inputs(&self) -> Vec { + vec![PyLogicalPlan::from((*self.dml.input).clone())] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyDmlStatement { + pub fn table_name(&self) -> PyResult { + Ok(self.dml.table_name.to_string()) + } + + pub fn target(&self) -> PyResult { + Ok(PyTableSource { + table_source: self.dml.target.clone(), + }) + } + + pub fn op(&self) -> PyWriteOp { + self.dml.op.clone().into() + } + + pub fn input(&self) -> PyLogicalPlan { + PyLogicalPlan { + plan: self.dml.input.clone(), + } + } + + pub fn output_schema(&self) -> PyDFSchema { + (*self.dml.output_schema).clone().into() + } + + fn __repr__(&self) -> PyResult { + Ok("DmlStatement".to_string()) + } + + fn __name__(&self) -> PyResult { + Ok("DmlStatement".to_string()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + eq, + eq_int, + name = "WriteOp", + module = "datafusion.expr" +)] +pub enum PyWriteOp { + Append, + Overwrite, + Replace, + Update, + Delete, + Ctas, + Truncate, +} + +impl From for PyWriteOp { + fn from(write_op: WriteOp) -> Self { + match write_op { + WriteOp::Insert(InsertOp::Append) => PyWriteOp::Append, + WriteOp::Insert(InsertOp::Overwrite) => PyWriteOp::Overwrite, + WriteOp::Insert(InsertOp::Replace) => PyWriteOp::Replace, + WriteOp::Update => PyWriteOp::Update, + WriteOp::Delete => PyWriteOp::Delete, + WriteOp::Ctas => PyWriteOp::Ctas, + WriteOp::Truncate => PyWriteOp::Truncate, + } + } +} + +impl From for WriteOp { + fn from(py: PyWriteOp) -> Self { + match py { + PyWriteOp::Append => WriteOp::Insert(InsertOp::Append), + PyWriteOp::Overwrite => WriteOp::Insert(InsertOp::Overwrite), + PyWriteOp::Replace => WriteOp::Insert(InsertOp::Replace), + PyWriteOp::Update => WriteOp::Update, + PyWriteOp::Delete => WriteOp::Delete, + PyWriteOp::Ctas => WriteOp::Ctas, + PyWriteOp::Truncate => WriteOp::Truncate, + } + } +} + +#[pymethods] +impl PyWriteOp { + fn name(&self) -> String { + let write_op: WriteOp = self.clone().into(); + write_op.name().to_string() + } +} diff --git a/crates/core/src/expr/drop_catalog_schema.rs b/crates/core/src/expr/drop_catalog_schema.rs new file mode 100644 index 000000000..fd5105332 --- /dev/null +++ b/crates/core/src/expr/drop_catalog_schema.rs @@ -0,0 +1,123 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::common::SchemaReference; +use datafusion::logical_expr::DropCatalogSchema; +use datafusion::sql::TableReference; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "DropCatalogSchema", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDropCatalogSchema { + drop: DropCatalogSchema, +} + +impl From for DropCatalogSchema { + fn from(drop: PyDropCatalogSchema) -> Self { + drop.drop + } +} + +impl From for PyDropCatalogSchema { + fn from(drop: DropCatalogSchema) -> PyDropCatalogSchema { + PyDropCatalogSchema { drop } + } +} + +impl Display for PyDropCatalogSchema { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "DropCatalogSchema") + } +} + +fn parse_schema_reference(name: String) -> PyResult { + match name.into() { + TableReference::Bare { table } => Ok(SchemaReference::Bare { schema: table }), + TableReference::Partial { schema, table } => Ok(SchemaReference::Full { + schema: table, + catalog: schema, + }), + TableReference::Full { + catalog: _, + schema: _, + table: _, + } => Err(PyErr::new::( + "Invalid schema specifier (has 3 parts)".to_string(), + )), + } +} + +#[pymethods] +impl PyDropCatalogSchema { + #[new] + fn new(name: String, schema: PyDFSchema, if_exists: bool, cascade: bool) -> PyResult { + let name = parse_schema_reference(name)?; + Ok(PyDropCatalogSchema { + drop: DropCatalogSchema { + name, + schema: Arc::new(schema.into()), + if_exists, + cascade, + }, + }) + } + + fn name(&self) -> PyResult { + Ok(self.drop.name.to_string()) + } + + fn schema(&self) -> PyDFSchema { + (*self.drop.schema).clone().into() + } + + fn if_exists(&self) -> PyResult { + Ok(self.drop.if_exists) + } + + fn cascade(&self) -> PyResult { + Ok(self.drop.cascade) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("DropCatalogSchema({self})")) + } +} + +impl LogicalNode for PyDropCatalogSchema { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/crates/core/src/expr/drop_function.rs b/crates/core/src/expr/drop_function.rs new file mode 100644 index 000000000..0599dd49e --- /dev/null +++ b/crates/core/src/expr/drop_function.rs @@ -0,0 +1,100 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::DropFunction; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "DropFunction", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDropFunction { + drop: DropFunction, +} + +impl From for DropFunction { + fn from(drop: PyDropFunction) -> Self { + drop.drop + } +} + +impl From for PyDropFunction { + fn from(drop: DropFunction) -> PyDropFunction { + PyDropFunction { drop } + } +} + +impl Display for PyDropFunction { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "DropFunction") + } +} + +#[pymethods] +impl PyDropFunction { + #[new] + fn new(name: String, schema: PyDFSchema, if_exists: bool) -> PyResult { + Ok(PyDropFunction { + drop: DropFunction { + name, + schema: Arc::new(schema.into()), + if_exists, + }, + }) + } + fn name(&self) -> PyResult { + Ok(self.drop.name.clone()) + } + + fn schema(&self) -> PyDFSchema { + (*self.drop.schema).clone().into() + } + + fn if_exists(&self) -> PyResult { + Ok(self.drop.if_exists) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("DropFunction({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("DropFunction".to_string()) + } +} + +impl LogicalNode for PyDropFunction { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/src/expr/drop_table.rs b/crates/core/src/expr/drop_table.rs similarity index 91% rename from src/expr/drop_table.rs rename to crates/core/src/expr/drop_table.rs index 96983c1cf..46fe67465 100644 --- a/src/expr/drop_table.rs +++ b/crates/core/src/expr/drop_table.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::logical_plan::DropTable; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::sql::logical::PyLogicalPlan; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "DropTable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "DropTable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyDropTable { drop: DropTable, @@ -70,7 +76,7 @@ impl PyDropTable { } fn __repr__(&self) -> PyResult { - Ok(format!("DropTable({})", self)) + Ok(format!("DropTable({self})")) } fn __name__(&self) -> PyResult { diff --git a/crates/core/src/expr/drop_view.rs b/crates/core/src/expr/drop_view.rs new file mode 100644 index 000000000..0d0c51f13 --- /dev/null +++ b/crates/core/src/expr/drop_view.rs @@ -0,0 +1,106 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use datafusion::logical_expr::DropView; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "DropView", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDropView { + drop: DropView, +} + +impl From for DropView { + fn from(drop: PyDropView) -> Self { + drop.drop + } +} + +impl From for PyDropView { + fn from(drop: DropView) -> PyDropView { + PyDropView { drop } + } +} + +impl Display for PyDropView { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "DropView: {name:?} if not exist:={if_exists}", + name = self.drop.name, + if_exists = self.drop.if_exists + ) + } +} + +#[pymethods] +impl PyDropView { + #[new] + fn new(name: String, schema: PyDFSchema, if_exists: bool) -> PyResult { + Ok(PyDropView { + drop: DropView { + name: name.into(), + schema: Arc::new(schema.into()), + if_exists, + }, + }) + } + + fn name(&self) -> PyResult { + Ok(self.drop.name.to_string()) + } + + fn schema(&self) -> PyDFSchema { + (*self.drop.schema).clone().into() + } + + fn if_exists(&self) -> PyResult { + Ok(self.drop.if_exists) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("DropView({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("DropView".to_string()) + } +} + +impl LogicalNode for PyDropView { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/src/expr/empty_relation.rs b/crates/core/src/expr/empty_relation.rs similarity index 89% rename from src/expr/empty_relation.rs rename to crates/core/src/expr/empty_relation.rs index a1534ac15..f3c237731 100644 --- a/src/expr/empty_relation.rs +++ b/crates/core/src/expr/empty_relation.rs @@ -15,14 +15,23 @@ // specific language governing permissions and limitations // under the License. -use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan}; -use datafusion::logical_expr::EmptyRelation; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::EmptyRelation; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "EmptyRelation", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "EmptyRelation", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyEmptyRelation { empty: EmptyRelation, @@ -65,7 +74,7 @@ impl PyEmptyRelation { /// Get a String representation of this column fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } fn __name__(&self) -> PyResult { diff --git a/src/expr/exists.rs b/crates/core/src/expr/exists.rs similarity index 91% rename from src/expr/exists.rs rename to crates/core/src/expr/exists.rs index 693357836..d2e816127 100644 --- a/src/expr/exists.rs +++ b/crates/core/src/expr/exists.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use super::subquery::PySubquery; -#[pyclass(name = "Exists", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Exists", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExists { exists: Exists, diff --git a/src/expr/explain.rs b/crates/core/src/expr/explain.rs similarity index 88% rename from src/expr/explain.rs rename to crates/core/src/expr/explain.rs index fc02fe2b5..6259951de 100644 --- a/src/expr/explain.rs +++ b/crates/core/src/expr/explain.rs @@ -17,14 +17,23 @@ use std::fmt::{self, Display, Formatter}; -use datafusion::logical_expr::{logical_plan::Explain, LogicalPlan}; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::{common::df_schema::PyDFSchema, errors::py_type_err, sql::logical::PyLogicalPlan}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::logical_expr::logical_plan::Explain; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; - -#[pyclass(name = "Explain", module = "datafusion.expr", subclass)] +use crate::common::df_schema::PyDFSchema; +use crate::errors::py_type_err; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "Explain", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExplain { explain: Explain, diff --git a/src/expr/extension.rs b/crates/core/src/expr/extension.rs similarity index 90% rename from src/expr/extension.rs rename to crates/core/src/expr/extension.rs index 1e3fbb199..a0b617565 100644 --- a/src/expr/extension.rs +++ b/crates/core/src/expr/extension.rs @@ -16,13 +16,19 @@ // under the License. use datafusion::logical_expr::Extension; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::sql::logical::PyLogicalPlan; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Extension", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Extension", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyExtension { pub node: Extension, diff --git a/src/expr/filter.rs b/crates/core/src/expr/filter.rs similarity index 92% rename from src/expr/filter.rs rename to crates/core/src/expr/filter.rs index 9bdb667cd..67426806d 100644 --- a/src/expr/filter.rs +++ b/crates/core/src/expr/filter.rs @@ -15,16 +15,24 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Filter; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::logical_plan::Filter; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Filter", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Filter", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyFilter { filter: Filter, @@ -72,7 +80,7 @@ impl PyFilter { } fn __repr__(&self) -> String { - format!("Filter({})", self) + format!("Filter({self})") } } diff --git a/src/expr/grouping_set.rs b/crates/core/src/expr/grouping_set.rs similarity index 91% rename from src/expr/grouping_set.rs rename to crates/core/src/expr/grouping_set.rs index 63a1c0b50..549a866ed 100644 --- a/src/expr/grouping_set.rs +++ b/crates/core/src/expr/grouping_set.rs @@ -18,7 +18,13 @@ use datafusion::logical_expr::GroupingSet; use pyo3::prelude::*; -#[pyclass(name = "GroupingSet", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "GroupingSet", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyGroupingSet { grouping_set: GroupingSet, diff --git a/src/expr/in_list.rs b/crates/core/src/expr/in_list.rs similarity index 92% rename from src/expr/in_list.rs rename to crates/core/src/expr/in_list.rs index 5dfd8d8eb..0612cc21e 100644 --- a/src/expr/in_list.rs +++ b/crates/core/src/expr/in_list.rs @@ -15,11 +15,18 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; use datafusion::logical_expr::expr::InList; use pyo3::prelude::*; -#[pyclass(name = "InList", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "InList", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyInList { in_list: InList, diff --git a/src/expr/in_subquery.rs b/crates/core/src/expr/in_subquery.rs similarity index 89% rename from src/expr/in_subquery.rs rename to crates/core/src/expr/in_subquery.rs index 306b68a6e..81a2c5794 100644 --- a/src/expr/in_subquery.rs +++ b/crates/core/src/expr/in_subquery.rs @@ -18,9 +18,16 @@ use datafusion::logical_expr::expr::InSubquery; use pyo3::prelude::*; -use super::{subquery::PySubquery, PyExpr}; +use super::PyExpr; +use super::subquery::PySubquery; -#[pyclass(name = "InSubquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "InSubquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyInSubquery { in_subquery: InSubquery, diff --git a/src/expr/indexed_field.rs b/crates/core/src/expr/indexed_field.rs similarity index 94% rename from src/expr/indexed_field.rs rename to crates/core/src/expr/indexed_field.rs index a22dc6b27..98a90d8d4 100644 --- a/src/expr/indexed_field.rs +++ b/crates/core/src/expr/indexed_field.rs @@ -15,14 +15,21 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; +use std::fmt::{Display, Formatter}; + use datafusion::logical_expr::expr::{GetFieldAccess, GetIndexedField}; use pyo3::prelude::*; -use std::fmt::{Display, Formatter}; use super::literal::PyLiteral; +use crate::expr::PyExpr; -#[pyclass(name = "GetIndexedField", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "GetIndexedField", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyGetIndexedField { indexed_field: GetIndexedField, diff --git a/src/expr/join.rs b/crates/core/src/expr/join.rs similarity index 87% rename from src/expr/join.rs rename to crates/core/src/expr/join.rs index 76ec532e7..b90f2f57d 100644 --- a/src/expr/join.rs +++ b/crates/core/src/expr/join.rs @@ -15,16 +15,20 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::{Join, JoinConstraint, JoinType}; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::common::NullEquality; +use datafusion::logical_expr::logical_plan::{Join, JoinConstraint, JoinType}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; -use crate::expr::{logical_node::LogicalNode, PyExpr}; +use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; #[derive(Debug, Clone, PartialEq, Eq, Hash)] -#[pyclass(name = "JoinType", module = "datafusion.expr")] +#[pyclass(from_py_object, frozen, name = "JoinType", module = "datafusion.expr")] pub struct PyJoinType { join_type: JoinType, } @@ -59,7 +63,12 @@ impl Display for PyJoinType { } #[derive(Debug, Clone, Copy)] -#[pyclass(name = "JoinConstraint", module = "datafusion.expr")] +#[pyclass( + from_py_object, + frozen, + name = "JoinConstraint", + module = "datafusion.expr" +)] pub struct PyJoinConstraint { join_constraint: JoinConstraint, } @@ -86,7 +95,13 @@ impl PyJoinConstraint { } } -#[pyclass(name = "Join", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Join", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyJoin { join: Join, @@ -116,7 +131,7 @@ impl Display for PyJoin { JoinType: {:?} JoinConstraint: {:?} Schema: {:?} - NullEqualsNull: {:?}", + NullEquality: {:?}", &self.join.left, &self.join.right, &self.join.on, @@ -124,7 +139,7 @@ impl Display for PyJoin { &self.join.join_type, &self.join.join_constraint, &self.join.schema, - &self.join.null_equals_null, + &self.join.null_equality, ) } } @@ -173,11 +188,14 @@ impl PyJoin { /// If null_equals_null is true, null == null else null != null fn null_equals_null(&self) -> PyResult { - Ok(self.join.null_equals_null) + match self.join.null_equality { + NullEquality::NullEqualsNothing => Ok(false), + NullEquality::NullEqualsNull => Ok(true), + } } fn __repr__(&self) -> PyResult { - Ok(format!("Join({})", self)) + Ok(format!("Join({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/like.rs b/crates/core/src/expr/like.rs similarity index 90% rename from src/expr/like.rs rename to crates/core/src/expr/like.rs index 2e1f060bd..417dc9182 100644 --- a/src/expr/like.rs +++ b/crates/core/src/expr/like.rs @@ -15,13 +15,20 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::logical_expr::expr::Like; use pyo3::prelude::*; -use std::fmt::{self, Display, Formatter}; use crate::expr::PyExpr; -#[pyclass(name = "Like", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Like", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyLike { like: Like, @@ -75,11 +82,17 @@ impl PyLike { } fn __repr__(&self) -> String { - format!("Like({})", self) + format!("Like({self})") } } -#[pyclass(name = "ILike", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ILike", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyILike { like: Like, @@ -133,11 +146,17 @@ impl PyILike { } fn __repr__(&self) -> String { - format!("Like({})", self) + format!("Like({self})") } } -#[pyclass(name = "SimilarTo", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SimilarTo", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySimilarTo { like: Like, @@ -191,6 +210,6 @@ impl PySimilarTo { } fn __repr__(&self) -> String { - format!("Like({})", self) + format!("Like({self})") } } diff --git a/src/expr/limit.rs b/crates/core/src/expr/limit.rs similarity index 93% rename from src/expr/limit.rs rename to crates/core/src/expr/limit.rs index c2a33ff89..c04b8bfa8 100644 --- a/src/expr/limit.rs +++ b/crates/core/src/expr/limit.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Limit; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::logical_plan::Limit; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Limit", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Limit", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyLimit { limit: Limit, @@ -81,7 +89,7 @@ impl PyLimit { } fn __repr__(&self) -> PyResult { - Ok(format!("Limit({})", self)) + Ok(format!("Limit({self})")) } } diff --git a/src/expr/literal.rs b/crates/core/src/expr/literal.rs similarity index 87% rename from src/expr/literal.rs rename to crates/core/src/expr/literal.rs index a660ac914..9db0f594b 100644 --- a/src/expr/literal.rs +++ b/crates/core/src/expr/literal.rs @@ -15,14 +15,30 @@ // specific language governing permissions and limitations // under the License. -use crate::errors::PyDataFusionError; use datafusion::common::ScalarValue; -use pyo3::{prelude::*, IntoPyObjectExt}; +use datafusion::logical_expr::expr::FieldMetadata; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use crate::errors::PyDataFusionError; -#[pyclass(name = "Literal", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + name = "Literal", + module = "datafusion.expr", + subclass, + frozen +)] #[derive(Clone)] pub struct PyLiteral { pub value: ScalarValue, + pub metadata: Option, +} + +impl PyLiteral { + pub fn new_with_metadata(value: ScalarValue, metadata: Option) -> PyLiteral { + Self { value, metadata } + } } impl From for ScalarValue { @@ -33,7 +49,10 @@ impl From for ScalarValue { impl From for PyLiteral { fn from(value: ScalarValue) -> PyLiteral { - PyLiteral { value } + PyLiteral { + value, + metadata: None, + } } } @@ -61,7 +80,7 @@ impl PyLiteral { extract_scalar_value!(self, Float64) } - pub fn value_decimal128(&mut self) -> PyResult<(Option, u8, i8)> { + pub fn value_decimal128(&self) -> PyResult<(Option, u8, i8)> { match &self.value { ScalarValue::Decimal128(value, precision, scale) => Ok((*value, *precision, *scale)), other => Err(unexpected_literal_value(other)), @@ -112,7 +131,7 @@ impl PyLiteral { extract_scalar_value!(self, Time64Nanosecond) } - pub fn value_timestamp(&mut self) -> PyResult<(Option, Option)> { + pub fn value_timestamp(&self) -> PyResult<(Option, Option)> { match &self.value { ScalarValue::TimestampNanosecond(iv, tz) | ScalarValue::TimestampMicrosecond(iv, tz) diff --git a/src/expr/logical_node.rs b/crates/core/src/expr/logical_node.rs similarity index 100% rename from src/expr/logical_node.rs rename to crates/core/src/expr/logical_node.rs diff --git a/src/expr/placeholder.rs b/crates/core/src/expr/placeholder.rs similarity index 76% rename from src/expr/placeholder.rs rename to crates/core/src/expr/placeholder.rs index 4ac2c47e3..6bd88321c 100644 --- a/src/expr/placeholder.rs +++ b/crates/core/src/expr/placeholder.rs @@ -15,12 +15,20 @@ // specific language governing permissions and limitations // under the License. +use arrow::datatypes::Field; +use arrow::pyarrow::PyArrowType; use datafusion::logical_expr::expr::Placeholder; use pyo3::prelude::*; use crate::common::data_type::PyDataType; -#[pyclass(name = "Placeholder", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Placeholder", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyPlaceholder { placeholder: Placeholder, @@ -40,8 +48,15 @@ impl PyPlaceholder { fn data_type(&self) -> Option { self.placeholder - .data_type + .field .as_ref() - .map(|e| e.clone().into()) + .map(|f| f.data_type().clone().into()) + } + + fn field(&self) -> Option> { + self.placeholder + .field + .as_ref() + .map(|f| f.as_ref().clone().into()) } } diff --git a/src/expr/projection.rs b/crates/core/src/expr/projection.rs similarity index 94% rename from src/expr/projection.rs rename to crates/core/src/expr/projection.rs index dc7e5e3c1..456e06412 100644 --- a/src/expr/projection.rs +++ b/crates/core/src/expr/projection.rs @@ -15,17 +15,25 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Projection; -use datafusion::logical_expr::Expr; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::Expr; +use datafusion::logical_expr::logical_plan::Projection; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; -use crate::expr::logical_node::LogicalNode; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Projection", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Projection", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyProjection { pub projection: Projection, @@ -85,7 +93,7 @@ impl PyProjection { } fn __repr__(&self) -> PyResult { - Ok(format!("Projection({})", self)) + Ok(format!("Projection({self})")) } fn __name__(&self) -> PyResult { diff --git a/crates/core/src/expr/recursive_query.rs b/crates/core/src/expr/recursive_query.rs new file mode 100644 index 000000000..e03137b80 --- /dev/null +++ b/crates/core/src/expr/recursive_query.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 std::fmt::{self, Display, Formatter}; + +use datafusion::logical_expr::RecursiveQuery; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "RecursiveQuery", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyRecursiveQuery { + query: RecursiveQuery, +} + +impl From for RecursiveQuery { + fn from(query: PyRecursiveQuery) -> Self { + query.query + } +} + +impl From for PyRecursiveQuery { + fn from(query: RecursiveQuery) -> PyRecursiveQuery { + PyRecursiveQuery { query } + } +} + +impl Display for PyRecursiveQuery { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "RecursiveQuery {name:?} is_distinct:={is_distinct}", + name = self.query.name, + is_distinct = self.query.is_distinct + ) + } +} + +#[pymethods] +impl PyRecursiveQuery { + #[new] + fn new( + name: String, + static_term: PyLogicalPlan, + recursive_term: PyLogicalPlan, + is_distinct: bool, + ) -> Self { + Self { + query: RecursiveQuery { + name, + static_term: static_term.plan(), + recursive_term: recursive_term.plan(), + is_distinct, + }, + } + } + + fn name(&self) -> PyResult { + Ok(self.query.name.clone()) + } + + fn static_term(&self) -> PyLogicalPlan { + PyLogicalPlan::from((*self.query.static_term).clone()) + } + + fn recursive_term(&self) -> PyLogicalPlan { + PyLogicalPlan::from((*self.query.recursive_term).clone()) + } + + fn is_distinct(&self) -> PyResult { + Ok(self.query.is_distinct) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("RecursiveQuery({self})")) + } + + fn __name__(&self) -> PyResult { + Ok("RecursiveQuery".to_string()) + } +} + +impl LogicalNode for PyRecursiveQuery { + fn inputs(&self) -> Vec { + vec![ + PyLogicalPlan::from((*self.query.static_term).clone()), + PyLogicalPlan::from((*self.query.recursive_term).clone()), + ] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} diff --git a/src/expr/repartition.rs b/crates/core/src/expr/repartition.rs similarity index 86% rename from src/expr/repartition.rs rename to crates/core/src/expr/repartition.rs index 3e782d6af..be39b9978 100644 --- a/src/expr/repartition.rs +++ b/crates/core/src/expr/repartition.rs @@ -17,20 +17,35 @@ use std::fmt::{self, Display, Formatter}; -use datafusion::logical_expr::{logical_plan::Repartition, Expr, Partitioning}; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::{errors::py_type_err, sql::logical::PyLogicalPlan}; - -use super::{logical_node::LogicalNode, PyExpr}; - -#[pyclass(name = "Repartition", module = "datafusion.expr", subclass)] +use datafusion::logical_expr::logical_plan::Repartition; +use datafusion::logical_expr::{Expr, Partitioning}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::PyExpr; +use super::logical_node::LogicalNode; +use crate::errors::py_type_err; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "Repartition", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyRepartition { repartition: Repartition, } -#[pyclass(name = "Partitioning", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Partitioning", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyPartitioning { partitioning: Partitioning, @@ -108,7 +123,7 @@ impl PyRepartition { } fn __repr__(&self) -> PyResult { - Ok(format!("Repartition({})", self)) + Ok(format!("Repartition({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/scalar_subquery.rs b/crates/core/src/expr/scalar_subquery.rs similarity index 91% rename from src/expr/scalar_subquery.rs rename to crates/core/src/expr/scalar_subquery.rs index 9d35f28a9..c7852a4c4 100644 --- a/src/expr/scalar_subquery.rs +++ b/crates/core/src/expr/scalar_subquery.rs @@ -20,7 +20,13 @@ use pyo3::prelude::*; use super::subquery::PySubquery; -#[pyclass(name = "ScalarSubquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ScalarSubquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyScalarSubquery { subquery: Subquery, diff --git a/src/expr/scalar_variable.rs b/crates/core/src/expr/scalar_variable.rs similarity index 76% rename from src/expr/scalar_variable.rs rename to crates/core/src/expr/scalar_variable.rs index 7b50ba241..2d3bc4b76 100644 --- a/src/expr/scalar_variable.rs +++ b/crates/core/src/expr/scalar_variable.rs @@ -15,22 +15,28 @@ // specific language governing permissions and limitations // under the License. -use datafusion::arrow::datatypes::DataType; +use arrow::datatypes::FieldRef; use pyo3::prelude::*; use crate::common::data_type::PyDataType; -#[pyclass(name = "ScalarVariable", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "ScalarVariable", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyScalarVariable { - data_type: DataType, + field: FieldRef, variables: Vec, } impl PyScalarVariable { - pub fn new(data_type: &DataType, variables: &[String]) -> Self { + pub fn new(field: &FieldRef, variables: &[String]) -> Self { Self { - data_type: data_type.to_owned(), + field: field.to_owned(), variables: variables.to_vec(), } } @@ -40,7 +46,7 @@ impl PyScalarVariable { impl PyScalarVariable { /// Get the data type fn data_type(&self) -> PyResult { - Ok(self.data_type.clone().into()) + Ok(self.field.data_type().clone().into()) } fn variables(&self) -> PyResult> { @@ -48,6 +54,6 @@ impl PyScalarVariable { } fn __repr__(&self) -> PyResult { - Ok(format!("{}{:?}", self.data_type, self.variables)) + Ok(format!("{}{:?}", self.field.data_type(), self.variables)) } } diff --git a/crates/core/src/expr/set_comparison.rs b/crates/core/src/expr/set_comparison.rs new file mode 100644 index 000000000..9f0c077e1 --- /dev/null +++ b/crates/core/src/expr/set_comparison.rs @@ -0,0 +1,59 @@ +// 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 datafusion::logical_expr::expr::SetComparison; +use pyo3::prelude::*; + +use super::subquery::PySubquery; +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "SetComparison", + module = "datafusion.set_comparison", + subclass +)] +#[derive(Clone)] +pub struct PySetComparison { + set_comparison: SetComparison, +} + +impl From for PySetComparison { + fn from(set_comparison: SetComparison) -> Self { + PySetComparison { set_comparison } + } +} + +#[pymethods] +impl PySetComparison { + fn expr(&self) -> PyExpr { + (*self.set_comparison.expr).clone().into() + } + + fn subquery(&self) -> PySubquery { + self.set_comparison.subquery.clone().into() + } + + fn op(&self) -> String { + format!("{}", self.set_comparison.op) + } + + fn quantifier(&self) -> String { + format!("{}", self.set_comparison.quantifier) + } +} diff --git a/src/expr/signature.rs b/crates/core/src/expr/signature.rs similarity index 91% rename from src/expr/signature.rs rename to crates/core/src/expr/signature.rs index e85763555..35268e3a9 100644 --- a/src/expr/signature.rs +++ b/crates/core/src/expr/signature.rs @@ -19,7 +19,13 @@ use datafusion::logical_expr::{TypeSignature, Volatility}; use pyo3::prelude::*; #[allow(dead_code)] -#[pyclass(name = "Signature", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Signature", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySignature { type_signature: TypeSignature, diff --git a/src/expr/sort.rs b/crates/core/src/expr/sort.rs similarity index 93% rename from src/expr/sort.rs rename to crates/core/src/expr/sort.rs index ed4947591..7c1e654c5 100644 --- a/src/expr/sort.rs +++ b/crates/core/src/expr/sort.rs @@ -15,17 +15,25 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::common::DataFusionError; use datafusion::logical_expr::logical_plan::Sort; -use pyo3::{prelude::*, IntoPyObjectExt}; -use std::fmt::{self, Display, Formatter}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::expr::sort_expr::PySortExpr; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Sort", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Sort", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySort { sort: Sort, @@ -87,7 +95,7 @@ impl PySort { } fn __repr__(&self) -> PyResult { - Ok(format!("Sort({})", self)) + Ok(format!("Sort({self})")) } } diff --git a/src/expr/sort_expr.rs b/crates/core/src/expr/sort_expr.rs similarity index 93% rename from src/expr/sort_expr.rs rename to crates/core/src/expr/sort_expr.rs index 12f74e4d8..3c3c86bc1 100644 --- a/src/expr/sort_expr.rs +++ b/crates/core/src/expr/sort_expr.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use crate::expr::PyExpr; +use std::fmt::{self, Display, Formatter}; + use datafusion::logical_expr::SortExpr; use pyo3::prelude::*; -use std::fmt::{self, Display, Formatter}; -#[pyclass(name = "SortExpr", module = "datafusion.expr", subclass)] +use crate::expr::PyExpr; + +#[pyclass( + from_py_object, + frozen, + name = "SortExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySortExpr { - sort: SortExpr, + pub(crate) sort: SortExpr, } impl From for SortExpr { @@ -85,6 +93,6 @@ impl PySortExpr { } fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } } diff --git a/crates/core/src/expr/statement.rs b/crates/core/src/expr/statement.rs new file mode 100644 index 000000000..5aa1e4e9c --- /dev/null +++ b/crates/core/src/expr/statement.rs @@ -0,0 +1,558 @@ +// 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 std::sync::Arc; + +use arrow::datatypes::Field; +use arrow::pyarrow::PyArrowType; +use datafusion::logical_expr::{ + Deallocate, Execute, Prepare, ResetVariable, SetVariable, TransactionAccessMode, + TransactionConclusion, TransactionEnd, TransactionIsolationLevel, TransactionStart, +}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use super::PyExpr; +use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "TransactionStart", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyTransactionStart { + transaction_start: TransactionStart, +} + +impl From for PyTransactionStart { + fn from(transaction_start: TransactionStart) -> PyTransactionStart { + PyTransactionStart { transaction_start } + } +} + +impl TryFrom for TransactionStart { + type Error = PyErr; + + fn try_from(py: PyTransactionStart) -> Result { + Ok(py.transaction_start) + } +} + +impl LogicalNode for PyTransactionStart { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "TransactionAccessMode", + module = "datafusion.expr" +)] +pub enum PyTransactionAccessMode { + ReadOnly, + ReadWrite, +} + +impl From for PyTransactionAccessMode { + fn from(access_mode: TransactionAccessMode) -> PyTransactionAccessMode { + match access_mode { + TransactionAccessMode::ReadOnly => PyTransactionAccessMode::ReadOnly, + TransactionAccessMode::ReadWrite => PyTransactionAccessMode::ReadWrite, + } + } +} + +impl TryFrom for TransactionAccessMode { + type Error = PyErr; + + fn try_from(py: PyTransactionAccessMode) -> Result { + match py { + PyTransactionAccessMode::ReadOnly => Ok(TransactionAccessMode::ReadOnly), + PyTransactionAccessMode::ReadWrite => Ok(TransactionAccessMode::ReadWrite), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "TransactionIsolationLevel", + module = "datafusion.expr" +)] +pub enum PyTransactionIsolationLevel { + ReadUncommitted, + ReadCommitted, + RepeatableRead, + Serializable, + Snapshot, +} + +impl From for PyTransactionIsolationLevel { + fn from(isolation_level: TransactionIsolationLevel) -> PyTransactionIsolationLevel { + match isolation_level { + TransactionIsolationLevel::ReadUncommitted => { + PyTransactionIsolationLevel::ReadUncommitted + } + TransactionIsolationLevel::ReadCommitted => PyTransactionIsolationLevel::ReadCommitted, + TransactionIsolationLevel::RepeatableRead => { + PyTransactionIsolationLevel::RepeatableRead + } + TransactionIsolationLevel::Serializable => PyTransactionIsolationLevel::Serializable, + TransactionIsolationLevel::Snapshot => PyTransactionIsolationLevel::Snapshot, + } + } +} + +impl TryFrom for TransactionIsolationLevel { + type Error = PyErr; + + fn try_from(value: PyTransactionIsolationLevel) -> Result { + match value { + PyTransactionIsolationLevel::ReadUncommitted => { + Ok(TransactionIsolationLevel::ReadUncommitted) + } + PyTransactionIsolationLevel::ReadCommitted => { + Ok(TransactionIsolationLevel::ReadCommitted) + } + PyTransactionIsolationLevel::RepeatableRead => { + Ok(TransactionIsolationLevel::RepeatableRead) + } + PyTransactionIsolationLevel::Serializable => { + Ok(TransactionIsolationLevel::Serializable) + } + PyTransactionIsolationLevel::Snapshot => Ok(TransactionIsolationLevel::Snapshot), + } + } +} + +#[pymethods] +impl PyTransactionStart { + #[new] + pub fn new( + access_mode: PyTransactionAccessMode, + isolation_level: PyTransactionIsolationLevel, + ) -> PyResult { + let access_mode = access_mode.try_into()?; + let isolation_level = isolation_level.try_into()?; + Ok(PyTransactionStart { + transaction_start: TransactionStart { + access_mode, + isolation_level, + }, + }) + } + + pub fn access_mode(&self) -> PyResult { + Ok(self.transaction_start.access_mode.clone().into()) + } + + pub fn isolation_level(&self) -> PyResult { + Ok(self.transaction_start.isolation_level.clone().into()) + } +} + +#[pyclass( + from_py_object, + frozen, + name = "TransactionEnd", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyTransactionEnd { + transaction_end: TransactionEnd, +} + +impl From for PyTransactionEnd { + fn from(transaction_end: TransactionEnd) -> PyTransactionEnd { + PyTransactionEnd { transaction_end } + } +} + +impl TryFrom for TransactionEnd { + type Error = PyErr; + + fn try_from(py: PyTransactionEnd) -> Result { + Ok(py.transaction_end) + } +} + +impl LogicalNode for PyTransactionEnd { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[pyclass( + from_py_object, + frozen, + eq, + eq_int, + name = "TransactionConclusion", + module = "datafusion.expr" +)] +pub enum PyTransactionConclusion { + Commit, + Rollback, +} + +impl From for PyTransactionConclusion { + fn from(value: TransactionConclusion) -> Self { + match value { + TransactionConclusion::Commit => PyTransactionConclusion::Commit, + TransactionConclusion::Rollback => PyTransactionConclusion::Rollback, + } + } +} + +impl TryFrom for TransactionConclusion { + type Error = PyErr; + + fn try_from(value: PyTransactionConclusion) -> Result { + match value { + PyTransactionConclusion::Commit => Ok(TransactionConclusion::Commit), + PyTransactionConclusion::Rollback => Ok(TransactionConclusion::Rollback), + } + } +} +#[pymethods] +impl PyTransactionEnd { + #[new] + pub fn new(conclusion: PyTransactionConclusion, chain: bool) -> PyResult { + let conclusion = conclusion.try_into()?; + Ok(PyTransactionEnd { + transaction_end: TransactionEnd { conclusion, chain }, + }) + } + + pub fn conclusion(&self) -> PyResult { + Ok(self.transaction_end.conclusion.clone().into()) + } + + pub fn chain(&self) -> bool { + self.transaction_end.chain + } +} + +#[pyclass( + from_py_object, + frozen, + name = "ResetVariable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyResetVariable { + reset_variable: ResetVariable, +} + +impl From for PyResetVariable { + fn from(reset_variable: ResetVariable) -> PyResetVariable { + PyResetVariable { reset_variable } + } +} + +impl TryFrom for ResetVariable { + type Error = PyErr; + + fn try_from(py: PyResetVariable) -> Result { + Ok(py.reset_variable) + } +} + +impl LogicalNode for PyResetVariable { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyResetVariable { + #[new] + pub fn new(variable: String) -> Self { + PyResetVariable { + reset_variable: ResetVariable { variable }, + } + } + + pub fn variable(&self) -> String { + self.reset_variable.variable.clone() + } +} + +#[pyclass( + from_py_object, + frozen, + name = "SetVariable", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PySetVariable { + set_variable: SetVariable, +} + +impl From for PySetVariable { + fn from(set_variable: SetVariable) -> PySetVariable { + PySetVariable { set_variable } + } +} + +impl TryFrom for SetVariable { + type Error = PyErr; + + fn try_from(py: PySetVariable) -> Result { + Ok(py.set_variable) + } +} + +impl LogicalNode for PySetVariable { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PySetVariable { + #[new] + pub fn new(variable: String, value: String) -> Self { + PySetVariable { + set_variable: SetVariable { variable, value }, + } + } + + pub fn variable(&self) -> String { + self.set_variable.variable.clone() + } + + pub fn value(&self) -> String { + self.set_variable.value.clone() + } +} + +#[pyclass( + from_py_object, + frozen, + name = "Prepare", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyPrepare { + prepare: Prepare, +} + +impl From for PyPrepare { + fn from(prepare: Prepare) -> PyPrepare { + PyPrepare { prepare } + } +} + +impl TryFrom for Prepare { + type Error = PyErr; + + fn try_from(py: PyPrepare) -> Result { + Ok(py.prepare) + } +} + +impl LogicalNode for PyPrepare { + fn inputs(&self) -> Vec { + vec![PyLogicalPlan::from((*self.prepare.input).clone())] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyPrepare { + #[new] + pub fn new(name: String, fields: Vec>, input: PyLogicalPlan) -> Self { + let input = input.plan().clone(); + let fields = fields.into_iter().map(|field| Arc::new(field.0)).collect(); + PyPrepare { + prepare: Prepare { + name, + fields, + input, + }, + } + } + + pub fn name(&self) -> String { + self.prepare.name.clone() + } + + pub fn fields(&self) -> Vec> { + self.prepare + .fields + .clone() + .into_iter() + .map(|f| f.as_ref().clone().into()) + .collect() + } + + pub fn input(&self) -> PyLogicalPlan { + PyLogicalPlan { + plan: self.prepare.input.clone(), + } + } +} + +#[pyclass( + from_py_object, + frozen, + name = "Execute", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyExecute { + execute: Execute, +} + +impl From for PyExecute { + fn from(execute: Execute) -> PyExecute { + PyExecute { execute } + } +} + +impl TryFrom for Execute { + type Error = PyErr; + + fn try_from(py: PyExecute) -> Result { + Ok(py.execute) + } +} + +impl LogicalNode for PyExecute { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyExecute { + #[new] + pub fn new(name: String, parameters: Vec) -> Self { + let parameters = parameters + .into_iter() + .map(|parameter| parameter.into()) + .collect(); + PyExecute { + execute: Execute { name, parameters }, + } + } + + pub fn name(&self) -> String { + self.execute.name.clone() + } + + pub fn parameters(&self) -> Vec { + self.execute + .parameters + .clone() + .into_iter() + .map(|t| t.into()) + .collect() + } +} + +#[pyclass( + from_py_object, + frozen, + name = "Deallocate", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyDeallocate { + deallocate: Deallocate, +} + +impl From for PyDeallocate { + fn from(deallocate: Deallocate) -> PyDeallocate { + PyDeallocate { deallocate } + } +} + +impl TryFrom for Deallocate { + type Error = PyErr; + + fn try_from(py: PyDeallocate) -> Result { + Ok(py.deallocate) + } +} + +impl LogicalNode for PyDeallocate { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyDeallocate { + #[new] + pub fn new(name: String) -> Self { + PyDeallocate { + deallocate: Deallocate { name }, + } + } + + pub fn name(&self) -> String { + self.deallocate.name.clone() + } +} diff --git a/src/expr/subquery.rs b/crates/core/src/expr/subquery.rs similarity index 91% rename from src/expr/subquery.rs rename to crates/core/src/expr/subquery.rs index 5ebfe6927..c6fa83db8 100644 --- a/src/expr/subquery.rs +++ b/crates/core/src/expr/subquery.rs @@ -18,13 +18,19 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::Subquery; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::sql::logical::PyLogicalPlan; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Subquery", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Subquery", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySubquery { subquery: Subquery, @@ -62,7 +68,7 @@ impl PySubquery { } fn __repr__(&self) -> PyResult { - Ok(format!("Subquery({})", self)) + Ok(format!("Subquery({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/subquery_alias.rs b/crates/core/src/expr/subquery_alias.rs similarity index 89% rename from src/expr/subquery_alias.rs rename to crates/core/src/expr/subquery_alias.rs index 267a4d485..a6b09e842 100644 --- a/src/expr/subquery_alias.rs +++ b/crates/core/src/expr/subquery_alias.rs @@ -18,13 +18,20 @@ use std::fmt::{self, Display, Formatter}; use datafusion::logical_expr::SubqueryAlias; -use pyo3::{prelude::*, IntoPyObjectExt}; - -use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "SubqueryAlias", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "SubqueryAlias", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PySubqueryAlias { subquery_alias: SubqueryAlias, @@ -72,7 +79,7 @@ impl PySubqueryAlias { } fn __repr__(&self) -> PyResult { - Ok(format!("SubqueryAlias({})", self)) + Ok(format!("SubqueryAlias({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/table_scan.rs b/crates/core/src/expr/table_scan.rs similarity index 94% rename from src/expr/table_scan.rs rename to crates/core/src/expr/table_scan.rs index 6a0d53f0f..8ba7e4a69 100644 --- a/src/expr/table_scan.rs +++ b/crates/core/src/expr/table_scan.rs @@ -15,16 +15,25 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::common::TableReference; use datafusion::logical_expr::logical_plan::TableScan; -use pyo3::{prelude::*, IntoPyObjectExt}; -use std::fmt::{self, Display, Formatter}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; +use crate::common::df_schema::PyDFSchema; +use crate::expr::PyExpr; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -use crate::{common::df_schema::PyDFSchema, expr::PyExpr}; -#[pyclass(name = "TableScan", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "TableScan", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyTableScan { table_scan: TableScan, @@ -136,7 +145,7 @@ impl PyTableScan { } fn __repr__(&self) -> PyResult { - Ok(format!("TableScan({})", self)) + Ok(format!("TableScan({self})")) } } diff --git a/src/expr/union.rs b/crates/core/src/expr/union.rs similarity index 92% rename from src/expr/union.rs rename to crates/core/src/expr/union.rs index 5a08ccc13..a3b9efe91 100644 --- a/src/expr/union.rs +++ b/crates/core/src/expr/union.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Union; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::logical_plan::Union; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Union", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Union", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnion { union_: Union, @@ -66,7 +74,7 @@ impl PyUnion { } fn __repr__(&self) -> PyResult { - Ok(format!("Union({})", self)) + Ok(format!("Union({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/unnest.rs b/crates/core/src/expr/unnest.rs similarity index 92% rename from src/expr/unnest.rs rename to crates/core/src/expr/unnest.rs index 8e70e0990..880d0a279 100644 --- a/src/expr/unnest.rs +++ b/crates/core/src/expr/unnest.rs @@ -15,15 +15,23 @@ // specific language governing permissions and limitations // under the License. -use datafusion::logical_expr::logical_plan::Unnest; -use pyo3::{prelude::*, IntoPyObjectExt}; use std::fmt::{self, Display, Formatter}; +use datafusion::logical_expr::logical_plan::Unnest; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + use crate::common::df_schema::PyDFSchema; use crate::expr::logical_node::LogicalNode; use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Unnest", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Unnest", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnnest { unnest_: Unnest, @@ -66,7 +74,7 @@ impl PyUnnest { } fn __repr__(&self) -> PyResult { - Ok(format!("Unnest({})", self)) + Ok(format!("Unnest({self})")) } fn __name__(&self) -> PyResult { diff --git a/src/expr/unnest_expr.rs b/crates/core/src/expr/unnest_expr.rs similarity index 91% rename from src/expr/unnest_expr.rs rename to crates/core/src/expr/unnest_expr.rs index 2234d24b1..97feef1d1 100644 --- a/src/expr/unnest_expr.rs +++ b/crates/core/src/expr/unnest_expr.rs @@ -15,13 +15,20 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::logical_expr::expr::Unnest; use pyo3::prelude::*; -use std::fmt::{self, Display, Formatter}; use super::PyExpr; -#[pyclass(name = "UnnestExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "UnnestExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyUnnestExpr { unnest: Unnest, @@ -58,7 +65,7 @@ impl PyUnnestExpr { } fn __repr__(&self) -> PyResult { - Ok(format!("UnnestExpr({})", self)) + Ok(format!("UnnestExpr({self})")) } fn __name__(&self) -> PyResult { diff --git a/crates/core/src/expr/values.rs b/crates/core/src/expr/values.rs new file mode 100644 index 000000000..d40b0e7cf --- /dev/null +++ b/crates/core/src/expr/values.rs @@ -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. + +use std::sync::Arc; + +use datafusion::logical_expr::Values; +use pyo3::prelude::*; +use pyo3::{IntoPyObjectExt, PyErr, PyResult, Python, pyclass}; + +use super::PyExpr; +use super::logical_node::LogicalNode; +use crate::common::df_schema::PyDFSchema; +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "Values", + module = "datafusion.expr", + subclass +)] +#[derive(Clone)] +pub struct PyValues { + values: Values, +} + +impl From for PyValues { + fn from(values: Values) -> PyValues { + PyValues { values } + } +} + +impl TryFrom for Values { + type Error = PyErr; + + fn try_from(py: PyValues) -> Result { + Ok(py.values) + } +} + +impl LogicalNode for PyValues { + fn inputs(&self) -> Vec { + vec![] + } + + fn to_variant<'py>(&self, py: Python<'py>) -> PyResult> { + self.clone().into_bound_py_any(py) + } +} + +#[pymethods] +impl PyValues { + #[new] + pub fn new(schema: PyDFSchema, values: Vec>) -> PyResult { + let values = values + .into_iter() + .map(|row| row.into_iter().map(|expr| expr.into()).collect()) + .collect(); + Ok(PyValues { + values: Values { + schema: Arc::new(schema.into()), + values, + }, + }) + } + + pub fn schema(&self) -> PyResult { + Ok((*self.values.schema).clone().into()) + } + + pub fn values(&self) -> Vec> { + self.values + .values + .clone() + .into_iter() + .map(|row| row.into_iter().map(|expr| expr.into()).collect()) + .collect() + } +} diff --git a/src/expr/window.rs b/crates/core/src/expr/window.rs similarity index 85% rename from src/expr/window.rs rename to crates/core/src/expr/window.rs index 13deaec25..92d909bfc 100644 --- a/src/expr/window.rs +++ b/crates/core/src/expr/window.rs @@ -15,31 +15,42 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{self, Display, Formatter}; + use datafusion::common::{DataFusionError, ScalarValue}; -use datafusion::logical_expr::expr::WindowFunction; use datafusion::logical_expr::{Expr, Window, WindowFrame, WindowFrameBound, WindowFrameUnits}; -use pyo3::{prelude::*, IntoPyObjectExt}; -use std::fmt::{self, Display, Formatter}; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyNotImplementedError; +use pyo3::prelude::*; +use super::py_expr_list; use crate::common::data_type::PyScalarValue; use crate::common::df_schema::PyDFSchema; -use crate::errors::{py_type_err, PyDataFusionResult}; -use crate::expr::logical_node::LogicalNode; -use crate::expr::sort_expr::{py_sort_expr_list, PySortExpr}; +use crate::errors::{PyDataFusionResult, py_type_err}; use crate::expr::PyExpr; +use crate::expr::logical_node::LogicalNode; +use crate::expr::sort_expr::{PySortExpr, py_sort_expr_list}; use crate::sql::logical::PyLogicalPlan; -use super::py_expr_list; - -use crate::errors::py_datafusion_err; - -#[pyclass(name = "WindowExpr", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowExpr", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyWindowExpr { window: Window, } -#[pyclass(name = "WindowFrame", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowFrame", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyWindowFrame { window_frame: WindowFrame, @@ -57,7 +68,13 @@ impl From for PyWindowFrame { } } -#[pyclass(name = "WindowFrameBound", module = "datafusion.expr", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowFrameBound", + module = "datafusion.expr", + subclass +)] #[derive(Clone)] pub struct PyWindowFrameBound { frame_bound: WindowFrameBound, @@ -118,7 +135,9 @@ impl PyWindowExpr { /// Returns order by columns in a window function expression pub fn get_sort_exprs(&self, expr: PyExpr) -> PyResult> { match expr.expr.unalias() { - Expr::WindowFunction(WindowFunction { order_by, .. }) => py_sort_expr_list(&order_by), + Expr::WindowFunction(boxed_window_fn) => { + py_sort_expr_list(&boxed_window_fn.params.order_by) + } other => Err(not_window_function_err(other)), } } @@ -126,8 +145,8 @@ impl PyWindowExpr { /// Return partition by columns in a window function expression pub fn get_partition_exprs(&self, expr: PyExpr) -> PyResult> { match expr.expr.unalias() { - Expr::WindowFunction(WindowFunction { partition_by, .. }) => { - py_expr_list(&partition_by) + Expr::WindowFunction(boxed_window_fn) => { + py_expr_list(&boxed_window_fn.params.partition_by) } other => Err(not_window_function_err(other)), } @@ -136,7 +155,7 @@ impl PyWindowExpr { /// Return input args for window function pub fn get_args(&self, expr: PyExpr) -> PyResult> { match expr.expr.unalias() { - Expr::WindowFunction(WindowFunction { args, .. }) => py_expr_list(&args), + Expr::WindowFunction(boxed_window_fn) => py_expr_list(&boxed_window_fn.params.args), other => Err(not_window_function_err(other)), } } @@ -144,7 +163,7 @@ impl PyWindowExpr { /// Return window function name pub fn window_func_name(&self, expr: PyExpr) -> PyResult { match expr.expr.unalias() { - Expr::WindowFunction(WindowFunction { fun, .. }) => Ok(fun.to_string()), + Expr::WindowFunction(boxed_window_fn) => Ok(boxed_window_fn.fun.to_string()), other => Err(not_window_function_err(other)), } } @@ -152,7 +171,9 @@ impl PyWindowExpr { /// Returns a Pywindow frame for a given window function expression pub fn get_frame(&self, expr: PyExpr) -> Option { match expr.expr.unalias() { - Expr::WindowFunction(WindowFunction { window_frame, .. }) => Some(window_frame.into()), + Expr::WindowFunction(boxed_window_fn) => { + Some(boxed_window_fn.params.window_frame.into()) + } _ => None, } } @@ -181,10 +202,7 @@ impl PyWindowFrame { "range" => WindowFrameUnits::Range, "groups" => WindowFrameUnits::Groups, _ => { - return Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - units, - )))); + return Err(PyNotImplementedError::new_err(format!("{units:?}"))); } }; let start_bound = match start_bound { @@ -193,10 +211,7 @@ impl PyWindowFrame { WindowFrameUnits::Range => WindowFrameBound::Preceding(ScalarValue::UInt64(None)), WindowFrameUnits::Rows => WindowFrameBound::Preceding(ScalarValue::UInt64(None)), WindowFrameUnits::Groups => { - return Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - units, - )))); + return Err(PyNotImplementedError::new_err(format!("{units:?}"))); } }, }; @@ -206,10 +221,7 @@ impl PyWindowFrame { WindowFrameUnits::Rows => WindowFrameBound::Following(ScalarValue::UInt64(None)), WindowFrameUnits::Range => WindowFrameBound::Following(ScalarValue::UInt64(None)), WindowFrameUnits::Groups => { - return Err(py_datafusion_err(DataFusionError::NotImplemented(format!( - "{:?}", - units, - )))); + return Err(PyNotImplementedError::new_err(format!("{units:?}"))); } }, }; @@ -233,7 +245,7 @@ impl PyWindowFrame { /// Get a String representation of this window frame fn __repr__(&self) -> String { - format!("{}", self) + format!("{self}") } } diff --git a/src/functions.rs b/crates/core/src/functions.rs similarity index 85% rename from src/functions.rs rename to crates/core/src/functions.rs index 8fac239b4..c32134054 100644 --- a/src/functions.rs +++ b/crates/core/src/functions.rs @@ -15,30 +15,27 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; + +use datafusion::common::{Column, ScalarValue, TableReference}; +use datafusion::execution::FunctionRegistry; use datafusion::functions_aggregate::all_default_aggregate_functions; use datafusion::functions_window::all_default_window_functions; -use datafusion::logical_expr::ExprFunctionExt; -use datafusion::logical_expr::WindowFrame; -use pyo3::{prelude::*, wrap_pyfunction}; - -use crate::common::data_type::NullTreatment; -use crate::common::data_type::PyScalarValue; +use datafusion::logical_expr::expr::{ + Alias, FieldMetadata, NullTreatment as DFNullTreatment, WindowFunction, WindowFunctionParams, +}; +use datafusion::logical_expr::{Expr, ExprFunctionExt, WindowFrame, WindowFunctionDefinition, lit}; +use datafusion::{functions, functions_aggregate, functions_window}; +use pyo3::prelude::*; +use pyo3::wrap_pyfunction; + +use crate::common::data_type::{NullTreatment, PyScalarValue}; use crate::context::PySessionContext; -use crate::errors::PyDataFusionError; -use crate::errors::PyDataFusionResult; +use crate::errors::{PyDataFusionError, PyDataFusionResult}; +use crate::expr::PyExpr; use crate::expr::conditional_expr::PyCaseBuilder; -use crate::expr::sort_expr::to_sort_expressions; -use crate::expr::sort_expr::PySortExpr; +use crate::expr::sort_expr::{PySortExpr, to_sort_expressions}; use crate::expr::window::PyWindowFrame; -use crate::expr::PyExpr; -use datafusion::common::{Column, ScalarValue, TableReference}; -use datafusion::execution::FunctionRegistry; -use datafusion::functions; -use datafusion::functions_aggregate; -use datafusion::functions_window; -use datafusion::logical_expr::expr::Alias; -use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment; -use datafusion::logical_expr::{expr::WindowFunction, lit, Expr, WindowFunctionDefinition}; fn add_builder_fns_to_aggregate( agg_fn: Expr, @@ -100,7 +97,7 @@ fn array_cat(exprs: Vec) -> PyExpr { #[pyo3(signature = (array, element, index=None))] fn array_position(array: PyExpr, element: PyExpr, index: Option) -> PyExpr { let index = ScalarValue::Int64(index); - let index = Expr::Literal(index); + let index = Expr::Literal(index, None); datafusion::functions_nested::expr_fn::array_position(array.into(), element.into(), index) .into() } @@ -192,6 +189,29 @@ fn regexp_count( .into()) } +#[pyfunction] +#[pyo3(signature = (values, regex, start=None, n=None, flags=None, subexpr=None))] +/// Returns the position in a string where the specified occurrence of a regular expression is located +fn regexp_instr( + values: PyExpr, + regex: PyExpr, + start: Option, + n: Option, + flags: Option, + subexpr: Option, +) -> PyResult { + Ok(functions::expr_fn::regexp_instr( + values.into(), + regex.into(), + start.map(|x| x.expr).or(Some(lit(1))), + n.map(|x| x.expr).or(Some(lit(1))), + None, + flags.map(|x| x.expr).or(Some(lit(""))), + subexpr.map(|x| x.expr).or(Some(lit(0))), + ) + .into()) +} + /// Creates a new Sort Expr #[pyfunction] fn order_by(expr: PyExpr, asc: bool, nulls_first: bool) -> PyResult { @@ -204,10 +224,14 @@ fn order_by(expr: PyExpr, asc: bool, nulls_first: bool) -> PyResult /// Creates a new Alias Expr #[pyfunction] -fn alias(expr: PyExpr, name: &str) -> PyResult { +#[pyo3(signature = (expr, name, metadata=None))] +fn alias(expr: PyExpr, name: &str, metadata: Option>) -> PyResult { let relation: Option = None; + let metadata = metadata.map(|m| FieldMetadata::new(m.into_iter().collect())); Ok(PyExpr { - expr: datafusion::logical_expr::Expr::Alias(Alias::new(expr.expr, relation, name)), + expr: datafusion::logical_expr::Expr::Alias( + Alias::new(expr.expr, relation, name).with_metadata(metadata), + ), }) } @@ -215,27 +239,20 @@ fn alias(expr: PyExpr, name: &str) -> PyResult { #[pyfunction] fn col(name: &str) -> PyResult { Ok(PyExpr { - expr: datafusion::logical_expr::Expr::Column(Column { - relation: None, - name: name.to_string(), - }), + expr: datafusion::logical_expr::Expr::Column(Column::new_unqualified(name)), }) } /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression. #[pyfunction] fn case(expr: PyExpr) -> PyResult { - Ok(PyCaseBuilder { - case_builder: datafusion::logical_expr::case(expr.expr), - }) + Ok(PyCaseBuilder::new(Some(expr))) } /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression. #[pyfunction] fn when(when: PyExpr, then: PyExpr) -> PyResult { - Ok(PyCaseBuilder { - case_builder: datafusion::logical_expr::when(when.expr, then.expr), - }) + Ok(PyCaseBuilder::new(None).when(when, then)) } /// Helper function to find the appropriate window function. @@ -314,14 +331,17 @@ fn find_window_fn( } /// Creates a new Window function expression +#[allow(clippy::too_many_arguments)] #[pyfunction] -#[pyo3(signature = (name, args, partition_by=None, order_by=None, window_frame=None, ctx=None))] +#[pyo3(signature = (name, args, partition_by=None, order_by=None, window_frame=None, filter=None, distinct=false, ctx=None))] fn window( name: &str, args: Vec, partition_by: Option>, order_by: Option>, window_frame: Option, + filter: Option, + distinct: bool, ctx: Option, ) -> PyResult { let fun = find_window_fn(name, ctx)?; @@ -329,24 +349,29 @@ fn window( let window_frame = window_frame .map(|w| w.into()) .unwrap_or(WindowFrame::new(order_by.as_ref().map(|v| !v.is_empty()))); + let filter = filter.map(|f| f.expr.into()); Ok(PyExpr { - expr: datafusion::logical_expr::Expr::WindowFunction(WindowFunction { + expr: datafusion::logical_expr::Expr::WindowFunction(Box::new(WindowFunction { fun, - args: args.into_iter().map(|x| x.expr).collect::>(), - partition_by: partition_by - .unwrap_or_default() - .into_iter() - .map(|x| x.expr) - .collect::>(), - order_by: order_by - .unwrap_or_default() - .into_iter() - .map(|x| x.into()) - .collect::>(), - window_frame, - null_treatment: None, - }), + params: WindowFunctionParams { + args: args.into_iter().map(|x| x.expr).collect::>(), + partition_by: partition_by + .unwrap_or_default() + .into_iter() + .map(|x| x.expr) + .collect::>(), + order_by: order_by + .unwrap_or_default() + .into_iter() + .map(|x| x.into()) + .collect::>(), + window_frame, + filter, + distinct, + null_treatment: None, + }, + })), }) } @@ -375,27 +400,6 @@ macro_rules! aggregate_function { }; } -macro_rules! aggregate_function_vec_args { - ($NAME: ident) => { - aggregate_function_vec_args!($NAME, expr); - }; - ($NAME: ident, $($arg:ident)*) => { - #[pyfunction] - #[pyo3(signature = ($($arg),*, distinct=None, filter=None, order_by=None, null_treatment=None))] - fn $NAME( - $($arg: PyExpr),*, - distinct: Option, - filter: Option, - order_by: Option>, - null_treatment: Option - ) -> PyDataFusionResult { - let agg_fn = functions_aggregate::expr_fn::$NAME(vec![$($arg.into()),*]); - - add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment) - } - }; -} - /// Generates a [pyo3] wrapper for [datafusion::functions::expr_fn] /// /// These functions have explicit named arguments. @@ -460,7 +464,11 @@ macro_rules! array_fn { expr_fn!(abs, num); expr_fn!(acos, num); expr_fn!(acosh, num); -expr_fn!(ascii, arg1, "Returns the numeric code of the first character of the argument. In UTF8 encoding, returns the Unicode code point of the character. In other multibyte encodings, the argument must be an ASCII character."); +expr_fn!( + ascii, + arg1, + "Returns the numeric code of the first character of the argument. In UTF8 encoding, returns the Unicode code point of the character. In other multibyte encodings, the argument must be an ASCII character." +); expr_fn!(asin, num); expr_fn!(asinh, num); expr_fn!(atan, num); @@ -471,7 +479,10 @@ expr_fn!( arg, "Returns number of bits in the string (8 times the octet_length)." ); -expr_fn_vec!(btrim, "Removes the longest string containing only characters in characters (a space by default) from the start and end of string."); +expr_fn_vec!( + btrim, + "Removes the longest string containing only characters in characters (a space by default) from the start and end of string." +); expr_fn!(cbrt, num); expr_fn!(ceil, num); expr_fn!( @@ -494,7 +505,11 @@ expr_fn!(exp, num); expr_fn!(factorial, num); expr_fn!(floor, num); expr_fn!(gcd, x y); -expr_fn!(initcap, string, "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters."); +expr_fn!( + initcap, + string, + "Converts the first letter of each word to upper case and the rest to lower case. Words are sequences of alphanumeric characters separated by non-alphanumeric characters." +); expr_fn!(isnan, num); expr_fn!(iszero, num); expr_fn!(levenshtein, string1 string2); @@ -505,8 +520,14 @@ expr_fn!(log, base num); expr_fn!(log10, num); expr_fn!(log2, num); expr_fn!(lower, arg1, "Converts the string to all lower case"); -expr_fn_vec!(lpad, "Extends the string to length length by prepending the characters fill (a space by default). If the string is already longer than length then it is truncated (on the right)."); -expr_fn_vec!(ltrim, "Removes the longest string containing only characters in characters (a space by default) from the start of string."); +expr_fn_vec!( + lpad, + "Extends the string to length length by prepending the characters fill (a space by default). If the string is already longer than length then it is truncated (on the right)." +); +expr_fn_vec!( + ltrim, + "Removes the longest string containing only characters in characters (a space by default) from the start of string." +); expr_fn!( md5, input_arg, @@ -523,7 +544,11 @@ expr_fn!( "Returns x if x is not NULL otherwise returns y." ); expr_fn!(nullif, arg_1 arg_2); -expr_fn!(octet_length, args, "Returns number of bytes in the string. Since this version of the function accepts type character directly, it will not strip trailing spaces."); +expr_fn!( + octet_length, + args, + "Returns number of bytes in the string. Since this version of the function accepts type character directly, it will not strip trailing spaces." +); expr_fn_vec!(overlay); expr_fn!(pi); expr_fn!(power, base exponent); @@ -541,8 +566,14 @@ expr_fn!( ); expr_fn!(right, string n, "Returns last n characters in the string, or when n is negative, returns all but first |n| characters."); expr_fn_vec!(round); -expr_fn_vec!(rpad, "Extends the string to length length by appending the characters fill (a space by default). If the string is already longer than length then it is truncated."); -expr_fn_vec!(rtrim, "Removes the longest string containing only characters in characters (a space by default) from the end of string."); +expr_fn_vec!( + rpad, + "Extends the string to length length by appending the characters fill (a space by default). If the string is already longer than length then it is truncated." +); +expr_fn_vec!( + rtrim, + "Removes the longest string containing only characters in characters (a space by default) from the end of string." +); expr_fn!(sha224, input_arg1); expr_fn!(sha256, input_arg1); expr_fn!(sha384, input_arg1); @@ -570,6 +601,9 @@ expr_fn!( "Converts the number to its equivalent hexadecimal representation." ); expr_fn!(now); +expr_fn_vec!(to_date); +expr_fn_vec!(to_local_time); +expr_fn_vec!(to_time); expr_fn_vec!(to_timestamp); expr_fn_vec!(to_timestamp_millis); expr_fn_vec!(to_timestamp_nanos); @@ -582,9 +616,13 @@ expr_fn!(date_part, part date); expr_fn!(date_trunc, part date); expr_fn!(date_bin, stride source origin); expr_fn!(make_date, year month day); +expr_fn!(to_char, datetime format); expr_fn!(translate, string from to, "Replaces each character in string that matches a character in the from set with the corresponding character in the to set. If from is longer than to, occurrences of the extra characters in from are deleted."); -expr_fn_vec!(trim, "Removes the longest string containing only characters in characters (a space by default) from the start, end, or both ends (BOTH is the default) of string."); +expr_fn_vec!( + trim, + "Removes the longest string containing only characters in characters (a space by default) from the start, end, or both ends (BOTH is the default) of string." +); expr_fn_vec!(trunc); expr_fn!(upper, arg1, "Converts the string to all upper case."); expr_fn!(uuid); @@ -663,43 +701,57 @@ aggregate_function!(approx_median); // aggregate_function!(grouping); #[pyfunction] -#[pyo3(signature = (expression, percentile, num_centroids=None, filter=None))] +#[pyo3(signature = (sort_expression, percentile, num_centroids=None, filter=None))] pub fn approx_percentile_cont( - expression: PyExpr, + sort_expression: PySortExpr, percentile: f64, num_centroids: Option, // enforces optional arguments at the end, currently filter: Option, ) -> PyDataFusionResult { - let args = if let Some(num_centroids) = num_centroids { - vec![expression.expr, lit(percentile), lit(num_centroids)] - } else { - vec![expression.expr, lit(percentile)] - }; - let udaf = functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf(); - let agg_fn = udaf.call(args); + let agg_fn = functions_aggregate::expr_fn::approx_percentile_cont( + sort_expression.sort, + lit(percentile), + num_centroids.map(lit), + ); add_builder_fns_to_aggregate(agg_fn, None, filter, None, None) } #[pyfunction] -#[pyo3(signature = (expression, weight, percentile, filter=None))] +#[pyo3(signature = (sort_expression, weight, percentile, num_centroids=None, filter=None))] pub fn approx_percentile_cont_with_weight( - expression: PyExpr, + sort_expression: PySortExpr, weight: PyExpr, percentile: f64, + num_centroids: Option, filter: Option, ) -> PyDataFusionResult { let agg_fn = functions_aggregate::expr_fn::approx_percentile_cont_with_weight( - expression.expr, + sort_expression.sort, weight.expr, lit(percentile), + num_centroids.map(lit), ); add_builder_fns_to_aggregate(agg_fn, None, filter, None, None) } -aggregate_function_vec_args!(last_value); +// We handle last_value explicitly because the signature expects an order_by +// https://github.com/apache/datafusion/issues/12376 +#[pyfunction] +#[pyo3(signature = (expr, distinct=None, filter=None, order_by=None, null_treatment=None))] +pub fn last_value( + expr: PyExpr, + distinct: Option, + filter: Option, + order_by: Option>, + null_treatment: Option, +) -> PyDataFusionResult { + // If we initialize the UDAF with order_by directly, then it gets over-written by the builder + let agg_fn = functions_aggregate::expr_fn::last_value(expr.expr, vec![]); + add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment) +} // We handle first_value explicitly because the signature expects an order_by // https://github.com/apache/datafusion/issues/12376 #[pyfunction] @@ -712,7 +764,7 @@ pub fn first_value( null_treatment: Option, ) -> PyDataFusionResult { // If we initialize the UDAF with order_by directly, then it gets over-written by the builder - let agg_fn = functions_aggregate::expr_fn::first_value(expr.expr, None); + let agg_fn = functions_aggregate::expr_fn::first_value(expr.expr, vec![]); add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment) } @@ -939,7 +991,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(left))?; m.add_wrapped(wrap_pyfunction!(length))?; m.add_wrapped(wrap_pyfunction!(ln))?; - m.add_wrapped(wrap_pyfunction!(log))?; + m.add_wrapped(wrap_pyfunction!(self::log))?; m.add_wrapped(wrap_pyfunction!(log10))?; m.add_wrapped(wrap_pyfunction!(log2))?; m.add_wrapped(wrap_pyfunction!(lower))?; @@ -963,6 +1015,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(radians))?; m.add_wrapped(wrap_pyfunction!(random))?; m.add_wrapped(wrap_pyfunction!(regexp_count))?; + m.add_wrapped(wrap_pyfunction!(regexp_instr))?; m.add_wrapped(wrap_pyfunction!(regexp_like))?; m.add_wrapped(wrap_pyfunction!(regexp_match))?; m.add_wrapped(wrap_pyfunction!(regexp_replace))?; @@ -996,6 +1049,10 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(tan))?; m.add_wrapped(wrap_pyfunction!(tanh))?; m.add_wrapped(wrap_pyfunction!(to_hex))?; + m.add_wrapped(wrap_pyfunction!(to_char))?; + m.add_wrapped(wrap_pyfunction!(to_date))?; + m.add_wrapped(wrap_pyfunction!(to_local_time))?; + m.add_wrapped(wrap_pyfunction!(to_time))?; m.add_wrapped(wrap_pyfunction!(to_timestamp))?; m.add_wrapped(wrap_pyfunction!(to_timestamp_millis))?; m.add_wrapped(wrap_pyfunction!(to_timestamp_nanos))?; diff --git a/src/lib.rs b/crates/core/src/lib.rs similarity index 80% rename from src/lib.rs rename to crates/core/src/lib.rs index ce93ff0c3..fc2d006d3 100644 --- a/src/lib.rs +++ b/crates/core/src/lib.rs @@ -15,19 +15,16 @@ // specific language governing permissions and limitations // under the License. -#[cfg(feature = "mimalloc")] -use mimalloc::MiMalloc; -use pyo3::prelude::*; - // Re-export Apache Arrow DataFusion dependencies -pub use datafusion; -pub use datafusion::common as datafusion_common; -pub use datafusion::logical_expr as datafusion_expr; -pub use datafusion::optimizer; -pub use datafusion::sql as datafusion_sql; - +pub use datafusion::{ + self, common as datafusion_common, logical_expr as datafusion_expr, optimizer, + sql as datafusion_sql, +}; #[cfg(feature = "substrait")] pub use datafusion_substrait; +#[cfg(feature = "mimalloc")] +use mimalloc::MiMalloc; +use pyo3::prelude::*; #[allow(clippy::borrow_deref_ref)] pub mod catalog; @@ -46,53 +43,63 @@ pub mod errors; pub mod expr; #[allow(clippy::borrow_deref_ref)] mod functions; +mod options; pub mod physical_plan; mod pyarrow_filter_expression; pub mod pyarrow_util; mod record_batch; pub mod sql; pub mod store; +pub mod table; +pub mod unparser; +mod array; #[cfg(feature = "substrait")] pub mod substrait; #[allow(clippy::borrow_deref_ref)] mod udaf; #[allow(clippy::borrow_deref_ref)] mod udf; +pub mod udtf; mod udwf; -pub mod utils; #[cfg(feature = "mimalloc")] #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; -// Used to define Tokio Runtime as a Python module attribute -pub(crate) struct TokioRuntime(tokio::runtime::Runtime); - /// Low-level DataFusion internal package. /// /// The higher-level public API is defined in pure python files under the /// datafusion directory. #[pymodule] fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { + // Initialize logging + pyo3_log::init(); + // Register the python classes - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + let catalog = PyModule::new(py, "catalog")?; + catalog::init_module(&catalog)?; + m.add_submodule(&catalog)?; + // Register `common` as a submodule. Matching `datafusion-common` https://docs.rs/datafusion-common/latest/datafusion_common/ let common = PyModule::new(py, "common")?; common::init_module(&common)?; @@ -103,6 +110,10 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { expr::init_module(&expr)?; m.add_submodule(&expr)?; + let unparser = PyModule::new(py, "unparser")?; + unparser::init_module(&unparser)?; + m.add_submodule(&unparser)?; + // Register the functions as a submodule let funcs = PyModule::new(py, "functions")?; functions::init_module(&funcs)?; @@ -112,6 +123,10 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { store::init_module(&store)?; m.add_submodule(&store)?; + let options = PyModule::new(py, "options")?; + options::init_module(&options)?; + m.add_submodule(&options)?; + // Register substrait as a submodule #[cfg(feature = "substrait")] setup_substrait_module(py, &m)?; diff --git a/crates/core/src/options.rs b/crates/core/src/options.rs new file mode 100644 index 000000000..6b6037695 --- /dev/null +++ b/crates/core/src/options.rs @@ -0,0 +1,159 @@ +// 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::datatypes::{DataType, Schema}; +use arrow::pyarrow::PyArrowType; +use datafusion::prelude::CsvReadOptions; +use pyo3::prelude::{PyModule, PyModuleMethods}; +use pyo3::{Bound, PyResult, pyclass, pymethods}; + +use crate::context::parse_file_compression_type; +use crate::errors::PyDataFusionError; +use crate::expr::sort_expr::PySortExpr; + +/// Options for reading CSV files +#[pyclass(name = "CsvReadOptions", module = "datafusion.options", frozen)] +pub struct PyCsvReadOptions { + pub has_header: bool, + pub delimiter: u8, + pub quote: u8, + pub terminator: Option, + pub escape: Option, + pub comment: Option, + pub newlines_in_values: bool, + pub schema: Option>, + pub schema_infer_max_records: usize, + pub file_extension: String, + pub table_partition_cols: Vec<(String, PyArrowType)>, + pub file_compression_type: String, + pub file_sort_order: Vec>, + pub null_regex: Option, + pub truncated_rows: bool, +} + +#[pymethods] +impl PyCsvReadOptions { + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = ( + has_header=true, + delimiter=b',', + quote=b'"', + terminator=None, + escape=None, + comment=None, + newlines_in_values=false, + schema=None, + schema_infer_max_records=1000, + file_extension=".csv".to_string(), + table_partition_cols=vec![], + file_compression_type="".to_string(), + file_sort_order=vec![], + null_regex=None, + truncated_rows=false + ))] + #[new] + fn new( + has_header: bool, + delimiter: u8, + quote: u8, + terminator: Option, + escape: Option, + comment: Option, + newlines_in_values: bool, + schema: Option>, + schema_infer_max_records: usize, + file_extension: String, + table_partition_cols: Vec<(String, PyArrowType)>, + file_compression_type: String, + file_sort_order: Vec>, + null_regex: Option, + truncated_rows: bool, + ) -> Self { + Self { + has_header, + delimiter, + quote, + terminator, + escape, + comment, + newlines_in_values, + schema, + schema_infer_max_records, + file_extension, + table_partition_cols, + file_compression_type, + file_sort_order, + null_regex, + truncated_rows, + } + } +} + +impl<'a> TryFrom<&'a PyCsvReadOptions> for CsvReadOptions<'a> { + type Error = PyDataFusionError; + + fn try_from(value: &'a PyCsvReadOptions) -> Result, Self::Error> { + let partition_cols: Vec<(String, DataType)> = value + .table_partition_cols + .iter() + .map(|(name, dtype)| (name.clone(), dtype.0.clone())) + .collect(); + + let compression = parse_file_compression_type(Some(value.file_compression_type.clone()))?; + + let sort_order: Vec> = value + .file_sort_order + .iter() + .map(|inner| { + inner + .iter() + .map(|sort_expr| sort_expr.sort.clone()) + .collect() + }) + .collect(); + + // Explicit struct initialization to catch upstream changes + let mut options = CsvReadOptions { + has_header: value.has_header, + delimiter: value.delimiter, + quote: value.quote, + terminator: value.terminator, + escape: value.escape, + comment: value.comment, + newlines_in_values: value.newlines_in_values, + schema: None, // Will be set separately due to lifetime constraints + schema_infer_max_records: value.schema_infer_max_records, + file_extension: value.file_extension.as_str(), + table_partition_cols: partition_cols, + file_compression_type: compression, + file_sort_order: sort_order, + null_regex: value.null_regex.clone(), + truncated_rows: value.truncated_rows, + }; + + // Set schema separately to handle the lifetime + options.schema = value.schema.as_ref().map(|s| &s.0); + + Ok(options) + } +} + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + + Ok(()) +} diff --git a/src/physical_plan.rs b/crates/core/src/physical_plan.rs similarity index 85% rename from src/physical_plan.rs rename to crates/core/src/physical_plan.rs index f0be45c6a..8674a8b55 100644 --- a/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -15,16 +15,25 @@ // specific language governing permissions and limitations // under the License. -use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties}; -use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; -use prost::Message; use std::sync::Arc; -use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyBytes}; - -use crate::{context::PySessionContext, errors::PyDataFusionResult}; - -#[pyclass(name = "ExecutionPlan", module = "datafusion", subclass)] +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; +use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use prost::Message; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use crate::context::PySessionContext; +use crate::errors::PyDataFusionResult; + +#[pyclass( + from_py_object, + frozen, + name = "ExecutionPlan", + module = "datafusion", + subclass +)] #[derive(Debug, Clone)] pub struct PyExecutionPlan { pub plan: Arc, @@ -74,17 +83,16 @@ impl PyExecutionPlan { ctx: PySessionContext, proto_msg: Bound<'_, PyBytes>, ) -> PyDataFusionResult { - let bytes: &[u8] = proto_msg.extract()?; + let bytes: &[u8] = proto_msg.extract().map_err(Into::::into)?; let proto_plan = datafusion_proto::protobuf::PhysicalPlanNode::decode(bytes).map_err(|e| { PyRuntimeError::new_err(format!( - "Unable to decode logical node from serialized bytes: {}", - e + "Unable to decode logical node from serialized bytes: {e}" )) })?; let codec = DefaultPhysicalExtensionCodec {}; - let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), &codec)?; + let plan = proto_plan.try_into_physical_plan(ctx.ctx.task_ctx().as_ref(), &codec)?; Ok(Self::new(plan)) } diff --git a/src/pyarrow_filter_expression.rs b/crates/core/src/pyarrow_filter_expression.rs similarity index 94% rename from src/pyarrow_filter_expression.rs rename to crates/core/src/pyarrow_filter_expression.rs index 4b4c86597..e3b4b6009 100644 --- a/src/pyarrow_filter_expression.rs +++ b/crates/core/src/pyarrow_filter_expression.rs @@ -15,21 +15,21 @@ // specific language governing permissions and limitations // under the License. -/// Converts a Datafusion logical plan expression (Expr) into a PyArrow compute expression -use pyo3::{prelude::*, IntoPyObjectExt}; - use std::convert::TryFrom; use std::result::Result; use datafusion::common::{Column, ScalarValue}; -use datafusion::logical_expr::{expr::InList, Between, BinaryExpr, Expr, Operator}; +use datafusion::logical_expr::expr::InList; +use datafusion::logical_expr::{Between, BinaryExpr, Expr, Operator}; +/// Converts a Datafusion logical plan expression (Expr) into a PyArrow compute expression +use pyo3::{IntoPyObjectExt, prelude::*}; use crate::errors::{PyDataFusionError, PyDataFusionResult}; use crate::pyarrow_util::scalar_to_pyarrow; #[derive(Debug)] #[repr(transparent)] -pub(crate) struct PyArrowFilterExpression(PyObject); +pub(crate) struct PyArrowFilterExpression(Py); fn operator_to_py<'py>( operator: &Operator, @@ -47,7 +47,7 @@ fn operator_to_py<'py>( _ => { return Err(PyDataFusionError::Common(format!( "Unsupported operator {operator:?}" - ))) + ))); } }; Ok(py_op) @@ -57,11 +57,11 @@ fn extract_scalar_list<'py>( exprs: &[Expr], py: Python<'py>, ) -> PyDataFusionResult>> { - let ret = exprs + exprs .iter() .map(|expr| match expr { // TODO: should we also leverage `ScalarValue::to_pyarrow` here? - Expr::Literal(v) => match v { + Expr::Literal(v, _) => match v { // The unwraps here are for infallible conversions ScalarValue::Boolean(Some(b)) => Ok(b.into_bound_py_any(py)?), ScalarValue::Int8(Some(i)) => Ok(i.into_bound_py_any(py)?), @@ -83,12 +83,11 @@ fn extract_scalar_list<'py>( "Only a list of Literals are supported got {expr:?}" ))), }) - .collect(); - ret + .collect() } impl PyArrowFilterExpression { - pub fn inner(&self) -> &PyObject { + pub fn inner(&self) -> &Py { &self.0 } } @@ -101,12 +100,12 @@ impl TryFrom<&Expr> for PyArrowFilterExpression { // isin, is_null, and is_valid (~is_null) are methods of pyarrow.dataset.Expression // https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Expression.html#pyarrow-dataset-expression fn try_from(expr: &Expr) -> Result { - Python::with_gil(|py| { + Python::attach(|py| { let pc = Python::import(py, "pyarrow.compute")?; let op_module = Python::import(py, "operator")?; let pc_expr: PyDataFusionResult> = match expr { Expr::Column(Column { name, .. }) => Ok(pc.getattr("field")?.call1((name,))?), - Expr::Literal(scalar) => Ok(scalar_to_pyarrow(scalar, py)?.into_bound(py)), + Expr::Literal(scalar, _) => Ok(scalar_to_pyarrow(scalar, py)?), Expr::BinaryExpr(BinaryExpr { left, op, right }) => { let operator = operator_to_py(op, &op_module)?; let left = PyArrowFilterExpression::try_from(left.as_ref())?.0; diff --git a/crates/core/src/pyarrow_util.rs b/crates/core/src/pyarrow_util.rs new file mode 100644 index 000000000..1401a4938 --- /dev/null +++ b/crates/core/src/pyarrow_util.rs @@ -0,0 +1,163 @@ +// 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. + +//! Conversions between PyArrow and DataFusion types + +use std::sync::Arc; + +use arrow::array::{Array, ArrayData, ArrayRef, ListArray, make_array}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::Field; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; +use datafusion::common::exec_err; +use datafusion::scalar::ScalarValue; +use pyo3::types::{PyAnyMethods, PyList}; +use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyErr, PyResult, Python}; + +use crate::common::data_type::PyScalarValue; +use crate::errors::PyDataFusionError; + +/// Helper function to turn an Array into a ScalarValue. If ``as_list_array`` is true, +/// the array will be turned into a ``ListArray``. Otherwise, we extract the first value +/// from the array. +fn array_to_scalar_value(array: ArrayRef, as_list_array: bool) -> PyResult { + if as_list_array { + let field = Arc::new(Field::new_list_field( + array.data_type().clone(), + array.nulls().is_some(), + )); + let offsets = OffsetBuffer::from_lengths(vec![array.len()]); + let list_array = ListArray::new(field, offsets, array, None); + Ok(PyScalarValue(ScalarValue::List(Arc::new(list_array)))) + } else { + let scalar = ScalarValue::try_from_array(&array, 0).map_err(PyDataFusionError::from)?; + Ok(PyScalarValue(scalar)) + } +} + +/// Helper function to take any Python object that contains an Arrow PyCapsule +/// interface and attempt to extract a scalar value from it. If `as_list_array` +/// is true, the array will be turned into a ``ListArray``. Otherwise, we extract +/// the first value from the array. +fn pyobj_extract_scalar_via_capsule( + value: &Bound<'_, PyAny>, + as_list_array: bool, +) -> PyResult { + let array_data = ArrayData::from_pyarrow_bound(value)?; + let array = make_array(array_data); + + array_to_scalar_value(array, as_list_array) +} + +impl FromPyArrow for PyScalarValue { + fn from_pyarrow_bound(value: &Bound<'_, PyAny>) -> PyResult { + let py = value.py(); + let pyarrow_mod = py.import("pyarrow"); + + // Is it a PyArrow object? + if let Ok(pa) = pyarrow_mod.as_ref() { + let scalar_type = pa.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + let typ = value.getattr("type")?; + + // construct pyarrow array from the python value and pyarrow type + let factory = py.import("pyarrow")?.getattr("array")?; + let args = PyList::new(py, [value])?; + let array = factory.call1((args, typ))?; + + return pyobj_extract_scalar_via_capsule(&array, false); + } + + let array_type = pa.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Is it a NanoArrow scalar? + if let Ok(na) = py.import("nanoarrow") { + let scalar_type = py.import("nanoarrow.array")?.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + return pyobj_extract_scalar_via_capsule(value, false); + } + let array_type = na.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Is it a arro3 scalar? + if let Ok(arro3) = py.import("arro3").and_then(|arro3| arro3.getattr("core")) { + let scalar_type = arro3.getattr("Scalar")?; + if value.is_instance(&scalar_type)? { + return pyobj_extract_scalar_via_capsule(value, false); + } + let array_type = arro3.getattr("Array")?; + if value.is_instance(&array_type)? { + return pyobj_extract_scalar_via_capsule(value, true); + } + } + + // Does it have a PyCapsule interface but isn't one of our known libraries? + // If so do our "best guess". Try checking type name, and if that fails + // return a single value if the length is 1 and return a List value otherwise + if value.hasattr("__arrow_c_array__")? { + let type_name = value.get_type().repr()?; + if type_name.contains("Scalar")? { + return pyobj_extract_scalar_via_capsule(value, false); + } + if type_name.contains("Array")? { + return pyobj_extract_scalar_via_capsule(value, true); + } + + let array_data = ArrayData::from_pyarrow_bound(value)?; + let array = make_array(array_data); + + let as_array_list = array.len() != 1; + return array_to_scalar_value(array, as_array_list); + } + + // Last attempt - try to create a PyArrow scalar from a plain Python object + if let Ok(pa) = pyarrow_mod.as_ref() { + let scalar = pa.call_method1("scalar", (value,))?; + + PyScalarValue::from_pyarrow_bound(&scalar) + } else { + exec_err!("Unable to import scalar value").map_err(PyDataFusionError::from)? + } + } +} + +impl<'source> FromPyObject<'_, 'source> for PyScalarValue { + type Error = PyErr; + + fn extract(value: Borrowed<'_, 'source, PyAny>) -> Result { + Self::from_pyarrow_bound(&value) + } +} + +pub fn scalar_to_pyarrow<'py>( + scalar: &ScalarValue, + py: Python<'py>, +) -> PyResult> { + let array = scalar.to_array().map_err(PyDataFusionError::from)?; + // convert to pyarrow array using C data interface + let pyarray = array.to_data().to_pyarrow(py)?; + let pyscalar = pyarray.call_method1("__getitem__", (0,))?; + + Ok(pyscalar) +} diff --git a/src/record_batch.rs b/crates/core/src/record_batch.rs similarity index 75% rename from src/record_batch.rs rename to crates/core/src/record_batch.rs index ec61c263f..0492c6c76 100644 --- a/src/record_batch.rs +++ b/crates/core/src/record_batch.rs @@ -17,25 +17,26 @@ use std::sync::Arc; -use crate::errors::PyDataFusionError; -use crate::utils::wait_for_future; use datafusion::arrow::pyarrow::ToPyArrow; use datafusion::arrow::record_batch::RecordBatch; use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion_python_util::wait_for_future; use futures::StreamExt; use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration}; use pyo3::prelude::*; -use pyo3::{pyclass, pymethods, PyObject, PyResult, Python}; +use pyo3::{PyAny, PyResult, Python, pyclass, pymethods}; use tokio::sync::Mutex; -#[pyclass(name = "RecordBatch", module = "datafusion", subclass)] +use crate::errors::PyDataFusionError; + +#[pyclass(name = "RecordBatch", module = "datafusion", subclass, frozen)] pub struct PyRecordBatch { batch: RecordBatch, } #[pymethods] impl PyRecordBatch { - fn to_pyarrow(&self, py: Python) -> PyResult { + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { self.batch.to_pyarrow(py) } } @@ -46,7 +47,7 @@ impl From for PyRecordBatch { } } -#[pyclass(name = "RecordBatchStream", module = "datafusion", subclass)] +#[pyclass(name = "RecordBatchStream", module = "datafusion", subclass, frozen)] pub struct PyRecordBatchStream { stream: Arc>, } @@ -61,12 +62,12 @@ impl PyRecordBatchStream { #[pymethods] impl PyRecordBatchStream { - fn next(&mut self, py: Python) -> PyResult { + fn next(&self, py: Python) -> PyResult { let stream = self.stream.clone(); - wait_for_future(py, next_stream(stream, true)) + wait_for_future(py, next_stream(stream, true))? } - fn __next__(&mut self, py: Python) -> PyResult { + fn __next__(&self, py: Python) -> PyResult { self.next(py) } @@ -84,15 +85,21 @@ impl PyRecordBatchStream { } } +/// Polls the next batch from a `SendableRecordBatchStream`, converting the `Option>` form. +pub(crate) async fn poll_next_batch( + stream: &mut SendableRecordBatchStream, +) -> datafusion::error::Result> { + stream.next().await.transpose() +} + async fn next_stream( stream: Arc>, sync: bool, ) -> PyResult { let mut stream = stream.lock().await; - match stream.next().await { - Some(Ok(batch)) => Ok(batch.into()), - Some(Err(e)) => Err(PyDataFusionError::from(e))?, - None => { + match poll_next_batch(&mut stream).await { + Ok(Some(batch)) => Ok(batch.into()), + Ok(None) => { // Depending on whether the iteration is sync or not, we raise either a // StopIteration or a StopAsyncIteration if sync { @@ -101,5 +108,6 @@ async fn next_stream( Err(PyStopAsyncIteration::new_err("stream exhausted")) } } + Err(e) => Err(PyDataFusionError::from(e))?, } } diff --git a/src/sql.rs b/crates/core/src/sql.rs similarity index 97% rename from src/sql.rs rename to crates/core/src/sql.rs index 9f1fe81be..dea9b566a 100644 --- a/src/sql.rs +++ b/crates/core/src/sql.rs @@ -17,3 +17,4 @@ pub mod exceptions; pub mod logical; +pub(crate) mod util; diff --git a/src/sql/exceptions.rs b/crates/core/src/sql/exceptions.rs similarity index 100% rename from src/sql/exceptions.rs rename to crates/core/src/sql/exceptions.rs diff --git a/src/sql/logical.rs b/crates/core/src/sql/logical.rs similarity index 54% rename from src/sql/logical.rs rename to crates/core/src/sql/logical.rs index 96561c434..631aa9b09 100644 --- a/src/sql/logical.rs +++ b/crates/core/src/sql/logical.rs @@ -17,33 +17,64 @@ use std::sync::Arc; +use datafusion::logical_expr::{DdlStatement, LogicalPlan, Statement}; +use datafusion_proto::logical_plan::{AsLogicalPlan, DefaultLogicalExtensionCodec}; +use prost::Message; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use crate::context::PySessionContext; use crate::errors::PyDataFusionResult; use crate::expr::aggregate::PyAggregate; use crate::expr::analyze::PyAnalyze; +use crate::expr::copy_to::PyCopyTo; +use crate::expr::create_catalog::PyCreateCatalog; +use crate::expr::create_catalog_schema::PyCreateCatalogSchema; +use crate::expr::create_external_table::PyCreateExternalTable; +use crate::expr::create_function::PyCreateFunction; +use crate::expr::create_index::PyCreateIndex; +use crate::expr::create_memory_table::PyCreateMemoryTable; +use crate::expr::create_view::PyCreateView; +use crate::expr::describe_table::PyDescribeTable; use crate::expr::distinct::PyDistinct; +use crate::expr::dml::PyDmlStatement; +use crate::expr::drop_catalog_schema::PyDropCatalogSchema; +use crate::expr::drop_function::PyDropFunction; +use crate::expr::drop_table::PyDropTable; +use crate::expr::drop_view::PyDropView; use crate::expr::empty_relation::PyEmptyRelation; use crate::expr::explain::PyExplain; use crate::expr::extension::PyExtension; use crate::expr::filter::PyFilter; use crate::expr::join::PyJoin; use crate::expr::limit::PyLimit; +use crate::expr::logical_node::LogicalNode; use crate::expr::projection::PyProjection; +use crate::expr::recursive_query::PyRecursiveQuery; +use crate::expr::repartition::PyRepartition; use crate::expr::sort::PySort; +use crate::expr::statement::{ + PyDeallocate, PyExecute, PyPrepare, PyResetVariable, PySetVariable, PyTransactionEnd, + PyTransactionStart, +}; use crate::expr::subquery::PySubquery; use crate::expr::subquery_alias::PySubqueryAlias; use crate::expr::table_scan::PyTableScan; +use crate::expr::union::PyUnion; use crate::expr::unnest::PyUnnest; +use crate::expr::values::PyValues; use crate::expr::window::PyWindowExpr; -use crate::{context::PySessionContext, errors::py_unsupported_variant_err}; -use datafusion::logical_expr::LogicalPlan; -use datafusion_proto::logical_plan::{AsLogicalPlan, DefaultLogicalExtensionCodec}; -use prost::Message; -use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyBytes}; - -use crate::expr::logical_node::LogicalNode; -#[pyclass(name = "LogicalPlan", module = "datafusion", subclass)] -#[derive(Debug, Clone)] +#[pyclass( + from_py_object, + frozen, + name = "LogicalPlan", + module = "datafusion", + subclass, + eq +)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct PyLogicalPlan { pub(crate) plan: Arc, } @@ -82,18 +113,57 @@ impl PyLogicalPlan { LogicalPlan::SubqueryAlias(plan) => PySubqueryAlias::from(plan.clone()).to_variant(py), LogicalPlan::Unnest(plan) => PyUnnest::from(plan.clone()).to_variant(py), LogicalPlan::Window(plan) => PyWindowExpr::from(plan.clone()).to_variant(py), - LogicalPlan::Repartition(_) - | LogicalPlan::Union(_) - | LogicalPlan::Statement(_) - | LogicalPlan::Values(_) - | LogicalPlan::Dml(_) - | LogicalPlan::Ddl(_) - | LogicalPlan::Copy(_) - | LogicalPlan::DescribeTable(_) - | LogicalPlan::RecursiveQuery(_) => Err(py_unsupported_variant_err(format!( - "Conversion of variant not implemented: {:?}", - self.plan - ))), + LogicalPlan::Repartition(plan) => PyRepartition::from(plan.clone()).to_variant(py), + LogicalPlan::Union(plan) => PyUnion::from(plan.clone()).to_variant(py), + LogicalPlan::Statement(plan) => match plan { + Statement::TransactionStart(plan) => { + PyTransactionStart::from(plan.clone()).to_variant(py) + } + Statement::TransactionEnd(plan) => { + PyTransactionEnd::from(plan.clone()).to_variant(py) + } + Statement::SetVariable(plan) => PySetVariable::from(plan.clone()).to_variant(py), + Statement::ResetVariable(plan) => { + PyResetVariable::from(plan.clone()).to_variant(py) + } + Statement::Prepare(plan) => PyPrepare::from(plan.clone()).to_variant(py), + Statement::Execute(plan) => PyExecute::from(plan.clone()).to_variant(py), + Statement::Deallocate(plan) => PyDeallocate::from(plan.clone()).to_variant(py), + }, + LogicalPlan::Values(plan) => PyValues::from(plan.clone()).to_variant(py), + LogicalPlan::Dml(plan) => PyDmlStatement::from(plan.clone()).to_variant(py), + LogicalPlan::Ddl(plan) => match plan { + DdlStatement::CreateExternalTable(plan) => { + PyCreateExternalTable::from(plan.clone()).to_variant(py) + } + DdlStatement::CreateMemoryTable(plan) => { + PyCreateMemoryTable::from(plan.clone()).to_variant(py) + } + DdlStatement::CreateView(plan) => PyCreateView::from(plan.clone()).to_variant(py), + DdlStatement::CreateCatalogSchema(plan) => { + PyCreateCatalogSchema::from(plan.clone()).to_variant(py) + } + DdlStatement::CreateCatalog(plan) => { + PyCreateCatalog::from(plan.clone()).to_variant(py) + } + DdlStatement::CreateIndex(plan) => PyCreateIndex::from(plan.clone()).to_variant(py), + DdlStatement::DropTable(plan) => PyDropTable::from(plan.clone()).to_variant(py), + DdlStatement::DropView(plan) => PyDropView::from(plan.clone()).to_variant(py), + DdlStatement::DropCatalogSchema(plan) => { + PyDropCatalogSchema::from(plan.clone()).to_variant(py) + } + DdlStatement::CreateFunction(plan) => { + PyCreateFunction::from(plan.clone()).to_variant(py) + } + DdlStatement::DropFunction(plan) => { + PyDropFunction::from(plan.clone()).to_variant(py) + } + }, + LogicalPlan::Copy(plan) => PyCopyTo::from(plan.clone()).to_variant(py), + LogicalPlan::DescribeTable(plan) => PyDescribeTable::from(plan.clone()).to_variant(py), + LogicalPlan::RecursiveQuery(plan) => { + PyRecursiveQuery::from(plan.clone()).to_variant(py) + } } } @@ -140,17 +210,16 @@ impl PyLogicalPlan { ctx: PySessionContext, proto_msg: Bound<'_, PyBytes>, ) -> PyDataFusionResult { - let bytes: &[u8] = proto_msg.extract()?; + let bytes: &[u8] = proto_msg.extract().map_err(Into::::into)?; let proto_plan = datafusion_proto::protobuf::LogicalPlanNode::decode(bytes).map_err(|e| { PyRuntimeError::new_err(format!( - "Unable to decode logical node from serialized bytes: {}", - e + "Unable to decode logical node from serialized bytes: {e}" )) })?; let codec = DefaultLogicalExtensionCodec {}; - let plan = proto_plan.try_into_logical_plan(&ctx.ctx, &codec)?; + let plan = proto_plan.try_into_logical_plan(&ctx.ctx.task_ctx(), &codec)?; Ok(Self::new(plan)) } } diff --git a/crates/core/src/sql/util.rs b/crates/core/src/sql/util.rs new file mode 100644 index 000000000..d1e8964f8 --- /dev/null +++ b/crates/core/src/sql/util.rs @@ -0,0 +1,87 @@ +// 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 std::collections::HashMap; + +use datafusion::common::{DataFusionError, exec_err, plan_datafusion_err}; +use datafusion::logical_expr::sqlparser::dialect::dialect_from_str; +use datafusion::sql::sqlparser::dialect::Dialect; +use datafusion::sql::sqlparser::parser::Parser; +use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer}; + +fn tokens_from_replacements( + placeholder: &str, + replacements: &HashMap>, +) -> Option> { + if let Some(pattern) = placeholder.strip_prefix("$") { + replacements.get(pattern).cloned() + } else { + None + } +} + +fn get_tokens_for_string_replacement( + dialect: &dyn Dialect, + replacements: HashMap, +) -> Result>, DataFusionError> { + replacements + .into_iter() + .map(|(name, value)| { + let tokens = Tokenizer::new(dialect, &value) + .tokenize() + .map_err(|err| DataFusionError::External(err.into()))?; + Ok((name, tokens)) + }) + .collect() +} + +pub(crate) fn replace_placeholders_with_strings( + query: &str, + dialect: &str, + replacements: HashMap, +) -> Result { + let dialect = dialect_from_str(dialect) + .ok_or_else(|| plan_datafusion_err!("Unsupported SQL dialect: {dialect}."))?; + + let replacements = get_tokens_for_string_replacement(dialect.as_ref(), replacements)?; + + let tokens = Tokenizer::new(dialect.as_ref(), query) + .tokenize() + .map_err(|err| DataFusionError::External(err.into()))?; + + let replaced_tokens = tokens + .into_iter() + .flat_map(|token| { + if let Token::Placeholder(placeholder) = &token { + tokens_from_replacements(placeholder, &replacements).unwrap_or(vec![token]) + } else { + vec![token] + } + }) + .collect::>(); + + let statement = Parser::new(dialect.as_ref()) + .with_tokens(replaced_tokens) + .parse_statements() + .map_err(|err| DataFusionError::External(Box::new(err)))?; + + if statement.len() != 1 { + return exec_err!("placeholder replacement should return exactly one statement"); + } + + Ok(statement[0].to_string()) +} diff --git a/src/store.rs b/crates/core/src/store.rs similarity index 87% rename from src/store.rs rename to crates/core/src/store.rs index 1e5fab472..8535e83b7 100644 --- a/src/store.rs +++ b/crates/core/src/store.rs @@ -17,14 +17,13 @@ use std::sync::Arc; -use pyo3::prelude::*; - use object_store::aws::{AmazonS3, AmazonS3Builder}; use object_store::azure::{MicrosoftAzure, MicrosoftAzureBuilder}; use object_store::gcp::{GoogleCloudStorage, GoogleCloudStorageBuilder}; use object_store::http::{HttpBuilder, HttpStore}; use object_store::local::LocalFileSystem; use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; use url::Url; #[derive(FromPyObject)] @@ -36,7 +35,13 @@ pub enum StorageContexts { HTTP(PyHttpContext), } -#[pyclass(name = "LocalFileSystem", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "LocalFileSystem", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyLocalFileSystemContext { pub inner: Arc, @@ -62,7 +67,13 @@ impl PyLocalFileSystemContext { } } -#[pyclass(name = "MicrosoftAzure", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "MicrosoftAzure", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyMicrosoftAzureContext { pub inner: Arc, @@ -72,7 +83,7 @@ pub struct PyMicrosoftAzureContext { #[pymethods] impl PyMicrosoftAzureContext { #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (container_name, account=None, access_key=None, bearer_token=None, client_id=None, client_secret=None, tenant_id=None, sas_query_pairs=None, use_emulator=None, allow_http=None))] + #[pyo3(signature = (container_name, account=None, access_key=None, bearer_token=None, client_id=None, client_secret=None, tenant_id=None, sas_query_pairs=None, use_emulator=None, allow_http=None, use_fabric_endpoint=None))] #[new] fn new( container_name: String, @@ -85,6 +96,7 @@ impl PyMicrosoftAzureContext { sas_query_pairs: Option>, use_emulator: Option, allow_http: Option, + use_fabric_endpoint: Option, ) -> Self { let mut builder = MicrosoftAzureBuilder::from_env().with_container_name(&container_name); @@ -123,6 +135,10 @@ impl PyMicrosoftAzureContext { builder = builder.with_allow_http(allow_http); } + if let Some(use_fabric_endpoint) = use_fabric_endpoint { + builder = builder.with_use_fabric_endpoint(use_fabric_endpoint); + } + Self { inner: Arc::new( builder @@ -134,7 +150,13 @@ impl PyMicrosoftAzureContext { } } -#[pyclass(name = "GoogleCloud", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "GoogleCloud", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyGoogleCloudContext { pub inner: Arc, @@ -164,7 +186,13 @@ impl PyGoogleCloudContext { } } -#[pyclass(name = "AmazonS3", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "AmazonS3", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyAmazonS3Context { pub inner: Arc, @@ -174,13 +202,14 @@ pub struct PyAmazonS3Context { #[pymethods] impl PyAmazonS3Context { #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (bucket_name, region=None, access_key_id=None, secret_access_key=None, endpoint=None, allow_http=false, imdsv1_fallback=false))] + #[pyo3(signature = (bucket_name, region=None, access_key_id=None, secret_access_key=None, session_token=None, endpoint=None, allow_http=false, imdsv1_fallback=false))] #[new] fn new( bucket_name: String, region: Option, access_key_id: Option, secret_access_key: Option, + session_token: Option, endpoint: Option, //retry_config: RetryConfig, allow_http: bool, @@ -201,6 +230,10 @@ impl PyAmazonS3Context { builder = builder.with_secret_access_key(secret_access_key); }; + if let Some(session_token) = session_token { + builder = builder.with_token(session_token); + } + if let Some(endpoint) = endpoint { builder = builder.with_endpoint(endpoint); }; @@ -223,7 +256,13 @@ impl PyAmazonS3Context { } } -#[pyclass(name = "Http", module = "datafusion.store", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Http", + module = "datafusion.store", + subclass +)] #[derive(Debug, Clone)] pub struct PyHttpContext { pub url: String, diff --git a/src/substrait.rs b/crates/core/src/substrait.rs similarity index 73% rename from src/substrait.rs rename to crates/core/src/substrait.rs index 1fefc0bbd..27e446f48 100644 --- a/src/substrait.rs +++ b/crates/core/src/substrait.rs @@ -15,19 +15,25 @@ // specific language governing permissions and limitations // under the License. -use pyo3::{prelude::*, types::PyBytes}; - -use crate::context::PySessionContext; -use crate::errors::{py_datafusion_err, PyDataFusionError, PyDataFusionResult}; -use crate::sql::logical::PyLogicalPlan; -use crate::utils::wait_for_future; - +use datafusion_python_util::wait_for_future; use datafusion_substrait::logical_plan::{consumer, producer}; use datafusion_substrait::serializer; use datafusion_substrait::substrait::proto::Plan; use prost::Message; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use crate::context::PySessionContext; +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err}; +use crate::sql::logical::PyLogicalPlan; -#[pyclass(name = "Plan", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Plan", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PyPlan { pub plan: Plan, @@ -35,13 +41,26 @@ pub struct PyPlan { #[pymethods] impl PyPlan { - fn encode(&self, py: Python) -> PyResult { + fn encode(&self, py: Python) -> PyResult> { let mut proto_bytes = Vec::::new(); self.plan .encode(&mut proto_bytes) .map_err(PyDataFusionError::EncodeError)?; Ok(PyBytes::new(py, &proto_bytes).into()) } + + /// Get the JSON representation of the substrait plan + fn to_json(&self) -> PyDataFusionResult { + let json = serde_json::to_string_pretty(&self.plan).map_err(to_datafusion_err)?; + Ok(json) + } + + /// Parse a Substrait Plan from its JSON representation + #[staticmethod] + fn from_json(json: &str) -> PyDataFusionResult { + let plan: Plan = serde_json::from_str(json).map_err(to_datafusion_err)?; + Ok(PyPlan { plan }) + } } impl From for Plan { @@ -59,7 +78,13 @@ impl From for PyPlan { /// A PySubstraitSerializer is a representation of a Serializer that is capable of both serializing /// a `LogicalPlan` instance to Substrait Protobuf bytes and also deserialize Substrait Protobuf bytes /// to a valid `LogicalPlan` instance. -#[pyclass(name = "Serde", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Serde", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitSerializer; @@ -72,7 +97,7 @@ impl PySubstraitSerializer { path: &str, py: Python, ) -> PyDataFusionResult<()> { - wait_for_future(py, serializer::serialize(sql, &ctx.ctx, path))?; + wait_for_future(py, serializer::serialize(sql, &ctx.ctx, path))??; Ok(()) } @@ -83,7 +108,7 @@ impl PySubstraitSerializer { py: Python, ) -> PyDataFusionResult { PySubstraitSerializer::serialize_bytes(sql, ctx, py).and_then(|proto_bytes| { - let proto_bytes = proto_bytes.bind(py).downcast::().unwrap(); + let proto_bytes = proto_bytes.bind(py).cast::().unwrap(); PySubstraitSerializer::deserialize_bytes(proto_bytes.as_bytes().to_vec(), py) }) } @@ -93,25 +118,32 @@ impl PySubstraitSerializer { sql: &str, ctx: PySessionContext, py: Python, - ) -> PyDataFusionResult { - let proto_bytes: Vec = wait_for_future(py, serializer::serialize_bytes(sql, &ctx.ctx))?; + ) -> PyDataFusionResult> { + let proto_bytes: Vec = + wait_for_future(py, serializer::serialize_bytes(sql, &ctx.ctx))??; Ok(PyBytes::new(py, &proto_bytes).into()) } #[staticmethod] pub fn deserialize(path: &str, py: Python) -> PyDataFusionResult { - let plan = wait_for_future(py, serializer::deserialize(path))?; + let plan = wait_for_future(py, serializer::deserialize(path))??; Ok(PyPlan { plan: *plan }) } #[staticmethod] pub fn deserialize_bytes(proto_bytes: Vec, py: Python) -> PyDataFusionResult { - let plan = wait_for_future(py, serializer::deserialize_bytes(proto_bytes))?; + let plan = wait_for_future(py, serializer::deserialize_bytes(proto_bytes))??; Ok(PyPlan { plan: *plan }) } } -#[pyclass(name = "Producer", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Producer", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitProducer; @@ -128,7 +160,13 @@ impl PySubstraitProducer { } } -#[pyclass(name = "Consumer", module = "datafusion.substrait", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "Consumer", + module = "datafusion.substrait", + subclass +)] #[derive(Debug, Clone)] pub struct PySubstraitConsumer; @@ -137,13 +175,13 @@ impl PySubstraitConsumer { /// Convert Substrait Plan to DataFusion DataFrame #[staticmethod] pub fn from_substrait_plan( - ctx: &mut PySessionContext, + ctx: &PySessionContext, plan: PyPlan, py: Python, ) -> PyDataFusionResult { let session_state = ctx.ctx.state(); let result = consumer::from_substrait_plan(&session_state, &plan.plan); - let logical_plan = wait_for_future(py, result)?; + let logical_plan = wait_for_future(py, result)??; Ok(PyLogicalPlan::new(logical_plan)) } } diff --git a/crates/core/src/table.rs b/crates/core/src/table.rs new file mode 100644 index 000000000..623349771 --- /dev/null +++ b/crates/core/src/table.rs @@ -0,0 +1,261 @@ +// 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 std::any::Any; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::pyarrow::ToPyArrow; +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProviderFactory}; +use datafusion::common::Column; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::logical_expr::{ + CreateExternalTable, Expr, LogicalPlanBuilder, TableProviderFilterPushDown, +}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::DataFrame; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_python_util::{create_logical_extension_capsule, table_provider_from_pycapsule}; +use pyo3::IntoPyObjectExt; +use pyo3::prelude::*; + +use crate::context::PySessionContext; +use crate::dataframe::PyDataFrame; +use crate::dataset::Dataset; +use crate::errors; +use crate::expr::create_external_table::PyCreateExternalTable; + +/// This struct is used as a common method for all TableProviders, +/// whether they refer to an FFI provider, an internally known +/// implementation, a dataset, or a dataframe view. +#[pyclass( + from_py_object, + frozen, + name = "RawTable", + module = "datafusion.catalog", + subclass +)] +#[derive(Clone)] +pub struct PyTable { + pub table: Arc, +} + +impl PyTable { + pub fn table(&self) -> Arc { + self.table.clone() + } +} + +#[pymethods] +impl PyTable { + /// Instantiate from any Python object that supports any of the table + /// types. We do not know a priori when using this method if the object + /// will be passed a wrapped or raw class. Here we handle all of the + /// following object types: + /// + /// - PyTable (essentially a clone operation), but either raw or wrapped + /// - DataFrame, either raw or wrapped + /// - FFI Table Providers via PyCapsule + /// - PyArrow Dataset objects + #[new] + pub fn new(obj: Bound<'_, PyAny>, session: Option>) -> PyResult { + let py = obj.py(); + if let Ok(py_table) = obj.extract::() { + Ok(py_table) + } else if let Ok(py_table) = obj + .getattr("_inner") + .and_then(|inner| inner.extract::().map_err(Into::::into)) + { + Ok(py_table) + } else if let Ok(py_df) = obj.extract::() { + let provider = py_df.inner_df().as_ref().clone().into_view(); + Ok(PyTable::from(provider)) + } else if let Ok(py_df) = obj + .getattr("df") + .and_then(|inner| inner.extract::().map_err(Into::::into)) + { + let provider = py_df.inner_df().as_ref().clone().into_view(); + Ok(PyTable::from(provider)) + } else if let Some(provider) = { + let session = match session { + Some(session) => session, + None => PySessionContext::global_ctx()?.into_bound_py_any(obj.py())?, + }; + table_provider_from_pycapsule(obj.clone(), session)? + } { + Ok(PyTable::from(provider)) + } else { + let provider = Arc::new(Dataset::new(&obj, py)?) as Arc; + Ok(PyTable::from(provider)) + } + } + + /// Get a reference to the schema for this table + #[getter] + fn schema<'py>(&self, py: Python<'py>) -> PyResult> { + self.table.schema().to_pyarrow(py) + } + + /// Get the type of this table for metadata/catalog purposes. + #[getter] + fn kind(&self) -> &str { + match self.table.table_type() { + TableType::Base => "physical", + TableType::View => "view", + TableType::Temporary => "temporary", + } + } + + fn __repr__(&self) -> PyResult { + let kind = self.kind(); + Ok(format!("Table(kind={kind})")) + } +} + +impl From> for PyTable { + fn from(table: Arc) -> Self { + Self { table } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct TempViewTable { + df: Arc, +} + +/// This is nearly identical to `DataFrameTableProvider` +/// except that it is for temporary tables. +/// Remove when https://github.com/apache/datafusion/issues/18026 +/// closes. +impl TempViewTable { + pub(crate) fn new(df: Arc) -> Self { + Self { df } + } +} + +#[async_trait] +impl TableProvider for TempViewTable { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + Arc::new(self.df.schema().as_arrow().clone()) + } + + fn table_type(&self) -> TableType { + TableType::Temporary + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> datafusion::common::Result> { + let filter = filters.iter().cloned().reduce(|acc, new| acc.and(new)); + let plan = self.df.logical_plan().clone(); + let mut plan = LogicalPlanBuilder::from(plan); + + if let Some(filter) = filter { + plan = plan.filter(filter)?; + } + + let mut plan = if let Some(projection) = projection { + // avoiding adding a redundant projection (e.g. SELECT * FROM view) + let current_projection = (0..plan.schema().fields().len()).collect::>(); + if projection == ¤t_projection { + plan + } else { + let fields: Vec = projection + .iter() + .map(|i| { + Expr::Column(Column::from( + self.df.logical_plan().schema().qualified_field(*i), + )) + }) + .collect(); + plan.project(fields)? + } + } else { + plan + }; + + if let Some(limit) = limit { + plan = plan.limit(0, Some(limit))?; + } + + state.create_physical_plan(&plan.build()?).await + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> datafusion::common::Result> { + Ok(vec![TableProviderFilterPushDown::Exact; filters.len()]) + } +} + +#[derive(Debug)] +pub(crate) struct RustWrappedPyTableProviderFactory { + pub(crate) table_provider_factory: Py, + pub(crate) codec: Arc, +} + +impl RustWrappedPyTableProviderFactory { + pub fn new(table_provider_factory: Py, codec: Arc) -> Self { + Self { + table_provider_factory, + codec, + } + } + + fn create_inner( + &self, + cmd: CreateExternalTable, + codec: Bound, + ) -> PyResult> { + Python::attach(|py| { + let provider = self.table_provider_factory.bind(py); + let cmd = PyCreateExternalTable::from(cmd); + + provider + .call_method1("create", (cmd,)) + .and_then(|t| PyTable::new(t, Some(codec))) + .map(|t| t.table()) + }) + } +} + +#[async_trait] +impl TableProviderFactory for RustWrappedPyTableProviderFactory { + async fn create( + &self, + _: &dyn Session, + cmd: &CreateExternalTable, + ) -> datafusion::common::Result> { + Python::attach(|py| { + let codec = create_logical_extension_capsule(py, self.codec.as_ref()) + .map_err(errors::to_datafusion_err)?; + + self.create_inner(cmd.clone(), codec.into_any()) + .map_err(errors::to_datafusion_err) + }) + } +} diff --git a/src/udaf.rs b/crates/core/src/udaf.rs similarity index 61% rename from src/udaf.rs rename to crates/core/src/udaf.rs index 34a9cd51d..ed26c79cc 100644 --- a/src/udaf.rs +++ b/crates/core/src/udaf.rs @@ -15,64 +15,67 @@ // specific language governing permissions and limitations // under the License. +use std::ptr::NonNull; use std::sync::Arc; -use pyo3::{prelude::*, types::PyTuple}; - -use datafusion::arrow::array::{Array, ArrayRef}; +use datafusion::arrow::array::ArrayRef; use datafusion::arrow::datatypes::DataType; use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow}; use datafusion::common::ScalarValue; use datafusion::error::{DataFusionError, Result}; use datafusion::logical_expr::{ - create_udaf, Accumulator, AccumulatorFactoryFunction, AggregateUDF, + Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, create_udaf, }; +use datafusion_ffi::udaf::FFI_AggregateUDF; +use datafusion_python_util::{parse_volatility, validate_pycapsule}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyTuple}; use crate::common::data_type::PyScalarValue; -use crate::errors::to_datafusion_err; +use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err}; use crate::expr::PyExpr; -use crate::utils::parse_volatility; #[derive(Debug)] struct RustAccumulator { - accum: PyObject, + accum: Py, } impl RustAccumulator { - fn new(accum: PyObject) -> Self { + fn new(accum: Py) -> Self { Self { accum } } } impl Accumulator for RustAccumulator { fn state(&mut self) -> Result> { - Python::with_gil(|py| { - self.accum - .bind(py) - .call_method0("state")? - .extract::>() + Python::attach(|py| -> PyResult> { + let values = self.accum.bind(py).call_method0("state")?; + let mut scalars = Vec::new(); + for item in values.try_iter()? { + let item: Bound<'_, PyAny> = item?; + let scalar = item.extract::()?.0; + scalars.push(scalar); + } + Ok(scalars) }) - .map(|v| v.into_iter().map(|x| x.0).collect()) .map_err(|e| DataFusionError::Execution(format!("{e}"))) } fn evaluate(&mut self) -> Result { - Python::with_gil(|py| { - self.accum - .bind(py) - .call_method0("evaluate")? - .extract::() + Python::attach(|py| -> PyResult { + let value = self.accum.bind(py).call_method0("evaluate")?; + value.extract::().map(|v| v.0) }) - .map(|v| v.0) .map_err(|e| DataFusionError::Execution(format!("{e}"))) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - Python::with_gil(|py| { + Python::attach(|py| { // 1. cast args to Pyarrow array let py_args = values .iter() - .map(|arg| arg.into_data().to_pyarrow(py).unwrap()) + .map(|arg| arg.to_data().to_pyarrow(py).unwrap()) .collect::>(); let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; @@ -87,13 +90,13 @@ impl Accumulator for RustAccumulator { } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - Python::with_gil(|py| { + Python::attach(|py| { // // 1. cast states to Pyarrow arrays - let py_states: Result> = states + let py_states: Result>> = states .iter() .map(|state| { state - .into_data() + .to_data() .to_pyarrow(py) .map_err(|e| DataFusionError::Execution(format!("{e}"))) }) @@ -114,11 +117,11 @@ impl Accumulator for RustAccumulator { } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - Python::with_gil(|py| { + Python::attach(|py| { // 1. cast args to Pyarrow array let py_args = values .iter() - .map(|arg| arg.into_data().to_pyarrow(py).unwrap()) + .map(|arg| arg.to_data().to_pyarrow(py).unwrap()) .collect::>(); let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; @@ -133,7 +136,7 @@ impl Accumulator for RustAccumulator { } fn supports_retract_batch(&self) -> bool { - Python::with_gil( + Python::attach( |py| match self.accum.bind(py).call_method0("supports_retract_batch") { Ok(x) => x.extract().unwrap_or(false), Err(_) => false, @@ -142,9 +145,9 @@ impl Accumulator for RustAccumulator { } } -pub fn to_rust_accumulator(accum: PyObject) -> AccumulatorFactoryFunction { - Arc::new(move |_| -> Result> { - let accum = Python::with_gil(|py| { +pub fn to_rust_accumulator(accum: Py) -> AccumulatorFactoryFunction { + Arc::new(move |_args| -> Result> { + let accum = Python::attach(|py| { accum .call0(py) .map_err(|e| DataFusionError::Execution(format!("{e}"))) @@ -153,8 +156,26 @@ pub fn to_rust_accumulator(accum: PyObject) -> AccumulatorFactoryFunction { }) } +fn aggregate_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult { + validate_pycapsule(capsule, "datafusion_aggregate_udf")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_aggregate_udf")))? + .cast(); + let udaf = unsafe { data.as_ref() }; + let udaf: Arc = udaf.into(); + + Ok(AggregateUDF::new_from_shared_impl(udaf)) +} + /// Represents an AggregateUDF -#[pyclass(name = "AggregateUDF", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "AggregateUDF", + module = "datafusion", + subclass +)] #[derive(Debug, Clone)] pub struct PyAggregateUDF { pub(crate) function: AggregateUDF, @@ -166,7 +187,7 @@ impl PyAggregateUDF { #[pyo3(signature=(name, accumulator, input_type, return_type, state_type, volatility))] fn new( name: &str, - accumulator: PyObject, + accumulator: Py, input_type: PyArrowType>, return_type: PyArrowType, state_type: PyArrowType>, @@ -183,6 +204,26 @@ impl PyAggregateUDF { Ok(Self { function }) } + #[staticmethod] + pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { + if func.is_instance_of::() { + let capsule = func.cast::().map_err(py_datafusion_err)?; + let function = aggregate_udf_from_capsule(capsule)?; + return Ok(Self { function }); + } + + if func.hasattr("__datafusion_aggregate_udf__")? { + let capsule = func.getattr("__datafusion_aggregate_udf__")?.call0()?; + let capsule = capsule.cast::().map_err(py_datafusion_err)?; + let function = aggregate_udf_from_capsule(capsule)?; + return Ok(Self { function }); + } + + Err(crate::errors::PyDataFusionError::Common( + "__datafusion_aggregate_udf__ does not exist on AggregateUDF object.".to_string(), + )) + } + /// creates a new PyExpr with the call of the udf #[pyo3(signature = (*args))] fn __call__(&self, args: Vec) -> PyResult { diff --git a/crates/core/src/udf.rs b/crates/core/src/udf.rs new file mode 100644 index 000000000..7543f96d4 --- /dev/null +++ b/crates/core/src/udf.rs @@ -0,0 +1,226 @@ +// 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 std::any::Any; +use std::hash::{Hash, Hasher}; +use std::ptr::NonNull; +use std::sync::Arc; + +use arrow::datatypes::{Field, FieldRef}; +use arrow::pyarrow::ToPyArrow; +use datafusion::arrow::array::{ArrayData, make_array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType}; +use datafusion::common::internal_err; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_ffi::udf::FFI_ScalarUDF; +use datafusion_python_util::{parse_volatility, validate_pycapsule}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyTuple}; + +use crate::array::PyArrowArrayExportable; +use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err}; +use crate::expr::PyExpr; + +/// This struct holds the Python written function that is a +/// ScalarUDF. +#[derive(Debug)] +struct PythonFunctionScalarUDF { + name: String, + func: Py, + signature: Signature, + return_field: FieldRef, +} + +impl PythonFunctionScalarUDF { + fn new( + name: String, + func: Py, + input_fields: Vec, + return_field: Field, + volatility: Volatility, + ) -> Self { + let input_types = input_fields.iter().map(|f| f.data_type().clone()).collect(); + let signature = Signature::exact(input_types, volatility); + Self { + name, + func, + signature, + return_field: Arc::new(return_field), + } + } +} + +impl Eq for PythonFunctionScalarUDF {} +impl PartialEq for PythonFunctionScalarUDF { + fn eq(&self, other: &Self) -> bool { + self.name == other.name + && self.signature == other.signature + && self.return_field == other.return_field + && Python::attach(|py| self.func.bind(py).eq(other.func.bind(py)).unwrap_or(false)) + } +} + +impl Hash for PythonFunctionScalarUDF { + fn hash(&self, state: &mut H) { + self.name.hash(state); + self.signature.hash(state); + self.return_field.hash(state); + + Python::attach(|py| { + let py_hash = self.func.bind(py).hash().unwrap_or(0); // Handle unhashable objects + + state.write_isize(py_hash); + }); + } +} + +impl ScalarUDFImpl for PythonFunctionScalarUDF { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> datafusion::common::Result { + internal_err!( + "return_field should not be called when return_field_from_args is implemented." + ) + } + + fn return_field_from_args( + &self, + _args: ReturnFieldArgs, + ) -> datafusion::common::Result { + Ok(Arc::clone(&self.return_field)) + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion::common::Result { + let num_rows = args.number_rows; + Python::attach(|py| { + // 1. cast args to Pyarrow arrays + let py_args = args + .args + .into_iter() + .zip(args.arg_fields) + .map(|(arg, field)| { + let array = arg.to_array(num_rows)?; + PyArrowArrayExportable::new(array, field) + .to_pyarrow(py) + .map_err(to_datafusion_err) + }) + .collect::, _>>()?; + let py_args = PyTuple::new(py, py_args).map_err(to_datafusion_err)?; + + // 2. call function + let value = self + .func + .call(py, py_args, None) + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + + // 3. cast to arrow::array::Array + let array_data = ArrayData::from_pyarrow_bound(value.bind(py)) + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + Ok(ColumnarValue::Array(make_array(array_data))) + }) + } +} + +/// Represents a PyScalarUDF +#[pyclass( + from_py_object, + frozen, + name = "ScalarUDF", + module = "datafusion", + subclass +)] +#[derive(Debug, Clone)] +pub struct PyScalarUDF { + pub(crate) function: ScalarUDF, +} + +#[pymethods] +impl PyScalarUDF { + #[new] + #[pyo3(signature=(name, func, input_types, return_type, volatility))] + fn new( + name: String, + func: Py, + input_types: PyArrowType>, + return_type: PyArrowType, + volatility: &str, + ) -> PyResult { + let py_function = PythonFunctionScalarUDF::new( + name, + func, + input_types.0, + return_type.0, + parse_volatility(volatility)?, + ); + let function = ScalarUDF::new_from_impl(py_function); + + Ok(Self { function }) + } + + #[staticmethod] + pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { + if func.hasattr("__datafusion_scalar_udf__")? { + let capsule = func.getattr("__datafusion_scalar_udf__")?.call0()?; + let capsule = capsule.cast::().map_err(py_datafusion_err)?; + validate_pycapsule(capsule, "datafusion_scalar_udf")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_scalar_udf")))? + .cast(); + let udf = unsafe { data.as_ref() }; + let udf: Arc = udf.into(); + + Ok(Self { + function: ScalarUDF::new_from_shared_impl(udf), + }) + } else { + Err(crate::errors::PyDataFusionError::Common( + "__datafusion_scalar_udf__ does not exist on ScalarUDF object.".to_string(), + )) + } + } + + /// creates a new PyExpr with the call of the udf + #[pyo3(signature = (*args))] + fn __call__(&self, args: Vec) -> PyResult { + let args = args.iter().map(|e| e.expr.clone()).collect(); + Ok(self.function.call(args).into()) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("ScalarUDF({})", self.function.name())) + } +} diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs new file mode 100644 index 000000000..77c5ffbbc --- /dev/null +++ b/crates/core/src/udtf.rs @@ -0,0 +1,138 @@ +// 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 std::ptr::NonNull; +use std::sync::Arc; + +use datafusion::catalog::{TableFunctionImpl, TableProvider}; +use datafusion::error::Result as DataFusionResult; +use datafusion::logical_expr::Expr; +use datafusion_ffi::udtf::FFI_TableFunction; +use datafusion_python_util::validate_pycapsule; +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyImportError, PyTypeError}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyTuple, PyType}; + +use crate::context::PySessionContext; +use crate::errors::{py_datafusion_err, to_datafusion_err}; +use crate::expr::PyExpr; +use crate::table::PyTable; + +/// Represents a user defined table function +#[pyclass(from_py_object, frozen, name = "TableFunction", module = "datafusion")] +#[derive(Debug, Clone)] +pub struct PyTableFunction { + pub(crate) name: String, + pub(crate) inner: PyTableFunctionInner, +} + +// TODO: Implement pure python based user defined table functions +#[derive(Debug, Clone)] +pub(crate) enum PyTableFunctionInner { + PythonFunction(Arc>), + FFIFunction(Arc), +} + +#[pymethods] +impl PyTableFunction { + #[new] + #[pyo3(signature=(name, func, session))] + pub fn new( + name: &str, + func: Bound<'_, PyAny>, + session: Option>, + ) -> PyResult { + let inner = if func.hasattr("__datafusion_table_function__")? { + let py = func.py(); + let session = match session { + Some(session) => session, + None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, + }; + let capsule = func + .getattr("__datafusion_table_function__")? + .call1((session,)).map_err(|err| { + if err.get_type(py).is(PyType::new::(py)) { + PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table functions. Either downgrade DataFusion or upgrade your function library.") + } else { + err + } + })?; + let capsule = capsule.cast::().map_err(py_datafusion_err)?; + validate_pycapsule(capsule, "datafusion_table_function")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_table_function")))? + .cast(); + let ffi_func = unsafe { data.as_ref() }; + let foreign_func: Arc = ffi_func.to_owned().into(); + + PyTableFunctionInner::FFIFunction(foreign_func) + } else { + let py_obj = Arc::new(func.unbind()); + PyTableFunctionInner::PythonFunction(py_obj) + }; + + Ok(Self { + name: name.to_string(), + inner, + }) + } + + #[pyo3(signature = (*args))] + pub fn __call__(&self, args: Vec) -> PyResult { + let args: Vec = args.iter().map(|e| e.expr.clone()).collect(); + let table_provider = self.call(&args).map_err(py_datafusion_err)?; + + Ok(PyTable::from(table_provider)) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("TableUDF({})", self.name)) + } +} + +#[allow(clippy::result_large_err)] +fn call_python_table_function( + func: &Arc>, + args: &[Expr], +) -> DataFusionResult> { + let args = args + .iter() + .map(|arg| PyExpr::from(arg.clone())) + .collect::>(); + + // move |args: &[ArrayRef]| -> Result { + Python::attach(|py| { + let py_args = PyTuple::new(py, args)?; + let provider_obj = func.call1(py, py_args)?; + let provider = provider_obj.bind(py).clone(); + + Ok::, PyErr>(PyTable::new(provider, None)?.table) + }) + .map_err(to_datafusion_err) +} + +impl TableFunctionImpl for PyTableFunction { + fn call(&self, args: &[Expr]) -> DataFusionResult> { + match &self.inner { + PyTableFunctionInner::FFIFunction(func) => func.call(args), + PyTableFunctionInner::PythonFunction(obj) => call_python_table_function(obj, args), + } + } +} diff --git a/src/udwf.rs b/crates/core/src/udwf.rs similarity index 81% rename from src/udwf.rs rename to crates/core/src/udwf.rs index defd9c522..ff7ab0352 100644 --- a/src/udwf.rs +++ b/crates/core/src/udwf.rs @@ -17,46 +17,50 @@ use std::any::Any; use std::ops::Range; +use std::ptr::NonNull; use std::sync::Arc; -use arrow::array::{make_array, Array, ArrayData, ArrayRef}; +use arrow::array::{Array, ArrayData, ArrayRef, make_array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}; +use datafusion::error::{DataFusionError, Result}; use datafusion::logical_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs}; +use datafusion::logical_expr::ptr_eq::PtrEq; use datafusion::logical_expr::window_state::WindowAggState; +use datafusion::logical_expr::{ + PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl, +}; use datafusion::scalar::ScalarValue; +use datafusion_ffi::udwf::FFI_WindowUDF; +use datafusion_python_util::{parse_volatility, validate_pycapsule}; use pyo3::exceptions::PyValueError; +use pyo3::ffi::c_str; use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyList, PyTuple}; use crate::common::data_type::PyScalarValue; -use crate::errors::to_datafusion_err; +use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err}; use crate::expr::PyExpr; -use crate::utils::parse_volatility; -use datafusion::arrow::datatypes::DataType; -use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}; -use datafusion::error::{DataFusionError, Result}; -use datafusion::logical_expr::{ - PartitionEvaluator, PartitionEvaluatorFactory, Signature, Volatility, WindowUDF, WindowUDFImpl, -}; -use pyo3::types::{PyList, PyTuple}; #[derive(Debug)] struct RustPartitionEvaluator { - evaluator: PyObject, + evaluator: Py, } impl RustPartitionEvaluator { - fn new(evaluator: PyObject) -> Self { + fn new(evaluator: Py) -> Self { Self { evaluator } } } impl PartitionEvaluator for RustPartitionEvaluator { fn memoize(&mut self, _state: &mut WindowAggState) -> Result<()> { - Python::with_gil(|py| self.evaluator.bind(py).call_method0("memoize").map(|_| ())) + Python::attach(|py| self.evaluator.bind(py).call_method0("memoize").map(|_| ())) .map_err(|e| DataFusionError::Execution(format!("{e}"))) } fn get_range(&self, idx: usize, n_rows: usize) -> Result> { - Python::with_gil(|py| { + Python::attach(|py| { let py_args = vec![idx.into_pyobject(py)?, n_rows.into_pyobject(py)?]; let py_args = PyTuple::new(py, py_args)?; @@ -82,7 +86,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn is_causal(&self) -> bool { - Python::with_gil(|py| { + Python::attach(|py| { self.evaluator .bind(py) .call_method0("is_causal") @@ -92,8 +96,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn evaluate_all(&mut self, values: &[ArrayRef], num_rows: usize) -> Result { - println!("evaluate all called with number of values {}", values.len()); - Python::with_gil(|py| { + Python::attach(|py| { let py_values = PyList::new( py, values @@ -115,7 +118,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn evaluate(&mut self, values: &[ArrayRef], range: &Range) -> Result { - Python::with_gil(|py| { + Python::attach(|py| { let py_values = PyList::new( py, values @@ -139,7 +142,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { num_rows: usize, ranks_in_partition: &[Range], ) -> Result { - Python::with_gil(|py| { + Python::attach(|py| { let ranks = ranks_in_partition .iter() .map(|r| PyTuple::new(py, vec![r.start, r.end])) @@ -166,7 +169,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn supports_bounded_execution(&self) -> bool { - Python::with_gil(|py| { + Python::attach(|py| { self.evaluator .bind(py) .call_method0("supports_bounded_execution") @@ -176,7 +179,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn uses_window_frame(&self) -> bool { - Python::with_gil(|py| { + Python::attach(|py| { self.evaluator .bind(py) .call_method0("uses_window_frame") @@ -186,7 +189,7 @@ impl PartitionEvaluator for RustPartitionEvaluator { } fn include_rank(&self) -> bool { - Python::with_gil(|py| { + Python::attach(|py| { self.evaluator .bind(py) .call_method0("include_rank") @@ -196,9 +199,9 @@ impl PartitionEvaluator for RustPartitionEvaluator { } } -pub fn to_rust_partition_evaluator(evaluator: PyObject) -> PartitionEvaluatorFactory { +pub fn to_rust_partition_evaluator(evaluator: Py) -> PartitionEvaluatorFactory { Arc::new(move || -> Result> { - let evaluator = Python::with_gil(|py| { + let evaluator = Python::attach(|py| { evaluator .call0(py) .map_err(|e| DataFusionError::Execution(e.to_string())) @@ -208,7 +211,13 @@ pub fn to_rust_partition_evaluator(evaluator: PyObject) -> PartitionEvaluatorFac } /// Represents an WindowUDF -#[pyclass(name = "WindowUDF", module = "datafusion", subclass)] +#[pyclass( + from_py_object, + frozen, + name = "WindowUDF", + module = "datafusion", + subclass +)] #[derive(Debug, Clone)] pub struct PyWindowUDF { pub(crate) function: WindowUDF, @@ -220,7 +229,7 @@ impl PyWindowUDF { #[pyo3(signature=(name, evaluator, input_types, return_type, volatility))] fn new( name: &str, - evaluator: PyObject, + evaluator: Py, input_types: Vec>, return_type: PyArrowType, volatility: &str, @@ -245,16 +254,39 @@ impl PyWindowUDF { Ok(self.function.call(args).into()) } + #[staticmethod] + pub fn from_pycapsule(func: Bound<'_, PyAny>) -> PyDataFusionResult { + let capsule = if func.hasattr("__datafusion_window_udf__")? { + func.getattr("__datafusion_window_udf__")?.call0()? + } else { + func + }; + + let capsule = capsule.cast::().map_err(py_datafusion_err)?; + validate_pycapsule(capsule, "datafusion_window_udf")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_window_udf")))? + .cast(); + let udwf = unsafe { data.as_ref() }; + let udwf: Arc = udwf.into(); + + Ok(Self { + function: WindowUDF::new_from_shared_impl(udwf), + }) + } + fn __repr__(&self) -> PyResult { Ok(format!("WindowUDF({})", self.function.name())) } } +#[derive(Hash, Eq, PartialEq)] pub struct MultiColumnWindowUDF { name: String, signature: Signature, return_type: DataType, - partition_evaluator_factory: PartitionEvaluatorFactory, + partition_evaluator_factory: PtrEq, } impl std::fmt::Debug for MultiColumnWindowUDF { @@ -282,7 +314,7 @@ impl MultiColumnWindowUDF { name, signature, return_type, - partition_evaluator_factory, + partition_evaluator_factory: partition_evaluator_factory.into(), } } } @@ -300,13 +332,9 @@ impl WindowUDFImpl for MultiColumnWindowUDF { &self.signature } - fn field(&self, field_args: WindowUDFFieldArgs) -> Result { + fn field(&self, field_args: WindowUDFFieldArgs) -> Result { // TODO: Should nullable always be `true`? - Ok(arrow::datatypes::Field::new( - field_args.name(), - self.return_type.clone(), - true, - )) + Ok(arrow::datatypes::Field::new(field_args.name(), self.return_type.clone(), true).into()) } // TODO: Enable passing partition_evaluator_args to python? diff --git a/crates/core/src/unparser/dialect.rs b/crates/core/src/unparser/dialect.rs new file mode 100644 index 000000000..52a2da00b --- /dev/null +++ b/crates/core/src/unparser/dialect.rs @@ -0,0 +1,69 @@ +// 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 std::sync::Arc; + +use datafusion::sql::unparser::dialect::{ + DefaultDialect, Dialect, DuckDBDialect, MySqlDialect, PostgreSqlDialect, SqliteDialect, +}; +use pyo3::prelude::*; + +#[pyclass( + from_py_object, + frozen, + name = "Dialect", + module = "datafusion.unparser", + subclass +)] +#[derive(Clone)] +pub struct PyDialect { + pub dialect: Arc, +} + +#[pymethods] +impl PyDialect { + #[staticmethod] + pub fn default() -> Self { + Self { + dialect: Arc::new(DefaultDialect {}), + } + } + #[staticmethod] + pub fn postgres() -> Self { + Self { + dialect: Arc::new(PostgreSqlDialect {}), + } + } + #[staticmethod] + pub fn mysql() -> Self { + Self { + dialect: Arc::new(MySqlDialect {}), + } + } + #[staticmethod] + pub fn sqlite() -> Self { + Self { + dialect: Arc::new(SqliteDialect {}), + } + } + #[staticmethod] + pub fn duckdb() -> Self { + Self { + dialect: Arc::new(DuckDBDialect::new()), + } + } +} diff --git a/crates/core/src/unparser/mod.rs b/crates/core/src/unparser/mod.rs new file mode 100644 index 000000000..5142b918e --- /dev/null +++ b/crates/core/src/unparser/mod.rs @@ -0,0 +1,74 @@ +// 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. + +mod dialect; + +use std::sync::Arc; + +use datafusion::sql::unparser::Unparser; +use datafusion::sql::unparser::dialect::Dialect; +use dialect::PyDialect; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +use crate::sql::logical::PyLogicalPlan; + +#[pyclass( + from_py_object, + frozen, + name = "Unparser", + module = "datafusion.unparser", + subclass +)] +#[derive(Clone)] +pub struct PyUnparser { + dialect: Arc, + pretty: bool, +} + +#[pymethods] +impl PyUnparser { + #[new] + pub fn new(dialect: PyDialect) -> Self { + Self { + dialect: dialect.dialect.clone(), + pretty: false, + } + } + + pub fn plan_to_sql(&self, plan: &PyLogicalPlan) -> PyResult { + let mut unparser = Unparser::new(self.dialect.as_ref()); + unparser = unparser.with_pretty(self.pretty); + let sql = unparser + .plan_to_sql(&plan.plan()) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(sql.to_string()) + } + + pub fn with_pretty(&self, pretty: bool) -> Self { + Self { + dialect: self.dialect.clone(), + pretty, + } + } +} + +pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/ffi-table-provider/Cargo.toml b/crates/util/Cargo.toml similarity index 64% rename from examples/ffi-table-provider/Cargo.toml rename to crates/util/Cargo.toml index f4e4fda79..00d5946a5 100644 --- a/examples/ffi-table-provider/Cargo.toml +++ b/crates/util/Cargo.toml @@ -16,21 +16,19 @@ # under the License. [package] -name = "ffi-table-provider" -version = "0.1.0" -edition = "2021" +name = "datafusion-python-util" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true [dependencies] -datafusion = { version = "45.0.0" } -datafusion-ffi = { version = "45.0.0" } -pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py39"] } -arrow = { version = "54" } -arrow-array = { version = "54" } -arrow-schema = { version = "54" } - -[build-dependencies] -pyo3-build-config = "0.23" - -[lib] -name = "ffi_table_provider" -crate-type = ["cdylib", "rlib"] +tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] } +pyo3 = { workspace = true } +datafusion = { workspace = true } +datafusion-ffi = { workspace = true } +arrow = { workspace = true } +prost = { workspace = true } diff --git a/src/errors.rs b/crates/util/src/errors.rs similarity index 84% rename from src/errors.rs rename to crates/util/src/errors.rs index f1d5aeb23..0d25c8847 100644 --- a/src/errors.rs +++ b/crates/util/src/errors.rs @@ -22,13 +22,14 @@ use std::fmt::Debug; use datafusion::arrow::error::ArrowError; use datafusion::error::DataFusionError as InnerDataFusionError; use prost::EncodeError; -use pyo3::{exceptions::PyException, PyErr}; +use pyo3::PyErr; +use pyo3::exceptions::{PyException, PyValueError}; pub type PyDataFusionResult = std::result::Result; #[derive(Debug)] pub enum PyDataFusionError { - ExecutionError(InnerDataFusionError), + ExecutionError(Box), ArrowError(ArrowError), Common(String), PythonError(PyErr), @@ -38,7 +39,7 @@ pub enum PyDataFusionError { impl fmt::Display for PyDataFusionError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - PyDataFusionError::ExecutionError(e) => write!(f, "DataFusion error: {e:?}"), + PyDataFusionError::ExecutionError(e) => write!(f, "DataFusion error: {e}"), PyDataFusionError::ArrowError(e) => write!(f, "Arrow error: {e:?}"), PyDataFusionError::PythonError(e) => write!(f, "Python error {e:?}"), PyDataFusionError::Common(e) => write!(f, "{e}"), @@ -55,7 +56,7 @@ impl From for PyDataFusionError { impl From for PyDataFusionError { fn from(err: InnerDataFusionError) -> PyDataFusionError { - PyDataFusionError::ExecutionError(err) + PyDataFusionError::ExecutionError(Box::new(err)) } } @@ -95,3 +96,13 @@ pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr { pub fn to_datafusion_err(e: impl Debug) -> InnerDataFusionError { InnerDataFusionError::Execution(format!("{e:?}")) } + +pub fn from_datafusion_error(err: InnerDataFusionError) -> PyErr { + match err { + InnerDataFusionError::External(boxed) => match boxed.downcast::() { + Ok(py_err) => *py_err, + Err(original_boxed) => PyValueError::new_err(format!("{original_boxed}")), + }, + _ => PyValueError::new_err(format!("{err}")), + } +} diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs new file mode 100644 index 000000000..2678a6b9a --- /dev/null +++ b/crates/util/src/lib.rs @@ -0,0 +1,231 @@ +// 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 std::future::Future; +use std::ptr::NonNull; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use datafusion::datasource::TableProvider; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::Volatility; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::table_provider::FFI_TableProvider; +use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; +use pyo3::ffi::c_str; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyType}; +use tokio::runtime::Runtime; +use tokio::task::JoinHandle; +use tokio::time::sleep; + +use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err}; + +pub mod errors; + +/// Utility to get the Tokio Runtime from Python +#[inline] +pub fn get_tokio_runtime() -> &'static Runtime { + // NOTE: Other pyo3 python libraries have had issues with using tokio + // behind a forking app-server like `gunicorn` + // If we run into that problem, in the future we can look to `delta-rs` + // which adds a check in that disallows calls from a forked process + // https://github.com/delta-io/delta-rs/blob/87010461cfe01563d91a4b9cd6fa468e2ad5f283/python/src/utils.rs#L10-L31 + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| Runtime::new().unwrap()) +} + +#[inline] +pub fn is_ipython_env(py: Python) -> &'static bool { + static IS_IPYTHON_ENV: OnceLock = OnceLock::new(); + IS_IPYTHON_ENV.get_or_init(|| { + py.import("IPython") + .and_then(|ipython| ipython.call_method0("get_ipython")) + .map(|ipython| !ipython.is_none()) + .unwrap_or(false) + }) +} + +/// Utility to get the Global Datafussion CTX +#[inline] +pub fn get_global_ctx() -> &'static Arc { + static CTX: OnceLock> = OnceLock::new(); + CTX.get_or_init(|| Arc::new(SessionContext::new())) +} + +/// Utility to collect rust futures with GIL released and respond to +/// Python interrupts such as ``KeyboardInterrupt``. If a signal is +/// received while the future is running, the future is aborted and the +/// corresponding Python exception is raised. +pub fn wait_for_future(py: Python, fut: F) -> PyResult +where + F: Future + Send, + F::Output: Send, +{ + let runtime: &Runtime = get_tokio_runtime(); + const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(1_000); + + // Some fast running processes that generate many `wait_for_future` calls like + // PartitionedDataFrameStreamReader::next require checking for interrupts early + py.run(cr"pass", None, None)?; + py.check_signals()?; + + py.detach(|| { + runtime.block_on(async { + tokio::pin!(fut); + loop { + tokio::select! { + res = &mut fut => break Ok(res), + _ = sleep(INTERVAL_CHECK_SIGNALS) => { + Python::attach(|py| { + // Execute a no-op Python statement to trigger signal processing. + // This is necessary because py.check_signals() alone doesn't + // actually check for signals - it only raises an exception if + // a signal was already set during a previous Python API call. + // Running even trivial Python code forces the interpreter to + // process any pending signals (like KeyboardInterrupt). + py.run(cr"pass", None, None)?; + py.check_signals() + })?; + } + } + } + }) + }) +} + +/// Spawn a [`Future`] on the Tokio runtime and wait for completion +/// while respecting Python signal handling. +pub fn spawn_future(py: Python, fut: F) -> PyDataFusionResult +where + F: Future> + Send + 'static, + T: Send + 'static, +{ + let rt = get_tokio_runtime(); + let handle: JoinHandle> = rt.spawn(fut); + // Wait for the join handle while respecting Python signal handling. + // We handle errors in two steps so `?` maps the error types correctly: + // 1) convert any Python-related error from `wait_for_future` into `PyDataFusionError` + // 2) convert any DataFusion error (inner result) into `PyDataFusionError` + let inner_result = wait_for_future(py, async { + // handle.await yields `Result, JoinError>` + // map JoinError into a DataFusion error so the async block returns + // `datafusion::common::Result` (i.e. Result) + match handle.await { + Ok(inner) => inner, + Err(join_err) => Err(to_datafusion_err(join_err)), + } + })?; // converts PyErr -> PyDataFusionError + + // `inner_result` is `datafusion::common::Result`; use `?` to convert + // the inner DataFusion error into `PyDataFusionError` via `From` and + // return the inner `T` on success. + Ok(inner_result?) +} + +pub fn parse_volatility(value: &str) -> PyDataFusionResult { + Ok(match value { + "immutable" => Volatility::Immutable, + "stable" => Volatility::Stable, + "volatile" => Volatility::Volatile, + value => { + return Err(PyDataFusionError::Common(format!( + "Unsupported volatility type: `{value}`, supported \ + values are: immutable, stable and volatile." + ))); + } + }) +} + +pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<()> { + let capsule_name = capsule.name()?; + if capsule_name.is_none() { + return Err(PyValueError::new_err(format!( + "Expected {name} PyCapsule to have name set." + ))); + } + + let capsule_name = unsafe { capsule_name.unwrap().as_cstr().to_str()? }; + if capsule_name != name { + return Err(PyValueError::new_err(format!( + "Expected name '{name}' in PyCapsule, instead got '{capsule_name}'" + ))); + } + + Ok(()) +} + +pub fn table_provider_from_pycapsule<'py>( + mut obj: Bound<'py, PyAny>, + session: Bound<'py, PyAny>, +) -> PyResult>> { + if obj.hasattr("__datafusion_table_provider__")? { + obj = obj + .getattr("__datafusion_table_provider__")? + .call1((session,)).map_err(|err| { + let py = obj.py(); + if err.get_type(py).is(PyType::new::(py)) { + PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table providers. Either downgrade DataFusion or upgrade your function library.") + } else { + err + } + })?; + } + + if let Ok(capsule) = obj.cast::().map_err(py_datafusion_err) { + validate_pycapsule(capsule, "datafusion_table_provider")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_table_provider")))? + .cast(); + let provider = unsafe { data.as_ref() }; + let provider: Arc = provider.into(); + + Ok(Some(provider)) + } else { + Ok(None) + } +} + +pub fn create_logical_extension_capsule<'py>( + py: Python<'py>, + codec: &FFI_LogicalExtensionCodec, +) -> PyResult> { + let name = cr"datafusion_logical_extension_codec".into(); + let codec = codec.clone(); + + PyCapsule::new(py, codec, Some(name)) +} + +pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult { + let attr_name = "__datafusion_logical_extension_codec__"; + let capsule = if obj.hasattr(attr_name)? { + obj.getattr(attr_name)?.call0()? + } else { + obj + }; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_logical_extension_codec")?; + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_logical_extension_codec")))? + .cast(); + let codec = unsafe { data.as_ref() }; + + Ok(codec.clone()) +} diff --git a/dev/changelog/46.0.0.md b/dev/changelog/46.0.0.md new file mode 100644 index 000000000..3e5768099 --- /dev/null +++ b/dev/changelog/46.0.0.md @@ -0,0 +1,73 @@ + + +# Apache DataFusion Python 46.0.0 Changelog + +This release consists of 21 commits from 11 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: reads using global ctx [#982](https://github.com/apache/datafusion-python/pull/982) (ion-elgreco) +- feat: Implementation of udf and udaf decorator [#1040](https://github.com/apache/datafusion-python/pull/1040) (CrystalZhou0529) +- feat: expose regex_count function [#1066](https://github.com/apache/datafusion-python/pull/1066) (nirnayroy) +- feat: Update DataFusion dependency to 46 [#1079](https://github.com/apache/datafusion-python/pull/1079) (timsaucer) + +**Fixed bugs:** + +- fix: add to_timestamp_nanos [#1020](https://github.com/apache/datafusion-python/pull/1020) (chenkovsky) +- fix: type checking [#993](https://github.com/apache/datafusion-python/pull/993) (chenkovsky) + +**Other:** + +- [infra] Fail Clippy on rust build warnings [#1029](https://github.com/apache/datafusion-python/pull/1029) (kevinjqliu) +- Add user documentation for the FFI approach [#1031](https://github.com/apache/datafusion-python/pull/1031) (timsaucer) +- build(deps): bump arrow from 54.1.0 to 54.2.0 [#1035](https://github.com/apache/datafusion-python/pull/1035) (dependabot[bot]) +- Chore: Release datafusion-python 45 [#1024](https://github.com/apache/datafusion-python/pull/1024) (timsaucer) +- Enable Dataframe to be converted into views which can be used in register_table [#1016](https://github.com/apache/datafusion-python/pull/1016) (kosiew) +- Add ruff check for missing futures import [#1052](https://github.com/apache/datafusion-python/pull/1052) (timsaucer) +- Enable take comments to assign issues to users [#1058](https://github.com/apache/datafusion-python/pull/1058) (timsaucer) +- Update python min version to 3.9 [#1043](https://github.com/apache/datafusion-python/pull/1043) (kevinjqliu) +- feat/improve ruff test coverage [#1055](https://github.com/apache/datafusion-python/pull/1055) (timsaucer) +- feat/making global context accessible for users [#1060](https://github.com/apache/datafusion-python/pull/1060) (jsai28) +- Renaming Internal Structs [#1059](https://github.com/apache/datafusion-python/pull/1059) (Spaarsh) +- test: add pytest asyncio tests [#1063](https://github.com/apache/datafusion-python/pull/1063) (jsai28) +- Add decorator for udwf [#1061](https://github.com/apache/datafusion-python/pull/1061) (kosiew) +- Add additional ruff suggestions [#1062](https://github.com/apache/datafusion-python/pull/1062) (Spaarsh) +- Improve collection during repr and repr_html [#1036](https://github.com/apache/datafusion-python/pull/1036) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 7 Tim Saucer + 2 Kevin Liu + 2 Spaarsh + 2 jsai28 + 2 kosiew + 1 Chen Chongchen + 1 Chongchen Chen + 1 Crystal Zhou + 1 Ion Koutsouris + 1 Nirnay Roy + 1 dependabot[bot] +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/47.0.0.md b/dev/changelog/47.0.0.md new file mode 100644 index 000000000..a7ed90313 --- /dev/null +++ b/dev/changelog/47.0.0.md @@ -0,0 +1,64 @@ + + +# Apache DataFusion Python 47.0.0 Changelog + +This release consists of 23 commits from 5 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: support unparser [#1088](https://github.com/apache/datafusion-python/pull/1088) (chenkovsky) +- feat: update datafusion dependency 47 [#1107](https://github.com/apache/datafusion-python/pull/1107) (timsaucer) +- feat: alias with metadata [#1111](https://github.com/apache/datafusion-python/pull/1111) (chenkovsky) +- feat: add missing PyLogicalPlan to_variant [#1085](https://github.com/apache/datafusion-python/pull/1085) (chenkovsky) +- feat: add user defined table function support [#1113](https://github.com/apache/datafusion-python/pull/1113) (timsaucer) + +**Fixed bugs:** + +- fix: recursive import [#1117](https://github.com/apache/datafusion-python/pull/1117) (chenkovsky) + +**Other:** + +- Update changelog and version number [#1089](https://github.com/apache/datafusion-python/pull/1089) (timsaucer) +- Documentation updates: mention correct dataset on basics page [#1081](https://github.com/apache/datafusion-python/pull/1081) (floscha) +- Add Configurable HTML Table Formatter for DataFusion DataFrames in Python [#1100](https://github.com/apache/datafusion-python/pull/1100) (kosiew) +- Add DataFrame usage guide with HTML rendering customization options [#1108](https://github.com/apache/datafusion-python/pull/1108) (kosiew) +- 1075/enhancement/Make col class with __getattr__ [#1076](https://github.com/apache/datafusion-python/pull/1076) (deanm0000) +- 1064/enhancement/add functions to Expr class [#1074](https://github.com/apache/datafusion-python/pull/1074) (deanm0000) +- ci: require approving review [#1122](https://github.com/apache/datafusion-python/pull/1122) (timsaucer) +- Partial fix for 1078: Enhance DataFrame Formatter Configuration with Memory and Display Controls [#1119](https://github.com/apache/datafusion-python/pull/1119) (kosiew) +- Add fill_null method to DataFrame API for handling missing values [#1019](https://github.com/apache/datafusion-python/pull/1019) (kosiew) +- minor: reduce error size [#1126](https://github.com/apache/datafusion-python/pull/1126) (timsaucer) +- Move the udf module to user_defined [#1112](https://github.com/apache/datafusion-python/pull/1112) (timsaucer) +- add unit tests for expression functions [#1121](https://github.com/apache/datafusion-python/pull/1121) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 12 Tim Saucer + 4 Chen Chongchen + 4 kosiew + 2 deanm0000 + 1 Florian Schäfer +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/48.0.0.md b/dev/changelog/48.0.0.md new file mode 100644 index 000000000..80bc61aca --- /dev/null +++ b/dev/changelog/48.0.0.md @@ -0,0 +1,59 @@ + + +# Apache DataFusion Python 48.0.0 Changelog + +This release consists of 15 commits from 6 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: upgrade df48 dependency [#1143](https://github.com/apache/datafusion-python/pull/1143) (timsaucer) +- feat: Support Parquet writer options [#1123](https://github.com/apache/datafusion-python/pull/1123) (nuno-faria) +- feat: dataframe string formatter [#1170](https://github.com/apache/datafusion-python/pull/1170) (timsaucer) +- feat: collect once during display() in jupyter notebooks [#1167](https://github.com/apache/datafusion-python/pull/1167) (timsaucer) +- feat: python based catalog and schema provider [#1156](https://github.com/apache/datafusion-python/pull/1156) (timsaucer) +- feat: add FFI support for user defined functions [#1145](https://github.com/apache/datafusion-python/pull/1145) (timsaucer) + +**Other:** + +- Release DataFusion 47.0.0 [#1130](https://github.com/apache/datafusion-python/pull/1130) (timsaucer) +- Add a documentation build step in CI [#1139](https://github.com/apache/datafusion-python/pull/1139) (crystalxyz) +- Add DataFrame API Documentation for DataFusion Python [#1132](https://github.com/apache/datafusion-python/pull/1132) (kosiew) +- Add Interruptible Query Execution in Jupyter via KeyboardInterrupt Support [#1141](https://github.com/apache/datafusion-python/pull/1141) (kosiew) +- Support types other than String and Int for partition columns [#1154](https://github.com/apache/datafusion-python/pull/1154) (miclegr) +- Fix signature of `__arrow_c_stream__` [#1168](https://github.com/apache/datafusion-python/pull/1168) (kylebarron) +- Consolidate DataFrame Docs: Merge HTML Rendering Section as Subpage [#1161](https://github.com/apache/datafusion-python/pull/1161) (kosiew) +- Add compression_level support to ParquetWriterOptions and enhance write_parquet to accept full options object [#1169](https://github.com/apache/datafusion-python/pull/1169) (kosiew) +- Simplify HTML Formatter Style Handling Using Script Injection [#1177](https://github.com/apache/datafusion-python/pull/1177) (kosiew) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 6 Tim Saucer + 5 kosiew + 1 Crystal Zhou + 1 Kyle Barron + 1 Michele Gregori + 1 Nuno Faria +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/49.0.0.md b/dev/changelog/49.0.0.md new file mode 100644 index 000000000..008bd43bc --- /dev/null +++ b/dev/changelog/49.0.0.md @@ -0,0 +1,61 @@ + + +# Apache DataFusion Python 49.0.0 Changelog + +This release consists of 16 commits from 7 contributors. See credits at the end of this changelog for more information. + +**Fixed bugs:** + +- fix(build): Include build.rs in published crates [#1199](https://github.com/apache/datafusion-python/pull/1199) (colinmarc) + +**Other:** + +- 48.0.0 Release [#1175](https://github.com/apache/datafusion-python/pull/1175) (timsaucer) +- Update CI rules [#1188](https://github.com/apache/datafusion-python/pull/1188) (timsaucer) +- Fix Python UDAF Accumulator Interface example to Properly Handle State and Updates with List[Array] Types [#1192](https://github.com/apache/datafusion-python/pull/1192) (kosiew) +- chore: Upgrade datafusion to version 49 [#1200](https://github.com/apache/datafusion-python/pull/1200) (nuno-faria) +- Update how to dev instructions [#1179](https://github.com/apache/datafusion-python/pull/1179) (ntjohnson1) +- build(deps): bump object_store from 0.12.2 to 0.12.3 [#1189](https://github.com/apache/datafusion-python/pull/1189) (dependabot[bot]) +- build(deps): bump uuid from 1.17.0 to 1.18.0 [#1202](https://github.com/apache/datafusion-python/pull/1202) (dependabot[bot]) +- build(deps): bump async-trait from 0.1.88 to 0.1.89 [#1203](https://github.com/apache/datafusion-python/pull/1203) (dependabot[bot]) +- build(deps): bump slab from 0.4.10 to 0.4.11 [#1205](https://github.com/apache/datafusion-python/pull/1205) (dependabot[bot]) +- Improved window and aggregate function signature [#1187](https://github.com/apache/datafusion-python/pull/1187) (timsaucer) +- Optional improvements in verification instructions [#1183](https://github.com/apache/datafusion-python/pull/1183) (paleolimbot) +- Improve `show()` output for empty DataFrames [#1208](https://github.com/apache/datafusion-python/pull/1208) (kosiew) +- build(deps): bump actions/download-artifact from 4 to 5 [#1201](https://github.com/apache/datafusion-python/pull/1201) (dependabot[bot]) +- build(deps): bump url from 2.5.4 to 2.5.7 [#1210](https://github.com/apache/datafusion-python/pull/1210) (dependabot[bot]) +- build(deps): bump actions/checkout from 4 to 5 [#1204](https://github.com/apache/datafusion-python/pull/1204) (dependabot[bot]) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 7 dependabot[bot] + 3 Tim Saucer + 2 kosiew + 1 Colin Marc + 1 Dewey Dunnington + 1 Nick + 1 Nuno Faria +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/50.0.0.md b/dev/changelog/50.0.0.md new file mode 100644 index 000000000..c3f09d180 --- /dev/null +++ b/dev/changelog/50.0.0.md @@ -0,0 +1,60 @@ + + +# Apache DataFusion Python 50.0.0 Changelog + +This release consists of 12 commits from 7 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: allow passing a slice to and expression with the [] indexing [#1215](https://github.com/apache/datafusion-python/pull/1215) (timsaucer) + +**Documentation updates:** + +- docs: fix CaseBuilder documentation example [#1225](https://github.com/apache/datafusion-python/pull/1225) (IndexSeek) +- docs: update link to user example for custom table provider [#1224](https://github.com/apache/datafusion-python/pull/1224) (IndexSeek) +- docs: add apache iceberg as datafusion data source [#1240](https://github.com/apache/datafusion-python/pull/1240) (kevinjqliu) + +**Other:** + +- 49.0.0 release [#1211](https://github.com/apache/datafusion-python/pull/1211) (timsaucer) +- Update development guide in README.md [#1213](https://github.com/apache/datafusion-python/pull/1213) (YKoustubhRao) +- Add benchmark script and documentation for maximizing CPU usage in DataFusion Python [#1216](https://github.com/apache/datafusion-python/pull/1216) (kosiew) +- Fixing a few Typos [#1220](https://github.com/apache/datafusion-python/pull/1220) (ntjohnson1) +- Set fail on warning for documentation generation [#1218](https://github.com/apache/datafusion-python/pull/1218) (timsaucer) +- chore: remove redundant error transformation [#1232](https://github.com/apache/datafusion-python/pull/1232) (mesejo) +- Support string column identifiers for sort/aggregate/window and stricter Expr validation [#1221](https://github.com/apache/datafusion-python/pull/1221) (kosiew) +- Prepare for DF50 [#1231](https://github.com/apache/datafusion-python/pull/1231) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 4 Tim Saucer + 2 Tyler White + 2 kosiew + 1 Daniel Mesejo + 1 Kevin Liu + 1 Koustubh Rao + 1 Nick +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/50.1.0.md b/dev/changelog/50.1.0.md new file mode 100644 index 000000000..3b9ff84ff --- /dev/null +++ b/dev/changelog/50.1.0.md @@ -0,0 +1,57 @@ + + +# Apache DataFusion Python 50.1.0 Changelog + +This release consists of 11 commits from 7 contributors. See credits at the end of this changelog for more information. + +**Breaking changes:** + +- Unify Table representations [#1256](https://github.com/apache/datafusion-python/pull/1256) (timsaucer) + +**Implemented enhancements:** + +- feat: expose DataFrame.write_table [#1264](https://github.com/apache/datafusion-python/pull/1264) (timsaucer) +- feat: expose` DataFrame.parse_sql_expr` [#1274](https://github.com/apache/datafusion-python/pull/1274) (milenkovicm) + +**Other:** + +- Update version number, add changelog [#1249](https://github.com/apache/datafusion-python/pull/1249) (timsaucer) +- Fix drop() method to handle quoted column names consistently [#1242](https://github.com/apache/datafusion-python/pull/1242) (H0TB0X420) +- Make Session Context `pyclass` frozen so interior mutability is only managed by rust [#1248](https://github.com/apache/datafusion-python/pull/1248) (ntjohnson1) +- macos-13 is deprecated [#1259](https://github.com/apache/datafusion-python/pull/1259) (kevinjqliu) +- Freeze PyO3 wrappers & introduce interior mutability to avoid PyO3 borrow errors [#1253](https://github.com/apache/datafusion-python/pull/1253) (kosiew) +- chore: update dependencies [#1269](https://github.com/apache/datafusion-python/pull/1269) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 4 Tim Saucer + 2 Siew Kam Onn + 1 H0TB0X420 + 1 Kevin Liu + 1 Marko Milenković + 1 Nick + 1 kosiew +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/51.0.0.md b/dev/changelog/51.0.0.md new file mode 100644 index 000000000..cc157eb0d --- /dev/null +++ b/dev/changelog/51.0.0.md @@ -0,0 +1,74 @@ + + +# Apache DataFusion Python 51.0.0 Changelog + +This release consists of 23 commits from 7 contributors. See credits at the end of this changelog for more information. + +**Breaking changes:** + +- feat: reduce duplicate fields on join [#1184](https://github.com/apache/datafusion-python/pull/1184) (timsaucer) + +**Implemented enhancements:** + +- feat: expose `select_exprs` method on DataFrame [#1271](https://github.com/apache/datafusion-python/pull/1271) (milenkovicm) +- feat: allow DataFrame.filter to accept SQL strings [#1276](https://github.com/apache/datafusion-python/pull/1276) (K-dash) +- feat: add temporary view option for into_view [#1267](https://github.com/apache/datafusion-python/pull/1267) (timsaucer) +- feat: support session token parameter for AmazonS3 [#1275](https://github.com/apache/datafusion-python/pull/1275) (GCHQDeveloper028) +- feat: `with_column` supports SQL expression [#1284](https://github.com/apache/datafusion-python/pull/1284) (milenkovicm) +- feat: Add SQL expression for `repartition_by_hash` [#1285](https://github.com/apache/datafusion-python/pull/1285) (milenkovicm) +- feat: Add SQL expression support for `with_columns` [#1286](https://github.com/apache/datafusion-python/pull/1286) (milenkovicm) + +**Fixed bugs:** + +- fix: use coalesce instead of drop_duplicate_keys for join [#1318](https://github.com/apache/datafusion-python/pull/1318) (mesejo) +- fix: Inconsistent schemas when converting to pyarrow [#1315](https://github.com/apache/datafusion-python/pull/1315) (nuno-faria) + +**Other:** + +- Release 50.1 [#1281](https://github.com/apache/datafusion-python/pull/1281) (timsaucer) +- Update python minimum version to 3.10 [#1296](https://github.com/apache/datafusion-python/pull/1296) (timsaucer) +- chore: update datafusion minor version [#1297](https://github.com/apache/datafusion-python/pull/1297) (timsaucer) +- Enable remaining pylints [#1298](https://github.com/apache/datafusion-python/pull/1298) (timsaucer) +- Add Arrow C streaming, DataFrame iteration, and OOM-safe streaming execution [#1222](https://github.com/apache/datafusion-python/pull/1222) (kosiew) +- Add PyCapsule Type Support and Type Hint Enhancements for AggregateUDF in DataFusion Python Bindings [#1277](https://github.com/apache/datafusion-python/pull/1277) (kosiew) +- Add collect_column to dataframe [#1302](https://github.com/apache/datafusion-python/pull/1302) (timsaucer) +- chore: apply cargo fmt with import organization [#1303](https://github.com/apache/datafusion-python/pull/1303) (timsaucer) +- Feat/parameterized sql queries [#964](https://github.com/apache/datafusion-python/pull/964) (timsaucer) +- Upgrade to Datafusion 51 [#1311](https://github.com/apache/datafusion-python/pull/1311) (nuno-faria) +- minor: resolve build errors after latest merge into main [#1325](https://github.com/apache/datafusion-python/pull/1325) (timsaucer) +- Update build workflow link [#1330](https://github.com/apache/datafusion-python/pull/1330) (timsaucer) +- Do not convert pyarrow scalar values to plain python types when passing as `lit` [#1319](https://github.com/apache/datafusion-python/pull/1319) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 12 Tim Saucer + 4 Marko Milenković + 2 Nuno Faria + 2 kosiew + 1 Daniel Mesejo + 1 GCHQDeveloper028 + 1 𝕂 +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/changelog/52.0.0.md b/dev/changelog/52.0.0.md new file mode 100644 index 000000000..3f848bb47 --- /dev/null +++ b/dev/changelog/52.0.0.md @@ -0,0 +1,78 @@ + + +# Apache DataFusion Python 52.0.0 Changelog + +This release consists of 26 commits from 9 contributors. See credits at the end of this changelog for more information. + +**Implemented enhancements:** + +- feat: add CatalogProviderList support [#1363](https://github.com/apache/datafusion-python/pull/1363) (timsaucer) +- feat: add support for generating JSON formatted substrait plan [#1376](https://github.com/apache/datafusion-python/pull/1376) (Prathamesh9284) +- feat: add regexp_instr function [#1382](https://github.com/apache/datafusion-python/pull/1382) (mesejo) + +**Fixed bugs:** + +- fix: mangled errors [#1377](https://github.com/apache/datafusion-python/pull/1377) (mesejo) + +**Documentation updates:** + +- docs: Clarify first_value usage in select vs aggregate [#1348](https://github.com/apache/datafusion-python/pull/1348) (AdMub) + +**Other:** + +- Release 51.0.0 [#1333](https://github.com/apache/datafusion-python/pull/1333) (timsaucer) +- Use explicit timer in unit test [#1338](https://github.com/apache/datafusion-python/pull/1338) (timsaucer) +- Add use_fabric_endpoint parameter to MicrosoftAzure class [#1357](https://github.com/apache/datafusion-python/pull/1357) (djouallah) +- Prepare for DF52 release [#1337](https://github.com/apache/datafusion-python/pull/1337) (timsaucer) +- build(deps): bump actions/checkout from 5 to 6 [#1310](https://github.com/apache/datafusion-python/pull/1310) (dependabot[bot]) +- build(deps): bump actions/download-artifact from 5 to 7 [#1321](https://github.com/apache/datafusion-python/pull/1321) (dependabot[bot]) +- build(deps): bump actions/upload-artifact from 4 to 6 [#1322](https://github.com/apache/datafusion-python/pull/1322) (dependabot[bot]) +- build(deps): bump actions/cache from 4 to 5 [#1323](https://github.com/apache/datafusion-python/pull/1323) (dependabot[bot]) +- Pass Field information back and forth when using scalar UDFs [#1299](https://github.com/apache/datafusion-python/pull/1299) (timsaucer) +- Update dependency minor versions to prepare for DF52 release [#1368](https://github.com/apache/datafusion-python/pull/1368) (timsaucer) +- Improve displayed error by using `DataFusionError`'s `Display` trait [#1370](https://github.com/apache/datafusion-python/pull/1370) (abey79) +- Enforce DataFrame display memory limits with `max_rows` + `min_rows` constraint (deprecate `repr_rows`) [#1367](https://github.com/apache/datafusion-python/pull/1367) (kosiew) +- Implement all CSV reader options [#1361](https://github.com/apache/datafusion-python/pull/1361) (timsaucer) +- chore: add confirmation before tarball is released [#1372](https://github.com/apache/datafusion-python/pull/1372) (milenkovicm) +- Build in debug mode for PRs [#1375](https://github.com/apache/datafusion-python/pull/1375) (timsaucer) +- minor: remove ffi test wheel from distribution artifact [#1378](https://github.com/apache/datafusion-python/pull/1378) (timsaucer) +- chore: update rust 2024 edition [#1371](https://github.com/apache/datafusion-python/pull/1371) (timsaucer) +- Fix Python UDAF list-of-timestamps return by enforcing list-valued scalars and caching PyArrow types [#1347](https://github.com/apache/datafusion-python/pull/1347) (kosiew) +- minor: update cargo dependencies [#1383](https://github.com/apache/datafusion-python/pull/1383) (timsaucer) +- chore: bump Python version for RAT checking [#1386](https://github.com/apache/datafusion-python/pull/1386) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 13 Tim Saucer + 4 dependabot[bot] + 2 Daniel Mesejo + 2 kosiew + 1 Adisa Mubarak (AdMub) + 1 Antoine Beyeler + 1 Dhanashri Prathamesh Iranna + 1 Marko Milenković + 1 Mimoune +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. + diff --git a/dev/check_crates_patch.py b/dev/check_crates_patch.py new file mode 100644 index 000000000..74e489e1f --- /dev/null +++ b/dev/check_crates_patch.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +# 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. + +"""Check that no Cargo.toml files contain [patch.crates-io] entries. + +Release builds must not depend on patched crates. During development it is +common to temporarily patch crates-io dependencies, but those patches must +be removed before creating a release. + +An empty [patch.crates-io] section is allowed. +""" + +import sys +from pathlib import Path + +import tomllib + + +def main() -> int: + errors: list[str] = [] + for cargo_toml in sorted(Path().rglob("Cargo.toml")): + if "target" in cargo_toml.parts: + continue + with Path.open(cargo_toml, "rb") as f: + data = tomllib.load(f) + patch = data.get("patch", {}).get("crates-io", {}) + if patch: + errors.append(str(cargo_toml)) + for name, spec in patch.items(): + errors.append(f" {name} = {spec}") + + if errors: + print("ERROR: Release builds must not contain [patch.crates-io] entries.") + print() + for line in errors: + print(line) + print() + print("Remove all [patch.crates-io] entries before creating a release.") + return 1 + + print("OK: No [patch.crates-io] entries found.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dev/create_license.py b/dev/create_license.py index 2a67cb8fd..acbf8587c 100644 --- a/dev/create_license.py +++ b/dev/create_license.py @@ -20,12 +20,11 @@ import json import subprocess +from pathlib import Path -subprocess.check_output(["cargo", "install", "cargo-license"]) data = subprocess.check_output( [ - "cargo", - "license", + "cargo-license", "--avoid-build-deps", "--avoid-dev-deps", "--do-not-bundle", @@ -248,5 +247,5 @@ result += "------------------\n\n" result += f"### {name} {version}\n* source: [{repository}]({repository})\n* license: {license}\n\n" -with open("LICENSE.txt", "w") as f: +with Path.open("LICENSE.txt", "w") as f: f.write(result) diff --git a/dev/release/README.md b/dev/release/README.md index f0b333999..ed28f4aa6 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -56,6 +56,8 @@ Before creating a new release: - a PR should be created and merged to update the major version number of the project - A new release branch should be created, such as `branch-0.8` +## Preparing a Release Candidate + ### Change Log We maintain a `CHANGELOG.md` so our users know what has been changed between releases. @@ -76,21 +78,20 @@ Categorizing pull requests Generating changelog content ``` -This process is not fully automated, so there are some additional manual steps: +### Update the version number -- Add the ASF header to the generated file -- Add a link to this changelog from the top-level `/datafusion/CHANGELOG.md` -- Add the following content (copy from the previous version's changelog and update as appropriate: +The only place you should need to update the version is in the root `Cargo.toml`. +After updating the toml file, run `cargo update` to update the cargo lock file. +If you do not want to update all the dependencies, you can instead run `cargo build` +which should only update the version number for `datafusion-python`. -``` -## [24.0.0](https://github.com/apache/datafusion-python/tree/24.0.0) (2023-05-06) +### Tag the Repository -[Full Changelog](https://github.com/apache/datafusion-python/compare/23.0.0...24.0.0) -``` +Commit the changes to the changelog and version. -### Preparing a Release Candidate - -### Tag the Repository +Assuming you have set up a remote to the `apache` repository rather than your personal fork, +you need to push a tag to start the CI process for release candidates. The following assumes +the upstream repository is called `apache`. ```bash git tag 0.8.0-rc1 @@ -103,7 +104,7 @@ git push apache 0.8.0-rc1 ./dev/release/create-tarball.sh 0.8.0 1 ``` -This will also create the email template to send to the mailing list. +This will also create the email template to send to the mailing list. Create a draft email using this content, but do not send until after completing the next step. @@ -153,13 +154,39 @@ This will create a file named `dist/datafusion-0.7.0.tar.gz`. Upload this to tes python3 -m twine upload --repository testpypi dist/datafusion-0.7.0.tar.gz ``` +### Run Verify Release Candidate Workflow + +Before sending the vote email, run the manually triggered GitHub Actions workflow +"Verify Release Candidate" and confirm all matrix jobs pass across the OS/architecture matrix +(for example, Linux, macOS, and Windows runners): + +1. Go to https://github.com/apache/datafusion-python/actions/workflows/verify-release-candidate.yml +2. Click "Run workflow" +3. Set `version` to the release version (for example, `52.0.0`) +4. Set `rc_number` to the RC number (for example, `0`) +5. Wait for all jobs to complete successfully + +Include a short note in the vote email template that this workflow was run across all OS/architecture +matrix entries and that all jobs passed. + +```text +Verification note: The manually triggered "Verify Release Candidate" workflow was run for version and rc_number across all configured OS/architecture matrix entries, and all matrix jobs completed successfully. +``` + ### Send the Email Send the email to start the vote. ## Verifying a Release -Running the unit tests against a testpypi release candidate: +Releases may be verified using `verify-release-candidate.sh`: + +```bash +git clone https://github.com/apache/datafusion-python.git +dev/release/verify-release-candidate.sh 48.0.0 1 +``` + +Alternatively, one can run unit tests against a testpypi release candidate: ```bash # clone a fresh repo @@ -178,11 +205,11 @@ source .venv/bin/activate # install release candidate pip install --extra-index-url https://test.pypi.org/simple/ datafusion==40.0.0 -# only dep needed to run tests is pytest -pip install pytest +# install test dependencies +pip install pytest numpy pytest-asyncio # run the tests -pytest --import-mode=importlib python/tests +pytest --import-mode=importlib python/tests -vv ``` Try running one of the examples from the top-level README, or write some custom Python code to query some available @@ -235,7 +262,7 @@ git push apache 0.8.0 Add the release to https://reporter.apache.org/addrelease.html?datafusion with a version name prefixed with `DATAFUSION-PYTHON`, for example `DATAFUSION-PYTHON-31.0.0`. -The release information is used to generate a template for a board report (see example from Apache Arrow +The release information is used to generate a template for a board report (see example from Apache Arrow [here](https://github.com/apache/arrow/pull/14357)). ### Delete old RCs and Releases diff --git a/dev/release/check-rat-report.py b/dev/release/check-rat-report.py index 0c9f4c326..72a35212e 100644 --- a/dev/release/check-rat-report.py +++ b/dev/release/check-rat-report.py @@ -21,6 +21,7 @@ import re import sys import xml.etree.ElementTree as ET +from pathlib import Path if len(sys.argv) != 3: sys.stderr.write("Usage: %s exclude_globs.lst rat_report.xml\n" % sys.argv[0]) @@ -29,7 +30,7 @@ exclude_globs_filename = sys.argv[1] xml_filename = sys.argv[2] -globs = [line.strip() for line in open(exclude_globs_filename)] +globs = [line.strip() for line in Path.open(exclude_globs_filename)] tree = ET.parse(xml_filename) root = tree.getroot() diff --git a/dev/release/release-tarball.sh b/dev/release/release-tarball.sh index 8c305a676..2b82d1bac 100755 --- a/dev/release/release-tarball.sh +++ b/dev/release/release-tarball.sh @@ -43,6 +43,13 @@ fi version=$1 rc=$2 +read -r -p "Proceed to release tarball for ${version}-rc${rc}? [y/N]: " answer +answer=${answer:-no} +if [ "${answer}" != "y" ]; then + echo "Cancelled tarball release!" + exit 1 +fi + tmp_dir=tmp-apache-datafusion-python-dist echo "Recreate temporary directory: ${tmp_dir}" diff --git a/dev/release/verify-release-candidate.sh b/dev/release/verify-release-candidate.sh index 2bfce0e2d..9591e0335 100755 --- a/dev/release/verify-release-candidate.sh +++ b/dev/release/verify-release-candidate.sh @@ -112,8 +112,17 @@ test_source_distribution() { curl https://sh.rustup.rs -sSf | sh -s -- -y --no-modify-path - export PATH=$RUSTUP_HOME/bin:$PATH - source $RUSTUP_HOME/env + # On Unix, rustup creates an env file. On Windows GitHub runners (MSYS bash), + # that file may not exist, so fall back to adding Cargo bin directly. + if [ -f "$CARGO_HOME/env" ]; then + # shellcheck disable=SC1090 + source "$CARGO_HOME/env" + elif [ -f "$RUSTUP_HOME/env" ]; then + # shellcheck disable=SC1090 + source "$RUSTUP_HOME/env" + else + export PATH="$CARGO_HOME/bin:$PATH" + fi # build and test rust @@ -126,10 +135,20 @@ test_source_distribution() { git clone https://github.com/apache/parquet-testing.git parquet-testing python3 -m venv .venv - source .venv/bin/activate - python3 -m pip install -U pip - python3 -m pip install -U maturin - maturin develop + if [ -x ".venv/bin/python" ]; then + VENV_PYTHON=".venv/bin/python" + elif [ -x ".venv/Scripts/python.exe" ]; then + VENV_PYTHON=".venv/Scripts/python.exe" + elif [ -x ".venv/Scripts/python" ]; then + VENV_PYTHON=".venv/Scripts/python" + else + echo "Unable to find python executable in virtual environment" + exit 1 + fi + + "$VENV_PYTHON" -m pip install -U pip + "$VENV_PYTHON" -m pip install -U maturin + "$VENV_PYTHON" -m maturin develop #TODO: we should really run tests here as well #python3 -m pytest diff --git a/docs/Makefile b/docs/Makefile index e65c8e250..49ebae372 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -35,4 +35,4 @@ help: # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) \ No newline at end of file + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) --fail-on-warning \ No newline at end of file diff --git a/docs/README.md b/docs/README.md index 2bffea9bd..502f1c2a1 100644 --- a/docs/README.md +++ b/docs/README.md @@ -59,7 +59,7 @@ firefox docs/build/html/index.html This documentation is hosted at https://datafusion.apache.org/python When the PR is merged to the `main` branch of the DataFusion -repository, a [github workflow](https://github.com/apache/datafusion-python/blob/main/.github/workflows/docs.yaml) which: +repository, a [github workflow](https://github.com/apache/datafusion-python/blob/main/.github/workflows/build.yml) which: 1. Builds the html content 2. Pushes the html content to the [`asf-site`](https://github.com/apache/datafusion-python/tree/asf-site) branch in this repository. @@ -67,4 +67,4 @@ repository, a [github workflow](https://github.com/apache/datafusion-python/blob The Apache Software Foundation provides https://arrow.apache.org/, which serves content based on the configuration in [.asf.yaml](https://github.com/apache/datafusion-python/blob/main/.asf.yaml), -which specifies the target as https://datafusion.apache.org/python. \ No newline at end of file +which specifies the target as https://datafusion.apache.org/python. diff --git a/docs/source/conf.py b/docs/source/conf.py index 0be03d81d..01813b032 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -71,6 +71,7 @@ autoapi_member_order = "groupwise" suppress_warnings = ["autoapi.python_import_resolution"] autoapi_python_class_content = "both" +autoapi_keep_files = False # set to True for debugging generated files def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa: ARG001 @@ -79,6 +80,9 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa ("class", "datafusion.DataFrame"), ("class", "datafusion.SessionContext"), ("module", "datafusion.common"), + # Duplicate modules (skip module-level docs to avoid duplication) + ("module", "datafusion.col"), + ("module", "datafusion.udf"), # Deprecated ("class", "datafusion.substrait.serde"), ("class", "datafusion.substrait.plan"), @@ -87,6 +91,13 @@ def autoapi_skip_member_fn(app, what, name, obj, skip, options) -> bool: # noqa ("method", "datafusion.context.SessionContext.tables"), ("method", "datafusion.dataframe.DataFrame.unnest_column"), ] + # Explicitly skip certain members listed above. These are either + # re-exports, duplicate module-level documentation, deprecated + # API surfaces, or private variables that would otherwise appear + # in the generated docs and cause confusing duplication. + # Keeping this explicit list avoids surprising entries in the + # AutoAPI output and gives us a single place to opt-out items + # when we intentionally hide them from the docs. if (what, name) in skip_contents: skip = True diff --git a/docs/source/contributor-guide/ffi.rst b/docs/source/contributor-guide/ffi.rst index c1f9806b3..e0158e0a2 100644 --- a/docs/source/contributor-guide/ffi.rst +++ b/docs/source/contributor-guide/ffi.rst @@ -15,6 +15,8 @@ .. specific language governing permissions and limitations .. under the License. +.. _ffi: + Python Extensions ================= @@ -34,7 +36,7 @@ as performant as possible and to utilize the features of DataFusion, you may dec your source in Rust and then expose it through `PyO3 `_ as a Python library. At first glance, it may appear the best way to do this is to add the ``datafusion-python`` -crate as a dependency, provide a ``PyTable``, and then to register it with the +crate as a dependency, provide a ``PyTable``, and then to register it with the ``SessionContext``. Unfortunately, this will not work. When you produce your code as a Python library and it needs to interact with the DataFusion @@ -137,6 +139,67 @@ and you want to create a sharable FFI counterpart, you could write: let my_provider = MyTableProvider::default(); let ffi_provider = FFI_TableProvider::new(Arc::new(my_provider), false, None); +.. _ffi_pyclass_mutability: + +PyO3 class mutability guidelines +-------------------------------- + +PyO3 bindings should present immutable wrappers whenever a struct stores shared or +interior-mutable state. In practice this means that any ``#[pyclass]`` containing an +``Arc>`` or similar synchronized primitive must opt into ``#[pyclass(frozen)]`` +unless there is a compelling reason not to. + +The :mod:`datafusion` configuration helpers illustrate the preferred pattern. The +``PyConfig`` class in :file:`src/config.rs` stores an ``Arc>`` and is +explicitly frozen so callers interact with configuration state through provided methods +instead of mutating the container directly: + +.. code-block:: rust + + #[pyclass(from_py_object, name = "Config", module = "datafusion", subclass, frozen)] + #[derive(Clone)] + pub(crate) struct PyConfig { + config: Arc>, + } + +The same approach applies to execution contexts. ``PySessionContext`` in +:file:`src/context.rs` stays frozen even though it shares mutable state internally via +``SessionContext``. This ensures PyO3 tracks borrows correctly while Python-facing APIs +clone the inner ``SessionContext`` or return new wrappers instead of mutating the +existing instance in place: + +.. code-block:: rust + + #[pyclass(from_py_object, frozen, name = "SessionContext", module = "datafusion", subclass)] + #[derive(Clone)] + pub struct PySessionContext { + pub ctx: SessionContext, + } + +Occasionally a type must remain mutable—for example when PyO3 attribute setters need to +update fields directly. In these rare cases add an inline justification so reviewers and +future contributors understand why ``frozen`` is unsafe to enable. ``DataTypeMap`` in +:file:`src/common/data_type.rs` includes such a comment because PyO3 still needs to track +field updates: + +.. code-block:: rust + + // TODO: This looks like this needs pyo3 tracking so leaving unfrozen for now + #[derive(Debug, Clone)] + #[pyclass(from_py_object, name = "DataTypeMap", module = "datafusion.common", subclass)] + pub struct DataTypeMap { + #[pyo3(get, set)] + pub arrow_type: PyDataType, + #[pyo3(get, set)] + pub python_type: PythonType, + #[pyo3(get, set)] + pub sql_type: SqlType, + } + +When reviewers encounter a mutable ``#[pyclass]`` without a comment, they should request +an explanation or ask that ``frozen`` be added. Keeping these wrappers frozen by default +helps avoid subtle bugs stemming from PyO3's interior mutability tracking. + If you were interfacing with a library that provided the above ``FFI_TableProvider`` and you needed to turn it back into an ``TableProvider``, you can turn it into a ``ForeignTableProvider`` with implements the ``TableProvider`` trait. @@ -169,14 +232,17 @@ can then be turned into a ``ForeignTableProvider`` the associated code is: .. code-block:: rust - let capsule = capsule.downcast::()?; - let provider = unsafe { capsule.reference::() }; + let capsule = capsule.cast::()?; + let data: NonNull = capsule + .pointer_checked(Some(name))? + .cast(); + let codec = unsafe { data.as_ref() }; By convention the ``datafusion-python`` library expects a Python object that has a ``TableProvider`` PyCapsule to have this capsule accessible by calling a function named ``__datafusion_table_provider__``. You can see a complete working example of how to share a ``TableProvider`` from one python library to DataFusion Python in the -`repository examples folder `_. +`repository examples folder `_. This section has been written using ``TableProvider`` as an example. It is the first extension that has been written using this approach and the most thoroughly implemented. @@ -195,7 +261,7 @@ optimization levels. If you wish to go down this route, there are two approaches have identified you can use. #. Re-export all of ``datafusion-python`` yourself with your extensions built in. -#. Carefully synchonize your software releases with the ``datafusion-python`` CI build +#. Carefully synchronize your software releases with the ``datafusion-python`` CI build system so that your libraries use the exact same compiler, features, and optimization level. diff --git a/docs/source/contributor-guide/introduction.rst b/docs/source/contributor-guide/introduction.rst index 2fba64111..33c2b274c 100644 --- a/docs/source/contributor-guide/introduction.rst +++ b/docs/source/contributor-guide/introduction.rst @@ -26,6 +26,10 @@ We welcome and encourage contributions of all kinds, such as: In addition to submitting new PRs, we have a healthy tradition of community members reviewing each other’s PRs. Doing so is a great way to help the community as well as get more familiar with Rust and the relevant codebases. +Before opening a pull request that touches PyO3 bindings, please review the +:ref:`PyO3 class mutability guidelines ` so you can flag missing +``#[pyclass(frozen)]`` annotations during development and review. + How to develop -------------- @@ -43,7 +47,7 @@ Bootstrap: # fetch this repo git clone git@github.com:apache/datafusion-python.git - # create the virtual enviornment + # create the virtual environment uv sync --dev --no-install-package datafusion # activate the environment source .venv/bin/activate diff --git a/docs/source/index.rst b/docs/source/index.rst index 558b2d572..134d41cb6 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -72,10 +72,12 @@ Example user-guide/introduction user-guide/basics user-guide/data-sources + user-guide/dataframe/index user-guide/common-operations/index user-guide/io/index user-guide/configuration user-guide/sql + user-guide/upgrade-guides .. _toc.contributor_guide: diff --git a/docs/source/user-guide/basics.rst b/docs/source/user-guide/basics.rst index f37378a41..7c6820461 100644 --- a/docs/source/user-guide/basics.rst +++ b/docs/source/user-guide/basics.rst @@ -20,8 +20,9 @@ Concepts ======== -In this section, we will cover a basic example to introduce a few key concepts. We will use the same -source file as described in the :ref:`Introduction `, the Pokemon data set. +In this section, we will cover a basic example to introduce a few key concepts. We will use the +2021 Yellow Taxi Trip Records (`download `_), +from the `TLC Trip Record Data `_. .. ipython:: python @@ -72,6 +73,8 @@ DataFrames are typically created by calling a method on :py:class:`~datafusion.c calling the transformation methods, such as :py:func:`~datafusion.dataframe.DataFrame.filter`, :py:func:`~datafusion.dataframe.DataFrame.select`, :py:func:`~datafusion.dataframe.DataFrame.aggregate`, and :py:func:`~datafusion.dataframe.DataFrame.limit` to build up a query definition. +For more details on working with DataFrames, including visualization options and conversion to other formats, see :doc:`dataframe/index`. + Expressions ----------- diff --git a/docs/source/user-guide/common-operations/expressions.rst b/docs/source/user-guide/common-operations/expressions.rst index e94e1a6b5..7848b4ee7 100644 --- a/docs/source/user-guide/common-operations/expressions.rst +++ b/docs/source/user-guide/common-operations/expressions.rst @@ -64,7 +64,7 @@ Arrays ------ For columns that contain arrays of values, you can access individual elements of the array by index -using bracket indexing. This is similar to callling the function +using bracket indexing. This is similar to calling the function :py:func:`datafusion.functions.array_element`, except that array indexing using brackets is 0 based, similar to Python arrays and ``array_element`` is 1 based indexing to be compatible with other SQL approaches. @@ -82,6 +82,13 @@ approaches. Indexing an element of an array via ``[]`` starts at index 0 whereas :py:func:`~datafusion.functions.array_element` starts at index 1. +Starting in DataFusion 49.0.0 you can also create slices of array elements using +slice syntax from Python. + +.. ipython:: python + + df.select(col("a")[1:3].alias("second_two_elements")) + To check if an array is empty, you can use the function :py:func:`datafusion.functions.array_empty` or `datafusion.functions.empty`. This function returns a boolean indicating whether the array is empty. diff --git a/docs/source/user-guide/common-operations/functions.rst b/docs/source/user-guide/common-operations/functions.rst index 12097be8f..ccb47a4e7 100644 --- a/docs/source/user-guide/common-operations/functions.rst +++ b/docs/source/user-guide/common-operations/functions.rst @@ -129,3 +129,24 @@ The function :py:func:`~datafusion.functions.in_list` allows to check a column f .limit(20) .to_pandas() ) + + +Handling Missing Values +======================= + +DataFusion provides methods to handle missing values in DataFrames: + +fill_null +--------- + +The ``fill_null()`` method replaces NULL values in specified columns with a provided value: + +.. code-block:: python + + # Fill all NULL values with 0 where possible + df = df.fill_null(0) + + # Fill NULL values only in specific string columns + df = df.fill_null("missing", subset=["name", "category"]) + +The fill value will be cast to match each column's type. If casting fails for a column, that column remains unchanged. diff --git a/docs/source/user-guide/common-operations/joins.rst b/docs/source/user-guide/common-operations/joins.rst index 40d922150..1d9d70385 100644 --- a/docs/source/user-guide/common-operations/joins.rst +++ b/docs/source/user-guide/common-operations/joins.rst @@ -101,4 +101,36 @@ the right table. .. ipython:: python - left.join(right, left_on="customer_id", right_on="id", how="anti") \ No newline at end of file + left.join(right, left_on="customer_id", right_on="id", how="anti") + +Duplicate Keys +-------------- + +It is common to join two DataFrames on a common column name. Starting in +version 51.0.0, ``datafusion-python``` will now coalesce on column with identical names by +default. This reduces problems with ambiguous column selection after joins. +You can disable this feature by setting the parameter ``coalesce_duplicate_keys`` +to ``False``. + +.. ipython:: python + + left = ctx.from_pydict( + { + "id": [1, 2, 3], + "customer": ["Alice", "Bob", "Charlie"], + } + ) + + right = ctx.from_pylist([ + {"id": 1, "name": "CityCabs"}, + {"id": 2, "name": "MetroRide"}, + {"id": 5, "name": "UrbanGo"}, + ]) + + left.join(right, "id", how="inner") + +In contrast to the above example, if we wish to get both columns: + +.. ipython:: python + + left.join(right, "id", how="inner", coalesce_duplicate_keys=False) diff --git a/docs/source/user-guide/common-operations/udf-and-udfa.rst b/docs/source/user-guide/common-operations/udf-and-udfa.rst index ffd7a05cb..f669721a3 100644 --- a/docs/source/user-guide/common-operations/udf-and-udfa.rst +++ b/docs/source/user-guide/common-operations/udf-and-udfa.rst @@ -26,7 +26,7 @@ Scalar Functions When writing a user-defined function that can operate on a row by row basis, these are called Scalar Functions. You can define your own scalar function by calling -:py:func:`~datafusion.udf.ScalarUDF.udf` . +:py:func:`~datafusion.user_defined.ScalarUDF.udf` . The basic definition of a scalar UDF is a python function that takes one or more `pyarrow `_ arrays and returns a single array as @@ -90,12 +90,23 @@ converting to Python objects to do the evaluation. df.select(col("a"), is_null_arr(col("a")).alias("is_null")).show() +In this example we passed the PyArrow ``DataType`` when we defined the function +by calling ``udf()``. If you need additional control, such as specifying +metadata or nullability of the input or output, you can instead specify a +PyArrow ``Field``. + +If you need to write a custom function but do not want to incur the performance +cost of converting to Python objects and back, a more advanced approach is to +write Rust based UDFs and to expose them to Python. There is an example in the +`DataFusion blog `_ +describing how to do this. + Aggregate Functions ------------------- -The :py:func:`~datafusion.udf.AggregateUDF.udaf` function allows you to define User-Defined +The :py:func:`~datafusion.user_defined.AggregateUDF.udaf` function allows you to define User-Defined Aggregate Functions (UDAFs). To use this you must implement an -:py:class:`~datafusion.udf.Accumulator` that determines how the aggregation is performed. +:py:class:`~datafusion.user_defined.Accumulator` that determines how the aggregation is performed. When defining a UDAF there are four methods you need to implement. The ``update`` function takes the array(s) of input and updates the internal state of the accumulator. You should define this function @@ -112,7 +123,7 @@ also see how the inputs to ``update`` and ``merge`` differ. .. code-block:: python - import pyarrow + import pyarrow as pa import pyarrow.compute import datafusion from datafusion import col, udaf, Accumulator @@ -125,16 +136,16 @@ also see how the inputs to ``update`` and ``merge`` differ. def __init__(self): self._sum = 0.0 - def update(self, values_a: pyarrow.Array, values_b: pyarrow.Array) -> None: + def update(self, values_a: pa.Array, values_b: pa.Array) -> None: self._sum = self._sum + pyarrow.compute.sum(values_a).as_py() - pyarrow.compute.sum(values_b).as_py() - def merge(self, states: List[pyarrow.Array]) -> None: + def merge(self, states: list[pa.Array]) -> None: self._sum = self._sum + pyarrow.compute.sum(states[0]).as_py() - def state(self) -> pyarrow.Array: - return pyarrow.array([self._sum]) + def state(self) -> list[pa.Scalar]: + return [pyarrow.scalar(self._sum)] - def evaluate(self) -> pyarrow.Scalar: + def evaluate(self) -> pa.Scalar: return pyarrow.scalar(self._sum) ctx = datafusion.SessionContext() @@ -145,16 +156,35 @@ also see how the inputs to ``update`` and ``merge`` differ. } ) - my_udaf = udaf(MyAccumulator, [pyarrow.float64(), pyarrow.float64()], pyarrow.float64(), [pyarrow.float64()], 'stable') + my_udaf = udaf(MyAccumulator, [pa.float64(), pa.float64()], pa.float64(), [pa.float64()], 'stable') df.aggregate([], [my_udaf(col("a"), col("b")).alias("col_diff")]) +FAQ +^^^ + +**How do I return a list from a UDAF?** + +Both the ``evaluate`` and the ``state`` functions expect to return scalar values. +If you wish to return a list array as a scalar value, the best practice is to +wrap the values in a ``pyarrow.Scalar`` object. For example, you can return a +timestamp list with ``pa.scalar([...], type=pa.list_(pa.timestamp("ms")))`` and +register the appropriate return or state types as +``return_type=pa.list_(pa.timestamp("ms"))`` and +``state_type=[pa.list_(pa.timestamp("ms"))]``, respectively. + +As of DataFusion 52.0.0 , you can pass return any Python object, including a +PyArrow array, as the return value(s) for these functions and DataFusion will +attempt to create a scalar type from the value. DataFusion has been tested to +convert PyArrow, nanoarrow, and arro3 objects as well as primitive data types +like integers, strings, and so on. + Window Functions ---------------- To implement a User-Defined Window Function (UDWF) you must call the -:py:func:`~datafusion.udf.WindowUDF.udwf` function using a class that implements the abstract -class :py:class:`~datafusion.udf.WindowEvaluator`. +:py:func:`~datafusion.user_defined.WindowUDF.udwf` function using a class that implements the abstract +class :py:class:`~datafusion.user_defined.WindowEvaluator`. There are three methods of evaluation of UDWFs. @@ -207,7 +237,7 @@ determine which evaluate functions are called. import pyarrow as pa from datafusion import udwf, col, SessionContext - from datafusion.udf import WindowEvaluator + from datafusion.user_defined import WindowEvaluator class ExponentialSmooth(WindowEvaluator): def __init__(self, alpha: float) -> None: @@ -242,3 +272,35 @@ determine which evaluate functions are called. }) df.select("a", exp_smooth(col("a")).alias("smooth_a")).show() + +Table Functions +--------------- + +User Defined Table Functions are slightly different than the other functions +described here. These functions take any number of `Expr` arguments, but only +literal expressions are supported. Table functions must return a Table +Provider as described in the ref:`_io_custom_table_provider` page. + +Once you have a table function, you can register it with the session context +by using :py:func:`datafusion.context.SessionContext.register_udtf`. + +There are examples of both rust backed and python based table functions in the +examples folder of the repository. If you have a rust backed table function +that you wish to expose via PyO3, you need to expose it as a ``PyCapsule``. + +.. code-block:: rust + + #[pymethods] + impl MyTableFunction { + fn __datafusion_table_function__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let name = cr"datafusion_table_function".into(); + + let func = self.clone(); + let provider = FFI_TableFunction::new(Arc::new(func), None); + + PyCapsule::new(py, provider, Some(name)) + } + } diff --git a/docs/source/user-guide/common-operations/windows.rst b/docs/source/user-guide/common-operations/windows.rst index 8225d125a..c8fdea8f4 100644 --- a/docs/source/user-guide/common-operations/windows.rst +++ b/docs/source/user-guide/common-operations/windows.rst @@ -24,14 +24,14 @@ In this section you will learn about window functions. A window function utilize multiple rows to produce a result for each individual row, unlike an aggregate function that provides a single value for multiple rows. -The window functions are availble in the :py:mod:`~datafusion.functions` module. +The window functions are available in the :py:mod:`~datafusion.functions` module. We'll use the pokemon dataset (from Ritchie Vink) in the following examples. .. ipython:: python from datafusion import SessionContext - from datafusion import col + from datafusion import col, lit from datafusion import functions as f ctx = SessionContext() @@ -99,8 +99,8 @@ If you do not specify a Window Frame, the frame will be set depending on the fol criteria. * If an ``order_by`` clause is set, the default window frame is defined as the rows between - unbounded preceeding and the current row. -* If an ``order_by`` is not set, the default frame is defined as the rows betwene unbounded + unbounded preceding and the current row. +* If an ``order_by`` is not set, the default frame is defined as the rows between unbounded and unbounded following (the entire partition). Window Frames are defined by three parameters: unit type, starting bound, and ending bound. @@ -116,20 +116,18 @@ The unit types available are: ``order_by`` clause. In this example we perform a "rolling average" of the speed of the current Pokemon and the -two preceeding rows. +two preceding rows. .. ipython:: python - from datafusion.expr import WindowFrame + from datafusion.expr import Window, WindowFrame df.select( col('"Name"'), col('"Speed"'), - f.window("avg", - [col('"Speed"')], - order_by=[col('"Speed"')], - window_frame=WindowFrame("rows", 2, 0) - ).alias("Previous Speed") + f.avg(col('"Speed"')) + .over(Window(window_frame=WindowFrame("rows", 2, 0), order_by=[col('"Speed"')])) + .alias("Previous Speed"), ) Null Treatment @@ -151,21 +149,27 @@ it's ``Type 2`` column that are null. from datafusion.common import NullTreatment - df.filter(col('"Type 1"') == lit("Bug")).select( + df.filter(col('"Type 1"') == lit("Bug")).select( '"Name"', '"Type 2"', - f.window("last_value", [col('"Type 2"')]) - .window_frame(WindowFrame("rows", None, 0)) - .order_by(col('"Speed"')) - .null_treatment(NullTreatment.IGNORE_NULLS) - .build() - .alias("last_wo_null"), - f.window("last_value", [col('"Type 2"')]) - .window_frame(WindowFrame("rows", None, 0)) - .order_by(col('"Speed"')) - .null_treatment(NullTreatment.RESPECT_NULLS) - .build() - .alias("last_with_null") + f.last_value(col('"Type 2"')) + .over( + Window( + window_frame=WindowFrame("rows", None, 0), + order_by=[col('"Speed"')], + null_treatment=NullTreatment.IGNORE_NULLS, + ) + ) + .alias("last_wo_null"), + f.last_value(col('"Type 2"')) + .over( + Window( + window_frame=WindowFrame("rows", None, 0), + order_by=[col('"Speed"')], + null_treatment=NullTreatment.RESPECT_NULLS, + ) + ) + .alias("last_with_null"), ) Aggregate Functions diff --git a/docs/source/user-guide/configuration.rst b/docs/source/user-guide/configuration.rst index db200a46a..f8e613cd4 100644 --- a/docs/source/user-guide/configuration.rst +++ b/docs/source/user-guide/configuration.rst @@ -15,6 +15,8 @@ .. specific language governing permissions and limitations .. under the License. +.. _configuration: + Configuration ============= @@ -46,6 +48,141 @@ a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.conte ctx = SessionContext(config, runtime) print(ctx) +Maximizing CPU Usage +-------------------- + +DataFusion uses partitions to parallelize work. For small queries the +default configuration (number of CPU cores) is often sufficient, but to +fully utilize available hardware you can tune how many partitions are +created and when DataFusion will repartition data automatically. + +Configure a ``SessionContext`` with a higher partition count: + +.. code-block:: python + + from datafusion import SessionConfig, SessionContext + + # allow up to 16 concurrent partitions + config = SessionConfig().with_target_partitions(16) + ctx = SessionContext(config) + +Automatic repartitioning for joins, aggregations, window functions and +other operations can be enabled to increase parallelism: + +.. code-block:: python + + config = ( + SessionConfig() + .with_target_partitions(16) + .with_repartition_joins(True) + .with_repartition_aggregations(True) + .with_repartition_windows(True) + ) + +Manual repartitioning is available on DataFrames when you need precise +control: + +.. code-block:: python + + from datafusion import col + + df = ctx.read_parquet("data.parquet") + + # Evenly divide into 16 partitions + df = df.repartition(16) + + # Or partition by the hash of a column + df = df.repartition_by_hash(col("a"), num=16) + + result = df.collect() + + +Benchmark Example +^^^^^^^^^^^^^^^^^ + +The repository includes a benchmark script that demonstrates how to maximize CPU usage +with DataFusion. The :code:`benchmarks/max_cpu_usage.py` script shows a practical example +of configuring DataFusion for optimal parallelism. + +You can run the benchmark script to see the impact of different configuration settings: + +.. code-block:: bash + + # Run with default settings (uses all CPU cores) + python benchmarks/max_cpu_usage.py + + # Run with specific number of rows and partitions + python benchmarks/max_cpu_usage.py --rows 5000000 --partitions 16 + + # See all available options + python benchmarks/max_cpu_usage.py --help + +Here's an example showing the performance difference between single and multiple partitions: + +.. code-block:: bash + + # Single partition - slower processing + $ python benchmarks/max_cpu_usage.py --rows=10000000 --partitions 1 + Processed 10000000 rows using 1 partitions in 0.107s + + # Multiple partitions - faster processing + $ python benchmarks/max_cpu_usage.py --rows=10000000 --partitions 10 + Processed 10000000 rows using 10 partitions in 0.038s + +This example demonstrates nearly 3x performance improvement (0.107s vs 0.038s) when using +10 partitions instead of 1, showcasing how proper partitioning can significantly improve +CPU utilization and query performance. + +The script demonstrates several key optimization techniques: + +1. **Higher target partition count**: Uses :code:`with_target_partitions()` to set the number of concurrent partitions +2. **Automatic repartitioning**: Enables repartitioning for joins, aggregations, and window functions +3. **Manual repartitioning**: Uses :code:`repartition()` to ensure all partitions are utilized +4. **CPU-intensive operations**: Performs aggregations that can benefit from parallelization + +The benchmark creates synthetic data and measures the time taken to perform a sum aggregation +across the specified number of partitions. This helps you understand how partition configuration +affects performance on your specific hardware. + +Important Considerations +"""""""""""""""""""""""" + +The provided benchmark script demonstrates partitioning concepts using synthetic in-memory data +and simple aggregation operations. While useful for understanding basic configuration principles, +actual performance in production environments may vary significantly based on numerous factors: + +**Data Sources and I/O Characteristics:** + +- **Table providers**: Performance differs greatly between Parquet files, CSV files, databases, and cloud storage +- **Storage type**: Local SSD, network-attached storage, and cloud storage have vastly different characteristics +- **Network latency**: Remote data sources introduce additional latency considerations +- **File sizes and distribution**: Large files may benefit differently from partitioning than many small files + +**Query and Workload Characteristics:** + +- **Operation complexity**: Simple aggregations versus complex joins, window functions, or nested queries +- **Data distribution**: Skewed data may not partition evenly, affecting parallel efficiency +- **Memory usage**: Large datasets may require different memory management strategies +- **Concurrent workloads**: Multiple queries running simultaneously affect resource allocation + +**Hardware and Environment Factors:** + +- **CPU architecture**: Different processors have varying parallel processing capabilities +- **Available memory**: Limited RAM may require different optimization strategies +- **System load**: Other applications competing for resources affect DataFusion performance + +**Recommendations for Production Use:** + +To optimize DataFusion for your specific use case, it is strongly recommended to: + +1. **Create custom benchmarks** using your actual data sources, formats, and query patterns +2. **Test with representative data volumes** that match your production workloads +3. **Measure end-to-end performance** including data loading, processing, and result handling +4. **Evaluate different configuration combinations** for your specific hardware and workload +5. **Monitor resource utilization** (CPU, memory, I/O) to identify bottlenecks in your environment + +This approach will provide more accurate insights into how DataFusion configuration options +will impact your particular applications and infrastructure. -You can read more about available :py:class:`~datafusion.context.SessionConfig` options in the `rust DataFusion Configuration guide `_, +For more information about available :py:class:`~datafusion.context.SessionConfig` options, see the `rust DataFusion Configuration guide `_, and about :code:`RuntimeEnvBuilder` options in the rust `online API documentation `_. diff --git a/docs/source/user-guide/data-sources.rst b/docs/source/user-guide/data-sources.rst index ba5967c97..26f1303c4 100644 --- a/docs/source/user-guide/data-sources.rst +++ b/docs/source/user-guide/data-sources.rst @@ -25,7 +25,7 @@ DataFusion provides a wide variety of ways to get data into a DataFrame to perfo Local file ---------- -DataFusion has the abilty to read from a variety of popular file formats, such as :ref:`Parquet `, +DataFusion has the ability to read from a variety of popular file formats, such as :ref:`Parquet `, :ref:`CSV `, :ref:`JSON `, and :ref:`AVRO `. .. ipython:: python @@ -120,7 +120,7 @@ DataFusion can import DataFrames directly from other libraries, such as `Polars `_ and `Pandas `_. Since DataFusion version 42.0.0, any DataFrame library that supports the Arrow FFI PyCapsule interface can be imported to DataFusion using the -:py:func:`~datafusion.context.SessionContext.from_arrow` function. Older verions of Polars may +:py:func:`~datafusion.context.SessionContext.from_arrow` function. Older versions of Polars may not support the arrow interface. In those cases, you can still import via the :py:func:`~datafusion.context.SessionContext.from_polars` function. @@ -154,11 +154,11 @@ as Delta Lake. This will require a recent version of from deltalake import DeltaTable delta_table = DeltaTable("path_to_table") - ctx.register_table_provider("my_delta_table", delta_table) + ctx.register_table("my_delta_table", delta_table) df = ctx.table("my_delta_table") df.show() -On older versions of ``deltalake`` (prior to 0.22) you can use the +On older versions of ``deltalake`` (prior to 0.22) you can use the `Arrow DataSet `_ interface to import to DataFusion, but this does not support features such as filter push down which can lead to a significant performance difference. @@ -172,10 +172,41 @@ which can lead to a significant performance difference. df = ctx.table("my_delta_table") df.show() -Iceberg -------- +Apache Iceberg +-------------- -Coming soon! +DataFusion 45.0.0 and later support the ability to register Apache Iceberg tables as table providers through the Custom Table Provider interface. + +This requires either the `pyiceberg `__ library (>=0.10.0) or the `pyiceberg-core `__ library (>=0.5.0). + +* The ``pyiceberg-core`` library exposes Iceberg Rust's implementation of the Custom Table Provider interface as python bindings. +* The ``pyiceberg`` library utilizes the ``pyiceberg-core`` python bindings under the hood and provides a native way for Python users to interact with the DataFusion. + +.. code-block:: python + + from datafusion import SessionContext + from pyiceberg.catalog import load_catalog + import pyarrow as pa + + # Load catalog and create/load a table + catalog = load_catalog("catalog", type="in-memory") + catalog.create_namespace_if_not_exists("default") + + # Create some sample data + data = pa.table({"x": [1, 2, 3], "y": [4, 5, 6]}) + iceberg_table = catalog.create_table("default.test", schema=data.schema) + iceberg_table.append(data) + + # Register the table with DataFusion + ctx = SessionContext() + ctx.register_table_provider("test", iceberg_table) + + # Query the table using DataFusion + ctx.table("test").show() + + +Note that the Datafusion integration rely on features from the `Iceberg Rust `_ implementation instead of the `PyIceberg `_ implementation. +Features that are available in PyIceberg but not yet in Iceberg Rust will not be available when using DataFusion. Custom Table Provider --------------------- @@ -183,5 +214,61 @@ Custom Table Provider You can implement a custom Data Provider in Rust and expose it to DataFusion through the the interface as describe in the :ref:`Custom Table Provider ` section. This is an advanced topic, but a -`user example `_ +`user example `_ is provided in the DataFusion repository. + +Catalog +======= + +A common technique for organizing tables is using a three level hierarchical approach. DataFusion +supports this form of organizing using the :py:class:`~datafusion.catalog.Catalog`, +:py:class:`~datafusion.catalog.Schema`, and :py:class:`~datafusion.catalog.Table`. By default, +a :py:class:`~datafusion.context.SessionContext` comes with a single Catalog and a single Schema +with the names ``datafusion`` and ``default``, respectively. + +The default implementation uses an in-memory approach to the catalog and schema. We have support +for adding additional in-memory catalogs and schemas. This can be done like in the following +example: + +.. code-block:: python + + from datafusion.catalog import Catalog, Schema + + my_catalog = Catalog.memory_catalog() + my_schema = Schema.memory_schema() + + my_catalog.register_schema("my_schema_name", my_schema) + + ctx.register_catalog("my_catalog_name", my_catalog) + +You could then register tables in ``my_schema`` and access them either through the DataFrame +API or via sql commands such as ``"SELECT * from my_catalog_name.my_schema_name.my_table"``. + +User Defined Catalog and Schema +------------------------------- + +If the in-memory catalogs are insufficient for your uses, there are two approaches you can take +to implementing a custom catalog and/or schema. In the below discussion, we describe how to +implement these for a Catalog, but the approach to implementing for a Schema is nearly +identical. + +DataFusion supports Catalogs written in either Rust or Python. If you write a Catalog in Rust, +you will need to export it as a Python library via PyO3. There is a complete example of a +catalog implemented this way in the +`examples folder `_ +of our repository. Writing catalog providers in Rust provides typically can lead to significant +performance improvements over the Python based approach. + +To implement a Catalog in Python, you will need to inherit from the abstract base class +:py:class:`~datafusion.catalog.CatalogProvider`. There are examples in the +`unit tests `_ of +implementing a basic Catalog in Python where we simply keep a dictionary of the +registered Schemas. + +One important note for developers is that when we have a Catalog defined in Python, we have +two different ways of accessing this Catalog. First, we register the catalog with a Rust +wrapper. This allows for any rust based code to call the Python functions as necessary. +Second, if the user access the Catalog via the Python API, we identify this and return back +the original Python object that implements the Catalog. This is an important distinction +for developers because we do *not* return a Python wrapper around the Rust wrapper of the +original Python object. diff --git a/docs/source/user-guide/dataframe/index.rst b/docs/source/user-guide/dataframe/index.rst new file mode 100644 index 000000000..510bcbc68 --- /dev/null +++ b/docs/source/user-guide/dataframe/index.rst @@ -0,0 +1,371 @@ +.. 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. + +DataFrames +========== + +Overview +-------- + +The ``DataFrame`` class is the core abstraction in DataFusion that represents tabular data and operations +on that data. DataFrames provide a flexible API for transforming data through various operations such as +filtering, projection, aggregation, joining, and more. + +A DataFrame represents a logical plan that is lazily evaluated. The actual execution occurs only when +terminal operations like ``collect()``, ``show()``, or ``to_pandas()`` are called. + +Creating DataFrames +------------------- + +DataFrames can be created in several ways: + +* From SQL queries via a ``SessionContext``: + + .. code-block:: python + + from datafusion import SessionContext + + ctx = SessionContext() + df = ctx.sql("SELECT * FROM your_table") + +* From registered tables: + + .. code-block:: python + + df = ctx.table("your_table") + +* From various data sources: + + .. code-block:: python + + # From CSV files (see :ref:`io_csv` for detailed options) + df = ctx.read_csv("path/to/data.csv") + + # From Parquet files (see :ref:`io_parquet` for detailed options) + df = ctx.read_parquet("path/to/data.parquet") + + # From JSON files (see :ref:`io_json` for detailed options) + df = ctx.read_json("path/to/data.json") + + # From Avro files (see :ref:`io_avro` for detailed options) + df = ctx.read_avro("path/to/data.avro") + + # From Pandas DataFrame + import pandas as pd + pandas_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + df = ctx.from_pandas(pandas_df) + + # From Arrow data + import pyarrow as pa + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3]), pa.array([4, 5, 6])], + names=["a", "b"] + ) + df = ctx.from_arrow(batch) + +For detailed information about reading from different data sources, see the :doc:`I/O Guide <../io/index>`. +For custom data sources, see :ref:`io_custom_table_provider`. + +Common DataFrame Operations +--------------------------- + +DataFusion's DataFrame API offers a wide range of operations: + +.. code-block:: python + + from datafusion import column, literal + + # Select specific columns + df = df.select("col1", "col2") + + # Select with expressions + df = df.select(column("a") + column("b"), column("a") - column("b")) + + # Filter rows (expressions or SQL strings) + df = df.filter(column("age") > literal(25)) + df = df.filter("age > 25") + + # Add computed columns + df = df.with_column("full_name", column("first_name") + literal(" ") + column("last_name")) + + # Multiple column additions + df = df.with_columns( + (column("a") + column("b")).alias("sum"), + (column("a") * column("b")).alias("product") + ) + + # Sort data + df = df.sort(column("age").sort(ascending=False)) + + # Join DataFrames + df = df1.join(df2, on="user_id", how="inner") + + # Aggregate data + from datafusion import functions as f + df = df.aggregate( + [], # Group by columns (empty for global aggregation) + [f.sum(column("amount")).alias("total_amount")] + ) + + # Limit rows + df = df.limit(100) + + # Drop columns + df = df.drop("temporary_column") + +Column Names as Function Arguments +---------------------------------- + +Some ``DataFrame`` methods accept column names when an argument refers to an +existing column. These include: + +* :py:meth:`~datafusion.DataFrame.select` +* :py:meth:`~datafusion.DataFrame.sort` +* :py:meth:`~datafusion.DataFrame.drop` +* :py:meth:`~datafusion.DataFrame.join` (``on`` argument) +* :py:meth:`~datafusion.DataFrame.aggregate` (grouping columns) + +See the full function documentation for details on any specific function. + +Note that :py:meth:`~datafusion.DataFrame.join_on` expects ``col()``/``column()`` expressions rather than plain strings. + +For such methods, you can pass column names directly: + +.. code-block:: python + + from datafusion import col, functions as f + + df.sort('id') + df.aggregate('id', [f.count(col('value'))]) + +The same operation can also be written with explicit column expressions, using either ``col()`` or ``column()``: + +.. code-block:: python + + from datafusion import col, column, functions as f + + df.sort(col('id')) + df.aggregate(column('id'), [f.count(col('value'))]) + +Note that ``column()`` is an alias of ``col()``, so you can use either name; the example above shows both in action. + +Whenever an argument represents an expression—such as in +:py:meth:`~datafusion.DataFrame.filter` or +:py:meth:`~datafusion.DataFrame.with_column`—use ``col()`` to reference +columns. The comparison and arithmetic operators on ``Expr`` will automatically +convert any non-``Expr`` value into a literal expression, so writing + +.. code-block:: python + + from datafusion import col + df.filter(col("age") > 21) + +is equivalent to using ``lit(21)`` explicitly. Use ``lit()`` (also available +as ``literal()``) when you need to construct a literal expression directly. + +Terminal Operations +------------------- + +To materialize the results of your DataFrame operations: + +.. code-block:: python + + # Collect all data as PyArrow RecordBatches + result_batches = df.collect() + + # Convert to various formats + pandas_df = df.to_pandas() # Pandas DataFrame + polars_df = df.to_polars() # Polars DataFrame + arrow_table = df.to_arrow_table() # PyArrow Table + py_dict = df.to_pydict() # Python dictionary + py_list = df.to_pylist() # Python list of dictionaries + + # Display results + df.show() # Print tabular format to console + + # Count rows + count = df.count() + + # Collect a single column of data as a PyArrow Array + arr = df.collect_column("age") + +Zero-copy streaming to Arrow-based Python libraries +--------------------------------------------------- + +DataFusion DataFrames implement the ``__arrow_c_stream__`` protocol, enabling +zero-copy, lazy streaming into Arrow-based Python libraries. With the streaming +protocol, batches are produced on demand. + +.. note:: + + The protocol is implementation-agnostic and works with any Python library + that understands the Arrow C streaming interface (for example, PyArrow + or other Arrow-compatible implementations). The sections below provide a + short PyArrow-specific example and general guidance for other + implementations. + +PyArrow +------- + +.. code-block:: python + + import pyarrow as pa + + # Create a PyArrow RecordBatchReader without materializing all batches + reader = pa.RecordBatchReader.from_stream(df) + for batch in reader: + ... # process each batch as it is produced + +DataFrames are also iterable, yielding :class:`datafusion.RecordBatch` +objects lazily so you can loop over results directly without importing +PyArrow: + +.. code-block:: python + + for batch in df: + ... # each batch is a ``datafusion.RecordBatch`` + +Each batch exposes ``to_pyarrow()``, allowing conversion to a PyArrow +table. ``pa.table(df)`` collects the entire DataFrame eagerly into a +PyArrow table: + +.. code-block:: python + + import pyarrow as pa + table = pa.table(df) + +Asynchronous iteration is supported as well, allowing integration with +``asyncio`` event loops: + +.. code-block:: python + + async for batch in df: + ... # process each batch as it is produced + +To work with the stream directly, use ``execute_stream()``, which returns a +:class:`~datafusion.RecordBatchStream`. + +.. code-block:: python + + stream = df.execute_stream() + for batch in stream: + ... + +Execute as Stream +^^^^^^^^^^^^^^^^^ + +For finer control over streaming execution, use +:py:meth:`~datafusion.DataFrame.execute_stream` to obtain a +:py:class:`datafusion.RecordBatchStream`: + +.. code-block:: python + + stream = df.execute_stream() + for batch in stream: + ... # process each batch as it is produced + +.. tip:: + + To get a PyArrow reader instead, call + + ``pa.RecordBatchReader.from_stream(df)``. + +When partition boundaries are important, +:py:meth:`~datafusion.DataFrame.execute_stream_partitioned` +returns an iterable of :py:class:`datafusion.RecordBatchStream` objects, one per +partition: + +.. code-block:: python + + for stream in df.execute_stream_partitioned(): + for batch in stream: + ... # each stream yields RecordBatches + +To process partitions concurrently, first collect the streams into a list +and then poll each one in a separate ``asyncio`` task: + +.. code-block:: python + + import asyncio + + async def consume(stream): + async for batch in stream: + ... + + streams = list(df.execute_stream_partitioned()) + await asyncio.gather(*(consume(s) for s in streams)) + +See :doc:`../io/arrow` for additional details on the Arrow interface. + +HTML Rendering +-------------- + +When working in Jupyter notebooks or other environments that support HTML rendering, DataFrames will +automatically display as formatted HTML tables. For detailed information about customizing HTML +rendering, formatting options, and advanced styling, see :doc:`rendering`. + +Core Classes +------------ + +**DataFrame** + The main DataFrame class for building and executing queries. + + See: :py:class:`datafusion.DataFrame` + +**SessionContext** + The primary entry point for creating DataFrames from various data sources. + + Key methods for DataFrame creation: + + * :py:meth:`~datafusion.SessionContext.read_csv` - Read CSV files + * :py:meth:`~datafusion.SessionContext.read_parquet` - Read Parquet files + * :py:meth:`~datafusion.SessionContext.read_json` - Read JSON files + * :py:meth:`~datafusion.SessionContext.read_avro` - Read Avro files + * :py:meth:`~datafusion.SessionContext.table` - Access registered tables + * :py:meth:`~datafusion.SessionContext.sql` - Execute SQL queries + * :py:meth:`~datafusion.SessionContext.from_pandas` - Create from Pandas DataFrame + * :py:meth:`~datafusion.SessionContext.from_arrow` - Create from Arrow data + + See: :py:class:`datafusion.SessionContext` + +Expression Classes +------------------ + +**Expr** + Represents expressions that can be used in DataFrame operations. + + See: :py:class:`datafusion.Expr` + +**Functions for creating expressions:** + +* :py:func:`datafusion.column` - Reference a column by name +* :py:func:`datafusion.literal` - Create a literal value expression + +Built-in Functions +------------------ + +DataFusion provides many built-in functions for data manipulation: + +* :py:mod:`datafusion.functions` - Mathematical, string, date/time, and aggregation functions + +For a complete list of available functions, see the :py:mod:`datafusion.functions` module documentation. + + +.. toctree:: + :maxdepth: 1 + + rendering diff --git a/docs/source/user-guide/dataframe/rendering.rst b/docs/source/user-guide/dataframe/rendering.rst new file mode 100644 index 000000000..9dea948bb --- /dev/null +++ b/docs/source/user-guide/dataframe/rendering.rst @@ -0,0 +1,223 @@ +.. 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. + +HTML Rendering in Jupyter +========================= + +When working in Jupyter notebooks or other environments that support rich HTML display, +DataFusion DataFrames automatically render as nicely formatted HTML tables. This functionality +is provided by the ``_repr_html_`` method, which is automatically called by Jupyter to provide +a richer visualization than plain text output. + +Basic HTML Rendering +-------------------- + +In a Jupyter environment, simply displaying a DataFrame object will trigger HTML rendering: + +.. code-block:: python + + # Will display as HTML table in Jupyter + df + + # Explicit display also uses HTML rendering + display(df) + +Customizing HTML Rendering +--------------------------- + +DataFusion provides extensive customization options for HTML table rendering through the +``datafusion.html_formatter`` module. + +Configuring the HTML Formatter +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +You can customize how DataFrames are rendered by configuring the formatter: + +.. code-block:: python + + from datafusion.html_formatter import configure_formatter + + # Change the default styling + configure_formatter( + max_cell_length=25, # Maximum characters in a cell before truncation + max_width=1000, # Maximum width in pixels + max_height=300, # Maximum height in pixels + max_memory_bytes=2097152, # Maximum memory for rendering (2MB) + min_rows=10, # Minimum number of rows to display + max_rows=10, # Maximum rows to display in __repr__ + enable_cell_expansion=True,# Allow expanding truncated cells + custom_css=None, # Additional custom CSS + show_truncation_message=True, # Show message when data is truncated + style_provider=None, # Custom styling provider + use_shared_styles=True # Share styles across tables + ) + +The formatter settings affect all DataFrames displayed after configuration. + +Custom Style Providers +----------------------- + +For advanced styling needs, you can create a custom style provider: + +.. code-block:: python + + from datafusion.html_formatter import StyleProvider, configure_formatter + + class MyStyleProvider(StyleProvider): + def get_table_styles(self): + return { + "table": "border-collapse: collapse; width: 100%;", + "th": "background-color: #007bff; color: white; padding: 8px; text-align: left;", + "td": "border: 1px solid #ddd; padding: 8px;", + "tr:nth-child(even)": "background-color: #f2f2f2;", + } + + def get_value_styles(self, dtype, value): + """Return custom styles for specific values""" + if dtype == "float" and value < 0: + return "color: red;" + return None + + # Apply the custom style provider + configure_formatter(style_provider=MyStyleProvider()) + +Performance Optimization with Shared Styles +-------------------------------------------- + +The ``use_shared_styles`` parameter (enabled by default) optimizes performance when displaying +multiple DataFrames in notebook environments: + +.. code-block:: python + + from datafusion.html_formatter import StyleProvider, configure_formatter + # Default: Use shared styles (recommended for notebooks) + configure_formatter(use_shared_styles=True) + + # Disable shared styles (each DataFrame includes its own styles) + configure_formatter(use_shared_styles=False) + +When ``use_shared_styles=True``: +- CSS styles and JavaScript are included only once per notebook session +- This reduces HTML output size and prevents style duplication +- Improves rendering performance with many DataFrames +- Applies consistent styling across all DataFrames + +Creating a Custom Formatter +---------------------------- + +For complete control over rendering, you can implement a custom formatter: + +.. code-block:: python + + from datafusion.html_formatter import Formatter, get_formatter + + class MyFormatter(Formatter): + def format_html(self, batches, schema, has_more=False, table_uuid=None): + # Create your custom HTML here + html = "
" + # ... formatting logic ... + html += "
" + return html + + # Set as the global formatter + configure_formatter(formatter_class=MyFormatter) + + # Or use the formatter just for specific operations + formatter = get_formatter() + custom_html = formatter.format_html(batches, schema) + +Managing Formatters +------------------- + +Reset to default formatting: + +.. code-block:: python + + from datafusion.html_formatter import reset_formatter + + # Reset to default settings + reset_formatter() + +Get the current formatter settings: + +.. code-block:: python + + from datafusion.html_formatter import get_formatter + + formatter = get_formatter() + print(formatter.max_rows) + print(formatter.theme) + +Contextual Formatting +---------------------- + +You can also use a context manager to temporarily change formatting settings: + +.. code-block:: python + + from datafusion.html_formatter import formatting_context + + # Default formatting + df.show() + + # Temporarily use different formatting + with formatting_context(max_rows=100, theme="dark"): + df.show() # Will use the temporary settings + + # Back to default formatting + df.show() + +Memory and Display Controls +--------------------------- + +You can control how much data is displayed and how much memory is used for rendering: + +.. code-block:: python + + configure_formatter( + max_memory_bytes=4 * 1024 * 1024, # 4MB maximum memory for display + min_rows=20, # Always show at least 20 rows + max_rows=50 # Show up to 50 rows in output + ) + +These parameters help balance comprehensive data display against performance considerations. + +Best Practices +-------------- + +1. **Global Configuration**: Use ``configure_formatter()`` at the beginning of your notebook to set up consistent formatting for all DataFrames. + +2. **Memory Management**: Set appropriate ``max_memory_bytes`` limits to prevent performance issues with large datasets. + +3. **Shared Styles**: Keep ``use_shared_styles=True`` (default) for better performance in notebooks with multiple DataFrames. + +4. **Reset When Needed**: Call ``reset_formatter()`` when you want to start fresh with default settings. + +5. **Cell Expansion**: Use ``enable_cell_expansion=True`` when cells might contain longer content that users may want to see in full. + +Additional Resources +-------------------- + +* :doc:`../dataframe/index` - Complete guide to using DataFrames +* :doc:`../io/index` - I/O Guide for reading data from various sources +* :doc:`../data-sources` - Comprehensive data sources guide +* :ref:`io_csv` - CSV file reading +* :ref:`io_parquet` - Parquet file reading +* :ref:`io_json` - JSON file reading +* :ref:`io_avro` - Avro file reading +* :ref:`io_custom_table_provider` - Custom table providers +* `API Reference `_ - Full API reference diff --git a/docs/source/user-guide/io/arrow.rst b/docs/source/user-guide/io/arrow.rst index d571aa99c..9196fcea7 100644 --- a/docs/source/user-guide/io/arrow.rst +++ b/docs/source/user-guide/io/arrow.rst @@ -60,14 +60,16 @@ Exporting from DataFusion DataFusion DataFrames implement ``__arrow_c_stream__`` PyCapsule interface, so any Python library that accepts these can import a DataFusion DataFrame directly. -.. warning:: - It is important to note that this will cause the DataFrame execution to happen, which may be - a time consuming task. That is, you will cause a - :py:func:`datafusion.dataframe.DataFrame.collect` operation call to occur. +Invoking ``__arrow_c_stream__`` triggers execution of the underlying query, but +batches are yielded incrementally rather than materialized all at once in memory. +Consumers can process the stream as it arrives. The stream executes lazily, +letting downstream readers pull batches on demand. .. ipython:: python + from datafusion import col, lit + df = df.select((col("a") * lit(1.5)).alias("c"), lit("df").alias("d")) pa.table(df) diff --git a/docs/source/user-guide/io/csv.rst b/docs/source/user-guide/io/csv.rst index 144b6615c..9c23c291b 100644 --- a/docs/source/user-guide/io/csv.rst +++ b/docs/source/user-guide/io/csv.rst @@ -36,3 +36,25 @@ An alternative is to use :py:func:`~datafusion.context.SessionContext.register_c ctx.register_csv("file", "file.csv") df = ctx.table("file") + +If you require additional control over how to read the CSV file, you can use +:py:class:`~datafusion.options.CsvReadOptions` to set a variety of options. + +.. code-block:: python + + from datafusion import CsvReadOptions + options = ( + CsvReadOptions() + .with_has_header(True) # File contains a header row + .with_delimiter(";") # Use ; as the delimiter instead of , + .with_comment("#") # Skip lines starting with # + .with_escape("\\") # Escape character + .with_null_regex(r"^(null|NULL|N/A)$") # Treat these as NULL + .with_truncated_rows(True) # Allow rows to have incomplete columns + .with_file_compression_type("gzip") # Read gzipped CSV + .with_file_extension(".gz") # File extension other than .csv + ) + df = ctx.read_csv("data.csv.gz", options=options) + +Details for all CSV reading options can be found on the +`DataFusion documentation site `_. diff --git a/docs/source/user-guide/io/table_provider.rst b/docs/source/user-guide/io/table_provider.rst index bd1d6b80f..29e5d9880 100644 --- a/docs/source/user-guide/io/table_provider.rst +++ b/docs/source/user-guide/io/table_provider.rst @@ -37,22 +37,26 @@ A complete example can be found in the `examples folder , ) -> PyResult> { - let name = CString::new("datafusion_table_provider").unwrap(); + let name = cr"datafusion_table_provider".into(); - let provider = Arc::new(self.clone()) - .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; - let provider = FFI_TableProvider::new(Arc::new(provider), false); + let provider = Arc::new(self.clone()); + let provider = FFI_TableProvider::new(provider, false, None); PyCapsule::new_bound(py, provider, Some(name.clone())) } } -Once you have this library available, in python you can register your table provider -to the ``SessionContext``. +Once you have this library available, you can construct a +:py:class:`~datafusion.Table` in Python and register it with the +``SessionContext``. .. code-block:: python + from datafusion import SessionContext, Table + + ctx = SessionContext() provider = MyTableProvider() - ctx.register_table_provider("my_table", provider) - ctx.table("my_table").show() + ctx.register_table("capsule_table", provider) + + ctx.table("capsule_table").show() diff --git a/docs/source/user-guide/sql.rst b/docs/source/user-guide/sql.rst index 6fa7f0c6a..b4bfb9611 100644 --- a/docs/source/user-guide/sql.rst +++ b/docs/source/user-guide/sql.rst @@ -23,17 +23,100 @@ DataFusion also offers a SQL API, read the full reference `here `_, +but allow passing named parameters into a SQL query. Consider this simple +example. + +.. ipython:: python + + def show_attacks(ctx: SessionContext, threshold: int) -> None: + ctx.sql( + 'SELECT "Name", "Attack" FROM pokemon WHERE "Attack" > $val', val=threshold + ).show(num=5) + show_attacks(ctx, 75) + +When passing parameters like the example above we convert the Python objects +into their string representation. We also have special case handling +for :py:class:`~datafusion.dataframe.DataFrame` objects, since they cannot simply +be turned into string representations for an SQL query. In these cases we +will register a temporary view in the :py:class:`~datafusion.context.SessionContext` +using a generated table name. + +The formatting for passing string replacement objects is to precede the +variable name with a single ``$``. This works for all dialects in +the SQL parser except ``hive`` and ``mysql``. Since these dialects do not +support named placeholders, we are unable to do this type of replacement. +We recommend either switching to another dialect or using Python +f-string style replacement. + +.. warning:: + + To support DataFrame parameterized queries, your session must support + registration of temporary views. The default + :py:class:`~datafusion.catalog.CatalogProvider` and + :py:class:`~datafusion.catalog.SchemaProvider` do have this capability. + If you have implemented custom providers, it is important that temporary + views do not persist across :py:class:`~datafusion.context.SessionContext` + or you may get unintended consequences. + +The following example shows passing in both a :py:class:`~datafusion.dataframe.DataFrame` +object as well as a Python object to be used in parameterized replacement. + +.. ipython:: python + + def show_column( + ctx: SessionContext, column: str, df: DataFrame, threshold: int + ) -> None: + ctx.sql( + 'SELECT "Name", $col FROM $df WHERE $col > $val', + col=column, + df=df, + val=threshold, + ).show(num=5) + df = ctx.table("pokemon") + show_column(ctx, '"Defense"', df, 75) + +The approach implemented for conversion of variables into a SQL query +relies on string conversion. This has the potential for data loss, +specifically for cases like floating point numbers. If you need to pass +variables into a parameterized query and it is important to maintain the +original value without conversion to a string, then you can use the +optional parameter ``param_values`` to specify these. This parameter +expects a dictionary mapping from the parameter name to a Python +object. Those objects will be cast into a +`PyArrow Scalar Value `_. + +Using ``param_values`` will rely on the SQL dialect you have configured +for your session. This can be set using the :ref:`configuration options ` +of your :py:class:`~datafusion.context.SessionContext`. Similar to how +`prepared statements `_ +work, these parameters are limited to places where you would pass in a +scalar value, such as a comparison. + +.. ipython:: python + + def param_attacks(ctx: SessionContext, threshold: int) -> None: + ctx.sql( + 'SELECT "Name", "Attack" FROM pokemon WHERE "Attack" > $val', + param_values={"val": threshold}, + ).show(num=5) + param_attacks(ctx, 75) diff --git a/docs/source/user-guide/upgrade-guides.rst b/docs/source/user-guide/upgrade-guides.rst new file mode 100644 index 000000000..e3d7c2d87 --- /dev/null +++ b/docs/source/user-guide/upgrade-guides.rst @@ -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. + +Upgrade Guides +============== + +DataFusion 53.0.0 +----------------- + +This version includes an upgraded version of ``pyo3``, which changed the way to extract an FFI +object. Example: + +Before: + +.. code-block:: rust + + let codec = unsafe { capsule.reference::() }; + +Now: + +.. code-block:: rust + + let data: NonNull = capsule + .pointer_checked(Some(c_str!("datafusion_logical_extension_codec")))? + .cast(); + let codec = unsafe { data.as_ref() }; + +DataFusion 52.0.0 +----------------- + +This version includes a major update to the :ref:`ffi` due to upgrades +to the `Foreign Function Interface `_. +Users who contribute their own ``CatalogProvider``, ``SchemaProvider``, +``TableProvider`` or ``TableFunction`` via FFI must now provide access to a +``LogicalExtensionCodec`` and a ``TaskContextProvider``. The function signatures +for the methods to get these ``PyCapsule`` objects now requires an additional +parameter, which is a Python object that can be used to extract the +``FFI_LogicalExtensionCodec`` that is necessary. + +A complete example can be found in the `FFI example `_. +Your methods need to be updated to take an additional parameter like in this +example. + +.. code-block:: rust + + #[pymethods] + impl MyCatalogProvider { + pub fn __datafusion_catalog_provider__<'py>( + &self, + py: Python<'py>, + session: Bound, + ) -> PyResult> { + let name = cr"datafusion_catalog_provider".into(); + + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_CatalogProvider::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new(py, provider, Some(name)) + } + } + +To extract the logical extension codec FFI object from the provided object you +can implement a helper method such as: + +.. code-block:: rust + + pub(crate) fn ffi_logical_codec_from_pycapsule( + obj: Bound, + ) -> PyResult { + let attr_name = "__datafusion_logical_extension_codec__"; + let capsule = if obj.hasattr(attr_name)? { + obj.getattr(attr_name)?.call0()? + } else { + obj + }; + + let capsule = capsule.downcast::()?; + validate_pycapsule(capsule, "datafusion_logical_extension_codec")?; + + let codec = unsafe { capsule.reference::() }; + + Ok(codec.clone()) + } + + +The DataFusion FFI interface updates no longer depend directly on the +``datafusion`` core crate. You can improve your build times and potentially +reduce your library binary size by removing this dependency and instead +using the specific datafusion project crates. + +For example, instead of including expressions like: + +.. code-block:: rust + + use datafusion::catalog::MemTable; + +Instead you can now write: + +.. code-block:: rust + + use datafusion_catalog::MemTable; diff --git a/examples/csv-read-options.py b/examples/csv-read-options.py new file mode 100644 index 000000000..a5952d950 --- /dev/null +++ b/examples/csv-read-options.py @@ -0,0 +1,96 @@ +# 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. + +"""Example demonstrating CsvReadOptions usage.""" + +from datafusion import CsvReadOptions, SessionContext + +# Create a SessionContext +ctx = SessionContext() + +# Example 1: Using CsvReadOptions with default values +print("Example 1: Default CsvReadOptions") +options = CsvReadOptions() +df = ctx.read_csv("data.csv", options=options) + +# Example 2: Using CsvReadOptions with custom parameters +print("\nExample 2: Custom CsvReadOptions") +options = CsvReadOptions( + has_header=True, + delimiter=",", + quote='"', + schema_infer_max_records=1000, + file_extension=".csv", +) +df = ctx.read_csv("data.csv", options=options) + +# Example 3: Using the builder pattern (recommended for readability) +print("\nExample 3: Builder pattern") +options = ( + CsvReadOptions() + .with_has_header(True) # noqa: FBT003 + .with_delimiter("|") + .with_quote("'") + .with_schema_infer_max_records(500) + .with_truncated_rows(False) # noqa: FBT003 + .with_newlines_in_values(True) # noqa: FBT003 +) +df = ctx.read_csv("data.csv", options=options) + +# Example 4: Advanced options +print("\nExample 4: Advanced options") +options = ( + CsvReadOptions() + .with_has_header(True) # noqa: FBT003 + .with_delimiter(",") + .with_comment("#") # Skip lines starting with # + .with_escape("\\") # Escape character + .with_null_regex(r"^(null|NULL|N/A)$") # Treat these as NULL + .with_truncated_rows(True) # noqa: FBT003 + .with_file_compression_type("gzip") # Read gzipped CSV + .with_file_extension(".gz") +) +df = ctx.read_csv("data.csv.gz", options=options) + +# Example 5: Register CSV table with options +print("\nExample 5: Register CSV table") +options = CsvReadOptions().with_has_header(True).with_delimiter(",") # noqa: FBT003 +ctx.register_csv("my_table", "data.csv", options=options) +df = ctx.sql("SELECT * FROM my_table") + +# Example 6: Backward compatibility (without options) +print("\nExample 6: Backward compatibility") +# Still works the old way! +df = ctx.read_csv("data.csv", has_header=True, delimiter=",") + +print("\nAll examples completed!") +print("\nFor all available options, see the CsvReadOptions documentation:") +print(" - has_header: bool") +print(" - delimiter: str") +print(" - quote: str") +print(" - terminator: str | None") +print(" - escape: str | None") +print(" - comment: str | None") +print(" - newlines_in_values: bool") +print(" - schema: pa.Schema | None") +print(" - schema_infer_max_records: int") +print(" - file_extension: str") +print(" - table_partition_cols: list[tuple[str, pa.DataType]]") +print(" - file_compression_type: str") +print(" - file_sort_order: list[list[SortExpr]]") +print(" - null_regex: str | None") +print(" - truncated_rows: bool") diff --git a/examples/datafusion-ffi-example/Cargo.toml b/examples/datafusion-ffi-example/Cargo.toml new file mode 100644 index 000000000..178dce9f9 --- /dev/null +++ b/examples/datafusion-ffi-example/Cargo.toml @@ -0,0 +1,53 @@ +# 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. + +[package] +name = "datafusion-ffi-example" +version.workspace = true +edition.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +datafusion-catalog = { workspace = true, default-features = false } +datafusion-common = { workspace = true, default-features = false } +datafusion-functions-aggregate = { workspace = true } +datafusion-functions-window = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-ffi = { workspace = true } + +arrow = { workspace = true } +arrow-array = { workspace = true } +arrow-schema = { workspace = true } +async-trait = { workspace = true } +datafusion-python-util.workspace = true +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "datafusion_ffi_example" +crate-type = ["cdylib", "rlib"] diff --git a/examples/ffi-table-provider/build.rs b/examples/datafusion-ffi-example/build.rs similarity index 100% rename from examples/ffi-table-provider/build.rs rename to examples/datafusion-ffi-example/build.rs diff --git a/examples/ffi-table-provider/pyproject.toml b/examples/datafusion-ffi-example/pyproject.toml similarity index 83% rename from examples/ffi-table-provider/pyproject.toml rename to examples/datafusion-ffi-example/pyproject.toml index 9cd25b423..7f85e9487 100644 --- a/examples/ffi-table-provider/pyproject.toml +++ b/examples/datafusion-ffi-example/pyproject.toml @@ -20,12 +20,12 @@ requires = ["maturin>=1.6,<2.0"] build-backend = "maturin" [project] -name = "ffi_table_provider" +name = "datafusion_ffi_example" requires-python = ">=3.9" classifiers = [ - "Programming Language :: Rust", - "Programming Language :: Python :: Implementation :: CPython", - "Programming Language :: Python :: Implementation :: PyPy", + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", ] dynamic = ["version"] diff --git a/examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py b/examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py new file mode 100644 index 000000000..7ea6b295c --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_aggregate_udf.py @@ -0,0 +1,77 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import SessionContext, col, udaf +from datafusion_ffi_example import MySumUDF + + +def setup_context_with_table(): + ctx = SessionContext() + + # Pick numbers here so we get the same value in both groups + # since we cannot be certain of the output order of batches + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, 2, 3, None], type=pa.int64()), + pa.array([1, 1, 2, 2], type=pa.int64()), + ], + names=["a", "b"], + ) + ctx.register_record_batches("test_table", [[batch]]) + return ctx + + +def test_ffi_aggregate_register(): + ctx = setup_context_with_table() + my_udaf = udaf(MySumUDF()) + ctx.register_udaf(my_udaf) + + result = ctx.sql("select my_custom_sum(a) from test_table group by b").collect() + + assert len(result) == 2 + assert result[0].num_columns == 1 + + result = [r.column(0) for r in result] + expected = [ + pa.array([3], type=pa.int64()), + pa.array([3], type=pa.int64()), + ] + + assert result == expected + + +def test_ffi_aggregate_call_directly(): + ctx = setup_context_with_table() + my_udaf = udaf(MySumUDF()) + + result = ( + ctx.table("test_table").aggregate([col("b")], [my_udaf(col("a"))]).collect() + ) + + assert len(result) == 2 + assert result[0].num_columns == 2 + + result = [r.column(1) for r in result] + expected = [ + pa.array([3], type=pa.int64()), + pa.array([3], type=pa.int64()), + ] + + assert result == expected diff --git a/examples/datafusion-ffi-example/python/tests/_test_catalog_provider.py b/examples/datafusion-ffi-example/python/tests/_test_catalog_provider.py new file mode 100644 index 000000000..a862b23ba --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_catalog_provider.py @@ -0,0 +1,136 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +import pyarrow.dataset as ds +import pytest +from datafusion import SessionContext, Table +from datafusion.catalog import Schema +from datafusion_ffi_example import MyCatalogProvider, MyCatalogProviderList + + +def create_test_dataset() -> Table: + """Create a simple test dataset.""" + batch = pa.RecordBatch.from_arrays( + [pa.array([100, 200, 300]), pa.array([1.1, 2.2, 3.3])], + names=["id", "value"], + ) + dataset = ds.dataset([batch]) + return Table(dataset) + + +@pytest.mark.parametrize("inner_capsule", [True, False]) +def test_ffi_catalog_provider_list(inner_capsule: bool) -> None: + """Test basic FFI CatalogProviderList functionality.""" + ctx = SessionContext() + + # Register FFI catalog + catalog_provider_list = MyCatalogProviderList() + if inner_capsule: + catalog_provider_list = ( + catalog_provider_list.__datafusion_catalog_provider_list__(ctx) + ) + + ctx.register_catalog_provider_list(catalog_provider_list) + + # Verify the catalog exists + catalog = ctx.catalog("auto_ffi_catalog") + schema_names = catalog.names() + assert "my_schema" in schema_names + + ctx.register_catalog_provider("second", MyCatalogProvider()) + + assert ctx.catalog_names() == {"auto_ffi_catalog", "second"} + + +@pytest.mark.parametrize("inner_capsule", [True, False]) +def test_ffi_catalog_provider_basic(inner_capsule: bool) -> None: + """Test basic FFI CatalogProvider functionality.""" + ctx = SessionContext() + + # Register FFI catalog + catalog_provider = MyCatalogProvider() + if inner_capsule: + catalog_provider = catalog_provider.__datafusion_catalog_provider__(ctx) + + ctx.register_catalog_provider("ffi_catalog", catalog_provider) + + # Verify the catalog exists + catalog = ctx.catalog("ffi_catalog") + schema_names = catalog.names() + assert "my_schema" in schema_names + + # Query the pre-populated table + result = ctx.sql("SELECT * FROM ffi_catalog.my_schema.my_table").collect() + assert len(result) == 2 + assert result[0].num_columns == 2 + + +def test_ffi_catalog_provider_register_schema(): + """Test registering additional schemas to FFI CatalogProvider.""" + ctx = SessionContext() + + catalog_provider = MyCatalogProvider() + ctx.register_catalog_provider("ffi_catalog", catalog_provider) + + catalog = ctx.catalog("ffi_catalog") + + # Register a new memory schema + new_schema = Schema.memory_schema() + catalog.register_schema("additional_schema", new_schema) + + # Verify the schema was registered + assert "additional_schema" in catalog.names() + + # Add a table to the new schema + new_schema.register_table("new_table", create_test_dataset()) + + # Query the new table + result = ctx.sql("SELECT * FROM ffi_catalog.additional_schema.new_table").collect() + assert len(result) == 1 + assert result[0].column(0) == pa.array([100, 200, 300]) + + +def test_ffi_catalog_provider_deregister_schema(): + """Test deregistering schemas from FFI CatalogProvider.""" + ctx = SessionContext() + + catalog_provider = MyCatalogProvider() + ctx.register_catalog_provider("ffi_catalog", catalog_provider) + + catalog = ctx.catalog("ffi_catalog") + + # Register two schemas + schema1 = Schema.memory_schema() + schema2 = Schema.memory_schema() + catalog.register_schema("temp_schema1", schema1) + catalog.register_schema("temp_schema2", schema2) + + # Verify both exist + names = catalog.names() + assert "temp_schema1" in names + assert "temp_schema2" in names + + # Deregister one schema + catalog.deregister_schema("temp_schema1") + + # Verify it's gone + names = catalog.names() + assert "temp_schema1" not in names + assert "temp_schema2" in names diff --git a/examples/datafusion-ffi-example/python/tests/_test_scalar_udf.py b/examples/datafusion-ffi-example/python/tests/_test_scalar_udf.py new file mode 100644 index 000000000..0c949c34a --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_scalar_udf.py @@ -0,0 +1,70 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import SessionContext, col, udf +from datafusion_ffi_example import IsNullUDF + + +def setup_context_with_table(): + ctx = SessionContext() + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3, None])], + names=["a"], + ) + ctx.register_record_batches("test_table", [[batch]]) + return ctx + + +def test_ffi_scalar_register(): + ctx = setup_context_with_table() + my_udf = udf(IsNullUDF()) + ctx.register_udf(my_udf) + + result = ctx.sql("select my_custom_is_null(a) from test_table").collect() + + assert len(result) == 1 + assert result[0].num_columns == 1 + print(result) + + result = [r.column(0) for r in result] + expected = [ + pa.array([False, False, False, True], type=pa.bool_()), + ] + + assert result == expected + + +def test_ffi_scalar_call_directly(): + ctx = setup_context_with_table() + my_udf = udf(IsNullUDF()) + + result = ctx.table("test_table").select(my_udf(col("a"))).collect() + + assert len(result) == 1 + assert result[0].num_columns == 1 + print(result) + + result = [r.column(0) for r in result] + expected = [ + pa.array([False, False, False, True], type=pa.bool_()), + ] + + assert result == expected diff --git a/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py b/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py new file mode 100644 index 000000000..93449c660 --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_schema_provider.py @@ -0,0 +1,232 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +import pyarrow.dataset as ds +import pytest +from datafusion import SessionContext, Table +from datafusion.catalog import Schema +from datafusion_ffi_example import FixedSchemaProvider, MyCatalogProvider + + +def create_test_dataset() -> Table: + """Create a simple test dataset.""" + batch = pa.RecordBatch.from_arrays( + [pa.array([100, 200, 300]), pa.array([1.1, 2.2, 3.3])], + names=["id", "value"], + ) + dataset = ds.dataset([batch]) + return Table(dataset) + + +@pytest.mark.parametrize("inner_capsule", [True, False]) +def test_schema_provider_extract_values(inner_capsule: bool) -> None: + ctx = SessionContext() + + my_schema_name = "my_schema" + + schema_provider = FixedSchemaProvider() + if inner_capsule: + schema_provider = schema_provider.__datafusion_schema_provider__(ctx) + + ctx.catalog().register_schema(my_schema_name, schema_provider) + + expected_schema_name = "my_schema" + expected_table_name = "my_table" + expected_table_columns = ["units", "price"] + + default_catalog = ctx.catalog() + + catalog_schemas = default_catalog.names() + assert expected_schema_name in catalog_schemas + my_schema = default_catalog.schema(expected_schema_name) + assert expected_table_name in my_schema.names() + my_table = my_schema.table(expected_table_name) + assert expected_table_columns == my_table.schema.names + + result = ctx.table(f"{expected_schema_name}.{expected_table_name}").collect() + assert len(result) == 2 + + col0_result = [r.column(0) for r in result] + col1_result = [r.column(1) for r in result] + expected_col0 = [ + pa.array([10, 20, 30], type=pa.int32()), + pa.array([5, 7], type=pa.int32()), + ] + expected_col1 = [ + pa.array([1, 2, 5], type=pa.float64()), + pa.array([1.5, 2.5], type=pa.float64()), + ] + assert col0_result == expected_col0 + assert col1_result == expected_col1 + + +def test_ffi_schema_provider_basic(): + """Test basic FFI SchemaProvider functionality.""" + ctx = SessionContext() + + # Register FFI schema + schema_provider = FixedSchemaProvider() + ctx.catalog().register_schema("ffi_schema", schema_provider) + + # Verify the schema exists + schema = ctx.catalog().schema("ffi_schema") + table_names = schema.names() + assert "my_table" in table_names + + # Query the pre-populated table + result = ctx.sql("SELECT * FROM ffi_schema.my_table").collect() + assert len(result) == 2 + assert result[0].num_columns == 2 + + +def test_ffi_schema_provider_register_table(): + """Test registering additional tables to FFI SchemaProvider.""" + ctx = SessionContext() + + schema_provider = FixedSchemaProvider() + ctx.catalog().register_schema("ffi_schema", schema_provider) + + schema = ctx.catalog().schema("ffi_schema") + + # Register a new table + schema.register_table("additional_table", create_test_dataset()) + + # Verify the table was registered + assert "additional_table" in schema.names() + + # Query the new table + result = ctx.sql("SELECT * FROM ffi_schema.additional_table").collect() + assert len(result) == 1 + assert result[0].column(0) == pa.array([100, 200, 300]) + assert result[0].column(1) == pa.array([1.1, 2.2, 3.3]) + + +def test_ffi_schema_provider_deregister_table(): + """Test deregistering tables from FFI SchemaProvider.""" + ctx = SessionContext() + + schema_provider = FixedSchemaProvider() + ctx.catalog().register_schema("ffi_schema", schema_provider) + + schema = ctx.catalog().schema("ffi_schema") + + # Register two tables + schema.register_table("temp_table1", create_test_dataset()) + schema.register_table("temp_table2", create_test_dataset()) + + # Verify both exist + names = schema.names() + assert "temp_table1" in names + assert "temp_table2" in names + + # Deregister one table + schema.deregister_table("temp_table1") + + # Verify it's gone + names = schema.names() + assert "temp_table1" not in names + assert "temp_table2" in names + + +def test_mixed_ffi_and_python_providers(): + """Test mixing FFI and Python providers in the same catalog/schema.""" + ctx = SessionContext() + + # Register FFI catalog + ffi_catalog = MyCatalogProvider() + ctx.register_catalog_provider("ffi_catalog", ffi_catalog) + + # Register Python memory schema to FFI catalog + python_schema = Schema.memory_schema() + ctx.catalog("ffi_catalog").register_schema("python_schema", python_schema) + + # Add table to Python schema + python_schema.register_table("python_table", create_test_dataset()) + + # Query both FFI table and Python table + result_ffi = ctx.sql("SELECT * FROM ffi_catalog.my_schema.my_table").collect() + assert len(result_ffi) == 2 + + result_python = ctx.sql( + "SELECT * FROM ffi_catalog.python_schema.python_table" + ).collect() + assert len(result_python) == 1 + assert result_python[0].column(0) == pa.array([100, 200, 300]) + + +def test_ffi_catalog_with_multiple_schemas(): + """Test FFI catalog with multiple schemas of different types.""" + ctx = SessionContext() + + catalog_provider = MyCatalogProvider() + ctx.register_catalog_provider("multi_catalog", catalog_provider) + + catalog = ctx.catalog("multi_catalog") + + # Register different types of schemas + ffi_schema = FixedSchemaProvider() + memory_schema = Schema.memory_schema() + + catalog.register_schema("ffi_schema", ffi_schema) + catalog.register_schema("memory_schema", memory_schema) + + # Add tables to memory schema + memory_schema.register_table("mem_table", create_test_dataset()) + + # Verify all schemas exist + names = catalog.names() + assert "my_schema" in names # Pre-populated + assert "ffi_schema" in names + assert "memory_schema" in names + + # Query tables from each schema + result = ctx.sql("SELECT * FROM multi_catalog.my_schema.my_table").collect() + assert len(result) == 2 + + result = ctx.sql("SELECT * FROM multi_catalog.ffi_schema.my_table").collect() + assert len(result) == 2 + + result = ctx.sql("SELECT * FROM multi_catalog.memory_schema.mem_table").collect() + assert len(result) == 1 + assert result[0].column(0) == pa.array([100, 200, 300]) + + +def test_ffi_schema_table_exist(): + """Test table_exist method on FFI SchemaProvider.""" + ctx = SessionContext() + + schema_provider = FixedSchemaProvider() + ctx.catalog().register_schema("ffi_schema", schema_provider) + + schema = ctx.catalog().schema("ffi_schema") + + # Check pre-populated table + assert schema.table_exist("my_table") + + # Check non-existent table + assert not schema.table_exist("nonexistent_table") + + # Register a new table and check + schema.register_table("new_table", create_test_dataset()) + assert schema.table_exist("new_table") + + # Deregister and check + schema.deregister_table("new_table") + assert not schema.table_exist("new_table") diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_function.py b/examples/datafusion-ffi-example/python/tests/_test_table_function.py new file mode 100644 index 000000000..bf5aae3bd --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_table_function.py @@ -0,0 +1,135 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pyarrow as pa +from datafusion import Expr, SessionContext, udtf +from datafusion_ffi_example import MyTableFunction, MyTableProvider + +if TYPE_CHECKING: + from datafusion.context import TableProviderExportable + + +def test_ffi_table_function_register() -> None: + ctx = SessionContext() + table_func = MyTableFunction() + + table_udtf = udtf(table_func, "my_table_func") + ctx.register_udtf(table_udtf) + result = ctx.sql("select * from my_table_func()").collect() + + assert len(result) == 2 + assert result[0].num_columns == 4 + print(result) + + result = [r.column(0) for r in result] + expected = [ + pa.array([0, 1, 2], type=pa.int32()), + pa.array([3, 4, 5, 6], type=pa.int32()), + ] + + assert result == expected + + +def test_ffi_table_function_call_directly(): + ctx = SessionContext() + table_func = MyTableFunction() + table_udtf = udtf(table_func, "my_table_func") + + my_table = table_udtf() + ctx.register_table("t", my_table) + result = ctx.table("t").collect() + + assert len(result) == 2 + assert result[0].num_columns == 4 + print(result) + + result = [r.column(0) for r in result] + expected = [ + pa.array([0, 1, 2], type=pa.int32()), + pa.array([3, 4, 5, 6], type=pa.int32()), + ] + + assert result == expected + + +class PythonTableFunction: + """Python based table function. + + This class is used as a Python implementation of a table function. + We use the existing TableProvider to create the underlying + provider, and this function takes no arguments + """ + + def __call__( + self, num_cols: Expr, num_rows: Expr, num_batches: Expr + ) -> TableProviderExportable: + args = [ + num_cols.to_variant().value_i64(), + num_rows.to_variant().value_i64(), + num_batches.to_variant().value_i64(), + ] + return MyTableProvider(*args) + + +def common_table_function_test(test_ctx: SessionContext) -> None: + result = test_ctx.sql("select * from my_table_func(3,2,4)").collect() + + assert len(result) == 4 + assert result[0].num_columns == 3 + print(result) + + result = [r.column(0) for r in result] + expected = [ + pa.array([0, 1], type=pa.int32()), + pa.array([2, 3, 4], type=pa.int32()), + pa.array([4, 5, 6, 7], type=pa.int32()), + pa.array([6, 7, 8, 9, 10], type=pa.int32()), + ] + + assert result == expected + + +def test_python_table_function(): + ctx = SessionContext() + table_func = PythonTableFunction() + table_udtf = udtf(table_func, "my_table_func") + ctx.register_udtf(table_udtf) + + common_table_function_test(ctx) + + +def test_python_table_function_decorator(): + ctx = SessionContext() + + @udtf("my_table_func") + def my_udtf( + num_cols: Expr, num_rows: Expr, num_batches: Expr + ) -> TableProviderExportable: + args = [ + num_cols.to_variant().value_i64(), + num_rows.to_variant().value_i64(), + num_batches.to_variant().value_i64(), + ] + return MyTableProvider(*args) + + ctx.register_udtf(my_udtf) + + common_table_function_test(ctx) diff --git a/examples/ffi-table-provider/python/tests/_test_table_provider.py b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py similarity index 74% rename from examples/ffi-table-provider/python/tests/_test_table_provider.py rename to examples/datafusion-ffi-example/python/tests/_test_table_provider.py index 0db3ec561..fc77d2d3b 100644 --- a/examples/ffi-table-provider/python/tests/_test_table_provider.py +++ b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py @@ -15,15 +15,22 @@ # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import pyarrow as pa +import pytest from datafusion import SessionContext -from ffi_table_provider import MyTableProvider +from datafusion_ffi_example import MyTableProvider -def test_table_loading(): +@pytest.mark.parametrize("inner_capsule", [True, False]) +def test_table_provider_ffi(inner_capsule: bool) -> None: ctx = SessionContext() table = MyTableProvider(3, 2, 4) - ctx.register_table_provider("t", table) + if inner_capsule: + table = table.__datafusion_table_provider__(ctx) + + ctx.register_table("t", table) result = ctx.table("t").collect() assert len(result) == 4 @@ -38,3 +45,7 @@ def test_table_loading(): ] assert result == expected + + result = ctx.read_table(table).collect() + result = [r.column(0) for r in result] + assert result == expected diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_provider_factory.py b/examples/datafusion-ffi-example/python/tests/_test_table_provider_factory.py new file mode 100644 index 000000000..b1e94ec73 --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_table_provider_factory.py @@ -0,0 +1,41 @@ +# 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. + +from __future__ import annotations + +from datafusion import SessionContext +from datafusion_ffi_example import MyTableProviderFactory + + +def test_table_provider_factory_ffi() -> None: + ctx = SessionContext() + table = MyTableProviderFactory() + + ctx.register_table_factory("MY_FORMAT", table) + + # Create a new external table + ctx.sql(""" + CREATE EXTERNAL TABLE + foo + STORED AS my_format + LOCATION ''; + """).collect() + + # Query the pre-populated table + result = ctx.sql("SELECT * FROM foo;").collect() + assert len(result) == 2 + assert result[0].num_columns == 2 diff --git a/examples/datafusion-ffi-example/python/tests/_test_window_udf.py b/examples/datafusion-ffi-example/python/tests/_test_window_udf.py new file mode 100644 index 000000000..7d96994b9 --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_window_udf.py @@ -0,0 +1,89 @@ +# 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. + +from __future__ import annotations + +import pyarrow as pa +from datafusion import SessionContext, col, udwf +from datafusion_ffi_example import MyRankUDF + + +def setup_context_with_table(): + ctx = SessionContext() + + # Pick numbers here so we get the same value in both groups + # since we cannot be certain of the output order of batches + batch = pa.RecordBatch.from_arrays( + [ + pa.array([40, 10, 30, 20], type=pa.int64()), + ], + names=["a"], + ) + ctx.register_record_batches("test_table", [[batch]]) + return ctx + + +def test_ffi_window_register(): + ctx = setup_context_with_table() + my_udwf = udwf(MyRankUDF()) + ctx.register_udwf(my_udwf) + + result = ctx.sql( + "select a, my_custom_rank() over (order by a) from test_table" + ).collect() + assert len(result) == 1 + assert result[0].num_columns == 2 + + results = [ + (result[0][0][idx].as_py(), result[0][1][idx].as_py()) for idx in range(4) + ] + results.sort() + + expected = [ + (10, 1), + (20, 2), + (30, 3), + (40, 4), + ] + assert results == expected + + +def test_ffi_window_call_directly(): + ctx = setup_context_with_table() + my_udwf = udwf(MyRankUDF()) + + result = ( + ctx.table("test_table") + .select(col("a"), my_udwf().order_by(col("a")).build()) + .collect() + ) + + assert len(result) == 1 + assert result[0].num_columns == 2 + + results = [ + (result[0][0][idx].as_py(), result[0][1][idx].as_py()) for idx in range(4) + ] + results.sort() + + expected = [ + (10, 1), + (20, 2), + (30, 3), + (40, 4), + ] + assert results == expected diff --git a/examples/datafusion-ffi-example/python/tests/conftest.py b/examples/datafusion-ffi-example/python/tests/conftest.py new file mode 100644 index 000000000..68f8057af --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/conftest.py @@ -0,0 +1,42 @@ +# 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. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from collections.abc import Generator + from typing import Any + + +class _FailOnWarning(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.levelno >= logging.WARNING: + err = f"Unexpected log warning from '{record.name}': {self.format(record)}" + raise AssertionError(err) + + +@pytest.fixture(autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) diff --git a/examples/datafusion-ffi-example/src/aggregate_udf.rs b/examples/datafusion-ffi-example/src/aggregate_udf.rs new file mode 100644 index 000000000..d5343ff91 --- /dev/null +++ b/examples/datafusion-ffi-example/src/aggregate_udf.rs @@ -0,0 +1,87 @@ +// 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 std::any::Any; +use std::sync::Arc; + +use arrow_schema::DataType; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_expr::function::AccumulatorArgs; +use datafusion_expr::{Accumulator, AggregateUDF, AggregateUDFImpl, Signature}; +use datafusion_ffi::udaf::FFI_AggregateUDF; +use datafusion_functions_aggregate::sum::Sum; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; + +#[pyclass( + from_py_object, + name = "MySumUDF", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +pub(crate) struct MySumUDF { + inner: Arc, +} + +#[pymethods] +impl MySumUDF { + #[new] + fn new() -> PyResult { + Ok(Self { + inner: Arc::new(Sum::new()), + }) + } + + fn __datafusion_aggregate_udf__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let name = cr"datafusion_aggregate_udf".into(); + + let func = Arc::new(AggregateUDF::from(self.clone())); + let provider = FFI_AggregateUDF::from(func); + + PyCapsule::new(py, provider, Some(name)) + } +} + +impl AggregateUDFImpl for MySumUDF { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "my_custom_sum" + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult { + self.inner.return_type(arg_types) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> DataFusionResult> { + self.inner.accumulator(acc_args) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult> { + self.inner.coerce_types(arg_types) + } +} diff --git a/examples/datafusion-ffi-example/src/catalog_provider.rs b/examples/datafusion-ffi-example/src/catalog_provider.rs new file mode 100644 index 000000000..bd5da1e4d --- /dev/null +++ b/examples/datafusion-ffi-example/src/catalog_provider.rs @@ -0,0 +1,272 @@ +// 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 std::any::Any; +use std::fmt::Debug; +use std::sync::Arc; + +use arrow::datatypes::Schema; +use async_trait::async_trait; +use datafusion_catalog::{ + CatalogProvider, CatalogProviderList, MemTable, MemoryCatalogProvider, + MemoryCatalogProviderList, MemorySchemaProvider, SchemaProvider, TableProvider, +}; +use datafusion_common::error::{DataFusionError, Result}; +use datafusion_ffi::catalog_provider::FFI_CatalogProvider; +use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList; +use datafusion_ffi::schema_provider::FFI_SchemaProvider; +use datafusion_python_util::ffi_logical_codec_from_pycapsule; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; + +pub fn my_table() -> Arc { + use arrow::datatypes::{DataType, Field}; + use datafusion_common::record_batch; + + let schema = Arc::new(Schema::new(vec![ + Field::new("units", DataType::Int32, true), + Field::new("price", DataType::Float64, true), + ])); + + let partitions = vec![ + record_batch!( + ("units", Int32, vec![10, 20, 30]), + ("price", Float64, vec![1.0, 2.0, 5.0]) + ) + .unwrap(), + record_batch!( + ("units", Int32, vec![5, 7]), + ("price", Float64, vec![1.5, 2.5]) + ) + .unwrap(), + ]; + + Arc::new(MemTable::try_new(schema, vec![partitions]).unwrap()) +} + +#[pyclass( + skip_from_py_object, + name = "FixedSchemaProvider", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug)] +pub struct FixedSchemaProvider { + inner: Arc, +} + +impl Default for FixedSchemaProvider { + fn default() -> Self { + let inner = Arc::new(MemorySchemaProvider::new()); + + let table = my_table(); + + let _ = inner.register_table("my_table".to_string(), table).unwrap(); + + Self { inner } + } +} + +#[pymethods] +impl FixedSchemaProvider { + #[new] + pub fn new() -> Self { + Self::default() + } + + pub fn __datafusion_schema_provider__<'py>( + &self, + py: Python<'py>, + session: Bound, + ) -> PyResult> { + let name = cr"datafusion_schema_provider".into(); + + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new(py, provider, Some(name)) + } +} + +#[async_trait] +impl SchemaProvider for FixedSchemaProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn table_names(&self) -> Vec { + self.inner.table_names() + } + + async fn table(&self, name: &str) -> Result>, DataFusionError> { + self.inner.table(name).await + } + + fn register_table( + &self, + name: String, + table: Arc, + ) -> Result>> { + self.inner.register_table(name, table) + } + + fn deregister_table(&self, name: &str) -> Result>> { + self.inner.deregister_table(name) + } + + fn table_exist(&self, name: &str) -> bool { + self.inner.table_exist(name) + } +} + +/// This catalog provider is intended only for unit tests. It prepopulates with one +/// schema and only allows for schemas named after four types of fruit. +#[pyclass( + skip_from_py_object, + name = "MyCatalogProvider", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct MyCatalogProvider { + inner: Arc, +} + +impl CatalogProvider for MyCatalogProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema_names(&self) -> Vec { + self.inner.schema_names() + } + + fn schema(&self, name: &str) -> Option> { + self.inner.schema(name) + } + + fn register_schema( + &self, + name: &str, + schema: Arc, + ) -> Result>> { + self.inner.register_schema(name, schema) + } + + fn deregister_schema( + &self, + name: &str, + cascade: bool, + ) -> Result>> { + self.inner.deregister_schema(name, cascade) + } +} + +#[pymethods] +impl MyCatalogProvider { + #[new] + pub fn new() -> PyResult { + let inner = Arc::new(MemoryCatalogProvider::new()); + + let schema_name: &str = "my_schema"; + let _ = inner.register_schema(schema_name, Arc::new(FixedSchemaProvider::default())); + + Ok(Self { inner }) + } + + pub fn __datafusion_catalog_provider__<'py>( + &self, + py: Python<'py>, + session: Bound, + ) -> PyResult> { + let name = cr"datafusion_catalog_provider".into(); + + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_CatalogProvider::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new(py, provider, Some(name)) + } +} + +/// This catalog provider list is intended only for unit tests. +/// It pre-populates with a single catalog. +#[pyclass( + skip_from_py_object, + name = "MyCatalogProviderList", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct MyCatalogProviderList { + inner: Arc, +} + +impl CatalogProviderList for MyCatalogProviderList { + fn as_any(&self) -> &dyn Any { + self + } + + fn catalog_names(&self) -> Vec { + self.inner.catalog_names() + } + + fn catalog(&self, name: &str) -> Option> { + self.inner.catalog(name) + } + + fn register_catalog( + &self, + name: String, + catalog: Arc, + ) -> Option> { + self.inner.register_catalog(name, catalog) + } +} + +#[pymethods] +impl MyCatalogProviderList { + #[new] + pub fn new() -> PyResult { + let inner = Arc::new(MemoryCatalogProviderList::new()); + + inner.register_catalog( + "auto_ffi_catalog".to_owned(), + Arc::new(MyCatalogProvider::new()?), + ); + + Ok(Self { inner }) + } + + pub fn __datafusion_catalog_provider_list__<'py>( + &self, + py: Python<'py>, + session: Bound, + ) -> PyResult> { + let name = cr"datafusion_catalog_provider_list".into(); + + let provider = Arc::clone(&self.inner) as Arc; + + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_CatalogProviderList::new_with_ffi_codec(provider, None, codec); + + PyCapsule::new(py, provider, Some(name)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs new file mode 100644 index 000000000..68120a4cd --- /dev/null +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -0,0 +1,50 @@ +// 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 pyo3::prelude::*; + +use crate::aggregate_udf::MySumUDF; +use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; +use crate::scalar_udf::IsNullUDF; +use crate::table_function::MyTableFunction; +use crate::table_provider::MyTableProvider; +use crate::table_provider_factory::MyTableProviderFactory; +use crate::window_udf::MyRankUDF; + +pub(crate) mod aggregate_udf; +pub(crate) mod catalog_provider; +pub(crate) mod scalar_udf; +pub(crate) mod table_function; +pub(crate) mod table_provider; +pub(crate) mod table_provider_factory; +pub(crate) mod window_udf; + +#[pymodule] +fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/datafusion-ffi-example/src/scalar_udf.rs b/examples/datafusion-ffi-example/src/scalar_udf.rs new file mode 100644 index 000000000..374924781 --- /dev/null +++ b/examples/datafusion-ffi-example/src/scalar_udf.rs @@ -0,0 +1,97 @@ +// 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 std::any::Any; +use std::sync::Arc; + +use arrow_array::{Array, BooleanArray}; +use arrow_schema::DataType; +use datafusion_common::ScalarValue; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; +use datafusion_ffi::udf::FFI_ScalarUDF; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; + +#[pyclass( + from_py_object, + name = "IsNullUDF", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct IsNullUDF { + signature: Signature, +} + +#[pymethods] +impl IsNullUDF { + #[new] + fn new() -> Self { + Self { + signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable), + } + } + + fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) -> PyResult> { + let name = cr"datafusion_scalar_udf".into(); + + let func = Arc::new(ScalarUDF::from(self.clone())); + let provider = FFI_ScalarUDF::from(func); + + PyCapsule::new(py, provider, Some(name)) + } +} + +impl ScalarUDFImpl for IsNullUDF { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "my_custom_is_null" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult { + let input = &args.args[0]; + + Ok(match input { + ColumnarValue::Array(arr) => match arr.is_nullable() { + true => { + let nulls = arr.nulls().unwrap(); + let nulls = BooleanArray::from_iter(nulls.iter().map(|x| Some(!x))); + ColumnarValue::Array(Arc::new(nulls)) + } + false => ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))), + }, + ColumnarValue::Scalar(sv) => { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(sv == &ScalarValue::Null))) + } + }) + } +} diff --git a/examples/datafusion-ffi-example/src/table_function.rs b/examples/datafusion-ffi-example/src/table_function.rs new file mode 100644 index 000000000..79c13f64d --- /dev/null +++ b/examples/datafusion-ffi-example/src/table_function.rs @@ -0,0 +1,66 @@ +// 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 std::sync::Arc; + +use datafusion_catalog::{TableFunctionImpl, TableProvider}; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_expr::Expr; +use datafusion_ffi::udtf::FFI_TableFunction; +use datafusion_python_util::ffi_logical_codec_from_pycapsule; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; + +use crate::table_provider::MyTableProvider; + +#[pyclass( + from_py_object, + name = "MyTableFunction", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone)] +pub(crate) struct MyTableFunction {} + +#[pymethods] +impl MyTableFunction { + #[new] + fn new() -> Self { + Self {} + } + + fn __datafusion_table_function__<'py>( + &self, + py: Python<'py>, + session: Bound, + ) -> PyResult> { + let name = cr"datafusion_table_function".into(); + + let func = self.clone(); + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec); + + PyCapsule::new(py, provider, Some(name)) + } +} + +impl TableFunctionImpl for MyTableFunction { + fn call(&self, _args: &[Expr]) -> DataFusionResult> { + let provider = MyTableProvider::new(4, 3, 2).create_table()?; + Ok(Arc::new(provider)) + } +} diff --git a/examples/ffi-table-provider/src/lib.rs b/examples/datafusion-ffi-example/src/table_provider.rs similarity index 64% rename from examples/ffi-table-provider/src/lib.rs rename to examples/datafusion-ffi-example/src/table_provider.rs index 88deeece2..358ef7402 100644 --- a/examples/ffi-table-provider/src/lib.rs +++ b/examples/datafusion-ffi-example/src/table_provider.rs @@ -15,25 +15,28 @@ // specific language governing permissions and limitations // under the License. -use std::{ffi::CString, sync::Arc}; +use std::sync::Arc; -use arrow_array::ArrayRef; -use datafusion::{ - arrow::{ - array::RecordBatch, - datatypes::{DataType, Field, Schema}, - }, - datasource::MemTable, - error::{DataFusionError, Result}, -}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; +use datafusion_catalog::MemTable; +use datafusion_common::error::{DataFusionError, Result as DataFusionResult}; use datafusion_ffi::table_provider::FFI_TableProvider; -use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyCapsule}; +use datafusion_python_util::ffi_logical_codec_from_pycapsule; +use pyo3::exceptions::PyRuntimeError; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; /// In order to provide a test that demonstrates different sized record batches, /// the first batch will have num_rows, the second batch num_rows+1, and so on. -#[pyclass(name = "MyTableProvider", module = "ffi_table_provider", subclass)] +#[pyclass( + from_py_object, + name = "MyTableProvider", + module = "datafusion_ffi_example", + subclass +)] #[derive(Clone)] -struct MyTableProvider { +pub(crate) struct MyTableProvider { num_cols: usize, num_rows: usize, num_batches: usize, @@ -44,21 +47,19 @@ fn create_record_batch( num_cols: usize, start_value: i32, num_values: usize, -) -> Result { +) -> DataFusionResult { let end_value = start_value + num_values as i32; let row_values: Vec = (start_value..end_value).collect(); let columns: Vec<_> = (0..num_cols) - .map(|_| { - std::sync::Arc::new(arrow::array::Int32Array::from(row_values.clone())) as ArrayRef - }) + .map(|_| Arc::new(arrow::array::Int32Array::from(row_values.clone())) as ArrayRef) .collect(); RecordBatch::try_new(Arc::clone(schema), columns).map_err(DataFusionError::from) } impl MyTableProvider { - fn create_table(&self) -> Result { + pub fn create_table(&self) -> DataFusionResult { let fields: Vec<_> = (0..self.num_cols) .map(|idx| (b'A' + idx as u8) as char) .map(|col_name| Field::new(col_name, DataType::Int32, true)) @@ -66,7 +67,7 @@ impl MyTableProvider { let schema = Arc::new(Schema::new(fields)); - let batches: Result> = (0..self.num_batches) + let batches: DataFusionResult> = (0..self.num_batches) .map(|batch_idx| { let start_value = batch_idx * self.num_rows; create_record_batch( @@ -85,7 +86,7 @@ impl MyTableProvider { #[pymethods] impl MyTableProvider { #[new] - fn new(num_cols: usize, num_rows: usize, num_batches: usize) -> Self { + pub fn new(num_cols: usize, num_rows: usize, num_batches: usize) -> Self { Self { num_cols, num_rows, @@ -93,23 +94,21 @@ impl MyTableProvider { } } - fn __datafusion_table_provider__<'py>( + pub fn __datafusion_table_provider__<'py>( &self, py: Python<'py>, + session: Bound, ) -> PyResult> { - let name = CString::new("datafusion_table_provider").unwrap(); + let name = cr"datafusion_table_provider".into(); let provider = self .create_table() - .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; - let provider = FFI_TableProvider::new(Arc::new(provider), false, None); + .map_err(|e: DataFusionError| PyRuntimeError::new_err(e.to_string()))?; - PyCapsule::new_bound(py, provider, Some(name.clone())) - } -} + let codec = ffi_logical_codec_from_pycapsule(session)?; + let provider = + FFI_TableProvider::new_with_ffi_codec(Arc::new(provider), false, None, codec); -#[pymodule] -fn ffi_table_provider(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - Ok(()) + PyCapsule::new(py, provider, Some(name)) + } } diff --git a/examples/datafusion-ffi-example/src/table_provider_factory.rs b/examples/datafusion-ffi-example/src/table_provider_factory.rs new file mode 100644 index 000000000..53248a905 --- /dev/null +++ b/examples/datafusion-ffi-example/src/table_provider_factory.rs @@ -0,0 +1,87 @@ +// 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 std::sync::Arc; + +use async_trait::async_trait; +use datafusion_catalog::{Session, TableProvider, TableProviderFactory}; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_expr::CreateExternalTable; +use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; +use datafusion_python_util::ffi_logical_codec_from_pycapsule; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods}; + +use crate::catalog_provider; + +#[derive(Debug)] +pub(crate) struct ExampleTableProviderFactory {} + +impl ExampleTableProviderFactory { + fn new() -> Self { + Self {} + } +} + +#[async_trait] +impl TableProviderFactory for ExampleTableProviderFactory { + async fn create( + &self, + _state: &dyn Session, + _cmd: &CreateExternalTable, + ) -> DataFusionResult> { + Ok(catalog_provider::my_table()) + } +} + +#[pyclass( + name = "MyTableProviderFactory", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug)] +pub struct MyTableProviderFactory { + inner: Arc, +} + +impl Default for MyTableProviderFactory { + fn default() -> Self { + let inner = Arc::new(ExampleTableProviderFactory::new()); + Self { inner } + } +} + +#[pymethods] +impl MyTableProviderFactory { + #[new] + pub fn new() -> Self { + Self::default() + } + + pub fn __datafusion_table_provider_factory__<'py>( + &self, + py: Python<'py>, + codec: Bound, + ) -> PyResult> { + let name = cr"datafusion_table_provider_factory".into(); + let codec = ffi_logical_codec_from_pycapsule(codec)?; + let factory = Arc::clone(&self.inner) as Arc; + let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec); + + PyCapsule::new(py, factory, Some(name)) + } +} diff --git a/examples/datafusion-ffi-example/src/window_udf.rs b/examples/datafusion-ffi-example/src/window_udf.rs new file mode 100644 index 000000000..cbf179a86 --- /dev/null +++ b/examples/datafusion-ffi-example/src/window_udf.rs @@ -0,0 +1,87 @@ +// 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 std::any::Any; +use std::sync::Arc; + +use arrow_schema::{DataType, FieldRef}; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_expr::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs}; +use datafusion_expr::{PartitionEvaluator, Signature, WindowUDF, WindowUDFImpl}; +use datafusion_ffi::udwf::FFI_WindowUDF; +use datafusion_functions_window::rank::rank_udwf; +use pyo3::types::PyCapsule; +use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; + +#[pyclass( + from_py_object, + name = "MyRankUDF", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +pub(crate) struct MyRankUDF { + inner: Arc, +} + +#[pymethods] +impl MyRankUDF { + #[new] + fn new() -> PyResult { + Ok(Self { inner: rank_udwf() }) + } + + fn __datafusion_window_udf__<'py>(&self, py: Python<'py>) -> PyResult> { + let name = cr"datafusion_window_udf".into(); + + let func = Arc::new(WindowUDF::from(self.clone())); + let provider = FFI_WindowUDF::from(func); + + PyCapsule::new(py, provider, Some(name)) + } +} + +impl WindowUDFImpl for MyRankUDF { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "my_custom_rank" + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn partition_evaluator( + &self, + partition_evaluator_args: PartitionEvaluatorArgs, + ) -> DataFusionResult> { + self.inner + .inner() + .partition_evaluator(partition_evaluator_args) + } + + fn field(&self, field_args: WindowUDFFieldArgs) -> DataFusionResult { + self.inner.inner().field(field_args) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult> { + self.inner.coerce_types(arg_types) + } +} diff --git a/examples/ffi-table-provider/.cargo/config.toml b/examples/ffi-table-provider/.cargo/config.toml deleted file mode 100644 index 91a099a61..000000000 --- a/examples/ffi-table-provider/.cargo/config.toml +++ /dev/null @@ -1,12 +0,0 @@ -[target.x86_64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - -[target.aarch64-apple-darwin] -rustflags = [ - "-C", "link-arg=-undefined", - "-C", "link-arg=dynamic_lookup", -] - diff --git a/examples/ffi-table-provider/Cargo.lock b/examples/ffi-table-provider/Cargo.lock deleted file mode 100644 index 8d0edd515..000000000 --- a/examples/ffi-table-provider/Cargo.lock +++ /dev/null @@ -1,3249 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 4 - -[[package]] -name = "abi_stable" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69d6512d3eb05ffe5004c59c206de7f99c34951504056ce23fc953842f12c445" -dependencies = [ - "abi_stable_derive", - "abi_stable_shared", - "const_panic", - "core_extensions", - "crossbeam-channel", - "generational-arena", - "libloading", - "lock_api", - "parking_lot", - "paste", - "repr_offset", - "rustc_version", - "serde", - "serde_derive", - "serde_json", -] - -[[package]] -name = "abi_stable_derive" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7178468b407a4ee10e881bc7a328a65e739f0863615cca4429d43916b05e898" -dependencies = [ - "abi_stable_shared", - "as_derive_utils", - "core_extensions", - "proc-macro2", - "quote", - "rustc_version", - "syn 1.0.109", - "typed-arena", -] - -[[package]] -name = "abi_stable_shared" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2b5df7688c123e63f4d4d649cba63f2967ba7f7861b1664fca3f77d3dad2b63" -dependencies = [ - "core_extensions", -] - -[[package]] -name = "addr2line" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1" -dependencies = [ - "gimli", -] - -[[package]] -name = "adler2" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627" - -[[package]] -name = "ahash" -version = "0.8.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" -dependencies = [ - "cfg-if", - "const-random", - "getrandom", - "once_cell", - "version_check", - "zerocopy", -] - -[[package]] -name = "aho-corasick" -version = "1.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" -dependencies = [ - "memchr", -] - -[[package]] -name = "alloc-no-stdlib" -version = "2.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" - -[[package]] -name = "alloc-stdlib" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" -dependencies = [ - "alloc-no-stdlib", -] - -[[package]] -name = "allocator-api2" -version = "0.2.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45862d1c77f2228b9e10bc609d5bc203d86ebc9b87ad8d5d5167a6c9abf739d9" - -[[package]] -name = "android-tzdata" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" - -[[package]] -name = "android_system_properties" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" -dependencies = [ - "libc", -] - -[[package]] -name = "anyhow" -version = "1.0.93" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c95c10ba0b00a02636238b814946408b1322d5ac4760326e6fb8ec956d85775" - -[[package]] -name = "arrayref" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" - -[[package]] -name = "arrayvec" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" - -[[package]] -name = "arrow" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6422e12ac345a0678d7a17e316238e3a40547ae7f92052b77bd86d5e0239f3fc" -dependencies = [ - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-csv", - "arrow-data", - "arrow-ipc", - "arrow-json", - "arrow-ord", - "arrow-row", - "arrow-schema", - "arrow-select", - "arrow-string", -] - -[[package]] -name = "arrow-arith" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23cf34bb1f48c41d3475927bcc7be498665b8e80b379b88f62a840337f8b8248" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "chrono", - "num", -] - -[[package]] -name = "arrow-array" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb4a06d507f54b70a277be22a127c8ffe0cec6cd98c0ad8a48e77779bbda8223" -dependencies = [ - "ahash", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "chrono", - "chrono-tz", - "half", - "hashbrown 0.15.1", - "num", -] - -[[package]] -name = "arrow-buffer" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d69d326d5ad1cb82dcefa9ede3fee8fdca98f9982756b16f9cb142f4aa6edc89" -dependencies = [ - "bytes", - "half", - "num", -] - -[[package]] -name = "arrow-cast" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "626e65bd42636a84a238bed49d09c8777e3d825bf81f5087a70111c2831d9870" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", - "atoi", - "base64", - "chrono", - "comfy-table", - "half", - "lexical-core", - "num", - "ryu", -] - -[[package]] -name = "arrow-csv" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71c8f959f7a1389b1dbd883cdcd37c3ed12475329c111912f7f69dad8195d8c6" -dependencies = [ - "arrow-array", - "arrow-cast", - "arrow-schema", - "chrono", - "csv", - "csv-core", - "lazy_static", - "regex", -] - -[[package]] -name = "arrow-data" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1858e7c7d01c44cf71c21a85534fd1a54501e8d60d1195d0d6fbcc00f4b10754" -dependencies = [ - "arrow-buffer", - "arrow-schema", - "half", - "num", -] - -[[package]] -name = "arrow-ipc" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6bb3f727f049884c7603f0364bc9315363f356b59e9f605ea76541847e06a1e" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "flatbuffers", - "lz4_flex", -] - -[[package]] -name = "arrow-json" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35de94f165ed8830aede72c35f238763794f0d49c69d30c44d49c9834267ff8c" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-schema", - "chrono", - "half", - "indexmap", - "lexical-core", - "num", - "serde", - "serde_json", -] - -[[package]] -name = "arrow-ord" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8aa06e5f267dc53efbacb933485c79b6fc1685d3ffbe870a16ce4e696fb429da" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", -] - -[[package]] -name = "arrow-row" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66f1144bb456a2f9d82677bd3abcea019217e572fc8f07de5a7bac4b2c56eb2c" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "half", -] - -[[package]] -name = "arrow-schema" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "105f01ec0090259e9a33a9263ec18ff223ab91a0ea9fbc18042f7e38005142f6" -dependencies = [ - "bitflags 2.6.0", -] - -[[package]] -name = "arrow-select" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f690752fdbd2dee278b5f1636fefad8f2f7134c85e20fd59c4199e15a39a6807" -dependencies = [ - "ahash", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "num", -] - -[[package]] -name = "arrow-string" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0fff9cd745a7039b66c47ecaf5954460f9fa12eed628f65170117ea93e64ee0" -dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", - "memchr", - "num", - "regex", - "regex-syntax", -] - -[[package]] -name = "as_derive_utils" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3c96645900a44cf11941c111bd08a6573b0e2f9f69bc9264b179d8fae753c4" -dependencies = [ - "core_extensions", - "proc-macro2", - "quote", - "syn 1.0.109", -] - -[[package]] -name = "async-compression" -version = "0.4.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cb8f1d480b0ea3783ab015936d2a55c87e219676f0c0b7dec61494043f21857" -dependencies = [ - "bzip2 0.4.4", - "flate2", - "futures-core", - "memchr", - "pin-project-lite", - "tokio", - "xz2", - "zstd", - "zstd-safe", -] - -[[package]] -name = "async-ffi" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" -dependencies = [ - "abi_stable", -] - -[[package]] -name = "async-trait" -version = "0.1.83" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "721cae7de5c34fbb2acd27e21e6d2cf7b886dce0c27388d46c4e6c47ea4318dd" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "atoi" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" -dependencies = [ - "num-traits", -] - -[[package]] -name = "autocfg" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" - -[[package]] -name = "backtrace" -version = "0.3.74" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a" -dependencies = [ - "addr2line", - "cfg-if", - "libc", - "miniz_oxide", - "object", - "rustc-demangle", - "windows-targets", -] - -[[package]] -name = "base64" -version = "0.22.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" - -[[package]] -name = "bigdecimal" -version = "0.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f31f3af01c5c65a07985c804d3366560e6fa7883d640a122819b14ec327482c" -dependencies = [ - "autocfg", - "libm", - "num-bigint", - "num-integer", - "num-traits", -] - -[[package]] -name = "bitflags" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" - -[[package]] -name = "bitflags" -version = "2.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" - -[[package]] -name = "blake2" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" -dependencies = [ - "digest", -] - -[[package]] -name = "blake3" -version = "1.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d82033247fd8e890df8f740e407ad4d038debb9eb1f40533fffb32e7d17dc6f7" -dependencies = [ - "arrayref", - "arrayvec", - "cc", - "cfg-if", - "constant_time_eq", -] - -[[package]] -name = "block-buffer" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" -dependencies = [ - "generic-array", -] - -[[package]] -name = "brotli" -version = "7.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc97b8f16f944bba54f0433f07e30be199b6dc2bd25937444bbad560bcea29bd" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", - "brotli-decompressor", -] - -[[package]] -name = "brotli-decompressor" -version = "4.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a45bd2e4095a8b518033b128020dd4a55aab1c0a381ba4404a472630f4bc362" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", -] - -[[package]] -name = "bumpalo" -version = "3.16.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" - -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - -[[package]] -name = "bytes" -version = "1.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ac0150caa2ae65ca5bd83f25c7de183dea78d4d366469f148435e2acfbad0da" - -[[package]] -name = "bzip2" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" -dependencies = [ - "bzip2-sys", - "libc", -] - -[[package]] -name = "bzip2" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bafdbf26611df8c14810e268ddceda071c297570a5fb360ceddf617fe417ef58" -dependencies = [ - "bzip2-sys", - "libc", -] - -[[package]] -name = "bzip2-sys" -version = "0.1.11+1.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "736a955f3fa7875102d57c82b8cac37ec45224a07fd32d58f9f7a186b6cd4cdc" -dependencies = [ - "cc", - "libc", - "pkg-config", -] - -[[package]] -name = "cc" -version = "1.1.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40545c26d092346d8a8dab71ee48e7685a7a9cba76e634790c215b41a4a7b4cf" -dependencies = [ - "jobserver", - "libc", - "shlex", -] - -[[package]] -name = "cfg-if" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" - -[[package]] -name = "chrono" -version = "0.4.38" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" -dependencies = [ - "android-tzdata", - "iana-time-zone", - "num-traits", - "windows-targets", -] - -[[package]] -name = "chrono-tz" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd6dd8046d00723a59a2f8c5f295c515b9bb9a331ee4f8f3d4dd49e428acd3b6" -dependencies = [ - "chrono", - "chrono-tz-build", - "phf", -] - -[[package]] -name = "chrono-tz-build" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94fea34d77a245229e7746bd2beb786cd2a896f306ff491fb8cecb3074b10a7" -dependencies = [ - "parse-zoneinfo", - "phf_codegen", -] - -[[package]] -name = "comfy-table" -version = "7.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b34115915337defe99b2aff5c2ce6771e5fbc4079f4b506301f5cf394c8452f7" -dependencies = [ - "strum", - "strum_macros", - "unicode-width", -] - -[[package]] -name = "const-random" -version = "0.1.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" -dependencies = [ - "const-random-macro", -] - -[[package]] -name = "const-random-macro" -version = "0.1.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" -dependencies = [ - "getrandom", - "once_cell", - "tiny-keccak", -] - -[[package]] -name = "const_panic" -version = "0.2.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "013b6c2c3a14d678f38cd23994b02da3a1a1b6a5d1eedddfe63a5a5f11b13a81" - -[[package]] -name = "constant_time_eq" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" - -[[package]] -name = "core-foundation-sys" -version = "0.8.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" - -[[package]] -name = "core_extensions" -version = "1.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92c71dc07c9721607e7a16108336048ee978c3a8b129294534272e8bac96c0ee" -dependencies = [ - "core_extensions_proc_macros", -] - -[[package]] -name = "core_extensions_proc_macros" -version = "1.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69f3b219d28b6e3b4ac87bc1fc522e0803ab22e055da177bff0068c4150c61a6" - -[[package]] -name = "cpufeatures" -version = "0.2.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "608697df725056feaccfa42cffdaeeec3fccc4ffc38358ecd19b243e716a78e0" -dependencies = [ - "libc", -] - -[[package]] -name = "crc32fast" -version = "1.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "crossbeam-channel" -version = "0.5.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33480d6946193aa8033910124896ca395333cae7e2d1113d1fef6c3272217df2" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-utils" -version = "0.8.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" - -[[package]] -name = "crunchy" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" - -[[package]] -name = "crypto-common" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" -dependencies = [ - "generic-array", - "typenum", -] - -[[package]] -name = "csv" -version = "1.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acdc4883a9c96732e4733212c01447ebd805833b7275a73ca3ee080fd77afdaf" -dependencies = [ - "csv-core", - "itoa", - "ryu", - "serde", -] - -[[package]] -name = "csv-core" -version = "0.1.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" -dependencies = [ - "memchr", -] - -[[package]] -name = "dashmap" -version = "6.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" -dependencies = [ - "cfg-if", - "crossbeam-utils", - "hashbrown 0.14.5", - "lock_api", - "once_cell", - "parking_lot_core", -] - -[[package]] -name = "datafusion" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eae420e7a5b0b7f1c39364cc76cbcd0f5fdc416b2514ae3847c2676bbd60702a" -dependencies = [ - "arrow", - "arrow-array", - "arrow-ipc", - "arrow-schema", - "async-compression", - "async-trait", - "bytes", - "bzip2 0.5.0", - "chrono", - "datafusion-catalog", - "datafusion-common", - "datafusion-common-runtime", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions", - "datafusion-functions-aggregate", - "datafusion-functions-nested", - "datafusion-functions-table", - "datafusion-functions-window", - "datafusion-optimizer", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-optimizer", - "datafusion-physical-plan", - "datafusion-sql", - "flate2", - "futures", - "glob", - "itertools 0.14.0", - "log", - "object_store", - "parking_lot", - "parquet", - "rand", - "regex", - "sqlparser", - "tempfile", - "tokio", - "tokio-util", - "url", - "uuid", - "xz2", - "zstd", -] - -[[package]] -name = "datafusion-catalog" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f27987bc22b810939e8dfecc55571e9d50355d6ea8ec1c47af8383a76a6d0e1" -dependencies = [ - "arrow", - "async-trait", - "dashmap", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "datafusion-physical-plan", - "datafusion-sql", - "futures", - "itertools 0.14.0", - "log", - "parking_lot", - "sqlparser", -] - -[[package]] -name = "datafusion-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3f6d5b8c9408cc692f7c194b8aa0c0f9b253e065a8d960ad9cdc2a13e697602" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ipc", - "arrow-schema", - "base64", - "half", - "hashbrown 0.14.5", - "indexmap", - "libc", - "log", - "object_store", - "parquet", - "paste", - "recursive", - "sqlparser", - "tokio", - "web-time", -] - -[[package]] -name = "datafusion-common-runtime" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d4603c8e8a4baf77660ab7074cc66fc15cc8a18f2ce9dfadb755fc6ee294e48" -dependencies = [ - "log", - "tokio", -] - -[[package]] -name = "datafusion-doc" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5bf4bc68623a5cf231eed601ed6eb41f46a37c4d15d11a0bff24cbc8396cd66" - -[[package]] -name = "datafusion-execution" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88b491c012cdf8e051053426013429a76f74ee3c2db68496c79c323ca1084d27" -dependencies = [ - "arrow", - "dashmap", - "datafusion-common", - "datafusion-expr", - "futures", - "log", - "object_store", - "parking_lot", - "rand", - "tempfile", - "url", -] - -[[package]] -name = "datafusion-expr" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5a181408d4fc5dc22f9252781a8f39f2d0e5d1b33ec9bde242844980a2689c1" -dependencies = [ - "arrow", - "chrono", - "datafusion-common", - "datafusion-doc", - "datafusion-expr-common", - "datafusion-functions-aggregate-common", - "datafusion-functions-window-common", - "datafusion-physical-expr-common", - "indexmap", - "paste", - "recursive", - "serde_json", - "sqlparser", -] - -[[package]] -name = "datafusion-expr-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1129b48e8534d8c03c6543bcdccef0b55c8ac0c1272a15a56c67068b6eb1885" -dependencies = [ - "arrow", - "datafusion-common", - "itertools 0.14.0", - "paste", -] - -[[package]] -name = "datafusion-ffi" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff47a79d442207c168c6e3e1d970c248589c148e4800e5b285ac1b2cb1a230f8" -dependencies = [ - "abi_stable", - "arrow", - "arrow-array", - "arrow-schema", - "async-ffi", - "async-trait", - "datafusion", - "datafusion-proto", - "futures", - "log", - "prost", - "semver", - "tokio", -] - -[[package]] -name = "datafusion-functions" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6125874e4856dfb09b59886784fcb74cde5cfc5930b3a80a1a728ef7a010df6b" -dependencies = [ - "arrow", - "arrow-buffer", - "base64", - "blake2", - "blake3", - "chrono", - "datafusion-common", - "datafusion-doc", - "datafusion-execution", - "datafusion-expr", - "datafusion-expr-common", - "datafusion-macros", - "hashbrown 0.14.5", - "hex", - "itertools 0.14.0", - "log", - "md-5", - "rand", - "regex", - "sha2", - "unicode-segmentation", - "uuid", -] - -[[package]] -name = "datafusion-functions-aggregate" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3add7b1d3888e05e7c95f2b281af900ca69ebdcb21069ba679b33bde8b3b9d6" -dependencies = [ - "ahash", - "arrow", - "arrow-buffer", - "arrow-schema", - "datafusion-common", - "datafusion-doc", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-aggregate-common", - "datafusion-macros", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "half", - "log", - "paste", -] - -[[package]] -name = "datafusion-functions-aggregate-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e18baa4cfc3d2f144f74148ed68a1f92337f5072b6dde204a0dbbdf3324989c" -dependencies = [ - "ahash", - "arrow", - "datafusion-common", - "datafusion-expr-common", - "datafusion-physical-expr-common", -] - -[[package]] -name = "datafusion-functions-nested" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ec5ee8cecb0dc370291279673097ddabec03a011f73f30d7f1096457127e03e" -dependencies = [ - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ord", - "arrow-schema", - "datafusion-common", - "datafusion-doc", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions", - "datafusion-functions-aggregate", - "datafusion-macros", - "datafusion-physical-expr-common", - "itertools 0.14.0", - "log", - "paste", -] - -[[package]] -name = "datafusion-functions-table" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c403ddd473bbb0952ba880008428b3c7febf0ed3ce1eec35a205db20efb2a36" -dependencies = [ - "arrow", - "async-trait", - "datafusion-catalog", - "datafusion-common", - "datafusion-expr", - "datafusion-physical-plan", - "parking_lot", - "paste", -] - -[[package]] -name = "datafusion-functions-window" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ab18c2fb835614d06a75f24a9e09136d3a8c12a92d97c95a6af316a1787a9c5" -dependencies = [ - "datafusion-common", - "datafusion-doc", - "datafusion-expr", - "datafusion-functions-window-common", - "datafusion-macros", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "log", - "paste", -] - -[[package]] -name = "datafusion-functions-window-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a77b73bc15e7d1967121fdc7a55d819bfb9d6c03766a6c322247dce9094a53a4" -dependencies = [ - "datafusion-common", - "datafusion-physical-expr-common", -] - -[[package]] -name = "datafusion-macros" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09369b8d962291e808977cf94d495fd8b5b38647232d7ef562c27ac0f495b0af" -dependencies = [ - "datafusion-expr", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "datafusion-optimizer" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2403a7e4a84637f3de7d8d4d7a9ccc0cc4be92d89b0161ba3ee5be82f0531c54" -dependencies = [ - "arrow", - "chrono", - "datafusion-common", - "datafusion-expr", - "datafusion-physical-expr", - "indexmap", - "itertools 0.14.0", - "log", - "recursive", - "regex", - "regex-syntax", -] - -[[package]] -name = "datafusion-physical-expr" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86ff72ac702b62dbf2650c4e1d715ebd3e4aab14e3885e72e8549e250307347c" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-schema", - "datafusion-common", - "datafusion-expr", - "datafusion-expr-common", - "datafusion-functions-aggregate-common", - "datafusion-physical-expr-common", - "half", - "hashbrown 0.14.5", - "indexmap", - "itertools 0.14.0", - "log", - "paste", - "petgraph", -] - -[[package]] -name = "datafusion-physical-expr-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60982b7d684e25579ee29754b4333057ed62e2cc925383c5f0bd8cab7962f435" -dependencies = [ - "ahash", - "arrow", - "arrow-buffer", - "datafusion-common", - "datafusion-expr-common", - "hashbrown 0.14.5", - "itertools 0.14.0", -] - -[[package]] -name = "datafusion-physical-optimizer" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac5e85c189d5238a5cf181a624e450c4cd4c66ac77ca551d6f3ff9080bac90bb" -dependencies = [ - "arrow", - "arrow-schema", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "datafusion-expr-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-plan", - "futures", - "itertools 0.14.0", - "log", - "recursive", - "url", -] - -[[package]] -name = "datafusion-physical-plan" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c36bf163956d7e2542657c78b3383fdc78f791317ef358a359feffcdb968106f" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ord", - "arrow-schema", - "async-trait", - "chrono", - "datafusion-common", - "datafusion-common-runtime", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-window-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "futures", - "half", - "hashbrown 0.14.5", - "indexmap", - "itertools 0.14.0", - "log", - "parking_lot", - "pin-project-lite", - "tokio", -] - -[[package]] -name = "datafusion-proto" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2db5d79f0c974041787b899d24dc91bdab2ff112d1942dd71356a4ce3b407e6c" -dependencies = [ - "arrow", - "chrono", - "datafusion", - "datafusion-common", - "datafusion-expr", - "datafusion-proto-common", - "object_store", - "prost", -] - -[[package]] -name = "datafusion-proto-common" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de21bde1603aac0ff32cf478e47081be6e3583c6861fe8f57034da911efe7578" -dependencies = [ - "arrow", - "datafusion-common", - "prost", -] - -[[package]] -name = "datafusion-sql" -version = "45.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13caa4daede211ecec53c78b13c503b592794d125f9a3cc3afe992edf9e7f43" -dependencies = [ - "arrow", - "arrow-array", - "arrow-schema", - "bigdecimal", - "datafusion-common", - "datafusion-expr", - "indexmap", - "log", - "recursive", - "regex", - "sqlparser", -] - -[[package]] -name = "digest" -version = "0.10.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" -dependencies = [ - "block-buffer", - "crypto-common", - "subtle", -] - -[[package]] -name = "displaydoc" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "either" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" - -[[package]] -name = "equivalent" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" - -[[package]] -name = "errno" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba" -dependencies = [ - "libc", - "windows-sys 0.52.0", -] - -[[package]] -name = "fastrand" -version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "486f806e73c5707928240ddc295403b1b93c96a02038563881c4a2fd84b81ac4" - -[[package]] -name = "ffi-table-provider" -version = "0.1.0" -dependencies = [ - "arrow", - "arrow-array", - "arrow-schema", - "datafusion", - "datafusion-ffi", - "pyo3", - "pyo3-build-config", -] - -[[package]] -name = "fixedbitset" -version = "0.5.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" - -[[package]] -name = "flatbuffers" -version = "24.12.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" -dependencies = [ - "bitflags 1.3.2", - "rustc_version", -] - -[[package]] -name = "flate2" -version = "1.0.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1b589b4dc103969ad3cf85c950899926ec64300a1a46d76c03a6072957036f0" -dependencies = [ - "crc32fast", - "miniz_oxide", -] - -[[package]] -name = "form_urlencoded" -version = "1.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" -dependencies = [ - "percent-encoding", -] - -[[package]] -name = "futures" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" -dependencies = [ - "futures-channel", - "futures-core", - "futures-executor", - "futures-io", - "futures-sink", - "futures-task", - "futures-util", -] - -[[package]] -name = "futures-channel" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" -dependencies = [ - "futures-core", - "futures-sink", -] - -[[package]] -name = "futures-core" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" - -[[package]] -name = "futures-executor" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" -dependencies = [ - "futures-core", - "futures-task", - "futures-util", -] - -[[package]] -name = "futures-io" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" - -[[package]] -name = "futures-macro" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "futures-sink" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" - -[[package]] -name = "futures-task" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" - -[[package]] -name = "futures-util" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" -dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-macro", - "futures-sink", - "futures-task", - "memchr", - "pin-project-lite", - "pin-utils", - "slab", -] - -[[package]] -name = "generational-arena" -version = "0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877e94aff08e743b651baaea359664321055749b398adff8740a7399af7796e7" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "generic-array" -version = "0.14.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" -dependencies = [ - "typenum", - "version_check", -] - -[[package]] -name = "getrandom" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" -dependencies = [ - "cfg-if", - "libc", - "wasi", -] - -[[package]] -name = "gimli" -version = "0.31.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" - -[[package]] -name = "glob" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" - -[[package]] -name = "half" -version = "2.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" -dependencies = [ - "cfg-if", - "crunchy", - "num-traits", -] - -[[package]] -name = "hashbrown" -version = "0.14.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" -dependencies = [ - "ahash", - "allocator-api2", -] - -[[package]] -name = "hashbrown" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a9bfc1af68b1726ea47d3d5109de126281def866b33970e10fbab11b5dafab3" - -[[package]] -name = "heck" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" - -[[package]] -name = "hex" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" - -[[package]] -name = "humantime" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" - -[[package]] -name = "iana-time-zone" -version = "0.1.61" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" -dependencies = [ - "android_system_properties", - "core-foundation-sys", - "iana-time-zone-haiku", - "js-sys", - "wasm-bindgen", - "windows-core", -] - -[[package]] -name = "iana-time-zone-haiku" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" -dependencies = [ - "cc", -] - -[[package]] -name = "icu_collections" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" -dependencies = [ - "displaydoc", - "yoke", - "zerofrom", - "zerovec", -] - -[[package]] -name = "icu_locid" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" -dependencies = [ - "displaydoc", - "litemap", - "tinystr", - "writeable", - "zerovec", -] - -[[package]] -name = "icu_locid_transform" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_locid_transform_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_locid_transform_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" - -[[package]] -name = "icu_normalizer" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" -dependencies = [ - "displaydoc", - "icu_collections", - "icu_normalizer_data", - "icu_properties", - "icu_provider", - "smallvec", - "utf16_iter", - "utf8_iter", - "write16", - "zerovec", -] - -[[package]] -name = "icu_normalizer_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" - -[[package]] -name = "icu_properties" -version = "1.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" -dependencies = [ - "displaydoc", - "icu_collections", - "icu_locid_transform", - "icu_properties_data", - "icu_provider", - "tinystr", - "zerovec", -] - -[[package]] -name = "icu_properties_data" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" - -[[package]] -name = "icu_provider" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" -dependencies = [ - "displaydoc", - "icu_locid", - "icu_provider_macros", - "stable_deref_trait", - "tinystr", - "writeable", - "yoke", - "zerofrom", - "zerovec", -] - -[[package]] -name = "icu_provider_macros" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "idna" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" -dependencies = [ - "idna_adapter", - "smallvec", - "utf8_iter", -] - -[[package]] -name = "idna_adapter" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" -dependencies = [ - "icu_normalizer", - "icu_properties", -] - -[[package]] -name = "indexmap" -version = "2.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652" -dependencies = [ - "equivalent", - "hashbrown 0.15.1", -] - -[[package]] -name = "indoc" -version = "2.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" - -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - -[[package]] -name = "itertools" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" -dependencies = [ - "either", -] - -[[package]] -name = "itoa" -version = "1.0.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" - -[[package]] -name = "jobserver" -version = "0.1.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0" -dependencies = [ - "libc", -] - -[[package]] -name = "js-sys" -version = "0.3.72" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a88f1bda2bd75b0452a14784937d796722fdebfe50df998aeb3f0b7603019a9" -dependencies = [ - "wasm-bindgen", -] - -[[package]] -name = "lazy_static" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" - -[[package]] -name = "lexical-core" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0431c65b318a590c1de6b8fd6e72798c92291d27762d94c9e6c37ed7a73d8458" -dependencies = [ - "lexical-parse-float", - "lexical-parse-integer", - "lexical-util", - "lexical-write-float", - "lexical-write-integer", -] - -[[package]] -name = "lexical-parse-float" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb17a4bdb9b418051aa59d41d65b1c9be5affab314a872e5ad7f06231fb3b4e0" -dependencies = [ - "lexical-parse-integer", - "lexical-util", - "static_assertions", -] - -[[package]] -name = "lexical-parse-integer" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5df98f4a4ab53bf8b175b363a34c7af608fe31f93cc1fb1bf07130622ca4ef61" -dependencies = [ - "lexical-util", - "static_assertions", -] - -[[package]] -name = "lexical-util" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85314db53332e5c192b6bca611fb10c114a80d1b831ddac0af1e9be1b9232ca0" -dependencies = [ - "static_assertions", -] - -[[package]] -name = "lexical-write-float" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e7c3ad4e37db81c1cbe7cf34610340adc09c322871972f74877a712abc6c809" -dependencies = [ - "lexical-util", - "lexical-write-integer", - "static_assertions", -] - -[[package]] -name = "lexical-write-integer" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb89e9f6958b83258afa3deed90b5de9ef68eef090ad5086c791cd2345610162" -dependencies = [ - "lexical-util", - "static_assertions", -] - -[[package]] -name = "libc" -version = "0.2.162" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18d287de67fe55fd7e1581fe933d965a5a9477b38e949cfa9f8574ef01506398" - -[[package]] -name = "libloading" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" -dependencies = [ - "cfg-if", - "winapi", -] - -[[package]] -name = "libm" -version = "0.2.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa" - -[[package]] -name = "linux-raw-sys" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" - -[[package]] -name = "litemap" -version = "0.7.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "643cb0b8d4fcc284004d5fd0d67ccf61dfffadb7f75e1e71bc420f4688a3a704" - -[[package]] -name = "lock_api" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" -dependencies = [ - "autocfg", - "scopeguard", -] - -[[package]] -name = "log" -version = "0.4.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" - -[[package]] -name = "lz4_flex" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75761162ae2b0e580d7e7c390558127e5f01b4194debd6221fd8c207fc80e3f5" -dependencies = [ - "twox-hash", -] - -[[package]] -name = "lzma-sys" -version = "0.1.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" -dependencies = [ - "cc", - "libc", - "pkg-config", -] - -[[package]] -name = "md-5" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" -dependencies = [ - "cfg-if", - "digest", -] - -[[package]] -name = "memchr" -version = "2.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" - -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - -[[package]] -name = "miniz_oxide" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2d80299ef12ff69b16a84bb182e3b9df68b5a91574d3d4fa6e41b65deec4df1" -dependencies = [ - "adler2", -] - -[[package]] -name = "num" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" -dependencies = [ - "num-bigint", - "num-complex", - "num-integer", - "num-iter", - "num-rational", - "num-traits", -] - -[[package]] -name = "num-bigint" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" -dependencies = [ - "num-integer", - "num-traits", -] - -[[package]] -name = "num-complex" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-integer" -version = "0.1.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-iter" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" -dependencies = [ - "autocfg", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-traits" -version = "0.2.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" -dependencies = [ - "autocfg", - "libm", -] - -[[package]] -name = "object" -version = "0.36.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e" -dependencies = [ - "memchr", -] - -[[package]] -name = "object_store" -version = "0.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eb4c22c6154a1e759d7099f9ffad7cc5ef8245f9efbab4a41b92623079c82f3" -dependencies = [ - "async-trait", - "bytes", - "chrono", - "futures", - "humantime", - "itertools 0.13.0", - "parking_lot", - "percent-encoding", - "snafu", - "tokio", - "tracing", - "url", - "walkdir", -] - -[[package]] -name = "once_cell" -version = "1.20.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" - -[[package]] -name = "ordered-float" -version = "2.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] - -[[package]] -name = "parking_lot" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" -dependencies = [ - "lock_api", - "parking_lot_core", -] - -[[package]] -name = "parking_lot_core" -version = "0.9.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" -dependencies = [ - "cfg-if", - "libc", - "redox_syscall", - "smallvec", - "windows-targets", -] - -[[package]] -name = "parquet" -version = "54.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a01a0efa30bbd601ae85b375c728efdb211ade54390281628a7b16708beb235" -dependencies = [ - "ahash", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-ipc", - "arrow-schema", - "arrow-select", - "base64", - "brotli", - "bytes", - "chrono", - "flate2", - "futures", - "half", - "hashbrown 0.15.1", - "lz4_flex", - "num", - "num-bigint", - "object_store", - "paste", - "seq-macro", - "simdutf8", - "snap", - "thrift", - "tokio", - "twox-hash", - "zstd", - "zstd-sys", -] - -[[package]] -name = "parse-zoneinfo" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f2a05b18d44e2957b88f96ba460715e295bc1d7510468a2f3d3b44535d26c24" -dependencies = [ - "regex", -] - -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - -[[package]] -name = "percent-encoding" -version = "2.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" - -[[package]] -name = "petgraph" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" -dependencies = [ - "fixedbitset", - "indexmap", -] - -[[package]] -name = "phf" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_codegen" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8d39688d359e6b34654d328e262234662d16cc0f60ec8dcbe5e718709342a5a" -dependencies = [ - "phf_generator", - "phf_shared", -] - -[[package]] -name = "phf_generator" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48e4cc64c2ad9ebe670cb8fd69dd50ae301650392e81c05f9bfcb2d5bdbc24b0" -dependencies = [ - "phf_shared", - "rand", -] - -[[package]] -name = "phf_shared" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" -dependencies = [ - "siphasher", -] - -[[package]] -name = "pin-project-lite" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff" - -[[package]] -name = "pin-utils" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" - -[[package]] -name = "pkg-config" -version = "0.3.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" - -[[package]] -name = "portable-atomic" -version = "1.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" - -[[package]] -name = "ppv-lite86" -version = "0.2.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" -dependencies = [ - "zerocopy", -] - -[[package]] -name = "proc-macro2" -version = "1.0.89" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "prost" -version = "0.13.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b0487d90e047de87f984913713b85c601c05609aad5b0df4b4573fbf69aa13f" -dependencies = [ - "bytes", - "prost-derive", -] - -[[package]] -name = "prost-derive" -version = "0.13.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9552f850d5f0964a4e4d0bf306459ac29323ddfbae05e35a7c0d35cb0803cc5" -dependencies = [ - "anyhow", - "itertools 0.13.0", - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "psm" -version = "0.1.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "200b9ff220857e53e184257720a14553b2f4aa02577d2ed9842d45d4b9654810" -dependencies = [ - "cc", -] - -[[package]] -name = "pyo3" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57fe09249128b3173d092de9523eaa75136bf7ba85e0d69eca241c7939c933cc" -dependencies = [ - "cfg-if", - "indoc", - "libc", - "memoffset", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", - "unindent", -] - -[[package]] -name = "pyo3-build-config" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cd3927b5a78757a0d71aa9dff669f903b1eb64b54142a9bd9f757f8fde65fd7" -dependencies = [ - "once_cell", - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dab6bb2102bd8f991e7749f130a70d05dd557613e39ed2deeee8e9ca0c4d548d" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-macros" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91871864b353fd5ffcb3f91f2f703a22a9797c91b9ab497b1acac7b07ae509c7" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43abc3b80bc20f3facd86cd3c60beed58c3e2aa26213f3cda368de39c60a27e4" -dependencies = [ - "heck", - "proc-macro2", - "pyo3-build-config", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "quote" -version = "1.0.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "rand" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" -dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", -] - -[[package]] -name = "rand_core" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] - -[[package]] -name = "recursive" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0786a43debb760f491b1bc0269fe5e84155353c67482b9e60d0cfb596054b43e" -dependencies = [ - "recursive-proc-macro-impl", - "stacker", -] - -[[package]] -name = "recursive-proc-macro-impl" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" -dependencies = [ - "quote", - "syn 2.0.87", -] - -[[package]] -name = "redox_syscall" -version = "0.5.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" -dependencies = [ - "bitflags 2.6.0", -] - -[[package]] -name = "regex" -version = "1.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" -dependencies = [ - "aho-corasick", - "memchr", - "regex-automata", - "regex-syntax", -] - -[[package]] -name = "regex-automata" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" -dependencies = [ - "aho-corasick", - "memchr", - "regex-syntax", -] - -[[package]] -name = "regex-syntax" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" - -[[package]] -name = "repr_offset" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb1070755bd29dffc19d0971cab794e607839ba2ef4b69a9e6fbc8733c1b72ea" -dependencies = [ - "tstr", -] - -[[package]] -name = "rustc-demangle" -version = "0.1.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" - -[[package]] -name = "rustc_version" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" -dependencies = [ - "semver", -] - -[[package]] -name = "rustix" -version = "0.38.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99e4ea3e1cdc4b559b8e5650f9c8e5998e3e5c1343b4eaf034565f32318d63c0" -dependencies = [ - "bitflags 2.6.0", - "errno", - "libc", - "linux-raw-sys", - "windows-sys 0.52.0", -] - -[[package]] -name = "rustversion" -version = "1.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e819f2bc632f285be6d7cd36e25940d45b2391dd6d9b939e79de557f7014248" - -[[package]] -name = "ryu" -version = "1.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" - -[[package]] -name = "same-file" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" -dependencies = [ - "winapi-util", -] - -[[package]] -name = "scopeguard" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" - -[[package]] -name = "semver" -version = "1.0.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f79dfe2d285b0488816f30e700a7438c5a73d816b5b7d3ac72fbc48b0d185e03" - -[[package]] -name = "seq-macro" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" - -[[package]] -name = "serde" -version = "1.0.214" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f55c3193aca71c12ad7890f1785d2b73e1b9f63a0bbc353c08ef26fe03fc56b5" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.214" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de523f781f095e28fa605cdce0f8307e451cc0fd14e2eb4cd2e98a355b147766" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "serde_json" -version = "1.0.132" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03" -dependencies = [ - "itoa", - "memchr", - "ryu", - "serde", -] - -[[package]] -name = "sha2" -version = "0.10.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" -dependencies = [ - "cfg-if", - "cpufeatures", - "digest", -] - -[[package]] -name = "shlex" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" - -[[package]] -name = "simdutf8" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" - -[[package]] -name = "siphasher" -version = "0.3.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" - -[[package]] -name = "slab" -version = "0.4.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" -dependencies = [ - "autocfg", -] - -[[package]] -name = "smallvec" -version = "1.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" - -[[package]] -name = "snafu" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "223891c85e2a29c3fe8fb900c1fae5e69c2e42415e3177752e8718475efa5019" -dependencies = [ - "snafu-derive", -] - -[[package]] -name = "snafu-derive" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03c3c6b7927ffe7ecaa769ee0e3994da3b8cafc8f444578982c83ecb161af917" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "snap" -version = "1.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" - -[[package]] -name = "sqlparser" -version = "0.53.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05a528114c392209b3264855ad491fcce534b94a38771b0a0b97a79379275ce8" -dependencies = [ - "log", - "sqlparser_derive", -] - -[[package]] -name = "sqlparser_derive" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da5fc6819faabb412da764b99d3b713bb55083c11e7e0c00144d386cd6a1939c" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "stable_deref_trait" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" - -[[package]] -name = "stacker" -version = "0.1.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "799c883d55abdb5e98af1a7b3f23b9b6de8ecada0ecac058672d7635eb48ca7b" -dependencies = [ - "cc", - "cfg-if", - "libc", - "psm", - "windows-sys 0.59.0", -] - -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - -[[package]] -name = "strum" -version = "0.26.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" - -[[package]] -name = "strum_macros" -version = "0.26.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.87", -] - -[[package]] -name = "subtle" -version = "2.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" - -[[package]] -name = "syn" -version = "1.0.109" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "syn" -version = "2.0.87" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "synstructure" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "target-lexicon" -version = "0.12.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" - -[[package]] -name = "tempfile" -version = "3.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28cce251fcbc87fac86a866eeb0d6c2d536fc16d06f184bb61aeae11aa4cee0c" -dependencies = [ - "cfg-if", - "fastrand", - "once_cell", - "rustix", - "windows-sys 0.59.0", -] - -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float", -] - -[[package]] -name = "tiny-keccak" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" -dependencies = [ - "crunchy", -] - -[[package]] -name = "tinystr" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" -dependencies = [ - "displaydoc", - "zerovec", -] - -[[package]] -name = "tokio" -version = "1.41.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22cfb5bee7a6a52939ca9224d6ac897bb669134078daa8735560897f69de4d33" -dependencies = [ - "backtrace", - "bytes", - "pin-project-lite", - "tokio-macros", -] - -[[package]] -name = "tokio-macros" -version = "2.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "tokio-util" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61e7c3654c13bcd040d4a03abee2c75b1d14a37b423cf5a813ceae1cc903ec6a" -dependencies = [ - "bytes", - "futures-core", - "futures-sink", - "pin-project-lite", - "tokio", -] - -[[package]] -name = "tracing" -version = "0.1.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" -dependencies = [ - "pin-project-lite", - "tracing-attributes", - "tracing-core", -] - -[[package]] -name = "tracing-attributes" -version = "0.1.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "tracing-core" -version = "0.1.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" -dependencies = [ - "once_cell", -] - -[[package]] -name = "tstr" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f8e0294f14baae476d0dd0a2d780b2e24d66e349a9de876f5126777a37bdba7" -dependencies = [ - "tstr_proc_macros", -] - -[[package]] -name = "tstr_proc_macros" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e78122066b0cb818b8afd08f7ed22f7fdbc3e90815035726f0840d0d26c0747a" - -[[package]] -name = "twox-hash" -version = "1.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" -dependencies = [ - "cfg-if", - "static_assertions", -] - -[[package]] -name = "typed-arena" -version = "2.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" - -[[package]] -name = "typenum" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" - -[[package]] -name = "unicode-ident" -version = "1.0.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" - -[[package]] -name = "unicode-segmentation" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" - -[[package]] -name = "unicode-width" -version = "0.1.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" - -[[package]] -name = "unindent" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" - -[[package]] -name = "url" -version = "2.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" -dependencies = [ - "form_urlencoded", - "idna", - "percent-encoding", -] - -[[package]] -name = "utf16_iter" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" - -[[package]] -name = "utf8_iter" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" - -[[package]] -name = "uuid" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c5f0a0af699448548ad1a2fbf920fb4bee257eae39953ba95cb84891a0446a" -dependencies = [ - "getrandom", -] - -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - -[[package]] -name = "walkdir" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" -dependencies = [ - "same-file", - "winapi-util", -] - -[[package]] -name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" - -[[package]] -name = "wasm-bindgen" -version = "0.2.95" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "128d1e363af62632b8eb57219c8fd7877144af57558fb2ef0368d0087bddeb2e" -dependencies = [ - "cfg-if", - "once_cell", - "wasm-bindgen-macro", -] - -[[package]] -name = "wasm-bindgen-backend" -version = "0.2.95" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb6dd4d3ca0ddffd1dd1c9c04f94b868c37ff5fac97c30b97cff2d74fce3a358" -dependencies = [ - "bumpalo", - "log", - "once_cell", - "proc-macro2", - "quote", - "syn 2.0.87", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-macro" -version = "0.2.95" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e79384be7f8f5a9dd5d7167216f022090cf1f9ec128e6e6a482a2cb5c5422c56" -dependencies = [ - "quote", - "wasm-bindgen-macro-support", -] - -[[package]] -name = "wasm-bindgen-macro-support" -version = "0.2.95" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26c6ab57572f7a24a4985830b120de1594465e5d500f24afe89e16b4e833ef68" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", - "wasm-bindgen-backend", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-shared" -version = "0.2.95" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65fc09f10666a9f147042251e0dda9c18f166ff7de300607007e96bdebc1068d" - -[[package]] -name = "web-time" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" -dependencies = [ - "js-sys", - "wasm-bindgen", -] - -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-util" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" -dependencies = [ - "windows-sys 0.59.0", -] - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - -[[package]] -name = "windows-core" -version = "0.52.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" -dependencies = [ - "windows-targets", -] - -[[package]] -name = "windows-sys" -version = "0.52.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" -dependencies = [ - "windows-targets", -] - -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets", -] - -[[package]] -name = "windows-targets" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" -dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", -] - -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" - -[[package]] -name = "windows_aarch64_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" - -[[package]] -name = "windows_i686_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" - -[[package]] -name = "windows_i686_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" - -[[package]] -name = "windows_i686_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" - -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.52.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" - -[[package]] -name = "write16" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" - -[[package]] -name = "writeable" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" - -[[package]] -name = "xz2" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" -dependencies = [ - "lzma-sys", -] - -[[package]] -name = "yoke" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c5b1314b079b0930c31e3af543d8ee1757b1951ae1e1565ec704403a7240ca5" -dependencies = [ - "serde", - "stable_deref_trait", - "yoke-derive", - "zerofrom", -] - -[[package]] -name = "yoke-derive" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28cc31741b18cb6f1d5ff12f5b7523e3d6eb0852bbbad19d73905511d9849b95" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", - "synstructure", -] - -[[package]] -name = "zerocopy" -version = "0.7.35" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" -dependencies = [ - "byteorder", - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.7.35" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "zerofrom" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91ec111ce797d0e0784a1116d0ddcdbea84322cd79e5d5ad173daeba4f93ab55" -dependencies = [ - "zerofrom-derive", -] - -[[package]] -name = "zerofrom-derive" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ea7b4a3637ea8669cedf0f1fd5c286a17f3de97b8dd5a70a6c167a1730e63a5" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", - "synstructure", -] - -[[package]] -name = "zerovec" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" -dependencies = [ - "yoke", - "zerofrom", - "zerovec-derive", -] - -[[package]] -name = "zerovec-derive" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.87", -] - -[[package]] -name = "zstd" -version = "0.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fcf2b778a664581e31e389454a7072dab1647606d44f7feea22cd5abb9c9f3f9" -dependencies = [ - "zstd-safe", -] - -[[package]] -name = "zstd-safe" -version = "7.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54a3ab4db68cea366acc5c897c7b4d4d1b8994a9cd6e6f841f8964566a419059" -dependencies = [ - "zstd-sys", -] - -[[package]] -name = "zstd-sys" -version = "2.0.13+zstd.1.5.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38ff0f21cfee8f97d94cef41359e0c89aa6113028ab0291aa8ca0038995a95aa" -dependencies = [ - "cc", - "pkg-config", -] diff --git a/examples/python-udf-comparisons.py b/examples/python-udf-comparisons.py index eb0825011..b870645a3 100644 --- a/examples/python-udf-comparisons.py +++ b/examples/python-udf-comparisons.py @@ -15,16 +15,16 @@ # specific language governing permissions and limitations # under the License. -import os import time +from pathlib import Path import pyarrow as pa import pyarrow.compute as pc from datafusion import SessionContext, col, lit, udf from datafusion import functions as F -path = os.path.dirname(os.path.abspath(__file__)) -filepath = os.path.join(path, "./tpch/data/lineitem.parquet") +path = Path(__file__).parent.resolve() +filepath = path / "./tpch/data/lineitem.parquet" # This example serves to demonstrate alternate approaches to answering the # question "return all of the rows that have a specific combination of these diff --git a/examples/python-udwf.py b/examples/python-udwf.py index 98d118bf2..645ded188 100644 --- a/examples/python-udwf.py +++ b/examples/python-udwf.py @@ -22,7 +22,7 @@ from datafusion import col, lit, udwf from datafusion import functions as f from datafusion.expr import WindowFrame -from datafusion.udf import WindowEvaluator +from datafusion.user_defined import WindowEvaluator # This example creates five different examples of user defined window functions in order # to demonstrate the variety of ways a user may need to implement. diff --git a/examples/sql-using-python-udaf.py b/examples/sql-using-python-udaf.py index 32ce38900..f42bbdc23 100644 --- a/examples/sql-using-python-udaf.py +++ b/examples/sql-using-python-udaf.py @@ -28,16 +28,16 @@ class MyAccumulator(Accumulator): def __init__(self) -> None: self._sum = pa.scalar(0.0) - def update(self, values: pa.Array) -> None: + def update(self, values: list[pa.Array]) -> None: # not nice since pyarrow scalars can't be summed yet. This breaks on `None` self._sum = pa.scalar(self._sum.as_py() + pa.compute.sum(values).as_py()) def merge(self, states: pa.Array) -> None: # not nice since pyarrow scalars can't be summed yet. This breaks on `None` - self._sum = pa.scalar(self._sum.as_py() + pa.compute.sum(states).as_py()) + self._sum = pa.scalar(self._sum.as_py() + pa.compute.sum(states[0]).as_py()) - def state(self) -> pa.Array: - return pa.array([self._sum.as_py()]) + def state(self) -> list[pa.Array]: + return [self._sum] def evaluate(self) -> pa.Scalar: return self._sum diff --git a/examples/tpch/_tests.py b/examples/tpch/_tests.py index 80ff80244..780fcf5e5 100644 --- a/examples/tpch/_tests.py +++ b/examples/tpch/_tests.py @@ -25,8 +25,10 @@ def df_selection(col_name, col_type): - if col_type == pa.float64() or isinstance(col_type, pa.Decimal128Type): + if col_type == pa.float64(): return F.round(col(col_name), lit(2)).alias(col_name) + if isinstance(col_type, pa.Decimal128Type): + return F.round(col(col_name).cast(pa.float64()), lit(2)).alias(col_name) if col_type == pa.string() or col_type == pa.string_view(): return F.trim(col(col_name)).alias(col_name) return col(col_name) diff --git a/examples/tpch/convert_data_to_parquet.py b/examples/tpch/convert_data_to_parquet.py index fd0fcca49..af554c39e 100644 --- a/examples/tpch/convert_data_to_parquet.py +++ b/examples/tpch/convert_data_to_parquet.py @@ -22,7 +22,7 @@ as will be generated by the script provided in this repository. """ -import os +from pathlib import Path import datafusion import pyarrow as pa @@ -116,7 +116,7 @@ ("S_COMMENT", pa.string()), ] -curr_dir = os.path.dirname(os.path.abspath(__file__)) +curr_dir = Path(__file__).resolve().parent for filename, curr_schema_val in all_schemas.items(): # For convenience, go ahead and convert the schema column names to lowercase curr_schema = [(s[0].lower(), s[1]) for s in curr_schema_val] @@ -132,10 +132,8 @@ schema = pa.schema(curr_schema) - source_file = os.path.abspath( - os.path.join(curr_dir, f"../../benchmarks/tpch/data/{filename}.csv") - ) - dest_file = os.path.abspath(os.path.join(curr_dir, f"./data/{filename}.parquet")) + source_file = (curr_dir / f"../../benchmarks/tpch/data/{filename}.csv").resolve() + dest_file = (curr_dir / f"./data/{filename}.parquet").resolve() df = ctx.read_csv(source_file, schema=schema, has_header=False, delimiter="|") diff --git a/examples/tpch/q07_volume_shipping.py b/examples/tpch/q07_volume_shipping.py index a84cf728a..ff2f891f1 100644 --- a/examples/tpch/q07_volume_shipping.py +++ b/examples/tpch/q07_volume_shipping.py @@ -80,7 +80,7 @@ # not match these will result in a null value and then get filtered out. # # To do the same using a simple filter would be: -# df_nation = df_nation.filter((F.col("n_name") == nation_1) | (F.col("n_name") == nation_2)) +# df_nation = df_nation.filter((F.col("n_name") == nation_1) | (F.col("n_name") == nation_2)) # noqa: ERA001 df_nation = df_nation.with_column( "n_name", F.case(col("n_name")) diff --git a/examples/tpch/q12_ship_mode_order_priority.py b/examples/tpch/q12_ship_mode_order_priority.py index f1d894940..9071597f0 100644 --- a/examples/tpch/q12_ship_mode_order_priority.py +++ b/examples/tpch/q12_ship_mode_order_priority.py @@ -73,7 +73,7 @@ # matches either of the two values, but we want to show doing some array operations in this # example. If you want to see this done with filters, comment out the above line and uncomment # this one. -# df = df.filter((col("l_shipmode") == lit(SHIP_MODE_1)) | (col("l_shipmode") == lit(SHIP_MODE_2))) +# df = df.filter((col("l_shipmode") == lit(SHIP_MODE_1)) | (col("l_shipmode") == lit(SHIP_MODE_2))) # noqa: ERA001 # We need order priority, so join order df to line item diff --git a/examples/tpch/util.py b/examples/tpch/util.py index 7e3d659dd..ec53bcd15 100644 --- a/examples/tpch/util.py +++ b/examples/tpch/util.py @@ -19,18 +19,16 @@ Common utilities for running TPC-H examples. """ -import os +from pathlib import Path -def get_data_path(filename: str) -> str: - path = os.path.dirname(os.path.abspath(__file__)) +def get_data_path(filename: str) -> Path: + path = Path(__file__).resolve().parent - return os.path.join(path, "data", filename) + return path / "data" / filename -def get_answer_file(answer_file: str) -> str: - path = os.path.dirname(os.path.abspath(__file__)) +def get_answer_file(answer_file: str) -> Path: + path = Path(__file__).resolve().parent - return os.path.join( - path, "../../benchmarks/tpch/data/answers", f"{answer_file}.out" - ) + return path / "../../benchmarks/tpch/data/answers" / f"{answer_file}.out" diff --git a/pyproject.toml b/pyproject.toml index d86b657ec..117aeefc2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,26 +24,30 @@ name = "datafusion" description = "Build and run queries against data" readme = "README.md" license = { file = "LICENSE.txt" } -requires-python = ">=3.9" -keywords = ["datafusion", "dataframe", "rust", "query-engine"] +requires-python = ">=3.10" +keywords = ["dataframe", "datafusion", "query-engine", "rust"] classifiers = [ - "Development Status :: 2 - Pre-Alpha", - "Intended Audience :: Developers", - "License :: OSI Approved :: Apache Software License", - "License :: OSI Approved", - "Operating System :: MacOS", - "Operating System :: Microsoft :: Windows", - "Operating System :: POSIX :: Linux", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python", - "Programming Language :: Rust", + "Development Status :: 2 - Pre-Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "License :: OSI Approved", + "Operating System :: MacOS", + "Operating System :: Microsoft :: Windows", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python", + "Programming Language :: Rust", +] +dependencies = [ + "pyarrow>=16.0.0;python_version<'3.14'", + "pyarrow>=22.0.0;python_version>='3.14'", + "typing-extensions;python_version<'3.13'", ] -dependencies = ["pyarrow>=11.0.0", "typing-extensions;python_version<'3.13'"] dynamic = ["version"] [project.urls] @@ -55,46 +59,41 @@ repository = "https://github.com/apache/datafusion-python" profile = "black" [tool.maturin] +manifest-path = "crates/core/Cargo.toml" python-source = "python" module-name = "datafusion._internal" include = [{ path = "Cargo.lock", format = "sdist" }] -exclude = [".github/**", "ci/**", ".asf.yaml"] +exclude = [".asf.yaml", ".github/**", "ci/**"] # Require Cargo.lock is up to date locked = true features = ["substrait"] +[tool.pytest.ini_options] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +addopts = "--doctest-modules" +doctest_optionflags = ["NORMALIZE_WHITESPACE", "ELLIPSIS"] +testpaths = ["python/tests", "python/datafusion"] + # Enable docstring linting using the google style guide [tool.ruff.lint] -select = ["ALL" ] +select = ["ALL"] ignore = [ - "A001", # Allow using words like min as variable names - "A002", # Allow using words like filter as variable names - "ANN401", # Allow Any for wrapper classes - "COM812", # Recommended to ignore these rules when using with ruff-format - "FIX002", # Allow TODO lines - consider removing at some point - "FBT001", # Allow boolean positional args - "FBT002", # Allow boolean positional args - "ISC001", # Recommended to ignore these rules when using with ruff-format - "SLF001", # Allow accessing private members - "TD002", - "TD003", # Allow TODO lines - "UP007", # Disallowing Union is pedantic - # TODO: Enable all of the following, but this PR is getting too large already - "PLR0913", - "TRY003", - "PLR2004", - "PD901", - "ERA001", - "ANN001", - "ANN202", - "PTH", - "N812", - "INP001", - "DTZ007", - "RUF015", - "A005", - "TC001", - "UP035", + "A001", # Allow using words like min as variable names + "A002", # Allow using words like filter as variable names + "A005", # Allow module named io + "ANN401", # Allow Any for wrapper classes + "COM812", # Recommended to ignore these rules when using with ruff-format + "FBT001", # Allow boolean positional args + "FBT002", # Allow boolean positional args + "FIX002", # Allow TODO lines - consider removing at some point + "ISC001", # Recommended to ignore these rules when using with ruff-format + "N812", # Allow importing functions as `F` + "PD901", # Allow variable name df + "PLR0913", # Allow many arguments in function definition + "SLF001", # Allow accessing private members + "TD002", # Do not require author names in TODO statements + "TD003", # Allow TODO lines ] [tool.ruff.lint.pydocstyle] @@ -103,46 +102,105 @@ convention = "google" [tool.ruff.lint.pycodestyle] max-doc-length = 88 +[tool.ruff.lint.flake8-boolean-trap] +extend-allowed-calls = ["datafusion.lit", "lit"] + # Disable docstring checking for these directories [tool.ruff.lint.per-file-ignores] "python/tests/*" = [ - "ANN", - "ARG", - "BLE001", - "D", - "S101", - "SLF", - "PD", - "PLR2004", - "PT011", - "RUF015", - "S608", - "PLR0913", - "PT004", + "ANN", + "ARG", + "BLE001", + "D", + "PD", + "PLC0415", + "PLR0913", + "PLR2004", + "PT004", + "PT011", + "RUF015", + "S101", + "S608", + "SLF", +] +"examples/*" = [ + "ANN001", + "ANN202", + "D", + "DTZ007", + "E501", + "INP001", + "PLR2004", + "RUF015", + "S101", + "T201", + "W505", +] +"dev/*" = [ + "ANN001", + "C", + "D", + "E", + "ERA001", + "EXE", + "N817", + "PLR", + "S", + "SIM", + "T", + "UP", +] +"benchmarks/*" = [ + "ANN001", + "BLE", + "D", + "E", + "ERA001", + "EXE", + "F", + "FURB", + "INP001", + "PLR", + "S", + "SIM", + "T", + "TD", + "TRY", + "UP", ] -"examples/*" = ["D", "W505", "E501", "T201", "S101"] -"dev/*" = ["D", "E", "T", "S", "PLR", "C", "SIM", "UP", "EXE", "N817"] -"benchmarks/*" = ["D", "F", "T", "BLE", "FURB", "PLR", "E", "TD", "TRY", "S", "SIM", "EXE", "UP"] "docs/*" = ["D"] +"docs/source/conf.py" = ["ANN001", "ERA001", "INP001"] + +[tool.codespell] +skip = ["./python/tests/test_functions.py", "./target", "uv.lock"] +count = true +ignore-words-list = ["IST", "ans"] [dependency-groups] dev = [ - "maturin>=1.8.1", - "numpy>1.25.0", - "pytest>=7.4.4", - "pytest-asyncio>=0.23.3", - "ruff>=0.9.1", - "toml>=0.10.2", - "pygithub==2.5.0", + "arro3-core==0.6.5", + "codespell==2.4.1", + "maturin>=1.8.1", + "nanoarrow==0.8.0", + "numpy>1.25.0;python_version<'3.14'", + "numpy>=2.3.2;python_version>='3.14'", + "pre-commit>=4.3.0", + "pyarrow>=19.0.0", + "pygithub==2.5.0", + "pytest-asyncio>=0.23.3", + "pytest>=7.4.4", + "pyyaml>=6.0.3", + "ruff>=0.9.1", + "toml>=0.10.2", ] docs = [ - "sphinx>=7.1.2", - "pydata-sphinx-theme==0.8.0", - "myst-parser>=3.0.1", - "jinja2>=3.1.5", - "ipython>=8.12.3", - "pandas>=2.0.3", - "pickleshare>=0.7.5", - "sphinx-autoapi>=3.4.0", - "setuptools>=75.3.0", + "ipython>=8.12.3", + "jinja2>=3.1.5", + "myst-parser>=3.0.1", + "pandas>=2.0.3", + "pickleshare>=0.7.5", + "pydata-sphinx-theme==0.8.0", + "setuptools>=75.3.0", + "sphinx-autoapi>=3.4.0", + "sphinx>=7.1.2", ] diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index d871fdb71..2e6f81166 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -21,34 +21,53 @@ See https://datafusion.apache.org/python for more information. """ +from __future__ import annotations + +from typing import Any + try: import importlib.metadata as importlib_metadata except ImportError: - import importlib_metadata + import importlib_metadata # type: ignore[import] -from . import functions, object_store, substrait +# Public submodules +from . import functions, object_store, substrait, unparser # The following imports are okay to remain as opaque to the user. from ._internal import Config from .catalog import Catalog, Database, Table -from .common import ( - DFSchema, -) +from .col import col, column +from .common import DFSchema from .context import ( RuntimeEnvBuilder, SessionConfig, SessionContext, SQLOptions, ) -from .dataframe import DataFrame -from .expr import ( - Expr, - WindowFrame, +from .dataframe import ( + DataFrame, + DataFrameWriteOptions, + InsertOp, + ParquetColumnOptions, + ParquetWriterOptions, ) +from .dataframe_formatter import configure_formatter +from .expr import Expr, WindowFrame from .io import read_avro, read_csv, read_json, read_parquet +from .options import CsvReadOptions from .plan import ExecutionPlan, LogicalPlan from .record_batch import RecordBatch, RecordBatchStream -from .udf import Accumulator, AggregateUDF, ScalarUDF, WindowUDF, udaf, udf, udwf +from .user_defined import ( + Accumulator, + AggregateUDF, + ScalarUDF, + TableFunction, + WindowUDF, + udaf, + udf, + udtf, + udwf, +) __version__ = importlib_metadata.version(__name__) @@ -57,12 +76,17 @@ "AggregateUDF", "Catalog", "Config", + "CsvReadOptions", "DFSchema", "DataFrame", + "DataFrameWriteOptions", "Database", "ExecutionPlan", "Expr", + "InsertOp", "LogicalPlan", + "ParquetColumnOptions", + "ParquetWriterOptions", "RecordBatch", "RecordBatchStream", "RuntimeEnvBuilder", @@ -71,16 +95,20 @@ "SessionConfig", "SessionContext", "Table", + "TableFunction", "WindowFrame", "WindowUDF", + "catalog", "col", "column", "common", + "configure_formatter", "expr", "functions", "lit", "literal", "object_store", + "options", "read_avro", "read_csv", "read_json", @@ -88,26 +116,18 @@ "substrait", "udaf", "udf", + "udtf", "udwf", + "unparser", ] -def column(value: str) -> Expr: - """Create a column expression.""" - return Expr.column(value) - - -def col(value: str) -> Expr: - """Create a column expression.""" - return Expr.column(value) - - -def literal(value) -> Expr: +def literal(value: Any) -> Expr: """Create a literal expression.""" return Expr.literal(value) -def string_literal(value): +def string_literal(value: str) -> Expr: """Create a UTF8 literal expression. It differs from `literal` which creates a UTF8view literal. @@ -115,11 +135,26 @@ def string_literal(value): return Expr.string_literal(value) -def str_lit(value): +def str_lit(value: str) -> Expr: """Alias for `string_literal`.""" return string_literal(value) -def lit(value) -> Expr: +def lit(value: Any) -> Expr: """Create a literal expression.""" return Expr.literal(value) + + +def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr: + """Creates a new expression representing a scalar value with metadata. + + Args: + value: A valid PyArrow scalar value or easily castable to one. + metadata: Metadata to attach to the expression. + """ + return Expr.literal_with_metadata(value, metadata) + + +def lit_with_metadata(value: Any, metadata: dict[str, str]) -> Expr: + """Alias for literal_with_metadata.""" + return literal_with_metadata(value, metadata) diff --git a/python/datafusion/catalog.py b/python/datafusion/catalog.py index 6c3f188cc..03c0ddc68 100644 --- a/python/datafusion/catalog.py +++ b/python/datafusion/catalog.py @@ -19,59 +19,372 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Protocol import datafusion._internal as df_internal if TYPE_CHECKING: import pyarrow as pa + from datafusion import DataFrame, SessionContext + from datafusion.context import TableProviderExportable + from datafusion.expr import CreateExternalTable + +try: + from warnings import deprecated # Python 3.13+ +except ImportError: + from typing_extensions import deprecated # Python 3.12 + + +__all__ = [ + "Catalog", + "CatalogList", + "CatalogProvider", + "CatalogProviderList", + "Schema", + "SchemaProvider", + "Table", +] + + +class CatalogList: + """DataFusion data catalog list.""" + + def __init__(self, catalog_list: df_internal.catalog.RawCatalogList) -> None: + """This constructor is not typically called by the end user.""" + self.catalog_list = catalog_list + + def __repr__(self) -> str: + """Print a string representation of the catalog list.""" + return self.catalog_list.__repr__() + + def names(self) -> set[str]: + """This is an alias for `catalog_names`.""" + return self.catalog_names() + + def catalog_names(self) -> set[str]: + """Returns the list of schemas in this catalog.""" + return self.catalog_list.catalog_names() + + @staticmethod + def memory_catalog(ctx: SessionContext | None = None) -> CatalogList: + """Create an in-memory catalog provider list.""" + catalog_list = df_internal.catalog.RawCatalogList.memory_catalog(ctx) + return CatalogList(catalog_list) + + def catalog(self, name: str = "datafusion") -> Catalog: + """Returns the catalog with the given ``name`` from this catalog.""" + catalog = self.catalog_list.catalog(name) + + return ( + Catalog(catalog) + if isinstance(catalog, df_internal.catalog.RawCatalog) + else catalog + ) + + def register_catalog( + self, + name: str, + catalog: Catalog | CatalogProvider | CatalogProviderExportable, + ) -> Catalog | None: + """Register a catalog with this catalog list.""" + if isinstance(catalog, Catalog): + return self.catalog_list.register_catalog(name, catalog.catalog) + return self.catalog_list.register_catalog(name, catalog) + class Catalog: """DataFusion data catalog.""" - def __init__(self, catalog: df_internal.Catalog) -> None: + def __init__(self, catalog: df_internal.catalog.RawCatalog) -> None: """This constructor is not typically called by the end user.""" self.catalog = catalog - def names(self) -> list[str]: - """Returns the list of databases in this catalog.""" - return self.catalog.names() + def __repr__(self) -> str: + """Print a string representation of the catalog.""" + return self.catalog.__repr__() + + def names(self) -> set[str]: + """This is an alias for `schema_names`.""" + return self.schema_names() + + def schema_names(self) -> set[str]: + """Returns the list of schemas in this catalog.""" + return self.catalog.schema_names() + + @staticmethod + def memory_catalog(ctx: SessionContext | None = None) -> Catalog: + """Create an in-memory catalog provider.""" + catalog = df_internal.catalog.RawCatalog.memory_catalog(ctx) + return Catalog(catalog) + + def schema(self, name: str = "public") -> Schema: + """Returns the database with the given ``name`` from this catalog.""" + schema = self.catalog.schema(name) + + return ( + Schema(schema) + if isinstance(schema, df_internal.catalog.RawSchema) + else schema + ) - def database(self, name: str = "public") -> Database: + @deprecated("Use `schema` instead.") + def database(self, name: str = "public") -> Schema: """Returns the database with the given ``name`` from this catalog.""" - return Database(self.catalog.database(name)) + return self.schema(name) + + def register_schema( + self, + name: str, + schema: Schema | SchemaProvider | SchemaProviderExportable, + ) -> Schema | None: + """Register a schema with this catalog.""" + if isinstance(schema, Schema): + return self.catalog.register_schema(name, schema._raw_schema) + return self.catalog.register_schema(name, schema) + def deregister_schema(self, name: str, cascade: bool = True) -> Schema | None: + """Deregister a schema from this catalog.""" + return self.catalog.deregister_schema(name, cascade) -class Database: - """DataFusion Database.""" - def __init__(self, db: df_internal.Database) -> None: +class Schema: + """DataFusion Schema.""" + + def __init__(self, schema: df_internal.catalog.RawSchema) -> None: """This constructor is not typically called by the end user.""" - self.db = db + self._raw_schema = schema + + def __repr__(self) -> str: + """Print a string representation of the schema.""" + return self._raw_schema.__repr__() + + @staticmethod + def memory_schema(ctx: SessionContext | None = None) -> Schema: + """Create an in-memory schema provider.""" + schema = df_internal.catalog.RawSchema.memory_schema(ctx) + return Schema(schema) def names(self) -> set[str]: - """Returns the list of all tables in this database.""" - return self.db.names() + """This is an alias for `table_names`.""" + return self.table_names() + + def table_names(self) -> set[str]: + """Returns the list of all tables in this schema.""" + return self._raw_schema.table_names def table(self, name: str) -> Table: - """Return the table with the given ``name`` from this database.""" - return Table(self.db.table(name)) + """Return the table with the given ``name`` from this schema.""" + return Table(self._raw_schema.table(name)) + + def register_table( + self, + name: str, + table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset, + ) -> None: + """Register a table in this schema.""" + return self._raw_schema.register_table(name, table) + + def deregister_table(self, name: str) -> None: + """Deregister a table provider from this schema.""" + return self._raw_schema.deregister_table(name) + + def table_exist(self, name: str) -> bool: + """Determines if a table exists in this schema.""" + return self._raw_schema.table_exist(name) + + +@deprecated("Use `Schema` instead.") +class Database(Schema): + """See `Schema`.""" class Table: - """DataFusion table.""" + """A DataFusion table. - def __init__(self, table: df_internal.Table) -> None: - """This constructor is not typically called by the end user.""" - self.table = table + Internally we currently support the following types of tables: + + - Tables created using built-in DataFusion methods, such as + reading from CSV or Parquet + - pyarrow datasets + - DataFusion DataFrames, which will be converted into a view + - Externally provided tables implemented with the FFI PyCapsule + interface (advanced) + """ + + __slots__ = ("_inner",) + + def __init__( + self, + table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset, + ctx: SessionContext | None = None, + ) -> None: + """Constructor.""" + self._inner = df_internal.catalog.RawTable(table, ctx) + + def __repr__(self) -> str: + """Print a string representation of the table.""" + return repr(self._inner) + + @staticmethod + @deprecated("Use Table() constructor instead.") + def from_dataset(dataset: pa.dataset.Dataset) -> Table: + """Turn a :mod:`pyarrow.dataset` ``Dataset`` into a :class:`Table`.""" + return Table(dataset) @property def schema(self) -> pa.Schema: """Returns the schema associated with this table.""" - return self.table.schema + return self._inner.schema @property def kind(self) -> str: """Returns the kind of table.""" - return self.table.kind + return self._inner.kind + + +class TableProviderFactory(ABC): + """Abstract class for defining a Python based Table Provider Factory.""" + + @abstractmethod + def create(self, cmd: CreateExternalTable) -> Table: + """Create a table using the :class:`CreateExternalTable`.""" + ... + + +class TableProviderFactoryExportable(Protocol): + """Type hint for object that has __datafusion_table_provider_factory__ PyCapsule. + + https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProviderFactory.html + """ + + def __datafusion_table_provider_factory__(self, session: Any) -> object: ... + + +class CatalogProviderList(ABC): + """Abstract class for defining a Python based Catalog Provider List.""" + + @abstractmethod + def catalog_names(self) -> set[str]: + """Set of the names of all catalogs in this catalog list.""" + ... + + @abstractmethod + def catalog( + self, name: str + ) -> CatalogProviderExportable | CatalogProvider | Catalog | None: + """Retrieve a specific catalog from this catalog list.""" + ... + + def register_catalog( # noqa: B027 + self, name: str, catalog: CatalogProviderExportable | CatalogProvider | Catalog + ) -> None: + """Add a catalog to this catalog list. + + This method is optional. If your catalog provides a fixed list of catalogs, you + do not need to implement this method. + """ + + +class CatalogProviderListExportable(Protocol): + """Type hint for object that has __datafusion_catalog_provider_list__ PyCapsule. + + https://docs.rs/datafusion/latest/datafusion/catalog/trait.CatalogProviderList.html + """ + + def __datafusion_catalog_provider_list__(self, session: Any) -> object: ... + + +class CatalogProvider(ABC): + """Abstract class for defining a Python based Catalog Provider.""" + + @abstractmethod + def schema_names(self) -> set[str]: + """Set of the names of all schemas in this catalog.""" + ... + + @abstractmethod + def schema(self, name: str) -> Schema | None: + """Retrieve a specific schema from this catalog.""" + ... + + def register_schema( # noqa: B027 + self, name: str, schema: SchemaProviderExportable | SchemaProvider | Schema + ) -> None: + """Add a schema to this catalog. + + This method is optional. If your catalog provides a fixed list of schemas, you + do not need to implement this method. + """ + + def deregister_schema(self, name: str, cascade: bool) -> None: # noqa: B027 + """Remove a schema from this catalog. + + This method is optional. If your catalog provides a fixed list of schemas, you + do not need to implement this method. + + Args: + name: The name of the schema to remove. + cascade: If true, deregister the tables within the schema. + """ + + +class CatalogProviderExportable(Protocol): + """Type hint for object that has __datafusion_catalog_provider__ PyCapsule. + + https://docs.rs/datafusion/latest/datafusion/catalog/trait.CatalogProvider.html + """ + + def __datafusion_catalog_provider__(self, session: Any) -> object: ... + + +class SchemaProvider(ABC): + """Abstract class for defining a Python based Schema Provider.""" + + def owner_name(self) -> str | None: + """Returns the owner of the schema. + + This is an optional method. The default return is None. + """ + return None + + @abstractmethod + def table_names(self) -> set[str]: + """Set of the names of all tables in this schema.""" + ... + + @abstractmethod + def table(self, name: str) -> Table | None: + """Retrieve a specific table from this schema.""" + ... + + def register_table( # noqa: B027 + self, name: str, table: Table | TableProviderExportable | Any + ) -> None: + """Add a table to this schema. + + This method is optional. If your schema provides a fixed list of tables, you do + not need to implement this method. + """ + + def deregister_table(self, name: str, cascade: bool) -> None: # noqa: B027 + """Remove a table from this schema. + + This method is optional. If your schema provides a fixed list of tables, you do + not need to implement this method. + """ + + @abstractmethod + def table_exist(self, name: str) -> bool: + """Returns true if the table exists in this schema.""" + ... + + +class SchemaProviderExportable(Protocol): + """Type hint for object that has __datafusion_schema_provider__ PyCapsule. + + https://docs.rs/datafusion/latest/datafusion/catalog/trait.SchemaProvider.html + """ + + def __datafusion_schema_provider__(self, session: Any) -> object: ... diff --git a/python/datafusion/col.py b/python/datafusion/col.py new file mode 100644 index 000000000..1141dc092 --- /dev/null +++ b/python/datafusion/col.py @@ -0,0 +1,45 @@ +# 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. + +"""Col class.""" + +from datafusion.expr import Expr + + +class Col: + """Create a column expression. + + This helper class allows an extra syntax of creating columns using the __getattr__ + method. + """ + + def __call__(self, value: str) -> Expr: + """Create a column expression.""" + return Expr.column(value) + + def __getattr__(self, value: str) -> Expr: + """Create a column using attribute syntax.""" + # For autocomplete to work with IPython + if value.startswith("__wrapped__"): + return getattr(type(self), value) + + return Expr.column(value) + + +col: Col = Col() +column: Col = Col() +__all__ = ["col", "column"] diff --git a/python/datafusion/common.py b/python/datafusion/common.py index e762a993b..c689a816d 100644 --- a/python/datafusion/common.py +++ b/python/datafusion/common.py @@ -33,8 +33,12 @@ SqlTable = common_internal.SqlTable SqlType = common_internal.SqlType SqlView = common_internal.SqlView +TableType = common_internal.TableType +TableSource = common_internal.TableSource +Constraints = common_internal.Constraints __all__ = [ + "Constraints", "DFSchema", "DataType", "DataTypeMap", @@ -47,6 +51,8 @@ "SqlTable", "SqlType", "SqlView", + "TableSource", + "TableType", ] diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 1429a4975..ba9290a58 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -19,6 +19,8 @@ from __future__ import annotations +import uuid +import warnings from typing import TYPE_CHECKING, Any, Protocol try: @@ -26,25 +28,49 @@ except ImportError: from typing_extensions import deprecated # Python 3.12 -from datafusion.catalog import Catalog, Table + +import pyarrow as pa + +from datafusion.catalog import ( + Catalog, + CatalogList, + CatalogProviderExportable, + CatalogProviderList, + CatalogProviderListExportable, + TableProviderFactory, + TableProviderFactoryExportable, +) from datafusion.dataframe import DataFrame -from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list +from datafusion.expr import sort_list_to_raw_sort_list +from datafusion.options import ( + DEFAULT_MAX_INFER_SCHEMA, + CsvReadOptions, + _convert_table_partition_cols, +) from datafusion.record_batch import RecordBatchStream -from datafusion.udf import AggregateUDF, ScalarUDF, WindowUDF from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal from ._internal import SessionConfig as SessionConfigInternal from ._internal import SessionContext as SessionContextInternal from ._internal import SQLOptions as SQLOptionsInternal +from ._internal import expr as expr_internal if TYPE_CHECKING: import pathlib + from collections.abc import Sequence import pandas as pd - import polars as pl - import pyarrow as pa + import polars as pl # type: ignore[import] + from datafusion.catalog import CatalogProvider, Table + from datafusion.expr import SortKey from datafusion.plan import ExecutionPlan, LogicalPlan + from datafusion.user_defined import ( + AggregateUDF, + ScalarUDF, + TableFunction, + WindowUDF, + ) class ArrowStreamExportable(Protocol): @@ -75,7 +101,7 @@ class TableProviderExportable(Protocol): https://datafusion.apache.org/python/user-guide/io/table_provider.html """ - def __datafusion_table_provider__(self) -> object: ... # noqa: D105 + def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 class SessionConfig: @@ -496,6 +522,10 @@ def __init__( self.ctx = SessionContextInternal(config, runtime) + def __repr__(self) -> str: + """Print a string representation of the Session Context.""" + return self.ctx.__repr__() + @classmethod def global_ctx(cls) -> SessionContext: """Retrieve the global context as a `SessionContext` wrapper. @@ -535,10 +565,10 @@ def register_listing_table( self, name: str, path: str | pathlib.Path, - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_extension: str = ".parquet", schema: pa.Schema | None = None, - file_sort_order: list[list[Expr | SortExpr]] | None = None, + file_sort_order: Sequence[Sequence[SortKey]] | None = None, ) -> None: """Register multiple files as a single table. @@ -552,27 +582,35 @@ def register_listing_table( table_partition_cols: Partition columns. file_extension: File extension of the provided table. schema: The data source schema. - file_sort_order: Sort order for the file. + file_sort_order: Sort order for the file. Each sort key can be + specified as a column name (``str``), an expression + (``Expr``), or a ``SortExpr``. """ if table_partition_cols is None: table_partition_cols = [] - file_sort_order_raw = ( - [sort_list_to_raw_sort_list(f) for f in file_sort_order] - if file_sort_order is not None - else None - ) + table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_listing_table( name, str(path), table_partition_cols, file_extension, schema, - file_sort_order_raw, + self._convert_file_sort_order(file_sort_order), ) - def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame: + def sql( + self, + query: str, + options: SQLOptions | None = None, + param_values: dict[str, Any] | None = None, + **named_params: Any, + ) -> DataFrame: """Create a :py:class:`~datafusion.DataFrame` from SQL query text. + See the online documentation for a description of how to perform + parameterized substitution via either the ``param_values`` option + or passing in ``named_params``. + Note: This API implements DDL statements such as ``CREATE TABLE`` and ``CREATE VIEW`` and DML statements such as ``INSERT INTO`` with in-memory default implementation.See @@ -581,15 +619,57 @@ def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame: Args: query: SQL query text. options: If provided, the query will be validated against these options. + param_values: Provides substitution of scalar values in the query + after parsing. + named_params: Provides string or DataFrame substitution in the query string. Returns: DataFrame representation of the SQL query. """ - if options is None: - return DataFrame(self.ctx.sql(query)) - return DataFrame(self.ctx.sql_with_options(query, options.options_internal)) - def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame: + def value_to_scalar(value: Any) -> pa.Scalar: + if isinstance(value, pa.Scalar): + return value + return pa.scalar(value) + + def value_to_string(value: Any) -> str: + if isinstance(value, DataFrame): + view_name = str(uuid.uuid4()).replace("-", "_") + view_name = f"view_{view_name}" + view = value.df.into_view(temporary=True) + self.ctx.register_table(view_name, view) + return view_name + return str(value) + + param_values = ( + {name: value_to_scalar(value) for (name, value) in param_values.items()} + if param_values is not None + else {} + ) + param_strings = ( + {name: value_to_string(value) for (name, value) in named_params.items()} + if named_params is not None + else {} + ) + + options_raw = options.options_internal if options is not None else None + + return DataFrame( + self.ctx.sql_with_options( + query, + options=options_raw, + param_values=param_values, + param_strings=param_strings, + ) + ) + + def sql_with_options( + self, + query: str, + options: SQLOptions, + param_values: dict[str, Any] | None = None, + **named_params: Any, + ) -> DataFrame: """Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text. This function will first validate that the query is allowed by the @@ -598,11 +678,16 @@ def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame: Args: query: SQL query text. options: SQL options. + param_values: Provides substitution of scalar values in the query + after parsing. + named_params: Provides string or DataFrame substitution in the query string. Returns: DataFrame representation of the SQL query. """ - return self.sql(query, options) + return self.sql( + query, options=options, param_values=param_values, **named_params + ) def create_dataframe( self, @@ -718,7 +803,7 @@ def from_polars(self, data: pl.DataFrame, name: str | None = None) -> DataFrame: # https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116 # is the discussion on how we arrived at adding register_view def register_view(self, name: str, df: DataFrame) -> None: - """Register a :py:class: `~datafusion.detaframe.DataFrame` as a view. + """Register a :py:class:`~datafusion.dataframe.DataFrame` as a view. Args: name (str): The name to register the view under. @@ -727,30 +812,80 @@ def register_view(self, name: str, df: DataFrame) -> None: view = df.into_view() self.ctx.register_table(name, view) - def register_table(self, name: str, table: Table) -> None: - """Register a :py:class: `~datafusion.catalog.Table` as a table. + def register_table( + self, + name: str, + table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset, + ) -> None: + """Register a :py:class:`~datafusion.Table` with this context. - The registered table can be referenced from SQL statement executed against. + The registered table can be referenced from SQL statements executed against + this context. Args: name: Name of the resultant table. - table: DataFusion table to add to the session context. + table: Any object that can be converted into a :class:`Table`. """ - self.ctx.register_table(name, table.table) + self.ctx.register_table(name, table) def deregister_table(self, name: str) -> None: """Remove a table from the session.""" self.ctx.deregister_table(name) + def register_table_factory( + self, + format: str, + factory: TableProviderFactory | TableProviderFactoryExportable, + ) -> None: + """Register a :py:class:`~datafusion.TableProviderFactoryExportable`. + + The registered factory can be referenced from SQL DDL statements executed + against this context. + + Args: + format: The value to be used in `STORED AS ${format}` clause. + factory: A PyCapsule that implements :class:`TableProviderFactoryExportable` + """ + self.ctx.register_table_factory(format, factory) + + def catalog_names(self) -> set[str]: + """Returns the list of catalogs in this context.""" + return self.ctx.catalog_names() + + def register_catalog_provider_list( + self, + provider: CatalogProviderListExportable | CatalogProviderList | CatalogList, + ) -> None: + """Register a catalog provider list.""" + if isinstance(provider, CatalogList): + self.ctx.register_catalog_provider_list(provider.catalog) + else: + self.ctx.register_catalog_provider_list(provider) + + def register_catalog_provider( + self, name: str, provider: CatalogProviderExportable | CatalogProvider | Catalog + ) -> None: + """Register a catalog provider.""" + if isinstance(provider, Catalog): + self.ctx.register_catalog_provider(name, provider.catalog) + else: + self.ctx.register_catalog_provider(name, provider) + + @deprecated("Use register_table() instead.") def register_table_provider( - self, name: str, provider: TableProviderExportable + self, + name: str, + provider: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset, ) -> None: """Register a table provider. - This table provider must have a method called ``__datafusion_table_provider__`` - which returns a PyCapsule that exposes a ``FFI_TableProvider``. + Deprecated: use :meth:`register_table` instead. """ - self.ctx.register_table_provider(name, provider) + self.register_table(name, provider) + + def register_udtf(self, func: TableFunction) -> None: + """Register a user defined table function.""" + self.ctx.register_udtf(func._udtf) def register_record_batches( self, name: str, partitions: list[list[pa.RecordBatch]] @@ -770,12 +905,12 @@ def register_parquet( self, name: str, path: str | pathlib.Path, - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, parquet_pruning: bool = True, file_extension: str = ".parquet", skip_metadata: bool = True, schema: pa.Schema | None = None, - file_sort_order: list[list[SortExpr]] | None = None, + file_sort_order: Sequence[Sequence[SortKey]] | None = None, ) -> None: """Register a Parquet file as a table. @@ -794,10 +929,13 @@ def register_parquet( that may be in the file schema. This can help avoid schema conflicts due to metadata. schema: The data source schema. - file_sort_order: Sort order for the file. + file_sort_order: Sort order for the file. Each sort key can be + specified as a column name (``str``), an expression + (``Expr``), or a ``SortExpr``. """ if table_partition_cols is None: table_partition_cols = [] + table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_parquet( name, str(path), @@ -806,9 +944,7 @@ def register_parquet( file_extension, skip_metadata, schema, - [sort_list_to_raw_sort_list(exprs) for exprs in file_sort_order] - if file_sort_order is not None - else None, + self._convert_file_sort_order(file_sort_order), ) def register_csv( @@ -818,9 +954,10 @@ def register_csv( schema: pa.Schema | None = None, has_header: bool = True, delimiter: str = ",", - schema_infer_max_records: int = 1000, + schema_infer_max_records: int = DEFAULT_MAX_INFER_SCHEMA, file_extension: str = ".csv", file_compression_type: str | None = None, + options: CsvReadOptions | None = None, ) -> None: """Register a CSV file as a table. @@ -840,18 +977,46 @@ def register_csv( file_extension: File extension; only files with this extension are selected for data input. file_compression_type: File compression type. + options: Set advanced options for CSV reading. This cannot be + combined with any of the other options in this method. """ - path = [str(p) for p in path] if isinstance(path, list) else str(path) + path_arg = [str(p) for p in path] if isinstance(path, list) else str(path) + + if options is not None and ( + schema is not None + or not has_header + or delimiter != "," + or schema_infer_max_records != DEFAULT_MAX_INFER_SCHEMA + or file_extension != ".csv" + or file_compression_type is not None + ): + message = ( + "Combining CsvReadOptions parameter with additional options " + "is not supported. Use CsvReadOptions to set parameters." + ) + warnings.warn( + message, + category=UserWarning, + stacklevel=2, + ) + + options = ( + options + if options is not None + else CsvReadOptions( + schema=schema, + has_header=has_header, + delimiter=delimiter, + schema_infer_max_records=schema_infer_max_records, + file_extension=file_extension, + file_compression_type=file_compression_type, + ) + ) self.ctx.register_csv( name, - path, - schema, - has_header, - delimiter, - schema_infer_max_records, - file_extension, - file_compression_type, + path_arg, + options.to_inner(), ) def register_json( @@ -861,7 +1026,7 @@ def register_json( schema: pa.Schema | None = None, schema_infer_max_records: int = 1000, file_extension: str = ".json", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, ) -> None: """Register a JSON file as a table. @@ -882,6 +1047,7 @@ def register_json( """ if table_partition_cols is None: table_partition_cols = [] + table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_json( name, str(path), @@ -898,7 +1064,7 @@ def register_avro( path: str | pathlib.Path, schema: pa.Schema | None = None, file_extension: str = ".avro", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, ) -> None: """Register an Avro file as a table. @@ -914,6 +1080,7 @@ def register_avro( """ if table_partition_cols is None: table_partition_cols = [] + table_partition_cols = _convert_table_partition_cols(table_partition_cols) self.ctx.register_avro( name, str(path), schema, file_extension, table_partition_cols ) @@ -973,7 +1140,7 @@ def read_json( schema: pa.Schema | None = None, schema_infer_max_records: int = 1000, file_extension: str = ".json", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, ) -> DataFrame: """Read a line-delimited JSON data source. @@ -993,6 +1160,7 @@ def read_json( """ if table_partition_cols is None: table_partition_cols = [] + table_partition_cols = _convert_table_partition_cols(table_partition_cols) return DataFrame( self.ctx.read_json( str(path), @@ -1010,10 +1178,11 @@ def read_csv( schema: pa.Schema | None = None, has_header: bool = True, delimiter: str = ",", - schema_infer_max_records: int = 1000, + schema_infer_max_records: int = DEFAULT_MAX_INFER_SCHEMA, file_extension: str = ".csv", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, + options: CsvReadOptions | None = None, ) -> DataFrame: """Read a CSV data source. @@ -1031,37 +1200,63 @@ def read_csv( selected for data input. table_partition_cols: Partition columns. file_compression_type: File compression type. + options: Set advanced options for CSV reading. This cannot be + combined with any of the other options in this method. Returns: DataFrame representation of the read CSV files """ - if table_partition_cols is None: - table_partition_cols = [] + path_arg = [str(p) for p in path] if isinstance(path, list) else str(path) + + if options is not None and ( + schema is not None + or not has_header + or delimiter != "," + or schema_infer_max_records != DEFAULT_MAX_INFER_SCHEMA + or file_extension != ".csv" + or table_partition_cols is not None + or file_compression_type is not None + ): + message = ( + "Combining CsvReadOptions parameter with additional options " + "is not supported. Use CsvReadOptions to set parameters." + ) + warnings.warn( + message, + category=UserWarning, + stacklevel=2, + ) - path = [str(p) for p in path] if isinstance(path, list) else str(path) + options = ( + options + if options is not None + else CsvReadOptions( + schema=schema, + has_header=has_header, + delimiter=delimiter, + schema_infer_max_records=schema_infer_max_records, + file_extension=file_extension, + table_partition_cols=table_partition_cols, + file_compression_type=file_compression_type, + ) + ) return DataFrame( self.ctx.read_csv( - path, - schema, - has_header, - delimiter, - schema_infer_max_records, - file_extension, - table_partition_cols, - file_compression_type, + path_arg, + options.to_inner(), ) ) def read_parquet( self, path: str | pathlib.Path, - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, parquet_pruning: bool = True, file_extension: str = ".parquet", skip_metadata: bool = True, schema: pa.Schema | None = None, - file_sort_order: list[list[Expr | SortExpr]] | None = None, + file_sort_order: Sequence[Sequence[SortKey]] | None = None, ) -> DataFrame: """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`. @@ -1078,18 +1273,17 @@ def read_parquet( schema: An optional schema representing the parquet files. If None, the parquet reader will try to infer it based on data in the file. - file_sort_order: Sort order for the file. + file_sort_order: Sort order for the file. Each sort key can be + specified as a column name (``str``), an expression + (``Expr``), or a ``SortExpr``. Returns: DataFrame representation of the read Parquet files """ if table_partition_cols is None: table_partition_cols = [] - file_sort_order = ( - [sort_list_to_raw_sort_list(f) for f in file_sort_order] - if file_sort_order is not None - else None - ) + table_partition_cols = _convert_table_partition_cols(table_partition_cols) + file_sort_order = self._convert_file_sort_order(file_sort_order) return DataFrame( self.ctx.read_parquet( str(path), @@ -1106,7 +1300,7 @@ def read_avro( self, path: str | pathlib.Path, schema: pa.Schema | None = None, - file_partition_cols: list[tuple[str, str]] | None = None, + file_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_extension: str = ".avro", ) -> DataFrame: """Create a :py:class:`DataFrame` for reading Avro data source. @@ -1122,19 +1316,89 @@ def read_avro( """ if file_partition_cols is None: file_partition_cols = [] + file_partition_cols = _convert_table_partition_cols(file_partition_cols) return DataFrame( self.ctx.read_avro(str(path), schema, file_partition_cols, file_extension) ) - def read_table(self, table: Table) -> DataFrame: - """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table. - - For a :py:class:`~datafusion.catalog.Table` such as a - :py:class:`~datafusion.catalog.ListingTable`, create a - :py:class:`~datafusion.dataframe.DataFrame`. - """ - return DataFrame(self.ctx.read_table(table.table)) + def read_table( + self, table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset + ) -> DataFrame: + """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table.""" + return DataFrame(self.ctx.read_table(table)) def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream: """Execute the ``plan`` and return the results.""" return RecordBatchStream(self.ctx.execute(plan._raw_plan, partitions)) + + @staticmethod + def _convert_file_sort_order( + file_sort_order: Sequence[Sequence[SortKey]] | None, + ) -> list[list[expr_internal.SortExpr]] | None: + """Convert nested ``SortKey`` sequences into raw sort expressions. + + Each ``SortKey`` can be a column name string, an ``Expr``, or a + ``SortExpr`` and will be converted using + :func:`datafusion.expr.sort_list_to_raw_sort_list`. + """ + # Convert each ``SortKey`` in the provided sort order to the low-level + # representation expected by the Rust bindings. + return ( + [sort_list_to_raw_sort_list(f) for f in file_sort_order] + if file_sort_order is not None + else None + ) + + @staticmethod + def _convert_table_partition_cols( + table_partition_cols: list[tuple[str, str | pa.DataType]], + ) -> list[tuple[str, pa.DataType]]: + warn = False + converted_table_partition_cols = [] + + for col, data_type in table_partition_cols: + if isinstance(data_type, str): + warn = True + if data_type == "string": + converted_data_type = pa.string() + elif data_type == "int": + converted_data_type = pa.int32() + else: + message = ( + f"Unsupported literal data type '{data_type}' for partition " + "column. Supported types are 'string' and 'int'" + ) + raise ValueError(message) + else: + converted_data_type = data_type + + converted_table_partition_cols.append((col, converted_data_type)) + + if warn: + message = ( + "using literals for table_partition_cols data types is deprecated," + "use pyarrow types instead" + ) + warnings.warn( + message, + category=DeprecationWarning, + stacklevel=2, + ) + + return converted_table_partition_cols + + def __datafusion_task_context_provider__(self) -> Any: + """Access the PyCapsule FFI_TaskContextProvider.""" + return self.ctx.__datafusion_task_context_provider__() + + def __datafusion_logical_extension_codec__(self) -> Any: + """Access the PyCapsule FFI_LogicalExtensionCodec.""" + return self.ctx.__datafusion_logical_extension_codec__() + + def with_logical_extension_codec(self, codec: Any) -> SessionContext: + """Create a new session context with specified codec. + + This only supports codecs that have been implemented using the + FFI interface. + """ + return self.ctx.with_logical_extension_codec(codec) diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index 26fe8f453..214d44a42 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -22,13 +22,11 @@ from __future__ import annotations import warnings +from collections.abc import AsyncIterator, Iterable, Iterator, Sequence from typing import ( TYPE_CHECKING, Any, - Iterable, Literal, - Optional, - Union, overload, ) @@ -37,24 +35,35 @@ except ImportError: from typing_extensions import deprecated # Python 3.12 +from datafusion._internal import DataFrame as DataFrameInternal +from datafusion._internal import DataFrameWriteOptions as DataFrameWriteOptionsInternal +from datafusion._internal import InsertOp as InsertOpInternal +from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal +from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal +from datafusion.expr import ( + Expr, + SortExpr, + SortKey, + ensure_expr, + ensure_expr_list, + expr_list_to_raw_expr_list, + sort_list_to_raw_sort_list, +) from datafusion.plan import ExecutionPlan, LogicalPlan -from datafusion.record_batch import RecordBatchStream +from datafusion.record_batch import RecordBatch, RecordBatchStream if TYPE_CHECKING: import pathlib - from typing import Callable, Sequence + from collections.abc import Callable import pandas as pd import polars as pl import pyarrow as pa - from datafusion._internal import DataFrame as DataFrameInternal - from datafusion._internal import expr as expr_internal + from datafusion.catalog import Table from enum import Enum -from datafusion.expr import Expr, SortExpr, sort_or_default - # excerpt from deltalake # https://github.com/apache/datafusion-python/pull/981#discussion_r1905619163 @@ -68,7 +77,7 @@ class Compression(Enum): LZ4 = "lz4" # lzo is not implemented yet # https://github.com/apache/arrow-rs/issues/6970 - # LZO = "lzo" + # LZO = "lzo" # noqa: ERA001 ZSTD = "zstd" LZ4_RAW = "lz4_raw" @@ -95,7 +104,7 @@ def from_str(cls: type[Compression], value: str) -> Compression: """ raise ValueError(error_msg) from err - def get_default_level(self) -> Optional[int]: + def get_default_level(self) -> int | None: """Get the default compression level for the compression type. Returns: @@ -114,9 +123,190 @@ def get_default_level(self) -> Optional[int]: return None +class ParquetWriterOptions: + """Advanced parquet writer options. + + Allows settings the writer options that apply to the entire file. Some options can + also be set on a column by column basis, with the field ``column_specific_options`` + (see ``ParquetColumnOptions``). + """ + + def __init__( + self, + data_pagesize_limit: int = 1024 * 1024, + write_batch_size: int = 1024, + writer_version: str = "1.0", + skip_arrow_metadata: bool = False, + compression: str | None = "zstd(3)", + compression_level: int | None = None, + dictionary_enabled: bool | None = True, + dictionary_page_size_limit: int = 1024 * 1024, + statistics_enabled: str | None = "page", + max_row_group_size: int = 1024 * 1024, + created_by: str = "datafusion-python", + column_index_truncate_length: int | None = 64, + statistics_truncate_length: int | None = None, + data_page_row_count_limit: int = 20_000, + encoding: str | None = None, + bloom_filter_on_write: bool = False, + bloom_filter_fpp: float | None = None, + bloom_filter_ndv: int | None = None, + allow_single_file_parallelism: bool = True, + maximum_parallel_row_group_writers: int = 1, + maximum_buffered_record_batches_per_stream: int = 2, + column_specific_options: dict[str, ParquetColumnOptions] | None = None, + ) -> None: + """Initialize the ParquetWriterOptions. + + Args: + data_pagesize_limit: Sets best effort maximum size of data page in bytes. + write_batch_size: Sets write_batch_size in bytes. + writer_version: Sets parquet writer version. Valid values are ``1.0`` and + ``2.0``. + skip_arrow_metadata: Skip encoding the embedded arrow metadata in the + KV_meta. + compression: Compression type to use. Default is ``zstd(3)``. + Available compression types are + + - ``uncompressed``: No compression. + - ``snappy``: Snappy compression. + - ``gzip(n)``: Gzip compression with level n. + - ``brotli(n)``: Brotli compression with level n. + - ``lz4``: LZ4 compression. + - ``lz4_raw``: LZ4_RAW compression. + - ``zstd(n)``: Zstandard compression with level n. + compression_level: Compression level to set. + dictionary_enabled: Sets if dictionary encoding is enabled. If ``None``, + uses the default parquet writer setting. + dictionary_page_size_limit: Sets best effort maximum dictionary page size, + in bytes. + statistics_enabled: Sets if statistics are enabled for any column Valid + values are ``none``, ``chunk``, and ``page``. If ``None``, uses the + default parquet writer setting. + max_row_group_size: Target maximum number of rows in each row group + (defaults to 1M rows). Writing larger row groups requires more memory + to write, but can get better compression and be faster to read. + created_by: Sets "created by" property. + column_index_truncate_length: Sets column index truncate length. + statistics_truncate_length: Sets statistics truncate length. If ``None``, + uses the default parquet writer setting. + data_page_row_count_limit: Sets best effort maximum number of rows in a data + page. + encoding: Sets default encoding for any column. Valid values are ``plain``, + ``plain_dictionary``, ``rle``, ``bit_packed``, ``delta_binary_packed``, + ``delta_length_byte_array``, ``delta_byte_array``, ``rle_dictionary``, + and ``byte_stream_split``. If ``None``, uses the default parquet writer + setting. + bloom_filter_on_write: Write bloom filters for all columns when creating + parquet files. + bloom_filter_fpp: Sets bloom filter false positive probability. If ``None``, + uses the default parquet writer setting + bloom_filter_ndv: Sets bloom filter number of distinct values. If ``None``, + uses the default parquet writer setting. + allow_single_file_parallelism: Controls whether DataFusion will attempt to + speed up writing parquet files by serializing them in parallel. Each + column in each row group in each output file are serialized in parallel + leveraging a maximum possible core count of + ``n_files * n_row_groups * n_columns``. + maximum_parallel_row_group_writers: By default parallel parquet writer is + tuned for minimum memory usage in a streaming execution plan. You may + see a performance benefit when writing large parquet files by increasing + ``maximum_parallel_row_group_writers`` and + ``maximum_buffered_record_batches_per_stream`` if your system has idle + cores and can tolerate additional memory usage. Boosting these values is + likely worthwhile when writing out already in-memory data, such as from + a cached data frame. + maximum_buffered_record_batches_per_stream: See + ``maximum_parallel_row_group_writers``. + column_specific_options: Overrides options for specific columns. If a column + is not a part of this dictionary, it will use the parameters provided + here. + """ + self.data_pagesize_limit = data_pagesize_limit + self.write_batch_size = write_batch_size + self.writer_version = writer_version + self.skip_arrow_metadata = skip_arrow_metadata + if compression_level is not None: + self.compression = f"{compression}({compression_level})" + else: + self.compression = compression + self.dictionary_enabled = dictionary_enabled + self.dictionary_page_size_limit = dictionary_page_size_limit + self.statistics_enabled = statistics_enabled + self.max_row_group_size = max_row_group_size + self.created_by = created_by + self.column_index_truncate_length = column_index_truncate_length + self.statistics_truncate_length = statistics_truncate_length + self.data_page_row_count_limit = data_page_row_count_limit + self.encoding = encoding + self.bloom_filter_on_write = bloom_filter_on_write + self.bloom_filter_fpp = bloom_filter_fpp + self.bloom_filter_ndv = bloom_filter_ndv + self.allow_single_file_parallelism = allow_single_file_parallelism + self.maximum_parallel_row_group_writers = maximum_parallel_row_group_writers + self.maximum_buffered_record_batches_per_stream = ( + maximum_buffered_record_batches_per_stream + ) + self.column_specific_options = column_specific_options + + +class ParquetColumnOptions: + """Parquet options for individual columns. + + Contains the available options that can be applied for an individual Parquet column, + replacing the global options in ``ParquetWriterOptions``. + """ + + def __init__( + self, + encoding: str | None = None, + dictionary_enabled: bool | None = None, + compression: str | None = None, + statistics_enabled: str | None = None, + bloom_filter_enabled: bool | None = None, + bloom_filter_fpp: float | None = None, + bloom_filter_ndv: int | None = None, + ) -> None: + """Initialize the ParquetColumnOptions. + + Args: + encoding: Sets encoding for the column path. Valid values are: ``plain``, + ``plain_dictionary``, ``rle``, ``bit_packed``, ``delta_binary_packed``, + ``delta_length_byte_array``, ``delta_byte_array``, ``rle_dictionary``, + and ``byte_stream_split``. These values are not case-sensitive. If + ``None``, uses the default parquet options + dictionary_enabled: Sets if dictionary encoding is enabled for the column + path. If `None`, uses the default parquet options + compression: Sets default parquet compression codec for the column path. + Valid values are ``uncompressed``, ``snappy``, ``gzip(level)``, ``lzo``, + ``brotli(level)``, ``lz4``, ``zstd(level)``, and ``lz4_raw``. These + values are not case-sensitive. If ``None``, uses the default parquet + options. + statistics_enabled: Sets if statistics are enabled for the column Valid + values are: ``none``, ``chunk``, and ``page`` These values are not case + sensitive. If ``None``, uses the default parquet options. + bloom_filter_enabled: Sets if bloom filter is enabled for the column path. + If ``None``, uses the default parquet options. + bloom_filter_fpp: Sets bloom filter false positive probability for the + column path. If ``None``, uses the default parquet options. + bloom_filter_ndv: Sets bloom filter number of distinct values. If ``None``, + uses the default parquet options. + """ + self.encoding = encoding + self.dictionary_enabled = dictionary_enabled + self.compression = compression + self.statistics_enabled = statistics_enabled + self.bloom_filter_enabled = bloom_filter_enabled + self.bloom_filter_fpp = bloom_filter_fpp + self.bloom_filter_ndv = bloom_filter_ndv + + class DataFrame: """Two dimensional table representation of data. + DataFrame objects are iterable; iterating over a DataFrame yields + :class:`datafusion.RecordBatch` instances lazily. + See :ref:`user_guide_concepts` in the online documentation for more information. """ @@ -128,12 +318,25 @@ def __init__(self, df: DataFrameInternal) -> None: """ self.df = df - def into_view(self) -> pa.Table: - """Convert DataFrame as a ViewTable which can be used in register_table.""" - return self.df.into_view() + def into_view(self, temporary: bool = False) -> Table: + """Convert ``DataFrame`` into a :class:`~datafusion.Table`. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> df = ctx.sql("SELECT 1 AS value") + >>> view = df.into_view() + >>> ctx.register_table("values_view", view) + >>> result = ctx.sql("SELECT value FROM values_view").collect() + >>> result[0].column("value").to_pylist() + [1] + """ + from datafusion.catalog import Table as _Table + + return _Table(self.df.into_view(temporary)) def __getitem__(self, key: str | list[str]) -> DataFrame: - """Return a new :py:class`DataFrame` with the specified column or columns. + """Return a new :py:class:`DataFrame` with the specified column or columns. Args: key: Column name or list of column names to select. @@ -154,6 +357,20 @@ def __repr__(self) -> str: def _repr_html_(self) -> str: return self.df._repr_html_() + @staticmethod + def default_str_repr( + batches: list[pa.RecordBatch], + schema: pa.Schema, + has_more: bool, + table_uuid: str | None = None, + ) -> str: + """Return the default string representation of a DataFrame. + + This method is used by the default formatter and implemented in Rust for + performance reasons. + """ + return DataFrameInternal.default_str_repr(batches, schema, has_more, table_uuid) + def describe(self) -> DataFrame: """Return the statistics for this DataFrame. @@ -189,6 +406,17 @@ def select_columns(self, *args: str) -> DataFrame: """ return self.select(*args) + def select_exprs(self, *args: str) -> DataFrame: + """Project arbitrary list of expression strings into a new DataFrame. + + This method will parse string expressions into logical plan expressions. + The output DataFrame has one column for each expression. + + Returns: + DataFrame only containing the specified columns. + """ + return self.df.select_exprs(*args) + def select(self, *exprs: Expr | str) -> DataFrame: """Project arbitrary expressions into a new :py:class:`DataFrame`. @@ -208,44 +436,94 @@ def select(self, *exprs: Expr | str) -> DataFrame: df = df.select("a", col("b"), col("a").alias("alternate_a")) """ - exprs_internal = [ - Expr.column(arg).expr if isinstance(arg, str) else arg.expr for arg in exprs - ] + exprs_internal = expr_list_to_raw_expr_list(exprs) return DataFrame(self.df.select(*exprs_internal)) def drop(self, *columns: str) -> DataFrame: """Drop arbitrary amount of columns. + Column names are case-sensitive and require double quotes to be dropped + if the original name is not strictly lower case. + Args: columns: Column names to drop from the dataframe. Returns: DataFrame with those columns removed in the projection. + + Example Usage:: + df.drop('a') # To drop a lower-cased column 'a' + df.drop('"a"') # To drop an upper-cased column 'A' """ return DataFrame(self.df.drop(*columns)) - def filter(self, *predicates: Expr) -> DataFrame: + def filter(self, *predicates: Expr | str) -> DataFrame: """Return a DataFrame for which ``predicate`` evaluates to ``True``. Rows for which ``predicate`` evaluates to ``False`` or ``None`` are filtered - out. If more than one predicate is provided, these predicates will be - combined as a logical AND. If more complex logic is required, see the - logical operations in :py:mod:`~datafusion.functions`. + out. If more than one predicate is provided, these predicates will be + combined as a logical AND. Each ``predicate`` can be an + :class:`~datafusion.expr.Expr` created using helper functions such as + :func:`datafusion.col` or :func:`datafusion.lit`, or a SQL expression string + that will be parsed against the DataFrame schema. If more complex logic is + required, see the logical operations in :py:mod:`~datafusion.functions`. + + Example:: + + from datafusion import col, lit + df.filter(col("a") > lit(1)) + df.filter("a > 1") Args: - predicates: Predicate expression(s) to filter the DataFrame. + predicates: Predicate expression(s) or SQL strings to filter the DataFrame. Returns: DataFrame after filtering. """ df = self.df - for p in predicates: - df = df.filter(p.expr) + for predicate in predicates: + expr = ( + self.parse_sql_expr(predicate) + if isinstance(predicate, str) + else predicate + ) + df = df.filter(ensure_expr(expr)) return DataFrame(df) - def with_column(self, name: str, expr: Expr) -> DataFrame: + def parse_sql_expr(self, expr: str) -> Expr: + """Creates logical expression from a SQL query text. + + The expression is created and processed against the current schema. + + Example:: + + from datafusion import col, lit + df.parse_sql_expr("a > 1") + + should produce: + + col("a") > lit(1) + + Args: + expr: Expression string to be converted to datafusion expression + + Returns: + Logical expression . + """ + return Expr(self.df.parse_sql_expr(expr)) + + def with_column(self, name: str, expr: Expr | str) -> DataFrame: """Add an additional column to the DataFrame. + The ``expr`` must be an :class:`~datafusion.expr.Expr` constructed with + :func:`datafusion.col` or :func:`datafusion.lit`, or a SQL expression + string that will be parsed against the DataFrame schema. + + Example:: + + from datafusion import col, lit + df.with_column("b", col("a") + lit(1)) + Args: name: Name of the column to add. expr: Expression to compute the column. @@ -253,49 +531,69 @@ def with_column(self, name: str, expr: Expr) -> DataFrame: Returns: DataFrame with the new column. """ - return DataFrame(self.df.with_column(name, expr.expr)) + expr = self.parse_sql_expr(expr) if isinstance(expr, str) else expr + + return DataFrame(self.df.with_column(name, ensure_expr(expr))) def with_columns( - self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr + self, *exprs: Expr | str | Iterable[Expr | str], **named_exprs: Expr | str ) -> DataFrame: """Add columns to the DataFrame. - By passing expressions, iteratables of expressions, or named expressions. To - pass named expressions use the form name=Expr. + By passing expressions, iterables of expressions, string SQL expressions, + or named expressions. + All expressions must be :class:`~datafusion.expr.Expr` objects created via + :func:`datafusion.col` or :func:`datafusion.lit`, or SQL expression strings. + To pass named expressions use the form ``name=Expr``. - Example usage: The following will add 4 columns labeled a, b, c, and d:: + Example usage: The following will add 4 columns labeled ``a``, ``b``, ``c``, + and ``d``:: + from datafusion import col, lit df = df.with_columns( - lit(0).alias('a'), - [lit(1).alias('b'), lit(2).alias('c')], + col("x").alias("a"), + [lit(1).alias("b"), col("y").alias("c")], d=lit(3) - ) + ) + + Equivalent example using just SQL strings: + + df = df.with_columns( + "x as a", + ["1 as b", "y as c"], + d="3" + ) Args: - exprs: Either a single expression or an iterable of expressions to add. + exprs: Either a single expression, an iterable of expressions to add or + SQL expression strings. named_exprs: Named expressions in the form of ``name=expr`` Returns: DataFrame with the new columns added. """ + expressions = [] + for expr in exprs: + if isinstance(expr, str): + expressions.append(self.parse_sql_expr(expr).expr) + elif isinstance(expr, Iterable) and not isinstance( + expr, Expr | str | bytes | bytearray + ): + expressions.extend( + [ + self.parse_sql_expr(e).expr + if isinstance(e, str) + else ensure_expr(e) + for e in expr + ] + ) + else: + expressions.append(ensure_expr(expr)) - def _simplify_expression( - *exprs: Expr | Iterable[Expr], **named_exprs: Expr - ) -> list[expr_internal.Expr]: - expr_list = [] - for expr in exprs: - if isinstance(expr, Expr): - expr_list.append(expr.expr) - elif isinstance(expr, Iterable): - expr_list.extend(inner_expr.expr for inner_expr in expr) - else: - raise NotImplementedError - if named_exprs: - for alias, expr in named_exprs.items(): - expr_list.append(expr.alias(alias).expr) - return expr_list - - expressions = _simplify_expression(*exprs, **named_exprs) + for alias, expr in named_exprs.items(): + e = self.parse_sql_expr(expr) if isinstance(expr, str) else expr + ensure_expr(e) + expressions.append(e.alias(alias).expr) return DataFrame(self.df.with_columns(expressions)) @@ -317,37 +615,47 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame: return DataFrame(self.df.with_column_renamed(old_name, new_name)) def aggregate( - self, group_by: list[Expr] | Expr, aggs: list[Expr] | Expr + self, + group_by: Sequence[Expr | str] | Expr | str, + aggs: Sequence[Expr] | Expr, ) -> DataFrame: """Aggregates the rows of the current DataFrame. Args: - group_by: List of expressions to group by. - aggs: List of expressions to aggregate. + group_by: Sequence of expressions or column names to group by. + aggs: Sequence of expressions to aggregate. Returns: DataFrame after aggregation. """ - group_by = group_by if isinstance(group_by, list) else [group_by] - aggs = aggs if isinstance(aggs, list) else [aggs] - - group_by = [e.expr for e in group_by] - aggs = [e.expr for e in aggs] - return DataFrame(self.df.aggregate(group_by, aggs)) - - def sort(self, *exprs: Expr | SortExpr) -> DataFrame: - """Sort the DataFrame by the specified sorting expressions. + group_by_list = ( + list(group_by) + if isinstance(group_by, Sequence) and not isinstance(group_by, Expr | str) + else [group_by] + ) + aggs_list = ( + list(aggs) + if isinstance(aggs, Sequence) and not isinstance(aggs, Expr) + else [aggs] + ) + + group_by_exprs = expr_list_to_raw_expr_list(group_by_list) + aggs_exprs = ensure_expr_list(aggs_list) + return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs)) + + def sort(self, *exprs: SortKey) -> DataFrame: + """Sort the DataFrame by the specified sorting expressions or column names. Note that any expression can be turned into a sort expression by - calling its` ``sort`` method. + calling its ``sort`` method. Args: - exprs: Sort expressions, applied in order. + exprs: Sort expressions or column names, applied in order. Returns: DataFrame after sorting. """ - exprs_raw = [sort_or_default(expr) for expr in exprs] + exprs_raw = sort_list_to_raw_sort_list(exprs) return DataFrame(self.df.sort(*exprs_raw)) def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame: @@ -402,7 +710,7 @@ def tail(self, n: int = 5) -> DataFrame: def collect(self) -> list[pa.RecordBatch]: """Execute this :py:class:`DataFrame` and collect results into memory. - Prior to calling ``collect``, modifying a DataFrme simply updates a plan + Prior to calling ``collect``, modifying a DataFrame simply updates a plan (no actual computation is performed). Calling ``collect`` triggers the computation. @@ -411,6 +719,10 @@ def collect(self) -> list[pa.RecordBatch]: """ return self.df.collect() + def collect_column(self, column_name: str) -> pa.Array | pa.ChunkedArray: + """Executes this :py:class:`DataFrame` for a single column.""" + return self.df.collect_column(column_name) + def cache(self) -> DataFrame: """Cache the DataFrame as a memory table. @@ -457,6 +769,7 @@ def join( left_on: None = None, right_on: None = None, join_keys: None = None, + coalesce_duplicate_keys: bool = True, ) -> DataFrame: ... @overload @@ -469,6 +782,7 @@ def join( left_on: str | Sequence[str], right_on: str | Sequence[str], join_keys: tuple[list[str], list[str]] | None = None, + coalesce_duplicate_keys: bool = True, ) -> DataFrame: ... @overload @@ -481,6 +795,7 @@ def join( join_keys: tuple[list[str], list[str]], left_on: None = None, right_on: None = None, + coalesce_duplicate_keys: bool = True, ) -> DataFrame: ... def join( @@ -492,6 +807,7 @@ def join( left_on: str | Sequence[str] | None = None, right_on: str | Sequence[str] | None = None, join_keys: tuple[list[str], list[str]] | None = None, + coalesce_duplicate_keys: bool = True, ) -> DataFrame: """Join this :py:class:`DataFrame` with another :py:class:`DataFrame`. @@ -504,33 +820,37 @@ def join( "right", "full", "semi", "anti". left_on: Join column of the left dataframe. right_on: Join column of the right dataframe. + coalesce_duplicate_keys: When True, coalesce the columns + from the right DataFrame and left DataFrame + that have identical names in the ``on`` fields. join_keys: Tuple of two lists of column names to join on. [Deprecated] Returns: DataFrame after join. """ + if join_keys is not None: + warnings.warn( + "`join_keys` is deprecated, use `on` or `left_on` with `right_on`", + category=DeprecationWarning, + stacklevel=2, + ) + left_on = join_keys[0] + right_on = join_keys[1] + # This check is to prevent breaking API changes where users prior to # DF 43.0.0 would pass the join_keys as a positional argument instead # of a keyword argument. if ( isinstance(on, tuple) - and len(on) == 2 + and len(on) == 2 # noqa: PLR2004 and isinstance(on[0], list) and isinstance(on[1], list) ): # We know this is safe because we've checked the types - join_keys = on # type: ignore[assignment] + left_on = on[0] + right_on = on[1] on = None - if join_keys is not None: - warnings.warn( - "`join_keys` is deprecated, use `on` or `left_on` with `right_on`", - category=DeprecationWarning, - stacklevel=2, - ) - left_on = join_keys[0] - right_on = join_keys[1] - if on is not None: if left_on is not None or right_on is not None: error_msg = "`left_on` or `right_on` should not provided with `on`" @@ -549,7 +869,9 @@ def join( if isinstance(right_on, str): right_on = [right_on] - return DataFrame(self.df.join(right.df, how, left_on, right_on)) + return DataFrame( + self.df.join(right.df, how, left_on, right_on, coalesce_duplicate_keys) + ) def join_on( self, @@ -559,8 +881,14 @@ def join_on( ) -> DataFrame: """Join two :py:class:`DataFrame` using the specified expressions. - On expressions are used to support in-equality predicates. Equality - predicates are correctly optimized + Join predicates must be :class:`~datafusion.expr.Expr` objects, typically + built with :func:`datafusion.col`. On expressions are used to support + in-equality predicates. Equality predicates are correctly optimized. + + Example:: + + from datafusion import col + df.join_on(other_df, col("id") == col("other_id")) Args: right: Other DataFrame to join with. @@ -571,22 +899,19 @@ def join_on( Returns: DataFrame after join. """ - exprs = [expr.expr for expr in on_exprs] + exprs = [ensure_expr(expr) for expr in on_exprs] return DataFrame(self.df.join_on(right.df, exprs, how)) - def explain(self, verbose: bool = False, analyze: bool = False) -> DataFrame: - """Return a DataFrame with the explanation of its plan so far. + def explain(self, verbose: bool = False, analyze: bool = False) -> None: + """Print an explanation of the DataFrame's plan so far. If ``analyze`` is specified, runs the plan and reports metrics. Args: verbose: If ``True``, more details will be included. - analyze: If ``Tru`e``, the plan will run and metrics reported. - - Returns: - DataFrame with the explanation of its plan. + analyze: If ``True``, the plan will run and metrics reported. """ - return DataFrame(self.df.explain(verbose, analyze)) + self.df.explain(verbose, analyze) def logical_plan(self) -> LogicalPlan: """Return the unoptimized ``LogicalPlan``. @@ -625,17 +950,20 @@ def repartition(self, num: int) -> DataFrame: """ return DataFrame(self.df.repartition(num)) - def repartition_by_hash(self, *exprs: Expr, num: int) -> DataFrame: + def repartition_by_hash(self, *exprs: Expr | str, num: int) -> DataFrame: """Repartition a DataFrame using a hash partitioning scheme. Args: - exprs: Expressions to evaluate and perform hashing on. + exprs: Expressions or a SQL expression string to evaluate + and perform hashing on. num: Number of partitions to repartition the DataFrame into. Returns: Repartitioned DataFrame. """ - exprs = [expr.expr for expr in exprs] + exprs = [self.parse_sql_expr(e) if isinstance(e, str) else e for e in exprs] + exprs = expr_list_to_raw_expr_list(exprs) + return DataFrame(self.df.repartition_by_hash(*exprs, num=num)) def union(self, other: DataFrame, distinct: bool = False) -> DataFrame: @@ -692,40 +1020,88 @@ def except_all(self, other: DataFrame) -> DataFrame: """ return DataFrame(self.df.except_all(other.df)) - def write_csv(self, path: str | pathlib.Path, with_header: bool = False) -> None: + def write_csv( + self, + path: str | pathlib.Path, + with_header: bool = False, + write_options: DataFrameWriteOptions | None = None, + ) -> None: """Execute the :py:class:`DataFrame` and write the results to a CSV file. Args: path: Path of the CSV file to write. with_header: If true, output the CSV header row. + write_options: Options that impact how the DataFrame is written. """ - self.df.write_csv(str(path), with_header) + raw_write_options = ( + write_options._raw_write_options if write_options is not None else None + ) + self.df.write_csv(str(path), with_header, raw_write_options) + + @overload + def write_parquet( + self, + path: str | pathlib.Path, + compression: str, + compression_level: int | None = None, + write_options: DataFrameWriteOptions | None = None, + ) -> None: ... + + @overload + def write_parquet( + self, + path: str | pathlib.Path, + compression: Compression = Compression.ZSTD, + compression_level: int | None = None, + write_options: DataFrameWriteOptions | None = None, + ) -> None: ... + + @overload + def write_parquet( + self, + path: str | pathlib.Path, + compression: ParquetWriterOptions, + compression_level: None = None, + write_options: DataFrameWriteOptions | None = None, + ) -> None: ... def write_parquet( self, path: str | pathlib.Path, - compression: Union[str, Compression] = Compression.ZSTD, + compression: str | Compression | ParquetWriterOptions = Compression.ZSTD, compression_level: int | None = None, + write_options: DataFrameWriteOptions | None = None, ) -> None: """Execute the :py:class:`DataFrame` and write the results to a Parquet file. + Available compression types are: + + - "uncompressed": No compression. + - "snappy": Snappy compression. + - "gzip": Gzip compression. + - "brotli": Brotli compression. + - "lz4": LZ4 compression. + - "lz4_raw": LZ4_RAW compression. + - "zstd": Zstandard compression. + + LZO compression is not yet implemented in arrow-rs and is therefore + excluded. + Args: path: Path of the Parquet file to write. compression: Compression type to use. Default is "ZSTD". - Available compression types are: - - "uncompressed": No compression. - - "snappy": Snappy compression. - - "gzip": Gzip compression. - - "brotli": Brotli compression. - - "lz4": LZ4 compression. - - "lz4_raw": LZ4_RAW compression. - - "zstd": Zstandard compression. - Note: LZO is not yet implemented in arrow-rs and is therefore excluded. compression_level: Compression level to use. For ZSTD, the recommended range is 1 to 22, with the default being 4. Higher levels provide better compression but slower speed. + write_options: Options that impact how the DataFrame is written. """ - # Convert string to Compression enum if necessary + if isinstance(compression, ParquetWriterOptions): + if compression_level is not None: + msg = "compression_level should be None when using ParquetWriterOptions" + raise ValueError(msg) + self.write_parquet_with_options(path, compression) + return + if isinstance(compression, str): compression = Compression.from_str(compression) @@ -735,15 +1111,105 @@ def write_parquet( ): compression_level = compression.get_default_level() - self.df.write_parquet(str(path), compression.value, compression_level) + raw_write_options = ( + write_options._raw_write_options if write_options is not None else None + ) + self.df.write_parquet( + str(path), + compression.value, + compression_level, + raw_write_options, + ) + + def write_parquet_with_options( + self, + path: str | pathlib.Path, + options: ParquetWriterOptions, + write_options: DataFrameWriteOptions | None = None, + ) -> None: + """Execute the :py:class:`DataFrame` and write the results to a Parquet file. + + Allows advanced writer options to be set with `ParquetWriterOptions`. + + Args: + path: Path of the Parquet file to write. + options: Sets the writer parquet options (see `ParquetWriterOptions`). + write_options: Options that impact how the DataFrame is written. + """ + options_internal = ParquetWriterOptionsInternal( + options.data_pagesize_limit, + options.write_batch_size, + options.writer_version, + options.skip_arrow_metadata, + options.compression, + options.dictionary_enabled, + options.dictionary_page_size_limit, + options.statistics_enabled, + options.max_row_group_size, + options.created_by, + options.column_index_truncate_length, + options.statistics_truncate_length, + options.data_page_row_count_limit, + options.encoding, + options.bloom_filter_on_write, + options.bloom_filter_fpp, + options.bloom_filter_ndv, + options.allow_single_file_parallelism, + options.maximum_parallel_row_group_writers, + options.maximum_buffered_record_batches_per_stream, + ) + + column_specific_options_internal = {} + for column, opts in (options.column_specific_options or {}).items(): + column_specific_options_internal[column] = ParquetColumnOptionsInternal( + bloom_filter_enabled=opts.bloom_filter_enabled, + encoding=opts.encoding, + dictionary_enabled=opts.dictionary_enabled, + compression=opts.compression, + statistics_enabled=opts.statistics_enabled, + bloom_filter_fpp=opts.bloom_filter_fpp, + bloom_filter_ndv=opts.bloom_filter_ndv, + ) - def write_json(self, path: str | pathlib.Path) -> None: + raw_write_options = ( + write_options._raw_write_options if write_options is not None else None + ) + self.df.write_parquet_with_options( + str(path), + options_internal, + column_specific_options_internal, + raw_write_options, + ) + + def write_json( + self, + path: str | pathlib.Path, + write_options: DataFrameWriteOptions | None = None, + ) -> None: """Execute the :py:class:`DataFrame` and write the results to a JSON file. Args: path: Path of the JSON file to write. + write_options: Options that impact how the DataFrame is written. + """ + raw_write_options = ( + write_options._raw_write_options if write_options is not None else None + ) + self.df.write_json(str(path), write_options=raw_write_options) + + def write_table( + self, table_name: str, write_options: DataFrameWriteOptions | None = None + ) -> None: + """Execute the :py:class:`DataFrame` and write the results to a table. + + The table must be registered with the session to perform this operation. + Not all table providers support writing operations. See the individual + implementations for details. """ - self.df.write_json(str(path)) + raw_write_options = ( + write_options._raw_write_options if write_options is not None else None + ) + self.df.write_table(table_name, raw_write_options) def to_arrow_table(self) -> pa.Table: """Execute the :py:class:`DataFrame` and convert it into an Arrow Table. @@ -832,22 +1298,55 @@ def unnest_columns(self, *columns: str, preserve_nulls: bool = True) -> DataFram columns = list(columns) return DataFrame(self.df.unnest_columns(columns, preserve_nulls=preserve_nulls)) - def __arrow_c_stream__(self, requested_schema: pa.Schema) -> Any: - """Export an Arrow PyCapsule Stream. + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """Export the DataFrame as an Arrow C Stream. + + The DataFrame is executed using DataFusion's streaming APIs and exposed via + Arrow's C Stream interface. Record batches are produced incrementally, so the + full result set is never materialized in memory. - This will execute and collect the DataFrame. We will attempt to respect the - requested schema, but only trivial transformations will be applied such as only - returning the fields listed in the requested schema if their data types match - those in the DataFrame. + When ``requested_schema`` is provided, DataFusion applies only simple + projections such as selecting a subset of existing columns or reordering + them. Column renaming, computed expressions, or type coercion are not + supported through this interface. Args: - requested_schema: Attempt to provide the DataFrame using this schema. + requested_schema: Either a :py:class:`pyarrow.Schema` or an Arrow C + Schema capsule (``PyCapsule``) produced by + ``schema._export_to_c_capsule()``. The DataFrame will attempt to + align its output with the fields and order specified by this schema. Returns: - Arrow PyCapsule object. + Arrow ``PyCapsule`` object representing an ``ArrowArrayStream``. + + For practical usage patterns, see the Apache Arrow streaming + documentation: https://arrow.apache.org/docs/python/ipc.html#streaming. + + For details on DataFusion's Arrow integration and DataFrame streaming, + see the user guide (user-guide/io/arrow and user-guide/dataframe/index). + + Notes: + The Arrow C Data Interface PyCapsule details are documented by Apache + Arrow and can be found at: + https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html """ + # ``DataFrame.__arrow_c_stream__`` in the Rust extension leverages + # ``execute_stream_partitioned`` under the hood to stream batches while + # preserving the original partition order. return self.df.__arrow_c_stream__(requested_schema) + def __iter__(self) -> Iterator[RecordBatch]: + """Return an iterator over this DataFrame's record batches.""" + return iter(self.execute_stream()) + + def __aiter__(self) -> AsyncIterator[RecordBatch]: + """Return an async iterator over this DataFrame's record batches. + + We're using __aiter__ because we support Python < 3.10 where aiter() is not + available. + """ + return self.execute_stream().__aiter__() + def transform(self, func: Callable[..., DataFrame], *args: Any) -> DataFrame: """Apply a function to the current DataFrame which returns another DataFrame. @@ -869,3 +1368,74 @@ def within_limit(df: DataFrame, limit: int) -> DataFrame: DataFrame: After applying func to the original dataframe. """ return func(self, *args) + + def fill_null(self, value: Any, subset: list[str] | None = None) -> DataFrame: + """Fill null values in specified columns with a value. + + Args: + value: Value to replace nulls with. Will be cast to match column type. + subset: Optional list of column names to fill. If None, fills all columns. + + Returns: + DataFrame with null values replaced where type casting is possible + + Examples: + >>> from datafusion import SessionContext, col + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, None, 3], "b": [None, 5, 6]}) + >>> filled = df.fill_null(0) + >>> filled.sort(col("a")).collect()[0].column("a").to_pylist() + [0, 1, 3] + + Notes: + - Only fills nulls in columns where the value can be cast to the column type + - For columns where casting fails, the original column is kept unchanged + - For columns not in subset, the original column is kept unchanged + """ + return DataFrame(self.df.fill_null(value, subset)) + + +class InsertOp(Enum): + """Insert operation mode. + + These modes are used by the table writing feature to define how record + batches should be written to a table. + """ + + APPEND = InsertOpInternal.APPEND + """Appends new rows to the existing table without modifying any existing rows.""" + + REPLACE = InsertOpInternal.REPLACE + """Replace existing rows that collide with the inserted rows. + + Replacement is typically based on a unique key or primary key. + """ + + OVERWRITE = InsertOpInternal.OVERWRITE + """Overwrites all existing rows in the table with the new rows.""" + + +class DataFrameWriteOptions: + """Writer options for DataFrame. + + There is no guarantee the table provider supports all writer options. + See the individual implementation and documentation for details. + """ + + def __init__( + self, + insert_operation: InsertOp | None = None, + single_file_output: bool = False, + partition_by: str | Sequence[str] | None = None, + sort_by: Expr | SortExpr | Sequence[Expr] | Sequence[SortExpr] | None = None, + ) -> None: + """Instantiate writer options for DataFrame.""" + if isinstance(partition_by, str): + partition_by = [partition_by] + + sort_by_raw = sort_list_to_raw_sort_list(sort_by) + insert_op = insert_operation.value if insert_operation is not None else None + + self._raw_write_options = DataFrameWriteOptionsInternal( + insert_op, single_file_output, partition_by, sort_by_raw + ) diff --git a/python/datafusion/dataframe_formatter.py b/python/datafusion/dataframe_formatter.py new file mode 100644 index 000000000..b8af45a1b --- /dev/null +++ b/python/datafusion/dataframe_formatter.py @@ -0,0 +1,843 @@ +# 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. +"""HTML formatting utilities for DataFusion DataFrames.""" + +from __future__ import annotations + +import warnings +from typing import ( + TYPE_CHECKING, + Any, + Protocol, + runtime_checkable, +) + +from datafusion._internal import DataFrame as DataFrameInternal + +if TYPE_CHECKING: + from collections.abc import Callable + + +def _validate_positive_int(value: Any, param_name: str) -> None: + """Validate that a parameter is a positive integer. + + Args: + value: The value to validate + param_name: Name of the parameter (used in error message) + + Raises: + ValueError: If the value is not a positive integer + """ + if not isinstance(value, int) or value <= 0: + msg = f"{param_name} must be a positive integer" + raise ValueError(msg) + + +def _validate_bool(value: Any, param_name: str) -> None: + """Validate that a parameter is a boolean. + + Args: + value: The value to validate + param_name: Name of the parameter (used in error message) + + Raises: + TypeError: If the value is not a boolean + """ + if not isinstance(value, bool): + msg = f"{param_name} must be a boolean" + raise TypeError(msg) + + +def _validate_formatter_parameters( + max_cell_length: int, + max_width: int, + max_height: int, + max_memory_bytes: int, + min_rows: int, + max_rows: int | None, + repr_rows: int | None, + enable_cell_expansion: bool, + show_truncation_message: bool, + use_shared_styles: bool, + custom_css: str | None, + style_provider: Any, +) -> int: + """Validate all formatter parameters and return resolved max_rows value. + + Args: + max_cell_length: Maximum cell length value to validate + max_width: Maximum width value to validate + max_height: Maximum height value to validate + max_memory_bytes: Maximum memory bytes value to validate + min_rows: Minimum rows to display value to validate + max_rows: Maximum rows value to validate (None means use default) + repr_rows: Deprecated repr_rows value to validate + enable_cell_expansion: Boolean expansion flag to validate + show_truncation_message: Boolean message flag to validate + use_shared_styles: Boolean styles flag to validate + custom_css: Custom CSS string to validate + style_provider: Style provider object to validate + + Returns: + The resolved max_rows value after handling repr_rows deprecation + + Raises: + ValueError: If any numeric parameter is invalid or constraints are violated + TypeError: If any parameter has invalid type + DeprecationWarning: If repr_rows parameter is used + """ + # Validate numeric parameters + _validate_positive_int(max_cell_length, "max_cell_length") + _validate_positive_int(max_width, "max_width") + _validate_positive_int(max_height, "max_height") + _validate_positive_int(max_memory_bytes, "max_memory_bytes") + _validate_positive_int(min_rows, "min_rows") + + # Handle deprecated repr_rows parameter + if repr_rows is not None: + warnings.warn( + "repr_rows parameter is deprecated, use max_rows instead", + DeprecationWarning, + stacklevel=4, + ) + _validate_positive_int(repr_rows, "repr_rows") + if max_rows is not None and repr_rows != max_rows: + msg = "Cannot specify both repr_rows and max_rows; use max_rows only" + raise ValueError(msg) + max_rows = repr_rows + + # Use default if max_rows was not provided + if max_rows is None: + max_rows = 10 + + _validate_positive_int(max_rows, "max_rows") + + # Validate constraint: min_rows <= max_rows + if min_rows > max_rows: + msg = "min_rows must be less than or equal to max_rows" + raise ValueError(msg) + + # Validate boolean parameters + _validate_bool(enable_cell_expansion, "enable_cell_expansion") + _validate_bool(show_truncation_message, "show_truncation_message") + _validate_bool(use_shared_styles, "use_shared_styles") + + # Validate custom_css + if custom_css is not None and not isinstance(custom_css, str): + msg = "custom_css must be None or a string" + raise TypeError(msg) + + # Validate style_provider + if style_provider is not None and not isinstance(style_provider, StyleProvider): + msg = "style_provider must implement the StyleProvider protocol" + raise TypeError(msg) + + return max_rows + + +@runtime_checkable +class CellFormatter(Protocol): + """Protocol for cell value formatters.""" + + def __call__(self, value: Any) -> str: + """Format a cell value to string representation.""" + ... + + +@runtime_checkable +class StyleProvider(Protocol): + """Protocol for HTML style providers.""" + + def get_cell_style(self) -> str: + """Get the CSS style for table cells.""" + ... + + def get_header_style(self) -> str: + """Get the CSS style for header cells.""" + ... + + +class DefaultStyleProvider: + """Default implementation of StyleProvider.""" + + def get_cell_style(self) -> str: + """Get the CSS style for table cells. + + Returns: + CSS style string + """ + return ( + "border: 1px solid black; padding: 8px; text-align: left; " + "white-space: nowrap;" + ) + + def get_header_style(self) -> str: + """Get the CSS style for header cells. + + Returns: + CSS style string + """ + return ( + "border: 1px solid black; padding: 8px; text-align: left; " + "background-color: #f2f2f2; white-space: nowrap; min-width: fit-content; " + "max-width: fit-content;" + ) + + +class DataFrameHtmlFormatter: + """Configurable HTML formatter for DataFusion DataFrames. + + This class handles the HTML rendering of DataFrames for display in + Jupyter notebooks and other rich display contexts. + + This class supports extension through composition. Key extension points: + - Provide a custom StyleProvider for styling cells and headers + - Register custom formatters for specific types + - Provide custom cell builders for specialized cell rendering + + Args: + max_cell_length: Maximum characters to display in a cell before truncation + max_width: Maximum width of the HTML table in pixels + max_height: Maximum height of the HTML table in pixels + max_memory_bytes: Maximum memory in bytes for rendered data (default: 2MB) + min_rows: Minimum number of rows to display (must be <= max_rows) + max_rows: Maximum number of rows to display in repr output + repr_rows: Deprecated alias for max_rows + enable_cell_expansion: Whether to add expand/collapse buttons for long cell + values + custom_css: Additional CSS to include in the HTML output + show_truncation_message: Whether to display a message when data is truncated + style_provider: Custom provider for cell and header styles + use_shared_styles: Whether to load styles and scripts only once per notebook + session + """ + + def __init__( + self, + max_cell_length: int = 25, + max_width: int = 1000, + max_height: int = 300, + max_memory_bytes: int = 2 * 1024 * 1024, # 2 MB + min_rows: int = 10, + max_rows: int | None = None, + repr_rows: int | None = None, + enable_cell_expansion: bool = True, + custom_css: str | None = None, + show_truncation_message: bool = True, + style_provider: StyleProvider | None = None, + use_shared_styles: bool = True, + ) -> None: + """Initialize the HTML formatter. + + Parameters + ---------- + max_cell_length + Maximum length of cell content before truncation. + max_width + Maximum width of the displayed table in pixels. + max_height + Maximum height of the displayed table in pixels. + max_memory_bytes + Maximum memory in bytes for rendered data. Helps prevent performance + issues with large datasets. + min_rows + Minimum number of rows to display even if memory limit is reached. + Must not exceed ``max_rows``. + max_rows + Maximum number of rows to display. Takes precedence over memory limits + when fewer rows are requested. + repr_rows + Deprecated alias for ``max_rows``. Use ``max_rows`` instead. + enable_cell_expansion + Whether to allow cells to expand when clicked. + custom_css + Custom CSS to apply to the HTML table. + show_truncation_message + Whether to show a message indicating that content has been truncated. + style_provider + Provider of CSS styles for the HTML table. If None, DefaultStyleProvider + is used. + use_shared_styles + Whether to use shared styles across multiple tables. This improves + performance when displaying many DataFrames in a single notebook. + + Raises: + ------ + ValueError + If max_cell_length, max_width, max_height, max_memory_bytes, + min_rows or max_rows is not a positive integer, or if min_rows + exceeds max_rows. + TypeError + If enable_cell_expansion, show_truncation_message, or use_shared_styles is + not a boolean, or if custom_css is provided but is not a string, or if + style_provider is provided but does not implement the StyleProvider + protocol. + """ + # Validate all parameters and get resolved max_rows + resolved_max_rows = _validate_formatter_parameters( + max_cell_length, + max_width, + max_height, + max_memory_bytes, + min_rows, + max_rows, + repr_rows, + enable_cell_expansion, + show_truncation_message, + use_shared_styles, + custom_css, + style_provider, + ) + + self.max_cell_length = max_cell_length + self.max_width = max_width + self.max_height = max_height + self.max_memory_bytes = max_memory_bytes + self.min_rows = min_rows + self._max_rows = resolved_max_rows + self.enable_cell_expansion = enable_cell_expansion + self.custom_css = custom_css + self.show_truncation_message = show_truncation_message + self.style_provider = style_provider or DefaultStyleProvider() + self.use_shared_styles = use_shared_styles + # Registry for custom type formatters + self._type_formatters: dict[type, CellFormatter] = {} + # Custom cell builders + self._custom_cell_builder: Callable[[Any, int, int, str], str] | None = None + self._custom_header_builder: Callable[[Any], str] | None = None + + @property + def max_rows(self) -> int: + """Get the maximum number of rows to display. + + Returns: + The maximum number of rows to display in repr output + """ + return self._max_rows + + @max_rows.setter + def max_rows(self, value: int) -> None: + """Set the maximum number of rows to display. + + Args: + value: The maximum number of rows + """ + self._max_rows = value + + @property + def repr_rows(self) -> int: + """Get the maximum number of rows (deprecated name). + + .. deprecated:: + Use :attr:`max_rows` instead. This property is provided for + backward compatibility. + + Returns: + The maximum number of rows to display + """ + return self._max_rows + + @repr_rows.setter + def repr_rows(self, value: int) -> None: + """Set the maximum number of rows using deprecated name. + + .. deprecated:: + Use :attr:`max_rows` setter instead. This property is provided for + backward compatibility. + + Args: + value: The maximum number of rows + """ + warnings.warn( + "repr_rows is deprecated, use max_rows instead", + DeprecationWarning, + stacklevel=2, + ) + self._max_rows = value + + def register_formatter(self, type_class: type, formatter: CellFormatter) -> None: + """Register a custom formatter for a specific data type. + + Args: + type_class: The type to register a formatter for + formatter: Function that takes a value of the given type and returns + a formatted string + """ + self._type_formatters[type_class] = formatter + + def set_custom_cell_builder( + self, builder: Callable[[Any, int, int, str], str] + ) -> None: + """Set a custom cell builder function. + + Args: + builder: Function that takes (value, row, col, table_id) and returns HTML + """ + self._custom_cell_builder = builder + + def set_custom_header_builder(self, builder: Callable[[Any], str]) -> None: + """Set a custom header builder function. + + Args: + builder: Function that takes a field and returns HTML + """ + self._custom_header_builder = builder + + def format_html( + self, + batches: list, + schema: Any, + has_more: bool = False, + table_uuid: str | None = None, + ) -> str: + """Format record batches as HTML. + + This method is used by DataFrame's _repr_html_ implementation and can be + called directly when custom HTML rendering is needed. + + Args: + batches: List of Arrow RecordBatch objects + schema: Arrow Schema object + has_more: Whether there are more batches not shown + table_uuid: Unique ID for the table, used for JavaScript interactions + + Returns: + HTML string representation of the data + + Raises: + TypeError: If schema is invalid and no batches are provided + """ + if not batches: + return "No data to display" + + # Validate schema + if schema is None or not hasattr(schema, "__iter__"): + msg = "Schema must be provided" + raise TypeError(msg) + + # Generate a unique ID if none provided + table_uuid = table_uuid or f"df-{id(batches)}" + + # Build HTML components + html = [] + + html.extend(self._build_html_header()) + + html.extend(self._build_table_container_start()) + + # Add table header and body + html.extend(self._build_table_header(schema)) + html.extend(self._build_table_body(batches, table_uuid)) + + html.append("") + html.append("") + + # Add footer (JavaScript and messages) + if self.enable_cell_expansion: + html.append(self._get_javascript()) + + # Always add truncation message if needed (independent of styles) + if has_more and self.show_truncation_message: + html.append("
Data truncated due to size.
") + + return "\n".join(html) + + def format_str( + self, + batches: list, + schema: Any, + has_more: bool = False, + table_uuid: str | None = None, + ) -> str: + """Format record batches as a string. + + This method is used by DataFrame's __repr__ implementation and can be + called directly when string rendering is needed. + + Args: + batches: List of Arrow RecordBatch objects + schema: Arrow Schema object + has_more: Whether there are more batches not shown + table_uuid: Unique ID for the table, used for JavaScript interactions + + Returns: + String representation of the data + + Raises: + TypeError: If schema is invalid and no batches are provided + """ + return DataFrameInternal.default_str_repr(batches, schema, has_more, table_uuid) + + def _build_html_header(self) -> list[str]: + """Build the HTML header with CSS styles.""" + default_css = self._get_default_css() if self.enable_cell_expansion else "" + script = f""" + +""" + html = [script] + if self.custom_css: + html.append(f"") + return html + + def _build_table_container_start(self) -> list[str]: + """Build the opening tags for the table container.""" + html = [] + html.append( + f'
' + ) + html.append('') + return html + + def _build_table_header(self, schema: Any) -> list[str]: + """Build the HTML table header with column names.""" + html = [] + html.append("") + html.append("") + for field in schema: + if self._custom_header_builder: + html.append(self._custom_header_builder(field)) + else: + html.append( + f"" + ) + html.append("") + html.append("") + return html + + def _build_table_body(self, batches: list, table_uuid: str) -> list[str]: + """Build the HTML table body with data rows.""" + html = [] + html.append("") + + row_count = 0 + for batch in batches: + for row_idx in range(batch.num_rows): + row_count += 1 + html.append("") + + for col_idx, column in enumerate(batch.columns): + # Get the raw value from the column + raw_value = self._get_cell_value(column, row_idx) + + # Always check for type formatters first to format the value + formatted_value = self._format_cell_value(raw_value) + + # Then apply either custom cell builder or standard cell formatting + if self._custom_cell_builder: + # Pass both the raw value and formatted value to let the + # builder decide + cell_html = self._custom_cell_builder( + raw_value, row_count, col_idx, table_uuid + ) + html.append(cell_html) + else: + # Standard cell formatting with formatted value + if ( + len(str(raw_value)) > self.max_cell_length + and self.enable_cell_expansion + ): + cell_html = self._build_expandable_cell( + formatted_value, row_count, col_idx, table_uuid + ) + else: + cell_html = self._build_regular_cell(formatted_value) + html.append(cell_html) + + html.append("") + + html.append("") + return html + + def _get_cell_value(self, column: Any, row_idx: int) -> Any: + """Extract a cell value from a column. + + Args: + column: Arrow array + row_idx: Row index + + Returns: + The raw cell value + """ + try: + value = column[row_idx] + + if hasattr(value, "as_py"): + return value.as_py() + except (AttributeError, TypeError): + pass + else: + return value + + def _format_cell_value(self, value: Any) -> str: + """Format a cell value for display. + + Uses registered type formatters if available. + + Args: + value: The cell value to format + + Returns: + Formatted cell value as string + """ + # Check for custom type formatters + for type_cls, formatter in self._type_formatters.items(): + if isinstance(value, type_cls): + return formatter(value) + + # If no formatter matched, return string representation + return str(value) + + def _build_expandable_cell( + self, formatted_value: str, row_count: int, col_idx: int, table_uuid: str + ) -> str: + """Build an expandable cell for long content.""" + short_value = str(formatted_value)[: self.max_cell_length] + return ( + f"" + ) + + def _build_regular_cell(self, formatted_value: str) -> str: + """Build a regular table cell.""" + return ( + f"" + ) + + def _build_html_footer(self, has_more: bool) -> list[str]: + """Build the HTML footer with JavaScript and messages.""" + html = [] + + # Add JavaScript for interactivity only if cell expansion is enabled + # and we're not using the shared styles approach + if self.enable_cell_expansion and not self.use_shared_styles: + html.append(self._get_javascript()) + + # Add truncation message if needed + if has_more and self.show_truncation_message: + html.append("
Data truncated due to size.
") + + return html + + def _get_default_css(self) -> str: + """Get default CSS styles for the HTML table.""" + return """ + .expandable-container { + display: inline-block; + max-width: 200px; + } + .expandable { + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + display: block; + } + .full-text { + display: none; + white-space: normal; + } + .expand-btn { + cursor: pointer; + color: blue; + text-decoration: underline; + border: none; + background: none; + font-size: inherit; + display: block; + margin-top: 5px; + } + """ + + def _get_javascript(self) -> str: + """Get JavaScript code for interactive elements.""" + return """ + +""" + + +class FormatterManager: + """Manager class for the global DataFrame HTML formatter instance.""" + + _default_formatter: DataFrameHtmlFormatter = DataFrameHtmlFormatter() + + @classmethod + def set_formatter(cls, formatter: DataFrameHtmlFormatter) -> None: + """Set the global DataFrame HTML formatter. + + Args: + formatter: The formatter instance to use globally + """ + cls._default_formatter = formatter + _refresh_formatter_reference() + + @classmethod + def get_formatter(cls) -> DataFrameHtmlFormatter: + """Get the current global DataFrame HTML formatter. + + Returns: + The global HTML formatter instance + """ + return cls._default_formatter + + +def get_formatter() -> DataFrameHtmlFormatter: + """Get the current global DataFrame HTML formatter. + + This function is used by the DataFrame._repr_html_ implementation to access + the shared formatter instance. It can also be used directly when custom + HTML rendering is needed. + + Returns: + The global HTML formatter instance + + Example: + >>> from datafusion.html_formatter import get_formatter + >>> formatter = get_formatter() + >>> formatter.max_cell_length = 50 # Increase cell length + """ + return FormatterManager.get_formatter() + + +def set_formatter(formatter: DataFrameHtmlFormatter) -> None: + """Set the global DataFrame HTML formatter. + + Args: + formatter: The formatter instance to use globally + + Example: + >>> from datafusion.html_formatter import get_formatter, set_formatter + >>> custom_formatter = DataFrameHtmlFormatter(max_cell_length=100) + >>> set_formatter(custom_formatter) + """ + FormatterManager.set_formatter(formatter) + + +def configure_formatter(**kwargs: Any) -> None: + """Configure the global DataFrame HTML formatter. + + This function creates a new formatter with the provided configuration + and sets it as the global formatter for all DataFrames. + + Args: + **kwargs: Formatter configuration parameters like max_cell_length, + max_width, max_height, enable_cell_expansion, etc. + + Raises: + ValueError: If any invalid parameters are provided + + Example: + >>> from datafusion.html_formatter import configure_formatter + >>> configure_formatter( + ... max_cell_length=50, + ... max_height=500, + ... enable_cell_expansion=True, + ... use_shared_styles=True + ... ) + """ + # Valid parameters accepted by DataFrameHtmlFormatter + valid_params = { + "max_cell_length", + "max_width", + "max_height", + "max_memory_bytes", + "min_rows", + "max_rows", + "repr_rows", + "enable_cell_expansion", + "custom_css", + "show_truncation_message", + "style_provider", + "use_shared_styles", + } + + # Check for invalid parameters + invalid_params = set(kwargs) - valid_params + if invalid_params: + msg = ( + f"Invalid formatter parameters: {', '.join(invalid_params)}. " + f"Valid parameters are: {', '.join(valid_params)}" + ) + raise ValueError(msg) + + # Create and set formatter with validated parameters + set_formatter(DataFrameHtmlFormatter(**kwargs)) + + +def reset_formatter() -> None: + """Reset the global DataFrame HTML formatter to default settings. + + This function creates a new formatter with default configuration + and sets it as the global formatter for all DataFrames. + + Example: + >>> from datafusion.html_formatter import reset_formatter + >>> reset_formatter() # Reset formatter to default settings + """ + formatter = DataFrameHtmlFormatter() + set_formatter(formatter) + + +def _refresh_formatter_reference() -> None: + """Refresh formatter reference in any modules using it. + + This helps ensure that changes to the formatter are reflected in existing + DataFrames that might be caching the formatter reference. + """ + # This is a no-op but signals modules to refresh their reference diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index 2697d8143..5760b8948 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -20,25 +20,38 @@ See :ref:`Expressions` in the online documentation for more details. """ -from __future__ import annotations +# ruff: noqa: PLC0415 -from typing import TYPE_CHECKING, Any, ClassVar, Optional +from __future__ import annotations -import pyarrow as pa +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING, Any, ClassVar try: from warnings import deprecated # Python 3.13+ except ImportError: from typing_extensions import deprecated # Python 3.12 -from datafusion.common import DataTypeMap, NullTreatment, RexType +import pyarrow as pa from ._internal import expr as expr_internal from ._internal import functions as functions_internal if TYPE_CHECKING: + from collections.abc import Sequence + + from datafusion.common import ( # type: ignore[import] + DataTypeMap, + NullTreatment, + RexType, + ) from datafusion.plan import LogicalPlan + +# Standard error message for invalid expression types +# Mention both alias forms of column and literal helpers +EXPR_TYPE_ERROR = "Use col()/column() or lit()/literal() to construct expressions" + # The following are imported from the internal representation. We may choose to # give these all proper wrappers, or to simply leave as is. These were added # in order to support passing the `test_imports` unit test. @@ -54,14 +67,29 @@ Case = expr_internal.Case Cast = expr_internal.Cast Column = expr_internal.Column +CopyTo = expr_internal.CopyTo +CreateCatalog = expr_internal.CreateCatalog +CreateCatalogSchema = expr_internal.CreateCatalogSchema +CreateExternalTable = expr_internal.CreateExternalTable +CreateFunction = expr_internal.CreateFunction +CreateFunctionBody = expr_internal.CreateFunctionBody +CreateIndex = expr_internal.CreateIndex CreateMemoryTable = expr_internal.CreateMemoryTable CreateView = expr_internal.CreateView +Deallocate = expr_internal.Deallocate +DescribeTable = expr_internal.DescribeTable Distinct = expr_internal.Distinct +DmlStatement = expr_internal.DmlStatement +DropCatalogSchema = expr_internal.DropCatalogSchema +DropFunction = expr_internal.DropFunction DropTable = expr_internal.DropTable +DropView = expr_internal.DropView EmptyRelation = expr_internal.EmptyRelation +Execute = expr_internal.Execute Exists = expr_internal.Exists Explain = expr_internal.Explain Extension = expr_internal.Extension +FileType = expr_internal.FileType Filter = expr_internal.Filter GroupingSet = expr_internal.GroupingSet Join = expr_internal.Join @@ -83,24 +111,35 @@ Literal = expr_internal.Literal Negative = expr_internal.Negative Not = expr_internal.Not +OperateFunctionArg = expr_internal.OperateFunctionArg Partitioning = expr_internal.Partitioning Placeholder = expr_internal.Placeholder +Prepare = expr_internal.Prepare Projection = expr_internal.Projection +RecursiveQuery = expr_internal.RecursiveQuery Repartition = expr_internal.Repartition ScalarSubquery = expr_internal.ScalarSubquery ScalarVariable = expr_internal.ScalarVariable +SetVariable = expr_internal.SetVariable SimilarTo = expr_internal.SimilarTo Sort = expr_internal.Sort Subquery = expr_internal.Subquery SubqueryAlias = expr_internal.SubqueryAlias TableScan = expr_internal.TableScan +TransactionAccessMode = expr_internal.TransactionAccessMode +TransactionConclusion = expr_internal.TransactionConclusion +TransactionEnd = expr_internal.TransactionEnd +TransactionIsolationLevel = expr_internal.TransactionIsolationLevel +TransactionStart = expr_internal.TransactionStart TryCast = expr_internal.TryCast Union = expr_internal.Union Unnest = expr_internal.Unnest UnnestExpr = expr_internal.UnnestExpr +Values = expr_internal.Values WindowExpr = expr_internal.WindowExpr __all__ = [ + "EXPR_TYPE_ERROR", "Aggregate", "AggregateFunction", "Alias", @@ -111,15 +150,30 @@ "CaseBuilder", "Cast", "Column", + "CopyTo", + "CreateCatalog", + "CreateCatalogSchema", + "CreateExternalTable", + "CreateFunction", + "CreateFunctionBody", + "CreateIndex", "CreateMemoryTable", "CreateView", + "Deallocate", + "DescribeTable", "Distinct", + "DmlStatement", + "DropCatalogSchema", + "DropFunction", "DropTable", + "DropView", "EmptyRelation", + "Execute", "Exists", "Explain", "Expr", "Extension", + "FileType", "Filter", "GroupingSet", "ILike", @@ -142,34 +196,125 @@ "Literal", "Negative", "Not", + "OperateFunctionArg", "Partitioning", "Placeholder", + "Prepare", "Projection", + "RecursiveQuery", "Repartition", "ScalarSubquery", "ScalarVariable", + "SetVariable", "SimilarTo", "Sort", "SortExpr", + "SortKey", "Subquery", "SubqueryAlias", "TableScan", + "TransactionAccessMode", + "TransactionConclusion", + "TransactionEnd", + "TransactionIsolationLevel", + "TransactionStart", "TryCast", "Union", "Unnest", "UnnestExpr", + "Values", "Window", "WindowExpr", "WindowFrame", "WindowFrameBound", + "ensure_expr", + "ensure_expr_list", ] +def ensure_expr(value: Expr | Any) -> expr_internal.Expr: + """Return the internal expression from ``Expr`` or raise ``TypeError``. + + This helper rejects plain strings and other non-:class:`Expr` values so + higher level APIs consistently require explicit :func:`~datafusion.col` or + :func:`~datafusion.lit` expressions. + + Args: + value: Candidate expression or other object. + + Returns: + The internal expression representation. + + Raises: + TypeError: If ``value`` is not an instance of :class:`Expr`. + """ + if not isinstance(value, Expr): + raise TypeError(EXPR_TYPE_ERROR) + return value.expr + + +def ensure_expr_list( + exprs: Iterable[Expr | Iterable[Expr]], +) -> list[expr_internal.Expr]: + """Flatten an iterable of expressions, validating each via ``ensure_expr``. + + Args: + exprs: Possibly nested iterable containing expressions. + + Returns: + A flat list of raw expressions. + + Raises: + TypeError: If any item is not an instance of :class:`Expr`. + """ + + def _iter( + items: Iterable[Expr | Iterable[Expr]], + ) -> Iterable[expr_internal.Expr]: + for expr in items: + if isinstance(expr, Iterable) and not isinstance( + expr, Expr | str | bytes | bytearray + ): + # Treat string-like objects as atomic to surface standard errors + yield from _iter(expr) + else: + yield ensure_expr(expr) + + return list(_iter(exprs)) + + +def _to_raw_expr(value: Expr | str) -> expr_internal.Expr: + """Convert a Python expression or column name to its raw variant. + + Args: + value: Candidate expression or column name. + + Returns: + The internal :class:`~datafusion._internal.expr.Expr` representation. + + Raises: + TypeError: If ``value`` is neither an :class:`Expr` nor ``str``. + """ + if isinstance(value, str): + return Expr.column(value).expr + if isinstance(value, Expr): + return value.expr + error = ( + "Expected Expr or column name, found:" + f" {type(value).__name__}. {EXPR_TYPE_ERROR}." + ) + raise TypeError(error) + + def expr_list_to_raw_expr_list( - expr_list: Optional[list[Expr]], -) -> Optional[list[expr_internal.Expr]]: - """Helper function to convert an optional list to raw expressions.""" - return [e.expr for e in expr_list] if expr_list is not None else None + expr_list: list[Expr] | Expr | None, +) -> list[expr_internal.Expr] | None: + """Convert a sequence of expressions or column names to raw expressions.""" + if isinstance(expr_list, Expr | str): + expr_list = [expr_list] + if expr_list is None: + return None + return [_to_raw_expr(e) for e in expr_list] def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr: @@ -180,10 +325,21 @@ def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr: def sort_list_to_raw_sort_list( - sort_list: Optional[list[Expr | SortExpr]], -) -> Optional[list[expr_internal.SortExpr]]: + sort_list: Sequence[SortKey] | SortKey | None, +) -> list[expr_internal.SortExpr] | None: """Helper function to return an optional sort list to raw variant.""" - return [sort_or_default(e) for e in sort_list] if sort_list is not None else None + if isinstance(sort_list, Expr | SortExpr | str): + sort_list = [sort_list] + if sort_list is None: + return None + raw_sort_list = [] + for item in sort_list: + if isinstance(item, SortExpr): + raw_sort_list.append(sort_or_default(item)) + else: + raw_expr = _to_raw_expr(item) # may raise ``TypeError`` + raw_sort_list.append(sort_or_default(Expr(raw_expr))) + return raw_sort_list class Expr: @@ -303,12 +459,39 @@ def __getitem__(self, key: str | int) -> Expr: If ``key`` is a string, returns the subfield of the struct. If ``key`` is an integer, retrieves the element in the array. Note that the - element index begins at ``0``, unlike `array_element` which begins at ``1``. + element index begins at ``0``, unlike + :py:func:`~datafusion.functions.array_element` which begins at ``1``. + If ``key`` is a slice, returns an array that contains a slice of the + original array. Similar to integer indexing, this follows Python convention + where the index begins at ``0`` unlike + :py:func:`~datafusion.functions.array_slice` which begins at ``1``. """ if isinstance(key, int): return Expr( functions_internal.array_element(self.expr, Expr.literal(key + 1).expr) ) + if isinstance(key, slice): + if isinstance(key.start, int): + start = Expr.literal(key.start + 1).expr + elif isinstance(key.start, Expr): + start = (key.start + Expr.literal(1)).expr + else: + # Default start at the first element, index 1 + start = Expr.literal(1).expr + + if isinstance(key.stop, int): + stop = Expr.literal(key.stop).expr + else: + stop = key.stop.expr + + if isinstance(key.step, int): + step = Expr.literal(key.step).expr + elif isinstance(key.step, Expr): + step = key.step.expr + else: + step = key.step + + return Expr(functions_internal.array_slice(self.expr, start, stop, step)) return Expr(self.expr.__getitem__(key)) def __eq__(self, rhs: object) -> Expr: @@ -381,10 +564,21 @@ def literal(value: Any) -> Expr: """ if isinstance(value, str): value = pa.scalar(value, type=pa.string_view()) - if not isinstance(value, pa.Scalar): - value = pa.scalar(value) return Expr(expr_internal.RawExpr.literal(value)) + @staticmethod + def literal_with_metadata(value: Any, metadata: dict[str, str]) -> Expr: + """Creates a new expression representing a scalar value with metadata. + + Args: + value: A valid PyArrow scalar value or easily castable to one. + metadata: Metadata to attach to the expression. + """ + if isinstance(value, str): + value = pa.scalar(value, type=pa.string_view()) + + return Expr(expr_internal.RawExpr.literal_with_metadata(value, metadata)) + @staticmethod def string_literal(value: str) -> Expr: """Creates a new expression representing a UTF8 literal value. @@ -406,9 +600,17 @@ def column(value: str) -> Expr: """Creates a new expression representing a column.""" return Expr(expr_internal.RawExpr.column(value)) - def alias(self, name: str) -> Expr: - """Assign a name to the expression.""" - return Expr(self.expr.alias(name)) + def alias(self, name: str, metadata: dict[str, str] | None = None) -> Expr: + """Assign a name to the expression. + + Args: + name: The name to assign to the expression. + metadata: Optional metadata to attach to the expression. + + Returns: + A new expression with the assigned name. + """ + return Expr(self.expr.alias(name, metadata)) def sort(self, ascending: bool = True, nulls_first: bool = True) -> SortExpr: """Creates a sort :py:class:`Expr` from an existing :py:class:`Expr`. @@ -446,7 +648,7 @@ def fill_null(self, value: Any | Expr | None = None) -> Expr: bool: pa.bool_(), } - def cast(self, to: pa.DataType[Any] | type[float | int | str | bool]) -> Expr: + def cast(self, to: pa.DataType[Any] | type) -> Expr: """Cast to a new data type.""" if not isinstance(to, pa.DataType): try: @@ -492,7 +694,7 @@ def types(self) -> DataTypeMap: return self.expr.types() def python_value(self) -> Any: - """Extracts the Expr value into a PyObject. + """Extracts the Expr value into `Any`. This is only valid for literal expressions. @@ -584,7 +786,7 @@ def over(self, window: Window) -> Expr: window: Window definition """ partition_by_raw = expr_list_to_raw_expr_list(window._partition_by) - order_by_raw = sort_list_to_raw_sort_list(window._order_by) + order_by_raw = window._order_by window_frame_raw = ( window._window_frame.window_frame if window._window_frame is not None @@ -603,6 +805,424 @@ def over(self, window: Window) -> Expr: ) ) + def asin(self) -> Expr: + """Returns the arc sine or inverse sine of a number.""" + from . import functions as F + + return F.asin(self) + + def array_pop_back(self) -> Expr: + """Returns the array without the last element.""" + from . import functions as F + + return F.array_pop_back(self) + + def reverse(self) -> Expr: + """Reverse the string argument.""" + from . import functions as F + + return F.reverse(self) + + def bit_length(self) -> Expr: + """Returns the number of bits in the string argument.""" + from . import functions as F + + return F.bit_length(self) + + def array_length(self) -> Expr: + """Returns the length of the array.""" + from . import functions as F + + return F.array_length(self) + + def array_ndims(self) -> Expr: + """Returns the number of dimensions of the array.""" + from . import functions as F + + return F.array_ndims(self) + + def to_hex(self) -> Expr: + """Converts an integer to a hexadecimal string.""" + from . import functions as F + + return F.to_hex(self) + + def array_dims(self) -> Expr: + """Returns an array of the array's dimensions.""" + from . import functions as F + + return F.array_dims(self) + + def from_unixtime(self) -> Expr: + """Converts an integer to RFC3339 timestamp format string.""" + from . import functions as F + + return F.from_unixtime(self) + + def array_empty(self) -> Expr: + """Returns a boolean indicating whether the array is empty.""" + from . import functions as F + + return F.array_empty(self) + + def sin(self) -> Expr: + """Returns the sine of the argument.""" + from . import functions as F + + return F.sin(self) + + def log10(self) -> Expr: + """Base 10 logarithm of the argument.""" + from . import functions as F + + return F.log10(self) + + def initcap(self) -> Expr: + """Set the initial letter of each word to capital. + + Converts the first letter of each word in ``string`` to uppercase and the + remaining characters to lowercase. + """ + from . import functions as F + + return F.initcap(self) + + def list_distinct(self) -> Expr: + """Returns distinct values from the array after removing duplicates. + + This is an alias for :py:func:`array_distinct`. + """ + from . import functions as F + + return F.list_distinct(self) + + def iszero(self) -> Expr: + """Returns true if a given number is +0.0 or -0.0 otherwise returns false.""" + from . import functions as F + + return F.iszero(self) + + def array_distinct(self) -> Expr: + """Returns distinct values from the array after removing duplicates.""" + from . import functions as F + + return F.array_distinct(self) + + def arrow_typeof(self) -> Expr: + """Returns the Arrow type of the expression.""" + from . import functions as F + + return F.arrow_typeof(self) + + def length(self) -> Expr: + """The number of characters in the ``string``.""" + from . import functions as F + + return F.length(self) + + def lower(self) -> Expr: + """Converts a string to lowercase.""" + from . import functions as F + + return F.lower(self) + + def acos(self) -> Expr: + """Returns the arc cosine or inverse cosine of a number. + + Returns: + -------- + Expr + A new expression representing the arc cosine of the input expression. + """ + from . import functions as F + + return F.acos(self) + + def ascii(self) -> Expr: + """Returns the numeric code of the first character of the argument.""" + from . import functions as F + + return F.ascii(self) + + def sha384(self) -> Expr: + """Computes the SHA-384 hash of a binary string.""" + from . import functions as F + + return F.sha384(self) + + def isnan(self) -> Expr: + """Returns true if a given number is +NaN or -NaN otherwise returns false.""" + from . import functions as F + + return F.isnan(self) + + def degrees(self) -> Expr: + """Converts the argument from radians to degrees.""" + from . import functions as F + + return F.degrees(self) + + def cardinality(self) -> Expr: + """Returns the total number of elements in the array.""" + from . import functions as F + + return F.cardinality(self) + + def sha224(self) -> Expr: + """Computes the SHA-224 hash of a binary string.""" + from . import functions as F + + return F.sha224(self) + + def asinh(self) -> Expr: + """Returns inverse hyperbolic sine.""" + from . import functions as F + + return F.asinh(self) + + def flatten(self) -> Expr: + """Flattens an array of arrays into a single array.""" + from . import functions as F + + return F.flatten(self) + + def exp(self) -> Expr: + """Returns the exponential of the argument.""" + from . import functions as F + + return F.exp(self) + + def abs(self) -> Expr: + """Return the absolute value of a given number. + + Returns: + -------- + Expr + A new expression representing the absolute value of the input expression. + """ + from . import functions as F + + return F.abs(self) + + def btrim(self) -> Expr: + """Removes all characters, spaces by default, from both sides of a string.""" + from . import functions as F + + return F.btrim(self) + + def md5(self) -> Expr: + """Computes an MD5 128-bit checksum for a string expression.""" + from . import functions as F + + return F.md5(self) + + def octet_length(self) -> Expr: + """Returns the number of bytes of a string.""" + from . import functions as F + + return F.octet_length(self) + + def cosh(self) -> Expr: + """Returns the hyperbolic cosine of the argument.""" + from . import functions as F + + return F.cosh(self) + + def radians(self) -> Expr: + """Converts the argument from degrees to radians.""" + from . import functions as F + + return F.radians(self) + + def sqrt(self) -> Expr: + """Returns the square root of the argument.""" + from . import functions as F + + return F.sqrt(self) + + def character_length(self) -> Expr: + """Returns the number of characters in the argument.""" + from . import functions as F + + return F.character_length(self) + + def tanh(self) -> Expr: + """Returns the hyperbolic tangent of the argument.""" + from . import functions as F + + return F.tanh(self) + + def atan(self) -> Expr: + """Returns inverse tangent of a number.""" + from . import functions as F + + return F.atan(self) + + def rtrim(self) -> Expr: + """Removes all characters, spaces by default, from the end of a string.""" + from . import functions as F + + return F.rtrim(self) + + def atanh(self) -> Expr: + """Returns inverse hyperbolic tangent.""" + from . import functions as F + + return F.atanh(self) + + def list_dims(self) -> Expr: + """Returns an array of the array's dimensions. + + This is an alias for :py:func:`array_dims`. + """ + from . import functions as F + + return F.list_dims(self) + + def sha256(self) -> Expr: + """Computes the SHA-256 hash of a binary string.""" + from . import functions as F + + return F.sha256(self) + + def factorial(self) -> Expr: + """Returns the factorial of the argument.""" + from . import functions as F + + return F.factorial(self) + + def acosh(self) -> Expr: + """Returns inverse hyperbolic cosine.""" + from . import functions as F + + return F.acosh(self) + + def floor(self) -> Expr: + """Returns the nearest integer less than or equal to the argument.""" + from . import functions as F + + return F.floor(self) + + def ceil(self) -> Expr: + """Returns the nearest integer greater than or equal to argument.""" + from . import functions as F + + return F.ceil(self) + + def list_length(self) -> Expr: + """Returns the length of the array. + + This is an alias for :py:func:`array_length`. + """ + from . import functions as F + + return F.list_length(self) + + def upper(self) -> Expr: + """Converts a string to uppercase.""" + from . import functions as F + + return F.upper(self) + + def chr(self) -> Expr: + """Converts the Unicode code point to a UTF8 character.""" + from . import functions as F + + return F.chr(self) + + def ln(self) -> Expr: + """Returns the natural logarithm (base e) of the argument.""" + from . import functions as F + + return F.ln(self) + + def tan(self) -> Expr: + """Returns the tangent of the argument.""" + from . import functions as F + + return F.tan(self) + + def array_pop_front(self) -> Expr: + """Returns the array without the first element.""" + from . import functions as F + + return F.array_pop_front(self) + + def cbrt(self) -> Expr: + """Returns the cube root of a number.""" + from . import functions as F + + return F.cbrt(self) + + def sha512(self) -> Expr: + """Computes the SHA-512 hash of a binary string.""" + from . import functions as F + + return F.sha512(self) + + def char_length(self) -> Expr: + """The number of characters in the ``string``.""" + from . import functions as F + + return F.char_length(self) + + def list_ndims(self) -> Expr: + """Returns the number of dimensions of the array. + + This is an alias for :py:func:`array_ndims`. + """ + from . import functions as F + + return F.list_ndims(self) + + def trim(self) -> Expr: + """Removes all characters, spaces by default, from both sides of a string.""" + from . import functions as F + + return F.trim(self) + + def cos(self) -> Expr: + """Returns the cosine of the argument.""" + from . import functions as F + + return F.cos(self) + + def sinh(self) -> Expr: + """Returns the hyperbolic sine of the argument.""" + from . import functions as F + + return F.sinh(self) + + def empty(self) -> Expr: + """This is an alias for :py:func:`array_empty`.""" + from . import functions as F + + return F.empty(self) + + def ltrim(self) -> Expr: + """Removes all characters, spaces by default, from the beginning of a string.""" + from . import functions as F + + return F.ltrim(self) + + def signum(self) -> Expr: + """Returns the sign of the argument (-1, 0, +1).""" + from . import functions as F + + return F.signum(self) + + def log2(self) -> Expr: + """Base 2 logarithm of the argument.""" + from . import functions as F + + return F.log2(self) + + def cot(self) -> Expr: + """Returns the cotangent of the argument.""" + from . import functions as F + + return F.cot(self) + class ExprFuncBuilder: def __init__(self, builder: expr_internal.ExprFuncBuilder) -> None: @@ -650,10 +1270,10 @@ class Window: def __init__( self, - partition_by: Optional[list[Expr]] = None, - window_frame: Optional[WindowFrame] = None, - order_by: Optional[list[SortExpr | Expr]] = None, - null_treatment: Optional[NullTreatment] = None, + partition_by: list[Expr] | Expr | None = None, + window_frame: WindowFrame | None = None, + order_by: list[SortExpr | Expr | str] | Expr | SortExpr | str | None = None, + null_treatment: NullTreatment | None = None, ) -> None: """Construct a window definition. @@ -665,7 +1285,7 @@ def __init__( """ self._partition_by = partition_by self._window_frame = window_frame - self._order_by = order_by + self._order_by = sort_list_to_raw_sort_list(order_by) self._null_treatment = null_treatment @@ -673,7 +1293,7 @@ class WindowFrame: """Defines a window frame for performing window operations.""" def __init__( - self, units: str, start_bound: Optional[Any], end_bound: Optional[Any] + self, units: str, start_bound: Any | None, end_bound: Any | None ) -> None: """Construct a window frame using the given parameters. @@ -696,6 +1316,10 @@ def __init__( end_bound = end_bound.cast(pa.uint64()) self.window_frame = expr_internal.WindowFrame(units, start_bound, end_bound) + def __repr__(self) -> str: + """Print a string representation of the window frame.""" + return self.window_frame.__repr__() + def get_frame_units(self) -> str: """Returns the window frame units for the bounds.""" return self.window_frame.get_frame_units() @@ -748,7 +1372,7 @@ class CaseBuilder: import datafusion.functions as f from datafusion import lit, col df.select( - f.case(col("column_a") + f.case(col("column_a")) .when(lit(1), lit("One")) .when(lit(2), lit("Two")) .otherwise(lit("Unknown")) @@ -801,3 +1425,6 @@ def nulls_first(self) -> bool: def __repr__(self) -> str: """Generate a string representation of this expression.""" return self.raw_sort.__repr__() + + +SortKey = Expr | SortExpr | str diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 5cf914e16..e85d710e7 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any import pyarrow as pa @@ -28,14 +28,20 @@ CaseBuilder, Expr, SortExpr, + SortKey, WindowFrame, expr_list_to_raw_expr_list, sort_list_to_raw_sort_list, + sort_or_default, ) +try: + from warnings import deprecated # Python 3.13+ +except ImportError: + from typing_extensions import deprecated # Python 3.12 + if TYPE_CHECKING: from datafusion.context import SessionContext - __all__ = [ "abs", "acos", @@ -218,6 +224,7 @@ "range", "rank", "regexp_count", + "regexp_instr", "regexp_like", "regexp_match", "regexp_replace", @@ -260,13 +267,18 @@ "sum", "tan", "tanh", + "to_char", + "to_date", "to_hex", + "to_local_time", + "to_time", "to_timestamp", "to_timestamp_micros", "to_timestamp_millis", "to_timestamp_nanos", "to_timestamp_seconds", "to_unixtime", + "today", "translate", "trim", "trunc", @@ -283,7 +295,15 @@ def isnan(expr: Expr) -> Expr: - """Returns true if a given number is +NaN or -NaN otherwise returns false.""" + """Returns true if a given number is +NaN or -NaN otherwise returns false. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, np.nan]}) + >>> result = df.select(dfn.functions.isnan(dfn.col("a")).alias("isnan")) + >>> result.collect_column("isnan")[1].as_py() + True + """ return Expr(f.isnan(expr.expr)) @@ -291,29 +311,65 @@ def nullif(expr1: Expr, expr2: Expr) -> Expr: """Returns NULL if expr1 equals expr2; otherwise it returns expr1. This can be used to perform the inverse operation of the COALESCE expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2], "b": [1, 3]}) + >>> result = df.select( + ... dfn.functions.nullif(dfn.col("a"), dfn.col("b")).alias("nullif")) + >>> result.collect_column("nullif").to_pylist() + [None, 2] """ return Expr(f.nullif(expr1.expr, expr2.expr)) def encode(expr: Expr, encoding: Expr) -> Expr: - """Encode the ``input``, using the ``encoding``. encoding can be base64 or hex.""" + """Encode the ``input``, using the ``encoding``. encoding can be base64 or hex. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.encode(dfn.col("a"), dfn.lit("base64")).alias("enc")) + >>> result.collect_column("enc")[0].as_py() + 'aGVsbG8' + """ return Expr(f.encode(expr.expr, encoding.expr)) def decode(expr: Expr, encoding: Expr) -> Expr: - """Decode the ``input``, using the ``encoding``. encoding can be base64 or hex.""" + """Decode the ``input``, using the ``encoding``. encoding can be base64 or hex. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["aGVsbG8="]}) + >>> result = df.select( + ... dfn.functions.decode(dfn.col("a"), dfn.lit("base64")).alias("dec")) + >>> result.collect_column("dec")[0].as_py() + b'hello' + """ return Expr(f.decode(expr.expr, encoding.expr)) def array_to_string(expr: Expr, delimiter: Expr) -> Expr: - """Converts each element to its text representation.""" + """Converts each element to its text representation. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_to_string(dfn.col("a"), dfn.lit(",")).alias("s")) + >>> result.collect_column("s")[0].as_py() + '1,2,3' + """ return Expr(f.array_to_string(expr.expr, delimiter.expr.cast(pa.string()))) def array_join(expr: Expr, delimiter: Expr) -> Expr: """Converts each element to its text representation. - This is an alias for :py:func:`array_to_string`. + See Also: + This is an alias for :py:func:`array_to_string`. """ return array_to_string(expr, delimiter) @@ -321,7 +377,8 @@ def array_join(expr: Expr, delimiter: Expr) -> Expr: def list_to_string(expr: Expr, delimiter: Expr) -> Expr: """Converts each element to its text representation. - This is an alias for :py:func:`array_to_string`. + See Also: + This is an alias for :py:func:`array_to_string`. """ return array_to_string(expr, delimiter) @@ -330,12 +387,27 @@ def list_join(expr: Expr, delimiter: Expr) -> Expr: """Converts each element to its text representation. This is an alias for :py:func:`array_to_string`. + + See Also: + This is an alias for :py:func:`array_to_string`. """ return array_to_string(expr, delimiter) def in_list(arg: Expr, values: list[Expr], negated: bool = False) -> Expr: - """Returns whether the argument is contained within the list ``values``.""" + """Returns whether the argument is contained within the list ``values``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.select( + ... dfn.functions.in_list( + ... dfn.col("a"), [dfn.lit(1), dfn.lit(3)] + ... ).alias("in") + ... ) + >>> result.collect_column("in").to_pylist() + [True, False, True] + """ values = [v.expr for v in values] return Expr(f.in_list(arg.expr, values, negated)) @@ -345,6 +417,14 @@ def digest(value: Expr, method: Expr) -> Expr: Standard algorithms are md5, sha224, sha256, sha384, sha512, blake2s, blake2b, and blake3. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.digest(dfn.col("a"), dfn.lit("md5")).alias("d")) + >>> len(result.collect_column("d")[0].as_py()) > 0 + True """ return Expr(f.digest(value.expr, method.expr)) @@ -353,6 +433,15 @@ def concat(*args: Expr) -> Expr: """Concatenates the text representations of all the arguments. NULL arguments are ignored. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"], "b": [" world"]}) + >>> result = df.select( + ... dfn.functions.concat(dfn.col("a"), dfn.col("b")).alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() + 'hello world' """ args = [arg.expr for arg in args] return Expr(f.concat(args)) @@ -362,27 +451,62 @@ def concat_ws(separator: str, *args: Expr) -> Expr: """Concatenates the list ``args`` with the separator. ``NULL`` arguments are ignored. ``separator`` should not be ``NULL``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"], "b": ["world"]}) + >>> result = df.select( + ... dfn.functions.concat_ws("-", dfn.col("a"), dfn.col("b")).alias("c")) + >>> result.collect_column("c")[0].as_py() + 'hello-world' """ args = [arg.expr for arg in args] return Expr(f.concat_ws(separator, args)) def order_by(expr: Expr, ascending: bool = True, nulls_first: bool = True) -> SortExpr: - """Creates a new sort expression.""" + """Creates a new sort expression. + + Examples: + >>> sort_expr = dfn.functions.order_by(dfn.col("a"), ascending=False) + >>> sort_expr.ascending() + False + """ return SortExpr(expr, ascending=ascending, nulls_first=nulls_first) -def alias(expr: Expr, name: str) -> Expr: - """Creates an alias expression.""" - return Expr(f.alias(expr.expr, name)) +def alias(expr: Expr, name: str, metadata: dict[str, str] | None = None) -> Expr: + """Creates an alias expression with an optional metadata dictionary. + + Args: + expr: The expression to alias + name: The alias name + metadata: Optional metadata to attach to the column + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2]}) + >>> df.select( + ... dfn.functions.alias(dfn.col("a"), "b") + ... ).collect_column("b")[0].as_py() + 1 + """ + return Expr(f.alias(expr.expr, name, metadata)) def col(name: str) -> Expr: - """Creates a column reference expression.""" + """Creates a column reference expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.select(dfn.functions.col("a")).collect_column("a")[0].as_py() + 1 + """ return Expr(f.col(name)) -def count_star(filter: Optional[Expr] = None) -> Expr: +def count_star(filter: Expr | None = None) -> Expr: """Create a COUNT(1) aggregate expression. This aggregate function will count all of the rows in the partition. @@ -392,6 +516,13 @@ def count_star(filter: Optional[Expr] = None) -> Expr: Args: filter: If provided, only count rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.count_star().alias("cnt")]) + >>> result.collect_column("cnt")[0].as_py() + 3 """ return count(Expr.literal(1), filter=filter) @@ -402,6 +533,15 @@ def case(expr: Expr) -> CaseBuilder: Create a :py:class:`~datafusion.expr.CaseBuilder` to match cases for the expression ``expr``. See :py:class:`~datafusion.expr.CaseBuilder` for detailed usage. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.select( + ... dfn.functions.case(dfn.col("a")).when(dfn.lit(1), + ... dfn.lit("one")).otherwise(dfn.lit("other")).alias("c")) + >>> result.collect_column("c")[0].as_py() + 'one' """ return CaseBuilder(f.case(expr.expr)) @@ -412,16 +552,28 @@ def when(when: Expr, then: Expr) -> CaseBuilder: Create a :py:class:`~datafusion.expr.CaseBuilder` to match cases for the expression ``expr``. See :py:class:`~datafusion.expr.CaseBuilder` for detailed usage. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.select( + ... dfn.functions.when(dfn.col("a") > dfn.lit(2), + ... dfn.lit("big")).otherwise(dfn.lit("small")).alias("c")) + >>> result.collect_column("c")[2].as_py() + 'big' """ return CaseBuilder(f.when(when.expr, then.expr)) +@deprecated("Prefer to call Expr.over() instead") def window( name: str, args: list[Expr], - partition_by: list[Expr] | None = None, - order_by: list[Expr | SortExpr] | None = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, window_frame: WindowFrame | None = None, + filter: Expr | None = None, + distinct: bool = False, ctx: SessionContext | None = None, ) -> Expr: """Creates a new Window function expression. @@ -431,23 +583,41 @@ def window( lag use:: df.select(functions.lag(col("a")).partition_by(col("b")).build()) + + The ``order_by`` parameter accepts column names or expressions, e.g.:: + + window("lag", [col("a")], order_by="ts") """ args = [a.expr for a in args] - partition_by = expr_list_to_raw_expr_list(partition_by) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) window_frame = window_frame.window_frame if window_frame is not None else None ctx = ctx.ctx if ctx is not None else None - return Expr(f.window(name, args, partition_by, order_by_raw, window_frame, ctx)) + filter_raw = filter.expr if filter is not None else None + return Expr( + f.window( + name, + args, + partition_by=partition_by_raw, + order_by=order_by_raw, + window_frame=window_frame, + ctx=ctx, + filter=filter_raw, + distinct=distinct, + ) + ) # scalar functions def abs(arg: Expr) -> Expr: """Return the absolute value of a given number. - Returns: - -------- - Expr - A new expression representing the absolute value of the input expression. + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [-1, 0, 1]}) + >>> result = df.select(dfn.functions.abs(dfn.col("a")).alias("abs")) + >>> result.collect_column("abs")[0].as_py() + 1 """ return Expr(f.abs(arg.expr)) @@ -455,127 +625,331 @@ def abs(arg: Expr) -> Expr: def acos(arg: Expr) -> Expr: """Returns the arc cosine or inverse cosine of a number. - Returns: - -------- - Expr - A new expression representing the arc cosine of the input expression. + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0]}) + >>> result = df.select(dfn.functions.acos(dfn.col("a")).alias("acos")) + >>> result.collect_column("acos")[0].as_py() + 0.0 """ return Expr(f.acos(arg.expr)) def acosh(arg: Expr) -> Expr: - """Returns inverse hyperbolic cosine.""" + """Returns inverse hyperbolic cosine. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0]}) + >>> result = df.select(dfn.functions.acosh(dfn.col("a")).alias("acosh")) + >>> result.collect_column("acosh")[0].as_py() + 0.0 + """ return Expr(f.acosh(arg.expr)) def ascii(arg: Expr) -> Expr: - """Returns the numeric code of the first character of the argument.""" + """Returns the numeric code of the first character of the argument. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["a","b","c"]}) + >>> ascii_df = df.select(dfn.functions.ascii(dfn.col("a")).alias("ascii")) + >>> ascii_df.collect_column("ascii")[0].as_py() + 97 + """ return Expr(f.ascii(arg.expr)) def asin(arg: Expr) -> Expr: - """Returns the arc sine or inverse sine of a number.""" + """Returns the arc sine or inverse sine of a number. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.asin(dfn.col("a")).alias("asin")) + >>> result.collect_column("asin")[0].as_py() + 0.0 + """ return Expr(f.asin(arg.expr)) def asinh(arg: Expr) -> Expr: - """Returns inverse hyperbolic sine.""" + """Returns inverse hyperbolic sine. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.asinh(dfn.col("a")).alias("asinh")) + >>> result.collect_column("asinh")[0].as_py() + 0.0 + """ return Expr(f.asinh(arg.expr)) def atan(arg: Expr) -> Expr: - """Returns inverse tangent of a number.""" + """Returns inverse tangent of a number. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.atan(dfn.col("a")).alias("atan")) + >>> result.collect_column("atan")[0].as_py() + 0.0 + """ return Expr(f.atan(arg.expr)) def atanh(arg: Expr) -> Expr: - """Returns inverse hyperbolic tangent.""" + """Returns inverse hyperbolic tangent. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.atanh(dfn.col("a")).alias("atanh")) + >>> result.collect_column("atanh")[0].as_py() + 0.0 + """ return Expr(f.atanh(arg.expr)) def atan2(y: Expr, x: Expr) -> Expr: - """Returns inverse tangent of a division given in the argument.""" + """Returns inverse tangent of a division given in the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [0.0], "x": [1.0]}) + >>> result = df.select( + ... dfn.functions.atan2(dfn.col("y"), dfn.col("x")).alias("atan2")) + >>> result.collect_column("atan2")[0].as_py() + 0.0 + """ return Expr(f.atan2(y.expr, x.expr)) def bit_length(arg: Expr) -> Expr: - """Returns the number of bits in the string argument.""" + """Returns the number of bits in the string argument. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["a","b","c"]}) + >>> bit_df = df.select(dfn.functions.bit_length(dfn.col("a")).alias("bit_len")) + >>> bit_df.collect_column("bit_len")[0].as_py() + 8 + """ return Expr(f.bit_length(arg.expr)) def btrim(arg: Expr) -> Expr: - """Removes all characters, spaces by default, from both sides of a string.""" + """Removes all characters, spaces by default, from both sides of a string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [" a "]}) + >>> trim_df = df.select(dfn.functions.btrim(dfn.col("a")).alias("trimmed")) + >>> trim_df.collect_column("trimmed")[0].as_py() + 'a' + """ return Expr(f.btrim(arg.expr)) def cbrt(arg: Expr) -> Expr: - """Returns the cube root of a number.""" + """Returns the cube root of a number. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [27]}) + >>> cbrt_df = df.select(dfn.functions.cbrt(dfn.col("a")).alias("cbrt")) + >>> cbrt_df.collect_column("cbrt")[0].as_py() + 3.0 + """ return Expr(f.cbrt(arg.expr)) def ceil(arg: Expr) -> Expr: - """Returns the nearest integer greater than or equal to argument.""" + """Returns the nearest integer greater than or equal to argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.9]}) + >>> ceil_df = df.select(dfn.functions.ceil(dfn.col("a")).alias("ceil")) + >>> ceil_df.collect_column("ceil")[0].as_py() + 2.0 + """ return Expr(f.ceil(arg.expr)) def character_length(arg: Expr) -> Expr: - """Returns the number of characters in the argument.""" + """Returns the number of characters in the argument. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["abc","b","c"]}) + >>> char_len_df = df.select( + ... dfn.functions.character_length(dfn.col("a")).alias("char_len")) + >>> char_len_df.collect_column("char_len")[0].as_py() + 3 + """ return Expr(f.character_length(arg.expr)) def length(string: Expr) -> Expr: - """The number of characters in the ``string``.""" + """The number of characters in the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.length(dfn.col("a")).alias("len")) + >>> result.collect_column("len")[0].as_py() + 5 + """ return Expr(f.length(string.expr)) def char_length(string: Expr) -> Expr: - """The number of characters in the ``string``.""" + """The number of characters in the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.char_length(dfn.col("a")).alias("len")) + >>> result.collect_column("len")[0].as_py() + 5 + """ return Expr(f.char_length(string.expr)) def chr(arg: Expr) -> Expr: - """Converts the Unicode code point to a UTF8 character.""" + """Converts the Unicode code point to a UTF8 character. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [65]}) + >>> result = df.select(dfn.functions.chr(dfn.col("a")).alias("chr")) + >>> result.collect_column("chr")[0].as_py() + 'A' + """ return Expr(f.chr(arg.expr)) def coalesce(*args: Expr) -> Expr: - """Returns the value of the first expr in ``args`` which is not NULL.""" + """Returns the value of the first expr in ``args`` which is not NULL. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [None, 1], "b": [2, 3]}) + >>> result = df.select( + ... dfn.functions.coalesce(dfn.col("a"), dfn.col("b")).alias("c")) + >>> result.collect_column("c")[0].as_py() + 2 + """ args = [arg.expr for arg in args] return Expr(f.coalesce(*args)) def cos(arg: Expr) -> Expr: - """Returns the cosine of the argument.""" + """Returns the cosine of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0,-1,1]}) + >>> cos_df = df.select(dfn.functions.cos(dfn.col("a")).alias("cos")) + >>> cos_df.collect_column("cos")[0].as_py() + 1.0 + """ return Expr(f.cos(arg.expr)) def cosh(arg: Expr) -> Expr: - """Returns the hyperbolic cosine of the argument.""" + """Returns the hyperbolic cosine of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0,-1,1]}) + >>> cosh_df = df.select(dfn.functions.cosh(dfn.col("a")).alias("cosh")) + >>> cosh_df.collect_column("cosh")[0].as_py() + 1.0 + """ return Expr(f.cosh(arg.expr)) def cot(arg: Expr) -> Expr: - """Returns the cotangent of the argument.""" + """Returns the cotangent of the argument. + + Examples: + >>> from math import pi + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [pi / 4]}) + >>> result = df.select( + ... dfn.functions.cot(dfn.col("a")).alias("cot") + ... ) + >>> result.collect_column("cot")[0].as_py() + 1.0... + """ return Expr(f.cot(arg.expr)) def degrees(arg: Expr) -> Expr: - """Converts the argument from radians to degrees.""" + """Converts the argument from radians to degrees. + + Examples: + >>> from math import pi + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0,pi,2*pi]}) + >>> deg_df = df.select(dfn.functions.degrees(dfn.col("a")).alias("deg")) + >>> deg_df.collect_column("deg")[2].as_py() + 360.0 + """ return Expr(f.degrees(arg.expr)) def ends_with(arg: Expr, suffix: Expr) -> Expr: - """Returns true if the ``string`` ends with the ``suffix``, false otherwise.""" + """Returns true if the ``string`` ends with the ``suffix``, false otherwise. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["abc","b","c"]}) + >>> ends_with_df = df.select( + ... dfn.functions.ends_with(dfn.col("a"), dfn.lit("c")).alias("ends_with")) + >>> ends_with_df.collect_column("ends_with")[0].as_py() + True + """ return Expr(f.ends_with(arg.expr, suffix.expr)) def exp(arg: Expr) -> Expr: - """Returns the exponential of the argument.""" + """Returns the exponential of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.exp(dfn.col("a")).alias("exp")) + >>> result.collect_column("exp")[0].as_py() + 1.0 + """ return Expr(f.exp(arg.expr)) def factorial(arg: Expr) -> Expr: - """Returns the factorial of the argument.""" + """Returns the factorial of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [3]}) + >>> result = df.select( + ... dfn.functions.factorial(dfn.col("a")).alias("factorial") + ... ) + >>> result.collect_column("factorial")[0].as_py() + 6 + """ return Expr(f.factorial(arg.expr)) @@ -586,17 +960,44 @@ def find_in_set(string: Expr, string_list: Expr) -> Expr: ``string_list`` consisting of N substrings. The string list is a string composed of substrings separated by ``,`` characters. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["b"]}) + >>> result = df.select( + ... dfn.functions.find_in_set(dfn.col("a"), dfn.lit("a,b,c")).alias("pos")) + >>> result.collect_column("pos")[0].as_py() + 2 """ return Expr(f.find_in_set(string.expr, string_list.expr)) def floor(arg: Expr) -> Expr: - """Returns the nearest integer less than or equal to the argument.""" + """Returns the nearest integer less than or equal to the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.9]}) + >>> floor_df = df.select(dfn.functions.floor(dfn.col("a")).alias("floor")) + >>> floor_df.collect_column("floor")[0].as_py() + 1.0 + """ return Expr(f.floor(arg.expr)) def gcd(x: Expr, y: Expr) -> Expr: - """Returns the greatest common divisor.""" + """Returns the greatest common divisor. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [12], "b": [8]}) + >>> result = df.select( + ... dfn.functions.gcd(dfn.col("a"), dfn.col("b")).alias("gcd") + ... ) + >>> result.collect_column("gcd")[0].as_py() + 4 + """ return Expr(f.gcd(x.expr, y.expr)) @@ -605,6 +1006,14 @@ def initcap(string: Expr) -> Expr: Converts the first letter of each word in ``string`` to uppercase and the remaining characters to lowercase. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["the cat"]}) + >>> cap_df = df.select(dfn.functions.initcap(dfn.col("a")).alias("cap")) + >>> cap_df.collect_column("cap")[0].as_py() + 'The Cat' """ return Expr(f.initcap(string.expr)) @@ -618,47 +1027,127 @@ def instr(string: Expr, substring: Expr) -> Expr: def iszero(arg: Expr) -> Expr: - """Returns true if a given number is +0.0 or -0.0 otherwise returns false.""" + """Returns true if a given number is +0.0 or -0.0 otherwise returns false. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0, 1.0]}) + >>> result = df.select(dfn.functions.iszero(dfn.col("a")).alias("iz")) + >>> result.collect_column("iz")[0].as_py() + True + """ return Expr(f.iszero(arg.expr)) def lcm(x: Expr, y: Expr) -> Expr: - """Returns the least common multiple.""" + """Returns the least common multiple. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [4], "b": [6]}) + >>> result = df.select( + ... dfn.functions.lcm(dfn.col("a"), dfn.col("b")).alias("lcm") + ... ) + >>> result.collect_column("lcm")[0].as_py() + 12 + """ return Expr(f.lcm(x.expr, y.expr)) def left(string: Expr, n: Expr) -> Expr: - """Returns the first ``n`` characters in the ``string``.""" + """Returns the first ``n`` characters in the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["the cat"]}) + >>> left_df = df.select(dfn.functions.left(dfn.col("a"), dfn.lit(3)).alias("left")) + >>> left_df.collect_column("left")[0].as_py() + 'the' + """ return Expr(f.left(string.expr, n.expr)) def levenshtein(string1: Expr, string2: Expr) -> Expr: - """Returns the Levenshtein distance between the two given strings.""" + """Returns the Levenshtein distance between the two given strings. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["kitten"]}) + >>> result = df.select( + ... dfn.functions.levenshtein(dfn.col("a"), dfn.lit("sitting")).alias("d")) + >>> result.collect_column("d")[0].as_py() + 3 + """ return Expr(f.levenshtein(string1.expr, string2.expr)) def ln(arg: Expr) -> Expr: - """Returns the natural logarithm (base e) of the argument.""" + """Returns the natural logarithm (base e) of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0]}) + >>> result = df.select(dfn.functions.ln(dfn.col("a")).alias("ln")) + >>> result.collect_column("ln")[0].as_py() + 0.0 + """ return Expr(f.ln(arg.expr)) def log(base: Expr, num: Expr) -> Expr: - """Returns the logarithm of a number for a particular ``base``.""" + """Returns the logarithm of a number for a particular ``base``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [100.0]}) + >>> result = df.select( + ... dfn.functions.log(dfn.lit(10.0), dfn.col("a")).alias("log") + ... ) + >>> result.collect_column("log")[0].as_py() + 2.0 + """ return Expr(f.log(base.expr, num.expr)) def log10(arg: Expr) -> Expr: - """Base 10 logarithm of the argument.""" + """Base 10 logarithm of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [100.0]}) + >>> result = df.select(dfn.functions.log10(dfn.col("a")).alias("log10")) + >>> result.collect_column("log10")[0].as_py() + 2.0 + """ return Expr(f.log10(arg.expr)) def log2(arg: Expr) -> Expr: - """Base 2 logarithm of the argument.""" + """Base 2 logarithm of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [8.0]}) + >>> result = df.select(dfn.functions.log2(dfn.col("a")).alias("log2")) + >>> result.collect_column("log2")[0].as_py() + 3.0 + """ return Expr(f.log2(arg.expr)) def lower(arg: Expr) -> Expr: - """Converts a string to lowercase.""" + """Converts a string to lowercase. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["THE CaT"]}) + >>> lower_df = df.select(dfn.functions.lower(dfn.col("a")).alias("lower")) + >>> lower_df.collect_column("lower")[0].as_py() + 'the cat' + """ return Expr(f.lower(arg.expr)) @@ -668,33 +1157,92 @@ def lpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: Extends the string to length length by prepending the characters fill (a space by default). If the string is already longer than length then it is truncated (on the right). + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["the cat", "a hat"]}) + >>> lpad_df = df.select(dfn.functions.lpad(dfn.col("a"), dfn.lit(6)).alias("lpad")) + >>> lpad_df.collect_column("lpad")[0].as_py() + 'the ca' + >>> lpad_df.collect_column("lpad")[1].as_py() + ' a hat' """ characters = characters if characters is not None else Expr.literal(" ") return Expr(f.lpad(string.expr, count.expr, characters.expr)) def ltrim(arg: Expr) -> Expr: - """Removes all characters, spaces by default, from the beginning of a string.""" + """Removes all characters, spaces by default, from the beginning of a string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [" a "]}) + >>> trim_df = df.select(dfn.functions.ltrim(dfn.col("a")).alias("trimmed")) + >>> trim_df.collect_column("trimmed")[0].as_py() + 'a ' + """ return Expr(f.ltrim(arg.expr)) def md5(arg: Expr) -> Expr: - """Computes an MD5 128-bit checksum for a string expression.""" + """Computes an MD5 128-bit checksum for a string expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.md5(dfn.col("a")).alias("md5")) + >>> result.collect_column("md5")[0].as_py() + '5d41402abc4b2a76b9719d911017c592' + """ return Expr(f.md5(arg.expr)) def nanvl(x: Expr, y: Expr) -> Expr: - """Returns ``x`` if ``x`` is not ``NaN``. Otherwise returns ``y``.""" + """Returns ``x`` if ``x`` is not ``NaN``. Otherwise returns ``y``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [np.nan, 1.0], "b": [0.0, 0.0]}) + >>> nanvl_df = df.select( + ... dfn.functions.nanvl(dfn.col("a"), dfn.col("b")).alias("nanvl")) + >>> nanvl_df.collect_column("nanvl")[0].as_py() + 0.0 + >>> nanvl_df.collect_column("nanvl")[1].as_py() + 1.0 + """ return Expr(f.nanvl(x.expr, y.expr)) def nvl(x: Expr, y: Expr) -> Expr: - """Returns ``x`` if ``x`` is not ``NULL``. Otherwise returns ``y``.""" + """Returns ``x`` if ``x`` is not ``NULL``. Otherwise returns ``y``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [None, 1], "b": [0, 0]}) + >>> nvl_df = df.select( + ... dfn.functions.nvl(dfn.col("a"), dfn.col("b")).alias("nvl") + ... ) + >>> nvl_df.collect_column("nvl")[0].as_py() + 0 + >>> nvl_df.collect_column("nvl")[1].as_py() + 1 + """ return Expr(f.nvl(x.expr, y.expr)) def octet_length(arg: Expr) -> Expr: - """Returns the number of bytes of a string.""" + """Returns the number of bytes of a string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.octet_length(dfn.col("a")).alias("len")) + >>> result.collect_column("len")[0].as_py() + 5 + """ return Expr(f.octet_length(arg.expr)) @@ -705,6 +1253,16 @@ def overlay( Replace the substring of string that starts at the ``start``'th character and extends for ``length`` characters with new substring. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["abcdef"]}) + >>> result = df.select( + ... dfn.functions.overlay(dfn.col("a"), dfn.lit("XY"), dfn.lit(3), + ... dfn.lit(2)).alias("o")) + >>> result.collect_column("o")[0].as_py() + 'abXYef' """ if length is None: return Expr(f.overlay(string.expr, substring.expr, start.expr)) @@ -712,7 +1270,20 @@ def overlay( def pi() -> Expr: - """Returns an approximate value of π.""" + """Returns an approximate value of π. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> import builtins + >>> result = df.select( + ... dfn.functions.pi().alias("pi") + ... ) + >>> builtins.round( + ... result.collect_column("pi")[0].as_py(), 5 + ... ) + 3.14159 + """ return Expr(f.pi()) @@ -725,7 +1296,17 @@ def position(string: Expr, substring: Expr) -> Expr: def power(base: Expr, exponent: Expr) -> Expr: - """Returns ``base`` raised to the power of ``exponent``.""" + """Returns ``base`` raised to the power of ``exponent``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [2.0]}) + >>> result = df.select( + ... dfn.functions.power(dfn.col("a"), dfn.lit(3.0)).alias("pow") + ... ) + >>> result.collect_column("pow")[0].as_py() + 8.0 + """ return Expr(f.power(base.expr, exponent.expr)) @@ -738,15 +1319,37 @@ def pow(base: Expr, exponent: Expr) -> Expr: def radians(arg: Expr) -> Expr: - """Converts the argument from degrees to radians.""" + """Converts the argument from degrees to radians. + + Examples: + >>> from math import pi + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [180.0]}) + >>> result = df.select( + ... dfn.functions.radians(dfn.col("a")).alias("rad") + ... ) + >>> result.collect_column("rad")[0].as_py() == pi + True + """ return Expr(f.radians(arg.expr)) def regexp_like(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: - """Find if any regular expression (regex) matches exist. + r"""Find if any regular expression (regex) matches exist. Tests a string using a regular expression returning true if at least one match, false otherwise. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello123"]}) + >>> result = df.select( + ... dfn.functions.regexp_like( + ... dfn.col("a"), dfn.lit("\\d+") + ... ).alias("m") + ... ) + >>> result.collect_column("m")[0].as_py() + True """ if flags is not None: flags = flags.expr @@ -754,10 +1357,21 @@ def regexp_like(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: def regexp_match(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: - """Perform regular expression (regex) matching. + r"""Perform regular expression (regex) matching. Returns an array with each element containing the leftmost-first match of the corresponding index in ``regex`` to string in ``string``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello 42 world"]}) + >>> result = df.select( + ... dfn.functions.regexp_match( + ... dfn.col("a"), dfn.lit("(\\d+)") + ... ).alias("m") + ... ) + >>> result.collect_column("m")[0].as_py() + ['42'] """ if flags is not None: flags = flags.expr @@ -767,13 +1381,25 @@ def regexp_match(string: Expr, regex: Expr, flags: Expr | None = None) -> Expr: def regexp_replace( string: Expr, pattern: Expr, replacement: Expr, flags: Expr | None = None ) -> Expr: - """Replaces substring(s) matching a PCRE-like regular expression. + r"""Replaces substring(s) matching a PCRE-like regular expression. The full list of supported features and syntax can be found at Supported flags with the addition of 'g' can be found at + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello 42"]}) + >>> result = df.select( + ... dfn.functions.regexp_replace( + ... dfn.col("a"), dfn.lit("\\d+"), + ... dfn.lit("XX") + ... ).alias("r") + ... ) + >>> result.collect_column("r")[0].as_py() + 'hello XX' """ if flags is not None: flags = flags.expr @@ -781,36 +1407,128 @@ def regexp_replace( def regexp_count( - string: Expr, pattern: Expr, start: Expr, flags: Expr | None = None + string: Expr, pattern: Expr, start: Expr | None = None, flags: Expr | None = None ) -> Expr: """Returns the number of matches in a string. Optional start position (the first position is 1) to search for the regular expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["abcabc"]}) + >>> result = df.select( + ... dfn.functions.regexp_count(dfn.col("a"), dfn.lit("abc")).alias("c")) + >>> result.collect_column("c")[0].as_py() + 2 """ if flags is not None: flags = flags.expr - start = start.expr if start is not None else Expr.expr + start = start.expr if start is not None else start return Expr(f.regexp_count(string.expr, pattern.expr, start, flags)) +def regexp_instr( + values: Expr, + regex: Expr, + start: Expr | None = None, + n: Expr | None = None, + flags: Expr | None = None, + sub_expr: Expr | None = None, +) -> Expr: + r"""Returns the position of a regular expression match in a string. + + Args: + values: Data to search for the regular expression match. + regex: Regular expression to search for. + start: Optional position to start the search (the first position is 1). + n: Optional occurrence of the match to find (the first occurrence is 1). + flags: Optional regular expression flags to control regex behavior. + sub_expr: Optionally capture group position instead of the entire match. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello 42 world"]}) + >>> result = df.select( + ... dfn.functions.regexp_instr( + ... dfn.col("a"), dfn.lit("\\d+") + ... ).alias("pos") + ... ) + >>> result.collect_column("pos")[0].as_py() + 7 + """ + start = start.expr if start is not None else None + n = n.expr if n is not None else None + flags = flags.expr if flags is not None else None + sub_expr = sub_expr.expr if sub_expr is not None else None + + return Expr( + f.regexp_instr( + values.expr, + regex.expr, + start, + n, + flags, + sub_expr, + ) + ) + + def repeat(string: Expr, n: Expr) -> Expr: - """Repeats the ``string`` to ``n`` times.""" + """Repeats the ``string`` to ``n`` times. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["ha"]}) + >>> result = df.select(dfn.functions.repeat(dfn.col("a"), dfn.lit(3)).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'hahaha' + """ return Expr(f.repeat(string.expr, n.expr)) def replace(string: Expr, from_val: Expr, to_val: Expr) -> Expr: - """Replaces all occurrences of ``from_val`` with ``to_val`` in the ``string``.""" + """Replaces all occurrences of ``from_val`` with ``to_val`` in the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello world"]}) + >>> result = df.select( + ... dfn.functions.replace(dfn.col("a"), dfn.lit("world"), + ... dfn.lit("there")).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'hello there' + """ return Expr(f.replace(string.expr, from_val.expr, to_val.expr)) def reverse(arg: Expr) -> Expr: - """Reverse the string argument.""" + """Reverse the string argument. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.reverse(dfn.col("a")).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'olleh' + """ return Expr(f.reverse(arg.expr)) def right(string: Expr, n: Expr) -> Expr: - """Returns the last ``n`` characters in the ``string``.""" + """Returns the last ``n`` characters in the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.right(dfn.col("a"), dfn.lit(3)).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'llo' + """ return Expr(f.right(string.expr, n.expr)) @@ -820,6 +1538,13 @@ def round(value: Expr, decimal_places: Expr | None = None) -> Expr: If the optional ``decimal_places`` is specified, round to the nearest number of decimal places. You can specify a negative number of decimal places. For example ``round(lit(125.2345), lit(-2))`` would yield a value of ``100.0``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.567]}) + >>> result = df.select(dfn.functions.round(dfn.col("a"), dfn.lit(2)).alias("r")) + >>> result.collect_column("r")[0].as_py() + 1.57 """ if decimal_places is None: decimal_places = Expr.literal(0) @@ -831,48 +1556,130 @@ def rpad(string: Expr, count: Expr, characters: Expr | None = None) -> Expr: Extends the string to length length by appending the characters fill (a space by default). If the string is already longer than length then it is truncated. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hi"]}) + >>> result = df.select( + ... dfn.functions.rpad(dfn.col("a"), dfn.lit(5), dfn.lit("!")).alias("r")) + >>> result.collect_column("r")[0].as_py() + 'hi!!!' """ characters = characters if characters is not None else Expr.literal(" ") return Expr(f.rpad(string.expr, count.expr, characters.expr)) def rtrim(arg: Expr) -> Expr: - """Removes all characters, spaces by default, from the end of a string.""" + """Removes all characters, spaces by default, from the end of a string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [" a "]}) + >>> trim_df = df.select(dfn.functions.rtrim(dfn.col("a")).alias("trimmed")) + >>> trim_df.collect_column("trimmed")[0].as_py() + ' a' + """ return Expr(f.rtrim(arg.expr)) def sha224(arg: Expr) -> Expr: - """Computes the SHA-224 hash of a binary string.""" + """Computes the SHA-224 hash of a binary string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.sha224(dfn.col("a")).alias("h") + ... ) + >>> result.collect_column("h")[0].as_py().hex() + 'ea09ae9cc6768c50fcee903ed054556e5bfc8347907f12598aa24193' + """ return Expr(f.sha224(arg.expr)) def sha256(arg: Expr) -> Expr: - """Computes the SHA-256 hash of a binary string.""" + """Computes the SHA-256 hash of a binary string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.sha256(dfn.col("a")).alias("h") + ... ) + >>> result.collect_column("h")[0].as_py().hex() + '2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' + """ return Expr(f.sha256(arg.expr)) def sha384(arg: Expr) -> Expr: - """Computes the SHA-384 hash of a binary string.""" + """Computes the SHA-384 hash of a binary string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.sha384(dfn.col("a")).alias("h") + ... ) + >>> result.collect_column("h")[0].as_py().hex() + '59e1748777448c69de6b800d7a33bbfb9ff1b... + """ return Expr(f.sha384(arg.expr)) def sha512(arg: Expr) -> Expr: - """Computes the SHA-512 hash of a binary string.""" + """Computes the SHA-512 hash of a binary string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.sha512(dfn.col("a")).alias("h") + ... ) + >>> result.collect_column("h")[0].as_py().hex() + '9b71d224bd62f3785d96d46ad3ea3d73319bfb... + """ return Expr(f.sha512(arg.expr)) def signum(arg: Expr) -> Expr: - """Returns the sign of the argument (-1, 0, +1).""" + """Returns the sign of the argument (-1, 0, +1). + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [-5.0, 0.0, 5.0]}) + >>> result = df.select(dfn.functions.signum(dfn.col("a")).alias("s")) + >>> result.collect_column("s").to_pylist() + [-1.0, 0.0, 1.0] + """ return Expr(f.signum(arg.expr)) def sin(arg: Expr) -> Expr: - """Returns the sine of the argument.""" + """Returns the sine of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.sin(dfn.col("a")).alias("sin")) + >>> result.collect_column("sin")[0].as_py() + 0.0 + """ return Expr(f.sin(arg.expr)) def sinh(arg: Expr) -> Expr: - """Returns the hyperbolic sine of the argument.""" + """Returns the hyperbolic sine of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.sinh(dfn.col("a")).alias("sinh")) + >>> result.collect_column("sinh")[0].as_py() + 0.0 + """ return Expr(f.sinh(arg.expr)) @@ -881,27 +1688,73 @@ def split_part(string: Expr, delimiter: Expr, index: Expr) -> Expr: Splits a string based on a delimiter and picks out the desired field based on the index. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["a,b,c"]}) + >>> result = df.select( + ... dfn.functions.split_part(dfn.col("a"), dfn.lit(","), dfn.lit(2)).alias("s")) + >>> result.collect_column("s")[0].as_py() + 'b' """ return Expr(f.split_part(string.expr, delimiter.expr, index.expr)) def sqrt(arg: Expr) -> Expr: - """Returns the square root of the argument.""" + """Returns the square root of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [9.0]}) + >>> result = df.select(dfn.functions.sqrt(dfn.col("a")).alias("sqrt")) + >>> result.collect_column("sqrt")[0].as_py() + 3.0 + """ return Expr(f.sqrt(arg.expr)) def starts_with(string: Expr, prefix: Expr) -> Expr: - """Returns true if string starts with prefix.""" + """Returns true if string starts with prefix. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello_from_datafusion"]}) + >>> result = df.select( + ... dfn.functions.starts_with(dfn.col("a"), dfn.lit("hello")).alias("sw")) + >>> result.collect_column("sw")[0].as_py() + True + """ return Expr(f.starts_with(string.expr, prefix.expr)) def strpos(string: Expr, substring: Expr) -> Expr: - """Finds the position from where the ``substring`` matches the ``string``.""" + """Finds the position from where the ``substring`` matches the ``string``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.strpos(dfn.col("a"), dfn.lit("llo")).alias("pos")) + >>> result.collect_column("pos")[0].as_py() + 3 + """ return Expr(f.strpos(string.expr, substring.expr)) def substr(string: Expr, position: Expr) -> Expr: - """Substring from the ``position`` to the end.""" + """Substring from the ``position`` to the end. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.substr(dfn.col("a"), dfn.lit(3)).alias("s")) + >>> result.collect_column("s")[0].as_py() + 'llo' + """ return Expr(f.substr(string.expr, position.expr)) @@ -910,27 +1763,72 @@ def substr_index(string: Expr, delimiter: Expr, count: Expr) -> Expr: The return will be the ``string`` from before ``count`` occurrences of ``delimiter``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["a.b.c"]}) + >>> result = df.select( + ... dfn.functions.substr_index(dfn.col("a"), dfn.lit("."), + ... dfn.lit(2)).alias("s")) + >>> result.collect_column("s")[0].as_py() + 'a.b' """ return Expr(f.substr_index(string.expr, delimiter.expr, count.expr)) def substring(string: Expr, position: Expr, length: Expr) -> Expr: - """Substring from the ``position`` with ``length`` characters.""" + """Substring from the ``position`` with ``length`` characters. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello world"]}) + >>> result = df.select( + ... dfn.functions.substring(dfn.col("a"), dfn.lit(1), dfn.lit(5)).alias("s")) + >>> result.collect_column("s")[0].as_py() + 'hello' + """ return Expr(f.substring(string.expr, position.expr, length.expr)) def tan(arg: Expr) -> Expr: - """Returns the tangent of the argument.""" + """Returns the tangent of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.tan(dfn.col("a")).alias("tan")) + >>> result.collect_column("tan")[0].as_py() + 0.0 + """ return Expr(f.tan(arg.expr)) def tanh(arg: Expr) -> Expr: - """Returns the hyperbolic tangent of the argument.""" + """Returns the hyperbolic tangent of the argument. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0]}) + >>> result = df.select(dfn.functions.tanh(dfn.col("a")).alias("tanh")) + >>> result.collect_column("tanh")[0].as_py() + 0.0 + """ return Expr(f.tanh(arg.expr)) def to_hex(arg: Expr) -> Expr: - """Converts an integer to a hexadecimal string.""" + """Converts an integer to a hexadecimal string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [255]}) + >>> result = df.select(dfn.functions.to_hex(dfn.col("a")).alias("hex")) + >>> result.collect_column("hex")[0].as_py() + 'ff' + """ return Expr(f.to_hex(arg.expr)) @@ -938,73 +1836,217 @@ def now() -> Expr: """Returns the current timestamp in nanoseconds. This will use the same value for all instances of now() in same statement. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.now().alias("now") + ... ) + + Use .value instead of .as_py() because nanosecond timestamps + require pandas to convert to Python datetime objects. + + >>> result.collect_column("now")[0].value > 0 + True """ return Expr(f.now()) +def to_char(arg: Expr, formatter: Expr) -> Expr: + """Returns a string representation of a date, time, timestamp or duration. + + For usage of ``formatter`` see the rust chrono package ``strftime`` package. + + [Documentation here.](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) + """ + return Expr(f.to_char(arg.expr, formatter.expr)) + + +def _unwrap_exprs(args: tuple[Expr, ...]) -> list: + return [arg.expr for arg in args] + + +def to_date(arg: Expr, *formatters: Expr) -> Expr: + """Converts a value to a date (YYYY-MM-DD). + + Supports strings, numeric and timestamp types as input. + Integers and doubles are interpreted as days since the unix epoch. + Strings are parsed as YYYY-MM-DD (e.g. '2023-07-20') + if ``formatters`` are not provided. + + For usage of ``formatters`` see the rust chrono package ``strftime`` package. + + [Documentation here.](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) + """ + return Expr(f.to_date(arg.expr, *_unwrap_exprs(formatters))) + + +def to_local_time(*args: Expr) -> Expr: + """Converts a timestamp with a timezone to a timestamp without a timezone. + + This function handles daylight saving time changes. + """ + return Expr(f.to_local_time(*_unwrap_exprs(args))) + + +def to_time(arg: Expr, *formatters: Expr) -> Expr: + """Converts a value to a time. Supports strings and timestamps as input. + + If ``formatters`` is not provided strings are parsed as HH:MM:SS, HH:MM or + HH:MM:SS.nnnnnnnnn; + + For usage of ``formatters`` see the rust chrono package ``strftime`` package. + + [Documentation here.](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) + """ + return Expr(f.to_time(arg.expr, *_unwrap_exprs(formatters))) + + def to_timestamp(arg: Expr, *formatters: Expr) -> Expr: """Converts a string and optional formats to a ``Timestamp`` in nanoseconds. For usage of ``formatters`` see the rust chrono package ``strftime`` package. [Documentation here.](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) - """ - if formatters is None: - return f.to_timestamp(arg.expr) - formatters = [f.expr for f in formatters] - return Expr(f.to_timestamp(arg.expr, *formatters)) + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-01-01T00:00:00"]}) + >>> result = df.select( + ... dfn.functions.to_timestamp( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '2021-01-01 00:00:00' + """ + return Expr(f.to_timestamp(arg.expr, *_unwrap_exprs(formatters))) def to_timestamp_millis(arg: Expr, *formatters: Expr) -> Expr: """Converts a string and optional formats to a ``Timestamp`` in milliseconds. See :py:func:`to_timestamp` for a description on how to use formatters. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-01-01T00:00:00"]}) + >>> result = df.select( + ... dfn.functions.to_timestamp_millis( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '2021-01-01 00:00:00' """ - formatters = [f.expr for f in formatters] - return Expr(f.to_timestamp_millis(arg.expr, *formatters)) + return Expr(f.to_timestamp_millis(arg.expr, *_unwrap_exprs(formatters))) def to_timestamp_micros(arg: Expr, *formatters: Expr) -> Expr: """Converts a string and optional formats to a ``Timestamp`` in microseconds. See :py:func:`to_timestamp` for a description on how to use formatters. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-01-01T00:00:00"]}) + >>> result = df.select( + ... dfn.functions.to_timestamp_micros( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '2021-01-01 00:00:00' """ - formatters = [f.expr for f in formatters] - return Expr(f.to_timestamp_micros(arg.expr, *formatters)) + return Expr(f.to_timestamp_micros(arg.expr, *_unwrap_exprs(formatters))) def to_timestamp_nanos(arg: Expr, *formatters: Expr) -> Expr: """Converts a string and optional formats to a ``Timestamp`` in nanoseconds. See :py:func:`to_timestamp` for a description on how to use formatters. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-01-01T00:00:00"]}) + >>> result = df.select( + ... dfn.functions.to_timestamp_nanos( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '2021-01-01 00:00:00' """ - formatters = [f.expr for f in formatters] - return Expr(f.to_timestamp_nanos(arg.expr, *formatters)) + return Expr(f.to_timestamp_nanos(arg.expr, *_unwrap_exprs(formatters))) def to_timestamp_seconds(arg: Expr, *formatters: Expr) -> Expr: """Converts a string and optional formats to a ``Timestamp`` in seconds. See :py:func:`to_timestamp` for a description on how to use formatters. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-01-01T00:00:00"]}) + >>> result = df.select( + ... dfn.functions.to_timestamp_seconds( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '2021-01-01 00:00:00' """ - formatters = [f.expr for f in formatters] - return Expr(f.to_timestamp_seconds(arg.expr, *formatters)) + return Expr(f.to_timestamp_seconds(arg.expr, *_unwrap_exprs(formatters))) def to_unixtime(string: Expr, *format_arguments: Expr) -> Expr: - """Converts a string and optional formats to a Unixtime.""" - args = [f.expr for f in format_arguments] - return Expr(f.to_unixtime(string.expr, *args)) + """Converts a string and optional formats to a Unixtime. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["1970-01-01T00:00:00"]}) + >>> result = df.select(dfn.functions.to_unixtime(dfn.col("a")).alias("u")) + >>> result.collect_column("u")[0].as_py() + 0 + """ + return Expr(f.to_unixtime(string.expr, *_unwrap_exprs(format_arguments))) def current_date() -> Expr: - """Returns current UTC date as a Date32 value.""" + """Returns current UTC date as a Date32 value. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.current_date().alias("d") + ... ) + >>> result.collect_column("d")[0].as_py() is not None + True + """ return Expr(f.current_date()) +today = current_date + + def current_time() -> Expr: - """Returns current UTC time as a Time64 value.""" + """Returns current UTC time as a Time64 value. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.current_time().alias("t") + ... ) + + Use .value instead of .as_py() because nanosecond timestamps + require pandas to convert to Python datetime objects. + + >>> result.collect_column("t")[0].value > 0 + True + """ return Expr(f.current_time()) @@ -1017,7 +2059,17 @@ def datepart(part: Expr, date: Expr) -> Expr: def date_part(part: Expr, date: Expr) -> Expr: - """Extracts a subfield from the date.""" + """Extracts a subfield from the date. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-07-15T00:00:00"]}) + >>> df = df.select(dfn.functions.to_timestamp(dfn.col("a")).alias("a")) + >>> result = df.select( + ... dfn.functions.date_part(dfn.lit("year"), dfn.col("a")).alias("y")) + >>> result.collect_column("y")[0].as_py() + 2021 + """ return Expr(f.date_part(part.expr, date.expr)) @@ -1030,7 +2082,20 @@ def extract(part: Expr, date: Expr) -> Expr: def date_trunc(part: Expr, date: Expr) -> Expr: - """Truncates the date to a specified level of precision.""" + """Truncates the date to a specified level of precision. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["2021-07-15T12:34:56"]}) + >>> df = df.select(dfn.functions.to_timestamp(dfn.col("a")).alias("a")) + >>> result = df.select( + ... dfn.functions.date_trunc( + ... dfn.lit("month"), dfn.col("a") + ... ).alias("t") + ... ) + >>> str(result.collect_column("t")[0].as_py()) + '2021-07-01 00:00:00' + """ return Expr(f.date_trunc(part.expr, date.expr)) @@ -1043,39 +2108,113 @@ def datetrunc(part: Expr, date: Expr) -> Expr: def date_bin(stride: Expr, source: Expr, origin: Expr) -> Expr: - """Coerces an arbitrary timestamp to the start of the nearest specified interval.""" + """Coerces an arbitrary timestamp to the start of the nearest specified interval. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"timestamp": ['2021-07-15 12:34:56', '2021-01-01']}) + >>> result = df.select( + ... dfn.functions.date_bin( + ... dfn.string_literal("15 minutes"), + ... dfn.col("timestamp"), + ... dfn.string_literal("2001-01-01 00:00:00") + ... ).alias("b") + ... ) + >>> str(result.collect_column("b")[0].as_py()) + '2021-07-15 12:30:00' + >>> str(result.collect_column("b")[1].as_py()) + '2021-01-01 00:00:00' + """ return Expr(f.date_bin(stride.expr, source.expr, origin.expr)) def make_date(year: Expr, month: Expr, day: Expr) -> Expr: - """Make a date from year, month and day component parts.""" + """Make a date from year, month and day component parts. + + Examples: + >>> from datetime import date + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [2024], "m": [1], "d": [15]}) + >>> result = df.select( + ... dfn.functions.make_date(dfn.col("y"), dfn.col("m"), + ... dfn.col("d")).alias("dt")) + >>> result.collect_column("dt")[0].as_py() + datetime.date(2024, 1, 15) + """ return Expr(f.make_date(year.expr, month.expr, day.expr)) def translate(string: Expr, from_val: Expr, to_val: Expr) -> Expr: - """Replaces the characters in ``from_val`` with the counterpart in ``to_val``.""" + """Replaces the characters in ``from_val`` with the counterpart in ``to_val``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select( + ... dfn.functions.translate(dfn.col("a"), dfn.lit("helo"), + ... dfn.lit("HELO")).alias("t")) + >>> result.collect_column("t")[0].as_py() + 'HELLO' + """ return Expr(f.translate(string.expr, from_val.expr, to_val.expr)) def trim(arg: Expr) -> Expr: - """Removes all characters, spaces by default, from both sides of a string.""" + """Removes all characters, spaces by default, from both sides of a string. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [" hello "]}) + >>> result = df.select(dfn.functions.trim(dfn.col("a")).alias("t")) + >>> result.collect_column("t")[0].as_py() + 'hello' + """ return Expr(f.trim(arg.expr)) def trunc(num: Expr, precision: Expr | None = None) -> Expr: - """Truncate the number toward zero with optional precision.""" + """Truncate the number toward zero with optional precision. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.567]}) + >>> result = df.select(dfn.functions.trunc(dfn.col("a")).alias("t")) + >>> result.collect_column("t")[0].as_py() + 1.0 + """ if precision is not None: return Expr(f.trunc(num.expr, precision.expr)) return Expr(f.trunc(num.expr)) def upper(arg: Expr) -> Expr: - """Converts a string to uppercase.""" + """Converts a string to uppercase. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["hello"]}) + >>> result = df.select(dfn.functions.upper(dfn.col("a")).alias("u")) + >>> result.collect_column("u")[0].as_py() + 'HELLO' + """ return Expr(f.upper(arg.expr)) def make_array(*args: Expr) -> Expr: - """Returns an array using the specified input expressions.""" + """Returns an array using the specified input expressions. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.make_array(dfn.lit(1), dfn.lit(2), dfn.lit(3)).alias("arr")) + >>> result.collect_column("arr")[0].as_py() + [1, 2, 3] + """ args = [arg.expr for arg in args] return Expr(f.make_array(args)) @@ -1097,23 +2236,71 @@ def array(*args: Expr) -> Expr: def range(start: Expr, stop: Expr, step: Expr) -> Expr: - """Create a list of values in the range between start and stop.""" + """Create a list of values in the range between start and stop. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.range(dfn.lit(0), dfn.lit(5), dfn.lit(2)).alias("r")) + >>> result.collect_column("r")[0].as_py() + [0, 2, 4] + """ return Expr(f.range(start.expr, stop.expr, step.expr)) def uuid() -> Expr: - """Returns uuid v4 as a string value.""" + """Returns uuid v4 as a string value. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.uuid().alias("u") + ... ) + >>> len(result.collect_column("u")[0].as_py()) == 36 + True + """ return Expr(f.uuid()) def struct(*args: Expr) -> Expr: - """Returns a struct with the given arguments.""" + """Returns a struct with the given arguments. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1], "b": [2]}) + >>> result = df.select( + ... dfn.functions.struct( + ... dfn.col("a"), dfn.col("b") + ... ).alias("s") + ... ) + + Children in the new struct will always be `c0`, ..., `cN-1` + for `N` children. + + >>> result.collect_column("s")[0].as_py() == {"c0": 1, "c1": 2} + True + """ args = [arg.expr for arg in args] return Expr(f.struct(*args)) def named_struct(name_pairs: list[tuple[str, Expr]]) -> Expr: - """Returns a struct with the given names and arguments pairs.""" + """Returns a struct with the given names and arguments pairs. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.named_struct( + ... [("x", dfn.lit(10)), ("y", dfn.lit(20))] + ... ).alias("s") + ... ) + >>> result.collect_column("s")[0].as_py() == {"x": 10, "y": 20} + True + """ name_pair_exprs = [ [Expr.literal(pa.scalar(pair[0], type=pa.string())), pair[1]] for pair in name_pairs @@ -1125,27 +2312,79 @@ def named_struct(name_pairs: list[tuple[str, Expr]]) -> Expr: def from_unixtime(arg: Expr) -> Expr: - """Converts an integer to RFC3339 timestamp format string.""" + """Converts an integer to RFC3339 timestamp format string. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0]}) + >>> result = df.select( + ... dfn.functions.from_unixtime( + ... dfn.col("a") + ... ).alias("ts") + ... ) + >>> str(result.collect_column("ts")[0].as_py()) + '1970-01-01 00:00:00' + """ return Expr(f.from_unixtime(arg.expr)) def arrow_typeof(arg: Expr) -> Expr: - """Returns the Arrow type of the expression.""" + """Returns the Arrow type of the expression. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select(dfn.functions.arrow_typeof(dfn.col("a")).alias("t")) + >>> result.collect_column("t")[0].as_py() + 'Int64' + """ return Expr(f.arrow_typeof(arg.expr)) def arrow_cast(expr: Expr, data_type: Expr) -> Expr: - """Casts an expression to a specified data type.""" + """Casts an expression to a specified data type. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> data_type = dfn.string_literal("Float64") + >>> result = df.select( + ... dfn.functions.arrow_cast(dfn.col("a"), data_type).alias("c") + ... ) + >>> result.collect_column("c")[0].as_py() + 1.0 + """ return Expr(f.arrow_cast(expr.expr, data_type.expr)) def random() -> Expr: - """Returns a random value in the range ``0.0 <= x < 1.0``.""" + """Returns a random value in the range ``0.0 <= x < 1.0``. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.random().alias("r") + ... ) + >>> val = result.collect_column("r")[0].as_py() + >>> 0.0 <= val < 1.0 + True + """ return Expr(f.random()) def array_append(array: Expr, element: Expr) -> Expr: - """Appends an element to the end of an array.""" + """Appends an element to the end of an array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_append(dfn.col("a"), dfn.lit(4)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3, 4] + """ return Expr(f.array_append(array.expr, element.expr)) @@ -1174,7 +2413,17 @@ def list_push_back(array: Expr, element: Expr) -> Expr: def array_concat(*args: Expr) -> Expr: - """Concatenates the input arrays.""" + """Concatenates the input arrays. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]], "b": [[3, 4]]}) + >>> result = df.select( + ... dfn.functions.array_concat(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3, 4] + """ args = [arg.expr for arg in args] return Expr(f.array_concat(args)) @@ -1188,12 +2437,36 @@ def array_cat(*args: Expr) -> Expr: def array_dims(array: Expr) -> Expr: - """Returns an array of the array's dimensions.""" + """Returns an array of the array's dimensions. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.array_dims(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [3] + """ return Expr(f.array_dims(array.expr)) def array_distinct(array: Expr) -> Expr: - """Returns distinct values from the array after removing duplicates.""" + """Returns distinct values from the array after removing duplicates. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_distinct( + ... dfn.col("a") + ... ).alias("result") + ... ) + >>> sorted( + ... result.collect_column("result")[0].as_py() + ... ) + [1, 2, 3] + """ return Expr(f.array_distinct(array.expr)) @@ -1230,12 +2503,31 @@ def list_dims(array: Expr) -> Expr: def array_element(array: Expr, n: Expr) -> Expr: - """Extracts the element with the index n from the array.""" + """Extracts the element with the index n from the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[10, 20, 30]]}) + >>> result = df.select( + ... dfn.functions.array_element(dfn.col("a"), dfn.lit(2)).alias("result")) + >>> result.collect_column("result")[0].as_py() + 20 + """ return Expr(f.array_element(array.expr, n.expr)) def array_empty(array: Expr) -> Expr: - """Returns a boolean indicating whether the array is empty.""" + """Returns a boolean indicating whether the array is empty. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]]}) + >>> result = df.select(dfn.functions.array_empty(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + False + """ return Expr(f.array_empty(array.expr)) @@ -1264,7 +2556,16 @@ def list_extract(array: Expr, n: Expr) -> Expr: def array_length(array: Expr) -> Expr: - """Returns the length of the array.""" + """Returns the length of the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.array_length(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 3 + """ return Expr(f.array_length(array.expr)) @@ -1277,7 +2578,17 @@ def list_length(array: Expr) -> Expr: def array_has(first_array: Expr, second_array: Expr) -> Expr: - """Returns true if the element appears in the first array, otherwise false.""" + """Returns true if the element appears in the first array, otherwise false. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select( + ... dfn.functions.array_has(dfn.col("a"), dfn.lit(2)).alias("result")) + >>> result.collect_column("result")[0].as_py() + True + """ return Expr(f.array_has(first_array.expr, second_array.expr)) @@ -1286,6 +2597,15 @@ def array_has_all(first_array: Expr, second_array: Expr) -> Expr: Returns true if each element of the second array appears in the first array. Otherwise, it returns false. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]], "b": [[1, 2]]}) + >>> result = df.select( + ... dfn.functions.array_has_all(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + True """ return Expr(f.array_has_all(first_array.expr, second_array.expr)) @@ -1295,12 +2615,31 @@ def array_has_any(first_array: Expr, second_array: Expr) -> Expr: Returns true if at least one element of the second array appears in the first array. Otherwise, it returns false. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]], "b": [[2, 5]]}) + >>> result = df.select( + ... dfn.functions.array_has_any(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + True """ return Expr(f.array_has_any(first_array.expr, second_array.expr)) def array_position(array: Expr, element: Expr, index: int | None = 1) -> Expr: - """Return the position of the first occurrence of ``element`` in ``array``.""" + """Return the position of the first occurrence of ``element`` in ``array``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[10, 20, 30]]}) + >>> result = df.select( + ... dfn.functions.array_position(dfn.col("a"), dfn.lit(20)).alias("result")) + >>> result.collect_column("result")[0].as_py() + 2 + """ return Expr(f.array_position(array.expr, element.expr, index)) @@ -1329,7 +2668,17 @@ def list_indexof(array: Expr, element: Expr, index: int | None = 1) -> Expr: def array_positions(array: Expr, element: Expr) -> Expr: - """Searches for an element in the array and returns all occurrences.""" + """Searches for an element in the array and returns all occurrences. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1]]}) + >>> result = df.select( + ... dfn.functions.array_positions(dfn.col("a"), dfn.lit(1)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 3] + """ return Expr(f.array_positions(array.expr, element.expr)) @@ -1342,7 +2691,16 @@ def list_positions(array: Expr, element: Expr) -> Expr: def array_ndims(array: Expr) -> Expr: - """Returns the number of dimensions of the array.""" + """Returns the number of dimensions of the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.array_ndims(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 1 + """ return Expr(f.array_ndims(array.expr)) @@ -1355,7 +2713,17 @@ def list_ndims(array: Expr) -> Expr: def array_prepend(element: Expr, array: Expr) -> Expr: - """Prepends an element to the beginning of an array.""" + """Prepends an element to the beginning of an array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]]}) + >>> result = df.select( + ... dfn.functions.array_prepend(dfn.lit(0), dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [0, 1, 2] + """ return Expr(f.array_prepend(element.expr, array.expr)) @@ -1384,17 +2752,45 @@ def list_push_front(element: Expr, array: Expr) -> Expr: def array_pop_back(array: Expr) -> Expr: - """Returns the array without the last element.""" + """Returns the array without the last element. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.array_pop_back(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2] + """ return Expr(f.array_pop_back(array.expr)) def array_pop_front(array: Expr) -> Expr: - """Returns the array without the first element.""" + """Returns the array without the first element. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.array_pop_front(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [2, 3] + """ return Expr(f.array_pop_front(array.expr)) def array_remove(array: Expr, element: Expr) -> Expr: - """Removes the first element from the array equal to the given value.""" + """Removes the first element from the array equal to the given value. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1]]}) + >>> result = df.select( + ... dfn.functions.array_remove(dfn.col("a"), dfn.lit(1)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [2, 1] + """ return Expr(f.array_remove(array.expr, element.expr)) @@ -1407,7 +2803,18 @@ def list_remove(array: Expr, element: Expr) -> Expr: def array_remove_n(array: Expr, element: Expr, max: Expr) -> Expr: - """Removes the first ``max`` elements from the array equal to the given value.""" + """Removes the first ``max`` elements from the array equal to the given value. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1, 1]]}) + >>> result = df.select( + ... dfn.functions.array_remove_n(dfn.col("a"), dfn.lit(1), + ... dfn.lit(2)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [2, 1] + """ return Expr(f.array_remove_n(array.expr, element.expr, max.expr)) @@ -1420,7 +2827,17 @@ def list_remove_n(array: Expr, element: Expr, max: Expr) -> Expr: def array_remove_all(array: Expr, element: Expr) -> Expr: - """Removes all elements from the array equal to the given value.""" + """Removes all elements from the array equal to the given value. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1]]}) + >>> result = df.select( + ... dfn.functions.array_remove_all(dfn.col("a"), dfn.lit(1)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [2] + """ return Expr(f.array_remove_all(array.expr, element.expr)) @@ -1433,7 +2850,17 @@ def list_remove_all(array: Expr, element: Expr) -> Expr: def array_repeat(element: Expr, count: Expr) -> Expr: - """Returns an array containing ``element`` ``count`` times.""" + """Returns an array containing ``element`` ``count`` times. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1]}) + >>> result = df.select( + ... dfn.functions.array_repeat(dfn.lit(3), dfn.lit(3)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [3, 3, 3] + """ return Expr(f.array_repeat(element.expr, count.expr)) @@ -1446,7 +2873,18 @@ def list_repeat(element: Expr, count: Expr) -> Expr: def array_replace(array: Expr, from_val: Expr, to_val: Expr) -> Expr: - """Replaces the first occurrence of ``from_val`` with ``to_val``.""" + """Replaces the first occurrence of ``from_val`` with ``to_val``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1]]}) + >>> result = df.select( + ... dfn.functions.array_replace(dfn.col("a"), dfn.lit(1), + ... dfn.lit(9)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [9, 2, 1] + """ return Expr(f.array_replace(array.expr, from_val.expr, to_val.expr)) @@ -1463,6 +2901,16 @@ def array_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr) -> Exp Replaces the first ``max`` occurrences of the specified element with another specified element. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1, 1]]}) + >>> result = df.select( + ... dfn.functions.array_replace_n(dfn.col("a"), dfn.lit(1), dfn.lit(9), + ... dfn.lit(2)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [9, 2, 9, 1] """ return Expr(f.array_replace_n(array.expr, from_val.expr, to_val.expr, max.expr)) @@ -1479,7 +2927,18 @@ def list_replace_n(array: Expr, from_val: Expr, to_val: Expr, max: Expr) -> Expr def array_replace_all(array: Expr, from_val: Expr, to_val: Expr) -> Expr: - """Replaces all occurrences of ``from_val`` with ``to_val``.""" + """Replaces all occurrences of ``from_val`` with ``to_val``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 1]]}) + >>> result = df.select( + ... dfn.functions.array_replace_all(dfn.col("a"), dfn.lit(1), + ... dfn.lit(9)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [9, 2, 9] + """ return Expr(f.array_replace_all(array.expr, from_val.expr, to_val.expr)) @@ -1498,6 +2957,14 @@ def array_sort(array: Expr, descending: bool = False, null_first: bool = False) array: The input array to sort. descending: If True, sorts in descending order. null_first: If True, nulls will be returned at the beginning of the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[3, 1, 2]]}) + >>> result = df.select(dfn.functions.array_sort(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3] """ desc = "DESC" if descending else "ASC" nulls_first = "NULLS FIRST" if null_first else "NULLS LAST" @@ -1518,7 +2985,18 @@ def list_sort(array: Expr, descending: bool = False, null_first: bool = False) - def array_slice( array: Expr, begin: Expr, end: Expr, stride: Expr | None = None ) -> Expr: - """Returns a slice of the array.""" + """Returns a slice of the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3, 4]]}) + >>> result = df.select( + ... dfn.functions.array_slice(dfn.col("a"), dfn.lit(2), + ... dfn.lit(3)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [2, 3] + """ if stride is not None: stride = stride.expr return Expr(f.array_slice(array.expr, begin.expr, end.expr, stride)) @@ -1533,7 +3011,22 @@ def list_slice(array: Expr, begin: Expr, end: Expr, stride: Expr | None = None) def array_intersect(array1: Expr, array2: Expr) -> Expr: - """Returns the intersection of ``array1`` and ``array2``.""" + """Returns the intersection of ``array1`` and ``array2``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]], "b": [[2, 3, 4]]}) + >>> result = df.select( + ... dfn.functions.array_intersect( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> sorted( + ... result.collect_column("result")[0].as_py() + ... ) + [2, 3] + """ return Expr(f.array_intersect(array1.expr, array2.expr)) @@ -1549,6 +3042,20 @@ def array_union(array1: Expr, array2: Expr) -> Expr: """Returns an array of the elements in the union of array1 and array2. Duplicate rows will not be returned. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]], "b": [[2, 3, 4]]}) + >>> result = df.select( + ... dfn.functions.array_union( + ... dfn.col("a"), dfn.col("b") + ... ).alias("result") + ... ) + >>> sorted( + ... result.collect_column("result")[0].as_py() + ... ) + [1, 2, 3, 4] """ return Expr(f.array_union(array1.expr, array2.expr)) @@ -1564,7 +3071,17 @@ def list_union(array1: Expr, array2: Expr) -> Expr: def array_except(array1: Expr, array2: Expr) -> Expr: - """Returns the elements that appear in ``array1`` but not in ``array2``.""" + """Returns the elements that appear in ``array1`` but not in ``array2``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]], "b": [[2, 3, 4]]}) + >>> result = df.select( + ... dfn.functions.array_except(dfn.col("a"), dfn.col("b")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1] + """ return Expr(f.array_except(array1.expr, array2.expr)) @@ -1581,6 +3098,16 @@ def array_resize(array: Expr, size: Expr, value: Expr) -> Expr: If ``size`` is greater than the ``array`` length, the additional entries will be filled with the given ``value``. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2]]}) + >>> result = df.select( + ... dfn.functions.array_resize(dfn.col("a"), dfn.lit(4), + ... dfn.lit(0)).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 0, 0] """ return Expr(f.array_resize(array.expr, size.expr, value.expr)) @@ -1595,12 +3122,30 @@ def list_resize(array: Expr, size: Expr, value: Expr) -> Expr: def flatten(array: Expr) -> Expr: - """Flattens an array of arrays into a single array.""" + """Flattens an array of arrays into a single array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[[1, 2], [3, 4]]]}) + >>> result = df.select(dfn.functions.flatten(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + [1, 2, 3, 4] + """ return Expr(f.flatten(array.expr)) def cardinality(array: Expr) -> Expr: - """Returns the total number of elements in the array.""" + """Returns the total number of elements in the array. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [[1, 2, 3]]}) + >>> result = df.select(dfn.functions.cardinality(dfn.col("a")).alias("result")) + >>> result.collect_column("result")[0].as_py() + 3 + """ return Expr(f.cardinality(array.expr)) @@ -1612,7 +3157,7 @@ def empty(array: Expr) -> Expr: # aggregate functions def approx_distinct( expression: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Returns the approximate number of distinct values. @@ -1626,13 +3171,22 @@ def approx_distinct( Args: expression: Values to check for distinct entries filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 1, 2, 3]}) + >>> result = df.aggregate( + ... [], [dfn.functions.approx_distinct(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() == 3 + True """ filter_raw = filter.expr if filter is not None else None return Expr(f.approx_distinct(expression.expr, filter=filter_raw)) -def approx_median(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def approx_median(expression: Expr, filter: Expr | None = None) -> Expr: """Returns the approximate median value. This aggregate function is similar to :py:func:`median`, but it will only @@ -1644,16 +3198,25 @@ def approx_median(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: Values to find the median for filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.approx_median(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.approx_median(expression.expr, filter=filter_raw)) def approx_percentile_cont( - expression: Expr, + sort_expression: Expr | SortExpr, percentile: float, - num_centroids: Optional[int] = None, - filter: Optional[Expr] = None, + num_centroids: int | None = None, + filter: Expr | None = None, ) -> Expr: """Returns the value that is approximately at a given percentile of ``expr``. @@ -1664,28 +3227,42 @@ def approx_percentile_cont( between two of the values. This function uses the [t-digest](https://arxiv.org/abs/1902.04023) algorithm to - compute the percentil. You can limit the number of bins used in this algorithm by + compute the percentile. You can limit the number of bins used in this algorithm by setting the ``num_centroids`` parameter. If using the builder functions described in ref:`_aggregation` this function ignores the options ``order_by``, ``null_treatment``, and ``distinct``. Args: - expression: Values for which to find the approximate percentile + sort_expression: Values for which to find the approximate percentile percentile: This must be between 0.0 and 1.0, inclusive num_centroids: Max bin size for the t-digest algorithm filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.approx_percentile_cont(dfn.col("a"), 0.5).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3.0 """ + sort_expr_raw = sort_or_default(sort_expression) filter_raw = filter.expr if filter is not None else None return Expr( f.approx_percentile_cont( - expression.expr, percentile, num_centroids=num_centroids, filter=filter_raw + sort_expr_raw, percentile, num_centroids=num_centroids, filter=filter_raw ) ) def approx_percentile_cont_with_weight( - expression: Expr, weight: Expr, percentile: float, filter: Optional[Expr] = None + sort_expression: Expr | SortExpr, + weight: Expr, + percentile: float, + num_centroids: int | None = None, + filter: Expr | None = None, ) -> Expr: """Returns the value of the weighted approximate percentile. @@ -1696,16 +3273,31 @@ def approx_percentile_cont_with_weight( the options ``order_by``, ``null_treatment``, and ``distinct``. Args: - expression: Values for which to find the approximate percentile + sort_expression: Values for which to find the approximate percentile weight: Relative weight for each of the values in ``expression`` percentile: This must be between 0.0 and 1.0, inclusive + num_centroids: Max bin size for the t-digest algorithm filter: If provided, only compute against rows for which the filter is True + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0], "w": [1.0, 1.0, 1.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.approx_percentile_cont_with_weight(dfn.col("a"), + ... dfn.col("w"), 0.5).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ + sort_expr_raw = sort_or_default(sort_expression) filter_raw = filter.expr if filter is not None else None return Expr( f.approx_percentile_cont_with_weight( - expression.expr, weight.expr, percentile, filter=filter_raw + sort_expr_raw, + weight.expr, + percentile, + num_centroids=num_centroids, + filter=filter_raw, ) ) @@ -1713,8 +3305,8 @@ def approx_percentile_cont_with_weight( def array_agg( expression: Expr, distinct: bool = False, - filter: Optional[Expr] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + filter: Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Aggregate values into an array. @@ -1729,7 +3321,15 @@ def array_agg( expression: Values to combine into an array distinct: If True, a single entry for each distinct value will be in the result filter: If provided, only compute against rows for which the filter is True - order_by: Order the resultant array values + order_by: Order the resultant array values. Accepts column names or expressions. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.array_agg(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + [1, 2, 3] """ order_by_raw = sort_list_to_raw_sort_list(order_by) filter_raw = filter.expr if filter is not None else None @@ -1743,7 +3343,7 @@ def array_agg( def avg( expression: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Returns the average value. @@ -1755,12 +3355,20 @@ def avg( Args: expression: Values to combine into an array filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate([], [dfn.functions.avg(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.avg(expression.expr, filter=filter_raw)) -def corr(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: +def corr(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr: """Returns the correlation coefficient between ``value1`` and ``value2``. This aggregate function expects both values to be numeric and will return a float. @@ -1772,6 +3380,14 @@ def corr(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: value_y: The dependent variable for correlation value_x: The independent variable for correlation filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0], "b": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.corr(dfn.col("a"), dfn.col("b")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.corr(value_y.expr, value_x.expr, filter=filter_raw)) @@ -1780,7 +3396,7 @@ def corr(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: def count( expressions: Expr | list[Expr] | None = None, distinct: bool = False, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Returns the number of rows that match the given arguments. @@ -1793,6 +3409,14 @@ def count( expressions: Argument to perform bitwise calculation on distinct: If True, a single entry for each distinct value will be in the result filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.count(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3 """ filter_raw = filter.expr if filter is not None else None @@ -1806,7 +3430,7 @@ def count( return Expr(f.count(*args, distinct=distinct, filter=filter_raw)) -def covar_pop(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: +def covar_pop(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr: """Computes the population covariance. This aggregate function expects both values to be numeric and will return a float. @@ -1818,12 +3442,24 @@ def covar_pop(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Ex value_y: The dependent variable for covariance value_x: The independent variable for covariance filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 5.0, 10.0], "b": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], + ... [dfn.functions.covar_pop( + ... dfn.col("a"), dfn.col("b") + ... ).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 3.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.covar_pop(value_y.expr, value_x.expr, filter=filter_raw)) -def covar_samp(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: +def covar_samp(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr: """Computes the sample covariance. This aggregate function expects both values to be numeric and will return a float. @@ -1835,20 +3471,29 @@ def covar_samp(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> E value_y: The dependent variable for covariance value_x: The independent variable for covariance filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.covar_samp(dfn.col("a"), dfn.col("b")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.covar_samp(value_y.expr, value_x.expr, filter=filter_raw)) -def covar(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr: +def covar(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr: """Computes the sample covariance. - This is an alias for :py:func:`covar_samp`. + See Also: + This is an alias for :py:func:`covar_samp`. """ return covar_samp(value_y, value_x, filter) -def max(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def max(expression: Expr, filter: Expr | None = None) -> Expr: """Aggregate function that returns the maximum value of the argument. If using the builder functions described in ref:`_aggregation` this function ignores @@ -1857,21 +3502,37 @@ def max(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The value to find the maximum of filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.max(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3 """ filter_raw = filter.expr if filter is not None else None return Expr(f.max(expression.expr, filter=filter_raw)) -def mean(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def mean(expression: Expr, filter: Expr | None = None) -> Expr: """Returns the average (mean) value of the argument. This is an alias for :py:func:`avg`. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate([], [dfn.functions.mean(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ return avg(expression, filter) def median( - expression: Expr, distinct: bool = False, filter: Optional[Expr] = None + expression: Expr, distinct: bool = False, filter: Expr | None = None ) -> Expr: """Computes the median of a set of numbers. @@ -1885,13 +3546,21 @@ def median( expression: The value to compute the median of distinct: If True, a single entry for each distinct value will be in the result filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate([], [dfn.functions.median(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.median(expression.expr, distinct=distinct, filter=filter_raw)) -def min(expression: Expr, filter: Optional[Expr] = None) -> Expr: - """Returns the minimum value of the argument. +def min(expression: Expr, filter: Expr | None = None) -> Expr: + """Aggregate function that returns the minimum value of the argument. If using the builder functions described in ref:`_aggregation` this function ignores the options ``order_by``, ``null_treatment``, and ``distinct``. @@ -1899,6 +3568,14 @@ def min(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The value to find the minimum of filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.min(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1 """ filter_raw = filter.expr if filter is not None else None return Expr(f.min(expression.expr, filter=filter_raw)) @@ -1906,7 +3583,7 @@ def min(expression: Expr, filter: Optional[Expr] = None) -> Expr: def sum( expression: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the sum of a set of numbers. @@ -1918,12 +3595,20 @@ def sum( Args: expression: Values to combine into an array filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.aggregate([], [dfn.functions.sum(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 6 """ filter_raw = filter.expr if filter is not None else None return Expr(f.sum(expression.expr, filter=filter_raw)) -def stddev(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def stddev(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the standard deviation of the argument. If using the builder functions described in ref:`_aggregation` this function ignores @@ -1932,12 +3617,19 @@ def stddev(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The value to find the minimum of filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [2.0, 4.0, 6.0]}) + >>> result = df.aggregate([], [dfn.functions.stddev(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.stddev(expression.expr, filter=filter_raw)) -def stddev_pop(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def stddev_pop(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the population standard deviation of the argument. If using the builder functions described in ref:`_aggregation` this function ignores @@ -1946,28 +3638,53 @@ def stddev_pop(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The value to find the minimum of filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.stddev_pop(dfn.col("a")).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.stddev_pop(expression.expr, filter=filter_raw)) -def stddev_samp(arg: Expr, filter: Optional[Expr] = None) -> Expr: +def stddev_samp(arg: Expr, filter: Expr | None = None) -> Expr: """Computes the sample standard deviation of the argument. This is an alias for :py:func:`stddev`. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [2.0, 4.0, 6.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.stddev_samp(dfn.col("a")).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 2.0 """ return stddev(arg, filter=filter) -def var(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def var(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the sample variance of the argument. This is an alias for :py:func:`var_samp`. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate([], [dfn.functions.var(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ return var_samp(expression, filter) -def var_pop(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def var_pop(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the population variance of the argument. If using the builder functions described in ref:`_aggregation` this function ignores @@ -1976,12 +3693,19 @@ def var_pop(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The variable to compute the variance for filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [0.0, 2.0]}) + >>> result = df.aggregate([], [dfn.functions.var_pop(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.var_pop(expression.expr, filter=filter_raw)) -def var_samp(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def var_samp(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the sample variance of the argument. If using the builder functions described in ref:`_aggregation` this function ignores @@ -1990,15 +3714,31 @@ def var_samp(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: The variable to compute the variance for filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate([], [dfn.functions.var_samp(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None return Expr(f.var_sample(expression.expr, filter=filter_raw)) -def var_sample(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def var_sample(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the sample variance of the argument. This is an alias for :py:func:`var_samp`. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.var_sample(dfn.col("a")).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 1.0 """ return var_samp(expression, filter) @@ -2006,7 +3746,7 @@ def var_sample(expression: Expr, filter: Optional[Expr] = None) -> Expr: def regr_avgx( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the average of the independent variable ``x``. @@ -2020,6 +3760,14 @@ def regr_avgx( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [4.0, 5.0, 6.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_avgx(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 5.0 """ filter_raw = filter.expr if filter is not None else None @@ -2029,7 +3777,7 @@ def regr_avgx( def regr_avgy( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the average of the dependent variable ``y``. @@ -2043,6 +3791,14 @@ def regr_avgy( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [4.0, 5.0, 6.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_avgy(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None @@ -2052,7 +3808,7 @@ def regr_avgy( def regr_count( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Counts the number of rows in which both expressions are not null. @@ -2066,6 +3822,14 @@ def regr_count( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [4.0, 5.0, 6.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_count(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3 """ filter_raw = filter.expr if filter is not None else None @@ -2075,7 +3839,7 @@ def regr_count( def regr_intercept( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the intercept from the linear regression. @@ -2089,6 +3853,15 @@ def regr_intercept( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [2.0, 4.0, 6.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], + ... [dfn.functions.regr_intercept(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 0.0 """ filter_raw = filter.expr if filter is not None else None @@ -2098,7 +3871,7 @@ def regr_intercept( def regr_r2( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the R-squared value from linear regression. @@ -2112,6 +3885,14 @@ def regr_r2( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [2.0, 4.0, 6.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_r2(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 1.0 """ filter_raw = filter.expr if filter is not None else None @@ -2121,7 +3902,7 @@ def regr_r2( def regr_slope( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the slope from linear regression. @@ -2135,6 +3916,14 @@ def regr_slope( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [2.0, 4.0, 6.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_slope(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None @@ -2144,7 +3933,7 @@ def regr_slope( def regr_sxx( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the sum of squares of the independent variable ``x``. @@ -2158,6 +3947,14 @@ def regr_sxx( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_sxx(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None @@ -2167,7 +3964,7 @@ def regr_sxx( def regr_sxy( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the sum of products of pairs of numbers. @@ -2181,6 +3978,14 @@ def regr_sxy( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_sxy(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None @@ -2190,7 +3995,7 @@ def regr_sxy( def regr_syy( y: Expr, x: Expr, - filter: Optional[Expr] = None, + filter: Expr | None = None, ) -> Expr: """Computes the sum of squares of the dependent variable ``y``. @@ -2204,6 +4009,14 @@ def regr_syy( y: The linear regression dependent variable x: The linear regression independent variable filter: If provided, only compute against rows for which the filter is True + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"y": [1.0, 2.0, 3.0], "x": [1.0, 2.0, 3.0]}) + >>> result = df.aggregate( + ... [], [dfn.functions.regr_syy(dfn.col("y"), dfn.col("x")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 2.0 """ filter_raw = filter.expr if filter is not None else None @@ -2212,8 +4025,8 @@ def regr_syy( def first_value( expression: Expr, - filter: Optional[Expr] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + filter: Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS, ) -> Expr: """Returns the first value in a group of values. @@ -2226,8 +4039,18 @@ def first_value( Args: expression: Argument to perform bitwise calculation on filter: If provided, only compute against rows for which the filter is True - order_by: Set the ordering of the expression to evaluate + order_by: Set the ordering of the expression to evaluate. Accepts + column names or expressions. null_treatment: Assign whether to respect or ignore null values. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> result = df.aggregate( + ... [], [dfn.functions.first_value(dfn.col("a")).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 10 """ order_by_raw = sort_list_to_raw_sort_list(order_by) filter_raw = filter.expr if filter is not None else None @@ -2244,8 +4067,8 @@ def first_value( def last_value( expression: Expr, - filter: Optional[Expr] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + filter: Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS, ) -> Expr: """Returns the last value in a group of values. @@ -2258,8 +4081,18 @@ def last_value( Args: expression: Argument to perform bitwise calculation on filter: If provided, only compute against rows for which the filter is True - order_by: Set the ordering of the expression to evaluate + order_by: Set the ordering of the expression to evaluate. Accepts + column names or expressions. null_treatment: Assign whether to respect or ignore null values. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> result = df.aggregate( + ... [], [dfn.functions.last_value(dfn.col("a")).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 30 """ order_by_raw = sort_list_to_raw_sort_list(order_by) filter_raw = filter.expr if filter is not None else None @@ -2277,8 +4110,8 @@ def last_value( def nth_value( expression: Expr, n: int, - filter: Optional[Expr] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + filter: Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS, ) -> Expr: """Returns the n-th value in a group of values. @@ -2292,8 +4125,18 @@ def nth_value( expression: Argument to perform bitwise calculation on n: Index of value to return. Starts at 1. filter: If provided, only compute against rows for which the filter is True - order_by: Set the ordering of the expression to evaluate + order_by: Set the ordering of the expression to evaluate. Accepts + column names or expressions. null_treatment: Assign whether to respect or ignore null values. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> result = df.aggregate( + ... [], [dfn.functions.nth_value(dfn.col("a"), 2).alias("v")] + ... ) + >>> result.collect_column("v")[0].as_py() + 20 """ order_by_raw = sort_list_to_raw_sort_list(order_by) filter_raw = filter.expr if filter is not None else None @@ -2309,7 +4152,7 @@ def nth_value( ) -def bit_and(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def bit_and(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the bitwise AND of the argument. This aggregate function will bitwise compare every value in the input partition. @@ -2320,12 +4163,20 @@ def bit_and(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: Argument to perform bitwise calculation on filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [7, 3]}) + >>> result = df.aggregate([], [dfn.functions.bit_and(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3 """ filter_raw = filter.expr if filter is not None else None return Expr(f.bit_and(expression.expr, filter=filter_raw)) -def bit_or(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def bit_or(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the bitwise OR of the argument. This aggregate function will bitwise compare every value in the input partition. @@ -2336,13 +4187,21 @@ def bit_or(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: Argument to perform bitwise calculation on filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2]}) + >>> result = df.aggregate([], [dfn.functions.bit_or(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 3 """ filter_raw = filter.expr if filter is not None else None return Expr(f.bit_or(expression.expr, filter=filter_raw)) def bit_xor( - expression: Expr, distinct: bool = False, filter: Optional[Expr] = None + expression: Expr, distinct: bool = False, filter: Expr | None = None ) -> Expr: """Computes the bitwise XOR of the argument. @@ -2355,12 +4214,20 @@ def bit_xor( expression: Argument to perform bitwise calculation on distinct: If True, evaluate each unique value of expression only once filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [5, 3]}) + >>> result = df.aggregate([], [dfn.functions.bit_xor(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + 6 """ filter_raw = filter.expr if filter is not None else None return Expr(f.bit_xor(expression.expr, distinct=distinct, filter=filter_raw)) -def bool_and(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def bool_and(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the boolean AND of the argument. This aggregate function will compare every value in the input partition. These are @@ -2372,12 +4239,20 @@ def bool_and(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: Argument to perform calculation on filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [True, True, False]}) + >>> result = df.aggregate([], [dfn.functions.bool_and(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + False """ filter_raw = filter.expr if filter is not None else None return Expr(f.bool_and(expression.expr, filter=filter_raw)) -def bool_or(expression: Expr, filter: Optional[Expr] = None) -> Expr: +def bool_or(expression: Expr, filter: Expr | None = None) -> Expr: """Computes the boolean OR of the argument. This aggregate function will compare every value in the input partition. These are @@ -2389,6 +4264,14 @@ def bool_or(expression: Expr, filter: Optional[Expr] = None) -> Expr: Args: expression: Argument to perform calculation on filter: If provided, only compute against rows for which the filter is True + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [False, False, True]}) + >>> result = df.aggregate([], [dfn.functions.bool_or(dfn.col("a")).alias("v")]) + >>> result.collect_column("v")[0].as_py() + True """ filter_raw = filter.expr if filter is not None else None return Expr(f.bool_or(expression.expr, filter=filter_raw)) @@ -2397,16 +4280,16 @@ def bool_or(expression: Expr, filter: Optional[Expr] = None) -> Expr: def lead( arg: Expr, shift_offset: int = 1, - default_value: Optional[Any] = None, - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + default_value: Any | None = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a lead window function. Lead operation will return the argument that is in the next shift_offset-th row in the partition. For example ``lead(col("b"), shift_offset=3, default_value=5)`` will return the 3rd following value in column ``b``. At the end of the partition, where - no futher values can be returned it will return the default value of 5. + no further values can be returned it will return the default value of 5. Here is an example of both the ``lead`` and :py:func:`datafusion.functions.lag` functions on a simple DataFrame:: @@ -2428,14 +4311,22 @@ def lead( shift_offset: Number of rows following the current row. default_value: Value to return if shift_offet row does not exist. partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.lead(dfn.col("a"), shift_offset=1, + ... default_value=0, order_by="a").alias("lead")) + >>> result.sort(dfn.col("a")).collect_column("lead").to_pylist() + [2, 3, 0] """ if not isinstance(default_value, pa.Scalar) and default_value is not None: default_value = pa.scalar(default_value) - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( @@ -2443,7 +4334,7 @@ def lead( arg.expr, shift_offset, default_value, - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) @@ -2452,15 +4343,15 @@ def lead( def lag( arg: Expr, shift_offset: int = 1, - default_value: Optional[Any] = None, - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + default_value: Any | None = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a lag window function. Lag operation will return the argument that is in the previous shift_offset-th row in the partition. For example ``lag(col("b"), shift_offset=3, default_value=5)`` - will return the 3rd previous value in column ``b``. At the beginnig of the + will return the 3rd previous value in column ``b``. At the beginning of the partition, where no values can be returned it will return the default value of 5. Here is an example of both the ``lag`` and :py:func:`datafusion.functions.lead` @@ -2480,14 +4371,22 @@ def lag( shift_offset: Number of rows before the current row. default_value: Value to return if shift_offet row does not exist. partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.lag(dfn.col("a"), shift_offset=1, + ... default_value=0, order_by="a").alias("lag")) + >>> result.sort(dfn.col("a")).collect_column("lag").to_pylist() + [0, 1, 2] """ if not isinstance(default_value, pa.Scalar): default_value = pa.scalar(default_value) - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( @@ -2495,15 +4394,15 @@ def lag( arg.expr, shift_offset, default_value, - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) def row_number( - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a row number window function. @@ -2522,30 +4421,37 @@ def row_number( Args: partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.row_number(order_by="a").alias("rn")) + >>> result.sort(dfn.col("a")).collect_column("rn").to_pylist() + [1, 2, 3] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.row_number( - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) def rank( - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a rank window function. Returns the rank based upon the window order. Consecutive equal values will receive the same rank, but the next different value will not be consecutive but rather the - number of rows that preceed it plus one. This is similar to Olympic medals. If two + number of rows that precede it plus one. This is similar to Olympic medals. If two people tie for gold, the next place is bronze. There would be no silver medal. Here is an example of a dataframe with a window ordered by descending ``points`` and the associated rank. @@ -2563,24 +4469,32 @@ def rank( Args: partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 10, 20]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.rank(order_by="a").alias("rnk") + ... ) + >>> result.sort(dfn.col("a")).collect_column("rnk").to_pylist() + [1, 1, 3] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.rank( - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) def dense_rank( - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a dense_rank window function. @@ -2599,24 +4513,31 @@ def dense_rank( Args: partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 10, 20]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.dense_rank(order_by="a").alias("dr")) + >>> result.sort(dfn.col("a")).collect_column("dr").to_pylist() + [1, 1, 2] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.dense_rank( - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) def percent_rank( - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a percent_rank window function. @@ -2636,29 +4557,37 @@ def percent_rank( Args: partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.percent_rank(order_by="a").alias("pr")) + >>> result.sort(dfn.col("a")).collect_column("pr").to_pylist() + [0.0, 0.5, 1.0] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.percent_rank( - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) def cume_dist( - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a cumulative distribution window function. This window function is similar to :py:func:`rank` except that the returned values - are the ratio of the row number to the total numebr of rows. Here is an example of a + are the ratio of the row number to the total number of rows. Here is an example of a dataframe with a window ordered by descending ``points`` and the associated cumulative distribution:: @@ -2673,16 +4602,27 @@ def cume_dist( Args: partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [1., 2., 2., 3.]}) + >>> result = df.select( + ... dfn.col("a"), + ... dfn.functions.cume_dist( + ... order_by="a" + ... ).alias("cd") + ... ) + >>> result.collect_column("cd").to_pylist() + [0.25..., 0.75..., 0.75..., 1.0...] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.cume_dist( - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) @@ -2690,8 +4630,8 @@ def cume_dist( def ntile( groups: int, - partition_by: Optional[list[Expr]] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + partition_by: list[Expr] | Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Create a n-tile window function. @@ -2714,17 +4654,24 @@ def ntile( Args: groups: Number of groups for the n-tile to be divided into. partition_by: Expressions to partition the window frame on. - order_by: Set ordering within the window frame. + order_by: Set ordering within the window frame. Accepts + column names or expressions. + + Examples: + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": [10, 20, 30, 40]}) + >>> result = df.select( + ... dfn.col("a"), dfn.functions.ntile(2, order_by="a").alias("nt")) + >>> result.sort(dfn.col("a")).collect_column("nt").to_pylist() + [1, 1, 2, 2] """ - partition_cols = ( - [col.expr for col in partition_by] if partition_by is not None else None - ) + partition_by_raw = expr_list_to_raw_expr_list(partition_by) order_by_raw = sort_list_to_raw_sort_list(order_by) return Expr( f.ntile( Expr.literal(groups).expr, - partition_by=partition_cols, + partition_by=partition_by_raw, order_by=order_by_raw, ) ) @@ -2733,13 +4680,13 @@ def ntile( def string_agg( expression: Expr, delimiter: str, - filter: Optional[Expr] = None, - order_by: Optional[list[Expr | SortExpr]] = None, + filter: Expr | None = None, + order_by: list[SortKey] | SortKey | None = None, ) -> Expr: """Concatenates the input strings. This aggregate function will concatenate input strings, ignoring null values, and - seperating them with the specified delimiter. Non-string values will be converted to + separating them with the specified delimiter. Non-string values will be converted to their string equivalents. If using the builder functions described in ref:`_aggregation` this function ignores @@ -2749,7 +4696,17 @@ def string_agg( expression: Argument to perform bitwise calculation on delimiter: Text to place between each value of expression filter: If provided, only compute against rows for which the filter is True - order_by: Set the ordering of the expression to evaluate + order_by: Set the ordering of the expression to evaluate. Accepts + column names or expressions. + + Examples: + --------- + >>> ctx = dfn.SessionContext() + >>> df = ctx.from_pydict({"a": ["x", "y", "z"]}) + >>> result = df.aggregate( + ... [], [dfn.functions.string_agg(dfn.col("a"), ",", order_by="a").alias("s")]) + >>> result.collect_column("s")[0].as_py() + 'x,y,z' """ order_by_raw = sort_list_to_raw_sort_list(order_by) filter_raw = filter.expr if filter is not None else None diff --git a/python/datafusion/html_formatter.py b/python/datafusion/html_formatter.py new file mode 100644 index 000000000..65eb1f042 --- /dev/null +++ b/python/datafusion/html_formatter.py @@ -0,0 +1,29 @@ +# 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. + +"""Deprecated module for dataframe formatting.""" + +import warnings + +from datafusion.dataframe_formatter import * # noqa: F403 + +warnings.warn( + "The module 'html_formatter' is deprecated and will be removed in the next release." + "Please use 'dataframe_formatter' instead.", + DeprecationWarning, + stacklevel=3, +) diff --git a/python/datafusion/input/location.py b/python/datafusion/input/location.py index 08d98d115..b804ac18b 100644 --- a/python/datafusion/input/location.py +++ b/python/datafusion/input/location.py @@ -17,7 +17,6 @@ """The default input source for DataFusion.""" -import glob from pathlib import Path from typing import Any @@ -84,6 +83,7 @@ def build_table( raise RuntimeError(msg) # Input could possibly be multiple files. Create a list if so - input_files = glob.glob(input_item) + input_path = Path(input_item) + input_files = [str(p) for p in input_path.parent.glob(input_path.name)] return SqlTable(table_name, columns, num_rows, input_files) diff --git a/python/datafusion/io.py b/python/datafusion/io.py index ef5ebf96f..4f9c3c516 100644 --- a/python/datafusion/io.py +++ b/python/datafusion/io.py @@ -22,19 +22,21 @@ from typing import TYPE_CHECKING from datafusion.context import SessionContext -from datafusion.dataframe import DataFrame if TYPE_CHECKING: import pathlib import pyarrow as pa + from datafusion.dataframe import DataFrame from datafusion.expr import Expr + from .options import CsvReadOptions + def read_parquet( path: str | pathlib.Path, - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, parquet_pruning: bool = True, file_extension: str = ".parquet", skip_metadata: bool = True, @@ -83,7 +85,7 @@ def read_json( schema: pa.Schema | None = None, schema_infer_max_records: int = 1000, file_extension: str = ".json", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, ) -> DataFrame: """Read a line-delimited JSON data source. @@ -124,8 +126,9 @@ def read_csv( delimiter: str = ",", schema_infer_max_records: int = 1000, file_extension: str = ".csv", - table_partition_cols: list[tuple[str, str]] | None = None, + table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, + options: CsvReadOptions | None = None, ) -> DataFrame: """Read a CSV data source. @@ -147,15 +150,12 @@ def read_csv( selected for data input. table_partition_cols: Partition columns. file_compression_type: File compression type. + options: Set advanced options for CSV reading. This cannot be + combined with any of the other options in this method. Returns: DataFrame representation of the read CSV files """ - if table_partition_cols is None: - table_partition_cols = [] - - path = [str(p) for p in path] if isinstance(path, list) else str(path) - return SessionContext.global_ctx().read_csv( path, schema, @@ -165,13 +165,14 @@ def read_csv( file_extension, table_partition_cols, file_compression_type, + options, ) def read_avro( path: str | pathlib.Path, schema: pa.Schema | None = None, - file_partition_cols: list[tuple[str, str]] | None = None, + file_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_extension: str = ".avro", ) -> DataFrame: """Create a :py:class:`DataFrame` for reading Avro data source. diff --git a/python/datafusion/options.py b/python/datafusion/options.py new file mode 100644 index 000000000..ec19f37d0 --- /dev/null +++ b/python/datafusion/options.py @@ -0,0 +1,284 @@ +# 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. + +"""Options for reading various file formats.""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import pyarrow as pa + +from datafusion.expr import sort_list_to_raw_sort_list + +if TYPE_CHECKING: + from datafusion.expr import SortExpr + +from ._internal import options + +__all__ = ["CsvReadOptions"] + +DEFAULT_MAX_INFER_SCHEMA = 1000 + + +class CsvReadOptions: + """Options for reading CSV files. + + This class provides a builder pattern for configuring CSV reading options. + All methods starting with ``with_`` return ``self`` to allow method chaining. + """ + + def __init__( + self, + *, + has_header: bool = True, + delimiter: str = ",", + quote: str = '"', + terminator: str | None = None, + escape: str | None = None, + comment: str | None = None, + newlines_in_values: bool = False, + schema: pa.Schema | None = None, + schema_infer_max_records: int = DEFAULT_MAX_INFER_SCHEMA, + file_extension: str = ".csv", + table_partition_cols: list[tuple[str, pa.DataType]] | None = None, + file_compression_type: str = "", + file_sort_order: list[list[SortExpr]] | None = None, + null_regex: str | None = None, + truncated_rows: bool = False, + ) -> None: + """Initialize CsvReadOptions. + + Args: + has_header: Does the CSV file have a header row? If schema inference + is run on a file with no headers, default column names are created. + delimiter: Column delimiter character. Must be a single ASCII character. + quote: Quote character for fields containing delimiters or newlines. + Must be a single ASCII character. + terminator: Optional line terminator character. If ``None``, uses CRLF. + Must be a single ASCII character. + escape: Optional escape character for quotes. Must be a single ASCII + character. + comment: If specified, lines beginning with this character are ignored. + Must be a single ASCII character. + newlines_in_values: Whether newlines in quoted values are supported. + Parsing newlines in quoted values may be affected by execution + behavior such as parallel file scanning. Setting this to ``True`` + ensures that newlines in values are parsed successfully, which may + reduce performance. + schema: Optional PyArrow schema representing the CSV files. If ``None``, + the CSV reader will try to infer it based on data in the file. + schema_infer_max_records: Maximum number of rows to read from CSV files + for schema inference if needed. + file_extension: File extension; only files with this extension are + selected for data input. + table_partition_cols: Partition columns as a list of tuples of + (column_name, data_type). + file_compression_type: File compression type. Supported values are + ``"gzip"``, ``"bz2"``, ``"xz"``, ``"zstd"``, or empty string for + uncompressed. + file_sort_order: Optional sort order of the files as a list of sort + expressions per file. + null_regex: Optional regex pattern to match null values in the CSV. + truncated_rows: Whether to allow truncated rows when parsing. By default + this is ``False`` and will error if the CSV rows have different + lengths. When set to ``True``, it will allow records with less than + the expected number of columns and fill the missing columns with + nulls. If the record's schema is not nullable, it will still return + an error. + """ + validate_single_character("delimiter", delimiter) + validate_single_character("quote", quote) + validate_single_character("terminator", terminator) + validate_single_character("escape", escape) + validate_single_character("comment", comment) + + self.has_header = has_header + self.delimiter = delimiter + self.quote = quote + self.terminator = terminator + self.escape = escape + self.comment = comment + self.newlines_in_values = newlines_in_values + self.schema = schema + self.schema_infer_max_records = schema_infer_max_records + self.file_extension = file_extension + self.table_partition_cols = table_partition_cols or [] + self.file_compression_type = file_compression_type + self.file_sort_order = file_sort_order or [] + self.null_regex = null_regex + self.truncated_rows = truncated_rows + + def with_has_header(self, has_header: bool) -> CsvReadOptions: + """Configure whether the CSV has a header row.""" + self.has_header = has_header + return self + + def with_delimiter(self, delimiter: str) -> CsvReadOptions: + """Configure the column delimiter.""" + self.delimiter = delimiter + return self + + def with_quote(self, quote: str) -> CsvReadOptions: + """Configure the quote character.""" + self.quote = quote + return self + + def with_terminator(self, terminator: str | None) -> CsvReadOptions: + """Configure the line terminator character.""" + self.terminator = terminator + return self + + def with_escape(self, escape: str | None) -> CsvReadOptions: + """Configure the escape character.""" + self.escape = escape + return self + + def with_comment(self, comment: str | None) -> CsvReadOptions: + """Configure the comment character.""" + self.comment = comment + return self + + def with_newlines_in_values(self, newlines_in_values: bool) -> CsvReadOptions: + """Configure whether newlines in values are supported.""" + self.newlines_in_values = newlines_in_values + return self + + def with_schema(self, schema: pa.Schema | None) -> CsvReadOptions: + """Configure the schema.""" + self.schema = schema + return self + + def with_schema_infer_max_records( + self, schema_infer_max_records: int + ) -> CsvReadOptions: + """Configure maximum records for schema inference.""" + self.schema_infer_max_records = schema_infer_max_records + return self + + def with_file_extension(self, file_extension: str) -> CsvReadOptions: + """Configure the file extension filter.""" + self.file_extension = file_extension + return self + + def with_table_partition_cols( + self, table_partition_cols: list[tuple[str, pa.DataType]] + ) -> CsvReadOptions: + """Configure table partition columns.""" + self.table_partition_cols = table_partition_cols + return self + + def with_file_compression_type(self, file_compression_type: str) -> CsvReadOptions: + """Configure file compression type.""" + self.file_compression_type = file_compression_type + return self + + def with_file_sort_order( + self, file_sort_order: list[list[SortExpr]] + ) -> CsvReadOptions: + """Configure file sort order.""" + self.file_sort_order = file_sort_order + return self + + def with_null_regex(self, null_regex: str | None) -> CsvReadOptions: + """Configure null value regex pattern.""" + self.null_regex = null_regex + return self + + def with_truncated_rows(self, truncated_rows: bool) -> CsvReadOptions: + """Configure whether to allow truncated rows.""" + self.truncated_rows = truncated_rows + return self + + def to_inner(self) -> options.CsvReadOptions: + """Convert this object into the underlying Rust structure. + + This is intended for internal use only. + """ + file_sort_order = ( + [] + if self.file_sort_order is None + else [ + sort_list_to_raw_sort_list(sort_list) + for sort_list in self.file_sort_order + ] + ) + + return options.CsvReadOptions( + has_header=self.has_header, + delimiter=ord(self.delimiter[0]) if self.delimiter else ord(","), + quote=ord(self.quote[0]) if self.quote else ord('"'), + terminator=ord(self.terminator[0]) if self.terminator else None, + escape=ord(self.escape[0]) if self.escape else None, + comment=ord(self.comment[0]) if self.comment else None, + newlines_in_values=self.newlines_in_values, + schema=self.schema, + schema_infer_max_records=self.schema_infer_max_records, + file_extension=self.file_extension, + table_partition_cols=_convert_table_partition_cols( + self.table_partition_cols + ), + file_compression_type=self.file_compression_type or "", + file_sort_order=file_sort_order, + null_regex=self.null_regex, + truncated_rows=self.truncated_rows, + ) + + +def validate_single_character(name: str, value: str | None) -> None: + if value is not None and len(value) != 1: + message = f"{name} must be a single character" + raise ValueError(message) + + +def _convert_table_partition_cols( + table_partition_cols: list[tuple[str, str | pa.DataType]], +) -> list[tuple[str, pa.DataType]]: + warn = False + converted_table_partition_cols = [] + + for col, data_type in table_partition_cols: + if isinstance(data_type, str): + warn = True + if data_type == "string": + converted_data_type = pa.string() + elif data_type == "int": + converted_data_type = pa.int32() + else: + message = ( + f"Unsupported literal data type '{data_type}' for partition " + "column. Supported types are 'string' and 'int'" + ) + raise ValueError(message) + else: + converted_data_type = data_type + + converted_table_partition_cols.append((col, converted_data_type)) + + if warn: + message = ( + "using literals for table_partition_cols data types is deprecated," + "use pyarrow types instead" + ) + warnings.warn( + message, + category=DeprecationWarning, + stacklevel=2, + ) + + return converted_table_partition_cols diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 0b7bebcb3..fb54fd624 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -98,6 +98,12 @@ def to_proto(self) -> bytes: """ return self._raw_plan.to_proto() + def __eq__(self, other: LogicalPlan) -> bool: + """Test equality.""" + if not isinstance(other, LogicalPlan): + return False + return self._raw_plan.__eq__(other._raw_plan) + class ExecutionPlan: """Represent nodes in the DataFusion Physical Plan.""" diff --git a/python/datafusion/record_batch.py b/python/datafusion/record_batch.py index 556eaa786..c24cde0ac 100644 --- a/python/datafusion/record_batch.py +++ b/python/datafusion/record_batch.py @@ -46,6 +46,26 @@ def to_pyarrow(self) -> pa.RecordBatch: """Convert to :py:class:`pa.RecordBatch`.""" return self.record_batch.to_pyarrow() + def __arrow_c_array__( + self, requested_schema: object | None = None + ) -> tuple[object, object]: + """Export the record batch via the Arrow C Data Interface. + + This allows zero-copy interchange with libraries that support the + `Arrow PyCapsule interface `_. + + Args: + requested_schema: Attempt to provide the record batch using this + schema. Only straightforward projections such as column + selection or reordering are applied. + + Returns: + Two Arrow PyCapsule objects representing the ``ArrowArray`` and + ``ArrowSchema``. + """ + return self.record_batch.__arrow_c_array__(requested_schema) + class RecordBatchStream: """This class represents a stream of record batches. @@ -63,19 +83,19 @@ def next(self) -> RecordBatch: return next(self) async def __anext__(self) -> RecordBatch: - """Async iterator function.""" + """Return the next :py:class:`RecordBatch` in the stream asynchronously.""" next_batch = await self.rbs.__anext__() return RecordBatch(next_batch) def __next__(self) -> RecordBatch: - """Iterator function.""" + """Return the next :py:class:`RecordBatch` in the stream.""" next_batch = next(self.rbs) return RecordBatch(next_batch) def __aiter__(self) -> typing_extensions.Self: - """Async iterator function.""" + """Return an asynchronous iterator over record batches.""" return self def __iter__(self) -> typing_extensions.Self: - """Iterator function.""" + """Return an iterator over record batches.""" return self diff --git a/python/datafusion/substrait.py b/python/datafusion/substrait.py index f10adfb0c..3115238fa 100644 --- a/python/datafusion/substrait.py +++ b/python/datafusion/substrait.py @@ -67,6 +67,26 @@ def encode(self) -> bytes: """ return self.plan_internal.encode() + def to_json(self) -> str: + """Get the JSON representation of the Substrait plan. + + Returns: + A JSON representation of the Substrait plan. + """ + return self.plan_internal.to_json() + + @staticmethod + def from_json(json: str) -> Plan: + """Parse a plan from a JSON string representation. + + Args: + json: JSON representation of a Substrait plan. + + Returns: + Plan object representing the Substrait plan. + """ + return Plan(substrait_internal.Plan.from_json(json)) + @deprecated("Use `Plan` instead.") class plan(Plan): # noqa: N801 diff --git a/python/datafusion/udf.py b/python/datafusion/udf.py index e93a34ca5..c7265fa09 100644 --- a/python/datafusion/udf.py +++ b/python/datafusion/udf.py @@ -15,753 +15,15 @@ # specific language governing permissions and limitations # under the License. -"""Provides the user-defined functions for evaluation of dataframes.""" +"""Deprecated module for user defined functions.""" -from __future__ import annotations +import warnings -import functools -from abc import ABCMeta, abstractmethod -from enum import Enum -from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, overload +from datafusion.user_defined import * # noqa: F403 -import pyarrow as pa - -import datafusion._internal as df_internal -from datafusion.expr import Expr - -if TYPE_CHECKING: - _R = TypeVar("_R", bound=pa.DataType) - - -class Volatility(Enum): - """Defines how stable or volatile a function is. - - When setting the volatility of a function, you can either pass this - enumeration or a ``str``. The ``str`` equivalent is the lower case value of the - name (`"immutable"`, `"stable"`, or `"volatile"`). - """ - - Immutable = 1 - """An immutable function will always return the same output when given the - same input. - - DataFusion will attempt to inline immutable functions during planning. - """ - - Stable = 2 - """ - Returns the same value for a given input within a single queries. - - A stable function may return different values given the same input across - different queries but must return the same value for a given input within a - query. An example of this is the ``Now`` function. DataFusion will attempt to - inline ``Stable`` functions during planning, when possible. For query - ``select col1, now() from t1``, it might take a while to execute but ``now()`` - column will be the same for each output row, which is evaluated during - planning. - """ - - Volatile = 3 - """A volatile function may change the return value from evaluation to - evaluation. - - Multiple invocations of a volatile function may return different results - when used in the same query. An example of this is the random() function. - DataFusion can not evaluate such functions during planning. In the query - ``select col1, random() from t1``, ``random()`` function will be evaluated - for each output row, resulting in a unique random value for each row. - """ - - def __str__(self) -> str: - """Returns the string equivalent.""" - return self.name.lower() - - -class ScalarUDF: - """Class for performing scalar user-defined functions (UDF). - - Scalar UDFs operate on a row by row basis. See also :py:class:`AggregateUDF` for - operating on a group of rows. - """ - - def __init__( - self, - name: str, - func: Callable[..., _R], - input_types: pa.DataType | list[pa.DataType], - return_type: _R, - volatility: Volatility | str, - ) -> None: - """Instantiate a scalar user-defined function (UDF). - - See helper method :py:func:`udf` for argument details. - """ - if isinstance(input_types, pa.DataType): - input_types = [input_types] - self._udf = df_internal.ScalarUDF( - name, func, input_types, return_type, str(volatility) - ) - - def __call__(self, *args: Expr) -> Expr: - """Execute the UDF. - - This function is not typically called by an end user. These calls will - occur during the evaluation of the dataframe. - """ - args_raw = [arg.expr for arg in args] - return Expr(self._udf.__call__(*args_raw)) - - @overload - @staticmethod - def udf( - input_types: list[pa.DataType], - return_type: _R, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable[..., ScalarUDF]: ... - - @overload - @staticmethod - def udf( - func: Callable[..., _R], - input_types: list[pa.DataType], - return_type: _R, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> ScalarUDF: ... - - @staticmethod - def udf(*args: Any, **kwargs: Any): # noqa: D417 - """Create a new User-Defined Function (UDF). - - This class can be used both as a **function** and as a **decorator**. - - Usage: - - **As a function**: Call `udf(func, input_types, return_type, volatility, - name)`. - - **As a decorator**: Use `@udf(input_types, return_type, volatility, - name)`. In this case, do **not** pass `func` explicitly. - - Args: - func (Callable, optional): **Only needed when calling as a function.** - Skip this argument when using `udf` as a decorator. - input_types (list[pa.DataType]): The data types of the arguments - to `func`. This list must be of the same length as the number of - arguments. - return_type (_R): The data type of the return value from the function. - volatility (Volatility | str): See `Volatility` for allowed values. - name (Optional[str]): A descriptive name for the function. - - Returns: - A user-defined function that can be used in SQL expressions, - data aggregation, or window function calls. - - Example: - **Using `udf` as a function:** - ``` - def double_func(x): - return x * 2 - double_udf = udf(double_func, [pa.int32()], pa.int32(), - "volatile", "double_it") - ``` - - **Using `udf` as a decorator:** - ``` - @udf([pa.int32()], pa.int32(), "volatile", "double_it") - def double_udf(x): - return x * 2 - ``` - """ - - def _function( - func: Callable[..., _R], - input_types: list[pa.DataType], - return_type: _R, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> ScalarUDF: - if not callable(func): - msg = "`func` argument must be callable" - raise TypeError(msg) - if name is None: - if hasattr(func, "__qualname__"): - name = func.__qualname__.lower() - else: - name = func.__class__.__name__.lower() - return ScalarUDF( - name=name, - func=func, - input_types=input_types, - return_type=return_type, - volatility=volatility, - ) - - def _decorator( - input_types: list[pa.DataType], - return_type: _R, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable: - def decorator(func: Callable): - udf_caller = ScalarUDF.udf( - func, input_types, return_type, volatility, name - ) - - @functools.wraps(func) - def wrapper(*args: Any, **kwargs: Any): - return udf_caller(*args, **kwargs) - - return wrapper - - return decorator - - if args and callable(args[0]): - # Case 1: Used as a function, require the first parameter to be callable - return _function(*args, **kwargs) - # Case 2: Used as a decorator with parameters - return _decorator(*args, **kwargs) - - -class Accumulator(metaclass=ABCMeta): - """Defines how an :py:class:`AggregateUDF` accumulates values.""" - - @abstractmethod - def state(self) -> list[pa.Scalar]: - """Return the current state.""" - - @abstractmethod - def update(self, *values: pa.Array) -> None: - """Evaluate an array of values and update state.""" - - @abstractmethod - def merge(self, states: list[pa.Array]) -> None: - """Merge a set of states.""" - - @abstractmethod - def evaluate(self) -> pa.Scalar: - """Return the resultant value.""" - - -class AggregateUDF: - """Class for performing scalar user-defined functions (UDF). - - Aggregate UDFs operate on a group of rows and return a single value. See - also :py:class:`ScalarUDF` for operating on a row by row basis. - """ - - def __init__( - self, - name: str, - accumulator: Callable[[], Accumulator], - input_types: list[pa.DataType], - return_type: pa.DataType, - state_type: list[pa.DataType], - volatility: Volatility | str, - ) -> None: - """Instantiate a user-defined aggregate function (UDAF). - - See :py:func:`udaf` for a convenience function and argument - descriptions. - """ - self._udaf = df_internal.AggregateUDF( - name, - accumulator, - input_types, - return_type, - state_type, - str(volatility), - ) - - def __call__(self, *args: Expr) -> Expr: - """Execute the UDAF. - - This function is not typically called by an end user. These calls will - occur during the evaluation of the dataframe. - """ - args_raw = [arg.expr for arg in args] - return Expr(self._udaf.__call__(*args_raw)) - - @overload - @staticmethod - def udaf( - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - state_type: list[pa.DataType], - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable[..., AggregateUDF]: ... - - @overload - @staticmethod - def udaf( - accum: Callable[[], Accumulator], - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - state_type: list[pa.DataType], - volatility: Volatility | str, - name: Optional[str] = None, - ) -> AggregateUDF: ... - - @staticmethod - def udaf(*args: Any, **kwargs: Any): # noqa: D417 - """Create a new User-Defined Aggregate Function (UDAF). - - This class allows you to define an **aggregate function** that can be used in - data aggregation or window function calls. - - Usage: - - **As a function**: Call `udaf(accum, input_types, return_type, state_type, - volatility, name)`. - - **As a decorator**: Use `@udaf(input_types, return_type, state_type, - volatility, name)`. - When using `udaf` as a decorator, **do not pass `accum` explicitly**. - - **Function example:** - - If your `:py:class:Accumulator` can be instantiated with no arguments, you - can simply pass it's type as `accum`. If you need to pass additional - arguments to it's constructor, you can define a lambda or a factory method. - During runtime the `:py:class:Accumulator` will be constructed for every - instance in which this UDAF is used. The following examples are all valid. - ``` - import pyarrow as pa - import pyarrow.compute as pc - - class Summarize(Accumulator): - def __init__(self, bias: float = 0.0): - self._sum = pa.scalar(bias) - - def state(self) -> list[pa.Scalar]: - return [self._sum] - - def update(self, values: pa.Array) -> None: - self._sum = pa.scalar(self._sum.as_py() + pc.sum(values).as_py()) - - def merge(self, states: list[pa.Array]) -> None: - self._sum = pa.scalar(self._sum.as_py() + pc.sum(states[0]).as_py()) - - def evaluate(self) -> pa.Scalar: - return self._sum - - def sum_bias_10() -> Summarize: - return Summarize(10.0) - - udaf1 = udaf(Summarize, pa.float64(), pa.float64(), [pa.float64()], - "immutable") - udaf2 = udaf(sum_bias_10, pa.float64(), pa.float64(), [pa.float64()], - "immutable") - udaf3 = udaf(lambda: Summarize(20.0), pa.float64(), pa.float64(), - [pa.float64()], "immutable") - ``` - - **Decorator example:** - ``` - @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable") - def udf4() -> Summarize: - return Summarize(10.0) - ``` - - Args: - accum: The accumulator python function. **Only needed when calling as a - function. Skip this argument when using `udaf` as a decorator.** - input_types: The data types of the arguments to ``accum``. - return_type: The data type of the return value. - state_type: The data types of the intermediate accumulation. - volatility: See :py:class:`Volatility` for allowed values. - name: A descriptive name for the function. - - Returns: - A user-defined aggregate function, which can be used in either data - aggregation or window function calls. - """ - - def _function( - accum: Callable[[], Accumulator], - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - state_type: list[pa.DataType], - volatility: Volatility | str, - name: Optional[str] = None, - ) -> AggregateUDF: - if not callable(accum): - msg = "`func` must be callable." - raise TypeError(msg) - if not isinstance(accum(), Accumulator): - msg = "Accumulator must implement the abstract base class Accumulator" - raise TypeError(msg) - if name is None: - name = accum().__class__.__qualname__.lower() - if isinstance(input_types, pa.DataType): - input_types = [input_types] - return AggregateUDF( - name=name, - accumulator=accum, - input_types=input_types, - return_type=return_type, - state_type=state_type, - volatility=volatility, - ) - - def _decorator( - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - state_type: list[pa.DataType], - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable[..., Callable[..., Expr]]: - def decorator(accum: Callable[[], Accumulator]) -> Callable[..., Expr]: - udaf_caller = AggregateUDF.udaf( - accum, input_types, return_type, state_type, volatility, name - ) - - @functools.wraps(accum) - def wrapper(*args: Any, **kwargs: Any) -> Expr: - return udaf_caller(*args, **kwargs) - - return wrapper - - return decorator - - if args and callable(args[0]): - # Case 1: Used as a function, require the first parameter to be callable - return _function(*args, **kwargs) - # Case 2: Used as a decorator with parameters - return _decorator(*args, **kwargs) - - -class WindowEvaluator: - """Evaluator class for user-defined window functions (UDWF). - - It is up to the user to decide which evaluate function is appropriate. - - +------------------------+--------------------------------+------------------+---------------------------+ - | ``uses_window_frame`` | ``supports_bounded_execution`` | ``include_rank`` | function_to_implement | - +========================+================================+==================+===========================+ - | False (default) | False (default) | False (default) | ``evaluate_all`` | - +------------------------+--------------------------------+------------------+---------------------------+ - | False | True | False | ``evaluate`` | - +------------------------+--------------------------------+------------------+---------------------------+ - | False | True/False | True | ``evaluate_all_with_rank``| - +------------------------+--------------------------------+------------------+---------------------------+ - | True | True/False | True/False | ``evaluate`` | - +------------------------+--------------------------------+------------------+---------------------------+ - """ # noqa: W505, E501 - - def memoize(self) -> None: - """Perform a memoize operation to improve performance. - - When the window frame has a fixed beginning (e.g UNBOUNDED - PRECEDING), some functions such as FIRST_VALUE and - NTH_VALUE do not need the (unbounded) input once they have - seen a certain amount of input. - - `memoize` is called after each input batch is processed, and - such functions can save whatever they need - """ - - def get_range(self, idx: int, num_rows: int) -> tuple[int, int]: # noqa: ARG002 - """Return the range for the window fuction. - - If `uses_window_frame` flag is `false`. This method is used to - calculate required range for the window function during - stateful execution. - - Generally there is no required range, hence by default this - returns smallest range(current row). e.g seeing current row is - enough to calculate window result (such as row_number, rank, - etc) - - Args: - idx:: Current index - num_rows: Number of rows. - """ - return (idx, idx + 1) - - def is_causal(self) -> bool: - """Get whether evaluator needs future data for its result.""" - return False - - def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array: - """Evaluate a window function on an entire input partition. - - This function is called once per input *partition* for window functions that - *do not use* values from the window frame, such as - :py:func:`~datafusion.functions.row_number`, - :py:func:`~datafusion.functions.rank`, - :py:func:`~datafusion.functions.dense_rank`, - :py:func:`~datafusion.functions.percent_rank`, - :py:func:`~datafusion.functions.cume_dist`, - :py:func:`~datafusion.functions.lead`, - and :py:func:`~datafusion.functions.lag`. - - It produces the result of all rows in a single pass. It - expects to receive the entire partition as the ``value`` and - must produce an output column with one output row for every - input row. - - ``num_rows`` is required to correctly compute the output in case - ``len(values) == 0`` - - Implementing this function is an optimization. Certain window - functions are not affected by the window frame definition or - the query doesn't have a frame, and ``evaluate`` skips the - (costly) window frame boundary calculation and the overhead of - calling ``evaluate`` for each output row. - - For example, the `LAG` built in window function does not use - the values of its window frame (it can be computed in one shot - on the entire partition with ``Self::evaluate_all`` regardless of the - window defined in the ``OVER`` clause) - - .. code-block:: text - - lag(x, 1) OVER (ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING) - - However, ``avg()`` computes the average in the window and thus - does use its window frame. - - .. code-block:: text - - avg(x) OVER (PARTITION BY y ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING) - """ # noqa: W505, E501 - - def evaluate( - self, values: list[pa.Array], eval_range: tuple[int, int] - ) -> pa.Scalar: - """Evaluate window function on a range of rows in an input partition. - - This is the simplest and most general function to implement - but also the least performant as it creates output one row at - a time. It is typically much faster to implement stateful - evaluation using one of the other specialized methods on this - trait. - - Returns a [`ScalarValue`] that is the value of the window - function within `range` for the entire partition. Argument - `values` contains the evaluation result of function arguments - and evaluation results of ORDER BY expressions. If function has a - single argument, `values[1..]` will contain ORDER BY expression results. - """ - - def evaluate_all_with_rank( - self, num_rows: int, ranks_in_partition: list[tuple[int, int]] - ) -> pa.Array: - """Called for window functions that only need the rank of a row. - - Evaluate the partition evaluator against the partition using - the row ranks. For example, ``rank(col("a"))`` produces - - .. code-block:: text - - a | rank - - + ---- - A | 1 - A | 1 - C | 3 - D | 4 - D | 4 - - For this case, `num_rows` would be `5` and the - `ranks_in_partition` would be called with - - .. code-block:: text - - [ - (0,1), - (2,2), - (3,4), - ] - - The user must implement this method if ``include_rank`` returns True. - """ - - def supports_bounded_execution(self) -> bool: - """Can the window function be incrementally computed using bounded memory?""" - return False - - def uses_window_frame(self) -> bool: - """Does the window function use the values from the window frame?""" - return False - - def include_rank(self) -> bool: - """Can this function be evaluated with (only) rank?""" - return False - - -class WindowUDF: - """Class for performing window user-defined functions (UDF). - - Window UDFs operate on a partition of rows. See - also :py:class:`ScalarUDF` for operating on a row by row basis. - """ - - def __init__( - self, - name: str, - func: Callable[[], WindowEvaluator], - input_types: list[pa.DataType], - return_type: pa.DataType, - volatility: Volatility | str, - ) -> None: - """Instantiate a user-defined window function (UDWF). - - See :py:func:`udwf` for a convenience function and argument - descriptions. - """ - self._udwf = df_internal.WindowUDF( - name, func, input_types, return_type, str(volatility) - ) - - def __call__(self, *args: Expr) -> Expr: - """Execute the UDWF. - - This function is not typically called by an end user. These calls will - occur during the evaluation of the dataframe. - """ - args_raw = [arg.expr for arg in args] - return Expr(self._udwf.__call__(*args_raw)) - - @overload - @staticmethod - def udwf( - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable[..., WindowUDF]: ... - - @overload - @staticmethod - def udwf( - func: Callable[[], WindowEvaluator], - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> WindowUDF: ... - - @staticmethod - def udwf(*args: Any, **kwargs: Any): # noqa: D417 - """Create a new User-Defined Window Function (UDWF). - - This class can be used both as a **function** and as a **decorator**. - - Usage: - - **As a function**: Call `udwf(func, input_types, return_type, volatility, - name)`. - - **As a decorator**: Use `@udwf(input_types, return_type, volatility, - name)`. When using `udwf` as a decorator, **do not pass `func` - explicitly**. - - **Function example:** - ``` - import pyarrow as pa - - class BiasedNumbers(WindowEvaluator): - def __init__(self, start: int = 0) -> None: - self.start = start - - def evaluate_all(self, values: list[pa.Array], - num_rows: int) -> pa.Array: - return pa.array([self.start + i for i in range(num_rows)]) - - def bias_10() -> BiasedNumbers: - return BiasedNumbers(10) - - udwf1 = udwf(BiasedNumbers, pa.int64(), pa.int64(), "immutable") - udwf2 = udwf(bias_10, pa.int64(), pa.int64(), "immutable") - udwf3 = udwf(lambda: BiasedNumbers(20), pa.int64(), pa.int64(), "immutable") - - ``` - - **Decorator example:** - ``` - @udwf(pa.int64(), pa.int64(), "immutable") - def biased_numbers() -> BiasedNumbers: - return BiasedNumbers(10) - ``` - - Args: - func: **Only needed when calling as a function. Skip this argument when - using `udwf` as a decorator.** - input_types: The data types of the arguments. - return_type: The data type of the return value. - volatility: See :py:class:`Volatility` for allowed values. - name: A descriptive name for the function. - - Returns: - A user-defined window function that can be used in window function calls. - """ - if args and callable(args[0]): - # Case 1: Used as a function, require the first parameter to be callable - return WindowUDF._create_window_udf(*args, **kwargs) - # Case 2: Used as a decorator with parameters - return WindowUDF._create_window_udf_decorator(*args, **kwargs) - - @staticmethod - def _create_window_udf( - func: Callable[[], WindowEvaluator], - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> WindowUDF: - """Create a WindowUDF instance from function arguments.""" - if not callable(func): - msg = "`func` must be callable." - raise TypeError(msg) - if not isinstance(func(), WindowEvaluator): - msg = "`func` must implement the abstract base class WindowEvaluator" - raise TypeError(msg) - - name = name or func.__qualname__.lower() - input_types = ( - [input_types] if isinstance(input_types, pa.DataType) else input_types - ) - - return WindowUDF(name, func, input_types, return_type, volatility) - - @staticmethod - def _get_default_name(func: Callable) -> str: - """Get the default name for a function based on its attributes.""" - if hasattr(func, "__qualname__"): - return func.__qualname__.lower() - return func.__class__.__name__.lower() - - @staticmethod - def _normalize_input_types( - input_types: pa.DataType | list[pa.DataType], - ) -> list[pa.DataType]: - """Convert a single DataType to a list if needed.""" - if isinstance(input_types, pa.DataType): - return [input_types] - return input_types - - @staticmethod - def _create_window_udf_decorator( - input_types: pa.DataType | list[pa.DataType], - return_type: pa.DataType, - volatility: Volatility | str, - name: Optional[str] = None, - ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]: - """Create a decorator for a WindowUDF.""" - - def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]: - udwf_caller = WindowUDF._create_window_udf( - func, input_types, return_type, volatility, name - ) - - @functools.wraps(func) - def wrapper(*args: Any, **kwargs: Any) -> Expr: - return udwf_caller(*args, **kwargs) - - return wrapper - - return decorator - - -# Convenience exports so we can import instead of treating as -# variables at the package root -udf = ScalarUDF.udf -udaf = AggregateUDF.udaf -udwf = WindowUDF.udwf +warnings.warn( + "The module 'udf' is deprecated and will be removed in the next release. " + "Please use 'user_defined' instead.", + DeprecationWarning, + stacklevel=2, +) diff --git a/python/datafusion/unparser.py b/python/datafusion/unparser.py new file mode 100644 index 000000000..7ca5b9190 --- /dev/null +++ b/python/datafusion/unparser.py @@ -0,0 +1,80 @@ +# 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 provides support for unparsing datafusion plans to SQL. + +For additional information about unparsing, see https://docs.rs/datafusion-sql/latest/datafusion_sql/unparser/index.html +""" + +from ._internal import unparser as unparser_internal +from .plan import LogicalPlan + + +class Dialect: + """DataFusion data catalog.""" + + def __init__(self, dialect: unparser_internal.Dialect) -> None: + """This constructor is not typically called by the end user.""" + self.dialect = dialect + + @staticmethod + def default() -> "Dialect": + """Create a new default dialect.""" + return Dialect(unparser_internal.Dialect.default()) + + @staticmethod + def mysql() -> "Dialect": + """Create a new MySQL dialect.""" + return Dialect(unparser_internal.Dialect.mysql()) + + @staticmethod + def postgres() -> "Dialect": + """Create a new PostgreSQL dialect.""" + return Dialect(unparser_internal.Dialect.postgres()) + + @staticmethod + def sqlite() -> "Dialect": + """Create a new SQLite dialect.""" + return Dialect(unparser_internal.Dialect.sqlite()) + + @staticmethod + def duckdb() -> "Dialect": + """Create a new DuckDB dialect.""" + return Dialect(unparser_internal.Dialect.duckdb()) + + +class Unparser: + """DataFusion unparser.""" + + def __init__(self, dialect: Dialect) -> None: + """This constructor is not typically called by the end user.""" + self.unparser = unparser_internal.Unparser(dialect.dialect) + + def plan_to_sql(self, plan: LogicalPlan) -> str: + """Convert a logical plan to a SQL string.""" + return self.unparser.plan_to_sql(plan._raw_plan) + + def with_pretty(self, pretty: bool) -> "Unparser": + """Set the pretty flag.""" + self.unparser = self.unparser.with_pretty(pretty) + return self + + +__all__ = [ + "Dialect", + "Unparser", +] diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py new file mode 100644 index 000000000..eef23e741 --- /dev/null +++ b/python/datafusion/user_defined.py @@ -0,0 +1,1044 @@ +# 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. + +"""Provides the user-defined functions for evaluation of dataframes.""" + +from __future__ import annotations + +import functools +from abc import ABCMeta, abstractmethod +from enum import Enum +from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, TypeVar, cast, overload + +import pyarrow as pa + +import datafusion._internal as df_internal +from datafusion import SessionContext +from datafusion.expr import Expr + +if TYPE_CHECKING: + from _typeshed import CapsuleType as _PyCapsule + + _R = TypeVar("_R", bound=pa.DataType) + from collections.abc import Callable, Sequence + + +class Volatility(Enum): + """Defines how stable or volatile a function is. + + When setting the volatility of a function, you can either pass this + enumeration or a ``str``. The ``str`` equivalent is the lower case value of the + name (`"immutable"`, `"stable"`, or `"volatile"`). + """ + + Immutable = 1 + """An immutable function will always return the same output when given the + same input. + + DataFusion will attempt to inline immutable functions during planning. + """ + + Stable = 2 + """ + Returns the same value for a given input within a single queries. + + A stable function may return different values given the same input across + different queries but must return the same value for a given input within a + query. An example of this is the ``Now`` function. DataFusion will attempt to + inline ``Stable`` functions during planning, when possible. For query + ``select col1, now() from t1``, it might take a while to execute but ``now()`` + column will be the same for each output row, which is evaluated during + planning. + """ + + Volatile = 3 + """A volatile function may change the return value from evaluation to + evaluation. + + Multiple invocations of a volatile function may return different results + when used in the same query. An example of this is the random() function. + DataFusion can not evaluate such functions during planning. In the query + ``select col1, random() from t1``, ``random()`` function will be evaluated + for each output row, resulting in a unique random value for each row. + """ + + def __str__(self) -> str: + """Returns the string equivalent.""" + return self.name.lower() + + +def data_type_or_field_to_field(value: pa.DataType | pa.Field, name: str) -> pa.Field: + """Helper function to return a Field from either a Field or DataType.""" + if isinstance(value, pa.Field): + return value + return pa.field(name, type=value) + + +def data_types_or_fields_to_field_list( + inputs: Sequence[pa.Field | pa.DataType] | pa.Field | pa.DataType, +) -> list[pa.Field]: + """Helper function to return a list of Fields.""" + if isinstance(inputs, pa.DataType): + return [pa.field("value", type=inputs)] + if isinstance(inputs, pa.Field): + return [inputs] + + return [ + data_type_or_field_to_field(v, f"value_{idx}") for (idx, v) in enumerate(inputs) + ] + + +class ScalarUDFExportable(Protocol): + """Type hint for object that has __datafusion_scalar_udf__ PyCapsule.""" + + def __datafusion_scalar_udf__(self) -> object: ... # noqa: D105 + + +def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: + """Return ``True`` when ``value`` is a CPython ``PyCapsule``.""" + return value.__class__.__name__ == "PyCapsule" + + +class ScalarUDF: + """Class for performing scalar user-defined functions (UDF). + + Scalar UDFs operate on a row by row basis. See also :py:class:`AggregateUDF` for + operating on a group of rows. + """ + + def __init__( + self, + name: str, + func: Callable[..., _R], + input_fields: list[pa.Field], + return_field: _R, + volatility: Volatility | str, + ) -> None: + """Instantiate a scalar user-defined function (UDF). + + See helper method :py:func:`udf` for argument details. + """ + if hasattr(func, "__datafusion_scalar_udf__"): + self._udf = df_internal.ScalarUDF.from_pycapsule(func) + return + if isinstance(input_fields, pa.DataType): + input_fields = [input_fields] + self._udf = df_internal.ScalarUDF( + name, func, input_fields, return_field, str(volatility) + ) + + def __repr__(self) -> str: + """Print a string representation of the Scalar UDF.""" + return self._udf.__repr__() + + def __call__(self, *args: Expr) -> Expr: + """Execute the UDF. + + This function is not typically called by an end user. These calls will + occur during the evaluation of the dataframe. + """ + args_raw = [arg.expr for arg in args] + return Expr(self._udf.__call__(*args_raw)) + + @overload + @staticmethod + def udf( + input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field, + return_field: pa.DataType | pa.Field, + volatility: Volatility | str, + name: str | None = None, + ) -> Callable[..., ScalarUDF]: ... + + @overload + @staticmethod + def udf( + func: Callable[..., _R], + input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field, + return_field: pa.DataType | pa.Field, + volatility: Volatility | str, + name: str | None = None, + ) -> ScalarUDF: ... + + @overload + @staticmethod + def udf(func: ScalarUDFExportable) -> ScalarUDF: ... + + @staticmethod + def udf(*args: Any, **kwargs: Any): # noqa: D417 + """Create a new User-Defined Function (UDF). + + This class can be used both as either a function or a decorator. + + Usage: + - As a function: ``udf(func, input_fields, return_field, volatility, name)``. + - As a decorator: ``@udf(input_fields, return_field, volatility, name)``. + When used a decorator, do **not** pass ``func`` explicitly. + + In lieu of passing a PyArrow Field, you can pass a DataType for simplicity. + When you do so, it will be assumed that the nullability of the inputs and + output are True and that they have no metadata. + + Args: + func (Callable, optional): Only needed when calling as a function. + Skip this argument when using `udf` as a decorator. If you have a Rust + backed ScalarUDF within a PyCapsule, you can pass this parameter + and ignore the rest. They will be determined directly from the + underlying function. See the online documentation for more information. + input_fields (list[pa.Field | pa.DataType]): The data types or Fields + of the arguments to ``func``. This list must be of the same length + as the number of arguments. + return_field (_R): The field of the return value from the function. + volatility (Volatility | str): See `Volatility` for allowed values. + name (Optional[str]): A descriptive name for the function. + + Returns: + A user-defined function that can be used in SQL expressions, + data aggregation, or window function calls. + + Example: Using ``udf`` as a function:: + + def double_func(x): + return x * 2 + double_udf = udf(double_func, [pa.int32()], pa.int32(), + "volatile", "double_it") + + Example: Using ``udf`` as a decorator:: + + @udf([pa.int32()], pa.int32(), "volatile", "double_it") + def double_udf(x): + return x * 2 + """ # noqa: W505 E501 + + def _function( + func: Callable[..., _R], + input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field, + return_field: pa.DataType | pa.Field, + volatility: Volatility | str, + name: str | None = None, + ) -> ScalarUDF: + if not callable(func): + msg = "`func` argument must be callable" + raise TypeError(msg) + if name is None: + if hasattr(func, "__qualname__"): + name = func.__qualname__.lower() + else: + name = func.__class__.__name__.lower() + input_fields = data_types_or_fields_to_field_list(input_fields) + return_field = data_type_or_field_to_field(return_field, "value") + return ScalarUDF( + name=name, + func=func, + input_fields=input_fields, + return_field=return_field, + volatility=volatility, + ) + + def _decorator( + input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field, + return_field: _R, + volatility: Volatility | str, + name: str | None = None, + ) -> Callable: + def decorator(func: Callable) -> Callable: + udf_caller = ScalarUDF.udf( + func, input_fields, return_field, volatility, name + ) + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Callable: + return udf_caller(*args, **kwargs) + + return wrapper + + return decorator + + if hasattr(args[0], "__datafusion_scalar_udf__"): + return ScalarUDF.from_pycapsule(args[0]) + + if args and callable(args[0]): + # Case 1: Used as a function, require the first parameter to be callable + return _function(*args, **kwargs) + # Case 2: Used as a decorator with parameters + return _decorator(*args, **kwargs) + + @staticmethod + def from_pycapsule(func: ScalarUDFExportable) -> ScalarUDF: + """Create a Scalar UDF from ScalarUDF PyCapsule object. + + This function will instantiate a Scalar UDF that uses a DataFusion + ScalarUDF that is exported via the FFI bindings. + """ + name = str(func.__class__) + return ScalarUDF( + name=name, + func=func, + input_fields=None, + return_field=None, + volatility=None, + ) + + +class Accumulator(metaclass=ABCMeta): + """Defines how an :py:class:`AggregateUDF` accumulates values.""" + + @abstractmethod + def state(self) -> list[pa.Scalar]: + """Return the current state. + + While this function template expects PyArrow Scalar values return type, + you can return any value that can be converted into a Scalar. This + includes basic Python data types such as integers and strings. In + addition to primitive types, we currently support PyArrow, nanoarrow, + and arro3 objects in addition to primitive data types. Other objects + that support the Arrow FFI standard will be given a "best attempt" at + conversion to scalar objects. + """ + + @abstractmethod + def update(self, *values: pa.Array) -> None: + """Evaluate an array of values and update state.""" + + @abstractmethod + def merge(self, states: list[pa.Array]) -> None: + """Merge a set of states.""" + + @abstractmethod + def evaluate(self) -> pa.Scalar: + """Return the resultant value. + + While this function template expects a PyArrow Scalar value return type, + you can return any value that can be converted into a Scalar. This + includes basic Python data types such as integers and strings. In + addition to primitive types, we currently support PyArrow, nanoarrow, + and arro3 objects in addition to primitive data types. Other objects + that support the Arrow FFI standard will be given a "best attempt" at + conversion to scalar objects. + """ + + +class AggregateUDFExportable(Protocol): + """Type hint for object that has __datafusion_aggregate_udf__ PyCapsule.""" + + def __datafusion_aggregate_udf__(self) -> object: ... # noqa: D105 + + +class AggregateUDF: + """Class for performing scalar user-defined functions (UDF). + + Aggregate UDFs operate on a group of rows and return a single value. See + also :py:class:`ScalarUDF` for operating on a row by row basis. + """ + + @overload + def __init__( + self, + name: str, + accumulator: Callable[[], Accumulator], + input_types: list[pa.DataType], + return_type: pa.DataType, + state_type: list[pa.DataType], + volatility: Volatility | str, + ) -> None: ... + + @overload + def __init__( + self, + name: str, + accumulator: AggregateUDFExportable, + input_types: None = ..., + return_type: None = ..., + state_type: None = ..., + volatility: None = ..., + ) -> None: ... + + def __init__( + self, + name: str, + accumulator: Callable[[], Accumulator] | AggregateUDFExportable, + input_types: list[pa.DataType] | None, + return_type: pa.DataType | None, + state_type: list[pa.DataType] | None, + volatility: Volatility | str | None, + ) -> None: + """Instantiate a user-defined aggregate function (UDAF). + + See :py:func:`udaf` for a convenience function and argument + descriptions. + """ + if hasattr(accumulator, "__datafusion_aggregate_udf__"): + self._udaf = df_internal.AggregateUDF.from_pycapsule(accumulator) + return + if ( + input_types is None + or return_type is None + or state_type is None + or volatility is None + ): + msg = ( + "`input_types`, `return_type`, `state_type`, and `volatility` " + "must be provided when `accumulator` is callable." + ) + raise TypeError(msg) + + self._udaf = df_internal.AggregateUDF( + name, + accumulator, + input_types, + return_type, + state_type, + str(volatility), + ) + + def __repr__(self) -> str: + """Print a string representation of the Aggregate UDF.""" + return self._udaf.__repr__() + + def __call__(self, *args: Expr) -> Expr: + """Execute the UDAF. + + This function is not typically called by an end user. These calls will + occur during the evaluation of the dataframe. + """ + args_raw = [arg.expr for arg in args] + return Expr(self._udaf.__call__(*args_raw)) + + @overload + @staticmethod + def udaf( + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + state_type: list[pa.DataType], + volatility: Volatility | str, + name: str | None = None, + ) -> Callable[..., AggregateUDF]: ... + + @overload + @staticmethod + def udaf( + accum: Callable[[], Accumulator], + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + state_type: list[pa.DataType], + volatility: Volatility | str, + name: str | None = None, + ) -> AggregateUDF: ... + + @overload + @staticmethod + def udaf(accum: AggregateUDFExportable) -> AggregateUDF: ... + + @overload + @staticmethod + def udaf(accum: _PyCapsule) -> AggregateUDF: ... + + @staticmethod + def udaf(*args: Any, **kwargs: Any): # noqa: D417, C901 + """Create a new User-Defined Aggregate Function (UDAF). + + This class allows you to define an aggregate function that can be used in + data aggregation or window function calls. + + Usage: + - As a function: ``udaf(accum, input_types, return_type, state_type, volatility, name)``. + - As a decorator: ``@udaf(input_types, return_type, state_type, volatility, name)``. + When using ``udaf`` as a decorator, do not pass ``accum`` explicitly. + + Function example: + + If your :py:class:`Accumulator` can be instantiated with no arguments, you + can simply pass it's type as `accum`. If you need to pass additional + arguments to it's constructor, you can define a lambda or a factory method. + During runtime the :py:class:`Accumulator` will be constructed for every + instance in which this UDAF is used. The following examples are all valid:: + + import pyarrow as pa + import pyarrow.compute as pc + + class Summarize(Accumulator): + def __init__(self, bias: float = 0.0): + self._sum = pa.scalar(bias) + + def state(self) -> list[pa.Scalar]: + return [self._sum] + + def update(self, values: pa.Array) -> None: + self._sum = pa.scalar(self._sum.as_py() + pc.sum(values).as_py()) + + def merge(self, states: list[pa.Array]) -> None: + self._sum = pa.scalar(self._sum.as_py() + pc.sum(states[0]).as_py()) + + def evaluate(self) -> pa.Scalar: + return self._sum + + def sum_bias_10() -> Summarize: + return Summarize(10.0) + + udaf1 = udaf(Summarize, pa.float64(), pa.float64(), [pa.float64()], + "immutable") + udaf2 = udaf(sum_bias_10, pa.float64(), pa.float64(), [pa.float64()], + "immutable") + udaf3 = udaf(lambda: Summarize(20.0), pa.float64(), pa.float64(), + [pa.float64()], "immutable") + + Decorator example::: + + @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable") + def udf4() -> Summarize: + return Summarize(10.0) + + Args: + accum: The accumulator python function. Only needed when calling as a + function. Skip this argument when using ``udaf`` as a decorator. + If you have a Rust backed AggregateUDF within a PyCapsule, you can + pass this parameter and ignore the rest. They will be determined + directly from the underlying function. See the online documentation + for more information. + input_types: The data types of the arguments to ``accum``. + return_type: The data type of the return value. + state_type: The data types of the intermediate accumulation. + volatility: See :py:class:`Volatility` for allowed values. + name: A descriptive name for the function. + + Returns: + A user-defined aggregate function, which can be used in either data + aggregation or window function calls. + """ # noqa: E501 W505 + + def _function( + accum: Callable[[], Accumulator], + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + state_type: list[pa.DataType], + volatility: Volatility | str, + name: str | None = None, + ) -> AggregateUDF: + if not callable(accum): + msg = "`func` must be callable." + raise TypeError(msg) + if not isinstance(accum(), Accumulator): + msg = "Accumulator must implement the abstract base class Accumulator" + raise TypeError(msg) + if name is None: + name = accum().__class__.__qualname__.lower() + if isinstance(input_types, pa.DataType): + input_types = [input_types] + return AggregateUDF( + name=name, + accumulator=accum, + input_types=input_types, + return_type=return_type, + state_type=state_type, + volatility=volatility, + ) + + def _decorator( + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + state_type: list[pa.DataType], + volatility: Volatility | str, + name: str | None = None, + ) -> Callable[..., Callable[..., Expr]]: + def decorator(accum: Callable[[], Accumulator]) -> Callable[..., Expr]: + udaf_caller = AggregateUDF.udaf( + accum, input_types, return_type, state_type, volatility, name + ) + + @functools.wraps(accum) + def wrapper(*args: Any, **kwargs: Any) -> Expr: + return udaf_caller(*args, **kwargs) + + return wrapper + + return decorator + + if hasattr(args[0], "__datafusion_aggregate_udf__") or _is_pycapsule(args[0]): + return AggregateUDF.from_pycapsule(args[0]) + + if args and callable(args[0]): + # Case 1: Used as a function, require the first parameter to be callable + return _function(*args, **kwargs) + # Case 2: Used as a decorator with parameters + return _decorator(*args, **kwargs) + + @staticmethod + def from_pycapsule(func: AggregateUDFExportable | _PyCapsule) -> AggregateUDF: + """Create an Aggregate UDF from AggregateUDF PyCapsule object. + + This function will instantiate a Aggregate UDF that uses a DataFusion + AggregateUDF that is exported via the FFI bindings. + """ + if _is_pycapsule(func): + aggregate = cast("AggregateUDF", object.__new__(AggregateUDF)) + aggregate._udaf = df_internal.AggregateUDF.from_pycapsule(func) + return aggregate + + capsule = cast("AggregateUDFExportable", func) + name = str(capsule.__class__) + return AggregateUDF( + name=name, + accumulator=capsule, + input_types=None, + return_type=None, + state_type=None, + volatility=None, + ) + + +class WindowEvaluator: + """Evaluator class for user-defined window functions (UDWF). + + It is up to the user to decide which evaluate function is appropriate. + + +------------------------+--------------------------------+------------------+---------------------------+ + | ``uses_window_frame`` | ``supports_bounded_execution`` | ``include_rank`` | function_to_implement | + +========================+================================+==================+===========================+ + | False (default) | False (default) | False (default) | ``evaluate_all`` | + +------------------------+--------------------------------+------------------+---------------------------+ + | False | True | False | ``evaluate`` | + +------------------------+--------------------------------+------------------+---------------------------+ + | False | True/False | True | ``evaluate_all_with_rank``| + +------------------------+--------------------------------+------------------+---------------------------+ + | True | True/False | True/False | ``evaluate`` | + +------------------------+--------------------------------+------------------+---------------------------+ + """ # noqa: W505, E501 + + def memoize(self) -> None: + """Perform a memoize operation to improve performance. + + When the window frame has a fixed beginning (e.g UNBOUNDED + PRECEDING), some functions such as FIRST_VALUE and + NTH_VALUE do not need the (unbounded) input once they have + seen a certain amount of input. + + `memoize` is called after each input batch is processed, and + such functions can save whatever they need + """ + + def get_range(self, idx: int, num_rows: int) -> tuple[int, int]: # noqa: ARG002 + """Return the range for the window function. + + If `uses_window_frame` flag is `false`. This method is used to + calculate required range for the window function during + stateful execution. + + Generally there is no required range, hence by default this + returns smallest range(current row). e.g seeing current row is + enough to calculate window result (such as row_number, rank, + etc) + + Args: + idx:: Current index + num_rows: Number of rows. + """ + return (idx, idx + 1) + + def is_causal(self) -> bool: + """Get whether evaluator needs future data for its result.""" + return False + + def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array: + """Evaluate a window function on an entire input partition. + + This function is called once per input *partition* for window functions that + *do not use* values from the window frame, such as + :py:func:`~datafusion.functions.row_number`, + :py:func:`~datafusion.functions.rank`, + :py:func:`~datafusion.functions.dense_rank`, + :py:func:`~datafusion.functions.percent_rank`, + :py:func:`~datafusion.functions.cume_dist`, + :py:func:`~datafusion.functions.lead`, + and :py:func:`~datafusion.functions.lag`. + + It produces the result of all rows in a single pass. It + expects to receive the entire partition as the ``value`` and + must produce an output column with one output row for every + input row. + + ``num_rows`` is required to correctly compute the output in case + ``len(values) == 0`` + + Implementing this function is an optimization. Certain window + functions are not affected by the window frame definition or + the query doesn't have a frame, and ``evaluate`` skips the + (costly) window frame boundary calculation and the overhead of + calling ``evaluate`` for each output row. + + For example, the `LAG` built in window function does not use + the values of its window frame (it can be computed in one shot + on the entire partition with ``Self::evaluate_all`` regardless of the + window defined in the ``OVER`` clause) + + .. code-block:: text + + lag(x, 1) OVER (ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING) + + However, ``avg()`` computes the average in the window and thus + does use its window frame. + + .. code-block:: text + + avg(x) OVER (PARTITION BY y ORDER BY z ROWS BETWEEN 2 PRECEDING AND 3 FOLLOWING) + """ # noqa: W505, E501 + + def evaluate( + self, values: list[pa.Array], eval_range: tuple[int, int] + ) -> pa.Scalar: + """Evaluate window function on a range of rows in an input partition. + + This is the simplest and most general function to implement + but also the least performant as it creates output one row at + a time. It is typically much faster to implement stateful + evaluation using one of the other specialized methods on this + trait. + + Returns a [`ScalarValue`] that is the value of the window + function within `range` for the entire partition. Argument + `values` contains the evaluation result of function arguments + and evaluation results of ORDER BY expressions. If function has a + single argument, `values[1..]` will contain ORDER BY expression results. + """ + + def evaluate_all_with_rank( + self, num_rows: int, ranks_in_partition: list[tuple[int, int]] + ) -> pa.Array: + """Called for window functions that only need the rank of a row. + + Evaluate the partition evaluator against the partition using + the row ranks. For example, ``rank(col("a"))`` produces + + .. code-block:: text + + a | rank + - + ---- + A | 1 + A | 1 + C | 3 + D | 4 + D | 4 + + For this case, `num_rows` would be `5` and the + `ranks_in_partition` would be called with + + .. code-block:: text + + [ + (0,1), + (2,2), + (3,4), + ] + + The user must implement this method if ``include_rank`` returns True. + """ + + def supports_bounded_execution(self) -> bool: + """Can the window function be incrementally computed using bounded memory?""" + return False + + def uses_window_frame(self) -> bool: + """Does the window function use the values from the window frame?""" + return False + + def include_rank(self) -> bool: + """Can this function be evaluated with (only) rank?""" + return False + + +class WindowUDFExportable(Protocol): + """Type hint for object that has __datafusion_window_udf__ PyCapsule.""" + + def __datafusion_window_udf__(self) -> object: ... # noqa: D105 + + +class WindowUDF: + """Class for performing window user-defined functions (UDF). + + Window UDFs operate on a partition of rows. See + also :py:class:`ScalarUDF` for operating on a row by row basis. + """ + + def __init__( + self, + name: str, + func: Callable[[], WindowEvaluator], + input_types: list[pa.DataType], + return_type: pa.DataType, + volatility: Volatility | str, + ) -> None: + """Instantiate a user-defined window function (UDWF). + + See :py:func:`udwf` for a convenience function and argument + descriptions. + """ + if hasattr(func, "__datafusion_window_udf__"): + self._udwf = df_internal.WindowUDF.from_pycapsule(func) + return + self._udwf = df_internal.WindowUDF( + name, func, input_types, return_type, str(volatility) + ) + + def __repr__(self) -> str: + """Print a string representation of the Window UDF.""" + return self._udwf.__repr__() + + def __call__(self, *args: Expr) -> Expr: + """Execute the UDWF. + + This function is not typically called by an end user. These calls will + occur during the evaluation of the dataframe. + """ + args_raw = [arg.expr for arg in args] + return Expr(self._udwf.__call__(*args_raw)) + + @overload + @staticmethod + def udwf( + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + volatility: Volatility | str, + name: str | None = None, + ) -> Callable[..., WindowUDF]: ... + + @overload + @staticmethod + def udwf( + func: Callable[[], WindowEvaluator], + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + volatility: Volatility | str, + name: str | None = None, + ) -> WindowUDF: ... + + @staticmethod + def udwf(*args: Any, **kwargs: Any): # noqa: D417 + """Create a new User-Defined Window Function (UDWF). + + This class can be used both as either a function or a decorator. + + Usage: + - As a function: ``udwf(func, input_types, return_type, volatility, name)``. + - As a decorator: ``@udwf(input_types, return_type, volatility, name)``. + When using ``udwf`` as a decorator, do not pass ``func`` explicitly. + + Function example:: + + import pyarrow as pa + + class BiasedNumbers(WindowEvaluator): + def __init__(self, start: int = 0) -> None: + self.start = start + + def evaluate_all(self, values: list[pa.Array], + num_rows: int) -> pa.Array: + return pa.array([self.start + i for i in range(num_rows)]) + + def bias_10() -> BiasedNumbers: + return BiasedNumbers(10) + + udwf1 = udwf(BiasedNumbers, pa.int64(), pa.int64(), "immutable") + udwf2 = udwf(bias_10, pa.int64(), pa.int64(), "immutable") + udwf3 = udwf(lambda: BiasedNumbers(20), pa.int64(), pa.int64(), "immutable") + + + Decorator example:: + + @udwf(pa.int64(), pa.int64(), "immutable") + def biased_numbers() -> BiasedNumbers: + return BiasedNumbers(10) + + Args: + func: Only needed when calling as a function. Skip this argument when + using ``udwf`` as a decorator. If you have a Rust backed WindowUDF + within a PyCapsule, you can pass this parameter and ignore the rest. + They will be determined directly from the underlying function. See + the online documentation for more information. + input_types: The data types of the arguments. + return_type: The data type of the return value. + volatility: See :py:class:`Volatility` for allowed values. + name: A descriptive name for the function. + + Returns: + A user-defined window function that can be used in window function calls. + """ + if hasattr(args[0], "__datafusion_window_udf__"): + return WindowUDF.from_pycapsule(args[0]) + + if args and callable(args[0]): + # Case 1: Used as a function, require the first parameter to be callable + return WindowUDF._create_window_udf(*args, **kwargs) + # Case 2: Used as a decorator with parameters + return WindowUDF._create_window_udf_decorator(*args, **kwargs) + + @staticmethod + def _create_window_udf( + func: Callable[[], WindowEvaluator], + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + volatility: Volatility | str, + name: str | None = None, + ) -> WindowUDF: + """Create a WindowUDF instance from function arguments.""" + if not callable(func): + msg = "`func` must be callable." + raise TypeError(msg) + if not isinstance(func(), WindowEvaluator): + msg = "`func` must implement the abstract base class WindowEvaluator" + raise TypeError(msg) + + name = name or func.__qualname__.lower() + input_types = ( + [input_types] if isinstance(input_types, pa.DataType) else input_types + ) + + return WindowUDF(name, func, input_types, return_type, volatility) + + @staticmethod + def _get_default_name(func: Callable) -> str: + """Get the default name for a function based on its attributes.""" + if hasattr(func, "__qualname__"): + return func.__qualname__.lower() + return func.__class__.__name__.lower() + + @staticmethod + def _normalize_input_types( + input_types: pa.DataType | list[pa.DataType], + ) -> list[pa.DataType]: + """Convert a single DataType to a list if needed.""" + if isinstance(input_types, pa.DataType): + return [input_types] + return input_types + + @staticmethod + def _create_window_udf_decorator( + input_types: pa.DataType | list[pa.DataType], + return_type: pa.DataType, + volatility: Volatility | str, + name: str | None = None, + ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]: + """Create a decorator for a WindowUDF.""" + + def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]: + udwf_caller = WindowUDF._create_window_udf( + func, input_types, return_type, volatility, name + ) + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Expr: + return udwf_caller(*args, **kwargs) + + return wrapper + + return decorator + + @staticmethod + def from_pycapsule(func: WindowUDFExportable) -> WindowUDF: + """Create a Window UDF from WindowUDF PyCapsule object. + + This function will instantiate a Window UDF that uses a DataFusion + WindowUDF that is exported via the FFI bindings. + """ + name = str(func.__class__) + return WindowUDF( + name=name, + func=func, + input_types=None, + return_type=None, + volatility=None, + ) + + +class TableFunction: + """Class for performing user-defined table functions (UDTF). + + Table functions generate new table providers based on the + input expressions. + """ + + def __init__( + self, name: str, func: Callable[[], any], ctx: SessionContext | None = None + ) -> None: + """Instantiate a user-defined table function (UDTF). + + See :py:func:`udtf` for a convenience function and argument + descriptions. + """ + self._udtf = df_internal.TableFunction(name, func, ctx) + + def __call__(self, *args: Expr) -> Any: + """Execute the UDTF and return a table provider.""" + args_raw = [arg.expr for arg in args] + return self._udtf.__call__(*args_raw) + + @overload + @staticmethod + def udtf( + name: str, + ) -> Callable[..., Any]: ... + + @overload + @staticmethod + def udtf( + func: Callable[[], Any], + name: str, + ) -> TableFunction: ... + + @staticmethod + def udtf(*args: Any, **kwargs: Any): + """Create a new User-Defined Table Function (UDTF).""" + if args and callable(args[0]): + # Case 1: Used as a function, require the first parameter to be callable + return TableFunction._create_table_udf(*args, **kwargs) + if args and hasattr(args[0], "__datafusion_table_function__"): + # Case 2: We have a datafusion FFI provided function + return TableFunction(args[1], args[0]) + # Case 3: Used as a decorator with parameters + return TableFunction._create_table_udf_decorator(*args, **kwargs) + + @staticmethod + def _create_table_udf( + func: Callable[..., Any], + name: str, + ) -> TableFunction: + """Create a TableFunction instance from function arguments.""" + if not callable(func): + msg = "`func` must be callable." + raise TypeError(msg) + + return TableFunction(name, func) + + @staticmethod + def _create_table_udf_decorator( + name: str | None = None, + ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]: + """Create a decorator for a WindowUDF.""" + + def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]: + return TableFunction._create_table_udf(func, name) + + return decorator + + def __repr__(self) -> str: + """User printable representation.""" + return self._udtf.__repr__() + + +# Convenience exports so we can import instead of treating as +# variables at the package root +udf = ScalarUDF.udf +udaf = AggregateUDF.udaf +udwf = WindowUDF.udwf +udtf = TableFunction.udtf diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 9548fbfe4..26ed7281d 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -17,7 +17,7 @@ import pyarrow as pa import pytest -from datafusion import SessionContext +from datafusion import DataFrame, SessionContext from pyarrow.csv import write_csv @@ -49,3 +49,12 @@ def database(ctx, tmp_path): delimiter=",", schema_infer_max_records=10, ) + + +@pytest.fixture +def fail_collect(monkeypatch): + def _fail_collect(self, *args, **kwargs): # pragma: no cover - failure path + msg = "collect should not be called" + raise AssertionError(msg) + + monkeypatch.setattr(DataFrame, "collect", _fail_collect) diff --git a/python/tests/test_aggregation.py b/python/tests/test_aggregation.py index 61b1c7d80..240332848 100644 --- a/python/tests/test_aggregation.py +++ b/python/tests/test_aggregation.py @@ -88,7 +88,7 @@ def df_aggregate_100(): f.covar_samp(column("b"), column("c")), lambda a, b, c, d: np.array(np.cov(b, c, ddof=1)[0][1]), ), - # f.grouping(col_a), # No physical plan implemented yet + # f.grouping(col_a), # noqa: ERA001 No physical plan implemented yet (f.max(column("a")), lambda a, b, c, d: np.array(np.max(a))), (f.mean(column("b")), lambda a, b, c, d: np.array(np.mean(b))), (f.median(column("b")), lambda a, b, c, d: np.array(np.median(b))), @@ -130,9 +130,25 @@ def test_aggregation_stats(df, agg_expr, calc_expected): (f.median(column("b"), filter=column("a") != 2), pa.array([5]), False), (f.approx_median(column("b"), filter=column("a") != 2), pa.array([5]), False), (f.approx_percentile_cont(column("b"), 0.5), pa.array([4]), False), + ( + f.approx_percentile_cont( + column("b").sort(ascending=True, nulls_first=False), + 0.5, + num_centroids=2, + ), + pa.array([4]), + False, + ), ( f.approx_percentile_cont_with_weight(column("b"), lit(0.6), 0.5), - pa.array([6], type=pa.float64()), + pa.array([4], type=pa.float64()), + False, + ), + ( + f.approx_percentile_cont_with_weight( + column("b").sort(ascending=False, nulls_first=False), lit(0.6), 0.5 + ), + pa.array([4], type=pa.float64()), False, ), ( @@ -154,6 +170,11 @@ def test_aggregation_stats(df, agg_expr, calc_expected): pa.array([[6, 4, 4]]), False, ), + ( + f.array_agg(column("b"), order_by=column("c")), + pa.array([[6, 4, 4]]), + False, + ), (f.avg(column("b"), filter=column("a") != lit(1)), pa.array([5.0]), False), (f.sum(column("b"), filter=column("a") != lit(1)), pa.array([10]), False), (f.count(column("b"), distinct=True), pa.array([2]), False), @@ -329,6 +350,15 @@ def test_bit_and_bool_fns(df, name, expr, result): ), [None, None], ), + ( + "first_value_no_list_order_by", + f.first_value( + column("b"), + order_by=column("b"), + null_treatment=NullTreatment.RESPECT_NULLS, + ), + [None, None], + ), ( "first_value_ignore_null", f.first_value( @@ -338,12 +368,16 @@ def test_bit_and_bool_fns(df, name, expr, result): ), [7, 9], ), - ("last_value", f.last_value(column("a")), [3, 6]), ( "last_value_ordered", f.last_value(column("a"), order_by=[column("a").sort(ascending=False)]), [0, 4], ), + ( + "last_value_no_list_ordered", + f.last_value(column("a"), order_by=column("a")), + [3, 6], + ), ( "last_value_with_null", f.last_value( @@ -367,6 +401,11 @@ def test_bit_and_bool_fns(df, name, expr, result): f.nth_value(column("a"), 2, order_by=[column("a").sort(ascending=False)]), [2, 5], ), + ( + "nth_value_no_list_ordered", + f.nth_value(column("a"), 2, order_by=column("a").sort(ascending=False)), + [2, 5], + ), ( "nth_value_with_null", f.nth_value( @@ -415,6 +454,11 @@ def test_first_last_value(df_partitioned, name, expr, result) -> None: f.string_agg(column("a"), ",", order_by=[column("b")]), "one,three,two,two", ), + ( + "string_agg", + f.string_agg(column("a"), ",", order_by=column("b")), + "one,three,two,two", + ), ], ) def test_string_agg(name, expr, result) -> None: diff --git a/python/tests/test_catalog.py b/python/tests/test_catalog.py index 23b328458..c89da36bf 100644 --- a/python/tests/test_catalog.py +++ b/python/tests/test_catalog.py @@ -14,9 +14,18 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations +from typing import TYPE_CHECKING + +import datafusion as dfn import pyarrow as pa +import pyarrow.dataset as ds import pytest +from datafusion import Catalog, SessionContext, Table, udtf + +if TYPE_CHECKING: + from datafusion.catalog import CatalogProvider, CatalogProviderExportable # Note we take in `database` as a variable even though we don't use @@ -27,9 +36,9 @@ def test_basic(ctx, database): ctx.catalog("non-existent") default = ctx.catalog() - assert default.names() == ["public"] + assert default.names() == {"public"} - for db in [default.database("public"), default.database()]: + for db in [default.schema("public"), default.schema()]: assert db.names() == {"csv1", "csv", "csv2"} table = db.table("csv") @@ -41,3 +50,294 @@ def test_basic(ctx, database): pa.field("float", pa.float64(), nullable=True), ] ) + + +def create_dataset() -> Table: + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3]), pa.array([4, 5, 6])], + names=["a", "b"], + ) + dataset = ds.dataset([batch]) + return Table(dataset) + + +class CustomSchemaProvider(dfn.catalog.SchemaProvider): + def __init__(self): + self.tables = {"table1": create_dataset()} + + def table_names(self) -> set[str]: + return set(self.tables.keys()) + + def register_table(self, name: str, table: Table): + self.tables[name] = table + + def deregister_table(self, name, cascade: bool = True): + del self.tables[name] + + def table(self, name: str) -> Table | None: + return self.tables[name] + + def table_exist(self, name: str) -> bool: + return name in self.tables + + +class CustomErrorSchemaProvider(CustomSchemaProvider): + def table(self, name: str) -> Table | None: + message = f"{name} is not an acceptable name" + raise ValueError(message) + + +class CustomCatalogProvider(dfn.catalog.CatalogProvider): + def __init__(self): + self.schemas = {"my_schema": CustomSchemaProvider()} + + def schema_names(self) -> set[str]: + return set(self.schemas.keys()) + + def schema(self, name: str): + return self.schemas[name] + + def register_schema(self, name: str, schema: dfn.catalog.Schema): + self.schemas[name] = schema + + def deregister_schema(self, name, cascade: bool): + del self.schemas[name] + + +class CustomCatalogProviderList(dfn.catalog.CatalogProviderList): + def __init__(self): + self.catalogs = {"my_catalog": CustomCatalogProvider()} + + def catalog_names(self) -> set[str]: + return set(self.catalogs.keys()) + + def catalog(self, name: str) -> Catalog | None: + return self.catalogs[name] + + def register_catalog( + self, name: str, catalog: CatalogProviderExportable | CatalogProvider | Catalog + ) -> None: + self.catalogs[name] = catalog + + +class CustomTableProviderFactory(dfn.catalog.TableProviderFactory): + def create(self, cmd: dfn.expr.CreateExternalTable): + assert cmd.name() == "test_table_factory" + return create_dataset() + + +def test_python_catalog_provider_list(ctx: SessionContext): + ctx.register_catalog_provider_list(CustomCatalogProviderList()) + + # Ensure `datafusion` catalog does not exist since + # we replaced the catalog list + assert ctx.catalog_names() == {"my_catalog"} + + # Ensure registering works + ctx.register_catalog_provider("second_catalog", Catalog.memory_catalog()) + assert ctx.catalog_names() == {"my_catalog", "second_catalog"} + + +def test_python_catalog_provider(ctx: SessionContext): + ctx.register_catalog_provider("my_catalog", CustomCatalogProvider()) + + # Check the default catalog provider + assert ctx.catalog("datafusion").names() == {"public"} + + my_catalog = ctx.catalog("my_catalog") + assert my_catalog.names() == {"my_schema"} + + my_catalog.register_schema("second_schema", CustomSchemaProvider()) + assert my_catalog.schema_names() == {"my_schema", "second_schema"} + + my_catalog.deregister_schema("my_schema") + assert my_catalog.schema_names() == {"second_schema"} + + +def test_in_memory_providers(ctx: SessionContext): + catalog = dfn.catalog.Catalog.memory_catalog() + ctx.register_catalog_provider("in_mem_catalog", catalog) + + assert ctx.catalog_names() == {"datafusion", "in_mem_catalog"} + + schema = dfn.catalog.Schema.memory_schema() + catalog.register_schema("in_mem_schema", schema) + + schema.register_table("my_table", create_dataset()) + + batches = ctx.sql("select * from in_mem_catalog.in_mem_schema.my_table").collect() + + assert len(batches) == 1 + assert batches[0].column(0) == pa.array([1, 2, 3]) + assert batches[0].column(1) == pa.array([4, 5, 6]) + + +def test_python_schema_provider(ctx: SessionContext): + catalog = ctx.catalog() + + catalog.deregister_schema("public") + + catalog.register_schema("test_schema1", CustomSchemaProvider()) + assert catalog.names() == {"test_schema1"} + + catalog.register_schema("test_schema2", CustomSchemaProvider()) + catalog.deregister_schema("test_schema1") + assert catalog.names() == {"test_schema2"} + + +def test_python_table_provider(ctx: SessionContext): + catalog = ctx.catalog() + + catalog.register_schema("custom_schema", CustomSchemaProvider()) + schema = catalog.schema("custom_schema") + + assert schema.table_names() == {"table1"} + + schema.deregister_table("table1") + schema.register_table("table2", create_dataset()) + assert schema.table_names() == {"table2"} + + # Use the default schema instead of our custom schema + + schema = catalog.schema() + + schema.register_table("table3", create_dataset()) + assert schema.table_names() == {"table3"} + + schema.deregister_table("table3") + schema.register_table("table4", create_dataset()) + assert schema.table_names() == {"table4"} + + +def test_schema_register_table_with_pyarrow_dataset(ctx: SessionContext): + schema = ctx.catalog().schema() + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3]), pa.array([4, 5, 6])], + names=["a", "b"], + ) + dataset = ds.dataset([batch]) + table_name = "pa_dataset" + + try: + schema.register_table(table_name, dataset) + assert table_name in schema.table_names() + + result = ctx.sql(f"SELECT a, b FROM {table_name}").collect() + + assert len(result) == 1 + assert result[0].column(0) == pa.array([1, 2, 3]) + assert result[0].column(1) == pa.array([4, 5, 6]) + finally: + schema.deregister_table(table_name) + + +def test_exception_not_mangled(ctx: SessionContext): + """Test registering all python providers and running a query against them.""" + + catalog_name = "custom_catalog" + schema_name = "custom_schema" + + ctx.register_catalog_provider(catalog_name, CustomCatalogProvider()) + + catalog = ctx.catalog(catalog_name) + + # Clean out previous schemas if they exist so we can start clean + for schema_name in catalog.schema_names(): + catalog.deregister_schema(schema_name, cascade=False) + + catalog.register_schema(schema_name, CustomErrorSchemaProvider()) + + schema = catalog.schema(schema_name) + + for table_name in schema.table_names(): + schema.deregister_table(table_name) + + schema.register_table("test_table", create_dataset()) + + with pytest.raises(ValueError, match=r"^test_table is not an acceptable name$"): + ctx.sql(f"select * from {catalog_name}.{schema_name}.test_table") + + +def test_in_end_to_end_python_providers(ctx: SessionContext): + """Test registering all python providers and running a query against them.""" + + all_catalog_names = [ + "datafusion", + "custom_catalog", + "in_mem_catalog", + ] + + all_schema_names = [ + "custom_schema", + "in_mem_schema", + ] + + ctx.register_catalog_provider(all_catalog_names[1], CustomCatalogProvider()) + ctx.register_catalog_provider( + all_catalog_names[2], dfn.catalog.Catalog.memory_catalog() + ) + + for catalog_name in all_catalog_names: + catalog = ctx.catalog(catalog_name) + + # Clean out previous schemas if they exist so we can start clean + for schema_name in catalog.schema_names(): + catalog.deregister_schema(schema_name, cascade=False) + + catalog.register_schema(all_schema_names[0], CustomSchemaProvider()) + catalog.register_schema(all_schema_names[1], dfn.catalog.Schema.memory_schema()) + + for schema_name in all_schema_names: + schema = catalog.schema(schema_name) + + for table_name in schema.table_names(): + schema.deregister_table(table_name) + + schema.register_table("test_table", create_dataset()) + + for catalog_name in all_catalog_names: + for schema_name in all_schema_names: + table_full_name = f"{catalog_name}.{schema_name}.test_table" + + batches = ctx.sql(f"select * from {table_full_name}").collect() + + assert len(batches) == 1 + assert batches[0].column(0) == pa.array([1, 2, 3]) + assert batches[0].column(1) == pa.array([4, 5, 6]) + + +def test_register_python_function_as_udtf(ctx: SessionContext): + basic_table = Table(ctx.sql("SELECT 3 AS value")) + + @udtf("my_table_function") + def my_table_function_udtf() -> Table: + return basic_table + + ctx.register_udtf(my_table_function_udtf) + + result = ctx.sql("SELECT * FROM my_table_function()").collect() + assert len(result) == 1 + assert len(result[0]) == 1 + assert len(result[0][0]) == 1 + assert result[0][0][0].as_py() == 3 + + +def test_register_python_table_provider_factory(ctx: SessionContext): + ctx.register_table_factory("CUSTOM_FACTORY", CustomTableProviderFactory()) + + ctx.sql(""" + CREATE EXTERNAL TABLE test_table_factory + STORED AS CUSTOM_FACTORY + LOCATION foo; + """).collect() + + result = ctx.sql("SELECT * FROM test_table_factory;").collect() + + expect = [ + pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3]), pa.array([4, 5, 6])], + names=["a", "b"], + ) + ] + + assert result == expect diff --git a/python/tests/test_concurrency.py b/python/tests/test_concurrency.py new file mode 100644 index 000000000..f790f9473 --- /dev/null +++ b/python/tests/test_concurrency.py @@ -0,0 +1,126 @@ +# 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. + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import pyarrow as pa +from datafusion import Config, SessionContext, col, lit +from datafusion import functions as f +from datafusion.common import SqlSchema + + +def _run_in_threads(fn, count: int = 8) -> None: + with ThreadPoolExecutor(max_workers=count) as executor: + futures = [executor.submit(fn, i) for i in range(count)] + for future in futures: + # Propagate any exception raised in the worker thread. + future.result() + + +def test_concurrent_access_to_shared_structures() -> None: + """Exercise SqlSchema, Config, and DataFrame concurrently.""" + + schema = SqlSchema("concurrency") + config = Config() + ctx = SessionContext() + + batch = pa.record_batch([pa.array([1, 2, 3], type=pa.int32())], names=["value"]) + df = ctx.create_dataframe([[batch]]) + + config_key = "datafusion.execution.batch_size" + expected_rows = batch.num_rows + + def worker(index: int) -> None: + schema.name = f"concurrency-{index}" + assert schema.name.startswith("concurrency-") + # Exercise getters that use internal locks. + assert isinstance(schema.tables, list) + assert isinstance(schema.views, list) + assert isinstance(schema.functions, list) + + config.set(config_key, str(1024 + index)) + assert config.get(config_key) is not None + # Access the full config map to stress lock usage. + assert config_key in config.get_all() + + batches = df.collect() + assert sum(batch.num_rows for batch in batches) == expected_rows + + _run_in_threads(worker, count=12) + + +def test_config_set_during_get_all() -> None: + """Ensure config writes proceed while another thread reads all entries.""" + + config = Config() + key = "datafusion.execution.batch_size" + + def reader() -> None: + for _ in range(200): + # get_all should not hold the lock while converting to Python objects + config.get_all() + + def writer() -> None: + for index in range(200): + config.set(key, str(1024 + index)) + + with ThreadPoolExecutor(max_workers=2) as executor: + reader_future = executor.submit(reader) + writer_future = executor.submit(writer) + reader_future.result(timeout=10) + writer_future.result(timeout=10) + + assert config.get(key) is not None + + +def test_case_builder_reuse_from_multiple_threads() -> None: + """Ensure the case builder can be safely reused across threads.""" + + ctx = SessionContext() + values = pa.array([0, 1, 2, 3, 4], type=pa.int32()) + df = ctx.create_dataframe([[pa.record_batch([values], names=["value"])]]) + + base_builder = f.case(col("value")) + + def add_case(i: int) -> None: + nonlocal base_builder + base_builder = base_builder.when(lit(i), lit(f"value-{i}")) + + _run_in_threads(add_case, count=8) + + with ThreadPoolExecutor(max_workers=2) as executor: + otherwise_future = executor.submit(base_builder.otherwise, lit("default")) + case_expr = otherwise_future.result() + + result = df.select(case_expr.alias("label")).collect() + assert sum(batch.num_rows for batch in result) == len(values) + + predicate_builder = f.when(col("value") == lit(0), lit("zero")) + + def add_predicate(i: int) -> None: + predicate_builder.when(col("value") == lit(i + 1), lit(f"value-{i + 1}")) + + _run_in_threads(add_predicate, count=4) + + with ThreadPoolExecutor(max_workers=2) as executor: + end_future = executor.submit(predicate_builder.end) + predicate_expr = end_future.result() + + result = df.select(predicate_expr.alias("label")).collect() + assert sum(batch.num_rows for batch in result) == len(values) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 4a15ac9cf..5df6ed20f 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -22,11 +22,13 @@ import pyarrow.dataset as ds import pytest from datafusion import ( + CsvReadOptions, DataFrame, RuntimeEnvBuilder, SessionConfig, SessionContext, SQLOptions, + Table, column, literal, ) @@ -57,7 +59,7 @@ def test_runtime_configs(tmp_path, path_to_str): ctx = SessionContext(config, runtime) assert ctx is not None - db = ctx.catalog("foo").database("bar") + db = ctx.catalog("foo").schema("bar") assert db is not None @@ -70,7 +72,7 @@ def test_temporary_files(tmp_path, path_to_str): ctx = SessionContext(config, runtime) assert ctx is not None - db = ctx.catalog("foo").database("bar") + db = ctx.catalog("foo").schema("bar") assert db is not None @@ -91,7 +93,7 @@ def test_create_context_with_all_valid_args(): ctx = SessionContext(config, runtime) # verify that at least some of the arguments worked - ctx.catalog("foo").database("bar") + ctx.catalog("foo").schema("bar") with pytest.raises(KeyError): ctx.catalog("datafusion") @@ -105,7 +107,7 @@ def test_register_record_batches(ctx): ctx.register_record_batches("t", [[batch]]) - assert ctx.catalog().database().names() == {"t"} + assert ctx.catalog().schema().names() == {"t"} result = ctx.sql("SELECT a+b, a-b FROM t").collect() @@ -121,7 +123,7 @@ def test_create_dataframe_registers_unique_table_name(ctx): ) df = ctx.create_dataframe([[batch]]) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -141,7 +143,7 @@ def test_create_dataframe_registers_with_defined_table_name(ctx): ) df = ctx.create_dataframe([[batch]], name="tbl") - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -155,7 +157,7 @@ def test_from_arrow_table(ctx): # convert to DataFrame df = ctx.from_arrow(table) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -200,7 +202,7 @@ def test_from_arrow_table_with_name(ctx): # convert to DataFrame with optional name df = ctx.from_arrow(table, name="tbl") - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert tables[0] == "tbl" @@ -213,7 +215,7 @@ def test_from_arrow_table_empty(ctx): # convert to DataFrame df = ctx.from_arrow(table) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -228,7 +230,7 @@ def test_from_arrow_table_empty_no_schema(ctx): # convert to DataFrame df = ctx.from_arrow(table) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -246,7 +248,7 @@ def test_from_pylist(ctx): ] df = ctx.from_pylist(data) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -260,7 +262,7 @@ def test_from_pydict(ctx): data = {"a": [1, 2, 3], "b": [4, 5, 6]} df = ctx.from_pydict(data) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -276,7 +278,7 @@ def test_from_pandas(ctx): pandas_df = pd.DataFrame(data) df = ctx.from_pandas(pandas_df) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -292,7 +294,7 @@ def test_from_polars(ctx): polars_df = pd.DataFrame(data) df = ctx.from_polars(polars_df) - tables = list(ctx.catalog().database().names()) + tables = list(ctx.catalog().schema().names()) assert df assert len(tables) == 1 @@ -303,7 +305,7 @@ def test_from_polars(ctx): def test_register_table(ctx, database): default = ctx.catalog() - public = default.database("public") + public = default.schema("public") assert public.names() == {"csv", "csv1", "csv2"} table = public.table("csv") @@ -311,9 +313,9 @@ def test_register_table(ctx, database): assert public.names() == {"csv", "csv1", "csv2", "csv3"} -def test_read_table(ctx, database): +def test_read_table_from_catalog(ctx, database): default = ctx.catalog() - public = default.database("public") + public = default.schema("public") assert public.names() == {"csv", "csv1", "csv2"} table = public.table("csv") @@ -321,15 +323,74 @@ def test_read_table(ctx, database): table_df.show() +def test_read_table_from_df(ctx): + df = ctx.from_pydict({"a": [1, 2]}) + result = ctx.read_table(df).collect() + assert [b.to_pydict() for b in result] == [{"a": [1, 2]}] + + +def test_read_table_from_dataset(ctx): + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3]), pa.array([4, 5, 6])], + names=["a", "b"], + ) + dataset = ds.dataset([batch]) + + result = ctx.read_table(dataset).collect() + + assert result[0].column(0) == pa.array([1, 2, 3]) + assert result[0].column(1) == pa.array([4, 5, 6]) + + def test_deregister_table(ctx, database): default = ctx.catalog() - public = default.database("public") + public = default.schema("public") assert public.names() == {"csv", "csv1", "csv2"} ctx.deregister_table("csv") assert public.names() == {"csv1", "csv2"} +def test_register_table_from_dataframe(ctx): + df = ctx.from_pydict({"a": [1, 2]}) + ctx.register_table("df_tbl", df) + result = ctx.sql("SELECT * FROM df_tbl").collect() + assert [b.to_pydict() for b in result] == [{"a": [1, 2]}] + + +@pytest.mark.parametrize("temporary", [True, False]) +def test_register_table_from_dataframe_into_view(ctx, temporary): + df = ctx.from_pydict({"a": [1, 2]}) + table = df.into_view(temporary=temporary) + assert isinstance(table, Table) + if temporary: + assert table.kind == "temporary" + else: + assert table.kind == "view" + + ctx.register_table("view_tbl", table) + result = ctx.sql("SELECT * FROM view_tbl").collect() + assert [b.to_pydict() for b in result] == [{"a": [1, 2]}] + + +def test_table_from_dataframe(ctx): + df = ctx.from_pydict({"a": [1, 2]}) + table = Table(df) + assert isinstance(table, Table) + ctx.register_table("from_dataframe_tbl", table) + result = ctx.sql("SELECT * FROM from_dataframe_tbl").collect() + assert [b.to_pydict() for b in result] == [{"a": [1, 2]}] + + +def test_table_from_dataframe_internal(ctx): + df = ctx.from_pydict({"a": [1, 2]}) + table = Table(df.df) + assert isinstance(table, Table) + ctx.register_table("from_internal_dataframe_tbl", table) + result = ctx.sql("SELECT * FROM from_internal_dataframe_tbl").collect() + assert [b.to_pydict() for b in result] == [{"a": [1, 2]}] + + def test_register_dataset(ctx): # create a RecordBatch and register it as a pyarrow.dataset.Dataset batch = pa.RecordBatch.from_arrays( @@ -339,7 +400,7 @@ def test_register_dataset(ctx): dataset = ds.dataset([batch]) ctx.register_dataset("t", dataset) - assert ctx.catalog().database().names() == {"t"} + assert ctx.catalog().schema().names() == {"t"} result = ctx.sql("SELECT a+b, a-b FROM t").collect() @@ -356,7 +417,7 @@ def test_dataset_filter(ctx, capfd): dataset = ds.dataset([batch]) ctx.register_dataset("t", dataset) - assert ctx.catalog().database().names() == {"t"} + assert ctx.catalog().schema().names() == {"t"} df = ctx.sql("SELECT a+b, a-b FROM t WHERE a BETWEEN 2 and 3 AND b > 5") # Make sure the filter was pushed down in Physical Plan @@ -455,7 +516,7 @@ def test_dataset_filter_nested_data(ctx): dataset = ds.dataset([batch]) ctx.register_dataset("t", dataset) - assert ctx.catalog().database().names() == {"t"} + assert ctx.catalog().schema().names() == {"t"} df = ctx.table("t") @@ -566,6 +627,8 @@ def test_read_csv_list(ctx): def test_read_csv_compressed(ctx, tmp_path): test_data_path = pathlib.Path("testing/data/csv/aggregate_test_100.csv") + expected = ctx.read_csv(test_data_path).collect() + # File compression type gzip_path = tmp_path / "aggregate_test_100.csv.gz" @@ -576,7 +639,13 @@ def test_read_csv_compressed(ctx, tmp_path): gzipped_file.writelines(csv_file) csv_df = ctx.read_csv(gzip_path, file_extension=".gz", file_compression_type="gz") - csv_df.select(column("c1")).show() + assert csv_df.collect() == expected + + csv_df = ctx.read_csv( + gzip_path, + options=CsvReadOptions(file_extension=".gz", file_compression_type="gz"), + ) + assert csv_df.collect() == expected def test_read_parquet(ctx): @@ -650,3 +719,154 @@ def test_create_dataframe_with_global_ctx(batch): result = df.collect()[0].column(0) assert result == pa.array([4, 5, 6]) + + +def test_csv_read_options_builder_pattern(): + """Test CsvReadOptions builder pattern.""" + from datafusion import CsvReadOptions + + options = ( + CsvReadOptions() + .with_has_header(False) # noqa: FBT003 + .with_delimiter("|") + .with_quote("'") + .with_schema_infer_max_records(2000) + .with_truncated_rows(True) # noqa: FBT003 + .with_newlines_in_values(True) # noqa: FBT003 + .with_file_extension(".tsv") + ) + assert options.has_header is False + assert options.delimiter == "|" + assert options.quote == "'" + assert options.schema_infer_max_records == 2000 + assert options.truncated_rows is True + assert options.newlines_in_values is True + assert options.file_extension == ".tsv" + + +def read_csv_with_options_inner( + tmp_path: pathlib.Path, + csv_content: str, + options: CsvReadOptions, + expected: pa.RecordBatch, + as_read: bool, + global_ctx: bool, +) -> None: + from datafusion import SessionContext + + # Create a test CSV file + group_dir = tmp_path / "group=a" + group_dir.mkdir(exist_ok=True) + + csv_path = group_dir / "test.csv" + csv_path.write_text(csv_content, newline="\n") + + ctx = SessionContext() + + if as_read: + if global_ctx: + from datafusion.io import read_csv + + df = read_csv(str(tmp_path), options=options) + else: + df = ctx.read_csv(str(tmp_path), options=options) + else: + ctx.register_csv("test_table", str(tmp_path), options=options) + df = ctx.sql("SELECT * FROM test_table") + df.show() + + # Verify the data + result = df.collect() + assert len(result) == 1 + assert result[0] == expected + + +@pytest.mark.parametrize( + ("as_read", "global_ctx"), + [ + (True, True), + (True, False), + (False, False), + ], +) +def test_read_csv_with_options(tmp_path, as_read, global_ctx): + """Test reading CSV with CsvReadOptions.""" + + csv_content = "Alice;30;|New York; NY|\nBob;25\n#Charlie;35;Paris\nPhil;75;Detroit' MI\nKarin;50;|Stockholm\nSweden|" # noqa: E501 + + # Some of the read options are difficult to test in combination + # such as schema and schema_infer_max_records so run multiple tests + # file_sort_order doesn't impact reading, but included here to ensure + # all options parse correctly + options = CsvReadOptions( + has_header=False, + delimiter=";", + quote="|", + terminator="\n", + escape="\\", + comment="#", + newlines_in_values=True, + schema_infer_max_records=1, + null_regex="[pP]+aris", + truncated_rows=True, + file_sort_order=[[column("column_1").sort(), column("column_2")], ["column_3"]], + ) + + expected = pa.RecordBatch.from_arrays( + [ + pa.array(["Alice", "Bob", "Phil", "Karin"]), + pa.array([30, 25, 75, 50]), + pa.array(["New York; NY", None, "Detroit' MI", "Stockholm\nSweden"]), + ], + names=["column_1", "column_2", "column_3"], + ) + + read_csv_with_options_inner( + tmp_path, csv_content, options, expected, as_read, global_ctx + ) + + schema = pa.schema( + [ + pa.field("name", pa.string(), nullable=False), + pa.field("age", pa.float32(), nullable=False), + pa.field("location", pa.string(), nullable=True), + ] + ) + options.with_schema(schema) + + expected = pa.RecordBatch.from_arrays( + [ + pa.array(["Alice", "Bob", "Phil", "Karin"]), + pa.array([30.0, 25.0, 75.0, 50.0]), + pa.array(["New York; NY", None, "Detroit' MI", "Stockholm\nSweden"]), + ], + schema=schema, + ) + + read_csv_with_options_inner( + tmp_path, csv_content, options, expected, as_read, global_ctx + ) + + csv_content = "name,age\nAlice,30\nBob,25\nCharlie,35\nDiego,40\nEmily,15" + + expected = pa.RecordBatch.from_arrays( + [ + pa.array(["Alice", "Bob", "Charlie", "Diego", "Emily"]), + pa.array([30, 25, 35, 40, 15]), + pa.array(["a", "a", "a", "a", "a"]), + ], + schema=pa.schema( + [ + pa.field("name", pa.string(), nullable=True), + pa.field("age", pa.int64(), nullable=True), + pa.field("group", pa.string(), nullable=False), + ] + ), + ) + options = CsvReadOptions( + table_partition_cols=[("group", pa.string())], + ) + + read_csv_with_options_inner( + tmp_path, csv_content, options, expected, as_read, global_ctx + ) diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index 384b17878..759d6278c 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -14,7 +14,14 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import ctypes +import datetime +import itertools import os +import re +import threading +import time +from pathlib import Path from typing import Any import pyarrow as pa @@ -22,15 +29,36 @@ import pytest from datafusion import ( DataFrame, + InsertOp, + ParquetColumnOptions, + ParquetWriterOptions, + RecordBatch, SessionContext, WindowFrame, column, literal, + udf, ) -from datafusion import functions as f -from datafusion.expr import Window +from datafusion import ( + col as df_col, +) +from datafusion import ( + functions as f, +) +from datafusion.dataframe import DataFrameWriteOptions +from datafusion.dataframe_formatter import ( + DataFrameHtmlFormatter, + configure_formatter, + get_formatter, + reset_formatter, +) +from datafusion.expr import EXPR_TYPE_ERROR, Window from pyarrow.csv import write_csv +pa_cffi = pytest.importorskip("pyarrow.cffi") + +MB = 1024 * 1024 + @pytest.fixture def ctx(): @@ -38,9 +66,7 @@ def ctx(): @pytest.fixture -def df(): - ctx = SessionContext() - +def df(ctx): # create a RecordBatch and a new DataFrame from it batch = pa.RecordBatch.from_arrays( [pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])], @@ -50,6 +76,54 @@ def df(): return ctx.from_arrow(batch) +@pytest.fixture +def large_df(): + ctx = SessionContext() + + rows = 100000 + data = { + "a": list(range(rows)), + "b": [f"s-{i}" for i in range(rows)], + "c": [float(i + 0.1) for i in range(rows)], + } + batch = pa.record_batch(data) + + return ctx.from_arrow(batch) + + +@pytest.fixture +def large_multi_batch_df(): + """Create a DataFrame with multiple record batches for testing stream behavior. + + This fixture creates 10 batches of 10,000 rows each (100,000 rows total), + ensuring the DataFrame spans multiple batches. This is essential for testing + that memory limits actually cause early stream termination rather than + truncating all collected data. + """ + ctx = SessionContext() + + # Create multiple batches, each with 10,000 rows + batches = [] + rows_per_batch = 10000 + num_batches = 10 + + for batch_idx in range(num_batches): + start_row = batch_idx * rows_per_batch + end_row = start_row + rows_per_batch + data = { + "a": list(range(start_row, end_row)), + "b": [f"s-{i}" for i in range(start_row, end_row)], + "c": [float(i + 0.1) for i in range(start_row, end_row)], + } + batch = pa.record_batch(data) + batches.append(batch) + + # Register as record batches to maintain multi-batch structure + # Using [batches] wraps list in another list as required by register_record_batches + ctx.register_record_batches("large_multi_batch_data", [batches]) + return ctx.table("large_multi_batch_data") + + @pytest.fixture def struct_df(): ctx = SessionContext() @@ -101,6 +175,69 @@ def partitioned_df(): return ctx.create_dataframe([[batch]]) +@pytest.fixture +def clean_formatter_state(): + """Reset the HTML formatter after each test.""" + reset_formatter() + + +@pytest.fixture +def null_df(): + """Create a DataFrame with null values of different types.""" + ctx = SessionContext() + + # Create a RecordBatch with nulls across different types + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, None, 3, None], type=pa.int64()), + pa.array([4.5, 6.7, None, None], type=pa.float64()), + pa.array(["a", None, "c", None], type=pa.string()), + pa.array([True, None, False, None], type=pa.bool_()), + pa.array( + [10957, None, 18993, None], type=pa.date32() + ), # 2000-01-01, null, 2022-01-01, null + pa.array( + [946684800000, None, 1640995200000, None], type=pa.date64() + ), # 2000-01-01, null, 2022-01-01, null + ], + names=[ + "int_col", + "float_col", + "str_col", + "bool_col", + "date32_col", + "date64_col", + ], + ) + + return ctx.create_dataframe([[batch]]) + + +# custom style for testing with html formatter +class CustomStyleProvider: + def get_cell_style(self) -> str: + return ( + "background-color: #f5f5f5; color: #333; padding: 8px; border: " + "1px solid #ddd;" + ) + + def get_header_style(self) -> str: + return ( + "background-color: #4285f4; color: white; font-weight: bold; " + "padding: 10px; border: 1px solid #3367d6;" + ) + + +def count_table_rows(html_content: str) -> int: + """Count the number of table rows in HTML content. + Args: + html_content: HTML string to analyze + Returns: + Number of table rows found (number of tags) + """ + return len(re.findall(r" literal(2)).select( column("a") + column("b"), @@ -157,6 +351,66 @@ def test_filter(df): assert result.column(2) == pa.array([5]) +def test_filter_string_predicates(df): + df_str = df.filter("a > 2") + result = df_str.collect()[0] + + assert result.column(0) == pa.array([3]) + assert result.column(1) == pa.array([6]) + assert result.column(2) == pa.array([8]) + + df_mixed = df.filter("a > 1", column("b") != literal(6)) + result_mixed = df_mixed.collect()[0] + + assert result_mixed.column(0) == pa.array([2]) + assert result_mixed.column(1) == pa.array([5]) + assert result_mixed.column(2) == pa.array([5]) + + df_strings = df.filter("a > 1", "b < 6") + result_strings = df_strings.collect()[0] + + assert result_strings.column(0) == pa.array([2]) + assert result_strings.column(1) == pa.array([5]) + assert result_strings.column(2) == pa.array([5]) + + +def test_parse_sql_expr(df): + plan1 = df.filter(df.parse_sql_expr("a > 2")).logical_plan() + plan2 = df.filter(column("a") > literal(2)).logical_plan() + # object equality not implemented but string representation should match + assert str(plan1) == str(plan2) + + df1 = df.filter(df.parse_sql_expr("a > 2")).select( + column("a") + column("b"), + column("a") - column("b"), + ) + + # execute and collect the first (and only) batch + result = df1.collect()[0] + + assert result.column(0) == pa.array([9]) + assert result.column(1) == pa.array([-3]) + + df.show() + # verify that if there is no filter applied, internal dataframe is unchanged + df2 = df.filter() + assert df.df == df2.df + + df3 = df.filter(df.parse_sql_expr("a > 1"), df.parse_sql_expr("b != 6")) + result = df3.collect()[0] + + assert result.column(0) == pa.array([2]) + assert result.column(1) == pa.array([5]) + assert result.column(2) == pa.array([5]) + + +def test_show_empty(df, capsys): + df_empty = df.filter(column("a") > literal(3)) + df_empty.show() + captured = capsys.readouterr() + assert "DataFrame has no rows" in captured.out + + def test_sort(df): df = df.sort(column("b").sort(ascending=False)) @@ -166,6 +420,54 @@ def test_sort(df): assert table.to_pydict() == expected +def test_sort_string_and_expression_equivalent(df): + from datafusion import col + + result_str = df.sort("a").to_pydict() + result_expr = df.sort(col("a")).to_pydict() + assert result_str == result_expr + + +def test_sort_unsupported(df): + with pytest.raises( + TypeError, + match=f"Expected Expr or column name.*{re.escape(EXPR_TYPE_ERROR)}", + ): + df.sort(1) + + +def test_aggregate_string_and_expression_equivalent(df): + from datafusion import col + + result_str = df.aggregate("a", [f.count()]).sort("a").to_pydict() + result_expr = df.aggregate(col("a"), [f.count()]).sort("a").to_pydict() + assert result_str == result_expr + + +def test_aggregate_tuple_group_by(df): + result_list = df.aggregate(["a"], [f.count()]).sort("a").to_pydict() + result_tuple = df.aggregate(("a",), [f.count()]).sort("a").to_pydict() + assert result_tuple == result_list + + +def test_aggregate_tuple_aggs(df): + result_list = df.aggregate("a", [f.count()]).sort("a").to_pydict() + result_tuple = df.aggregate("a", (f.count(),)).sort("a").to_pydict() + assert result_tuple == result_list + + +def test_filter_string_equivalent(df): + df1 = df.filter("a > 1").to_pydict() + df2 = df.filter(column("a") > literal(1)).to_pydict() + assert df1 == df2 + + +def test_filter_string_invalid(df): + with pytest.raises(Exception) as excinfo: + df.filter("this is not valid sql").collect() + assert "Expected Expr" not in str(excinfo.value) + + def test_drop(df): df = df.drop("c") @@ -220,6 +522,21 @@ def test_tail(df): assert result.column(2) == pa.array([8]) +def test_with_column_sql_expression(df): + df = df.with_column("c", "a + b") + + # execute and collect the first (and only) batch + result = df.collect()[0] + + assert result.schema.field(0).name == "a" + assert result.schema.field(1).name == "b" + assert result.schema.field(2).name == "c" + + assert result.column(0) == pa.array([1, 2, 3]) + assert result.column(1) == pa.array([4, 5, 6]) + assert result.column(2) == pa.array([5, 7, 9]) + + def test_with_column(df): df = df.with_column("c", column("a") + column("b")) @@ -266,6 +583,37 @@ def test_with_columns(df): assert result.column(6) == pa.array([5, 7, 9]) +def test_with_columns_str(df): + df = df.with_columns( + "a + b as c", + "a + b as d", + [ + "a + b as e", + "a + b as f", + ], + g="a + b", + ) + + # execute and collect the first (and only) batch + result = df.collect()[0] + + assert result.schema.field(0).name == "a" + assert result.schema.field(1).name == "b" + assert result.schema.field(2).name == "c" + assert result.schema.field(3).name == "d" + assert result.schema.field(4).name == "e" + assert result.schema.field(5).name == "f" + assert result.schema.field(6).name == "g" + + assert result.column(0) == pa.array([1, 2, 3]) + assert result.column(1) == pa.array([4, 5, 6]) + assert result.column(2) == pa.array([5, 7, 9]) + assert result.column(3) == pa.array([5, 7, 9]) + assert result.column(4) == pa.array([5, 7, 9]) + assert result.column(5) == pa.array([5, 7, 9]) + assert result.column(6) == pa.array([5, 7, 9]) + + def test_cast(df): df = df.cast({"a": pa.float16(), "b": pa.list_(pa.uint32())}) expected = pa.schema( @@ -275,6 +623,41 @@ def test_cast(df): assert df.schema() == expected +def test_iter_batches(df): + batches = [] + for batch in df: + batches.append(batch) # noqa: PERF402 + + # Delete DataFrame to ensure RecordBatches remain valid + del df + + assert len(batches) == 1 + + batch = batches[0] + assert isinstance(batch, RecordBatch) + pa_batch = batch.to_pyarrow() + assert pa_batch.column(0).to_pylist() == [1, 2, 3] + assert pa_batch.column(1).to_pylist() == [4, 5, 6] + assert pa_batch.column(2).to_pylist() == [8, 5, 8] + + +def test_iter_returns_datafusion_recordbatch(df): + for batch in df: + assert isinstance(batch, RecordBatch) + + +def test_execute_stream_basic(df): + stream = df.execute_stream() + batches = list(stream) + + assert len(batches) == 1 + assert isinstance(batches[0], RecordBatch) + pa_batch = batches[0].to_pyarrow() + assert pa_batch.column(0).to_pylist() == [1, 2, 3] + assert pa_batch.column(1).to_pylist() == [4, 5, 6] + assert pa_batch.column(2).to_pylist() == [8, 5, 8] + + def test_with_column_renamed(df): df = df.with_column("c", column("a") + column("b")).with_column_renamed("c", "sum") @@ -305,7 +688,6 @@ def test_unnest_without_nulls(nested_df): assert result.column(1) == pa.array([7, 8, 8, 9, 9, 9]) -@pytest.mark.filterwarnings("ignore:`join_keys`:DeprecationWarning") def test_join(): ctx = SessionContext() @@ -322,26 +704,41 @@ def test_join(): df1 = ctx.create_dataframe([[batch]], "r") df2 = df.join(df1, on="a", how="inner") - df2.show() - df2 = df2.sort(column("l.a")) + df2 = df2.sort(column("a")) table = pa.Table.from_batches(df2.collect()) expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]} assert table.to_pydict() == expected - df2 = df.join(df1, left_on="a", right_on="a", how="inner") - df2.show() + # Test the default behavior for dropping duplicate keys + # Since we may have a duplicate column name and pa.Table() + # hides the fact, instead we need to explicitly check the + # resultant arrays. + df2 = df.join( + df1, left_on="a", right_on="a", how="inner", coalesce_duplicate_keys=True + ) + df2 = df2.sort(column("a")) + result = df2.collect()[0] + assert result.num_columns == 3 + assert result.column(0) == pa.array([1, 2], pa.int64()) + assert result.column(1) == pa.array([4, 5], pa.int64()) + assert result.column(2) == pa.array([8, 10], pa.int64()) + + df2 = df.join( + df1, left_on="a", right_on="a", how="inner", coalesce_duplicate_keys=False + ) df2 = df2.sort(column("l.a")) - table = pa.Table.from_batches(df2.collect()) - - expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]} - assert table.to_pydict() == expected + result = df2.collect()[0] + assert result.num_columns == 4 + assert result.column(0) == pa.array([1, 2], pa.int64()) + assert result.column(1) == pa.array([4, 5], pa.int64()) + assert result.column(2) == pa.array([1, 2], pa.int64()) + assert result.column(3) == pa.array([8, 10], pa.int64()) # Verify we don't make a breaking change to pre-43.0.0 # where users would pass join_keys as a positional argument df2 = df.join(df1, (["a"], ["a"]), how="inner") - df2.show() - df2 = df2.sort(column("l.a")) + df2 = df2.sort(column("a")) table = pa.Table.from_batches(df2.collect()) expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]} @@ -366,7 +763,7 @@ def test_join_invalid_params(): with pytest.deprecated_call(): df2 = df.join(df1, join_keys=(["a"], ["a"]), how="inner") df2.show() - df2 = df2.sort(column("l.a")) + df2 = df2.sort(column("a")) table = pa.Table.from_batches(df2.collect()) expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]} @@ -424,6 +821,58 @@ def test_join_on(): assert table.to_pydict() == expected +def test_join_full_with_drop_duplicate_keys(): + ctx = SessionContext() + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 3, 5, 7, 9]), pa.array([True, True, True, True, True])], + names=["log_time", "key_frame"], + ) + key_frame = ctx.create_dataframe([[batch]]) + + batch = pa.RecordBatch.from_arrays( + [pa.array([2, 4, 6, 8, 10])], + names=["log_time"], + ) + query_times = ctx.create_dataframe([[batch]]) + + merged = query_times.join( + key_frame, + left_on="log_time", + right_on="log_time", + how="full", + coalesce_duplicate_keys=True, + ) + merged = merged.sort(column("log_time")) + result = merged.collect()[0] + + assert result.num_columns == 2 + assert result.column(0).to_pylist() == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + +def test_join_on_invalid_expr(): + ctx = SessionContext() + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2]), pa.array([4, 5])], + names=["a", "b"], + ) + df = ctx.create_dataframe([[batch]], "l") + df1 = ctx.create_dataframe([[batch]], "r") + + with pytest.raises( + TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)" + ): + df.join_on(df1, "a") + + +def test_aggregate_invalid_aggs(df): + with pytest.raises( + TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)" + ): + df.aggregate([], "a") + + def test_distinct(): ctx = SessionContext() @@ -456,12 +905,25 @@ def test_distinct(): ), [2, 1, 3, 4, 2, 1, 3], ), + ( + "row_w_params_no_lists", + f.row_number( + order_by=column("b"), + partition_by=column("c"), + ), + [2, 1, 3, 4, 2, 1, 3], + ), ("rank", f.rank(order_by=[column("b")]), [3, 1, 3, 5, 6, 1, 6]), ( "rank_w_params", f.rank(order_by=[column("b"), column("a")], partition_by=[column("c")]), [2, 1, 3, 4, 2, 1, 3], ), + ( + "rank_w_params_no_lists", + f.rank(order_by=column("a"), partition_by=column("c")), + [1, 2, 3, 4, 1, 2, 3], + ), ( "dense_rank", f.dense_rank(order_by=[column("b")]), @@ -472,6 +934,11 @@ def test_distinct(): f.dense_rank(order_by=[column("b"), column("a")], partition_by=[column("c")]), [2, 1, 3, 4, 2, 1, 3], ), + ( + "dense_rank_w_params_no_lists", + f.dense_rank(order_by=column("a"), partition_by=column("c")), + [1, 2, 3, 4, 1, 2, 3], + ), ( "percent_rank", f.round(f.percent_rank(order_by=[column("b")]), literal(3)), @@ -487,6 +954,14 @@ def test_distinct(): ), [0.333, 0.0, 0.667, 1.0, 0.5, 0.0, 1.0], ), + ( + "percent_rank_w_params_no_lists", + f.round( + f.percent_rank(order_by=column("a"), partition_by=column("c")), + literal(3), + ), + [0.0, 0.333, 0.667, 1.0, 0.0, 0.5, 1.0], + ), ( "cume_dist", f.round(f.cume_dist(order_by=[column("b")]), literal(3)), @@ -502,6 +977,14 @@ def test_distinct(): ), [0.5, 0.25, 0.75, 1.0, 0.667, 0.333, 1.0], ), + ( + "cume_dist_w_params_no_lists", + f.round( + f.cume_dist(order_by=column("a"), partition_by=column("c")), + literal(3), + ), + [0.25, 0.5, 0.75, 1.0, 0.333, 0.667, 1.0], + ), ( "ntile", f.ntile(2, order_by=[column("b")]), @@ -512,6 +995,11 @@ def test_distinct(): f.ntile(2, order_by=[column("b"), column("a")], partition_by=[column("c")]), [1, 1, 2, 2, 1, 1, 2], ), + ( + "ntile_w_params_no_lists", + f.ntile(2, order_by=column("b"), partition_by=column("c")), + [1, 1, 2, 2, 1, 1, 2], + ), ("lead", f.lead(column("b"), order_by=[column("b")]), [7, None, 8, 9, 9, 7, None]), ( "lead_w_params", @@ -524,6 +1012,17 @@ def test_distinct(): ), [8, 7, -1, -1, -1, 9, -1], ), + ( + "lead_w_params_no_lists", + f.lead( + column("b"), + shift_offset=2, + default_value=-1, + order_by=column("b"), + partition_by=column("c"), + ), + [8, 7, -1, -1, -1, 9, -1], + ), ("lag", f.lag(column("b"), order_by=[column("b")]), [None, None, 7, 7, 8, None, 9]), ( "lag_w_params", @@ -537,9 +1036,34 @@ def test_distinct(): [-1, -1, None, 7, -1, -1, None], ), ( - "first_value", - f.first_value(column("a")).over( - Window(partition_by=[column("c")], order_by=[column("b")]) + "lag_w_params_no_lists", + f.lag( + column("b"), + shift_offset=2, + default_value=-1, + order_by=column("b"), + partition_by=column("c"), + ), + [-1, -1, None, 7, -1, -1, None], + ), + ( + "first_value", + f.first_value(column("a")).over( + Window(partition_by=[column("c")], order_by=[column("b")]) + ), + [1, 1, 1, 1, 5, 5, 5], + ), + ( + "first_value_without_list_args", + f.first_value(column("a")).over( + Window(partition_by=column("c"), order_by=column("b")) + ), + [1, 1, 1, 1, 5, 5, 5], + ), + ( + "first_value_order_by_string", + f.first_value(column("a")).over( + Window(partition_by=[column("c")], order_by="b") ), [1, 1, 1, 1, 5, 5, 5], ), @@ -585,6 +1109,27 @@ def test_window_functions(partitioned_df, name, expr, result): assert table.sort_by("a").to_pydict() == expected +@pytest.mark.parametrize("partition", ["c", df_col("c")]) +def test_rank_partition_by_accepts_string(partitioned_df, partition): + """Passing a string to partition_by should match using col().""" + df = partitioned_df.select( + f.rank(order_by=column("a"), partition_by=partition).alias("r") + ) + table = pa.Table.from_batches(df.sort(column("a")).collect()) + assert table.column("r").to_pylist() == [1, 2, 3, 4, 1, 2, 3] + + +@pytest.mark.parametrize("partition", ["c", df_col("c")]) +def test_window_partition_by_accepts_string(partitioned_df, partition): + """Window.partition_by accepts string identifiers.""" + expr = f.first_value(column("a")).over( + Window(partition_by=partition, order_by=column("b")) + ) + df = partitioned_df.select(expr.alias("fv")) + table = pa.Table.from_batches(df.sort(column("a")).collect()) + assert table.column("fv").to_pylist() == [1, 1, 1, 1, 5, 5, 5] + + @pytest.mark.parametrize( ("units", "start_bound", "end_bound"), [ @@ -613,38 +1158,19 @@ def test_valid_window_frame(units, start_bound, end_bound): ], ) def test_invalid_window_frame(units, start_bound, end_bound): - with pytest.raises(RuntimeError): + with pytest.raises(NotImplementedError, match=f"(?i){units}"): WindowFrame(units, start_bound, end_bound) def test_window_frame_defaults_match_postgres(partitioned_df): - # ref: https://github.com/apache/datafusion-python/issues/688 - - window_frame = WindowFrame("rows", None, None) - col_a = column("a") - # Using `f.window` with or without an unbounded window_frame produces the same - # results. These tests are included as a regression check but can be removed when - # f.window() is deprecated in favor of using the .over() approach. - no_frame = f.window("avg", [col_a]).alias("no_frame") - with_frame = f.window("avg", [col_a], window_frame=window_frame).alias("with_frame") - df_1 = partitioned_df.select(col_a, no_frame, with_frame) - - expected = { - "a": [0, 1, 2, 3, 4, 5, 6], - "no_frame": [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0], - "with_frame": [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0], - } - - assert df_1.sort(col_a).to_pydict() == expected - - # When order is not set, the default frame should be unounded preceeding to - # unbounded following. When order is set, the default frame is unbounded preceeding + # When order is not set, the default frame should be unbounded preceding to + # unbounded following. When order is set, the default frame is unbounded preceding # to current row. no_order = f.avg(col_a).over(Window()).alias("over_no_order") with_order = f.avg(col_a).over(Window(order_by=[col_a])).alias("over_with_order") - df_2 = partitioned_df.select(col_a, no_order, with_order) + df = partitioned_df.select(col_a, no_order, with_order) expected = { "a": [0, 1, 2, 3, 4, 5, 6], @@ -652,7 +1178,583 @@ def test_window_frame_defaults_match_postgres(partitioned_df): "over_with_order": [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0], } - assert df_2.sort(col_a).to_pydict() == expected + assert df.sort(col_a).to_pydict() == expected + + +def _build_last_value_df(df): + return df.select( + f.last_value(column("a")) + .over( + Window( + partition_by=[column("c")], + order_by=[column("b")], + window_frame=WindowFrame("rows", None, None), + ) + ) + .alias("expr"), + f.last_value(column("a")) + .over( + Window( + partition_by=[column("c")], + order_by="b", + window_frame=WindowFrame("rows", None, None), + ) + ) + .alias("str"), + ) + + +def _build_nth_value_df(df): + return df.select( + f.nth_value(column("b"), 3).over(Window(order_by=[column("a")])).alias("expr"), + f.nth_value(column("b"), 3).over(Window(order_by="a")).alias("str"), + ) + + +def _build_rank_df(df): + return df.select( + f.rank(order_by=[column("b")]).alias("expr"), + f.rank(order_by="b").alias("str"), + ) + + +def _build_array_agg_df(df): + return df.aggregate( + [column("c")], + [ + f.array_agg(column("a"), order_by=[column("a")]).alias("expr"), + f.array_agg(column("a"), order_by="a").alias("str"), + ], + ).sort(column("c")) + + +@pytest.mark.parametrize( + ("builder", "expected"), + [ + pytest.param(_build_last_value_df, [3, 3, 3, 3, 6, 6, 6], id="last_value"), + pytest.param(_build_nth_value_df, [None, None, 7, 7, 7, 7, 7], id="nth_value"), + pytest.param(_build_rank_df, [1, 1, 3, 3, 5, 6, 6], id="rank"), + pytest.param(_build_array_agg_df, [[0, 1, 2, 3], [4, 5, 6]], id="array_agg"), + ], +) +def test_order_by_string_equivalence(partitioned_df, builder, expected): + df = builder(partitioned_df) + table = pa.Table.from_batches(df.collect()) + assert table.column("expr").to_pylist() == expected + assert table.column("expr").to_pylist() == table.column("str").to_pylist() + + +def test_html_formatter_cell_dimension(df, clean_formatter_state): + """Test configuring the HTML formatter with different options.""" + # Configure with custom settings + configure_formatter( + max_width=500, + max_height=200, + enable_cell_expansion=False, + ) + + html_output = df._repr_html_() + + # Verify our configuration was applied + assert "max-height: 200px" in html_output + assert "max-width: 500px" in html_output + # With cell expansion disabled, we shouldn't see expandable-container elements + assert "expandable-container" not in html_output + + +def test_html_formatter_custom_style_provider(df, clean_formatter_state): + """Test using custom style providers with the HTML formatter.""" + + # Configure with custom style provider + configure_formatter(style_provider=CustomStyleProvider()) + + html_output = df._repr_html_() + + # Verify our custom styles were applied + assert "background-color: #4285f4" in html_output + assert "color: white" in html_output + assert "background-color: #f5f5f5" in html_output + + +def test_html_formatter_type_formatters(df, clean_formatter_state): + """Test registering custom type formatters for specific data types.""" + + # Get current formatter and register custom formatters + formatter = get_formatter() + + # Format integers with color based on value + # Using int as the type for the formatter will work since we convert + # Arrow scalar values to Python native types in _get_cell_value + def format_int(value): + return f' 2 else "blue"}">{value}' + + formatter.register_formatter(int, format_int) + + html_output = df._repr_html_() + + # Our test dataframe has values 1,2,3 so we should see: + assert '1' in html_output + + +def test_html_formatter_custom_cell_builder(df, clean_formatter_state): + """Test using a custom cell builder function.""" + + # Create a custom cell builder with distinct styling for different value ranges + def custom_cell_builder(value, row, col, table_id): + try: + num_value = int(value) + if num_value > 5: # Values > 5 get green background with indicator + return ( + '' + ) + if num_value < 3: # Values < 3 get blue background with indicator + return ( + '' + ) + except (ValueError, TypeError): + pass + + # Default styling for other cells (3, 4, 5) + return f'' + + # Set our custom cell builder + formatter = get_formatter() + formatter.set_custom_cell_builder(custom_cell_builder) + + html_output = df._repr_html_() + + # Extract cells with specific styling using regex + low_cells = re.findall( + r'', html_output + ) + mid_cells = re.findall( + r'', html_output + ) + high_cells = re.findall( + r'', html_output + ) + + # Sort the extracted values for consistent comparison + low_cells = sorted(map(int, low_cells)) + mid_cells = sorted(map(int, mid_cells)) + high_cells = sorted(map(int, high_cells)) + + # Verify specific values have the correct styling applied + assert low_cells == [1, 2] # Values < 3 + assert mid_cells == [3, 4, 5, 5] # Values 3-5 + assert high_cells == [6, 8, 8] # Values > 5 + + # Verify the exact content with styling appears in the output + assert ( + '' + in html_output + ) + assert ( + '' + in html_output + ) + assert ( + '' in html_output + ) + assert ( + '' in html_output + ) + assert ( + '' + in html_output + ) + assert ( + '' + in html_output + ) + + # Count occurrences to ensure all cells are properly styled + assert html_output.count("-low") == 2 # Two low values (1, 2) + assert html_output.count("-mid") == 4 # Four mid values (3, 4, 5, 5) + assert html_output.count("-high") == 3 # Three high values (6, 8, 8) + + # Create a custom cell builder that changes background color based on value + def custom_cell_builder(value, row, col, table_id): + # Handle numeric values regardless of their exact type + try: + num_value = int(value) + if num_value > 5: # Values > 5 get green background + return f'' + if num_value < 3: # Values < 3 get light blue background + return f'' + except (ValueError, TypeError): + pass + + # Default styling for other cells + return f'' + + # Set our custom cell builder + formatter = get_formatter() + formatter.set_custom_cell_builder(custom_cell_builder) + + html_output = df._repr_html_() + + # Verify our custom cell styling was applied + assert "background-color: #d3e9f0" in html_output # For values 1,2 + + +def test_html_formatter_custom_header_builder(df, clean_formatter_state): + """Test using a custom header builder function.""" + + # Create a custom header builder with tooltips + def custom_header_builder(field): + tooltips = { + "a": "Primary key column", + "b": "Secondary values", + "c": "Additional data", + } + tooltip = tooltips.get(field.name, "") + return ( + f'' + ) + + # Set our custom header builder + formatter = get_formatter() + formatter.set_custom_header_builder(custom_header_builder) + + html_output = df._repr_html_() + + # Verify our custom headers were applied + assert 'title="Primary key column"' in html_output + assert 'title="Secondary values"' in html_output + assert "background-color: #333; color: white" in html_output + + +def test_html_formatter_complex_customization(df, clean_formatter_state): + """Test combining multiple customization options together.""" + + # Create a dark mode style provider + class DarkModeStyleProvider: + def get_cell_style(self) -> str: + return ( + "background-color: #222; color: #eee; " + "padding: 8px; border: 1px solid #444;" + ) + + def get_header_style(self) -> str: + return ( + "background-color: #111; color: #fff; padding: 10px; " + "border: 1px solid #333;" + ) + + # Configure with dark mode style + configure_formatter( + max_cell_length=10, + style_provider=DarkModeStyleProvider(), + custom_css=""" + .datafusion-table { + font-family: monospace; + border-collapse: collapse; + } + .datafusion-table tr:hover td { + background-color: #444 !important; + } + """, + ) + + # Add type formatters for special formatting - now working with native int values + formatter = get_formatter() + formatter.register_formatter( + int, + lambda n: f'{n}', + ) + + html_output = df._repr_html_() + + # Verify our customizations were applied + assert "background-color: #222" in html_output + assert "background-color: #111" in html_output + assert ".datafusion-table" in html_output + assert "color: #5af" in html_output # Even numbers + + +def test_html_formatter_memory(df, clean_formatter_state): + """Test the memory and row control parameters in DataFrameHtmlFormatter.""" + configure_formatter(max_memory_bytes=10, min_rows=1) + html_output = df._repr_html_() + + # Count the number of table rows in the output + tr_count = count_table_rows(html_output) + # With a tiny memory limit of 10 bytes, the formatter should display + # the minimum number of rows (1) plus a message about truncation + assert tr_count == 2 # 1 for header row, 1 for data row + assert "data truncated" in html_output.lower() + + configure_formatter(max_memory_bytes=10 * MB, min_rows=1) + html_output = df._repr_html_() + # With larger memory limit and min_rows=2, should display all rows + tr_count = count_table_rows(html_output) + # Table should have header row (1) + 3 data rows = 4 rows + assert tr_count == 4 + # No truncation message should appear + assert "data truncated" not in html_output.lower() + + +def test_html_formatter_memory_boundary_conditions(large_df, clean_formatter_state): + """Test memory limit behavior at boundary conditions with large dataset. + + This test validates that the formatter correctly handles edge cases when + the memory limit is reached with a large dataset (100,000 rows), ensuring + that min_rows constraint is properly respected while respecting memory limits. + Uses large_df to actually test memory limit behavior with realistic data sizes. + """ + + # Get the raw size of the data to test boundary conditions + # First, capture output with no limits + # NOTE: max_rows=200000 is set well above the dataset size (100k rows) to ensure + # we're testing memory limits, not row limits. Default max_rows=10 would + # truncate before memory limit is reached. + configure_formatter(max_memory_bytes=10 * MB, min_rows=1, max_rows=200000) + unrestricted_output = large_df._repr_html_() + unrestricted_rows = count_table_rows(unrestricted_output) + + # Test 1: Very small memory limit should still respect min_rows + # With large dataset, this should definitely hit memory limit before min_rows + configure_formatter(max_memory_bytes=10, min_rows=1) + html_output = large_df._repr_html_() + tr_count = count_table_rows(html_output) + assert tr_count >= 2 # At least header + 1 data row (minimum) + # Should show truncation since we limited memory so aggressively + assert "data truncated" in html_output.lower() + + # Test 2: Memory limit at default size (2MB) should truncate the large dataset + # Default max_rows would truncate at 10 rows, so we don't set it here to test + # that memory limit is respected even with default row limit + configure_formatter(max_memory_bytes=2 * MB, min_rows=1) + html_output = large_df._repr_html_() + tr_count = count_table_rows(html_output) + assert tr_count >= 2 # At least header + min_rows + # Should be truncated since full dataset is much larger than 2MB + assert tr_count < unrestricted_rows + + # Test 3: Very large memory limit should show much more data + # NOTE: max_rows=200000 is critical here - without it, default max_rows=10 + # would limit output to 10 rows even though we have 100MB of memory available + configure_formatter(max_memory_bytes=100 * MB, min_rows=1, max_rows=200000) + html_output = large_df._repr_html_() + tr_count = count_table_rows(html_output) + # Should show significantly more rows, possibly all + assert tr_count > 100 # Should show substantially more rows + + # Test 4: Min rows should override memory limit + # With tiny memory and larger min_rows, min_rows should win + configure_formatter(max_memory_bytes=10, min_rows=2) + html_output = large_df._repr_html_() + tr_count = count_table_rows(html_output) + assert tr_count >= 3 # At least header + 2 data rows (min_rows) + # Should show truncation message despite min_rows being satisfied + assert "data truncated" in html_output.lower() + + # Test 5: With reasonable memory and min_rows settings + # NOTE: max_rows=200000 ensures we test memory limit behavior, not row limit + configure_formatter(max_memory_bytes=2 * MB, min_rows=10, max_rows=200000) + html_output = large_df._repr_html_() + tr_count = count_table_rows(html_output) + assert tr_count >= 11 # header + at least 10 data rows (min_rows) + # Should be truncated due to memory limit + assert tr_count < unrestricted_rows + + +def test_html_formatter_stream_early_termination( + large_multi_batch_df, clean_formatter_state +): + """Test that memory limits cause early stream termination with multi-batch data. + + This test specifically validates that the formatter stops collecting data when + the memory limit is reached, rather than collecting all data and then truncating. + The large_multi_batch_df fixture creates 10 record batches, allowing us to verify + that not all batches are consumed when memory limit is hit. + + Key difference from test_html_formatter_memory_boundary_conditions: + - Uses multi-batch DataFrame to verify stream termination behavior + - Tests with memory limit exceeded by 2-3 batches but not 1 batch + - Verifies partial data + truncation message + respects min_rows + """ + + # Get baseline: how much data fits without memory limit + configure_formatter(max_memory_bytes=100 * MB, min_rows=1, max_rows=200000) + unrestricted_output = large_multi_batch_df._repr_html_() + unrestricted_rows = count_table_rows(unrestricted_output) + + # Test 1: Memory limit exceeded by ~2 batches (each batch ~10k rows) + # With 1 batch (~1-2MB), we should have space. With 2-3 batches, we exceed limit. + # Set limit to ~3MB to ensure we collect ~1 batch before hitting limit + configure_formatter(max_memory_bytes=3 * MB, min_rows=1, max_rows=200000) + html_output = large_multi_batch_df._repr_html_() + tr_count = count_table_rows(html_output) + + # Should show significant truncation (not all 100k rows) + assert tr_count < unrestricted_rows, "Should be truncated by memory limit" + assert tr_count >= 2, "Should respect min_rows" + assert "data truncated" in html_output.lower(), "Should indicate truncation" + + # Test 2: Very tight memory limit should still respect min_rows + # Even with tiny memory (10 bytes), should show at least min_rows + configure_formatter(max_memory_bytes=10, min_rows=5, max_rows=200000) + html_output = large_multi_batch_df._repr_html_() + tr_count = count_table_rows(html_output) + + assert tr_count >= 6, "Should show header + at least min_rows (5)" + assert "data truncated" in html_output.lower(), "Should indicate truncation" + + # Test 3: Memory limit should take precedence over max_rows in early termination + # With max_rows=100 but small memory limit, should terminate early due to memory + configure_formatter(max_memory_bytes=2 * MB, min_rows=1, max_rows=100) + html_output = large_multi_batch_df._repr_html_() + tr_count = count_table_rows(html_output) + + # Should be truncated by memory limit (showing more than max_rows would suggest + # but less than unrestricted) + assert tr_count >= 2, "Should respect min_rows" + assert tr_count < unrestricted_rows, "Should be truncated" + # Output should indicate why truncation occurred + assert "data truncated" in html_output.lower() + + +def test_html_formatter_max_rows(df, clean_formatter_state): + configure_formatter(min_rows=2, max_rows=2) + html_output = df._repr_html_() + + tr_count = count_table_rows(html_output) + # Table should have header row (1) + 2 data rows = 3 rows + assert tr_count == 3 + + configure_formatter(min_rows=2, max_rows=3) + html_output = df._repr_html_() + + tr_count = count_table_rows(html_output) + # Table should have header row (1) + 3 data rows = 4 rows + assert tr_count == 4 + + +def test_html_formatter_validation(): + # Test validation for invalid parameters + + with pytest.raises(ValueError, match="max_cell_length must be a positive integer"): + DataFrameHtmlFormatter(max_cell_length=0) + + with pytest.raises(ValueError, match="max_width must be a positive integer"): + DataFrameHtmlFormatter(max_width=0) + + with pytest.raises(ValueError, match="max_height must be a positive integer"): + DataFrameHtmlFormatter(max_height=0) + + with pytest.raises(ValueError, match="max_memory_bytes must be a positive integer"): + DataFrameHtmlFormatter(max_memory_bytes=0) + + with pytest.raises(ValueError, match="max_memory_bytes must be a positive integer"): + DataFrameHtmlFormatter(max_memory_bytes=-100) + + with pytest.raises(ValueError, match="min_rows must be a positive integer"): + DataFrameHtmlFormatter(min_rows=0) + + with pytest.raises(ValueError, match="min_rows must be a positive integer"): + DataFrameHtmlFormatter(min_rows=-5) + + with pytest.raises(ValueError, match="max_rows must be a positive integer"): + DataFrameHtmlFormatter(max_rows=0) + + with pytest.raises(ValueError, match="max_rows must be a positive integer"): + DataFrameHtmlFormatter(max_rows=-10) + + with pytest.raises( + ValueError, match="min_rows must be less than or equal to max_rows" + ): + DataFrameHtmlFormatter(min_rows=5, max_rows=4) + + +def test_repr_rows_backward_compatibility(clean_formatter_state): + """Test that repr_rows parameter still works as deprecated alias.""" + # Should work when not conflicting with max_rows + with pytest.warns(DeprecationWarning, match="repr_rows parameter is deprecated"): + formatter = DataFrameHtmlFormatter(repr_rows=15, min_rows=10) + assert formatter.max_rows == 15 + assert formatter.repr_rows == 15 + + # Should fail when conflicting with max_rows + with pytest.raises(ValueError, match="Cannot specify both repr_rows and max_rows"): + DataFrameHtmlFormatter(repr_rows=5, max_rows=10) + + # Setting repr_rows via property should warn + formatter2 = DataFrameHtmlFormatter() + with pytest.warns(DeprecationWarning, match="repr_rows is deprecated"): + formatter2.repr_rows = 7 + assert formatter2.max_rows == 7 + assert formatter2.repr_rows == 7 + + +def test_configure_formatter(df, clean_formatter_state): + """Test using custom style providers with the HTML formatter and configured + parameters.""" + + # these are non-default values + max_cell_length = 10 + max_width = 500 + max_height = 30 + max_memory_bytes = 3 * MB + min_rows = 2 + max_rows = 2 + enable_cell_expansion = False + show_truncation_message = False + use_shared_styles = False + + reset_formatter() + formatter_default = get_formatter() + + assert formatter_default.max_cell_length != max_cell_length + assert formatter_default.max_width != max_width + assert formatter_default.max_height != max_height + assert formatter_default.max_memory_bytes != max_memory_bytes + assert formatter_default.min_rows != min_rows + assert formatter_default.max_rows != max_rows + assert formatter_default.enable_cell_expansion != enable_cell_expansion + assert formatter_default.show_truncation_message != show_truncation_message + assert formatter_default.use_shared_styles != use_shared_styles + + # Configure with custom style provider and additional parameters + configure_formatter( + max_cell_length=max_cell_length, + max_width=max_width, + max_height=max_height, + max_memory_bytes=max_memory_bytes, + min_rows=min_rows, + max_rows=max_rows, + enable_cell_expansion=enable_cell_expansion, + show_truncation_message=show_truncation_message, + use_shared_styles=use_shared_styles, + ) + formatter_custom = get_formatter() + assert formatter_custom.max_cell_length == max_cell_length + assert formatter_custom.max_width == max_width + assert formatter_custom.max_height == max_height + assert formatter_custom.max_memory_bytes == max_memory_bytes + assert formatter_custom.min_rows == min_rows + assert formatter_custom.max_rows == max_rows + assert formatter_custom.enable_cell_expansion == enable_cell_expansion + assert formatter_custom.show_truncation_message == show_truncation_message + assert formatter_custom.use_shared_styles == use_shared_styles + + +def test_configure_formatter_invalid_params(clean_formatter_state): + """Test that configure_formatter rejects invalid parameters.""" + with pytest.raises(ValueError, match="Invalid formatter parameters"): + configure_formatter(invalid_param=123) + + # Test with multiple parameters, one valid and one invalid + with pytest.raises(ValueError, match="Invalid formatter parameters"): + configure_formatter(max_width=500, not_a_real_param="test") + + # Test with multiple invalid parameters + with pytest.raises(ValueError, match="Invalid formatter parameters"): + configure_formatter(fake_param1="test", fake_param2=456) def test_get_dataframe(tmp_path): @@ -750,9 +1852,9 @@ def test_execution_plan(aggregate_df): # indent plan will be different for everyone due to absolute path # to filename, so we just check for some expected content assert "AggregateExec:" in indent - assert "CoalesceBatchesExec:" in indent assert "RepartitionExec:" in indent - assert "CsvExec:" in indent + assert "DataSourceExec:" in indent + assert "file_type=csv" in indent ctx = SessionContext() rows_returned = 0 @@ -774,7 +1876,7 @@ def test_execution_plan(aggregate_df): @pytest.mark.asyncio async def test_async_iteration_of_df(aggregate_df): rows_returned = 0 - async for batch in aggregate_df.execute_stream(): + async for batch in aggregate_df: assert batch is not None rows_returned += len(batch.to_pyarrow()[0]) @@ -789,6 +1891,14 @@ def test_repartition_by_hash(df): df.repartition_by_hash(column("a"), num=2) +def test_repartition_by_hash_sql_expression(df): + df.repartition_by_hash("a", num=2) + + +def test_repartition_by_hash_mix(df): + df.repartition_by_hash(column("a"), "b", num=2) + + def test_intersect(): ctx = SessionContext() @@ -852,6 +1962,18 @@ def test_collect_partitioned(): assert [[batch]] == ctx.create_dataframe([[batch]]).collect_partitioned() +def test_collect_column(ctx: SessionContext): + batch_1 = pa.RecordBatch.from_pydict({"a": [1, 2, 3]}) + batch_2 = pa.RecordBatch.from_pydict({"a": [4, 5, 6]}) + batch_3 = pa.RecordBatch.from_pydict({"a": [7, 8, 9]}) + + ctx.register_record_batches("t", [[batch_1, batch_2], [batch_3]]) + + result = ctx.table("t").sort(column("a")).collect_column("a") + expected = pa.array([1, 2, 3, 4, 5, 6, 7, 8, 9]) + assert result == expected + + def test_union(ctx): batch = pa.RecordBatch.from_arrays( [pa.array([1, 2, 3]), pa.array([4, 5, 6])], @@ -962,6 +2084,53 @@ def test_to_arrow_table(df): assert set(pyarrow_table.column_names) == {"a", "b", "c"} +def test_parquet_non_null_column_to_pyarrow(ctx, tmp_path): + path = tmp_path.joinpath("t.parquet") + + ctx.sql("create table t_(a int not null)").collect() + ctx.sql("insert into t_ values (1), (2), (3)").collect() + ctx.sql(f"copy (select * from t_) to '{path}'").collect() + + ctx.register_parquet("t", path) + pyarrow_table = ctx.sql("select max(a) as m from t").to_arrow_table() + assert pyarrow_table.to_pydict() == {"m": [3]} + + +def test_parquet_empty_batch_to_pyarrow(ctx, tmp_path): + path = tmp_path.joinpath("t.parquet") + + ctx.sql("create table t_(a int not null)").collect() + ctx.sql("insert into t_ values (1), (2), (3)").collect() + ctx.sql(f"copy (select * from t_) to '{path}'").collect() + + ctx.register_parquet("t", path) + pyarrow_table = ctx.sql("select * from t limit 0").to_arrow_table() + assert pyarrow_table.schema == pa.schema( + [ + pa.field("a", pa.int32(), nullable=False), + ] + ) + + +def test_parquet_null_aggregation_to_pyarrow(ctx, tmp_path): + path = tmp_path.joinpath("t.parquet") + + ctx.sql("create table t_(a int not null)").collect() + ctx.sql("insert into t_ values (1), (2), (3)").collect() + ctx.sql(f"copy (select * from t_) to '{path}'").collect() + + ctx.register_parquet("t", path) + pyarrow_table = ctx.sql( + "select max(a) as m from (select * from t where a < 0)" + ).to_arrow_table() + assert pyarrow_table.to_pydict() == {"m": [None]} + assert pyarrow_table.schema == pa.schema( + [ + pa.field("m", pa.int32(), nullable=True), + ] + ) + + def test_execute_stream(df): stream = df.execute_stream() assert all(batch is not None for batch in stream) @@ -1044,6 +2213,121 @@ def test_empty_to_arrow_table(df): assert set(pyarrow_table.column_names) == {"a", "b", "c"} +def test_iter_batches_dataframe(fail_collect): + ctx = SessionContext() + + batch1 = pa.record_batch([pa.array([1])], names=["a"]) + batch2 = pa.record_batch([pa.array([2])], names=["a"]) + df = ctx.create_dataframe([[batch1], [batch2]]) + + expected = [batch1, batch2] + results = [b.to_pyarrow() for b in df] + + assert len(results) == len(expected) + for exp in expected: + assert any(got.equals(exp) for got in results) + + +def test_arrow_c_stream_to_table_and_reader(fail_collect): + ctx = SessionContext() + + # Create a DataFrame with two separate record batches + batch1 = pa.record_batch([pa.array([1])], names=["a"]) + batch2 = pa.record_batch([pa.array([2])], names=["a"]) + df = ctx.create_dataframe([[batch1], [batch2]]) + + table = pa.Table.from_batches(batch.to_pyarrow() for batch in df) + batches = table.to_batches() + + assert len(batches) == 2 + expected = [batch1, batch2] + for exp in expected: + assert any(got.equals(exp) for got in batches) + assert table.schema == df.schema() + assert table.column("a").num_chunks == 2 + + reader = pa.RecordBatchReader.from_stream(df) + assert isinstance(reader, pa.RecordBatchReader) + reader_table = pa.Table.from_batches(reader) + expected = pa.Table.from_batches([batch1, batch2]) + assert reader_table.equals(expected) + + +def test_arrow_c_stream_order(): + ctx = SessionContext() + + batch1 = pa.record_batch([pa.array([1])], names=["a"]) + batch2 = pa.record_batch([pa.array([2])], names=["a"]) + + df = ctx.create_dataframe([[batch1, batch2]]) + + table = pa.Table.from_batches(batch.to_pyarrow() for batch in df) + expected = pa.Table.from_batches([batch1, batch2]) + + assert table.equals(expected) + col = table.column("a") + assert col.chunk(0)[0].as_py() == 1 + assert col.chunk(1)[0].as_py() == 2 + + +def test_arrow_c_stream_schema_selection(fail_collect): + ctx = SessionContext() + + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, 2]), + pa.array([3, 4]), + pa.array([5, 6]), + ], + names=["a", "b", "c"], + ) + df = ctx.create_dataframe([[batch]]) + + requested_schema = pa.schema([("c", pa.int64()), ("a", pa.int64())]) + + c_schema = pa_cffi.ffi.new("struct ArrowSchema*") + address = int(pa_cffi.ffi.cast("uintptr_t", c_schema)) + requested_schema._export_to_c(address) + capsule_new = ctypes.pythonapi.PyCapsule_New + capsule_new.restype = ctypes.py_object + capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + + reader = pa.RecordBatchReader.from_stream(df, schema=requested_schema) + + assert reader.schema == requested_schema + + batches = list(reader) + + assert len(batches) == 1 + expected_batch = pa.record_batch( + [pa.array([5, 6]), pa.array([1, 2])], names=["c", "a"] + ) + assert batches[0].equals(expected_batch) + + +def test_arrow_c_stream_schema_mismatch(fail_collect): + ctx = SessionContext() + + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2]), pa.array([3, 4])], names=["a", "b"] + ) + df = ctx.create_dataframe([[batch]]) + + bad_schema = pa.schema([("a", pa.string())]) + + c_schema = pa_cffi.ffi.new("struct ArrowSchema*") + address = int(pa_cffi.ffi.cast("uintptr_t", c_schema)) + bad_schema._export_to_c(address) + + capsule_new = ctypes.pythonapi.PyCapsule_New + capsule_new.restype = ctypes.py_object + capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + bad_capsule = capsule_new(ctypes.c_void_p(address), b"arrow_schema", None) + + with pytest.raises(Exception, match="Fail to merge schema"): + df.__arrow_c_stream__(bad_capsule) + + def test_to_pylist(df): # Convert datafusion dataframe to Python list pylist = df.to_pylist() @@ -1098,6 +2382,69 @@ def test_write_csv(ctx, df, tmp_path, path_to_str): assert result == expected +def generate_test_write_params() -> list[tuple]: + # Overwrite and Replace are not implemented for many table writers + insert_ops = [InsertOp.APPEND, None] + sort_by_cases = [ + (None, [1, 2, 3], "unsorted"), + (column("c"), [2, 1, 3], "single_column_expr"), + (column("a").sort(ascending=False), [3, 2, 1], "single_sort_expr"), + ([column("c"), column("b")], [2, 1, 3], "list_col_expr"), + ( + [column("c").sort(ascending=False), column("b").sort(ascending=False)], + [3, 1, 2], + "list_sort_expr", + ), + ] + + formats = ["csv", "json", "parquet", "table"] + + return [ + pytest.param( + output_format, + insert_op, + sort_by, + expected_a, + id=f"{output_format}_{test_id}", + ) + for output_format, insert_op, ( + sort_by, + expected_a, + test_id, + ) in itertools.product(formats, insert_ops, sort_by_cases) + ] + + +@pytest.mark.parametrize( + ("output_format", "insert_op", "sort_by", "expected_a"), + generate_test_write_params(), +) +def test_write_files_with_options( + ctx, df, tmp_path, output_format, insert_op, sort_by, expected_a +) -> None: + write_options = DataFrameWriteOptions(insert_operation=insert_op, sort_by=sort_by) + + if output_format == "csv": + df.write_csv(tmp_path, with_header=True, write_options=write_options) + ctx.register_csv("test_table", tmp_path) + elif output_format == "json": + df.write_json(tmp_path, write_options=write_options) + ctx.register_json("test_table", tmp_path) + elif output_format == "parquet": + df.write_parquet(tmp_path, write_options=write_options) + ctx.register_parquet("test_table", tmp_path) + elif output_format == "table": + batch = pa.RecordBatch.from_arrays([[], [], []], schema=df.schema()) + ctx.register_record_batches("test_table", [[batch]]) + ctx.table("test_table").show() + df.write_table("test_table", write_options=write_options) + + result = ctx.table("test_table").to_pydict()["a"] + ctx.table("test_table").show() + + assert result == expected_a + + @pytest.mark.parametrize("path_to_str", [True, False]) def test_write_json(ctx, df, tmp_path, path_to_str): path = str(tmp_path) if path_to_str else tmp_path @@ -1185,21 +2532,443 @@ def test_write_compressed_parquet_default_compression_level(df, tmp_path, compre df.write_parquet(str(path), compression=compression) -def test_dataframe_export(df) -> None: - # Guarantees that we have the canonical implementation - # reading our dataframe export - table = pa.table(df) - assert table.num_columns == 3 - assert table.num_rows == 3 +def test_write_parquet_with_options_default_compression(df, tmp_path): + """Test that the default compression is ZSTD.""" + df.write_parquet(tmp_path) - desired_schema = pa.schema([("a", pa.int64())]) + for file in tmp_path.rglob("*.parquet"): + metadata = pq.ParquetFile(file).metadata.to_dict() + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + assert col["compression"].lower() == "zstd" - # Verify we can request a schema - table = pa.table(df, schema=desired_schema) - assert table.num_columns == 1 - assert table.num_rows == 3 - # Expect a table of nulls if the schema don't overlap +@pytest.mark.parametrize( + "compression", + ["gzip(6)", "brotli(7)", "zstd(15)", "snappy", "uncompressed"], +) +def test_write_parquet_with_options_compression(df, tmp_path, compression): + import re + + path = tmp_path + df.write_parquet_with_options( + str(path), ParquetWriterOptions(compression=compression) + ) + + # test that the actual compression scheme is the one written + for _root, _dirs, files in os.walk(path): + for file in files: + if file.endswith(".parquet"): + metadata = pq.ParquetFile(tmp_path / file).metadata.to_dict() + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + assert col["compression"].lower() == re.sub( + r"\(\d+\)", "", compression + ) + + result = pq.read_table(str(path)).to_pydict() + expected = df.to_pydict() + + assert result == expected + + +@pytest.mark.parametrize( + "compression", + ["gzip(12)", "brotli(15)", "zstd(23)"], +) +def test_write_parquet_with_options_wrong_compression_level(df, tmp_path, compression): + path = tmp_path + + with pytest.raises(Exception, match=r"valid compression range .*? exceeded."): + df.write_parquet_with_options( + str(path), ParquetWriterOptions(compression=compression) + ) + + +@pytest.mark.parametrize("compression", ["wrong", "wrong(12)"]) +def test_write_parquet_with_options_invalid_compression(df, tmp_path, compression): + path = tmp_path + + with pytest.raises(Exception, match="Unknown or unsupported parquet compression"): + df.write_parquet_with_options( + str(path), ParquetWriterOptions(compression=compression) + ) + + +@pytest.mark.parametrize( + ("writer_version", "format_version"), + [("1.0", "1.0"), ("2.0", "2.6"), (None, "1.0")], +) +def test_write_parquet_with_options_writer_version( + df, tmp_path, writer_version, format_version +): + """Test the Parquet writer version. Note that writer_version=2.0 results in + format_version=2.6""" + if writer_version is None: + df.write_parquet_with_options(tmp_path, ParquetWriterOptions()) + else: + df.write_parquet_with_options( + tmp_path, ParquetWriterOptions(writer_version=writer_version) + ) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + assert metadata["format_version"] == format_version + + +@pytest.mark.parametrize("writer_version", ["1.2.3", "custom-version", "0"]) +def test_write_parquet_with_options_wrong_writer_version(df, tmp_path, writer_version): + """Test that invalid writer versions in Parquet throw an exception.""" + with pytest.raises(Exception, match="Invalid parquet writer version"): + df.write_parquet_with_options( + tmp_path, ParquetWriterOptions(writer_version=writer_version) + ) + + +@pytest.mark.parametrize("dictionary_enabled", [True, False, None]) +def test_write_parquet_with_options_dictionary_enabled( + df, tmp_path, dictionary_enabled +): + """Test enabling/disabling the dictionaries in Parquet.""" + df.write_parquet_with_options( + tmp_path, ParquetWriterOptions(dictionary_enabled=dictionary_enabled) + ) + # by default, the dictionary is enabled, so None results in True + result = dictionary_enabled if dictionary_enabled is not None else True + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + assert col["has_dictionary_page"] == result + + +@pytest.mark.parametrize( + ("statistics_enabled", "has_statistics"), + [("page", True), ("chunk", True), ("none", False), (None, True)], +) +def test_write_parquet_with_options_statistics_enabled( + df, tmp_path, statistics_enabled, has_statistics +): + """Test configuring the statistics in Parquet. In pyarrow we can only check for + column-level statistics, so "page" and "chunk" are tested in the same way.""" + df.write_parquet_with_options( + tmp_path, ParquetWriterOptions(statistics_enabled=statistics_enabled) + ) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + if has_statistics: + assert col["statistics"] is not None + else: + assert col["statistics"] is None + + +@pytest.mark.parametrize("max_row_group_size", [1000, 5000, 10000, 100000]) +def test_write_parquet_with_options_max_row_group_size( + large_df, tmp_path, max_row_group_size +): + """Test configuring the max number of rows per group in Parquet. These test cases + guarantee that the number of rows for each row group is max_row_group_size, given + the total number of rows is a multiple of max_row_group_size.""" + path = f"{tmp_path}/t.parquet" + large_df.write_parquet_with_options( + path, ParquetWriterOptions(max_row_group_size=max_row_group_size) + ) + + parquet = pq.ParquetFile(path) + metadata = parquet.metadata.to_dict() + for row_group in metadata["row_groups"]: + assert row_group["num_rows"] == max_row_group_size + + +@pytest.mark.parametrize("created_by", ["datafusion", "datafusion-python", "custom"]) +def test_write_parquet_with_options_created_by(df, tmp_path, created_by): + """Test configuring the created by metadata in Parquet.""" + df.write_parquet_with_options(tmp_path, ParquetWriterOptions(created_by=created_by)) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + assert metadata["created_by"] == created_by + + +@pytest.mark.parametrize("statistics_truncate_length", [5, 25, 50]) +def test_write_parquet_with_options_statistics_truncate_length( + df, tmp_path, statistics_truncate_length +): + """Test configuring the truncate limit in Parquet's row-group-level statistics.""" + ctx = SessionContext() + data = { + "a": [ + "a_the_quick_brown_fox_jumps_over_the_lazy_dog", + "m_the_quick_brown_fox_jumps_over_the_lazy_dog", + "z_the_quick_brown_fox_jumps_over_the_lazy_dog", + ], + "b": ["a_smaller", "m_smaller", "z_smaller"], + } + df = ctx.from_arrow(pa.record_batch(data)) + df.write_parquet_with_options( + tmp_path, + ParquetWriterOptions(statistics_truncate_length=statistics_truncate_length), + ) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + statistics = col["statistics"] + assert len(statistics["min"]) <= statistics_truncate_length + assert len(statistics["max"]) <= statistics_truncate_length + + +def test_write_parquet_with_options_default_encoding(tmp_path): + """Test that, by default, Parquet files are written with dictionary encoding. + Note that dictionary encoding is not used for boolean values, so it is not tested + here.""" + ctx = SessionContext() + data = { + "a": [1, 2, 3], + "b": ["1", "2", "3"], + "c": [1.01, 2.02, 3.03], + } + df = ctx.from_arrow(pa.record_batch(data)) + df.write_parquet_with_options(tmp_path, ParquetWriterOptions()) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + assert col["encodings"] == ("PLAIN", "RLE", "RLE_DICTIONARY") + + +@pytest.mark.parametrize( + ("encoding", "data_types", "result"), + [ + ("plain", ["int", "float", "str", "bool"], ("PLAIN", "RLE")), + ("rle", ["bool"], ("RLE",)), + ("delta_binary_packed", ["int"], ("RLE", "DELTA_BINARY_PACKED")), + ("delta_length_byte_array", ["str"], ("RLE", "DELTA_LENGTH_BYTE_ARRAY")), + ("delta_byte_array", ["str"], ("RLE", "DELTA_BYTE_ARRAY")), + ("byte_stream_split", ["int", "float"], ("RLE", "BYTE_STREAM_SPLIT")), + ], +) +def test_write_parquet_with_options_encoding(tmp_path, encoding, data_types, result): + """Test different encodings in Parquet in their respective support column types.""" + ctx = SessionContext() + + data = {} + for data_type in data_types: + if data_type == "int": + data["int"] = [1, 2, 3] + elif data_type == "float": + data["float"] = [1.01, 2.02, 3.03] + elif data_type == "str": + data["str"] = ["a", "b", "c"] + elif data_type == "bool": + data["bool"] = [True, False, True] + + df = ctx.from_arrow(pa.record_batch(data)) + df.write_parquet_with_options( + tmp_path, ParquetWriterOptions(encoding=encoding, dictionary_enabled=False) + ) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + assert col["encodings"] == result + + +@pytest.mark.parametrize("encoding", ["bit_packed"]) +def test_write_parquet_with_options_unsupported_encoding(df, tmp_path, encoding): + """Test that unsupported Parquet encodings do not work.""" + # BaseException is used since this throws a Rust panic: https://github.com/PyO3/pyo3/issues/3519 + with pytest.raises(BaseException, match=r"Encoding .*? is not supported"): + df.write_parquet_with_options(tmp_path, ParquetWriterOptions(encoding=encoding)) + + +@pytest.mark.parametrize("encoding", ["non_existent", "unknown", "plain123"]) +def test_write_parquet_with_options_invalid_encoding(df, tmp_path, encoding): + """Test that invalid Parquet encodings do not work.""" + with pytest.raises(Exception, match="Unknown or unsupported parquet encoding"): + df.write_parquet_with_options(tmp_path, ParquetWriterOptions(encoding=encoding)) + + +@pytest.mark.parametrize("encoding", ["plain_dictionary", "rle_dictionary"]) +def test_write_parquet_with_options_dictionary_encoding_fallback( + df, tmp_path, encoding +): + """Test that the dictionary encoding cannot be used as fallback in Parquet.""" + # BaseException is used since this throws a Rust panic: https://github.com/PyO3/pyo3/issues/3519 + with pytest.raises( + BaseException, match="Dictionary encoding can not be used as fallback encoding" + ): + df.write_parquet_with_options(tmp_path, ParquetWriterOptions(encoding=encoding)) + + +def test_write_parquet_with_options_bloom_filter(df, tmp_path): + """Test Parquet files with and without (default) bloom filters. Since pyarrow does + not expose any information about bloom filters, the easiest way to confirm that they + are actually written is to compare the file size.""" + path_no_bloom_filter = tmp_path / "1" + path_bloom_filter = tmp_path / "2" + + df.write_parquet_with_options(path_no_bloom_filter, ParquetWriterOptions()) + df.write_parquet_with_options( + path_bloom_filter, ParquetWriterOptions(bloom_filter_on_write=True) + ) + + size_no_bloom_filter = 0 + for file in path_no_bloom_filter.rglob("*.parquet"): + size_no_bloom_filter += Path(file).stat().st_size + + size_bloom_filter = 0 + for file in path_bloom_filter.rglob("*.parquet"): + size_bloom_filter += Path(file).stat().st_size + + assert size_no_bloom_filter < size_bloom_filter + + +def test_write_parquet_with_options_column_options(df, tmp_path): + """Test writing Parquet files with different options for each column, which replace + the global configs (when provided).""" + data = { + "a": [1, 2, 3], + "b": ["a", "b", "c"], + "c": [False, True, False], + "d": [1.01, 2.02, 3.03], + "e": [4, 5, 6], + } + + column_specific_options = { + "a": ParquetColumnOptions(statistics_enabled="none"), + "b": ParquetColumnOptions(encoding="plain", dictionary_enabled=False), + "c": ParquetColumnOptions( + compression="snappy", encoding="rle", dictionary_enabled=False + ), + "d": ParquetColumnOptions( + compression="zstd(6)", + encoding="byte_stream_split", + dictionary_enabled=False, + statistics_enabled="none", + ), + # column "e" will use the global configs + } + + results = { + "a": { + "statistics": False, + "compression": "brotli", + "encodings": ("PLAIN", "RLE", "RLE_DICTIONARY"), + }, + "b": { + "statistics": True, + "compression": "brotli", + "encodings": ("PLAIN", "RLE"), + }, + "c": { + "statistics": True, + "compression": "snappy", + "encodings": ("RLE",), + }, + "d": { + "statistics": False, + "compression": "zstd", + "encodings": ("RLE", "BYTE_STREAM_SPLIT"), + }, + "e": { + "statistics": True, + "compression": "brotli", + "encodings": ("PLAIN", "RLE", "RLE_DICTIONARY"), + }, + } + + ctx = SessionContext() + df = ctx.from_arrow(pa.record_batch(data)) + df.write_parquet_with_options( + tmp_path, + ParquetWriterOptions( + compression="brotli(8)", column_specific_options=column_specific_options + ), + ) + + for file in tmp_path.rglob("*.parquet"): + parquet = pq.ParquetFile(file) + metadata = parquet.metadata.to_dict() + + for row_group in metadata["row_groups"]: + for col in row_group["columns"]: + column_name = col["path_in_schema"] + result = results[column_name] + assert (col["statistics"] is not None) == result["statistics"] + assert col["compression"].lower() == result["compression"].lower() + assert col["encodings"] == result["encodings"] + + +def test_write_parquet_options(df, tmp_path): + options = ParquetWriterOptions(compression="gzip", compression_level=6) + df.write_parquet(str(tmp_path), options) + + result = pq.read_table(str(tmp_path)).to_pydict() + expected = df.to_pydict() + + assert result == expected + + +def test_write_parquet_options_error(df, tmp_path): + options = ParquetWriterOptions(compression="gzip", compression_level=6) + with pytest.raises(ValueError): + df.write_parquet(str(tmp_path), options, compression_level=1) + + +def test_write_table(ctx, df): + batch = pa.RecordBatch.from_arrays( + [pa.array([1, 2, 3])], + names=["a"], + ) + + ctx.register_record_batches("t", [[batch]]) + + df = ctx.table("t").with_column("a", column("a") * literal(-1)) + + ctx.table("t").show() + + df.write_table("t") + result = ctx.table("t").sort(column("a")).collect()[0][0].to_pylist() + expected = [-3, -2, -1, 1, 2, 3] + + assert result == expected + + +def test_dataframe_export(df) -> None: + # Guarantees that we have the canonical implementation + # reading our dataframe export + table = pa.table(df) + assert table.num_columns == 3 + assert table.num_rows == 3 + + desired_schema = pa.schema([("a", pa.int64())]) + + # Verify we can request a schema + table = pa.table(df, schema=desired_schema) + assert table.num_columns == 1 + assert table.num_rows == 3 + + # Expect a table of nulls if the schema don't overlap desired_schema = pa.schema([("g", pa.string())]) table = pa.table(df, schema=desired_schema) assert table.num_columns == 1 @@ -1242,16 +3011,561 @@ def add_with_parameter(df_internal, value: Any) -> DataFrame: assert result["new_col"] == [3 for _i in range(3)] -def test_dataframe_repr_html(df) -> None: +def test_dataframe_repr_html_structure(df, clean_formatter_state) -> None: + """Test that DataFrame._repr_html_ produces expected HTML output structure.""" + output = df._repr_html_() - ref_html = """
" + f"{field.name}
" + f"
" + "" + "" + f"{formatted_value}" + f"" + f"
" + f"
{formatted_value}
{value}-high{value}-low{value}-mid]*>(\d+)-low]*>(\d+)-mid]*>(\d+)-high1-low2-low3-mid4-mid6-high8-high{value}{value}{value}{field.name}
- - - - -
abc
148
255
368
- """ + # Since we've added a fair bit of processing to the html output, lets just verify + # the values we are expecting in the table exist. Use regex and ignore everything + # between the and . We also don't want the closing > on the + # td and th segments because that is where the formatting data is written. + + headers = ["a", "b", "c"] + headers = [f"{v}" for v in headers] + header_pattern = "(.*?)".join(headers) + header_matches = re.findall(header_pattern, output, re.DOTALL) + assert len(header_matches) == 1 + + # Update the pattern to handle values that may be wrapped in spans + body_data = [[1, 4, 8], [2, 5, 5], [3, 6, 8]] + + body_lines = [ + f"(?:]*?>)?{v}(?:)?" + for inner in body_data + for v in inner + ] + body_pattern = "(.*?)".join(body_lines) + + body_matches = re.findall(body_pattern, output, re.DOTALL) + + assert len(body_matches) == 1, "Expected pattern of values not found in HTML output" + + +def test_dataframe_repr_html_values(df, clean_formatter_state): + """Test that DataFrame._repr_html_ contains the expected data values.""" + html = df._repr_html_() + assert html is not None + + # Create a more flexible pattern that handles values being wrapped in spans + # This pattern will match the sequence of values 1,4,8,2,5,5,3,6,8 regardless + # of formatting + pattern = re.compile( + r"]*?>(?:]*?>)?1(?:)?.*?" + r"]*?>(?:]*?>)?4(?:)?.*?" + r"]*?>(?:]*?>)?8(?:)?.*?" + r"]*?>(?:]*?>)?2(?:)?.*?" + r"]*?>(?:]*?>)?5(?:)?.*?" + r"]*?>(?:]*?>)?5(?:)?.*?" + r"]*?>(?:]*?>)?3(?:)?.*?" + r"]*?>(?:]*?>)?6(?:)?.*?" + r"]*?>(?:]*?>)?8(?:)?", + re.DOTALL, + ) + + # Print debug info if the test fails + matches = re.findall(pattern, html) + if not matches: + print(f"HTML output snippet: {html[:500]}...") # noqa: T201 + + assert len(matches) > 0, "Expected pattern of values not found in HTML output" + + +def test_html_formatter_shared_styles(df, clean_formatter_state): + """Test that shared styles work correctly across multiple tables.""" + + # First, ensure we're using shared styles + configure_formatter(use_shared_styles=True) + + html_first = df._repr_html_() + html_second = df._repr_html_() + + assert "