diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000..15699a19 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,50 @@ +name: Build documentation + +on: + push: + branches: + - main + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: "pages" + cancel-in-progress: true + +defaults: + run: + shell: bash + +jobs: + build: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Build documentation + run: uvx --with "mkdocstrings[python]" zensical build + + - name: Upload artifact + uses: actions/upload-pages-artifact@v4 + with: + path: ./site + + deploy: + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + needs: build + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v5 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 6063e4db..d6b1deb5 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -32,7 +32,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.9", "3.13"] + python-version: ["3.10", "3.13"] steps: - uses: actions/checkout@v4 - name: Install uv and set the python version @@ -40,4 +40,4 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Run pre-commit (with uvx) - run: uvx pre-commit run --all-files + run: uvx pre-commit run --all-files --verbose diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a061ad1f..ab567ebd 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -40,6 +40,7 @@ name: Publish to PyPi on: release: types: [published] + workflow_dispatch: jobs: build-artifacts-mac-win: @@ -49,7 +50,7 @@ jobs: matrix: os: [ macos-latest, macos-15-intel, windows-latest ] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - uses: dtolnay/rust-toolchain@stable @@ -70,7 +71,7 @@ jobs: uv sync --dev --no-install-package xarray-sql uv run --no-project maturin build --release --strip - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: name: dist-${{ matrix.os }} path: target/wheels/* @@ -78,7 +79,7 @@ jobs: build-artifacts-manylinux-x86_64: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Build wheels for manylinux x86_64 uses: PyO3/maturin-action@v1 @@ -86,21 +87,23 @@ jobs: RUST_BACKTRACE: 1 with: rust-toolchain: stable - target: x86_64 - manylinux: 2014 + target: x86_64-unknown-linux-gnu + manylinux: 2_28 rustup-components: rust-std rustfmt sccache: 'true' - args: --release + # abi3 (see Cargo.toml) produces one wheel for all CPython >= 3.10, + # so a single interpreter is enough to build it. + args: --release --strip --out dist -i python3.10 - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: name: dist-manylinux-x86_64 - path: target/wheels/* + path: dist/* build-artifacts-manylinux-arm64: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04-arm steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Build wheels for manylinux arm64 uses: PyO3/maturin-action@v1 @@ -108,33 +111,46 @@ jobs: RUST_BACKTRACE: 1 with: rust-toolchain: stable - target: aarch64 - # Use manylinux_2_28-cross because the manylinux2014-cross has GCC 4.8.5, which causes the build to fail + target: aarch64-unknown-linux-gnu manylinux: 2_28 rustup-components: rust-std rustfmt sccache: 'true' - args: --release + # abi3 (see Cargo.toml) produces one wheel for all CPython >= 3.10, + # so a single interpreter is enough to build it. + args: --release --strip --out dist -i python3.10 - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: name: dist-manylinux-aarch64 - path: target/wheels/* + path: dist/* build-sdist: name: Source distribution runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Build sdist uses: PyO3/maturin-action@v1 with: rust-toolchain: stable - manylinux: auto rustup-components: rust-std rustfmt + # No `manylinux:` — sdist is platform-independent and runs on the + # host runner, which has python3 available. Running this in a + # manylinux container fails because maturin verifies the sdist by + # building a wheel from it, and the container has no interpreter. args: --release --sdist --out dist - - uses: actions/upload-artifact@v4 + - name: Assert sdist build does not generate wheels + run: | + if [ -d "target/wheels" ] && [ "$(ls -A target/wheels)" ]; then + echo "Error: Sdist build generated wheels" + exit 1 + else + echo "Directory is clean" + fi + + - uses: actions/upload-artifact@v6 with: name: dist-sdist path: dist/* @@ -148,7 +164,7 @@ 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-* @@ -157,7 +173,7 @@ jobs: needs: merge-build-artifacts runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v6 with: name: dist path: dist @@ -179,8 +195,12 @@ jobs: upload-to-pypi: needs: verify-built-dist runs-on: ubuntu-latest + # Only publish for real releases. `workflow_dispatch` runs build and + # verify the wheels (uploaded as artifacts) without pushing to PyPI, so + # the workflow can be exercised manually without cutting a release. + if: github.event_name == 'release' steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v6 with: name: dist path: dist @@ -190,4 +210,4 @@ jobs: enable-cache: true - name: Publish package to PyPI - run: uv publish --token ${{ secrets.PYPI_TOKEN }} --username __token__ + run: uv publish --token ${{ secrets.PYPI_TOKEN }} diff --git a/.gitignore b/.gitignore index bc7aa065..7e5bb314 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,8 @@ __pycache__ target test_data *.so +.chainlink +.claude +CHANGELOG.md +*.ipynb +/site diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 73875b9c..bacf88ac 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,17 +3,18 @@ repos: rev: v6.0.0 hooks: - id: trailing-whitespace + exclude: README.md - id: end-of-file-fixer - id: check-yaml - id: check-added-large-files - id: check-json - id: check-toml - - repo: https://github.com/google/pyink - rev: 24.10.1 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.10 hooks: - - id: pyink - # Configuration is read from pyproject.toml [tool.pyink] + - id: ruff-format + # Configuration is read from pyproject.toml [tool.ruff] - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.11.2 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..fc419a15 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,37 @@ +# AGENTS.md + +Guidance for contributors (including AI assistants) working on `xarray-sql`. It +summarizes recurring maintainer review feedback so changes land clean. + +## Documentation and comments + +- Keep docstrings and comments self-contained. Do **not** put GitHub issue or PR + numbers in docstrings or code comments; a reader should not need the issue + tracker to understand the code. Issue references belong in the commit message + and PR description (e.g. `Closes #189`), not in the source. +- Do not reference the review conversation, chat, or "the reporter" in comments. + Describe the behavior, not how it came up. + +## API surface + +- Mark internal helpers private with a leading underscore when they are not part + of the public API. + +## Tests + +- Test the public contract (values, dims, coords, attrs), not internal call + counts or private classes, so the suite survives refactors. +- Avoid redundant tests: if a public-path test already covers a behavior, do not + add a second lower-level test for the same thing. +- Make query results deterministic with `ORDER BY` so assertions do not have to + re-sort the output. +- Do not pass `dims=` to `to_dataset()` when inference already resolves them. + Reserve explicit `dims=` / `template=` for genuinely ambiguous cases (multiple + registered Datasets, or a test that is specifically exercising those + arguments). + +## Imports + +- Keep imports at the top of the file. Assume transitive dependencies are safe + to import non-locally, rather than deferring imports into functions to avoid + a dependency. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d173793d..5e39cb77 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,7 +2,7 @@ ## Where to start? -Please check out the [issues tab](https://github.com/alxmrs/xarray-sql/issues). +Please check out the [issues tab](https://github.com/xqlsystems/xarray-sql/issues). Let's have a discussion over there before proceeding with any changes. Great minds think alike -- someone may have already created an issue related to your inquiry. If there's a bug, please let us know. @@ -12,32 +12,84 @@ reading [Xarray's contributing guide](https://docs.xarray.dev/en/stable/contribu ## Developer setup -0. We use `uv` to manage the project: https://docs.astral.sh/uv/getting-started/installation/ -1. Clone the repository (bonus: [via SSH](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/adding-a-new-ssh-key-to-your-github-account)) - and `cd xarray_sql` (the project root). -2. Install dev dependencies via: `uv sync --dev` -3. Install pre-commit hooks: `uv run pre-commit install` +We use [uv](https://docs.astral.sh/uv/) to manage the project. This project +also contains a Rust extension (built with [maturin](https://www.maturin.rs/)), +so a Rust toolchain is required. - This will automatically run code formatting (pyink) and type checking (mypy) before each commit. - You can also run the hooks manually with: `uv run pre-commit run --all-files` +0. Install Rust: https://rustup.rs/ +1. Install uv: https://docs.astral.sh/uv/getting-started/installation/ +2. Clone the repository (bonus: [via SSH](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/adding-a-new-ssh-key-to-your-github-account)) + and `cd xarray-sql` (the project root). +3. Install Python dev dependencies (without building the Rust extension yet): + ```shell + uv sync --dev --no-install-package xarray-sql + ``` +4. Build and install the Rust extension into the virtual environment: + ```shell + uv run --no-project maturin develop --uv + ``` + This compiles the native code and links it so that `import xarray_sql` works. + Re-run this step whenever you modify any Rust source files under `src/`. +5. Run the test suite to verify your setup: + ```shell + uv run --no-project pytest -v . -m "not integration" + ``` +6. Install pre-commit hooks: `uvx pre-commit install` + + This will automatically run code formatting and type checking before each commit. + You can also run the hooks manually with: `uvx pre-commit run --all-files` +7. Build and serve docs locally: `uvx zensical serve` ## Before submitting a pull request... Thanks so much for your contribution! For a volunteer led project, we so appreciate your help. A few things to keep in mind: + - Please be nice. We assume good intent from you, and we ask you to do the same for us. - Development in this project will be slow if not sporadic. Reviews will come as time allows. - Every contribution, big or small, matters and deserves credit. Here are a few requests for your development process: + - We require all code to be formatted with `pyink` and type-checked with `mypy`. These checks run automatically via pre-commit hooks (see Developer setup above). If you need to run them manually: - - Formatting: `uv run pre-commit run pyink --all-files` or `uvx pyink .` - - Type checking: `uv run pre-commit run mypy --all-files` or `uv run mypy xarray_sql/` + - Formatting: `uvx pre-commit run pyink --all-files` or `uvx pyink .` + - Type checking: `uvx pre-commit run mypy --all-files` or `uvx mypy xarray_sql/` - Please include unit tests, if possible, and performance tests when you touch the core functionality (see `perf_tests/`). - It's polite to do a self review before asking for one from a maintainer. Don't stress if you forget; we all do sometimes. - Please add (or update) documentation when adding new code. We use [Google Style docstrings](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html). - We are thrilled to get documentation-only PRs -- especially spelling and typo fixes (I am a bad speller). If writing tutorials excites you, it would be to everyone's benefit. + +## Versioning Guidelines + +We're using an "experimenter's" [SemVer](https://semver.org/): We're figuring +out what a solid API should be, working towards finality and stability in the +`1.0.0` release. Until then, new features will be introduced under the minor +version (`XX.MINOR.ZZ`), and incremental (non-API surface) changes will live +under the patch version (`XX.YY.PATCH`). + +## Releasing + +To create a release, please do the following: + +1. Increment the version in the `Cargo.toml` file manually to whatever the next release will be. This needs to be merged. You can make a PR, but I often just make a quick push to main. +2. Git tag the release version: `git tag -a vXX.YY.ZZ -m 'Headline description goes here'`. +3. Push the release to the remote: `git push origin vXX.YY.ZZ` +4. In the GitHub, go to the [Releases page](https://github.com/xqlsystems/xarray-sql/releases). Please click "Draft new release." +5. On that page, select the tag that you just pushed. Add a title that follows the pattern of all other releases: (Something like: `vXX.YY.ZZ: Headline description goes here`) +6. Generate the release notes and maybe add a one line description to accompany it. +7. Click "Publish Release". This will kick of a GitHub action to build the project and push the binaries + wheels to PyPI. +8. Celebrate a successful release! + +## Undoing a bad release + +We all mess up sometimes. For example, I have often forgotten to do one of the steps (often, step 1) in the above process, and it leads to a failed release (i.e. an unsuccessful push to PyPI.) +To recover from this, please do the following and then try the above steps again: + +1. Go to the [Releases page](https://github.com/xqlsystems/xarray-sql/releases). Click into the release that didn't go so well. +2. Click the red delete button (a trash can). +3. Delete the tag in the remote: `git push --delete origin vXX.YY.ZZ` +4. Delete your tag locally with `git tag -d vXX.YY.ZZ` diff --git a/Cargo.lock b/Cargo.lock index a4bde02f..0022dae1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,54 +2,6 @@ # 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 = "adler2" version = "2.0.1" @@ -129,9 +81,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a2b10dcb159faf30d3f81f6d56c1211a5bea2ca424eabe477648a44b993320e" +checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" dependencies = [ "arrow-arith", "arrow-array", @@ -151,9 +103,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "288015089e7931843c80ed4032c5274f02b37bcb720c4a42096d50b390e70372" +checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" dependencies = [ "arrow-array", "arrow-buffer", @@ -165,9 +117,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65ca404ea6191e06bf30956394173337fa9c35f445bd447fe6c21ab944e1a23c" +checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" dependencies = [ "ahash", "arrow-buffer", @@ -176,7 +128,7 @@ dependencies = [ "chrono", "chrono-tz", "half", - "hashbrown 0.16.0", + "hashbrown 0.17.1", "num-complex", "num-integer", "num-traits", @@ -184,9 +136,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36356383099be0151dacc4245309895f16ba7917d79bdb71a7148659c9206c56" +checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" dependencies = [ "bytes", "half", @@ -196,9 +148,9 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8e372ed52bd4ee88cc1e6c3859aa7ecea204158ac640b10e187936e7e87074" +checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" dependencies = [ "arrow-array", "arrow-buffer", @@ -218,9 +170,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e4100b729fe656f2e4fb32bc5884f14acf9118d4ad532b7b33c1132e4dce896" +checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" dependencies = [ "arrow-array", "arrow-cast", @@ -233,9 +185,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf87f4ff5fc13290aa47e499a8b669a82c5977c6a1fedce22c7f542c1fd5a597" +checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" dependencies = [ "arrow-buffer", "arrow-schema", @@ -246,9 +198,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb3ca63edd2073fcb42ba112f8ae165df1de935627ead6e203d07c99445f2081" +checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" dependencies = [ "arrow-array", "arrow-buffer", @@ -262,15 +214,16 @@ dependencies = [ [[package]] name = "arrow-json" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a36b2332559d3310ebe3e173f75b29989b4412df4029a26a30cc3f7da0869297" +checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" dependencies = [ "arrow-array", "arrow-buffer", "arrow-cast", - "arrow-data", + "arrow-ord", "arrow-schema", + "arrow-select", "chrono", "half", "indexmap", @@ -286,9 +239,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13c4e0530272ca755d6814218dffd04425c5b7854b87fa741d5ff848bf50aa39" +checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" dependencies = [ "arrow-array", "arrow-buffer", @@ -299,9 +252,9 @@ dependencies = [ [[package]] name = "arrow-pyarrow" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f45c7989cb70214b2f362eaa10266d15e1a433692f2ea1514018be3aace679f4" +checksum = "d29abdf672a81c1aeb57fd2661457f9918964d49aed0e9f18932535f2a9e49ce" dependencies = [ "arrow-array", "arrow-data", @@ -311,9 +264,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b07f52788744cc71c4628567ad834cadbaeb9f09026ff1d7a4120f69edf7abd3" +checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" dependencies = [ "arrow-array", "arrow-buffer", @@ -324,9 +277,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6bb63203e8e0e54b288d0d8043ca8fa1013820822a27692ef1b78a977d879f2c" +checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ "bitflags", "serde_core", @@ -335,9 +288,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c96d8a1c180b44ecf2e66c9a2f2bbcb8b1b6f14e165ce46ac8bde211a363411b" +checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" dependencies = [ "ahash", "arrow-array", @@ -349,9 +302,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8ad6a81add9d3ea30bf8374ee8329992c7fd246ffd8b7e2f48a3cea5aa0cc9a" +checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" dependencies = [ "arrow-array", "arrow-buffer", @@ -364,33 +317,16 @@ dependencies = [ "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.19" +version = "0.4.41" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" +checksum = "d0f9ee0f6e02ffd7ad5816e9464499fba7b3effd01123b515c41d1697c43dad1" dependencies = [ - "bzip2 0.5.2", - "flate2", - "futures-core", - "memchr", + "compression-codecs", + "compression-core", "pin-project-lite", "tokio", - "xz2", - "zstd", - "zstd-safe", ] [[package]] @@ -398,9 +334,6 @@ 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-stream" @@ -421,7 +354,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -432,7 +365,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -481,7 +414,7 @@ version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" dependencies = [ - "digest", + "digest 0.10.7", ] [[package]] @@ -506,6 +439,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "brotli" version = "8.0.2" @@ -541,18 +483,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.1" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" - -[[package]] -name = "bzip2" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" -dependencies = [ - "bzip2-sys", -] +checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" [[package]] name = "bzip2" @@ -563,16 +496,6 @@ dependencies = [ "libbz2-rs-sys", ] -[[package]] -name = "bzip2-sys" -version = "0.1.13+1.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" -dependencies = [ - "cc", - "pkg-config", -] - [[package]] name = "cc" version = "1.2.39" @@ -593,9 +516,9 @@ checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" [[package]] name = "chrono" -version = "0.4.42" +version = "0.4.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", "num-traits", @@ -622,6 +545,33 @@ dependencies = [ "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-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "const-random" version = "0.1.18" @@ -642,15 +592,6 @@ dependencies = [ "tiny-keccak", ] -[[package]] -name = "const_panic" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e262cdaac42494e3ae34c43969f9cdeb7da178bdb4b66fa6a1ea2edb4c8ae652" -dependencies = [ - "typewit", -] - [[package]] name = "constant_time_eq" version = "0.3.1" @@ -663,26 +604,11 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core_extensions" -version = "1.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42bb5e5d0269fd4f739ea6cedaf29c16d81c27a7ce7582008e90eb50dcd57003" -dependencies = [ - "core_extensions_proc_macros", -] - -[[package]] -name = "core_extensions_proc_macros" -version = "1.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533d38ecd2709b7608fb8e18e4504deb99e9a72879e6aa66373a76d8dc4259ea" - [[package]] name = "cpufeatures" -version = "0.2.17" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" dependencies = [ "libc", ] @@ -696,15 +622,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "crossbeam-channel" -version = "0.5.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" -dependencies = [ - "crossbeam-utils", -] - [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -727,6 +644,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + [[package]] name = "csv" version = "1.3.1" @@ -764,15 +690,14 @@ dependencies = [ [[package]] name = "datafusion" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ba7cb113e9c0bedf9e9765926031e132fa05a1b09ba6e93a6d1a4d7044457b8" +checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040" dependencies = [ "arrow", "arrow-schema", "async-trait", - "bytes", - "bzip2 0.6.1", + "bzip2", "chrono", "datafusion-catalog", "datafusion-catalog-listing", @@ -801,28 +726,26 @@ dependencies = [ "datafusion-sql", "flate2", "futures", + "indexmap", "itertools", + "liblzma", "log", "object_store", "parking_lot", "parquet", - "rand", - "regex", - "rstest", "sqlparser", "tempfile", "tokio", "url", "uuid", - "xz2", "zstd", ] [[package]] name = "datafusion-catalog" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66a3a799f914a59b1ea343906a0486f17061f39509af74e874a866428951130d" +checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139" dependencies = [ "arrow", "async-trait", @@ -845,9 +768,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db1b113c80d7a0febcd901476a57aef378e717c54517a163ed51417d87621b0" +checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891" dependencies = [ "arrow", "async-trait", @@ -864,38 +787,39 @@ dependencies = [ "itertools", "log", "object_store", - "tokio", ] [[package]] name = "datafusion-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c10f7659e96127d25e8366be7c8be4109595d6a2c3eac70421f380a7006a1b0" +checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65" dependencies = [ - "ahash", "arrow", "arrow-ipc", + "arrow-schema", "chrono", + "foldhash 0.2.0", "half", - "hashbrown 0.14.5", + "hashbrown 0.17.1", "indexmap", + "itertools", "libc", "log", "object_store", "parquet", - "paste", "recursive", "sqlparser", "tokio", + "uuid", "web-time", ] [[package]] name = "datafusion-common-runtime" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b92065bbc6532c6651e2f7dd30b55cba0c7a14f860c7e1d15f165c41a1868d95" +checksum = "981b9dae74f78ee3d9f714fb49b01919eab975461b56149510c3ba9ea11287d1" dependencies = [ "futures", "log", @@ -904,15 +828,15 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fde13794244bc7581cd82f6fff217068ed79cdc344cafe4ab2c3a1c3510b38d6" +checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe" dependencies = [ "arrow", "async-compression", "async-trait", "bytes", - "bzip2 0.6.1", + "bzip2", "chrono", "datafusion-common", "datafusion-common-runtime", @@ -927,21 +851,22 @@ dependencies = [ "futures", "glob", "itertools", + "liblzma", "log", "object_store", + "parking_lot", "rand", "tokio", "tokio-util", "url", - "xz2", "zstd", ] [[package]] name = "datafusion-datasource-arrow" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804fa9b4ecf3157982021770617200ef7c1b2979d57bec9044748314775a9aea" +checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f" dependencies = [ "arrow", "arrow-ipc", @@ -963,9 +888,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61a1641a40b259bab38131c5e6f48fac0717bedb7dc93690e604142a849e0568" +checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada" dependencies = [ "arrow", "async-trait", @@ -986,9 +911,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adeacdb00c1d37271176f8fb6a1d8ce096baba16ea7a4b2671840c5c9c64fe85" +checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f" dependencies = [ "arrow", "async-trait", @@ -1004,13 +929,14 @@ dependencies = [ "futures", "object_store", "tokio", + "tokio-stream", ] [[package]] name = "datafusion-datasource-parquet" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43d0b60ffd66f28bfb026565d62b0a6cbc416da09814766a3797bba7d85a3cd9" +checksum = "ffd2499c1bee0eeccf6a57156105700eeeb17bc701899ac719183c4e74231450" dependencies = [ "arrow", "async-trait", @@ -1020,6 +946,7 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-physical-expr", "datafusion-physical-expr-adapter", @@ -1038,21 +965,23 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b99e13947667b36ad713549237362afb054b2d8f8cc447751e23ec61202db07" +checksum = "cb9e7e5d11130c48c8bd4e80c79a9772dd28ce6dc330baca9246205d245b9e2e" [[package]] name = "datafusion-execution" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63695643190679037bc946ad46a263b62016931547bf119859c511f7ff2f5178" +checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0" dependencies = [ "arrow", + "arrow-buffer", "async-trait", "dashmap", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr-common", "futures", "log", "object_store", @@ -1064,11 +993,12 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9a4787cbf5feb1ab351f789063398f67654a6df75c4d37d7f637dc96f951a91" +checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0" dependencies = [ "arrow", + "arrow-schema", "async-trait", "chrono", "datafusion-common", @@ -1079,7 +1009,6 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap", "itertools", - "paste", "recursive", "serde_json", "sqlparser", @@ -1087,45 +1016,54 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ce2fb1b8c15c9ac45b0863c30b268c69dc9ee7a1ee13ecf5d067738338173dc" +checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2" dependencies = [ "arrow", "datafusion-common", "indexmap", "itertools", - "paste", ] [[package]] name = "datafusion-ffi" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec510e7787641279b0336e8b79e4b7bd1385d5976875ff9b97f4269ce5231a67" +checksum = "e5660e8fa79fd51e29ce46f3026b67317ef738ebd633e106beb1a1907a406152" dependencies = [ - "abi_stable", "arrow", "arrow-schema", "async-ffi", "async-trait", - "datafusion", + "chrono", + "datafusion-catalog", "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", "datafusion-functions-aggregate-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-optimizer", + "datafusion-physical-plan", "datafusion-proto", "datafusion-proto-common", + "datafusion-session", "futures", + "libloading", "log", "prost", "semver", + "stabby", "tokio", ] [[package]] name = "datafusion-functions" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "794a9db7f7b96b3346fc007ff25e994f09b8f0511b4cf7dff651fadfe3ebb28f" +checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" dependencies = [ "arrow", "arrow-buffer", @@ -1133,31 +1071,32 @@ dependencies = [ "blake2", "blake3", "chrono", + "chrono-tz", "datafusion-common", "datafusion-doc", "datafusion-execution", "datafusion-expr", "datafusion-expr-common", "datafusion-macros", + "datafusion-physical-expr-common", "hex", "itertools", "log", "md-5", + "memchr", "num-traits", "rand", "regex", "sha2", - "unicode-segmentation", "uuid", ] [[package]] name = "datafusion-functions-aggregate" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c25210520a9dcf9c2b2cbbce31ebd4131ef5af7fc60ee92b266dc7d159cb305" +checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-doc", @@ -1167,18 +1106,18 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr", "datafusion-physical-expr-common", + "foldhash 0.2.0", "half", "log", - "paste", + "num-traits", ] [[package]] name = "datafusion-functions-aggregate-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f4a66f3b87300bb70f4124b55434d2ae3fe80455f3574701d0348da040b55d" +checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr-common", @@ -1187,9 +1126,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae5c06eed03918dc7fe7a9f082a284050f0e9ecf95d72f57712d1496da03b8c4" +checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5" dependencies = [ "arrow", "arrow-ord", @@ -1203,32 +1142,34 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", + "hashbrown 0.17.1", "itertools", + "itoa", "log", - "paste", + "memchr", ] [[package]] name = "datafusion-functions-table" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db4fed1d71738fbe22e2712d71396db04c25de4111f1ec252b8f4c6d3b25d7f5" +checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26" dependencies = [ "arrow", "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", "parking_lot", - "paste", ] [[package]] name = "datafusion-functions-window" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d92206aa5ae21892f1552b4d61758a862a70956e6fd7a95cb85db1de74bc6d1" +checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78" dependencies = [ "arrow", "datafusion-common", @@ -1239,14 +1180,13 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "log", - "paste", ] [[package]] name = "datafusion-functions-window-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53ae9bcc39800820d53a22d758b3b8726ff84a5a3e24cecef04ef4e5fdf1c7cc" +checksum = "fa84836dc2392df6f43d6a29d37fb56a8ebdc8b3f4e10ae8dc15861fd20278fb" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1254,20 +1194,20 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1063ad4c9e094b3f798acee16d9a47bd7372d9699be2de21b05c3bd3f34ab848" +checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5" dependencies = [ "datafusion-doc", "quote", - "syn 2.0.114", + "syn", ] [[package]] name = "datafusion-optimizer" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f35f9ec5d08b87fd1893a30c2929f2559c2f9806ca072d8fefca5009dc0f06a" +checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281" dependencies = [ "arrow", "chrono", @@ -1285,11 +1225,10 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c30cc8012e9eedcb48bbe112c6eff4ae5ed19cf3003cb0f505662e88b7014c5d" +checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr", @@ -1297,19 +1236,20 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", "half", - "hashbrown 0.14.5", + "hashbrown 0.17.1", "indexmap", "itertools", "parking_lot", - "paste", "petgraph", + "recursive", + "tokio", ] [[package]] name = "datafusion-physical-expr-adapter" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f9ff2dbd476221b1f67337699eff432781c4e6e1713d2aefdaa517dfbf79768" +checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235" dependencies = [ "arrow", "datafusion-common", @@ -1322,23 +1262,26 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90da43e1ec550b172f34c87ec68161986ced70fd05c8d2a2add66eef9c276f03" +checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8" dependencies = [ - "ahash", "arrow", + "chrono", "datafusion-common", "datafusion-expr-common", - "hashbrown 0.14.5", + "hashbrown 0.17.1", + "indexmap", "itertools", + "parking_lot", + "pin-project", ] [[package]] name = "datafusion-physical-optimizer" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce9804f799acd7daef3be7aaffe77c0033768ed8fdbf5fb82fc4c5f2e6bc14e6" +checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36" dependencies = [ "arrow", "datafusion-common", @@ -1355,30 +1298,32 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0acf0ad6b6924c6b1aa7d213b181e012e2d3ec0a64ff5b10ee6282ab0f8532ac" +checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc" dependencies = [ - "ahash", "arrow", + "arrow-data", + "arrow-ipc", "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.17.1", "indexmap", "itertools", "log", + "num-traits", "parking_lot", "pin-project-lite", "tokio", @@ -1386,9 +1331,9 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d368093a98a17d1449b1083ac22ed16b7128e4c67789991869480d8c4a40ecb9" +checksum = "9dd15a1ba5d3af93808241065c6c44dbca8296a189845e8a587c45c07bf0ffae" dependencies = [ "arrow", "chrono", @@ -1413,9 +1358,9 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b6aef3d5e5c1d2bc3114c4876730cb76a9bdc5a8df31ef1b6db48f0c1671895" +checksum = "90042982cf9462eb06a0b81f92efa4188dae871e7ea3ab8dc61aa9c9349b2530" dependencies = [ "arrow", "datafusion-common", @@ -1424,9 +1369,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac2c2498a1f134a9e11a9f5ed202a2a7d7e9774bd9249295593053ea3be999db" +checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3" dependencies = [ "arrow", "datafusion-common", @@ -1435,15 +1380,14 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "itertools", "log", ] [[package]] name = "datafusion-session" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f96eebd17555386f459037c65ab73aae8df09f464524c709d6a3134ad4f4776" +checksum = "15ce715fa2a61f4623cc234bcc14a3ef6a91f189128d5b14b468a6a17cdfc417" dependencies = [ "async-trait", "datafusion-common", @@ -1455,15 +1399,16 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "51.0.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fc195fe60634b2c6ccfd131b487de46dc30eccae8a3c35a13f136e7f440414f" +checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099" dependencies = [ "arrow", "bigdecimal", "chrono", "datafusion-common", "datafusion-expr", + "datafusion-functions-nested", "indexmap", "log", "recursive", @@ -1477,11 +1422,22 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.6", "subtle", ] +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -1490,7 +1446,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -1545,9 +1501,9 @@ dependencies = [ [[package]] name = "flate2" -version = "1.1.8" +version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b375d6465b98090a5f25b1c7703f3859783755aa9a80433b36e0379a3ec2f369" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", "miniz_oxide", @@ -1566,6 +1522,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" 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.2" @@ -1631,7 +1593,7 @@ checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -1646,12 +1608,6 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" -[[package]] -name = "futures-timer" -version = "3.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" - [[package]] name = "futures-util" version = "0.3.31" @@ -1670,15 +1626,6 @@ dependencies = [ "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" @@ -1708,10 +1655,21 @@ checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", "wasi 0.14.7+wasi-0.2.4", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "glob" version = "0.3.3" @@ -1735,10 +1693,6 @@ 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" @@ -1746,14 +1700,19 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "foldhash", + "foldhash 0.1.5", ] [[package]] name = "hashbrown" -version = "0.16.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] [[package]] name = "heck" @@ -1784,6 +1743,15 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "iana-time-zone" version = "0.1.64" @@ -1917,20 +1885,14 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.12.0" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6717a8d2a5a929a1a2eb43a12812498ed141a0bcfb7e8f7844fbdbe4303bba9f" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.16.0", + "hashbrown 0.17.1", ] -[[package]] -name = "indoc" -version = "2.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd" - [[package]] name = "integer-encoding" version = "3.0.4" @@ -2037,18 +1999,38 @@ checksum = "2c4a545a15244c7d945065b5d392b2d2d7f21526fba56ce51467b06ed445e8f7" [[package]] name = "libc" -version = "0.2.180" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libloading" -version = "0.7.4" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67380fd3b2fbe7527a606e18729d21c6f3951633d0500574c4dc22d2d638b9f" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" dependencies = [ "cfg-if", - "winapi", + "windows-link", +] + +[[package]] +name = "liblzma" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" +dependencies = [ + "liblzma-sys", +] + +[[package]] +name = "liblzma-sys" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f2db66f3268487b5033077f266da6777d057949b8f93c8ad82e441df25e6186" +dependencies = [ + "cc", + "libc", + "pkg-config", ] [[package]] @@ -2087,48 +2069,28 @@ checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" [[package]] name = "lz4_flex" -version = "0.12.0" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab6473172471198271ff72e9379150e9dfd70d8e533e0752a27e515b48dd375e" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" 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" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" dependencies = [ "cfg-if", - "digest", + "digest 0.11.3", ] [[package]] name = "memchr" -version = "2.7.6" +version = "2.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" - -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" [[package]] name = "miniz_oxide" @@ -2180,14 +2142,16 @@ dependencies = [ [[package]] name = "object_store" -version = "0.12.4" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c1be0c6c22ec0817cdc77d3842f721a17fd30ab6965001415b5402a74e6b740" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" dependencies = [ "async-trait", "bytes", "chrono", - "futures", + "futures-channel", + "futures-core", + "futures-util", "http", "humantime", "itertools", @@ -2242,14 +2206,13 @@ dependencies = [ [[package]] name = "parquet" -version = "57.2.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f6a2926a30477c0b95fea6c28c3072712b139337a242c2cc64817bdc20a8854" +checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" dependencies = [ "ahash", "arrow-array", "arrow-buffer", - "arrow-cast", "arrow-data", "arrow-ipc", "arrow-schema", @@ -2261,7 +2224,7 @@ dependencies = [ "flate2", "futures", "half", - "hashbrown 0.16.0", + "hashbrown 0.17.1", "lz4_flex", "num-bigint", "num-integer", @@ -2319,6 +2282,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pin-project-lite" version = "0.2.16" @@ -2363,9 +2346,9 @@ dependencies = [ [[package]] name = "proc-macro-crate" -version = "3.4.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" dependencies = [ "toml_edit", ] @@ -2399,7 +2382,7 @@ dependencies = [ "itertools", "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2413,35 +2396,32 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.26.0" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" dependencies = [ - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] name = "pyo3-build-config" -version = "0.26.0" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.26.0" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" dependencies = [ "libc", "pyo3-build-config", @@ -2449,27 +2429,27 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.26.0" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.114", + "syn", ] [[package]] name = "pyo3-macros-backend" -version = "0.26.0" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" dependencies = [ "heck", "proc-macro2", "pyo3-build-config", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2487,6 +2467,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.9.2" @@ -2533,7 +2519,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2570,53 +2556,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" - -[[package]] -name = "relative-path" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" - -[[package]] -name = "repr_offset" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb1070755bd29dffc19d0971cab794e607839ba2ef4b69a9e6fbc8733c1b72ea" -dependencies = [ - "tstr", -] - -[[package]] -name = "rstest" -version = "0.26.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" -dependencies = [ - "futures-timer", - "futures-util", - "rstest_macros", -] - -[[package]] -name = "rstest_macros" -version = "0.26.1" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" -dependencies = [ - "cfg-if", - "glob", - "proc-macro-crate", - "proc-macro2", - "quote", - "regex", - "relative-path", - "rustc_version", - "syn 2.0.114", - "unicode-ident", -] +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "rustc_version" @@ -2669,9 +2611,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "semver" -version = "1.0.27" +version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "seq-macro" @@ -2706,7 +2648,7 @@ checksum = "51e694923b8824cf0e9b382adf0f60d4e05f348f357b38833a3fa5ed7c2ede04" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2724,15 +2666,21 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.9" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", "cpufeatures", - "digest", + "digest 0.11.3", ] +[[package]] +name = "sha2-const-stable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f179d4e11094a893b82fff208f74d448a7512f99f5a0acbd5c679b705f83ed9" + [[package]] name = "shlex" version = "1.3.0" @@ -2777,9 +2725,9 @@ checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" [[package]] name = "sqlparser" -version = "0.59.0" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4591acadbcf52f0af60eafbb2c003232b2b4cd8de5f0e9437cb8b1b59046cc0f" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" dependencies = [ "log", "recursive", @@ -2788,13 +2736,47 @@ dependencies = [ [[package]] name = "sqlparser_derive" -version = "0.3.0" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "stabby" +version = "72.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7b834ec7ced12095fea1e4b07dcb7e8cf2b59b18afa3eac52494d835965a5ec" +dependencies = [ + "rustversion", + "stabby-abi", +] + +[[package]] +name = "stabby-abi" +version = "72.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da5fc6819faabb412da764b99d3b713bb55083c11e7e0c00144d386cd6a1939c" +checksum = "ff1a4f477858a5bdf927c9fab7f579899de9b13e39f8b3b3b300c89fbab632f4" dependencies = [ + "rustc_version", + "rustversion", + "sha2-const-stable", + "stabby-macros", +] + +[[package]] +name = "stabby-macros" +version = "72.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b31c4b2434980b67ad83f300a58088ba14d59454dcd79ba3d87419bbd924d31e" +dependencies = [ + "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2824,20 +2806,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "1.0.109" +version = "2.0.118" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "syn" -version = "2.0.114" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" dependencies = [ "proc-macro2", "quote", @@ -2852,7 +2823,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2891,7 +2862,7 @@ checksum = "6c5e1be1c48b9172ee610da68fd9cd2770e7a4056cb3fc98710ee6906f0c7960" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -2926,9 +2897,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.49.0" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "pin-project-lite", @@ -2937,20 +2908,32 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.6.0" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", +] + +[[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.16" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14307c986784f72ef81c89db7d9e28d6ac26d16213b109ea501696195e6e3ce5" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", @@ -2961,18 +2944,18 @@ dependencies = [ [[package]] name = "toml_datetime" -version = "0.7.5+spec-1.1.0" +version = "1.1.0+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +checksum = "97251a7c317e03ad83774a8752a7e81fb6067740609f75ea2b585b569a59198f" dependencies = [ "serde_core", ] [[package]] name = "toml_edit" -version = "0.23.10+spec-1.0.0" +version = "0.25.8+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269" +checksum = "16bff38f1d86c47f9ff0647e6838d7bb362522bdf44006c7068c2b1e606f1f3c" dependencies = [ "indexmap", "toml_datetime", @@ -2982,9 +2965,9 @@ dependencies = [ [[package]] name = "toml_parser" -version = "1.0.6+spec-1.1.0" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" dependencies = [ "winnow", ] @@ -3008,7 +2991,7 @@ checksum = "81383ab64e72a7a8b8e13130c49e3dab29def6d0c7d76a03087b3cf71c5c6903" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -3020,44 +3003,17 @@ 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 = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" -[[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.18.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" - -[[package]] -name = "typewit" -version = "1.14.2" +version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8c1ae7cc0fdb8b842d65d127cb981574b0d2b249b74d1c7a2986863dc134f71" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" [[package]] name = "unicode-ident" @@ -3077,12 +3033,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "url" version = "2.5.7" @@ -3103,11 +3053,11 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.18.1" +version = "1.23.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" +checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" dependencies = [ - "getrandom 0.3.3", + "getrandom 0.4.3", "js-sys", "wasm-bindgen", ] @@ -3175,7 +3125,7 @@ dependencies = [ "log", "proc-macro2", "quote", - "syn 2.0.114", + "syn", "wasm-bindgen-shared", ] @@ -3210,7 +3160,7 @@ checksum = "9f07d2f20d4da7b26400c9f4a0511e6e0345b040694e8a75bd41d578fa4421d7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -3244,22 +3194,6 @@ dependencies = [ "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.11" @@ -3269,12 +3203,6 @@ dependencies = [ "windows-sys 0.61.1", ] -[[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.62.1" @@ -3296,7 +3224,7 @@ checksum = "edb307e42a74fb6de9bf3a02d9712678b22399c87e6fa869d6dfcd8c1b7754e0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -3307,7 +3235,7 @@ checksum = "c0abd1ddbc6964ac14db11c7213d6532ef34bd9aa042c2e5935f59d7908b46a5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -3418,9 +3346,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "0.7.14" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" dependencies = [ "memchr", ] @@ -3439,10 +3367,11 @@ checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] name = "xarray_sql" -version = "0.1.0" +version = "0.3.3" dependencies = [ "arrow", "async-stream", + "async-trait", "datafusion", "datafusion-ffi", "futures", @@ -3451,15 +3380,6 @@ dependencies = [ "tokio", ] -[[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.8.0" @@ -3480,7 +3400,7 @@ checksum = "38da3c9736e16c5d3c8c597a9aaa5d1fa565d0532ae05e27c24aa62fb32c0ab6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", "synstructure", ] @@ -3501,7 +3421,7 @@ checksum = "88d2b8d9c68ad2b9e4340d7832716a4d21a22a1154777ad56ea55c51a9cf3831" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] @@ -3521,7 +3441,7 @@ checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", "synstructure", ] @@ -3555,14 +3475,14 @@ checksum = "5b96237efa0c878c64bd89c436f661be4e46b2f3eff1ebb976f7ef2321d2f58f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.114", + "syn", ] [[package]] name = "zlib-rs" -version = "0.5.5" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40990edd51aae2c2b6907af74ffb635029d5788228222c4bb811e9351c0caad3" +checksum = "5431d5661c32445236631278f27946e444ddafe4684cac70b185272d4f9c52d5" [[package]] name = "zstd" diff --git a/Cargo.toml b/Cargo.toml index 86358f3e..3c1de54a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,22 +1,39 @@ [package] name = "xarray_sql" -version = "0.1.0" +version = "0.3.3" authors = ["Alex Merose"] edition = "2021" -exclude = ["perf_tests/*"] +exclude = [ + "perf_tests/*", + "demo/*", + ".github/*", + ".idea/*", + ".claude/*", + ".pre-commit-config.yaml", + ".python-version", + "CONTRIBUTING.md", + "MANIFEST.in", + "uv.lock", + "xarray_sql/*_test.py", +] [dependencies] -arrow = { version = "57.2.0", features = ["pyarrow"] } +arrow = { version = "58", features = ["pyarrow"] } async-stream = "0.3" -datafusion = { version = "51.0.0" } -datafusion-ffi = { version = "51.0.0" } +async-trait = "0.1" +datafusion = { version = "54.0.0" } +datafusion-ffi = { version = "54.0.0" } futures = { version = "0.3" } -pyo3 = { version = "0.26.0", features = ["extension-module"] } +# `abi3-py310` builds against CPython's stable ABI, so a single wheel per +# platform works on all CPython >= 3.10 (matching `requires-python`). This +# lets the release workflow ship pre-built wheels for every interpreter +# without compiling per-version, avoiding local rebuilds on install. +pyo3 = { version = "0.28.0", features = ["extension-module", "abi3-py310"] } tokio = { version = "1.46.1", features = ["rt"] } [build-dependencies] -pyo3-build-config = "0.26" +pyo3-build-config = "0.28" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [lib] diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..b10f7da9 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,30 @@ +include Cargo.toml +include Cargo.lock +include pyproject.toml +include README.md +include LICENSE + +recursive-include src *.rs +recursive-include xarray_sql *.py +recursive-exclude xarray_sql *_test.py + +prune demo +prune perf_tests +prune target +prune .github +prune .idea +prune .venv +prune .mypy_cache +prune .pytest_cache +prune .claude +prune xarray_sql.egg-info + +recursive-exclude xarray_sql *.so +recursive-exclude * __pycache__ + +exclude .DS_Store +exclude .pre-commit-config.yaml +exclude .python-version +exclude .gitignore +exclude CONTRIBUTING.md +exclude uv.lock diff --git a/README.md b/README.md index edc53d30..794605ef 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,14 @@ # xarray-sql -_Query Xarray with SQL_ +_Query [Xarray](https://xarray.dev/) with SQL_ -[![ci](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/ci.yml) -[![lint](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/alxmrs/xarray-sql/actions/workflows/lint.yml) +![PyPI Version](https://img.shields.io/pypi/v/xarray-sql?color=green) +[![ci](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci.yml) +[![lint](https://github.com/xqlsystems/xarray-sql/actions/workflows/lint.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/lint.yml) +[![ci-build](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-build.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-build.yml) +[![ci-rust](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-rust.yml/badge.svg)](https://github.com/xqlsystems/xarray-sql/actions/workflows/ci-rust.yml) +[![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/xarray-sql) +[![PyPI Downloads](https://static.pepy.tech/personalized-badge/xarray-sql?period=monthly&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads%2Fmonth)](https://pepy.tech/projects/xarray-sql) ```shell pip install xarray-sql @@ -12,72 +17,167 @@ pip install xarray-sql ## What is this? This is an experiment to provide a SQL interface for array datasets. +Succinctly, we "pivot" Xarray Datasets to treat them like tables so we can run +SQL queries against them — on the query engine of your choice. xarray-sql +translates data, not queries: it registers a lazy Dataset as a table on +DataFusion (built in), DuckDB, or Polars, and turns any engine's Arrow result +back into a labeled Dataset. Dialects, geometry functions, and optimizers stay +with the engine. + +## Quickstart + +Open a Dataset, register it as a table with `from_dataset`, compute a +climatology in SQL, then write the result back to Xarray and plot it: + +> **Note:** this example also needs `pooch` and a netCDF backend (for the +> tutorial download) and `matplotlib` (for the plot): +> `pip install pooch netCDF4 matplotlib`. ```python import xarray as xr import xarray_sql as xql +# 4x-daily surface air temperature on a lat/lon grid, 2013-2014. ds = xr.tutorial.open_dataset('air_temperature') -# The same as a dask-sql Context; i.e. an Apache DataFusion Context. ctx = xql.XarrayContext() -ctx.from_dataset('air', ds, chunks=dict(time=24)) # the dataset needs to be chunked! -# DataFrame() -# +------+---------------------+-------+--------------------+ -# | lat | time | lon | air | -# +------+---------------------+-------+--------------------+ -# | 75.0 | 2013-01-01T00:00:00 | 200.0 | 241.20000000000002 | -# | 75.0 | 2013-01-01T00:00:00 | 202.5 | 242.5 | -# | 75.0 | 2013-01-01T00:00:00 | 205.0 | 243.5 | -# | 75.0 | 2013-01-01T00:00:00 | 207.5 | 244.0 | -# | 75.0 | 2013-01-01T00:00:00 | 210.0 | 244.1 | -# | 75.0 | 2013-01-01T00:00:00 | 212.5 | 243.89000000000001 | -# | 75.0 | 2013-01-01T00:00:00 | 215.0 | 243.6 | -# | 75.0 | 2013-01-01T00:00:00 | 217.5 | 243.1 | -# | 75.0 | 2013-01-01T00:00:00 | 220.0 | 242.5 | -# | 75.0 | 2013-01-01T00:00:00 | 222.5 | 241.89000000000001 | -# +------+---------------------+-------+--------------------+ -# Data truncated. +ctx.from_dataset('air', ds, chunks=dict(time=100)) -result = ctx.sql(''' +# A climatology — the mean annual cycle — computed in SQL: average air +# temperature for each month of the year, over all grid cells and years. +clim = ctx.sql(''' SELECT - "lat", "lon", AVG("air") as air_avg - FROM - "air" - GROUP BY - "lat", "lon" + CAST(date_part('month', "time") AS INTEGER) AS month, + AVG("air") AS air + FROM "air" + GROUP BY CAST(date_part('month', "time") AS INTEGER) + ORDER BY month +''') + +# Round-trip the result back to Xarray. `month` is a derived column, so name +# it as the dimension. +clim_ds = clim.to_dataset(dims=["month"]) + +# Plot the annual cycle as a time series. +clim_ds["air"].plot() # in a script, call matplotlib.pyplot.show() to display +``` + +That's the round trip — Xarray in, SQL in the middle, Xarray (and a plot) back +out. + +The same Dataset registers on other engines with one call — DuckDB gets a +native lazy table with predicate pushdown, Polars scans the same object: + +```python +import duckdb + +con = duckdb.connect() +xql.register(con, 'air', ds, chunks=dict(time=100)) +rel = con.sql('SELECT time, AVG("air") AS air FROM air GROUP BY time ORDER BY time') +xql.to_dataset(rel, template=ds) # any engine's Arrow result round-trips +``` + +See [Engines](https://xqlsystems.github.io/xarray-sql/engines/) for the support matrix, DuckDB/Polars details, +and the lazy chunked round-trip. + +## A bigger example: ARCO-ERA5 + +The same interface scales to cloud-native datasets with hundreds of variables, +like [ARCO-ERA5](https://github.com/google-research/arco-era5). + +> **Note:** reading from `gs://` requires `gcsfs` (`pip install gcsfs`). + +```python +import xarray as xr +import xarray_sql as xql + + +# Open ARCO-ERA5 — a weather dataset with 273 variables since 1940. +# Turning off dask means we don't have to wait to construct a task graph. +ds = xr.open_zarr( + 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', + chunks=None, # Turn dask off + storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required. +) + +ctx = xql.XarrayContext() +# Make sure to pass `chunks`! +ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Registration takes ~10s on my machine. + +# Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library +# pushes column projection down to Zarr, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). +# ---> Always SELECT specific columns. <--- + +# Average 2m-temperature over NYC on the morning of 2020-01-01. The library +# pushes WHERE clauses on dimension columns down to partition pruning. +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes +''').to_pandas() +# avg_c +# 0 8.640069 + +# Average temperature per pressure level, globally. +result = ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC ''') # DataFrame() -# +------+-------+--------------------+ -# | lat | lon | air_avg | -# +------+-------+--------------------+ -# | 75.0 | 205.0 | 259.88662671232834 | -# | 75.0 | 207.5 | 259.48268150684896 | -# | 75.0 | 230.0 | 258.9192123287667 | -# | 75.0 | 275.0 | 257.07574315068456 | -# | 75.0 | 322.5 | 250.11792123287654 | -# | 75.0 | 325.0 | 250.81590068493125 | -# | 72.5 | 205.0 | 262.74933904109537 | -# | 72.5 | 207.5 | 262.5384315068488 | -# | 72.5 | 230.0 | 260.8287945205475 | -# | 72.5 | 275.0 | 257.30633219178037 | -# +------+-------+--------------------+ -# Data truncated. -# - -# A table of the average temperature for each location across time. -df = result.to_pandas() -df.head() -# lat lon air_total -# 0 75.0 210.0 259.016562 -# 1 75.0 222.5 258.362212 -# 2 75.0 237.5 258.318240 -# 3 75.0 267.5 256.928497 -# 4 75.0 285.0 261.614103 +# +-------+----------------------+ +# | level | avg_c | +# +-------+----------------------+ +# | 1000 | 6.6210120796502565 | +# | 975 | 5.185637919348153 | +# | 950 | 4.028428657263021 | +# | 925 | 3.0828117974912743 | +# | 900 | 2.2109172992531967 | +# | 875 | 1.395017610194202 | +# | 850 | 0.6342670572626616 | +# | 825 | -0.21037158786759846 | +# | 800 | -1.1810754318269687 | +# | 775 | -2.3064649711534457 | +# +-------+----------------------+ + +# `latitude`/`longitude` are inferred from the registered table's surviving +# dims; `template` is kept only to recover metadata (attrs, encoding). +ctx.sql(''' + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude +''').to_dataset(template=ds) +# Size: 8MB +# Dimensions: (latitude: 721, longitude: 1440) +# Coordinates: +# * latitude (latitude) float32 3kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0 +# * longitude (longitude) float32 6kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8 +# Data variables: +# avg_c (latitude, longitude) float64 8MB -26.84 -26.84 ... -27.38 -27.38 +# Attributes: +# last_updated: 2026-06-20 02:33:34.265980+00:00 +# valid_time_start: 1940-01-01 +# valid_time_stop: 2025-12-31 +# valid_time_stop_era5t: 2026-06-14 ``` -Succinctly, we "pivot" Xarray Datasets (with consistent dimensions) to treat them like tables so we can run -SQL queries against them. +_(A runnable version of this example lives at +[`perf_tests/era5_temp_profile.py`](https://github.com/xqlsystems/xarray-sql/blob/main/perf_tests/era5_temp_profile.py).)_ ## Why build this? @@ -89,7 +189,7 @@ A few reasons: easy. * There are many cloud-native, Xarray-openable datasets, from [Google Earth Engine](https://github.com/google/Xee) - to [Pangeo Forge](https://pangeo-forge.org/). Wouldn’t it be great if these + to the [Source Cooperative](https://source.coop/products?tags=zarr). Wouldn’t it be great if these were also SQL-accessible? How can the bridge be built with minimal effort? This is a light-weight way to prove the value of the interface. @@ -107,21 +207,63 @@ That's it! _2025 update_: This library now implements a Dask-like `from_map` interface in pure DataFusion and PyArrow, but works with the same principle! +_2026 update_: Instead of `from_map()`, we create a way to translate Xarray chunks +into Arrow RecordBatches. We pass a Python callback into a DataFusion `TableProvider` +that lets the DB engine translate the underlying Dataset arrays into DataFusion partitions. +The same chunks-to-batches translation is also exposed as a +`pyarrow.dataset.Dataset` with predicate and projection pushdown, which is how +DuckDB and Polars consume registered Datasets with no engine-specific code. +Ultimately, the initial insight of the `pivot()` function -- that any ndarray can be +translated into a 2D table -- underlies this performant query mechanism. + +## Does it work? + +Yes. The recurring worry is that the SQL interface is a toy — fine for `SELECT`s, +but not for the operations geoscience actually runs. So we wrote a suite that +takes the staples of geospatial and climate analysis — the ones we assume *need* +an array library — and expresses each one in SQL, then **checks the SQL answer +against an xarray/array reference** to floating-point tolerance: + +* **Spectral indices** (NDVI) — column arithmetic over a real Sentinel-2 scene. +* **Climatology, anomalies, zonal means** — `GROUP BY` and self-`JOIN` against + the 0.25° **ARCO-ERA5** archive registered as a lazy table. Each query is + bounded to a small window (a few days over a region) and reads only that + slice — the point is that you can aim a query at a multi-decade archive and + pay only for the data it asks for, not that the query scans the whole record. +* **Forecast skill** — scoring the **Pangu-Weather** and **GraphCast** ML models + against ERA5 (WeatherBench 2) as a `JOIN` on `valid_time = init + lead`; it + reproduces the published result that GraphCast beats Pangu at every lead. +* **Raster × vector zonal stats** — a range `JOIN` of the ERA5 grid against a + table of regions. +* **Reprojection and regridding** — a `reproject(x, y, src_crs, dst_crs)` + scalar PROJ UDF, shipped as the optional geo extension + (`pip install xarray-sql[geo]`, validated against Earth Engine's own + geodesy via [Xee](https://github.com/google/Xee)) and a + sparse-weight-table `JOIN` (regridding real SRTM terrain). + +Every case matches its array reference. The headline finding: these operations +are not really "array" operations at all — they are `GROUP BY`, `JOIN`, window +functions, and `CASE` in disguise, and a query engine runs them at scale. See +[`benchmarks/geospatial/`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial/) and the write-up, +[Geospatial operations are relational operations](https://xqlsystems.github.io/xarray-sql/geospatial/). + ## Why does this work? Underneath Xarray, Dask, and Pandas, there are NumPy arrays. These are paged in chunks and represented contiguously in memory. It is only a matter of metadata -that breaks them up into ndarrays. `to_dataframe()` +that breaks them up into ndarrays. `pivot()`, which uses `to_dataframe()`, just changes this metadata (via a `ravel()`/`reshape()`), back into a column -amenable to a DataFrame. We take advantage of this light weight metadata change to -make chunked information scannable by a DB engine (DataFusion). +amenable to a DataFrame. We take advantage of this lightweight metadata change to +make chunked information scannable by a DB engine (DataFusion, DuckDB, Polars — +anything that speaks Arrow). ## What are the current limitations? -_2025 update_: TBD, DataFusion provides a whole new world! Currently, we're looking for +The sharp edges we know about — per engine and fundamental — are cataloged in +[Known issues & limitations](https://xqlsystems.github.io/xarray-sql/limitations/). Currently, we're looking for early users – "tire kickers", if you will. We'd love your input to shape the direction of this -project! Please, give this a try and [file issues](https://github.com/alxmrs/xarray-sql/issues) as -you see fit. Check out our [contributing guide](CONTRIBUTING.md), too 😉. +project! Please, give this a try and [file issues](https://github.com/xqlsystems/xarray-sql/issues) as +you see fit. Check out our [contributing guide](https://xqlsystems.github.io/xarray-sql/contributing/), too 😉. ## What would a deeper integration look like? @@ -134,23 +276,28 @@ a [virtual](https://fsspec.github.io/kerchunk/) filesystem for parquet that would internally map to Zarr. Raster-backed virtual parquet would open up integrations to numerous tools like dask, pyarrow, duckdb, and BigQuery. More thoughts on this -in [#4](https://github.com/alxmrs/xarray-sql/issues/4). +in [#4](https://github.com/xqlsystems/xarray-sql/issues/4). _2025 update_: Something like this is being built across a few projects! The ones I know about are: + - [CartoDB's Raquet](https://github.com/CartoDB/raquet) - The DataFusion community's [arrow-zarr](https://github.com/datafusion-contrib/arrow-zarr) -As of writing, this project is [amid integrating](https://github.com/alxmrs/xarray-sql/pull/69) a -rust-based DataFusion backend provided by arrow-zarr. +_2026 update_: A colleague and I are experimenting with native Zarr RDBMS engines. Check out: + +- [Zarr-Datafusion](https://lib.rs/crates/zarr-datafusion) +- [DuckDB-Zarr](https://github.com/xqlsystems/duckdb-zarr) ## Roadmap -- [ ] [@RohanDisa](https://github.com/RohanDisa) Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/alxmrs/xarray-sql/issues/93). -- [ ] Translate a single Zarr to a collection of tables via DataFusion's catalog interface [#85](https://github.com/alxmrs/xarray-sql/issues/85). -- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/alxmrs/xarray-sql/issues/68). -- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/alxmrs/xarray-sql/issues/36). -- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/alxmrs/xarray-sql/issues/4). -- [ ] (To be formally announced eventually): The 100 Trillion Row Challenge [#34](https://github.com/alxmrs/xarray-sql/issues/34). +- [x] ~Lazy evaluation via the pyarrow Dataset interface [#93](https://github.com/xqlsystems/xarray-sql/issues/93).~ _Implemented in [#100](https://github.com/xqlsystems/xarray-sql/pull/100)_ +- [x] Support proper parallelism via proper partition handling on the rust/datafusion side. [#106](https://github.com/xqlsystems/xarray-sql/issues/106) +- [x] Support core datafusion optimizations to scan less data, like [#104](https://github.com/xqlsystems/xarray-sql/issues/104), ... +- [x] Translate a single Zarr to a collection of tables [#85](https://github.com/xqlsystems/xarray-sql/issues/85). +- [ ] Distributed beyond a single node through the DataFusion integration with Ray Datasets [#68](https://github.com/xqlsystems/xarray-sql/issues/68) or Apache Ballista [#98](https://github.com/xqlsystems/xarray-sql/issues/98). +- [ ] Demo: calculate Sea Surface Temperature from 1940 - Present in SQL [#36](https://github.com/xqlsystems/xarray-sql/issues/36). +- [ ] Provide an option to integrate DataFusion directly to Zarr via Rust [#4](https://github.com/xqlsystems/xarray-sql/issues/4). +- [ ] (To be formally announced eventually): The 100 Trillion Row Challenge [#34](https://github.com/xqlsystems/xarray-sql/issues/34). ## Sponsors & Contributors @@ -164,6 +311,13 @@ I want to give a special thanks to the following folks and institutions: - Tom Nichols, Kyle Barron, Tom White, and Maxime Dion for the [Array Working Group](https://discourse.pangeo.io/t/new-working-group-for-distributed-array-computing/2734) and DataFusion-specific collaboration. +- The gracious volunteer data science students at [UCSD's DS3](https://www.ds3atucsd.com/) org, + who are working to make this library better. +- Andrew Huang for the sense of taste he brings to the project and consummate code + changes. +- Aman Kumar for spending a considerable amount of his GSoC internship + contributing to this project. + ## License @@ -183,8 +337,4 @@ See the License for the specific language governing permissions and limitations under the License. ``` -Some sources are re-distributed from Google LLC -via https://github.com/google/Xee (also Apache-2.0 License) with and without -modification (specifically, Github Actions workflows). These files are subject -to the original copyright; they include the original license header comment as -well as a note to indicate modifications (when appropriate). +All vendored code has proper license attribution. diff --git a/benchmarks/duckdb_pushdown.py b/benchmarks/duckdb_pushdown.py new file mode 100644 index 00000000..e1db700b --- /dev/null +++ b/benchmarks/duckdb_pushdown.py @@ -0,0 +1,103 @@ +"""Benchmark: DuckDB re-scannable stream vs pushdown dataset vs ceiling. + +Times the three ways DuckDB can consume the same 10M-row synthetic +dataset — the re-scannable stream (no pushdown), the default +``register()`` pushdown dataset, and an in-memory ``pyarrow.dataset`` +as the ceiling — and asserts at the end that all three returned the +same answers. Cross-engine comparisons live in +``benchmarks/geospatial/``; this measures the adapter paths within one +engine. + +Usage: python benchmarks/duckdb_pushdown.py (needs duckdb installed) +""" + +import math +import statistics +import time + +import duckdb +import numpy as np +import pandas as pd +import pyarrow.dataset as pads +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.duckdb import XarrayArrowStream + +np.random.seed(0) +N_TIME, N_LAT, N_LON = 1000, 100, 100 # 10M rows +ds = xr.Dataset( + { + "temperature": ( + ["time", "lat", "lon"], + np.random.rand(N_TIME, N_LAT, N_LON), + ), + "humidity": ( + ["time", "lat", "lon"], + np.random.rand(N_TIME, N_LAT, N_LON), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=N_TIME, freq="h"), + "lat": np.linspace(-90, 90, N_LAT), + "lon": np.linspace(-180, 180, N_LON), + }, +).chunk({"time": 50}) # 20 partitions + +con = duckdb.connect() + +QUERIES = { + "full AVG scan": "SELECT AVG(temperature) FROM {t}", + "1pct time filter": ( + "SELECT AVG(temperature) FROM {t} WHERE time < '2020-01-01 10:00:00'" + ), + "bbox filter": ( + "SELECT AVG(temperature) FROM {t} " + "WHERE lat BETWEEN 0 AND 10 AND lon BETWEEN 0 AND 20" + ), + "projection (1 of 2 vars)": "SELECT AVG(humidity) FROM {t}", + "count only": "SELECT COUNT(*) FROM {t}", +} + + +def bench(table, label, n=5): + """Times each query; returns {query: answer} for equivalence checks.""" + print(f"\n== {label} ==") + answers = {} + for qname, q in QUERIES.items(): + sql = q.format(t=table) + times = [] + for _ in range(n): + t0 = time.perf_counter() + r = con.sql(sql).fetchall() + times.append(time.perf_counter() - t0) + answers[qname] = r[0][0] + med = statistics.median(times) + print( + f" {qname:28s} {med:8.3f}s " + f"(min {min(times):.3f} / max {max(times):.3f}) -> {r[0][0]:.6g}" + ) + return answers + + +# re-scannable stream, registered via the stream wrapper explicitly: +# DuckDB scans every row, no filter/projection pushdown +con.register("t_stream", XarrayArrowStream(ds)) +stream = bench("t_stream", "stream (no pushdown)") + +# default register(): the pushdown pyarrow-dataset path +xql.register(con, "t_pushdown", ds) +pushdown = bench("t_pushdown", "register() [pushdown]") + +# ceiling: materialized pa.Table via pyarrow.dataset +table = xql.read_xarray(ds).read_all() +con.register("t_ceiling", pads.dataset(table)) +ceiling = bench("t_ceiling", "ceiling: in-memory pyarrow.dataset") + +# The timings are only meaningful if every path computed the same thing. +for qname in QUERIES: + a, b, c = stream[qname], pushdown[qname], ceiling[qname] + assert math.isclose(a, b, rel_tol=1e-9) and math.isclose( + a, c, rel_tol=1e-9 + ), f"{qname}: paths disagree — stream={a} pushdown={b} ceiling={c}" +print("\nall paths agree") diff --git a/benchmarks/geospatial/01_ndvi.py b/benchmarks/geospatial/01_ndvi.py new file mode 100644 index 00000000..7b8272f6 --- /dev/null +++ b/benchmarks/geospatial/01_ndvi.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "aiohttp", +# "requests", +# "pystac-client", +# "zarr>=3", +# "numpy", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""NDVI — "apply_ufunc over a raster" is just column arithmetic. + +The Normalized Difference Vegetation Index is the workhorse of optical remote +sensing: ``NDVI = (NIR - Red) / (NIR + Red)``, computed per pixel. The array +paradigm reaches for ``xarray.apply_ufunc`` (the coiled/benchmarks #1545 +"vectorized operations" case) to broadcast this over a whole scene. + +But a per-pixel formula over two bands is just *column arithmetic over two +columns*:: + + SELECT x, y, (nir - red) / (nir + red) AS ndvi + FROM scene + ORDER BY y, x + +Each pixel is one row; the ufunc is the SELECT expression. Invalid pixels are +already NaN (xarray decodes the band's ``_FillValue`` on open), and NaN +propagates through the arithmetic on both sides — so the masking is free, no +``CASE`` required. + +Dataset: a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample +service, discovered with ``pystac-client`` and opened the canonical way with +``xarray`` — ``xr.open_datatree`` yields the reflectance bands (B04=red, +B08=NIR at 10 m) already scaled to reflectance and carrying their ``x``/``y`` +coordinates. We read one window so the case stays bounded. Requires network; +skips cleanly if the service is offline. +""" + +from __future__ import annotations + +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, +) + +# EOPF sample-service STAC catalog; an agricultural AOI near Torino, Italy, in +# early May (peak spring growth). The search is deterministic — it resolves to +# a specific archived Sentinel-2 product. +_STAC = "https://stac.core.eopf.eodc.eu" +_BBOX = [7.2, 44.5, 7.4, 44.7] +_DATETIME = "2025-04-25/2025-05-05" + +# A 1024×1024 (~105 km²) window over vegetated valley floor. +_Y0, _X0, _N = 4_000, 6_000, 1_024 + + +def _load_scene() -> tuple[xr.Dataset, str]: + """Discover a Sentinel-2 L2A product and open its 10 m red/NIR bands. + + Idiomatic end to end: ``pystac-client`` finds the product, ``open_datatree`` + opens the hierarchical EOPF Zarr, and the ``reflectance/r10m`` node already + carries B04/B08 scaled to reflectance (nodata decoded to NaN) with + ``x``/``y`` coordinates — no manual scaling or coordinate reconstruction. + """ + try: + from pystac_client import Client + + catalog = Client.open(_STAC) + search = catalog.search( + collections=["sentinel-2-l2a"], + bbox=_BBOX, + datetime=_DATETIME, + max_items=1, + ) + item = next(search.items()) + tree = xr.open_datatree( + item.assets["product"].href, engine="zarr", chunks={} + ) + except StopIteration as exc: + raise CaseSkipped("no Sentinel-2 product found for the query") from exc + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"EOPF Sentinel-2 unavailable ({exc})") from exc + + r10m = tree["measurements/reflectance/r10m"].to_dataset() + scene = ( + r10m[["b04", "b08"]] + .rename(b04="red", b08="nir") + .isel(y=slice(_Y0, _Y0 + _N), x=slice(_X0, _X0 + _N)) + ) + return scene, item.id + + +def main() -> None: + scene, item_id = _load_scene() + n = scene.sizes["y"] * scene.sizes["x"] + print(f" Sentinel-2 L2A {item_id}") + print( + f" scene window: {dict(scene.sizes)} ({n:,} pixels, B04=red/B08=NIR)" + ) + + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + ctx.from_dataset("scene", scene, chunks={"y": 256, "x": 256}) + + sql = """ + SELECT x, y, (nir - red) / (nir + red) AS ndvi + FROM scene + ORDER BY y, x + """ + show_sql(sql) + + for _ in measured("SQL NDVI"): + got = ctx.sql_to_dataset(sql, dims=["y", "x"]).ndvi + + # Array reference: the same formula in pure xarray. ``.compute()`` reads the + # window and evaluates it here (the scene is lazy), so this measures the same + # read-and-compute the SQL side does — not just graph construction. + for _ in measured("xarray reference"): + ref = ((scene.nir - scene.red) / (scene.nir + scene.red)).compute() + + # Compare the xarray way — aligned by coordinate label, so the ORDER BY + # above is enough and neither side needs an explicit sort. + assert_grid_close("NDVI (per-pixel)", got, ref, rtol=1e-6) + + show_result(got) + + valid = ref.notnull() + print( + f"\n NDVI over {int(valid.sum()):,} valid pixels: " + f"min {float(ref.min()):.3f}, " + f"mean {float(ref.mean()):.3f}, " + f"max {float(ref.max()):.3f}" + ) + + +if __name__ == "__main__": + raise SystemExit(run_case(main, "NDVI: per-pixel column arithmetic")) diff --git a/benchmarks/geospatial/02_climatology.py b/benchmarks/geospatial/02_climatology.py new file mode 100644 index 00000000..96135ab5 --- /dev/null +++ b/benchmarks/geospatial/02_climatology.py @@ -0,0 +1,139 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Diurnal climatology — the "rechunk + grouped reduction" that is a GROUP BY. + +A *climatology* is the average value for each time-of-cycle, computed +independently at every location: "what is the typical temperature here at +06:00?" In the array paradigm (and in the coiled/benchmarks #1545 write-up) +this is the canonical painful workload — load native Zarr chunks, *rechunk* to +put all of time in one chunk ("pencils"), run a grouped reduction over the +calendar, then rechunk back to "pancakes" for output. + +The rechunking exists only to serve the array layout. The *operation* is:: + + SELECT latitude, longitude, hour_of_day, AVG("2m_temperature") + GROUP BY latitude, longitude, hour_of_day + +Group by location and time-of-cycle, average the rest — the same answer as +``da.groupby("time.hour").mean()``. ERA5 is hourly, so grouping by hour of day +gives a clean 24-bin **diurnal cycle**, one sample per day in the window. + +We register the full ARCO-ERA5 archive as a lazy table, but the climatology here +is computed over a *bounded window* — a few summer days over a CONUS-ish box. The +``WHERE`` prunes the read, so the query touches only ``2m_temperature`` over that +window and never scans the rest of the archive. The point is not that we reduce +the whole record; it is that you can aim a query at a multi-decade archive and pay +only for the slice it asks for. +""" + +from __future__ import annotations + +import datetime + +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, + timed, +) + +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" +# A few days over a CONUS-ish box (ERA5 latitude descends; lon is 0–360°E). +_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23) +_LAT_N, _LAT_S = 50.0, 25.0 +_LON_W, _LON_E = 235.0, 290.0 +_PARAMS = { + "start": _START, + "end": _END, + "lat_s": _LAT_S, + "lat_n": _LAT_N, + "lon_w": _LON_W, + "lon_e": _LON_E, +} + + +def main() -> None: + # Open the full ARCO-ERA5 archive lazily — no data is read here. ERA5 mixes + # surface (time, lat, lon) and atmospheric (… level …) variables, so register + # it as two tables under an ``era5`` schema; the query below touches only the + # surface table's 2m_temperature. + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + with timed("register full ERA5 (lazy)"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + + sql = """ + SELECT latitude, + longitude, + date_part('hour', time) AS hour, + AVG("2m_temperature") - 273.15 AS clim_c + FROM era5.surface + WHERE time BETWEEN $start AND $end + AND latitude BETWEEN $lat_s AND $lat_n + AND longitude BETWEEN $lon_w AND $lon_e + GROUP BY latitude, longitude, date_part('hour', time) + ORDER BY latitude DESC, longitude, hour + """ + show_sql(sql) + + # A climatology is a gridded product: round-trip the result back to an + # xarray Dataset keyed by (latitude, longitude, hour) — how it is used. + for _ in measured("SQL diurnal climatology (lazy read)"): + got = ctx.sql_to_dataset( + sql, + dims=["latitude", "longitude", "hour"], + param_values=_PARAMS, + ) + + # Array reference: the textbook groupby-over-the-cycle reduction, in °C — + # the same lazy window, materialized only on demand. + for _ in measured("xarray reference"): + window = ds["2m_temperature"].sel( + time=slice(_START, _END), + latitude=slice(_LAT_N, _LAT_S), + longitude=slice(_LON_W, _LON_E), + ) + ref = window.groupby("time.hour").mean("time") - 273.15 + + assert_grid_close( + "diurnal climatology (°C)", got.clim_c, ref, rtol=1e-4, atol=1e-2 + ) + + show_result(got) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Climatology: GROUP BY lat, lon, hour (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/03_zonal_mean.py b/benchmarks/geospatial/03_zonal_mean.py new file mode 100644 index 00000000..581c25f2 --- /dev/null +++ b/benchmarks/geospatial/03_zonal_mean.py @@ -0,0 +1,134 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Zonal mean — the array reduction that is secretly a GROUP BY. + +A *zonal mean* averages a field around each circle of latitude (over all +longitudes, and here over a day of hours too), collapsing a 3-D field to a 1-D +profile of value-vs-latitude — the classic pole-to-pole temperature curve. In +the array paradigm this is ``da.mean(dim=["longitude", "time"])``, a reduction +over two axes. + +Relationally it is nothing more than:: + + SELECT latitude, AVG("2m_temperature") GROUP BY latitude + +The "axes" we reduce over are just the columns we *don't* group by. Same answer, +and the SQL reads like the plain-English definition of a zonal mean. + +Dataset: the full **ARCO-ERA5** archive (0.25° global, 1.3M hourly timesteps). +The table is the whole reanalysis; ``WHERE time …`` prunes it to one day, and +the GROUP BY produces a 721-point global temperature profile. +""" + +from __future__ import annotations + +import datetime + +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, + timed, +) + +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" +# One day of hourly data, global; the WHERE below prunes ERA5 to this window. +_DAY = "2020-06-01" +_START, _END = ( + datetime.datetime(2020, 6, 1, 0), + datetime.datetime(2020, 6, 1, 23), +) + + +def main() -> None: + # Open the full ARCO-ERA5 archive (lazy, dask off) — no slicing here; the + # SQL WHERE clause prunes it to the window we ask for. + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + + print( + f" ARCO-ERA5: {ds.sizes['time']:,} hourly timesteps, " + f"{ds.sizes['latitude']}×{ds.sizes['longitude']} grid, " + f"{len(ds.data_vars)} variables (no pre-slicing)" + ) + + # ERA5 mixes surface (time, lat, lon) and atmospheric (… level …) variables, + # so register it as two tables under an ``era5`` schema. + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + with timed("register full ERA5"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + + # Pass the day's bounds as query parameters; the query still reads only that + # one day out of the whole archive. + sql = """ + SELECT latitude, + AVG("2m_temperature") - 273.15 AS air_mean_c + FROM era5.surface + WHERE time BETWEEN $start AND $end + GROUP BY latitude + ORDER BY latitude DESC + """ + show_sql(sql) + + # Round-trip the profile back to an xarray Dataset keyed by latitude. + for _ in measured("SQL zonal mean (reads one day)"): + got = ctx.sql_to_dataset( + sql, + dims=["latitude"], + param_values={"start": _START, "end": _END}, + ) + + # Array reference: reduce the same day over the two un-grouped axes. + for _ in measured("xarray reference"): + ref = ( + ds["2m_temperature"].sel(time=_DAY).mean(["longitude", "time"]) + - 273.15 + ) + + assert_grid_close( + "zonal mean (2m_temp vs latitude, °C)", + got.air_mean_c, + ref, + rtol=1e-4, + atol=1e-3, + ) + + show_result(got) + + print("\n Global temperature profile (every 72nd parallel, °C):") + print(got.air_mean_c.isel(latitude=slice(None, None, 72)).to_series()) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Zonal mean: GROUP BY latitude (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/04_anomaly.py b/benchmarks/geospatial/04_anomaly.py new file mode 100644 index 00000000..ff497489 --- /dev/null +++ b/benchmarks/geospatial/04_anomaly.py @@ -0,0 +1,146 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Temperature anomaly — "broadcast-subtract the climatology" is a self-JOIN. + +An *anomaly* is the departure of each observation from its climatological +normal: ``anomaly(t) = T(t) − climatology(hour-of-day(t))`` at each cell. The +array paradigm computes the climatology, then leans on xarray's grouped +broadcasting to line it back up with every timestep: +``ds.groupby("time.hour") - climatology``. + +That broadcast — "attach each cell's normal back onto every matching timestep" — +is exactly a relational **JOIN** on the grouping key. So the anomaly is a +climatology CTE joined back to the raw observations:: + + WITH clim AS (SELECT latitude, longitude, hour, AVG(T) ... GROUP BY ...) + SELECT a.T - c.clim_t AS anomaly + FROM era5 a JOIN clim c + ON (a.latitude, a.longitude, hour(a.time)) = (c.latitude, c.longitude, c.hour) + +We register the full ARCO-ERA5 archive as a lazy table, but the anomaly here is +computed over a *bounded window* (a few summer days over a CONUS-ish box): both +the climatology CTE and the outer scan read only ``2m_temperature``, and only +over the window the ``WHERE`` asks for — never the rest of the archive. You can +aim a query at the whole archive and pay only for the slice it asks for. +""" + +from __future__ import annotations + +import datetime + +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, + timed, +) + +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" +_START, _END = datetime.datetime(2020, 6, 1), datetime.datetime(2020, 6, 3, 23) +_LAT_N, _LAT_S = 50.0, 25.0 +_LON_W, _LON_E = 235.0, 290.0 +_PARAMS = { + "start": _START, + "end": _END, + "lat_s": _LAT_S, + "lat_n": _LAT_N, + "lon_w": _LON_W, + "lon_e": _LON_E, +} + + +def main() -> None: + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + with timed("register full ERA5 (lazy)"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + + sql = """ + WITH clim AS ( + SELECT latitude, longitude, + date_part('hour', time) AS hour, + AVG("2m_temperature") AS clim_t + FROM era5.surface + WHERE time BETWEEN $start AND $end + AND latitude BETWEEN $lat_s AND $lat_n + AND longitude BETWEEN $lon_w AND $lon_e + GROUP BY latitude, longitude, date_part('hour', time) + ) + SELECT a.time, a.latitude, a.longitude, + a."2m_temperature" - c.clim_t AS anomaly + FROM era5.surface a + JOIN clim c + ON a.latitude = c.latitude + AND a.longitude = c.longitude + AND date_part('hour', a.time) = c.hour + WHERE a.time BETWEEN $start AND $end + AND a.latitude BETWEEN $lat_s AND $lat_n + AND a.longitude BETWEEN $lon_w AND $lon_e + ORDER BY a.time, a.latitude DESC, a.longitude + """ + show_sql(sql) + + # The anomaly is a gridded field; round-trip it to (time, lat, lon). + for _ in measured("SQL anomaly (climatology CTE self-join, lazy read)"): + got = ctx.sql_to_dataset( + sql, + dims=["time", "latitude", "longitude"], + param_values=_PARAMS, + ) + + # Array reference: grouped broadcast-subtract, in pure xarray (lazy window). + for _ in measured("xarray reference"): + window = ds["2m_temperature"].sel( + time=slice(_START, _END), + latitude=slice(_LAT_N, _LAT_S), + longitude=slice(_LON_W, _LON_E), + ) + grouped = window.groupby("time.hour") + ref = grouped - grouped.mean("time") + + assert_grid_close( + "anomaly (T − diurnal climatology)", + got.anomaly, + ref, + rtol=1e-3, + atol=1e-2, + ) + + show_result(got) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Anomaly: climatology CTE self-JOIN (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/05_forecast_skill.py b/benchmarks/geospatial/05_forecast_skill.py new file mode 100644 index 00000000..dac68816 --- /dev/null +++ b/benchmarks/geospatial/05_forecast_skill.py @@ -0,0 +1,200 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "pandas", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Forecast skill — scoring ML weather models against ERA5 is a JOIN + aggregate. + +Scoring the **Pangu-Weather** and **GraphCast** machine-learning forecast models +against ERA5 ground truth is the headline workload of +[WeatherBench 2](https://weatherbench2.readthedocs.io/). A forecast is indexed by +*initialization time* and *lead time* (``prediction_timedelta``); the truth is +indexed by *valid time*. Evaluation aligns them by ``valid_time = init + lead`` +and reduces the error to RMSE as a function of lead — the classic "error grows +with forecast horizon" curve. + +That alignment is a relational **JOIN**, and ``valid_time = init + lead`` is just +timestamp + duration arithmetic the engine does natively:: + + SELECT f.model, f.prediction_timedelta AS lead, + SQRT(AVG(POWER(f.t - e.t, 2))) AS rmse + FROM forecasts f + JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude + GROUP BY f.model, f.prediction_timedelta + +We stack the two models along a ``model`` dimension into a single forecast +table, so one query scores them together, grouped by the ``model`` column. The +forecasts and ERA5 are opened lazily, and the JOIN reads only what it needs. + +Datasets: WeatherBench 2 **Pangu**, **GraphCast**, and **ERA5** at a coarse +64×32 grid (so the demo is small and fast), read from the public ``gs:// +weatherbench2`` bucket. Requires network; skips cleanly offline. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, +) + +_GRID = "64x32_equiangular_conservative" +_ERA5 = f"gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-{_GRID}.zarr" +_PANGU = f"gs://weatherbench2/datasets/pangu/2018-2022_0012_{_GRID}.zarr" +_GRAPHCAST = ( + "gs://weatherbench2/datasets/graphcast/2020/" + f"date_range_2019-11-16_2021-02-01_12_hours-{_GRID}.zarr" +) +_VAR = "2m_temperature" +_INIT = slice("2020-01-01", "2020-01-10") # 20 init times (12-hourly) + + +def _open(url: str) -> xr.Dataset: + try: + import gcsfs # noqa: F401 + + # decode_timedelta=True: forecasts store prediction_timedelta as a + # real duration (and it silences xarray's decode-timedelta warning). + return xr.open_zarr( + url, + chunks=None, + storage_options={"token": "anon"}, + decode_timedelta=True, + ) + except Exception as exc: # noqa: BLE001 + raise CaseSkipped(f"WeatherBench2 unavailable ({exc})") from exc + + +def _reference_rmse(forecasts: xr.Dataset, truth: xr.Dataset) -> xr.DataArray: + """xarray reference: per (model, lead), align truth at valid_time, take RMSE. + + The 64×32 windows are tiny, so the reference reads them into memory and + reduces there; the SQL side above stays lazy. We use ``.compute()`` rather + than ``.load()`` deliberately: ``.load()`` caches the data *in place* on the + shared ``forecasts``/``truth`` objects (which the SQL table also reads from), + which would let a profiled reference serve a warm read — ``.compute()`` + returns a fresh array and leaves the inputs lazy, so each measurement is cold. + """ + f = forecasts[_VAR].compute() + e = truth[_VAR].compute() + leads = f.prediction_timedelta.values + per_lead = [] + for lead in leads: + e_at_valid = e.sel(time=f.time.values + lead) # (init, lat, lon) + diff = f.sel(prediction_timedelta=lead) - e_at_valid.values + per_lead.append( + np.sqrt((diff**2).mean(["time", "latitude", "longitude"])) + ) + return ( + xr.concat(per_lead, dim="lead") + .assign_coords(lead=leads) + .transpose("model", "lead") + ) + + +def main() -> None: + # Open everything lazily — no .load() here. + era5 = _open(_ERA5) + + # The two models store different pressure-level sets (Pangu 13, GraphCast + # 37), so we keep the common surface field 2m_temperature and stack the + # models along a `model` dimension into one forecast table. Snap the grid + # onto ERA5's exact coordinates (same 64×32 grid) so the join on latitude and + # longitude lines up exactly across the two Zarr stores. + pangu = _open(_PANGU)[[_VAR]].sel(time=_INIT) + graphcast = _open(_GRAPHCAST)[[_VAR]].sel(time=_INIT) + forecasts = xr.concat([pangu, graphcast], dim="model").assign_coords( + model=["pangu", "graphcast"], + latitude=era5.latitude.values, + longitude=era5.longitude.values, + ) + + # ERA5 truth must span every valid time (last init + longest lead); bound it + # lazily so the JOIN does not scan the whole 1959–2023 record. + valid_max = ( + pangu.time.values.max() + pangu.prediction_timedelta.values.max() + ) + truth = era5[[_VAR]].sel(time=slice(_INIT.start, pd.Timestamp(valid_max))) + + print( + f" 64×32 2m_temperature | init {_INIT.start}…{_INIT.stop} " + f"({pangu.sizes['time']} inits × {pangu.sizes['prediction_timedelta']} " + f"leads × 2 models)" + ) + + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + # chunks here is the Arrow batch (partition) size each table streams in, not a + # filter — no data is dropped. Both windows are small, so one partition each is + # fastest (fewer partitions = fewer Python→Arrow round-trips for the same + # rows); time:100 covers both the ~40 forecast inits and the ~79 truth steps. + # Empirically the truth chunk is what matters — splitting it small costs ~3×, + # while the forecasts chunk is in the noise — and a chunk *mismatch* costs + # nothing, so there is no need to keep them different. + ctx.from_dataset("forecasts", forecasts, chunks={"time": 100}) + ctx.from_dataset("era5", truth, chunks={"time": 100}) + + sql = """ + SELECT f.model, + f.prediction_timedelta AS lead, + SQRT(AVG(POWER( + CAST(f."2m_temperature" AS DOUBLE) - e."2m_temperature", 2 + ))) AS rmse + FROM forecasts f + JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude + GROUP BY f.model, f.prediction_timedelta + ORDER BY f.model, lead + """ + show_sql(sql) + + for _ in measured("SQL RMSE by (model, lead) — lazy JOIN"): + got = ctx.sql_to_dataset(sql, dims=["model", "lead"]).rmse + + for _ in measured("xarray reference"): + ref = _reference_rmse(forecasts, truth) + + assert_grid_close("RMSE(model, lead)", got, ref, rtol=1e-4, atol=1e-3) + + show_result(got) + + # Headline: error growth with forecast horizon, both models. The gridded SQL + # result round-trips to a pandas table directly — index is lead (in days), + # one column per model. + table = ( + got.assign_coords(lead=got["lead"].values / np.timedelta64(1, "D")) + .to_pandas() + .T + ) + table.index.name = "lead (days)" + print("\n 2m-temperature RMSE (K) vs lead — lower is better:\n") + print(table.iloc[::4].round(3).to_string()) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Forecast skill: Pangu vs GraphCast vs ERA5 (WB2)") + ) diff --git a/benchmarks/geospatial/06_zonal_vector.py b/benchmarks/geospatial/06_zonal_vector.py new file mode 100644 index 00000000..0ccc7f12 --- /dev/null +++ b/benchmarks/geospatial/06_zonal_vector.py @@ -0,0 +1,177 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Zonal statistics over regions — "rasterize the polygons, then mask" is a JOIN. + +"What is the average temperature inside each region?" is the canonical +*raster × vector* operation. The array paradigm rasterizes each region to a +mask and reduces the raster under it, one region at a time. But a region is +just a row in a table of bounds, and "pixel falls inside region" is a **range +predicate** — so zonal statistics is a JOIN between the raster table and the +regions table, plus a GROUP BY:: + + SELECT r.region, AVG(a."2m_temperature") - 273.15 AS avg_c + FROM era5.surface a JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max + GROUP BY r.region + +This is exactly the README's promise — *joining tabular data with raster data* — +made concrete: the raster is the full **ARCO-ERA5** archive (``WHERE time …`` +prunes it to one day), the regions are a second SQL table, and the spatial +relationship is an ordinary ``BETWEEN``. + +Dataset: the full ARCO-ERA5 archive opened *lazily* — the table spans the whole +record, but the query aggregates only one day's window (the ``WHERE`` prunes the +read; it is not a scan of the full archive) — plus a handful of continental-scale +bounding boxes (longitudes in ERA5's 0–360°E convention). +""" + +from __future__ import annotations + +import datetime + +import numpy as np +import xarray as xr + +from _engines import EngineContext +from _harness import ( + CaseSkipped, + assert_grid_close, + measured, + run_case, + show_result, + show_sql, + timed, +) + +_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" +_DAY = "2020-06-01" +_START, _END = ( + datetime.datetime(2020, 6, 1, 0), + datetime.datetime(2020, 6, 1, 23), +) + +# Continental-scale boxes (name, lat_min, lat_max, lon_min, lon_max), lon 0–360°E. +_REGIONS = [ + ("Sahara", 18.0, 30.0, 0.0, 30.0), + ("Amazon", -10.0, 5.0, 290.0, 310.0), + ("Australia_Outback", -30.0, -20.0, 125.0, 140.0), + ("Greenland", 65.0, 80.0, 300.0, 340.0), + ("SE_Asia", 5.0, 20.0, 95.0, 110.0), +] + + +def _regions_dataset() -> xr.Dataset: + """A vector layer as an xarray Dataset: one row per region, bounds as vars.""" + bounds = np.array([r[1:] for r in _REGIONS], dtype="float64") + return xr.Dataset( + { + "lat_min": (["region"], bounds[:, 0]), + "lat_max": (["region"], bounds[:, 1]), + "lon_min": (["region"], bounds[:, 2]), + "lon_max": (["region"], bounds[:, 3]), + }, + coords={"region": np.arange(len(_REGIONS))}, + ).chunk({"region": len(_REGIONS)}) + + +def main() -> None: + try: + import gcsfs # noqa: F401 — required by the gs:// protocol + + ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"}) + except Exception as exc: # noqa: BLE001 — any failure → skip, not crash + raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc + + print( + f" raster: full ARCO-ERA5 ({ds.sizes['time']:,} timesteps, " + f"{ds.sizes['latitude']}×{ds.sizes['longitude']}) " + f"vector: {len(_REGIONS)} continental boxes" + ) + + ctx = EngineContext() + print(f" engine: {ctx.flavor}") + with timed("register full ERA5 + regions"): + ctx.from_dataset( + "era5", + ds, + chunks={"time": 6}, + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + ctx.from_dataset( + "regions", _regions_dataset(), chunks={"region": len(_REGIONS)} + ) + + sql = """ + SELECT r.region AS region_id, + AVG(a."2m_temperature") - 273.15 AS avg_c, + COUNT(*) AS n_obs + FROM era5.surface a + JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max + WHERE a.time BETWEEN $start AND $end + GROUP BY r.region + ORDER BY r.region + """ + show_sql(sql) + + for _ in measured("SQL zonal stats (raster × vector range JOIN)"): + got = ctx.sql_to_dataset( + sql, + dims=["region_id"], + param_values={"start": _START, "end": _END}, + ) + + # Array reference: one lazy pass — stack the region masks and reduce. No + # .load(): the day's field is read inside this timed block (exactly like the + # SQL side), and reading it once for all regions is a single masked reduction + # rather than a Python loop that would re-read the field per region. + for _ in measured("xarray reference"): + day = xr.open_zarr( + _URL, chunks=None, storage_options={"token": "anon"} + )["2m_temperature"].sel(time=_DAY) + in_region = xr.concat( + [ + (day.latitude >= lat_min) + & (day.latitude <= lat_max) + & (day.longitude >= lon_min) + & (day.longitude <= lon_max) + for _, lat_min, lat_max, lon_min, lon_max in _REGIONS + ], + dim="region_id", + ) + ref = ( + day.where(in_region).mean(["time", "latitude", "longitude"]) + - 273.15 + ).assign_coords(region_id=got.region_id) + + assert_grid_close( + "zonal mean per region (°C)", got.avg_c, ref, rtol=1e-4, atol=1e-2 + ) + + show_result(got) + + print("\n Region avg °C n_obs") + for (name, *_), avg, n in zip(_REGIONS, got.avg_c.values, got.n_obs.values): + print(f" {name:<20} {avg:7.2f} {int(n):>10,}") + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Zonal stats: raster × vector range JOIN (ARCO-ERA5)") + ) diff --git a/benchmarks/geospatial/07_reproject_udf.py b/benchmarks/geospatial/07_reproject_udf.py new file mode 100644 index 00000000..29a2576e --- /dev/null +++ b/benchmarks/geospatial/07_reproject_udf.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[geo]", +# "xarray", +# "numpy", +# "pyarrow", +# "xee", +# "earthengine-api", +# "shapely", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Reprojection — a per-pixel CRS transform is a scalar UDF (à la ST_Transform). + +Reprojection moves coordinates from one CRS to another (here UTM zone 10N, +EPSG:32610, → lon/lat, EPSG:4326). Crucially it is **row-independent**: each +pixel's new coordinate depends only on its own old coordinate. That is exactly +the shape of a SQL *scalar UDF*, and it is precisely how the geospatial SQL +world already does it — PostGIS ``ST_Transform`` and DuckDB-spatial +``ST_Transform`` are scalar PROJ wrappers. + +xarray-sql ships that UDF as its geo extension (``xarray_sql.proj``): +with pyproj installed, every ``XarrayContext`` speaks CRS out of the box, +and the CRS pair is part of the query rather than baked into the UDF:: + + SELECT x, y, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat + FROM grid + +**The reference is Earth Engine itself.** There is *one* dataset: a single UTM +grid opened through [Xee](https://github.com/google/Xee) carrying +``ee.Image.pixelLonLat()``. Each pixel arrives with two things — its UTM ``x``/ +``y`` (the grid coordinates, our SQL input) and Earth Engine's *own* per-pixel +``longitude``/``latitude`` (data variables, the reference). So we are not +opening the same image twice in two CRS; we feed the UTM coordinates to the PROJ +UDF and check the lon/lat it returns against EE's independently-computed lon/lat +for the *same* pixels. The reference is a different geodesy engine, not PROJ +again, and they agree to sub-metre precision. + +The extension returns *both* coordinates from one struct-returning call +(one PROJ transform per row) and runs all PROJ work on its own worker +pool, so the query parallelizes across partitions safely. + +Requires Earth Engine access: ``earthengine authenticate`` once, then an +initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise. +""" + +from __future__ import annotations + +import xarray as xr + +import xarray_sql as xql + +from _harness import ( + CaseSkipped, + assert_grid_close, + initialize_earth_engine, + measured, + run_case, + show_result, + show_sql, +) + +_SRC_CRS, _DST_CRS = "EPSG:32610", "EPSG:4326" # UTM zone 10N → lon/lat +# A 1° box over the San Francisco Bay area, well inside UTM zone 10N. +_AOI = (-122.6, 37.4, -121.6, 38.4) +_SCALE_M = 2_000 # 2 km pixels → a ~50×60 grid + + +def _open_ee_lonlat_grid() -> xr.Dataset: + """Open ``ee.Image.pixelLonLat()`` on a UTM grid via Xee. + + Earth Engine evaluates ``pixelLonLat`` on the requested UTM grid, so each + pixel carries its UTM ``x``/``y`` (coordinates) and EE's own ``longitude`` / + ``latitude`` (data variables) — the independent reprojection reference. + """ + try: + import shapely.geometry as sgeom + from xee import helpers + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api xee`" + ) from exc + + ee = initialize_earth_engine() + + # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's + # backend expects — here a UTM grid at _SCALE_M metres covering the AOI. + grid = helpers.fit_geometry( + sgeom.box(*_AOI), + geometry_crs="EPSG:4326", + grid_crs=_SRC_CRS, + grid_scale=(float(_SCALE_M), float(_SCALE_M)), + ) + ic = ee.ImageCollection([ee.Image.pixelLonLat()]) + ds = xr.open_dataset(ic, engine="ee", **grid) + # One image → a length-1 time axis; drop it. Xee gives x/y coordinates (UTM + # metres) and longitude/latitude data variables (EE's per-pixel geodesy). + return ds.isel(time=0).load() + + +def main() -> None: + ds = _open_ee_lonlat_grid() + n = ds.sizes["y"] * ds.sizes["x"] + print( + f" EE pixelLonLat on UTM grid {dict(ds.sizes)} ({n:,} pixels) " + f"{_SRC_CRS} → {_DST_CRS}" + ) + + # XarrayContext registers reproject() automatically (the geo + # extension). The chunking deliberately splits the ~60-row grid into + # 15-row chunks → 4 partitions, forcing DataFusion to evaluate the UDF + # concurrently: the extension runs PROJ on its own worker pool, so + # parallel partitions are safe (previously this required one chunk → + # one partition → a serial UDF). + ctx = xql.XarrayContext() + ctx.from_dataset("grid", ds, chunks={"y": 15, "x": ds.sizes["x"]}) + + sql = f""" + SELECT x, y, + reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['x'] AS lon, + reproject(x, y, '{_SRC_CRS}', '{_DST_CRS}')['y'] AS lat + FROM grid + ORDER BY y, x + """ + show_sql(sql) + + for _ in measured("SQL reprojection (PROJ scalar UDF)"): + got = ctx.sql(sql).to_dataset(dims=["y", "x"]) + + # Reference: Earth Engine's own per-pixel lon/lat (independent of PROJ). + # EE and PROJ are separate implementations, so compare at ~1e-5° (~1 m). + assert_grid_close( + "reprojected longitude", got.lon, ds.longitude, rtol=0, atol=1e-5 + ) + assert_grid_close( + "reprojected latitude", got.lat, ds.latitude, rtol=0, atol=1e-5 + ) + + show_result(got) + + corner = got.isel(x=0, y=0) + print( + f"\n Corner check: UTM ({float(corner.x):.0f}, {float(corner.y):.0f}) → " + f"lon {float(corner.lon):.4f}, lat {float(corner.lat):.4f}" + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Reprojection: PROJ scalar UDF vs Earth Engine") + ) diff --git a/benchmarks/geospatial/08_regrid_weights.py b/benchmarks/geospatial/08_regrid_weights.py new file mode 100644 index 00000000..bf0ecd9a --- /dev/null +++ b/benchmarks/geospatial/08_regrid_weights.py @@ -0,0 +1,228 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql", +# "xarray", +# "numpy", +# "scipy", +# "xee", +# "earthengine-api", +# "shapely", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Regridding — interpolation to a new grid is a sparse matmul, i.e. a JOIN. + +Regridding (resampling a field from one grid onto another) is the operation we +most associate with the *array* paradigm — xESMF/ESMF, ``apply_ufunc``, +``.interp()``. But every linear regridding scheme (bilinear, conservative, +nearest) is mathematically a **sparse matrix–vector product**: each output cell +is a weighted sum of a few input cells. And a sparse matrix is just a table of +``(target, source, weight)`` rows. So *applying* a regridding is:: + + SELECT w.dst_lat, w.dst_lon, SUM(s.value * w.weight) AS regridded + FROM weights w JOIN src s ON s.lat = w.src_lat AND s.lon = w.src_lon + GROUP BY w.dst_lat, w.dst_lon + +— a JOIN against the weight table plus a weighted GROUP BY. This is the most +relational the "array" paradigm ever gets: the operation we reach for xESMF to +do is a join. + +**Where the array paradigm still earns its keep:** *generating* the weights is +the genuinely geometric part (cell overlaps, interpolation stencils, spherical +coordinates). Here we build bilinear weights with a few lines of numpy; for +conservative remapping on real grids you would let ESMF/xESMF compute them once +and hand the resulting sparse matrix to SQL as a table. SQL *applies* the +weights; it does not invent the geometry. + +The field is real **SRTM elevation** (terrain over the Sierra Nevada), opened +from the Earth Engine catalog through [Xee](https://github.com/google/Xee). We +regrid it coarse → fine in SQL and validate against xarray's own bilinear +``.interp()`` on the same source field. + +Requires Earth Engine access: ``earthengine authenticate`` once, then an +initialized project (set ``EARTHENGINE_PROJECT``). Skips cleanly otherwise. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +import xarray_sql as xql + +from _harness import ( + CaseSkipped, + assert_grid_close, + initialize_earth_engine, + measured, + run_case, + show_result, + show_sql, + timed, +) + +# A 1° box over the Sierra Nevada — real terrain with strong relief. +_AOI = (-119.6, 37.0, -118.6, 38.0) +_SRC_SCALE_DEG = 0.02 # ~2 km source pixels (a coarse DEM to upsample) + + +def _linear_weights( + src: np.ndarray, dst: np.ndarray +) -> list[tuple[int, int, float]]: + """1-D linear-interpolation weights: (dst_index, src_index, weight) triples. + + Each target point falls between two source points and borrows from both, + with weights summing to 1 — the 1-D building block of bilinear regridding. + """ + triples = [] + for t, x in enumerate(dst): + i = int(np.clip(np.searchsorted(src, x) - 1, 0, len(src) - 2)) + span = src[i + 1] - src[i] + hi = (x - src[i]) / span + triples.append((t, i, 1.0 - hi)) + triples.append((t, i + 1, hi)) + return triples + + +def _bilinear_weight_table( + slat: np.ndarray, slon: np.ndarray, tlat: np.ndarray, tlon: np.ndarray +) -> xr.Dataset: + """Build the sparse bilinear weight matrix as a weight table. + + The 2-D weight is the outer product of the 1-D lat and lon weights. Each + nonzero is one row naming the target cell by its ``(dst_lat, dst_lon)`` and + the source cell by its ``(src_lat, src_lon)`` — so the regrid SQL joins the + source grid on its coordinates (no pre-raveled cell id), lets the engine read + the source lazily, and rounds the result straight back to a (lat, lon) grid. + """ + lat_w = _linear_weights(slat, tlat) + lon_w = _linear_weights(slon, tlon) + dst_lats, dst_lons, src_lats, src_lons, weights = [], [], [], [], [] + for tj, si, wlat in lat_w: + for tk, sj, wlon in lon_w: + dst_lats.append(tlat[tj]) + dst_lons.append(tlon[tk]) + src_lats.append(slat[si]) + src_lons.append(slon[sj]) + weights.append(wlat * wlon) + n = len(weights) + return xr.Dataset( + { + "dst_lat": (["pair"], np.array(dst_lats, dtype="float64")), + "dst_lon": (["pair"], np.array(dst_lons, dtype="float64")), + "src_lat": (["pair"], np.array(src_lats, dtype="float64")), + "src_lon": (["pair"], np.array(src_lons, dtype="float64")), + "weight": (["pair"], np.array(weights, dtype="float64")), + }, + coords={"pair": np.arange(n)}, + ).chunk({"pair": n}) + + +def _open_srtm() -> xr.DataArray: + """Open SRTM elevation over the AOI as a coarse (lat, lon) field via Xee.""" + try: + import shapely.geometry as sgeom + from xee import helpers + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api xee`" + ) from exc + + ee = initialize_earth_engine() + + # fit_geometry builds the pixel grid (crs, crs_transform, shape_2d) Xee's + # backend expects — here a geographic grid at _SRC_SCALE_DEG° over the AOI. + grid = helpers.fit_geometry( + sgeom.box(*_AOI), + grid_crs="EPSG:4326", + grid_scale=(_SRC_SCALE_DEG, _SRC_SCALE_DEG), + ) + ic = ee.ImageCollection([ee.Image("USGS/SRTMGL1_003")]) # band: elevation + ds = xr.open_dataset(ic, engine="ee", **grid) + da = ds["elevation"].isel(time=0) + # Normalize Xee's spatial coordinate names to lat/lon and sort ascending so + # the 1-D weight construction (searchsorted) sees increasing coordinates. + rename = {} + for d in da.dims: + dl = d.lower() + if dl in ("y", "lat", "latitude"): + rename[d] = "lat" + elif dl in ("x", "lon", "longitude"): + rename[d] = "lon" + da = da.rename(rename).sortby("lat").sortby("lon") + # Stay lazy (no .load()): the source is read on demand by both the SQL table + # and the .interp reference, so each pays its own read. Force float64 coords + # so the weight table's src lat/lon match the source grid's exactly in the + # join. + return da.assign_coords( + lat=da.lat.astype("float64"), lon=da.lon.astype("float64") + ) + + +def main() -> None: + with timed("open SRTM via Xee (lazy)"): + src_da = _open_srtm() + slat = src_da.lat.values + slon = src_da.lon.values + print(f" SRTM elevation source grid {len(slat)}×{len(slon)} (read lazily)") + + # Finer target grid strictly inside the source extent (bilinear upsampling). + tlat = np.linspace(slat[1], slat[-2], 60) + tlon = np.linspace(slon[1], slon[-2], 72) + print( + f" regrid {len(slat)}×{len(slon)} → {len(tlat)}×{len(tlon)} (bilinear)" + ) + + weights = _bilinear_weight_table(slat, slon, tlat, tlon) + print( + f" weight matrix: {weights.sizes['pair']:,} nonzeros " + f"({len(tlat) * len(tlon)} targets × 4 corners)" + ) + + ctx = xql.XarrayContext() + # Register the source grid itself (lazy) — the join reads it on demand, the + # same source the .interp reference reads, so both pay an equal lazy read. + ctx.from_dataset( + "src", + src_da.to_dataset(name="value"), + chunks={"lat": len(slat), "lon": len(slon)}, + ) + ctx.from_dataset("weights", weights, chunks={"pair": weights.sizes["pair"]}) + + sql = """ + SELECT w.dst_lat AS lat, + w.dst_lon AS lon, + SUM(s.value * w.weight) AS regridded + FROM weights w + JOIN src s ON s.lat = w.src_lat AND s.lon = w.src_lon + GROUP BY w.dst_lat, w.dst_lon + ORDER BY w.dst_lat, w.dst_lon + """ + show_sql(sql) + + # The weights name each target cell by its (lat, lon), so the result rounds + # straight back to the (lat, lon) field it represents — no reshape. + for _ in measured("SQL regrid (weight-table JOIN + weighted SUM)"): + got = ctx.sql(sql).to_dataset(dims=["lat", "lon"]).regridded + + # Array reference: xarray's own bilinear interpolation of the same lazy field. + for _ in measured("xarray .interp reference"): + ref = src_da.interp(lat=tlat, lon=tlon, method="linear") + + assert_grid_close("bilinear regrid", got, ref, rtol=1e-9, atol=1e-9) + + show_result(got) + + print( + f"\n {got.size:,} target cells regridded; " + f"elevation range [{float(got.min()):.0f}, {float(got.max()):.0f}] m." + ) + + +if __name__ == "__main__": + raise SystemExit( + run_case(main, "Regridding: sparse weight-table JOIN (SRTM)") + ) diff --git a/benchmarks/geospatial/09_warp.py b/benchmarks/geospatial/09_warp.py new file mode 100644 index 00000000..09e051c7 --- /dev/null +++ b/benchmarks/geospatial/09_warp.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[geo]", +# "xarray", +# "numpy", +# "pyarrow", +# "scipy", +# "xee", +# "earthengine-api", +# "shapely", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Warp — reprojecting *and* resampling a raster is case 07's UDF + case 08's JOIN. + +A *warp* moves a raster from one CRS onto a grid in another — the everyday GDAL/ +rasterio ``reproject`` that GIS runs constantly. It is exactly the composition of +the two "hard" cases: + +* **case 07** — reproject coordinates with a scalar PROJ **UDF**; and +* **case 08** — resample values with a sparse-weight **JOIN**. + +The pipeline reads as that composition, and it shows the division of labor cleanly: + +1. **SQL reprojects the target grid** (the 07 UDF): for every target ``(lon, lat)`` + cell, ``reproject()`` returns where it falls in the source CRS (UTM ``x``/``y``). +2. **Arrays build the bilinear weights** (the geometry): each reprojected target + point lands between four source pixels; we compute its four bilinear weights — + the genuinely geometric step the array world owns. (This is the same + "arrays compute the weights, SQL applies them" boundary as case 08, except the + target points are *scattered* in source space because they were reprojected, + so the weights are a per-point stencil rather than a separable lat×lon grid.) +3. **SQL applies the weights** (the 08 JOIN): join the source values onto the + weight table and ``SUM(value * weight)`` per target cell. + +**References.** The exact check is the array paradigm doing the same warp — plain +``pyproj`` + xarray ``.interp`` at the reprojected points — which the SQL result +matches to floating-point tolerance. As an *independent* real-world cross-check we +also open the **same SRTM** directly on the lon/lat grid through Xee (Earth +Engine's own warp) and report the agreement; it is close (a few metres median) but +not bit-exact, because EE resamples from native 30 m while our source is the coarse +UTM grid — which is exactly why the deterministic warp, not EE, is the tolerance +reference. + +Data: real **SRTM elevation** (Northern California terrain) via [Xee](https://github.com/google/Xee), +opened once on a UTM grid (the source) and once on a lon/lat grid (the EE +cross-check). Requires Earth Engine access; skips cleanly otherwise. +""" + +from __future__ import annotations + +import numpy as np +import pyproj +import shapely.geometry as sgeom +import xarray as xr + +import xarray_sql as xql + +from _harness import ( + CaseSkipped, + assert_grid_close, + initialize_earth_engine, + measured, + run_case, + show_result, + show_sql, + timed, +) + +_SRC_CRS = "EPSG:32610" # UTM zone 10N — the source raster's CRS +_DST_CRS = "EPSG:4326" # lon/lat — the target grid's CRS +_AOI = (-122.6, 37.4, -121.6, 38.4) # ~1° box of Northern California terrain +_SRC_SCALE_M = 2_000.0 # ~2 km source pixels +_DST_SCALE_DEG = 0.02 # ~2 km target cells + + +def _open_srtm( + grid_crs: str, scale: tuple[float, float], xy_names +) -> xr.DataArray: + """Open SRTM elevation over the AOI on the requested grid via Xee (lazy).""" + try: + from xee import helpers + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api xee`" + ) from exc + + ee = initialize_earth_engine() + grid = helpers.fit_geometry( + sgeom.box(*_AOI), + geometry_crs="EPSG:4326", + grid_crs=grid_crs, + grid_scale=scale, + ) + ic = ee.ImageCollection([ee.Image("USGS/SRTMGL1_003")]) + da = xr.open_dataset(ic, engine="ee", **grid)["elevation"].isel(time=0) + a, b = xy_names + rename = {} + for d in da.dims: + dl = d.lower() + if dl in ("y", "lat", "latitude"): + rename[d] = a + elif dl in ("x", "lon", "longitude"): + rename[d] = b + da = da.rename(rename).sortby(a).sortby(b) + return da.assign_coords( + {a: da[a].astype("float64"), b: da[b].astype("float64")} + ) + + +def _warp_weight_table( + sx: np.ndarray, + sy: np.ndarray, + dst_lon: np.ndarray, + dst_lat: np.ndarray, + px: np.ndarray, + py: np.ndarray, +) -> xr.Dataset: + """Bilinear weights for reprojected target points — the geometry step. + + Each target cell ``(dst_lat, dst_lon)`` was reprojected to source coordinates + ``(px, py)``; here we find the four surrounding source pixels and their + bilinear weights. One row per (target cell, source corner). Targets that fall + outside the source footprint contribute no rows (and are dropped). + """ + dst_lats, dst_lons, src_xs, src_ys, weights = [], [], [], [], [] + for k in range(len(px)): + x, y = px[k], py[k] + if not (sx[0] <= x <= sx[-1] and sy[0] <= y <= sy[-1]): + continue + i = int(np.clip(np.searchsorted(sx, x) - 1, 0, len(sx) - 2)) + j = int(np.clip(np.searchsorted(sy, y) - 1, 0, len(sy) - 2)) + tx = (x - sx[i]) / (sx[i + 1] - sx[i]) + ty = (y - sy[j]) / (sy[j + 1] - sy[j]) + for ii, wx in ((i, 1.0 - tx), (i + 1, tx)): + for jj, wy in ((j, 1.0 - ty), (j + 1, ty)): + dst_lats.append(dst_lat[k]) + dst_lons.append(dst_lon[k]) + src_xs.append(sx[ii]) + src_ys.append(sy[jj]) + weights.append(wx * wy) + n = len(weights) + return xr.Dataset( + { + "dst_lat": (["pair"], np.array(dst_lats, "float64")), + "dst_lon": (["pair"], np.array(dst_lons, "float64")), + "src_x": (["pair"], np.array(src_xs, "float64")), + "src_y": (["pair"], np.array(src_ys, "float64")), + "weight": (["pair"], np.array(weights, "float64")), + }, + coords={"pair": np.arange(n)}, + ).chunk({"pair": n}) + + +def main() -> None: + with timed("open SRTM on UTM + lon/lat grids via Xee"): + src = _open_srtm(_SRC_CRS, (_SRC_SCALE_M, _SRC_SCALE_M), ("y", "x")) + ref_ee = _open_srtm( + _DST_CRS, (_DST_SCALE_DEG, _DST_SCALE_DEG), ("lat", "lon") + ) + sx, sy = src.x.values, src.y.values + + # Target lon/lat grid strictly inside the source UTM footprint, so every + # target cell reprojects to a point with four source corners (no edge cells + # to drop). Inscribe a lon/lat box in the reprojected UTM rectangle. + inv = pyproj.Transformer.from_crs(_SRC_CRS, _DST_CRS, always_xy=True) + cx = [sx[0], sx[-1], sx[0], sx[-1]] + cy = [sy[0], sy[0], sy[-1], sy[-1]] + clon, clat = inv.transform(cx, cy) + lon0, lon1 = max(clon[0], clon[2]) + 0.01, min(clon[1], clon[3]) - 0.01 + lat0, lat1 = max(clat[0], clat[1]) + 0.01, min(clat[2], clat[3]) - 0.01 + tlon = np.linspace(lon0, lon1, 60) + tlat = np.linspace(lat0, lat1, 60) + print( + f" source UTM grid {len(sy)}×{len(sx)} → target lon/lat grid " + f"{len(tlat)}×{len(tlon)} ({_SRC_CRS} → {_DST_CRS})" + ) + + # XarrayContext registers reproject() automatically (the geo + # extension) — the direction is spelled in the query itself. + ctx = xql.XarrayContext() + + # The target grid as a (dst_lat, dst_lon) table. + LON, LAT = np.meshgrid(tlon, tlat) + target = xr.Dataset( + { + "dst_lon": (["cell"], LON.ravel()), + "dst_lat": (["cell"], LAT.ravel()), + }, + coords={"cell": np.arange(LON.size)}, + ).chunk({"cell": LON.size}) + ctx.from_dataset("target", target, chunks={"cell": LON.size}) + + # 1) SQL reprojects the target grid into the source CRS (case 07's UDF). + reproj_sql = f""" + SELECT dst_lat, dst_lon, + reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['x'] AS sx, + reproject(dst_lon, dst_lat, '{_DST_CRS}', '{_SRC_CRS}')['y'] AS sy + FROM target + """ + show_sql(reproj_sql, label="SQL — reproject target grid (PROJ UDF)") + rp = ctx.sql(reproj_sql).to_pandas() + px, py = rp["sx"].to_numpy(), rp["sy"].to_numpy() + + # 2) Arrays turn the reprojected points into a bilinear weight table. + weights = _warp_weight_table( + sx, sy, rp["dst_lon"].to_numpy(), rp["dst_lat"].to_numpy(), px, py + ) + ctx.from_dataset( + "src", src.to_dataset(name="value"), chunks={"y": len(sy), "x": len(sx)} + ) + ctx.from_dataset("weights", weights, chunks={"pair": weights.sizes["pair"]}) + + # 3) SQL applies the weights (case 08's JOIN). + apply_sql = """ + SELECT w.dst_lat AS lat, w.dst_lon AS lon, + SUM(s.value * w.weight) AS warped + FROM weights w + JOIN src s ON s.x = w.src_x AND s.y = w.src_y + GROUP BY w.dst_lat, w.dst_lon + ORDER BY w.dst_lat, w.dst_lon + """ + show_sql(apply_sql, label="SQL — apply bilinear weights (JOIN)") + for _ in measured("SQL warp (reproject UDF + regrid JOIN)"): + got = ctx.sql(apply_sql).to_dataset(dims=["lat", "lon"]).warped + + # Reference: the array paradigm doing the same warp — pyproj reproject of the + # target grid, then xarray's own bilinear .interp at those source points. + for _ in measured("xarray reference (pyproj + .interp)"): + tr = pyproj.Transformer.from_crs(_DST_CRS, _SRC_CRS, always_xy=True) + rx, ry = tr.transform(LON.ravel(), LAT.ravel()) + warped = src.interp( + x=xr.DataArray(rx, dims="cell"), + y=xr.DataArray(ry, dims="cell"), + method="linear", + ).values.reshape(len(tlat), len(tlon)) + ref = xr.DataArray( + warped, dims=["lat", "lon"], coords={"lat": tlat, "lon": tlon} + ) + + assert_grid_close("warped elevation (m)", got, ref, rtol=1e-6, atol=1e-4) + show_result(got) + + # Independent cross-check: EE's own SRTM on the lon/lat grid (a real warp). + ee_on_grid = ref_ee.interp(lat=got.lat, lon=got.lon, method="linear").values + a, b = got.values.ravel(), ee_on_grid.ravel() + m = np.isfinite(a) & np.isfinite(b) + corr = float(np.corrcoef(a[m], b[m])[0, 1]) + print( + f"\n vs Earth Engine's own lon/lat SRTM: median |Δ| " + f"{np.nanmedian(np.abs(a[m] - b[m])):.1f} m, correlation {corr:.4f} " + f"(EE resamples native 30 m; ours warps the {_SRC_SCALE_M:.0f} m UTM grid)" + ) + + +if __name__ == "__main__": + raise SystemExit(run_case(main, "Warp: reproject UDF + regrid JOIN (SRTM)")) diff --git a/benchmarks/geospatial/README.md b/benchmarks/geospatial/README.md new file mode 100644 index 00000000..c046adc3 --- /dev/null +++ b/benchmarks/geospatial/README.md @@ -0,0 +1,106 @@ +# Geospatial SQL benchmarks + +**Thesis:** the core geospatial operations we assume require an *array* paradigm +are, underneath, **relational** operations — `GROUP BY`, `JOIN`, window +functions, and `CASE`. Each script here takes one such operation, expresses it +in SQL against [`xarray-sql`](../../README.md), and **proves the SQL answer +matches a plain-xarray reference** to floating-point tolerance. Wall-clock and +peak memory are reported too, but the headline is correctness + clarity of the +SQL. + +This suite is *expressibility-first*: the point is that the SQL reads like the +plain-English definition of the operation, and computes the same numbers. + +## The cases + +| # | Case | Array mental model | Relational reality | +|---|------|--------------------|--------------------| +| 01 | `01_ndvi.py` | `apply_ufunc` over a raster | column arithmetic | +| 02 | `02_climatology.py` | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | +| 03 | `03_zonal_mean.py` | reduce over lon/time axes | `GROUP BY latitude` | +| 04 | `04_anomaly.py` | climatology broadcast-subtract | climatology CTE self-`JOIN` | +| 05 | `05_forecast_skill.py` | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` + aggregate | +| 06 | `06_zonal_vector.py` | rasterize + mask per region | range `JOIN` raster↔regions | +| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()` from the geo extension), à la PostGIS `ST_Transform` | +| 08 | `08_regrid_weights.py` | interpolation to a new grid | sparse-weight table `JOIN` + weighted `GROUP BY` | +| 09 | `09_warp.py` | reproject **and** resample (warp) | reproject **UDF** (07) → weight table `JOIN` (08) | + +Cases 01–06 show operations that are *natively* relational. Cases 07–08 are the +"hardest" array operations — reprojection and regridding — and show where a UDF +fits (a per-row coordinate transform) versus where the operation is really a +sparse matrix multiply expressed as a `JOIN`. Case 09 composes the two into a full +**warp** (GDAL/rasterio `reproject`): the 07 UDF reprojects the target grid, arrays +turn the reprojected points into bilinear weights, and the 08 `JOIN` applies them. +See +[`docs/geospatial.md`](../../docs/geospatial.md) for the full narrative, +including *where the array paradigm still earns its keep* (generating the +interpolation weights — the geometry — which SQL applies but does not compute). + +## Datasets + +- **01 NDVI** — a real Sentinel-2 L2A scene in **Zarr** from the ESA EOPF sample + service, discovered with `pystac-client` and opened with `xr.open_datatree` + (bands B04/B08). Requires network; skips cleanly if offline. +- **02–06** — the full **[ARCO-ERA5](https://github.com/google-research/arco-era5)** + archive (0.25° global, ~1.3M hourly timesteps, 273 variables) read anonymously + from a public GCS bucket. Each case opens the *whole* archive lazily, so a query + reads only the variable and the window it asks for — never the other 272 + variables or the rest of the timesteps. All require network (`gcsfs`) and skip + cleanly offline; each takes roughly one to a few minutes, dominated by the read. +- **05 forecast skill** — the **[WeatherBench 2](https://weatherbench2.readthedocs.io/)** + Pangu-Weather, GraphCast, and ERA5 datasets at a coarse 64×32 grid, scoring + both ML models against ERA5 ground truth. Network-backed; runs in seconds + because the grid is small. +- **07–09** — the **Earth Engine** catalog via [Xee](https://github.com/google/Xee). + 07 reprojects a UTM grid and validates the SQL transform against Earth Engine's + *own* per-pixel lon/lat (`ee.Image.pixelLonLat()`) — an independent reprojection + reference, not PROJ-vs-PROJ. 08 regrids real **SRTM elevation** (Sierra Nevada) + and validates against xarray's bilinear `.interp()`. 09 warps SRTM from a UTM + grid onto a lon/lat grid (07's reproject UDF feeding 08's weight `JOIN`) and + validates against xarray's `.interp()` at the reprojected points, with Earth + Engine's own lon/lat SRTM as a second, cross-CRS check. All three run against + Earth Engine using your existing `gcloud` login, and skip cleanly without it. + +## Running + +Run a single case, or the whole suite, from any directory: + +```shell +uv run benchmarks/geospatial/03_zonal_mean.py # one case +benchmarks/geospatial/run_all.sh # all of them +``` + +Each script carries [PEP 723 / `uv` inline metadata](https://docs.astral.sh/uv/guides/scripts/) +and runs against the `xarray-sql` in this checkout. + +A passing case prints a `✅ … SQL matches xarray reference` line and the result +as an xarray repr; a mismatch raises `AssertionError` and exits non-zero. Cases +that need data or credentials you don't have print `⏭ SKIPPED` and exit 0. + +Shared helpers — timing, peak memory, the result check and its printout, SQL +echo — live in [`_harness.py`](_harness.py). + +## Profiling + +For a performance table, use `run_perf.sh`. It runs each case **once per fresh +process**, with no warmup, repeated `GEOBENCH_REPS` times, and aggregates the +runs into one CSV (and a markdown table on stdout): + +```shell +GEOBENCH_REPS=5 benchmarks/geospatial/run_perf.sh perf.csv +``` + +A fresh process per repetition is deliberate, and it's the only way the SQL and +xarray sides compare fairly. `xr.open_zarr(chunks=None)` caches each variable in +memory after its first read, so an in-process warm loop would let the xarray +reference serve later repetitions from RAM while the SQL side re-reads the +store — flattering the reference. One process per rep makes **both sides pay a +cold read every time**. The columns are `case, title, step, reps, t_min_s, +t_median_s, t_mean_s, t_stdev_s, t_max_s, peak_mb`. Run it close to the data (a +VM in the bucket's region) against a release build of `xarray-sql`; pass +`GEOBENCH_PYRUN="python"` to use an already-built venv instead of `uv run`. + +Under the hood each repeatable step is wrapped in `for _ in measured(...)` +(rather than `with timed(...)`); with `GEOBENCH_PROFILE=1` set, `measured` times +the step and, with `GEOBENCH_CSV`, records it. `run_perf.sh` drives that one cold +run at a time; everything else in the cases is the ordinary xarray/SQL. diff --git a/benchmarks/geospatial/_engines.py b/benchmarks/geospatial/_engines.py new file mode 100644 index 00000000..66d769bc --- /dev/null +++ b/benchmarks/geospatial/_engines.py @@ -0,0 +1,291 @@ +"""Engine-portable SQL layer for the geospatial suite. + +``GEOBENCH_ENGINE`` selects which SQL engine executes each case's query, +so the same case scripts (same SQL, same datasets, same correctness +assertions) can be measured across engines: + +* ``datafusion`` (default) — ``xql.XarrayContext`` over the native + DataFusion table provider, the suite's original path. Requires the + compiled ``xarray_sql._native`` module; raises at startup when it is + missing instead of falling back. +* ``datafusion-arrow`` — a plain ``datafusion.SessionContext`` scanning + ``xql.arrow_dataset`` (pure Python). +* ``duckdb`` — DuckDB over the same pyarrow pushdown datasets. +* ``polars`` — ``polars.SQLContext`` over ``scan_pyarrow_dataset`` frames. + +Every case builds one :class:`EngineContext`, registers datasets exactly +as it always registered them on ``XarrayContext``, and calls +:meth:`EngineContext.sql_to_dataset`. On the ``datafusion`` path this +is byte-for-byte the original behavior (``from_dataset`` + ``sql`` + +``XarrayDataFrame.to_dataset``); the other engines register one pyarrow +dataset per dimension group under flattened table names +(``era5.surface`` → ``era5_surface`` — rewritten in the SQL text) and the +result rows are round-tripped to an ``xr.Dataset`` through pandas. + +The DataFusion-only UDF cases (07 and the UDF half of 09) build +``xql.XarrayContext`` directly rather than through this layer; the suite +runner records them as n/a for every engine except ``datafusion``. +""" + +from __future__ import annotations + +import datetime +import os +import re +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr + + +def engine_name() -> str: + """The engine selected for this process (``GEOBENCH_ENGINE``).""" + engine = os.environ.get("GEOBENCH_ENGINE", "datafusion") + if engine not in _ENGINES: + raise ValueError(f"GEOBENCH_ENGINE={engine!r}; expected {_ENGINES}") + return engine + + +def _group_tables(name, ds, table_names): + """Split ``ds`` into per-dimension-group tables like XarrayContext does. + + Returns ``[(flat_name, dotted_name, sub_dataset)]``; a uniform dataset + keeps its plain name (flat == dotted == name). + """ + groups: dict[tuple, list] = {} + for var, v in ds.data_vars.items(): + groups.setdefault(tuple(v.dims), []).append(var) + if len(groups) == 1: + return [(name, name, ds)] + out = [] + for dims, variables in groups.items(): + sub = (table_names or {}).get(dims) or "_".join(dims) + out.append((f"{name}_{sub}", f"{name}.{sub}", ds[variables])) + return out + + +def _literal(value: Any) -> str: + """Render a parameter value as a SQL literal (for engines without binds).""" + if isinstance(value, (datetime.datetime, pd.Timestamp, np.datetime64)): + return ( + f"TIMESTAMP '{pd.Timestamp(value).strftime('%Y-%m-%d %H:%M:%S')}'" + ) + if isinstance(value, str): + escaped = value.replace("'", "''") + return f"'{escaped}'" + return repr(value) + + +def _to_ns(pdf: pd.DataFrame, dims: list[str]) -> pd.DataFrame: + """Normalize datetime/timedelta dim columns to ns for label alignment.""" + for col in dims: + dtype = pdf[col].dtype + if pd.api.types.is_datetime64_any_dtype(dtype): + pdf[col] = pdf[col].astype("datetime64[ns]") + elif pd.api.types.is_timedelta64_dtype(dtype): + pdf[col] = pdf[col].astype("timedelta64[ns]") + return pdf + + +def _pandas_to_dataset(pdf: pd.DataFrame, dims: list[str]) -> xr.Dataset: + """Round-trip a SQL result table to a gridded ``xr.Dataset`` by ``dims``.""" + pdf = _to_ns(pdf.copy(), dims) + return xr.Dataset.from_dataframe(pdf.set_index(dims).sort_index()) + + +class EngineContext: + """Uniform register-and-query facade over the suite's SQL engines. + + ``EngineContext(engine)`` instantiates the subclass ``_IMPLS`` maps + the engine name to (default: :func:`engine_name`). Subclasses set + ``flavor`` and implement three hooks: ``_connect`` (open the + engine's connection/context), ``_register`` (attach one pyarrow + dataset under a flat table name), and ``_execute`` (run SQL, + returning a ``pandas.DataFrame``). Engines that bypass the shared + pyarrow-dataset path override :meth:`from_dataset` / + :meth:`sql_to_dataset` instead. + """ + + flavor = "" + + def __new__(cls, engine: str | None = None): + if cls is EngineContext: + cls = _IMPLS[engine or engine_name()] + return super().__new__(cls) + + def __init__(self, engine: str | None = None): + self.engine = engine or engine_name() + self._renames: dict[str, str] = {} + self._connect() + + def _connect(self) -> None: + raise NotImplementedError + + def _register(self, flat: str, dataset) -> None: + raise NotImplementedError + + def _execute(self, sql: str, param_values) -> pd.DataFrame: + raise NotImplementedError + + # -- registration ----------------------------------------------------- + + def from_dataset(self, name, ds, *, chunks=None, table_names=None): + """Register ``ds`` as SQL table(s), mirroring XarrayContext naming.""" + import xarray_sql as xql + + for flat, dotted, sub in _group_tables(name, ds, table_names): + if dotted != flat: + self._renames[dotted] = flat + sub_chunks = ( + {d: c for d, c in chunks.items() if d in sub.dims} + if isinstance(chunks, dict) + else chunks + ) or None + self._register(flat, xql.arrow_dataset(sub, sub_chunks)) + + # -- querying ---------------------------------------------------------- + + def _rewrite(self, sql: str, param_values) -> str: + for dotted, flat in self._renames.items(): + sql = re.sub(rf"\b{re.escape(dotted)}\b", flat, sql) + return sql + + def sql_to_dataset( + self, sql: str, *, dims: list[str], param_values=None + ) -> xr.Dataset: + """Run ``sql`` and round-trip the result to an ``xr.Dataset``.""" + pdf = self._execute(self._rewrite(sql, param_values), param_values) + return _pandas_to_dataset(pdf, dims) + + +class _DataFusionNative(EngineContext): + """``xql.XarrayContext`` over the native DataFusion table provider.""" + + flavor = "datafusion (XarrayContext, native)" + + def _connect(self): + try: + import xarray_sql._native # noqa: F401 + except ImportError as exc: + raise RuntimeError( + "GEOBENCH_ENGINE=datafusion requires the compiled " + "xarray_sql._native module (`maturin develop`); use " + "GEOBENCH_ENGINE=datafusion-arrow for the pure-Python " + "pyarrow-dataset path." + ) from exc + import xarray_sql as xql + + self._ctx = xql.XarrayContext() + + def from_dataset(self, name, ds, *, chunks=None, table_names=None): + self._ctx.from_dataset(name, ds, chunks=chunks, table_names=table_names) + + def sql_to_dataset(self, sql, *, dims, param_values=None): + df = ( + self._ctx.sql(sql, param_values=param_values) + if param_values + else self._ctx.sql(sql) + ) + return df.to_dataset(dims=dims) + + +class _DataFusionArrow(EngineContext): + """Plain ``datafusion.SessionContext`` over ``xql.arrow_dataset``.""" + + flavor = "datafusion-arrow (pyarrow dataset, pure Python)" + + def _connect(self): + from datafusion import SessionContext + + self._ctx = SessionContext() + + def _register(self, flat, dataset): + self._ctx.register_dataset(flat, dataset) + + def _execute(self, sql, param_values): + df = ( + self._ctx.sql(sql, param_values=param_values) + if param_values + else self._ctx.sql(sql) + ) + return df.to_pandas() + + +class _DuckDB(EngineContext): + """DuckDB over the same pyarrow pushdown datasets.""" + + flavor = "duckdb" + + def _connect(self): + import duckdb + + self._con = duckdb.connect() + + def _register(self, flat, dataset): + self._con.register(flat, dataset) + + def _execute(self, sql, param_values): + return self._con.execute(sql, param_values or {}).df() + + +class _Polars(EngineContext): + """``polars.SQLContext`` over ``scan_pyarrow_dataset`` frames. + + Keeps the pyarrow datasets and builds the SQLContext per query. + Polars' SQL layer renders TIMESTAMP literals as strptime-plus-cast + expressions it cannot convert to pyarrow filters, so a WHERE over + the full archive would scan everything; the same bounds applied as + native expressions *do* push down. ``_execute`` therefore + pre-filters each frame with the query's window parameters + (identical predicate to the SQL WHERE, which still runs on top). + """ + + flavor = "polars (SQLContext + expression window pushdown)" + + # The window bounds a query passes as parameters, as (column, low + # param, high param); applied per registered frame when the column + # exists — the same inclusive predicate the SQL WHERE states. + _BOUND_PARAMS = ( + ("time", "start", "end"), + ("latitude", "lat_s", "lat_n"), + ("longitude", "lon_w", "lon_e"), + ) + + def _connect(self): + self._tables: dict[str, Any] = {} + + def _register(self, flat, dataset): + self._tables[flat] = dataset + + def _rewrite(self, sql, param_values): + sql = super()._rewrite(sql, param_values) + for key, value in (param_values or {}).items(): + sql = re.sub(rf"\${key}\b", _literal(value), sql) + return sql + + def _execute(self, sql, param_values): + import polars as pl + + ctx = pl.SQLContext() + params = param_values or {} + for flat, dataset in self._tables.items(): + lf = pl.scan_pyarrow_dataset(dataset) + names = set(dataset.schema.names) + for col, lo, hi in self._BOUND_PARAMS: + if col in names and lo in params and hi in params: + lf = lf.filter( + (pl.col(col) >= params[lo]) + & (pl.col(col) <= params[hi]) + ) + ctx.register(flat, lf) + return ctx.execute(sql, eager=True).to_pandas() + + +_IMPLS = { + "datafusion": _DataFusionNative, + "datafusion-arrow": _DataFusionArrow, + "duckdb": _DuckDB, + "polars": _Polars, +} +_ENGINES = tuple(_IMPLS) diff --git a/benchmarks/geospatial/_harness.py b/benchmarks/geospatial/_harness.py new file mode 100644 index 00000000..317d8b19 --- /dev/null +++ b/benchmarks/geospatial/_harness.py @@ -0,0 +1,276 @@ +"""Shared harness for the geospatial SQL benchmarks. + +The suite is *expressibility-first*: each case states a geospatial operation we +normally reach for an array library to perform, expresses it in SQL against +``xarray-sql``, and proves the SQL answer matches an xarray/array reference +implementation. Wall-clock and peak memory are reported too, but the headline +is correctness + clarity of the SQL. + +These helpers keep each case script short and uniform: + +* :func:`banner` / :func:`show_sql` — readable section headers and SQL echo. +* :func:`timed` — a context manager that reports elapsed time and peak memory, + for one-time steps (opening data, registering tables). +* :func:`measured` — a loop wrapper (``for _ in measured(label): ...``) for a + repeatable step (a query, a computation). It runs the body once normally, or — + under ``GEOBENCH_PROFILE`` — a warmup plus ``GEOBENCH_REPS`` timed repetitions, + writing a statistical summary to the ``GEOBENCH_CSV`` perf table. +* :func:`assert_grid_close` — assert a SQL result (round-tripped to an + ``xr.DataArray``) matches an xarray reference, aligned by coordinate label. + Raises ``AssertionError`` on mismatch (so a broken case fails loudly rather + than silently "passing"). +* :func:`run_case` — run a case's ``main()``, turning a raised + :class:`CaseSkipped` (e.g. an offline dataset) into a clean skip. +""" + +from __future__ import annotations + +import contextlib +import csv +import os +import statistics +import sys +import time +import tracemalloc +from collections.abc import Callable, Iterator +from typing import Any + +import xarray as xr + +_WIDTH = 72 + +# Performance profiling, opt-in via environment variables. With GEOBENCH_PROFILE +# set, a ``for _ in measured(label):`` block runs GEOBENCH_WARMUP + GEOBENCH_REPS +# times instead of once; GEOBENCH_CSV= collects one summary row per such +# block into a shared CSV — the perf table. Without the flag, runs are unchanged. +_CSV_HEADER = [ + "case", + "title", + "step", + "reps", + "t_min_s", + "t_median_s", + "t_mean_s", + "t_stdev_s", + "t_max_s", + "peak_mb", +] +_current_case = "" +_current_title = "" + +_EE_SCOPES = [ + "https://www.googleapis.com/auth/earthengine", + "https://www.googleapis.com/auth/cloud-platform", +] + + +class CaseSkipped(Exception): + """Raised by a case when it cannot run in this environment (e.g. offline).""" + + +def initialize_earth_engine() -> Any: + """Initialize Earth Engine from Application Default Credentials, or skip. + + Uses the credentials from ``gcloud auth application-default login`` (with the + Earth Engine scope) and the ADC project — so no separate ``earthengine + authenticate`` OAuth flow is needed, which also sidesteps the "this app is + blocked" error some org policies raise. Override the project with the + ``EARTHENGINE_PROJECT`` environment variable. Returns the initialized ``ee`` + module; raises :class:`CaseSkipped` if EE is unavailable or unauthenticated. + """ + try: + import ee + import google.auth + except ImportError as exc: # pragma: no cover + raise CaseSkipped( + "Earth Engine support needs `pip install earthengine-api`" + ) from exc + try: + credentials, adc_project = google.auth.default(scopes=_EE_SCOPES) + ee.Initialize( + credentials, + project=os.environ.get("EARTHENGINE_PROJECT") or adc_project, + opt_url="https://earthengine-highvolume.googleapis.com", + ) + except Exception as exc: # noqa: BLE001 — not authenticated → skip + raise CaseSkipped( + f"Earth Engine not initialized ({exc}); run " + "`gcloud auth application-default login` (or set EARTHENGINE_PROJECT)" + ) from exc + return ee + + +def banner(text: str) -> None: + """Print a titled section divider.""" + print(f"\n{'─' * _WIDTH}") + print(f" {text}") + print(f"{'─' * _WIDTH}") + + +def show_sql(sql: str, *, label: str = "SQL") -> None: + """Echo a SQL statement so the reader sees exactly what ran.""" + print(f"\n {label}:") + for line in sql.strip("\n").splitlines(): + print(f" │ {line}") + print() + + +@contextlib.contextmanager +def timed(label: str) -> Iterator[None]: + """Time a block and report elapsed wall-clock and peak memory. + + Peak memory is the Python-allocator peak during the block (via + ``tracemalloc``); it captures the materialized result and intermediate + buffers, which is what we care about for "did this blow up memory". + """ + tracemalloc.start() + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield + finally: + elapsed = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + print(f" ⏱ {label}: {elapsed:.3f}s (peak {peak / 1e6:.1f} MB)") + + +def _append_csv(step: str, times: list[float], peak_bytes: int) -> None: + """Append one step's summary stats to the GEOBENCH_CSV perf table, if set.""" + path = os.environ.get("GEOBENCH_CSV", "") + if not path: + return + row = { + "case": _current_case, + "title": _current_title, + "step": step, + "reps": len(times), + "t_min_s": round(min(times), 6), + "t_median_s": round(statistics.median(times), 6), + "t_mean_s": round(statistics.fmean(times), 6), + "t_stdev_s": round(statistics.stdev(times), 6) + if len(times) > 1 + else 0.0, + "t_max_s": round(max(times), 6), + "peak_mb": round(peak_bytes / 1e6, 1), + } + fresh = not os.path.exists(path) or os.path.getsize(path) == 0 + with open(path, "a", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=_CSV_HEADER) + if fresh: + writer.writeheader() + writer.writerow(row) + + +def measured(label: str) -> Iterator[None]: + """Time a repeatable block, optionally repeating it for a perf profile. + + Use it as a loop — ``for _ in measured("SQL …"): got = ``. Without + profiling it runs the body once and prints a ``⏱`` line, exactly like + :func:`timed`. Under ``GEOBENCH_PROFILE`` it runs a warmup pass plus + ``GEOBENCH_REPS`` measured passes, times each, and appends one row of summary + statistics to the ``GEOBENCH_CSV`` perf table. The body must be safe to + repeat — a query or pure computation, not one-time setup such as table + registration (which stays in :func:`timed`). + """ + if not os.environ.get("GEOBENCH_PROFILE"): + with timed(label): + yield + return + reps = max(1, int(os.environ.get("GEOBENCH_REPS", "5"))) + warmup = max(0, int(os.environ.get("GEOBENCH_WARMUP", "1"))) + times: list[float] = [] + peak_max = 0 + for i in range(warmup + reps): + tracemalloc.start() + tracemalloc.reset_peak() + t0 = time.perf_counter() + try: + yield + finally: + elapsed = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + if i >= warmup: + times.append(elapsed) + peak_max = max(peak_max, peak) + _append_csv(label, times, peak_max) + print( + f" 📊 {label}: median {statistics.median(times):.3f}s " + f"[min {min(times):.3f}, max {max(times):.3f}, " + f"n={len(times)}, peak {peak_max / 1e6:.0f} MB]" + ) + + +def assert_grid_close( + name: str, + got: xr.DataArray, + ref: xr.DataArray, + *, + rtol: float = 1e-5, + atol: float = 1e-6, +) -> None: + """Assert two gridded ``DataArray`` results match, then print PASS. + + For cases whose SQL result is round-tripped back to an ``xr.DataArray`` + (via ``XarrayDataFrame.to_dataset``), compare it to the array reference the + xarray way: align ``ref`` onto ``got``'s coordinates and dimension order, + then ``xr.testing.assert_allclose``. This aligns by *label*, so neither side + needs an explicit sort, and NaNs in matching cells compare equal. + + Helper coordinates xarray attaches along the way (e.g. the ``hour`` label a + ``groupby("time.hour")`` leaves behind) are dropped before comparing. + + ``reindex_like`` would quietly align away any cells missing from ``got``, so + a query that returns *fewer* cells than it should would still "pass". Guard + against that first — the suite's whole point is the same numbers, all of them. + """ + short = { + d: (got.sizes[d], ref.sizes[d]) + for d in ref.dims + if d in got.sizes and got.sizes[d] != ref.sizes[d] + } + if short: + raise AssertionError(f"{name}: result misses grid cells {short}") + aligned = ref.reindex_like(got).transpose(*got.dims) + extra = [c for c in aligned.coords if c not in got.coords] + aligned = aligned.drop_vars(extra) + xr.testing.assert_allclose(got, aligned, rtol=rtol, atol=atol) + print( + f" ✅ {name}: SQL matches xarray reference " + f"(n={got.size:,}, coordinate-aligned)" + ) + + +def show_result( + result: xr.DataArray | xr.Dataset, *, label: str = "Result (SQL → xarray)" +) -> None: + """Print the SQL result as an xarray object, using its standard repr. + + Called after the match is verified, so a run shows *what* it computed — the + gridded answer round-tripped back out of SQL as an ``xarray`` object. + """ + print(f"\n {label}:\n") + print(result) + + +def run_case(main: Callable[[], None], title: str) -> int: + """Run a case ``main()``; turn :class:`CaseSkipped` into a clean skip. + + Returns a process exit code: 0 on success or skip, 1 on failure. Use as + ``if __name__ == '__main__': raise SystemExit(run_case(main, '...'))``. + """ + global _current_case, _current_title + _current_title = title + _current_case = os.path.splitext(os.path.basename(sys.argv[0]))[0] + banner(title) + try: + main() + except CaseSkipped as exc: + print(f"\n ⏭ SKIPPED: {exc}") + return 0 + except Exception as exc: # noqa: BLE001 — surface any failure as exit 1 + print(f"\n ❌ FAILED: {type(exc).__name__}: {exc}", file=sys.stderr) + raise + print(f"\n 🎉 {title}: done.") + return 0 diff --git a/benchmarks/geospatial/engine_suite.py b/benchmarks/geospatial/engine_suite.py new file mode 100644 index 00000000..b8cb6e49 --- /dev/null +++ b/benchmarks/geospatial/engine_suite.py @@ -0,0 +1,699 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "aiohttp", +# "coiled", +# "datafusion", +# "duckdb", +# "earthengine-api", +# "gcsfs", +# "numpy", +# "pandas", +# "polars", +# "psutil", +# "pyarrow", +# "pyproj", +# "pystac-client", +# "requests", +# "scipy", +# "shapely", +# "xarray", +# "xarray-sql", +# "xee", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""The geospatial suite across engines and VM sizes, via Coiled Functions. + +Runs the nine geospatial cases (``01_ndvi`` … ``09_warp``) under every +SQL engine the suite supports — DataFusion over the native table +provider (``datafusion``, the original path; requires the compiled +native module), DataFusion over the pure-Python pyarrow dataset +(``datafusion-arrow``), DuckDB, and Polars, selected per process +through ``GEOBENCH_ENGINE`` and the ``_engines`` facade — on one reused +Coiled VM per machine size, driven in parallel across sizes. + +``datafusion`` cells need the compiled native module; the first such +cell on a VM provisions it (see :func:`_ensure_native`) and records the +outcome under ``native`` in its result, so a failed build surfaces as +that cell's error rather than a VM startup failure. The driver's own +build is copied in when it imports on that platform; otherwise rustup +(minimal profile) is installed and the shipped crate is built via the +project's maturin build backend, once per source digest. + +The measurement protocol is exactly ``run_perf.sh``'s: every repetition +is a **fresh process** with no warm-up (``GEOBENCH_PROFILE=1 +GEOBENCH_WARMUP=0 GEOBENCH_REPS=1``), so the SQL side and the xarray +reference each pay a cold read on every rep, and each case's own +correctness assertion (SQL answer == array reference) must pass for the +timing to count. The xarray-reference timings are engine-independent; +the tables report the reference column from the DataFusion runs. + +Coverage notes, recorded rather than hidden: cases 07 and 09 build +DataFusion scalar UDFs on ``xql.XarrayContext`` directly, so every +engine except ``datafusion`` is marked n/a; case 08 reads +through Earth Engine and is left on the original context (EE-gated); +cases 07–09 skip cleanly wherever Earth Engine auth is unavailable +(e.g. on the benchmark VMs) with the reason recorded. + +Each (vm, case, engine) cell returns a plain dict; every completed cell +is appended to a local ``--jsonl`` file immediately, and the driver +prints one timestamped line per event. + +Usage:: + + uv run benchmarks/geospatial/engine_suite.py --local --reps 1 \ + --cases 02_climatology --vms local # in-process check + cd /tmp && uv run ~/path/to/xarray-sql/benchmarks/geospatial/engine_suite.py + +(or ``python .../engine_suite.py`` from an environment that already has +the dependencies; the inline metadata above is for ``uv run``.) + +Remote runs (the second form) must be launched from a working directory +outside the repository: Coiled's package sync resolves the repo's +``uv.lock`` when it finds one at the cwd, and that lock does not carry +the driver's dependencies. ``--local`` runs work from anywhere. +""" + +from __future__ import annotations + +import argparse +import csv +import datetime +import io +import json +import os +import statistics +import subprocess +import sys +import tarfile +import tempfile +import threading +import time +from pathlib import Path +from typing import Any + +REGION = "us-central1" + +CASES = [ + "01_ndvi", + "02_climatology", + "03_zonal_mean", + "04_anomaly", + "05_forecast_skill", + "06_zonal_vector", + "07_reproject_udf", + "08_regrid_weights", + "09_warp", +] +ENGINES = ["datafusion", "datafusion-arrow", "duckdb", "polars"] +# Cases whose SQL builds DataFusion scalar UDFs on XarrayContext directly +# (07, and the UDF half of 09): they run only under ``datafusion``. Case +# 08 is portable SQL but Earth-Engine-gated, so it stays on the original +# context. +NOT_PORTABLE = { + "07_reproject_udf": "n/a (DataFusion scalar UDF)", + "09_warp": "n/a (DataFusion scalar UDF)", + "08_regrid_weights": "not ported (Earth-Engine-gated case)", +} +VM_SIZES = ["e2-standard-8", "e2-standard-16", "e2-standard-32"] + + +def cluster_name(vm: str) -> str: + return "xql-geo-" + vm.replace("standard-", "") + + +# -------------------------------------------------------------------------- +# Remote side (runs inside the coiled function, or locally with --local) +# -------------------------------------------------------------------------- + + +def _install_src(src_targz: bytes | None) -> tuple[str, str]: + """Unpack the shipped source tree; returns (sys.path root, geo dir). + + The root is keyed by the tarball's hash so a reused warm VM never + serves a stale tree from an earlier driver run. + """ + import hashlib + + digest = hashlib.md5(src_targz or b"local").hexdigest()[:10] + root = f"/tmp/xql_geo_src_{digest}" + # Written only after extractall returns, so an interrupted + # extraction is retried instead of reused as a corrupt tree. + marker = os.path.join(root, ".extraction-complete") + if src_targz is not None and not os.path.exists(marker): + os.makedirs(root, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(src_targz), mode="r:gz") as tf: + tf.extractall(root) # noqa: S202 — our own tarball + open(marker, "w").close() + return root, os.path.join(root, "benchmarks", "geospatial") + + +def _run_logged(cmd, **kwargs) -> None: + """subprocess.run(check=True) that surfaces stderr on failure.""" + proc = subprocess.run(cmd, capture_output=True, text=True, **kwargs) + if proc.returncode != 0: + raise RuntimeError( + f"{cmd if isinstance(cmd, str) else ' '.join(cmd)} failed:\n" + f"{proc.stderr[-800:]}" + ) + + +def _ensure_native(src_root: str) -> str: + """Make ``xarray_sql._native`` importable from ``src_root``. + + Tries, in order: the module already present in the tree; the one + installed in this interpreter's environment (copied in, when built + for this platform); a from-source build of the shipped crate — + rustup (minimal profile) plus ``pip wheel``, which drives the + project's maturin build backend — installed over the pure-Python + copy. The built module lands in ``src_root``, which is keyed by + source digest, so a warm VM builds at most once per source state. + + Returns a status string for the run log. + """ + import glob + import importlib.util + import shutil + + # cwd well inside the tree, so `-c` resolves xarray_sql only through + # PYTHONPATH=src_root — the same view the case subprocesses get. + geo_dir = os.path.join(src_root, "benchmarks", "geospatial") + env = dict(os.environ, PYTHONPATH=src_root) + + def _importable() -> bool: + return ( + subprocess.run( + [sys.executable, "-c", "import xarray_sql._native"], + env=env, + cwd=geo_dir, + capture_output=True, + ).returncode + == 0 + ) + + if _importable(): + return "importable" + + try: + spec = importlib.util.find_spec("xarray_sql._native") + except ImportError: + spec = None + if spec is not None and spec.origin: + shutil.copy2( + spec.origin, + os.path.join(src_root, "xarray_sql", os.path.basename(spec.origin)), + ) + if _importable(): + return "copied from driver environment" + + t0 = time.monotonic() + build_env = dict(env) + cargo_bin = os.path.expanduser("~/.cargo/bin") + build_env["PATH"] = cargo_bin + os.pathsep + build_env.get("PATH", "") + if shutil.which("cargo", path=build_env["PATH"]) is None: + _run_logged( + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs " + "| sh -s -- -y --profile minimal --default-toolchain stable", + shell=True, + env=build_env, + ) + wheel_dir = os.path.join(src_root, "wheelhouse") + _run_logged( + [ + sys.executable, + "-m", + "pip", + "wheel", + "--no-deps", + "-w", + wheel_dir, + src_root, + ], + env=build_env, + ) + wheel = sorted(glob.glob(os.path.join(wheel_dir, "xarray_sql-*.whl")))[-1] + _run_logged( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--upgrade", + "--target", + src_root, + wheel, + ], + ) + if not _importable(): + raise RuntimeError(f"built {wheel} but xarray_sql._native still fails") + return f"built from source in {time.monotonic() - t0:.0f}s" + + +def run_case_cell( + case: str, + engine: str, + reps: int, + src_targz: bytes | None = None, + rep_timeout: float = 600.0, +) -> dict: + """One (case, engine) cell: ``reps`` fresh-process cold runs.""" + result: dict[str, Any] = { + "case": case, + "engine": engine, + "status": "ok", + "reps": [], + } + try: + src_root, geo_dir = _install_src(src_targz) + if engine == "datafusion": + try: + result["native"] = _ensure_native(src_root) + except Exception: + result["native"] = "provisioning failed" + raise + env = dict( + os.environ, + GEOBENCH_ENGINE=engine, + GEOBENCH_PROFILE="1", + GEOBENCH_WARMUP="0", + GEOBENCH_REPS="1", + PYTHONUNBUFFERED="1", + PYTHONPATH=src_root, + ) + rows: list[dict] = [] + for rep in range(1, reps + 1): + with tempfile.NamedTemporaryFile(suffix=".csv") as csv_file: + env["GEOBENCH_CSV"] = csv_file.name + t0 = time.perf_counter() + try: + proc = subprocess.run( + [sys.executable, f"{case}.py"], + cwd=geo_dir, + env=env, + capture_output=True, + text=True, + timeout=rep_timeout, + ) + except subprocess.TimeoutExpired: + result["reps"].append({"rep": rep, "status": "timeout"}) + result["status"] = "timeout" + break + wall = round(time.perf_counter() - t0, 3) + out = proc.stdout + # Exit status outranks the skip marker: a case that + # prints SKIPPED and then crashes is an error, not a skip. + if proc.returncode == 0 and "SKIPPED" in out: + reason = next( + ( + line.split("SKIPPED:", 1)[1].strip() + for line in out.splitlines() + if "SKIPPED:" in line + ), + "skipped", + ) + result.update(status="skip", reason=reason[:300]) + break + if proc.returncode != 0: + result.update( + status="error", + error=(proc.stderr.strip() or out.strip())[-600:], + ) + break + flavor = next( + ( + line.split("engine:", 1)[1].strip() + for line in out.splitlines() + if "engine:" in line + ), + engine, + ) + result["flavor"] = flavor + with open(csv_file.name) as fh: + for row in csv.DictReader(fh): + row["rep"] = rep + rows.append(row) + result["reps"].append( + {"rep": rep, "status": "ok", "wall_s": wall} + ) + print(f"[vm] {case} x {engine}: rep {rep} {wall}s", flush=True) + steps: dict[str, dict] = {} + for row in rows: + step = steps.setdefault( + row["step"], {"times_s": [], "peak_mb": 0.0} + ) + step["times_s"].append(float(row["t_median_s"])) + step["peak_mb"] = max(step["peak_mb"], float(row["peak_mb"])) + for step in steps.values(): + times = step["times_s"] + step["median_s"] = round(statistics.median(times), 3) + step["min_s"] = round(min(times), 3) + step["max_s"] = round(max(times), 3) + step["n"] = len(times) + result["steps"] = steps + if result["status"] == "ok" and not steps: + result.update(status="error", error="no perf rows produced") + except Exception as exc: # noqa: BLE001 — cell errors are data + result.update(status="error", error=f"{type(exc).__name__}: {exc}") + return result + + +def probe_environment(src_targz: bytes | None = None) -> dict: + """Machine spec + package versions, gathered where the cells run.""" + import platform + + _install_src(src_targz) + info = { + "platform": platform.platform(), + "python": platform.python_version(), + "cpus": os.cpu_count(), + "node": platform.node(), + } + try: + import psutil + + info["mem_gb"] = round(psutil.virtual_memory().total / 2**30, 1) + except Exception: # noqa: BLE001 + pass + versions = {} + for pkg in ["duckdb", "polars", "datafusion", "pyarrow", "xarray"]: + try: + from importlib import metadata + + versions[pkg] = metadata.version(pkg) + except Exception: # noqa: BLE001 + versions[pkg] = "missing" + return {"machine": info, "versions": versions} + + +# -------------------------------------------------------------------------- +# Driver +# -------------------------------------------------------------------------- + +_PRINT_LOCK = threading.Lock() + + +def log(vm: str, msg: str) -> None: + now = datetime.datetime.now().strftime("%H:%M:%S") + with _PRINT_LOCK: + print(f"[{now}][{vm}] {msg}", flush=True) + + +def _pack_src() -> bytes: + """gzip tar of xarray_sql, benchmarks/geospatial, and the Rust crate. + + Byte-identical for identical file contents (gzip and tar metadata + normalized): _install_src keys its extraction root — and therefore + _ensure_native's build cache on a warm VM — on the digest of these + bytes. + """ + import gzip + + repo = Path(__file__).resolve().parents[2] + + def _normalize(info: tarfile.TarInfo) -> tarfile.TarInfo: + info.mtime = 0 + info.mode = 0o644 + info.uid = info.gid = 0 + info.uname = info.gname = "" + return info + + paths: list[Path] = [] + for rel in ["xarray_sql", "benchmarks/geospatial"]: + paths += [ + p + for p in sorted((repo / rel).rglob("*.py")) + if "__pycache__" not in p.parts + ] + # The crate sources, so `datafusion` cells can build the native + # module where it is not already importable (see _ensure_native). + for rel in [ + "src", + "Cargo.toml", + "Cargo.lock", + "pyproject.toml", + "README.md", + ]: + target = repo / rel + paths += ( + [p for p in sorted(target.rglob("*")) if p.is_file()] + if target.is_dir() + else [target] + ) + buf = io.BytesIO() + with gzip.GzipFile(fileobj=buf, mode="wb", mtime=0) as gz: + # GzipFile is a binary stream at runtime; typeshed wants IO[bytes]. + with tarfile.open(fileobj=gz, mode="w") as tf: # type: ignore[arg-type] + for path in paths: + tf.add( + path, + arcname=str(path.relative_to(repo)), + filter=_normalize, + ) + return buf.getvalue() + + +def _drive_vm(vm, cells, args, src, results, jsonl_lock): + """Run every (case, engine) cell for one VM size, sequentially.""" + if vm == "local": + remote_cell, remote_probe = run_case_cell, probe_environment + submit = None + else: + import coiled + + deco = coiled.function( + name=cluster_name(vm), + vm_type=vm, + region=REGION, + keepalive="10m", + idle_timeout="20 minutes", + spot_policy="on-demand", + package_sync_ignore=["xarray_sql", "xarray-sql"], + environ={"PYTHONUNBUFFERED": "1"}, + ) + remote_cell, remote_probe = deco(run_case_cell), deco(probe_environment) + submit = remote_cell.submit + + log(vm, "probing environment (provisions the VM on first call)...") + meta = None + for attempt in range(1, 4): + try: + meta = remote_probe(src) + break + except Exception as exc: # noqa: BLE001 — transient control plane + log(vm, f"probe attempt {attempt} failed: {exc}"[:200]) + if attempt < 3: + time.sleep(30 * attempt) + if meta is None: + log(vm, "giving up: VM never came up") + return + log(vm, f"machine: {json.dumps(meta['machine'])}") + total = len(cells) + for k, (case, engine) in enumerate(cells, 1): + tag = f"cell {k}/{total} {case} x {engine}" + if case in NOT_PORTABLE and engine != "datafusion": + rec = { + "case": case, + "engine": engine, + "status": "n/a", + "reason": NOT_PORTABLE[case], + } + else: + log(vm, f"{tag}: submitted") + t0 = time.monotonic() + try: + if submit is None: + rec = run_case_cell(case, engine, args.reps, src) + else: + fut = submit(case, engine, args.reps, src) + rec = fut.result(timeout=args.cell_timeout) + except Exception as exc: # noqa: BLE001 + rec = { + "case": case, + "engine": engine, + "status": "error", + "error": f"{type(exc).__name__}: {exc}"[:500], + } + rec["cell_wall_s"] = round(time.monotonic() - t0, 1) + rec["vm"] = vm + if rec.get("native", "importable") != "importable": + log(vm, f"{tag}: native module {rec['native']}") + results.append(rec) + with jsonl_lock, open(args.jsonl, "a") as fh: + fh.write(json.dumps(rec) + "\n") + if rec["status"] == "ok": + sql_step = next( + ( + s + for name, s in rec.get("steps", {}).items() + if name.startswith("SQL") + ), + None, + ) + brief = ( + f"SQL median {sql_step['median_s']}s (n={sql_step['n']})" + if sql_step + else "ok" + ) + log(vm, f"{tag}: ok {brief} [{rec.get('flavor', engine)}]") + else: + detail = rec.get("reason") or rec.get("error", "") + log(vm, f"{tag}: {rec['status']} {detail[:200]}") + meta_rec = {"vm": vm, "case": "_meta", "engine": "", **meta} + results.append(meta_rec) + with jsonl_lock, open(args.jsonl, "a") as fh: + fh.write(json.dumps(meta_rec) + "\n") + if submit is not None: + # Shut the VM down the moment its last cell finishes — don't + # leave the teardown to keepalive expiry. + try: + remote_cell.cluster.shutdown() + log(vm, "cluster shut down") + except Exception as exc: # noqa: BLE001 — teardown best-effort + log(vm, f"cluster shutdown failed: {exc}"[:200]) + + +def _markdown(results: list[dict]) -> str: + """One case x engine table per VM (SQL median s; reference column).""" + out = [] + vms = list(dict.fromkeys(r["vm"] for r in results)) + for vm in vms: + rows = [r for r in results if r["vm"] == vm and r["case"] != "_meta"] + if not rows: + continue + cases = list(dict.fromkeys(r["case"] for r in rows)) + out.append(f"\n### {vm}\n") + out.append("| Case | " + " | ".join(ENGINES) + " | xarray reference |") + out.append("|---|" + "---|" * (len(ENGINES) + 1)) + by = {(r["case"], r["engine"]): r for r in rows} + for case in cases: + cells = [] + for engine in ENGINES: + r = by.get((case, engine)) + if r is None: + cells.append("-") + elif r["status"] != "ok": + detail = r.get("reason") or r.get("error", "") + cells.append(f"{r['status']}: {detail[:40]}") + else: + s = next( + ( + v + for k, v in r["steps"].items() + if k.startswith("SQL") + ), + None, + ) + cells.append( + f"{s['median_s']:.3f}s (n={s['n']}, " + f"{s['peak_mb']:.0f} MB)" + if s + else "?" + ) + df_run = by.get((case, "datafusion"), {}) + ref = (df_run.get("steps") or {}).get("xarray reference") + ref_text = ( + f"{ref['median_s']:.3f}s ({ref['peak_mb']:.0f} MB)" + if ref + else "-" + ) + out.append(f"| {case} | " + " | ".join(cells) + f" | {ref_text} |") + return "\n".join(out) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--local", action="store_true") + ap.add_argument("--reps", type=int, default=5) + ap.add_argument("--cases", default=",".join(CASES)) + ap.add_argument("--engines", default=",".join(ENGINES)) + ap.add_argument("--vms", default=",".join(VM_SIZES)) + ap.add_argument("--cell-timeout", type=float, default=1800.0) + ap.add_argument("--out", default="engine_suite_results.json") + ap.add_argument("--jsonl", default="engine_suite_results.jsonl") + args = ap.parse_args() + + cases = [c for c in args.cases.split(",") if c] + engines = [e for e in args.engines.split(",") if e] + # dict.fromkeys: duplicate VM names would share one cluster and defeat + # the incomplete-run detection, which matches records by VM name. + vms = ( + ["local"] + if args.local + else list(dict.fromkeys(v for v in args.vms.split(",") if v)) + ) + cells = [(c, e) for c in cases for e in engines] + log("plan", f"{len(vms)} VMs x {len(cells)} cells, reps={args.reps}") + for c, e in cells: + note = ( + f" [{NOT_PORTABLE[c]}]" + if c in NOT_PORTABLE and e != "datafusion" + else "" + ) + log("plan", f" {c} x {e}{note}") + src = _pack_src() + log("plan", f"packed source: {len(src) / 1024:.0f} KiB") + + open(args.jsonl, "w").close() + results: list[dict] = [] + jsonl_lock = threading.Lock() + threads = [ + threading.Thread( + target=_drive_vm, + args=(vm, cells, args, src, results, jsonl_lock), + name=vm, + ) + for vm in vms + ] + for i, t in enumerate(threads): + if i: # stagger: concurrent package-sync scans trip the server + time.sleep(20) + t.start() + for t in threads: + t.join() + + payload = { + "meta": { + "region": REGION, + "reps": args.reps, + "protocol": "fresh process per rep, no warmup, cold reads", + }, + "results": results, + } + with open(args.out, "w") as fh: + json.dump(payload, fh, indent=2) + md = _markdown([r for r in results if r.get("case")]) + md_path = os.path.splitext(args.out)[0] + ".md" + with open(md_path, "w") as fh: + fh.write(md + "\n") + print(md) + log("done", f"wrote {args.out}, {md_path}, {args.jsonl}") + # A VM that never produced its _meta record never ran its cells; + # exit nonzero so partial runs cannot pass for complete ones. + incomplete = [ + vm + for vm in vms + if not any( + r.get("vm") == vm and r.get("case") == "_meta" for r in results + ) + ] + if incomplete: + log("done", f"incomplete run: no results from {', '.join(incomplete)}") + sys.exit(1) + failed = [ + f"{r['vm']}/{r['case']}/{r['engine']}" + for r in results + if r.get("status") in ("error", "timeout") + ] + if failed: + log("done", f"failed cells: {', '.join(failed)}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/geospatial/perf_summary.py b/benchmarks/geospatial/perf_summary.py new file mode 100755 index 00000000..9c5fad0f --- /dev/null +++ b/benchmarks/geospatial/perf_summary.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Aggregate the cold-run perf CSV into a per-step summary and a markdown table. + +``run_perf.sh`` runs each case once per fresh process and appends one row per +measured step (the SQL operation, the xarray reference) to a raw CSV — so every +row is an *independent cold measurement*. This reads those rows and reports, per +(case, step), the median and spread across the cold runs, writes a summary CSV, +and prints a markdown table. + + perf_summary.py RAW.csv [SUMMARY.csv] +""" + +from __future__ import annotations + +import csv +import statistics +import sys + +_HEADER = [ + "case", + "title", + "step", + "reps", + "t_min_s", + "t_median_s", + "t_mean_s", + "t_stdev_s", + "t_max_s", + "peak_mb", +] + + +def main() -> None: + raw_path = sys.argv[1] + summary_path = sys.argv[2] if len(sys.argv) > 2 else None + + with open(raw_path, newline="") as fh: + rows = list(csv.DictReader(fh)) + + # Each raw row is one cold run (reps=1), so its t_median_s == the sample. + groups: dict[tuple[str, str, str], list[tuple[float, float]]] = {} + for r in rows: + key = (r["case"], r["title"], r["step"]) + groups.setdefault(key, []).append( + (float(r["t_median_s"]), float(r["peak_mb"])) + ) + + summary = [] + for (case, title, step), vals in groups.items(): + times = [t for t, _ in vals] + summary.append( + { + "case": case, + "title": title, + "step": step, + "reps": len(times), + "t_min_s": round(min(times), 6), + "t_median_s": round(statistics.median(times), 6), + "t_mean_s": round(statistics.fmean(times), 6), + "t_stdev_s": round(statistics.stdev(times), 6) + if len(times) > 1 + else 0.0, + "t_max_s": round(max(times), 6), + "peak_mb": round(max(p for _, p in vals), 1), + } + ) + + summary.sort( + key=lambda r: ( + str(r["case"]), + 0 if str(r["step"]).upper().startswith("SQL") else 1, + ) + ) + + if summary_path: + with open(summary_path, "w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=_HEADER) + writer.writeheader() + writer.writerows(summary) + + print( + "| Case | Step | reps | median (s) | stdev (s) | min (s) | max (s) | peak (MB) |" + ) + print("|---|---|--:|--:|--:|--:|--:|--:|") + seen: set[str] = set() + for r in summary: + case = str(r["case"]) + cell = str(r["title"]) if case not in seen else "" + seen.add(case) + step = ( + "SQL" + if str(r["step"]).upper().startswith("SQL") + else "xarray reference" + ) + print( + f"| {cell} | {step} | {r['reps']} | {r['t_median_s']:.3f} | " + f"{r['t_stdev_s']:.3f} | {r['t_min_s']:.3f} | {r['t_max_s']:.3f} | " + f"{r['peak_mb']:.1f} |" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/geospatial/run_all.sh b/benchmarks/geospatial/run_all.sh new file mode 100755 index 00000000..2429b2e0 --- /dev/null +++ b/benchmarks/geospatial/run_all.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# +# Run every geospatial benchmark case with `uv run` (each script declares its +# own dependencies via PEP 723 inline metadata). Works from any directory: it +# resolves its own location, so the cases are found and the paths handed to +# `uv run` are absolute. +# +# ./run_all.sh # from anywhere +# bash benchmarks/geospatial/run_all.sh +# +# Each script's metadata points xarray-sql at this local checkout +# ([tool.uv.sources] path = "../../"), so uv uses the in-repo build (which has +# features newer than the latest PyPI release) — relative to the script, so it +# resolves no matter the working directory. +# +# Network/credential-gated cases (ERA5, WeatherBench2, Earth Engine) skip +# cleanly when their data is unavailable. Exits non-zero if any case fails +# (a skip is not a failure). + +set -uo pipefail + +# Directory this script lives in, regardless of the caller's working directory. +DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +status=0 +for script in "$DIR"/[0-9][0-9]_*.py; do + name="$(basename "$script")" + echo "════════════════════════════════════════ ${name}" + if uv run "$script"; then + echo "✅ ${name}" + else + echo "❌ ${name} (exit $?)" + status=1 + fi +done + +exit "$status" diff --git a/benchmarks/geospatial/run_perf.sh b/benchmarks/geospatial/run_perf.sh new file mode 100755 index 00000000..e7294b1e --- /dev/null +++ b/benchmarks/geospatial/run_perf.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Cold-vs-cold performance benchmark. +# +# Runs each case once per fresh process, with no warmup, repeated GEOBENCH_REPS +# times. A fresh process per repetition is deliberate: it makes the SQL operation +# AND the xarray reference each pay a *cold* read on every measurement. An +# in-process warm loop is unfair here — `xr.open_zarr(chunks=None)` caches each +# variable in memory after the first read, so the xarray reference would serve +# later reps from RAM while the SQL side re-reads the store. One process per rep +# defeats that (and the OS/connection reuse), so both sides are measured cold. +# +# Each run appends one row per step to a raw CSV; this script then aggregates the +# median/spread across the independent cold runs into a summary CSV + markdown. +# +# GEOBENCH_REPS=5 benchmarks/geospatial/run_perf.sh [summary.csv] +# +# For representative numbers use a release build of xarray-sql and run close to +# the data (a VM in the bucket's region). Override the launcher with +# GEOBENCH_PYRUN (e.g. `GEOBENCH_PYRUN="python"` to use an already-built venv +# instead of the default `uv run`, which builds an unoptimized editable install). +set -u + +DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPS="${GEOBENCH_REPS:-5}" +SUMMARY="${1:-$DIR/perf.csv}" +RAW="$(mktemp)" +read -r -a PYRUN <<<"${GEOBENCH_PYRUN:-uv run}" + +for f in "$DIR"/0[1-9]_*.py; do + name="$(basename "$f")" + for i in $(seq 1 "$REPS"); do + if GEOBENCH_PROFILE=1 GEOBENCH_WARMUP=0 GEOBENCH_REPS=1 GEOBENCH_CSV="$RAW" \ + "${PYRUN[@]}" "$f" >/dev/null 2>&1; then + echo " $name rep $i/$REPS ok" + else + echo " $name rep $i/$REPS skip/fail" + fi + done +done + +# Aggregate the per-process cold runs (one row each) into a per-step summary and +# a markdown table. +python3 "$DIR/perf_summary.py" "$RAW" "$SUMMARY" +echo "wrote $SUMMARY" diff --git a/docs/assets/logo.svg b/docs/assets/logo.svg new file mode 100644 index 00000000..d2406f27 --- /dev/null +++ b/docs/assets/logo.svg @@ -0,0 +1,104 @@ + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 00000000..ea38c9bf --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1 @@ +--8<-- "CONTRIBUTING.md" diff --git a/docs/engines.md b/docs/engines.md new file mode 100644 index 00000000..11dacac7 --- /dev/null +++ b/docs/engines.md @@ -0,0 +1,259 @@ +# Engines + +xarray-sql translates **data, not queries**. It does not own a SQL +dialect, a query IR, or a transpiler: you pick a query engine and write +that engine's native SQL, using that engine's extension ecosystem +(spatial, H3, …) directly. xarray-sql implements the two seams no engine +builds for itself: + +1. **register** — a lazy `xarray.Dataset` becomes a table on the + engine's own connection, streamed as Arrow record batches only while + a query executes. +2. **round-trip** — the engine's Arrow result plus the source Dataset as + a *template* becomes a labeled `xr.Dataset` again: attrs, non-dim + coordinates, and dtypes recovered. *SQL in, array out.* + +Everything between the seams — geometry functions, dialects, +optimizers — belongs to the engine. + +## DataFusion (default) + +DataFusion is the built-in engine, wrapped in a session: + +```python +import xarray_sql as xql + +ctx = xql.XarrayContext() +ctx.from_dataset("era5", ds, chunks={"time": 24}) +result = ctx.sql("SELECT ... FROM era5").to_dataset() +``` + +This is the deepest integration: the Rust `TableProvider` gives +partition pruning on dimension predicates, projection pushdown to the +storage layer, exact per-partition statistics for the optimizer, and a +lazy chunked round-trip (`to_dataset(chunks=...)`). + +The generic entry point dispatches here too: `xql.register(ctx, "era5", ds)` +works on any `datafusion.SessionContext`. + +### Relation to zarr-datafusion + +[zarr-datafusion](https://crates.io/crates/zarr-datafusion) extends +DataFusion with SQL over Zarr stores natively (early days — a single +0.1.0 release at the time of writing) — for plain-Zarr sources that is +the engine-native path, the same role duckdb-zarr plays for DuckDB. +This library's role is complementary there too: anything xarray can +open (NetCDF, GRIB, Earth Engine via Xee, CF-decoded/virtual datasets, +in-memory arrays), and the round-trip from a query result back to a +labeled Dataset, which no engine extension provides. + + +## DuckDB (adapter) + +```sh +pip install 'xarray-sql[duckdb]' +``` + +```python +import duckdb +import xarray_sql as xql + +con = duckdb.connect() +xql.register(con, "era5", ds) # seam 1 + +con.sql("INSTALL spatial; LOAD spatial;") # DuckDB's own shelf +rel = con.sql(""" + SELECT time, lat, lon, AVG(t2m) AS t2m + FROM era5 + WHERE lat BETWEEN 40 AND 41 + GROUP BY time, lat, lon +""") + +out = xql.to_dataset(rel, template=ds) # seam 2 +``` + +The adapter registers an `XarrayPushdownDataset` — a +`pyarrow.dataset.Dataset` subclass (the same pattern +[Lance](https://github.com/lancedb/lance) uses for `LanceDataset`), so +DuckDB hands each query's column list and pushed predicate to the +source. The scan then loads only the data variables the query mentions, +prunes chunks whose coordinate ranges cannot satisfy the predicate +(via Arrow's own guarantee simplification — sound for every predicate +shape), and prefetches surviving chunks on a thread pool. The table is +lazy, re-queryable, and a bounding-box query over a billions-of-pixels +raster answers in about a second because only the intersecting chunks +are ever read. + +Pushed comparison filters are a correctness contract in DuckDB (it +deletes them from its own plan), so the scanner always applies the +exact expression via pyarrow — pruning is only an optimization on top. +`XarrayArrowStream`, the dependency-light re-scannable C-stream wrapper +without pushdown, remains available as a fallback. + +Details that matter in production: + +- **Finely partitioned axes** (e.g. hourly-chunked reanalysis time with + hundreds of thousands of chunks) prune through a two-level shadow: + a coarse pass over at most 1024 buckets, refined per surviving + bucket — so pruning cost is bounded regardless of chunk count, and + refinement is skipped when a predicate matches most of the axis. +- **Tuning** via `xql.register(con, name, ds, batch_size=..., + prefetch=..., prefetch_bytes=..., coalesce_rows=...)`: `prefetch` + bounds concurrent chunk loads, `prefetch_bytes` caps estimated bytes + in flight, `coalesce_rows` merges runs of consecutive surviving + chunks into single reads, `batch_size` caps rows per Arrow batch. + See the [performance guide](performance.md#the-memory-contract). +- **Source parallelism matters as much as the adapter's**: rioxarray + serializes GDAL tile reads behind a lock by default, which caps any + scan at single-stream speed regardless of `prefetch`. Open rasters + with `rioxarray.open_rasterio(..., lock=False)` — measured 6× on + full scans of a 9-billion-pixel cloud GeoTIFF, making remote reads + as fast as a local copy. + + +### Relation to duckdb-zarr + +[duckdb-zarr](https://github.com/xqlsystems/duckdb-zarr) reads Zarr +stores natively inside DuckDB, with projection pushdown — for +plain-Zarr sources it is the engine-native path and will beat this +adapter. The adapter's role is complementary: anything xarray can open +(NetCDF, GRIB, Earth Engine via Xee, CF-decoded/virtual datasets, +in-memory arrays), and the round-trip from a DuckDB result back to a +labeled Dataset, which no engine extension provides. + +## Polars (via the pyarrow dataset protocol) + +```sh +pip install 'xarray-sql[polars]' +``` + +`xql.arrow_dataset(ds)` returns a real `pyarrow.dataset.Dataset`, so +any engine that consumes that protocol gets the same lazy scan with +projection pushdown and coordinate-range chunk pruning — no adapter +code at all. Polars works today: + +```python +import polars as pl +import xarray_sql as xql + +lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) +out = ( + lf.filter(pl.col("lat") > 0) + .group_by("time") + .agg(pl.col("t2m").mean()) + .collect() +) +xql.to_dataset(out, template=ds) # polars frames speak Arrow PyCapsule +``` + +Polars pushes its predicate and column selection into the dataset scan +(verified: a filtered group-by read 1 of 20 chunks and 3 of 5 columns), +and its results round-trip through `xql.to_dataset` unchanged. The +chunked round-trip is fully supported: windows re-execute on Polars' +streaming engine. + + +## Engine support matrix + +What each integration provides. Known issues and constraints live on +[Known issues & limitations](limitations.md). + +| | DataFusion | DuckDB | Polars | +|---|---|---|---| +| Register | `XarrayContext` / any `SessionContext` | `xql.register(con, name, ds)` | `pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` | +| Projection pushdown | yes | yes | yes | +| Chunk pruning on dim predicates | yes | yes | yes | +| Eager round-trip (`xql.to_dataset`) | yes | yes | yes | +| Chunked round-trip (`chunks=`) | re-execution | `spill=True` [^spill-only] | re-execution (streaming engine) | +| `geometry` column ([geospatial](geospatial.md#geoarrow-point-geometry-columns)) | annotated WKB passes through | native `GEOMETRY` (`"wkb"` encoding) | plain binary/struct | +| Mixed-dimension datasets | one schema, `name.group` tables | `_` tables | filter `data_vars` before `arrow_dataset` | +| Version floor | bundled (core dependency) | `duckdb >= 1.4` (tested on 1.5) | tested on `polars 1.42` | + +[^spill-only]: Why DuckDB relations do not re-execute — and two other + engine-specific issues worth knowing — is explained on + [Known issues & limitations](limitations.md#engine-specific-issues). + +## The lazy round-trip across engines + +`xql.to_dataset(result, chunks=...)` reconstructs a query result as a +*chunked, lazy* `xr.Dataset`: each output chunk re-executes the engine's +query narrowed to that chunk's coordinate window on first access. Over a +table registered through xarray-sql, the window's range predicate flows +back into chunk pruning at the source — accessing one output chunk reads +only the source chunks it maps onto. + +```mermaid +flowchart TB + R["xql.to_dataset(result, ...)"] --> K{"chunks=?"} + K -- "None (default)" --> E["eager: materialize once
max_result_bytes= guards both the
Arrow stream and the dense grid"] + K -- "mapping / auto / inherit" --> SP{"spill=?"} + SP -- "False (default)" --> HD{"result type"} + HD -- "Polars LazyFrame/DataFrame
DataFusion DataFrame" --> RX["re-execution: each window
re-runs the query narrowed to its
coordinate range (flows back into
chunk pruning at the source)"] + HD -- "DuckDB relation" --> NO["NotImplementedError
(upstream deadlock — see
Known issues)"] + HD -- "one-shot Arrow stream" --> NO2["TypeError
(nothing to re-execute)"] + SP -- "True / directory" --> SPL["one-pass spill: stream once
(bounded memory) → temp Parquet →
windows re-execute against the file
(row-group pruning); file deleted
with the Dataset"] +``` + +**Choosing:** re-execution pays per window — right when you'll touch a +few windows of a huge result. Spill pays one full pass plus temporary +disk — right when you'll touch most of the result, when the producer +is a DuckDB relation, or when all you have is a one-shot stream. + +Two knobs matter at scale: + +- `coords="template"` trusts the template's coordinate arrays instead of + running one `DISTINCT` query per dimension — construction then reads + nothing at all. Only valid when the result spans the template's full + extent (an unfiltered scan). On ARCO-ERA5 (1.32M hourly chunks) this + builds a lazy view over a 1.37-trillion-row table in ~0.3 s with zero + source reads; a one-day window then computes in ~2 s reading only the + source chunks under the window. +- Contiguous windows become two-literal range predicates the engine can + push and the source can prune on; stepped or fancy selections fall + back to explicit value lists (exact, just less prunable). + +With `spill=True`, the result is streamed **once** (bounded memory) +into a temporary Parquet file and windows re-execute against that +file — the right shape when most of the result will be touched, the +only chunked option for one-shot Arrow streams, and the required path +for DuckDB relations (see the DuckDB section above). Polars/DataFusion +re-execution remains the default for window-at-a-time access over huge +results. + +## Adding an engine + +An adapter implements one small contract +(`xarray_sql.backends.base.EngineAdapter`): `matches(con)` recognizes +the engine's connection object without importing the engine, and +`register(con, name, ds, chunks=...)` attaches the Dataset as a table. +Arrow C streams are the common wire; pushdown quality is where adapters +differ. The round-trip needs no per-engine work as long as the engine +can hand back Arrow. + +A complete adapter, modeled on the DuckDB one: + +```python +from xarray_sql.backends.base import register_adapter +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + +@register_adapter +class AcmeAdapter: + """Registers Datasets on acme.Connection objects.""" + + @staticmethod + def matches(con) -> bool: + # type inspection only, so `acme` stays an optional dependency + return type(con).__module__.split(".")[0] == "acme" + + @staticmethod + def register(con, name, ds, *, chunks=None, **kwargs): + dataset = XarrayPushdownDataset(ds, chunks, **kwargs) + con.register_arrow(name, dataset) # the engine's own API + return con +``` + +`xql.register(con, "t", ds)` then dispatches here whenever `matches` +recognizes the connection. Engines that consume the pyarrow dataset +protocol (DuckDB, Polars) get projection pushdown and chunk pruning for +free; an engine that only accepts Arrow streams can register +`XarrayArrowStream(ds)` instead, trading pushdown away. diff --git a/docs/examples.md b/docs/examples.md new file mode 100644 index 00000000..483fb2a5 --- /dev/null +++ b/docs/examples.md @@ -0,0 +1,158 @@ +# Examples + +A query result can be consumed two ways: as a flat pandas DataFrame +(`to_pandas`) or written back to an Xarray Dataset (`to_dataset`). This computes +a climatology — the mean annual cycle, one value per month of the year — and +shows both. + +> **Note:** this example also needs `pooch` and a netCDF backend (for the +> tutorial download) and `matplotlib` (for the plot): +> `pip install pooch netCDF4 matplotlib`. + +```python +import xarray as xr +import xarray_sql as xql + +ds = xr.tutorial.open_dataset('air_temperature') + +ctx = xql.XarrayContext() +ctx.from_dataset('air', ds, chunks=dict(time=100)) + +clim = ctx.sql(''' + SELECT + CAST(date_part('month', "time") AS INTEGER) AS month, + AVG("air") AS air + FROM + "air" + GROUP BY + CAST(date_part('month', "time") AS INTEGER) + ORDER BY + month +''') + +# Option 1: a flat pandas DataFrame. +clim.to_pandas().head() + +# Option 2: round-trip back to an Xarray Dataset and plot the annual cycle as +# a time series. `to_dataset()` infers dimensions from the registered table's +# surviving dims, so a GROUP BY on a real dimension needs no `dims=`. Here +# `month` is a derived column, not a registered dim, so name it explicitly; +# the variable's units are recovered from the registered table. +clim_ds = clim.to_dataset(dims=["month"]) +clim_ds["air"].plot() +``` + +## Mixed-dimension datasets: ARCO-ERA5 + +When a Dataset has variables with differing dimensions (e.g. surface fields on +`(time, latitude, longitude)` and atmospheric fields on +`(time, level, latitude, longitude)`), `from_dataset` splits them into one +table per dimension group, registered together under a SQL schema named after +the first argument. [ARCO-ERA5][arco-era5] is a good example: 262 of its +variables are surface fields and 11 are atmospheric. + +Open a year of ARCO-ERA5 and let SQL `WHERE` clauses do the filtering — the +library prunes time partitions and pushes dimension-column filters down. Use +the `table_names` kwarg to give each dimension group a friendly name: + +> **Note:** reading from `gs://` requires `gcsfs` (`pip install gcsfs`). + +```python +import xarray as xr +import xarray_sql as xql + +# Open ARCO-ERA5 directly from GCS (anonymous read). +url = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3' +full = xr.open_zarr(url, chunks=None, storage_options={'token': 'anon'}) + +# A full year of hourly ERA5 — all 273 variables. No spatial slicing on the +# xarray side; SQL WHERE clauses below express the filters. `chunks={'time': 1}` +# aligns Dask chunks to native Zarr chunks of shape (1, 37, 721, 1440) so +# chunk reads from GCS happen concurrently. +# +# Heads up: 262 of those variables are surface and 11 are atmospheric. The +# library pushes column projection down, so SELECT only fetches what you ask +# for — but `SELECT * FROM era5.surface` would try to pull every variable +# across the year (terabytes from GCS). Always SELECT specific columns. +ds = full.sel(time='2020').chunk({'time': 1}) + +ctx = xql.XarrayContext() +ctx.from_dataset('era5', ds, table_names={ + ('time', 'latitude', 'longitude'): 'surface', + ('time', 'level', 'latitude', 'longitude'): 'atmosphere', +}) +# Registers two tables under a SQL schema named 'era5': 'surface' and 'atmosphere'. + +# Average 2m-temperature over the NYC area on the morning of 2020-01-01. +ctx.sql(''' + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 +''').to_pandas() + +# Average temperature per pressure level, globally — the standard +# atmospheric temperature profile. Scans ~230M rows. +ctx.sql(''' + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC -- surface (1000 hPa) first +''').to_pandas() +``` + +If you omit `table_names`, each table is named by joining its dimension names +with underscores, e.g. `era5.time_latitude_longitude` and +`era5.time_level_latitude_longitude`. + +## GOES satellite imagery (scalar variables) + +Real-world stores often mix gridded data with scalar (0-dimensional) metadata. +GOES satellite imagery, for example, pairs `(y, x)` image bands with dozens of +scalar variables such as `goes_imager_projection`. `from_dataset` groups all the +scalars into a single one-row table named `scalar`: + +```python +import fsspec +import xarray as xr +from xarray_sql import XarrayContext + +# A real GOES-16 ABI cloud-and-moisture file from NOAA's public bucket: +# (y, x) image bands alongside dozens of scalar metadata variables. +url = ( + 'https://noaa-goes16.s3.amazonaws.com/ABI-L2-MCMIPM/2024/001/00/' + 'OR_ABI-L2-MCMIPM1-M6_G16_s20240010000281_e20240010000350_c20240010000426.nc' +) +ds = xr.open_dataset(fsspec.open_local(f'simplecache::{url}')).chunk( + {'y': 250, 'x': 250} +) + +ctx = XarrayContext() +ctx.from_dataset('goes', ds) + +# The gridded bands and the scalar metadata are separate tables. +ctx.sql('SELECT COUNT(*) AS n FROM goes.y_x').to_pandas()['n'][0] # -> 250000 +ctx.sql('SELECT * FROM goes.scalar').to_pandas().shape # -> (1, 89) +``` + +Override the default name like any other group with `table_names={(): 'metadata'}`. + +A runnable version of the ERA5 example lives at +[`perf_tests/era5_temp_profile.py`](https://github.com/xqlsystems/xarray-sql/blob/main/perf_tests/era5_temp_profile.py). + +[arco-era5]: https://github.com/google-research/arco-era5 + + +## The same tables on DuckDB and Polars + +Every example above registers through an `XarrayContext`, but the tables are +not DataFusion-specific: `xql.register(con, name, ds)` attaches the same lazy, +pushdown-scanned table to a DuckDB connection, and +`pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))` serves Polars — same +splitting rules for mixed-dimension Datasets, same round-trip through +`xql.to_dataset(result, template=ds)`. See [Engines](engines.md) for the +support matrix and per-engine details. diff --git a/docs/geospatial.md b/docs/geospatial.md new file mode 100644 index 00000000..2ae86258 --- /dev/null +++ b/docs/geospatial.md @@ -0,0 +1,595 @@ +# Geospatial operations are relational operations + +A working hypothesis, and a slightly radical one: **the core operations of +geospatial and climate analysis — the ones we reach for an array library to +perform — are, underneath, relational operations.** Climatologies, anomalies, +zonal means, spectral indices, forecast skill, even regridding: each maps onto +ordinary SQL — `GROUP BY`, `JOIN`, window functions, `CASE`, and the occasional +scalar UDF. + +The array paradigm (NumPy, Xarray, Dask) is a wonderful *interface* for these +operations. But it is not the only one, and for a large and growing audience — +the people fluent in SQL rather than in `apply_ufunc` and rechunking — it is not +the most accessible one. [`xarray-sql`](index.md) lets you pose these +questions in SQL and answers them with a real query engine (DataFusion here; +[the same tables serve DuckDB and Polars](engines.md)). The +datasets are opened *lazily*, so a query against the whole archive reads only the +variable and the slice it actually needs. And because a gridded result is still +gridded data, every query here round-trips its answer straight back to an +`xarray.Dataset` — SQL in, an array out, ready to plot or save. + +This page makes the argument case by case. Every claim below is backed by a +runnable script in [`benchmarks/geospatial/`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial/) that +poses the operation in SQL and **asserts the answer matches an xarray/array +reference** to floating-point tolerance. The point is not that "SQL is faster"; +the point is that the SQL reads like the *definition* of the operation and +computes the same numbers — at ERA5's real 0.25° global resolution. + +## Where this list comes from + +The operations here aren't a set we hand-picked to suit SQL. They're taken from +[**Large Scale Geospatial Benchmarks**](https://github.com/coiled/benchmarks/discussions/1545) +(coiled/benchmarks #1545), a discussion [James Bourbeau](https://github.com/jrbourbeau) +opened in 2024 asking the +geospatial and climate community a pointed question: what are the *end-to-end +workflows* the Xarray/Dask ecosystem needs to handle smoothly at the +100-terabyte scale? The replies are a representative survey of what geoscience +actually runs — and this suite works through nearly all of it: + +| #1545 workflow | Covered by | +|----------------|------------| +| Remote-sensing indices (NDVI/NDWI/NDSI over Sentinel-2 or Landsat) | case 01 | +| Vectorized functions (`apply_ufunc`-style per-cell math) | case 01 | +| Climatology (average weather for a time of year/day, per location) | case 02 | +| Transformed Eulerian Mean (circulation diagnostics — zonal means and anomalies) | cases 03, 04 | +| Forecast evaluation (scoring forecasts against ground truth) | case 05 | +| Regridding and reprojection (resolution and CRS changes) | cases 07, 08, 09 | +| Spatial joins (large polygon-to-polygon joins) | *not covered* — a vector-data problem; the closest analogue here is the raster × vector join in case 06 | + +So the claim isn't that a few cherry-picked operations happen to be relational. +It's that an independent survey of the operations geoscience runs at scale, run +through SQL one by one, turns out to be — almost entirely — queries. + +## The mapping + +| Operation | The "array" framing | The relational reality | Script | +|-----------|---------------------|------------------------|--------| +| Spectral index (NDVI) | `apply_ufunc` over a raster | column arithmetic | [`01_ndvi.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/01_ndvi.py) | +| Climatology | rechunk → grouped reduction | `GROUP BY lat, lon, hour-of-day` | [`02_climatology.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/02_climatology.py) | +| Zonal mean | reduce over lon/time axes | `GROUP BY lat` | [`03_zonal_mean.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/03_zonal_mean.py) | +| Anomaly | grouped broadcast-subtract | climatology CTE self-`JOIN` | [`04_anomaly.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/04_anomaly.py) | +| Forecast skill (RMSE) | align valid/init/lead, reduce | forecast↔truth `JOIN` on `valid_time` | [`05_forecast_skill.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/05_forecast_skill.py) | +| Zonal stats over regions | rasterize polygons + mask | raster × vector range `JOIN` | [`06_zonal_vector.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/06_zonal_vector.py) | +| Reprojection | per-pixel CRS transform | scalar **UDF** (`ST_Transform`-style) | [`07_reproject_udf.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/07_reproject_udf.py) | +| Regridding | interpolation to a new grid | sparse-weight table `JOIN` | [`08_regrid_weights.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/08_regrid_weights.py) | +| Warp (reproject + resample) | CRS transform *and* interpolation | reproject **UDF** → weight-table `JOIN` | [`09_warp.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/09_warp.py) | + +## 1. A pixel-wise formula is a column expression + +NDVI is `(NIR − Red) / (NIR + Red)`, per pixel. The array idiom broadcasts a +ufunc over the raster. But "one output per pixel, computed from that pixel's +bands" is the definition of a SQL projection: + +```sql +SELECT x, y, (nir - red) / (nir + red) AS ndvi +FROM scene +ORDER BY y, x +``` + +Invalid pixels need no special handling: xarray decodes the band's `_FillValue` +to `NaN` on open, and `NaN` propagates through the arithmetic on both sides, so +the masking is free. + +[`01_ndvi.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/01_ndvi.py) runs this against a **real +Sentinel-2 L2A scene in Zarr** — discovered with `pystac-client` and opened the +canonical way with `xr.open_datatree` (ESA's EOPF sample service) — and matches +xarray's `apply_ufunc`-style result over a million pixels. + +## 2. A climatology is a `GROUP BY` over the cycle + +A climatology is the average value for each time-of-cycle at each location. In +the array world this is the canonical painful workload — load native chunks, +*rechunk* so all of time lands in one chunk, reduce, rechunk back. The +rechunking serves the array layout, not the question. The question is: + +```sql +SELECT latitude, longitude, date_part('hour', time) AS hour, + AVG("2m_temperature") +FROM era5 GROUP BY latitude, longitude, date_part('hour', time) +``` + +The grouping keys are the dimensions you keep; everything else is reduced. No +layout to reason about. [`02_climatology.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/02_climatology.py) +computes the **diurnal cycle** of ERA5 2m-temperature over a region — averaging +each cell by hour of day — and matches `da.groupby("time.hour").mean()` across +~500k cells. + +A **zonal mean** ([`03_zonal_mean.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/03_zonal_mean.py)) +is the same idea with fewer keys: the axes you "reduce over" are simply the +columns you don't `GROUP BY`. + +## 3. Broadcasting a normal back onto observations is a `JOIN` + +An anomaly subtracts each cell's climatological normal from every matching +observation. Xarray expresses the realignment with grouped broadcasting +(`ds.groupby("time.hour") - climatology`). That realignment — *attach each +cell's normal to every timestep that shares its key* — is a JOIN on the +grouping key: + +```sql +WITH clim AS ( + SELECT latitude, longitude, date_part('hour', time) AS hour, + AVG("2m_temperature") AS clim_t + FROM era5 GROUP BY latitude, longitude, date_part('hour', time) +) +SELECT a.time, a.latitude, a.longitude, + a."2m_temperature" - c.clim_t AS anomaly +FROM era5 a JOIN clim c + ON a.latitude = c.latitude AND a.longitude = c.longitude + AND date_part('hour', a.time) = c.hour +``` + +[`04_anomaly.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/04_anomaly.py) computes the +climatology once (the CTE) and joins it back to every observation. + +## 4. Forecast evaluation is a `JOIN` on valid time + aggregate + +This is the real workload of [WeatherBench 2](https://weatherbench2.readthedocs.io/): +scoring machine-learning weather models — **Pangu-Weather** and **GraphCast** — +against ERA5 ground truth. A forecast is indexed by *initialization time* and +*lead time* (`prediction_timedelta`); the truth is indexed by *valid time*. +Evaluation aligns them by `valid_time = init + lead` and reduces the error to +RMSE as a function of lead. + +That alignment is a relational JOIN, and `valid_time = init + lead` is just +timestamp + duration arithmetic the engine does natively: + +```sql +SELECT f.model, f.prediction_timedelta AS lead, + SQRT(AVG(POWER(f."2m_temperature" - e."2m_temperature", 2))) AS rmse +FROM forecasts f +JOIN era5 e + ON e.time = f.time + f.prediction_timedelta -- valid_time = init + lead + AND e.latitude = f.latitude + AND e.longitude = f.longitude +GROUP BY f.model, f.prediction_timedelta +``` + +Both models are stacked along a `model` dimension into one forecast table, so a +single query scores them together, grouped by the `model` column. The entire +evaluation — temporal alignment across three time axes, spatial matching, and the +score — is one JOIN and one aggregate. +[`05_forecast_skill.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/05_forecast_skill.py) runs it +for both models, matches an xarray reference, and reproduces the published result +that GraphCast edges out Pangu at every lead — the classic "error grows with +horizon" curve (≈0.3 K at 6 h rising to ≈2.5 K at 9 days). + +The result round-trips to a `pandas` table directly (`got.to_pandas()`), RMSE in +kelvin by lead time: + +``` +model graphcast pangu +lead (days) +0.25 0.296 0.336 +1.25 0.464 0.554 +2.25 0.608 0.734 +3.25 0.780 0.936 +4.25 0.988 1.191 +5.25 1.228 1.469 +6.25 1.470 1.747 +7.25 1.763 2.096 +8.25 2.092 2.489 +9.25 2.380 2.814 +``` + +## 5. Raster × vector zonal statistics is a range `JOIN` + +"Average the raster inside each region" is the canonical raster-meets-vector +task. The array idiom rasterizes each polygon to a mask and reduces under it. But +a region is a row in a table of bounds, and "pixel inside region" is a range +predicate — so zonal statistics is a JOIN: + +```sql +SELECT r.region, AVG(a."2m_temperature") - 273.15 AS avg_c +FROM era5.surface a JOIN regions r + ON a.latitude BETWEEN r.lat_min AND r.lat_max + AND a.longitude BETWEEN r.lon_min AND r.lon_max +WHERE a.time BETWEEN TIMESTAMP '2020-06-01' AND TIMESTAMP '2020-06-01 23:00:00' +GROUP BY r.region +``` + +This is the README's promise — *joining tabular data with raster data* — made +literal: the raster is the full ERA5 archive (the `WHERE` prunes it to a day), +the regions are a second SQL table, and the spatial relationship is an ordinary +`BETWEEN`. See [`06_zonal_vector.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/06_zonal_vector.py) +— it reports e.g. Sahara 33 °C vs Greenland −8 °C for a June day. (Rectangular +regions keep this simple; arbitrary polygons would follow the same shape, with a +point-in-polygon test in the join.) + +## 6. The hard cases: where a UDF fits, and where it doesn't + +Reprojection and regridding are the operations most wedded to the array +paradigm. They split cleanly along one line: **is the operation row-independent?** + +**Reprojection is.** Moving a coordinate from one CRS to another depends only on +that coordinate, so it is a *scalar function* — exactly what PostGIS and +DuckDB-spatial already ship as `ST_Transform`. xarray-sql ships it as an +optional geo extension (`pip install xarray-sql[geo]`): with pyproj +installed, every `XarrayContext` registers a PROJ-backed +`reproject(x, y, src_crs, dst_crs)` scalar UDF, so the CRS pair — any CRS +pyproj understands — is part of the query rather than baked into the function: + +```sql +SELECT x, y, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat +FROM grid +``` + +[`07_reproject_udf.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/07_reproject_udf.py) validates +this against **Earth Engine itself**: it opens a UTM grid through +[Xee](https://github.com/google/Xee) carrying `ee.Image.pixelLonLat()`, so EE's +own geodesy engine reports the true lon/lat of every pixel — an *independent* +reprojection reference, not PROJ-vs-PROJ. The SQL UDF and EE agree to sub-metre +precision. There is one practical gotcha — released pyproj mishandles threads +not created by Python, like DataFusion's runtime workers +([pyproj#1541](https://github.com/pyproj4/pyproj/pull/1541)), so the extension +runs all PROJ work on its own pool of Python threads, which also caches +transformers per thread, +keeping the UDF safe (and parallel) under DataFusion's concurrent partitions — +but the caveat that matters here is conceptual: reprojection +moves the coordinates without resampling the data onto a new grid — and *that* is +the next operation. + +**Regridding is not** row-independent: each output cell is a weighted blend of +several input cells. That is a *many-to-many* relationship — and a many-to-many +weighted blend is a sparse matrix–vector product, which is a `JOIN` against a +weight table plus a weighted `GROUP BY`: + +```sql +SELECT w.dst_id, SUM(s.value * w.weight) AS regridded +FROM weights w JOIN src s ON s.cell_id = w.src_id +GROUP BY w.dst_id +``` + +[`08_regrid_weights.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/08_regrid_weights.py) regrids +real **SRTM elevation** (Sierra Nevada terrain, opened from the Earth Engine +catalog through [Xee](https://github.com/google/Xee)) coarse → fine and matches +xarray's bilinear `.interp()` exactly. So regridding does not weaken the thesis — +it is the most relational operation of all. + +**A warp is just the two composed.** The full operation a GIS calls *warp* (GDAL +and rasterio's `reproject`) does both at once: change the CRS *and* resample onto +the new grid. [`09_warp.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/09_warp.py) writes it as the +two cases above run back to back — the 07 reproject UDF carries the target +lon/lat grid back into the source UTM space, arrays turn those reprojected points +into bilinear weights, and the 08 `JOIN` applies them: + +```sql +-- 1. reproject the target grid into source coordinates (the 07 UDF) +SELECT dst_lat, dst_lon, + reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['x'] AS sx, + reproject(dst_lon, dst_lat, 'EPSG:4326', 'EPSG:32610')['y'] AS sy +FROM target +-- 2. apply the bilinear weights built from those points (the 08 JOIN) +SELECT w.dst_lat AS lat, w.dst_lon AS lon, SUM(s.value * w.weight) AS warped +FROM weights w JOIN src s ON s.x = w.src_x AND s.y = w.src_y +GROUP BY w.dst_lat, w.dst_lon +``` + +It warps SRTM from a UTM grid onto a lon/lat grid and matches xarray's `.interp()` +at the reprojected points exactly, with Earth Engine's own lon/lat SRTM as a +second, cross-CRS sanity check (a loose match — EE resamples its native 30 m data, +we resample the 2 km source — so it is a corroboration, not the assertion). The +warp lands exactly where the split predicts: the row-independent half is a UDF, +the many-to-many half is a `JOIN`, and the only genuinely geometric step — turning +the reprojected points into weights — is the array work the next section is about. + +## GeoArrow point-geometry columns + +`register(..., geometry=("x", "y"))` derives a `geometry` point column +from two coordinate dims. With the default `"wkb"` encoding DuckDB +(spatial loaded) ingests it as a native `GEOMETRY` with the CRS +attached, so geometry predicates need no `ST_Point(x, y)` construction: + +```sql +SELECT avg(risk) FROM eri +WHERE y BETWEEN -29 AND -28 AND x BETWEEN -58 AND -57 -- prunes chunks + AND ST_Within(geometry, ST_GeomFromText('POLYGON (...)')) -- refines +``` + +**Always pair geometry predicates with bbox conjuncts on the coordinate +columns.** Engines do not push functions like `ST_Within` into the +scan, so a geometry-only predicate scans (and encodes) every chunk — +measured ~29x slower than the paired form on a 10M-row grid, where the +bbox prunes first and the exact polygon test is nearly free. +`xql.bbox_conjuncts(geom, x=..., y=...)` renders the conjuncts from any +geometry's envelope (shapely objects or plain +`(xmin, ymin, xmax, ymax)` tuples), with `pad=` for +`ST_DWithin`-style margins — so the idiom is one f-string. The full +reasoning lives in [Known issues & limitations](limitations.md#geometry-predicates-alone-cannot-prune). + +`geometry_encoding="point"` emits GeoArrow-native separated coordinates +instead (the struct children *are* the coordinate arrays): zero-parse +for GeoPandas 1.x (`GeoDataFrame.from_arrow`), lonboard, geoarrow-rs +and SedonaDB. DuckDB does not consume this encoding — pick per +destination. The CRS tag defaults to `OGC:CRS84`; pass +`geometry_crs=...` for anything else. + +## Where the array paradigm still earns its keep + +The boundary is **weight generation**. Applying a regridding is a join; +*computing* the weights — cell overlaps for conservative remapping, stencils and +spherical geometry for bilinear, the whole machinery of xESMF/ESMF — is genuinely +geometric work that arrays (and specialized libraries) do well. The relational +view does not replace that; it consumes its output. The division of labor is +clean and, we think, the right one: + +> **Arrays compute the geometry (the weights). SQL applies it (the join).** + +Likewise, the array libraries remain the right tool for building the inputs in +the first place — opening Zarr, decoding CF metadata, the numerics of generating +a weight matrix. `xarray-sql` sits downstream of all that as a query front-end: +once the data is openable as an `xarray.Dataset`, these everyday operations are +expressible — and accessible — as SQL. + +That is the qualitative boundary; the rest of this page puts numbers to it. The +**Results** below report what each operation costs in SQL versus the array +reference, **Analysis** explains *why* the relational form is slower and where the +time goes, and the **Conclusion** turns the whole thing into a when-to-use-which. + +## Running the suite + +```shell +python benchmarks/geospatial/02_climatology.py # inside the repo +uv run benchmarks/geospatial/02_climatology.py # standalone (PEP 723 deps) +``` + +Each script prints its SQL, runs the array reference, and asserts the two agree. +See [`benchmarks/geospatial/README.md`](https://github.com/xqlsystems/xarray-sql/tree/main/benchmarks/geospatial) for +the full list and dataset notes. + +## Results + +Correctness is the headline, but every case is also profiled. The numbers below +come from [`run_perf.sh`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/run_perf.sh) on a single Google +Compute Engine `e2-standard-8` (8 vCPU, 32 GB) in `us-central1` — in-region with the +ARCO-ERA5 and WeatherBench 2 buckets, so the cloud read is fast — with Earth Engine +reached from the same VM, so all nine cases share one machine and one release build. +Each case runs **once per fresh process**, with no warmup, repeated five times: the +SQL operation *and* the xarray reference each pay a **cold** read on every +measurement. + +Fairness here took some care, because the obvious trap is caching. A reference +that calls `.load()` caches its data *in place* on the very object the SQL table +also reads from, so a later read — even just running the reference after the SQL +query in the same process — could be served warm. We close that two ways. The one +case that loads shared objects (05, forecast skill) uses `.compute()` instead, +which returns a fresh array and leaves the inputs lazy, caching nothing; the other +references either reopen their data or recompute their reduction eagerly on every +read (`chunks=None` is NumPy, not Dask, so there is no graph to keep warm). And +`run_perf.sh` runs each case in a fresh process per repetition, ruling out any +carryover between reps. We verified the result directly: reading a window +repeatedly in one process stays flat, and running either side after the other +speeds up neither — the SQL query and the reference do not warm each other. + +| Case | Step | median (s) | stdev (s) | min (s) | max (s) | peak (MB) | +|---|---|--:|--:|--:|--:|--:| +| 01 · NDVI (per-pixel arithmetic) | SQL | 3.528 | 0.803 | 2.861 | 5.024 | 114.0 | +| | xarray reference | 0.304 | 0.104 | 0.282 | 0.496 | 42.0 | +| 02 · Climatology (`GROUP BY` lat, lon, hour) | SQL | 4.443 | 0.383 | 4.216 | 5.198 | 490.2 | +| | xarray reference | 1.867 | 0.106 | 1.844 | 2.053 | 43.7 | +| 03 · Zonal mean (`GROUP BY` latitude) | SQL | 2.406 | 0.122 | 2.333 | 2.631 | 236.9 | +| | xarray reference | 0.385 | 0.006 | 0.381 | 0.395 | 249.5 | +| 04 · Anomaly (climatology self-`JOIN`) | SQL | 7.027 | 0.123 | 6.950 | 7.239 | 511.5 | +| | xarray reference | 2.549 | 0.219 | 2.126 | 2.657 | 72.1 | +| 05 · Forecast skill (forecast↔truth `JOIN`) | SQL | 10.714 | 0.093 | 10.663 | 10.891 | 6.6 | +| | xarray reference | 0.248 | 0.013 | 0.220 | 0.254 | 2.2 | +| 06 · Zonal stats (raster × vector `JOIN`) | SQL | 4.308 | 0.053 | 4.299 | 4.401 | 509.1 | +| | xarray reference | 1.557 | 0.029 | 1.499 | 1.567 | 1262.1 | +| 07 · Reprojection (PROJ scalar UDF) | SQL | 0.029 | 0.003 | 0.024 | 0.031 | 0.3 | +| 08 · Regridding (weight-table `JOIN`) | SQL | 0.875 | 0.037 | 0.845 | 0.933 | 11.9 | +| | xarray reference | 0.850 | 0.658 | 0.809 | 2.310 | 13.3 | +| 09 · Warp (reproject UDF → regrid `JOIN`) | SQL | 0.281 | 0.038 | 0.250 | 0.353 | 0.8 | +| | xarray reference | 0.817 | 0.030 | 0.764 | 0.828 | 11.2 | + +Two patterns are visible before any analysis. SQL is slower on wall-clock wherever +a cloud read or a large relational expansion dominates — by ~2.5–6× on the +`GROUP BY` and `JOIN` cases against ARCO-ERA5, and ~43× on case 05, the smallest +grid but the biggest blow-up into rows — and its peak memory is highest on those +join/group-by cases (≈0.5 GB on 02, 04, 06). But the pattern is **not** universal. +On cases 08 and 09, where the interpolation *weights* are precomputed and SQL just +applies them, SQL is at parity with the array reference (08: 0.875 vs 0.850 s) or +**faster** (09: 0.281 vs 0.817 s — the reference pays for `pyproj` + `.interp`, +while SQL streams the prebuilt weight `JOIN`). The slow and the fast cases follow +from the same cause, which the next section pins down. (Case 01 reads Sentinel-2 +from Europe, the only non-US source, so its SQL time includes a cross-region read. +Cases 07–09 run against Earth Engine from the same VM: 07 times only the SQL +reproject transform, checked against Earth Engine's own `pixelLonLat`; 08 and 09 +read SRTM lazily on **both** the SQL and reference sides, so that comparison is +symmetric.) + +Case 05 is the suite's most hardware-sensitive number: its SQL time is CPU-bound on +the join and the (GIL-held) row production that feeds it, so it swings with the +machine — across three `e2-standard-8` runs it has measured ≈10.7 s, ≈12 s, and +≈23 s, while the read-bound *reference* stays near 0.25 s. So read the 05 ratio as +"the relational form costs real CPU here," not as a fixed multiplier. + +### The suite across engines and machine sizes + +The table above measures the DataFusion-native path. The same cases also run +through the suite's engine layer +([`_engines.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/_engines.py), selected per process +with `GEOBENCH_ENGINE`) as four engines: `datafusion` (the native table +provider — the harness builds the compiled module from the shipped crate on +each VM), `datafusion-arrow` (the pure-Python pyarrow-dataset path, through +`SessionContext.register_dataset(xql.arrow_dataset(ds))`), DuckDB, and +Polars. We ran the portable cases across all four on an `e2-standard-8` in +`us-central1`, in-region with the data, under the same protocol: **fresh +process per repetition, no warmup, five cold reps**, and every engine's +answer asserted against the xarray reference before its timing counts +([`engine_suite.py`](https://github.com/xqlsystems/xarray-sql/blob/main/benchmarks/geospatial/engine_suite.py) drives it). + +Scope, stated plainly. The `datafusion` column is the same code path as the +headline table, re-measured on this VM and day, so the two DataFusion columns +compare the native table provider and the pyarrow-dataset path on identical +hardware; DuckDB and Polars consume the same pyarrow pushdown datasets. +Polars executes the identical SQL via `polars.SQLContext`, with the query's +window bounds also applied as native scan-level expressions (its SQL +`TIMESTAMP` literals compile to `strptime` casts that never reach the pyarrow +scanner as filters); its SQL dialect cannot express case 06's range `JOIN` (a +`BETWEEN` join constraint), recorded as unsupported rather than worked +around. Cases 07 and 09 build DataFusion scalar UDFs (n/a on the other +engines), case 08 is Earth-Engine-gated, and 07–09 skip on these VMs (no +Earth Engine auth) — the skip reasons ride along in the results. + +The same grid was also measured on `e2-standard-16` and `e2-standard-32` in +the same run, with no practical difference: every case is dominated by a +single-stream cold cloud read plus a mostly single-threaded row pipeline, so +extra vCPUs buy nothing, and the spread across sizes is shared-core `e2` and +network variance, not engine behavior (the 16-vCPU VM measured *slower* than +the 8-vCPU one on most cells in this run). The full per-size tables live in +the `xarray-sql-notes` repository (`engine-matrix-results.md`). + +Software: CPython 3.12.8, duckdb 1.5.5, polars 1.42.1, datafusion 54.0.0, +pyarrow 25.0.0, xarray 2026.7.0, Linux (glibc 2.41). Medians of 5 cold reps; +peak is the Python-allocator peak per process. + +| Case | DataFusion (native) | DataFusion (pyarrow) | DuckDB | Polars | xarray reference | +|---|--:|--:|--:|--:|--:| +| 01 · NDVI | 4.308 s (114 MB) | 4.380 s (97 MB) | 5.325 s (106 MB) | 5.725 s (106 MB) | 0.444 s (42 MB) | +| 02 · Climatology | 7.235 s (1116 MB) | 8.746 s (1142 MB) | 4.622 s (627 MB) | 4.936 s (637 MB) | 2.367 s (44 MB) | +| 03 · Zonal mean | 4.763 s (406 MB) | 3.614 s (403 MB) | 2.958 s (403 MB) | 3.705 s (413 MB) | 0.829 s (250 MB) | +| 04 · Anomaly | 10.416 s (1117 MB) | 16.262 s (3003 MB) | 9.817 s (627 MB) | 13.837 s (936 MB) | 4.410 s (76 MB) | +| 05 · Forecast skill | 1.791 s (170 MB) | 1.814 s (175 MB) | 1.797 s (170 MB) | 2.405 s (184 MB) | 0.247 s (2 MB) | +| 06 · Zonal stats | 2.390 s (515 MB) | 6.131 s (513 MB) | 9.112 s (503 MB) | unsupported (range `JOIN`) | 1.813 s (1262 MB) | + +Cases 07–09 could not run on these VMs (no Earth Engine auth), so their rows +come from the headline run instead: the original `e2-standard-8` GCE VM with +Earth Engine access, DataFusion **native** path (07 and 09 are DataFusion +scalar UDFs, n/a on the other engines): + +| Case | DataFusion (native path) | xarray reference | +|---|--:|--:| +| 07 · Reprojection (PROJ scalar UDF) | 0.029 s (0.3 MB) | — | +| 08 · Regridding (weight-table `JOIN`) | 0.875 s (11.9 MB) | 0.850 s (13.3 MB) | +| 09 · Warp (reproject UDF → regrid `JOIN`) | 0.281 s (0.8 MB) | 0.817 s (11.2 MB) | + +The engine story, in three observations. **Native vs pyarrow on the same +engine:** the native table provider wins where rows are consumed in bulk — +case 06's range `JOIN` (2.4 vs 6.1 s) and the anomaly self-`JOIN` 04 (10.4 vs +16.3 s, at a third of the peak memory), ~1.2× on climatology 02 — while on +the read-bound cases (01, 05) the two are at parity, and on the plain zonal +mean 03 the pyarrow path is the faster one (3.6 vs 4.8 s). **Across +engines**, DuckDB is the fastest consumer of the shared pushdown scan on the +plain group-bys (02, 03) and narrowly beats native DataFusion on 04 (9.8 vs +10.4 s), while native DataFusion leads the join-heavy 06 outright; Polars +stays close to DuckDB on 02 and falls back on 04. The spread between engines is much smaller +on cases whose cost is the read itself (01, 05), which is the same lesson as +the headline table: the paradigm and the I/O set the floor, the engine sets +the constant. And a benchmark side-effect worth keeping: streaming case 05's +full window through the pyarrow protocol surfaced a real library bug — +`pa.array` returns a `ChunkedArray` for a large string dimension coordinate, +which the pivot's fast path passed straight into `RecordBatch.from_arrays`; +fixed, and pinned by a regression test in `tests/test_df.py` +(`test_iter_record_batches_large_string_dim_coord`). + +## Analysis: how a relational operation spends its time + +Why is SQL slower, and where does the time actually go? Profiling case 05 — the +forecast-skill `JOIN`, the widest gap — with `cProfile`, run cold then warm so that +`cold − warm` isolates the cloud read and the warm floor is ≈pure compute, +decomposes it cleanly. (These are single-process numbers from a laptop with a slow +cross-region read — a *different* machine from the in-region table above, on +purpose: it puts both sides' reads on equal, slow footing so the compute gap shows +through. The absolute seconds therefore differ from the table; the decomposition, +not the totals, is the point.) + +| | read (I/O) | compute | total (cold) | +|---|--:|--:|--:| +| SQL | ~0.95 s | **~0.71 s** | ~1.66 s | +| xarray reference | ~0.79 s | **~0.024 s** | ~0.81 s | + +The read is comparable on both sides — both open the same Zarr store cold. **The +gap is compute, and it is about 30×.** The SQL path explodes the 64×32×20×2 grid +into Arrow rows, runs a hash `JOIN` to align each forecast row with its truth row +on `(valid_time, latitude, longitude)`, aggregates, and streams the result batches +back. The array reference does the identical math as a handful of vectorized NumPy +reductions over contiguous buffers. Row materialization + hashing + the join probe +is simply heavier than dense arithmetic on a regular grid — and it is the same +work that inflates SQL's peak memory in the Results table: the join and group-by +cases hold the grid as rows. + +`cProfile` is unambiguous about *where* the SQL time sits. Essentially all of it is +in pulling record batches from the DataFusion execution stream; the SQL→xarray +round-trip that turns the query result back into a gridded `Dataset` +(`to_dataset`) is **sub-millisecond — under 1% of the query.** So the cost is the +relational engine doing row-oriented work, not the array reconstruction. The +paradigm itself is the price, paid where the relational algebra runs. + +This explains the shape of the whole table. Case 05 stands alone at ~43× not +because its join is exotic but because its *reference* is nearly free — a 64×32×20×2 +grid reduces in-memory in a quarter-second — while SQL still has to explode that +grid into rows and hash-join them; a huge ratio over a tiny denominator. The +ARCO-ERA5 cases (02, 03, 04, 06) instead cluster at ~2.5–6×, because there a large +cloud read is a cost *both* sides pay, compressing the ratio. And cases 08 and 09 +invert it entirely: once the geometry — the interpolation weights — is precomputed, +applying it is a `JOIN` that streams about as fast as (or faster than) the array +reference's `pyproj`/`.interp`. The relational *overhead* is constant; the *ratio* +you observe depends on how much non-relational work (the cloud read, the weight +generation) sits on the other side of the comparison. And it shifts with hardware +too: SQL is CPU-bound on the join while the array reference is read-bound, so the +two are gated by different resources. On a fast laptop with a slow cross-region +read the gap nearly closes; on an in-region VM with modest cores it widens. The +underlying cause is +constant — materialize rows, hash-join, aggregate — but which resource you are +waiting on is not. + +## Conclusion + +None of this is an argument that SQL is *faster*. On a single node, for the +reduction-shaped operations, it is not — it pays a real per-operation overhead to +express an array reduction as relational algebra. (The exceptions, cases 08 and 09, +prove the rule: once the array work — generating the weights — is already done, the +relational half that remains is competitive, because there is no dense reduction +left for arrays to win.) The honest tradeoff is about which property you are +optimizing for. + +**Reach for the array paradigm when the work is dense and grid-aligned.** Per-pixel +formulas, stencils, convolutions, FFTs, linear algebra — anything that stays in +contiguous typed buffers and treats the chunk grid as its unit of parallelism. The +array model has the lowest overhead here, and the lead is structural, not +incidental: there are no rows to materialize and nothing to shuffle. NDVI (case 01) +is the tell — column arithmetic expresses cleanly in SQL, but the array side is +~10× faster (part of which is case 01's cross-region read; the rest is that +per-pixel math is exactly what arrays are for). + +**Reach for SQL when the work is relationally shaped, or the audience is.** Joins, +group-bys, alignment across data with different indexes (case 05's three time +axes), raster-meets-vector predicates (case 06) — these are awkward to express and +to reason about as array operations, and they are the native vocabulary of a query +engine. The overhead buys you an operation that reads like its own definition, that +prunes its own reads (a query against the whole ERA5 archive touches only the +variable and window it asks for), and that is accessible to the large audience +fluent in SQL rather than in `apply_ufunc` and rechunking. + +There is also a payoff this single-node benchmark cannot show. The same overhead — +row materialization and a hash join — is what makes the operation a *first-class +citizen of a distributed query engine.* Cost-based query optimization (join +reordering, choosing broadcast vs. shuffle joins, predicate pushdown), mature +partitioned shuffle and spill-to-disk, partitioning driven by the query rather than +locked to a physical chunk grid — these are exactly the capabilities the +array/Dask ecosystem struggles to provide for join- and group-by-heavy workloads +at scale, and exactly what the relational framing puts within reach. Whether the +constant-factor overhead is worth paying flips as the data grows and the bottleneck +moves from per-element compute to data movement. `xarray-sql` is single-node today, +so that is a direction rather than a result — but it is the latent reason the +thesis matters beyond expressibility. + +So the division of labor from the section above generalizes past regridding. Arrays +own the dense numerics and the geometry; SQL owns the relational shape — the joins, +the alignment, the aggregation — and, increasingly, the path to running them at +scale. The point of this suite is not to crown a winner but to show that the line +between the two is exactly where the operation is dense versus where it is +relational, and that for a surprising share of geoscience, the operation is +relational. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 00000000..612c7a5e --- /dev/null +++ b/docs/index.md @@ -0,0 +1 @@ +--8<-- "README.md" diff --git a/docs/limitations.md b/docs/limitations.md new file mode 100644 index 00000000..8aab000c --- /dev/null +++ b/docs/limitations.md @@ -0,0 +1,157 @@ +# Known issues & limitations + +What does not work, why, and what to do instead. How the machinery +*works* — scan pipeline, tuning, cost model — lives in +[Engines](engines.md) and the [performance guide](performance.md); this +page is only the sharp edges. Everything here is pinned by tests. + +## Engine-specific issues + +Pick your engine: + +=== "DataFusion" + + No engine-specific known issues. DataFusion is the deepest + integration (native table provider, chunked round-trip via + re-execution); the constraints below apply as everywhere. + +=== "DuckDB" + + **Re-executing relations from worker threads deadlocks.** + + - *Symptom:* a chunked round-trip (`chunks=`) of a DuckDB relation + would hang intermittently (~50% of runs) when dask workers + re-execute the relation, whenever the query scans a + Python-backed table. + - *Scope:* duckdb-python 1.4–1.5 with CPython 3.12 (observed on + macOS); unaffected by `SET threads=1`, connection-level + serialization, or thread-pool pre-warming. The identical + topology through Polars never hangs. The deadlock is in the + interpreter/engine thread-state interaction, not in xarray-sql. + - *What the library does:* `chunks=` on a DuckDB relation raises + `NotImplementedError` immediately rather than hanging. + `spill=True` provides the chunked path without ever + re-executing: the result is streamed once (bounded memory, on + the handle's dedicated engine thread) into a temporary Parquet + file that windows re-execute against. The eager round-trip is + unaffected. + + **Derived relations break under concurrent materialization.** + + - *Symptom:* materializing two relations derived from the same + base concurrently raises `InvalidInputException` (they share + pending-query state upstream). + - *What the library does:* the round-trip handle serializes every + engine call on one dedicated thread. Nothing to do on your + side — documented so the serialization is not mistaken for a + missing optimization. + + **GeoArrow-native points are not consumed.** + + - *Symptom:* a `geometry` column registered with + `geometry_encoding="point"` binds as a plain struct; `ST_*` + functions reject it. + - *What to do:* use the default `"wkb"` encoding for DuckDB — it + binds as a native `GEOMETRY` with the CRS attached. See + [Geospatial in SQL](geospatial.md#geoarrow-point-geometry-columns). + +=== "Polars" + + **Float `is_in` literals lose precision.** + + - *Symptom:* `is_in` with **float** literals can silently match + nothing (reproducible without xarray-sql; integer and timestamp + value sets are unaffected). + - *What to do:* prefer `is_between` for float coordinates in your + own queries. + - *What the library does:* the lazy round-trip's window queries + render float value lists as degenerate ranges internally, so + reconstruction is immune. + + **No geometry types.** + + - *Symptom:* a registered `geometry` column arrives as plain + binary (WKB) or a plain struct; there are no `ST_*` functions. + - *What to do:* filter on the coordinate columns instead, and do + geometry work in DuckDB, DataFusion, or GeoPandas. + + **Single-threaded source pull.** + + Polars pulls the scan sequentially; source-side parallelism comes + from the adapter's prefetch pool (`prefetch`, `prefetch_bytes`), + not from the consumer. Not a bug — worth knowing when sizing + scans. + +## Constraints in any engine + +These follow from the data model — no engine or configuration avoids +them. + +### Geometry predicates alone cannot prune + +Engines never push function calls (`ST_Within`, casts, arithmetic) into +a scan — only plain column-vs-constant comparisons, `IN`, `IS NULL`, +and boolean combinations. A geometry-only `WHERE` therefore scans and +encodes every chunk (measured ~29x slower than the paired form on a +10M-row grid). Pair every geometry predicate with range conjuncts on +the coordinate columns; [`bbox_conjuncts`][xarray_sql.bbox_conjuncts] +renders them from the geometry's envelope. See +[Geospatial in SQL](geospatial.md#geoarrow-point-geometry-columns). + +### Filters on data variables always scan + +Chunk pruning and arithmetic counting rest on per-chunk *coordinate* +ranges. A predicate on a data variable (`t2m > 300`) carries no such +guarantee: every surviving chunk is scanned, and the filter is applied +row-exactly. No configuration changes this; it is what the data model +can prove. + +### NaN coordinates disable pruning for their chunks + +A NaN/NaT anywhere in a chunk's coordinate span poisons its min/max +guarantee, so that chunk is kept for every predicate. This is the +correct trade: a range that pretended to cover NaN would let engines +whose NaN ordering differs (DuckDB sorts NaN greatest) silently lose +rows. Chunks without NaN prune normally. + +### String, object, and cftime dimensions never prune + +Chunk guarantees are built for numeric and datetime coordinates only; +predicates on other dimension types conservatively scan every chunk +(row-exactly, as always). + +### Scan-path pruning is per-dimension + +On the scan path, each dimension's surviving chunks are computed +independently and combined as a product, so a predicate pairing +*specific* ranges across dims — `(t < a AND lat < b) OR (t > c AND +lat > d)` — also reads the cross combinations (sound, conservative; +per-dim indexes are what keep million-chunk axes cheap). `count_rows` +refines the crosses away with cross-dimension bucket analysis; ordinary +scans accept the extra reads. + +### Sparse results can explode the dense grid + +The eager round-trip reconstructs the coordinate-product grid: a +diagonal of n rows becomes an n×n array that can dwarf its Arrow +payload. `max_result_bytes=` raises cleanly at both danger points +(stream collection and dense allocation); it is opt-in and unlimited +by default. + +### One-shot Arrow streams cannot re-execute + +A materialized table or bare C-stream has no query behind it, so the +re-execution form of `chunks=` cannot serve it — `spill=True` (one +pass to a temporary Parquet file) is the chunked path for these. + +### Mixed-dimension datasets split into one DuckDB table per dim group + +DuckDB registration has no schema namespace, so variables with +different dims land in suffixed tables (`_`), sharing one +set of coordinate reads. DataFusion registers the same layout as +`name.group` tables inside one schema. + +### Pointwise indexers on lazy round-trip arrays are slower + +Vectorized (pointwise) selection goes through xarray's +outer-then-gather fallback — correct, but slower than slice windows. diff --git a/docs/performance.md b/docs/performance.md new file mode 100644 index 00000000..349589d7 --- /dev/null +++ b/docs/performance.md @@ -0,0 +1,245 @@ +# Performance guide + +How to get engine-limited speed out of registered xarray tables. Every +number below was measured on real cloud rasters (billions of pixels); +your mileage scales with network and core count, but the *ratios* are +structural. + +## How a scan decides what to read + +Every engine query over a registered table flows through one pipeline; +each tuning knob on this page acts on one of its stages: + +```mermaid +flowchart TB + Q["engine calls scanner(columns, filter)"] --> P["prune chunks
per-dim coordinate ranges +
Arrow guarantee simplification"] + Q --> J["project
only referenced variables are read"] + P --> C["coalesce (opt-in)
merge consecutive surviving chunks
into single reads"] + C --> F["prefetch pool
bounded by prefetch (threads)
and prefetch_bytes (memory)"] + J --> F + F --> X["exact filter
the pushed expression is applied
row-exactly — pruning is only
ever an optimization"] + X --> B["Arrow batches → engine"] +``` + +Two invariants hold everywhere: pruning never decides correctness (the +exact expression is always applied — engines delete pushed conjuncts +from their own plans), and only what reaches `scanner()` can prune +(engines push plain comparisons, never function calls). + +Everything on this page up to [Per-engine notes](#per-engine-notes) +applies whichever engine you query with. + +## Make the source read in parallel + +The single biggest lever is usually the reader, not the engine. + +**GeoTIFF / rioxarray**: `rioxarray.open_rasterio` serializes GDAL tile +reads behind a lock by default, capping every scan at single-stream +speed no matter how many threads the adapter runs. On GDAL ≥ 3.11 use +the natively thread-safe LIBERTIFF driver; on older GDAL pass +`lock=False`: + +```python +da = rioxarray.open_rasterio( + url, chunks={"x": 2048, "y": 2048}, + driver="LIBERTIFF", # GDAL >= 3.11; else keep lock=False only + lock=False, +) +``` + +Measured on a 9-billion-pixel public cloud GeoTIFF, full-table +aggregation: default open **277 s** → `lock=False` **43 s** → +LIBERTIFF + `GDAL_NUM_THREADS=ALL_CPUS` **24 s**. With parallel reads, +remote (`/vsicurl/`) matched a local copy of the same file — the +network was never the bottleneck, the lock was. + +Remote-read environment preset worth exporting for `/vsicurl/` sources: + +```python +os.environ.update( + GDAL_NUM_THREADS="ALL_CPUS", + GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", + VSI_CACHE="TRUE", +) +``` + +**Zarr**: zarr-python 3's async store defaults to only 10 concurrent +requests; raise it before opening remote stores: + +```python +zarr.config.set({"async.concurrency": 128}) +``` + +On a moderately sized windowed query (~40 chunks of 4096² uint8 per +variable, GCS) this was a modest gain (4.2 s → 3.7 s); it matters more +as chunk counts grow and chunks shrink. The obstore-backed +`zarr.storage.ObjectStore` is worth benchmarking for high-concurrency +workloads, but was not faster at this scale in our tests — measure +before switching. + +## Choose chunk sizes for the scan, not just the store + +Every chunk costs one prefetch task, one pivot call, and one shadow +fragment. Aim for **1–8 M rows per chunk** (e.g. 2048²–4096² pixels for +2-D grids). The same 10 M-row scan ran 1.7× faster in 4 chunks than in +20. Axes with hundreds of thousands of chunks still prune in +milliseconds (the shadow index is bucketed), but scanning them pays +per-chunk overhead. + +## Tune the adapter knobs + +```python +xql.register(con, "t", ds, prefetch=12, batch_size=262_144) +``` + +- `prefetch`: chunk loads kept in flight ahead of the engine. The + default (4) saturates local CPU work; raise to 8–12 for remote + sources where latency dominates. Memory scales with + `prefetch × pivoted chunk size`. +- `batch_size`: rows per Arrow batch. The default (64 Ki) is fine; + values between 64 Ki and 1 Mi measured within a few percent of each + other. + +## The memory contract + +Peak scan memory is bounded by `prefetch × pivoted-block-size` plus the +engine's own aggregation state — it does not grow with the amount of +data scanned. Measured on ARCO-ERA5 over anonymous GCS: a one-month +full-globe aggregation (772M rows) peaks at the same resident set size +as the one-week scan (174M rows), ~0.75 GB with the defaults. + +`prefetch_bytes` caps *estimated bytes* in flight instead of block +count — set it when `coalesce_rows` makes blocks large or ragged. +The block size is the source chunk size unless `coalesce_rows` is set, +in which case in-flight units are merged blocks: raising +`coalesce_rows` buys fewer round-trips at proportionally higher peak +memory (`prefetch=16, coalesce_rows=8_000_000` peaked at ~1.2 GB on the +same scan while cutting wall time ~1.5-2x). Size the two together. + +`count(*)`-shaped queries never pay scan memory at all: unfiltered +counts are pure chunk arithmetic, and filtered counts scan only the +boundary chunks the filter cannot prove — at any filter breadth; see +[What counting costs](#what-counting-costs). + +## Let pushdown do its job + +Selective queries are fast *because of their predicates*: bounding-box +`WHERE` clauses on dimension columns prune to intersecting chunks, and +only the variables a query references are read. Corollaries: + +- Prefer explicit column lists over `SELECT *` on wide datasets. +- Spatial functions (`ST_Within`, ...) are not pushed down — pair them + with a bounding-box predicate that is: the box prunes, the geometry + refines. +- A query with no `WHERE` on dimension columns is a full scan on any + engine; that's physics, not a missing optimization. + +## What counting costs + +`count(*)` never pays scan memory, and usually no I/O either: + +```mermaid +flowchart TB + C["count_rows(filter)"] --> U{"filter?"} + U -- none --> A["pure arithmetic
0 reads"] + U -- "coordinate ranges" --> H["hierarchical strictness:
bucket-products proven or pruned
whole; only mixed cells recurse"] + H --> E["boundary chunks scanned exactly
(usually 0-2 per range edge,
at any axis size)"] + U -- "data variables" --> S["every surviving chunk scanned
(values carry no coordinate
guarantee — see Known issues)"] +``` + +Coordinate-range counts stay arithmetic at any breadth (a +near-universal filter over a million single-row chunks counts with +zero reads), and the strictness pass applies cross-dimension +information, so paired-range predicates count without reading the +cross combinations. + +## Stop re-scanning: cache derived tables + +Registered tables are virtual — every query re-streams the source. +Statistics you ask repeatedly should pay the scan once: create a native +table from your query, sorted by the coordinate columns so the engine's +storage compresses the repetitive coordinates (DuckDB picks ALP/RLE on +sorted runs) and zone maps prune range predicates. + +```sql +CREATE OR REPLACE TABLE grid_cube AS +SELECT FLOOR(y) AS lat, FLOOR(x) AS lon, klass, COUNT(*) AS n +FROM grid GROUP BY 1, 2, 3 +ORDER BY lat, lon; + +SELECT * FROM grid_cube WHERE lat = -32; -- native speed +``` + +One engine quirk to know: on DataFusion, DDL is a lazy plan — collect +it or nothing happens: + +```python +ctx.sql("CREATE OR REPLACE TABLE grid_cube AS ...").collect() +``` + +## Round-trip faster with ORDER BY + +Results that arrive **grid-ordered** — sorted by the dimension columns, +outermost first — reconstruct with a single reshape; unordered results +(DuckDB's parallel scans return chunk order, not grid order) pay a +per-row positional scatter instead, measured ~2x slower on large +windows. When you will round-trip a large result, add +`ORDER BY ` to the query. + +And if what you want is a raw sub-array of a registered Dataset rather +than a relational answer, plain `ds.sel(...)` is the direct path — SQL +adds value when the question is relational. + +## Per-engine notes + +=== "DataFusion" + + **Two registration paths.** `XarrayContext.from_dataset` uses the + native Rust table provider — partition-parallel, with `chunks=` + controlling partition granularity; the `prefetch`/`coalesce_rows` + scanner knobs on this page apply to the *pyarrow-dataset* path + (`ctx.register_dataset(xql.arrow_dataset(ds))`), not to the native + provider. + + **DDL/DML is lazy.** `CREATE TABLE`/`INSERT` statements are plans — + `.collect()` them or nothing executes (the caching recipe above + shows this). + +=== "DuckDB" + + **Connections and threads.** Registered Python objects are + connection-local: `con.cursor()` does not inherit them, and one + connection's result slot is not thread-safe. For multithreaded + querying, give each thread its own cursor and register the *same* + dataset object on it: + + ```python + dataset = xql.arrow_dataset(ds) + def worker(): + cur = con.cursor() + cur.register("t", dataset) # cheap; shares the pruning index + ... + ``` + + The dataset object itself is safe to share across threads + (verified under concurrent query load). + + **Row order.** DuckDB's parallel scans return results in chunk + order, not grid order — add `ORDER BY ` before round-tripping + large results (see [above](#round-trip-faster-with-order-by)). + + **Geometry.** Register with the default `"wkb"` encoding; pair + `ST_*` predicates with bbox conjuncts so pruning still applies. + +=== "Polars" + + **Parallelism.** Polars pulls the scan single-threaded; source-side + parallelism comes entirely from the adapter's `prefetch` / + `prefetch_bytes`, so tune those rather than Polars settings. + + **Batch sizing.** `scan_pyarrow_dataset` passes its `batch_size` + through to the scanner (honored), so Polars morsel sizing works as + documented on their side. + + **Large results.** Collect with `engine="streaming"` to keep + memory bounded; the lazy round-trip's windows already do this. diff --git a/docs/reference/xarray_sql.md b/docs/reference/xarray_sql.md new file mode 100644 index 00000000..bef077cd --- /dev/null +++ b/docs/reference/xarray_sql.md @@ -0,0 +1,8 @@ +# xarray-sql + +::: xarray_sql + options: + show_root_heading: true + show_source: false + members: true + show_submodules: true diff --git a/perf_tests/compute_air.py b/perf_tests/compute_air.py index 476cefe5..2413370d 100755 --- a/perf_tests/compute_air.py +++ b/perf_tests/compute_air.py @@ -1,13 +1,13 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - air = air.chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + air = air.chunk(chunks) - df = qr.read_xarray(air).compute() + df = xql.read_xarray(air).read_pandas() - print(len(df)) + print(len(df)) diff --git a/perf_tests/era5_temp_profile.py b/perf_tests/era5_temp_profile.py new file mode 100644 index 00000000..55cb8be9 --- /dev/null +++ b/perf_tests/era5_temp_profile.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Surface and global-atmospheric temperatures on 2020-01-01, in SQL. + +Three queries against ARCO-ERA5 on the morning of January 1, 2020: + + * **Surface (local).** Average 2m-temperature over a small grid covering the + New York City area for the first six hours. + * **Atmosphere (global).** Average temperature per pressure level, computed + over the entire planet for the same six hours — a classic atmospheric + temperature profile (surface around 1000 hPa is warmest, tropopause near + 100 hPa is coldest). + * **Surface (global, gridded).** Average 2m-temperature per (lat, lon) cell + for the same six hours, returned as an xarray Dataset. + +All filters live in SQL: the dataset is opened with no time or spatial +slicing on the xarray side. The library's table provider prunes time +partitions for ``WHERE time …`` filters, and pushes ``WHERE +latitude/longitude …`` down to dimension columns. + +ARCO-ERA5's atmospheric variables are stored in native Zarr chunks of shape +``(1, 37, 721, 1440)`` — about 150 MB per hour. ``chunks=dict(time=6)`` groups +six native chunks per DataFusion partition: large enough to keep partition +count (and registration time) low, small enough that a 6-hour WHERE clause +hits a single partition with no wasted I/O. + +The Zarr is read anonymously from the public GCS bucket — no auth required. +""" + +import time + +import xarray as xr + +import xarray_sql as xql + + +URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" + + +def main() -> None: + # Open the full ARCO-ERA5 archive — all 273 variables since 1940. No + # time or spatial slicing on the xarray side; SQL WHERE clauses below + # express the filters. Turning dask off (chunks=None) skips task-graph + # construction at open time. + ds = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) + print( + "ARCO-ERA5 opened: " + f"{ds.sizes['time']:,} hourly time steps, " + f"{len(ds.data_vars)} variables (no pre-slicing)." + ) + + # Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The + # library pushes column projection down to Zarr, so SELECT only fetches + # what you ask for — but `SELECT * FROM era5.surface` would try to pull + # every variable across the archive (terabytes from GCS). + # ---> Always SELECT specific columns. <--- + ctx = xql.XarrayContext() + t0 = time.perf_counter() + # Make sure to pass `chunks`! + ctx.from_dataset( + "era5", + ds, + chunks=dict(time=6), + table_names={ + ("time", "latitude", "longitude"): "surface", + ("time", "level", "latitude", "longitude"): "atmosphere", + }, + ) + print(f"Registration: {time.perf_counter() - t0:.2f}s") + ctx.sql("SELECT 1").to_pandas() # warm the planner + + print("\nAverage 2m-temperature over NYC, 2020-01-01 00:00-05:00 UTC (°C):") + t0 = time.perf_counter() + surface = ctx.sql( + """ + SELECT AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + AND latitude BETWEEN 39 AND 40 + AND longitude BETWEEN 286 AND 287 -- ERA5 uses 0-360 longitudes + """ + ).to_pandas() + print(surface) + print(f" ({time.perf_counter() - t0:.2f}s)") + + print( + "\nAverage temperature per pressure level, globally, " + "2020-01-01 00:00-05:00 UTC (°C):" + ) + t0 = time.perf_counter() + profile = ctx.sql( + """ + SELECT level, AVG(temperature) - 273.15 AS avg_c + FROM era5.atmosphere + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY level + ORDER BY level DESC -- surface (1000 hPa) first, top of atmosphere last + """ + ).to_pandas() + print(profile.to_string(index=False)) + print(f" ({time.perf_counter() - t0:.2f}s)") + + print( + "\nAverage 2m-temperature per (lat, lon) cell, globally, " + "2020-01-01 00:00-05:00 UTC (°C):" + ) + t0 = time.perf_counter() + gridded = ctx.sql( + """ + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude + """ + ).to_dataset(dims=["latitude", "longitude"], template=ds) + print(gridded) + print(f" ({time.perf_counter() - t0:.2f}s)") + + +if __name__ == "__main__": + main() diff --git a/perf_tests/groupby_air.py b/perf_tests/groupby_air.py index c9f4d8ec..61006aeb 100755 --- a/perf_tests/groupby_air.py +++ b/perf_tests/groupby_air.py @@ -1,25 +1,25 @@ #!/usr/bin/env python3 +from datafusion import SessionContext import xarray as xr -import xarray_sql as qr -from dask_sql import Context +import xarray_sql as xql if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240, "lat": 5, "lon": 7} - air = air.chunk(chunks) - air_small = air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240, "lat": 5, "lon": 7} + air = air.chunk(chunks) + air_small = air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(chunks) - df = qr.read_xarray(air_small) + df = xql.read_xarray_table(air_small) - c = Context() - c.create_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = c.sql( - """ + query = ctx.sql( + """ SELECT "lat", "lon", SUM("air") as air_total FROM @@ -27,12 +27,11 @@ GROUP BY "lat", "lon" """ - ) + ) - result = query.compute() + result = query.collect() - expected = air_small.dims["lat"] * air_small.dims["lon"] - assert ( - len(result) == expected - ), f"Length must be {expected}, but was {len(result)}." - print(expected) + expected = air_small.sizes["lat"] * air_small.sizes["lon"] + actual = sum(len(batch) for batch in result) + assert actual == expected, f"Length must be {expected}, but was {actual}." + print(expected) diff --git a/perf_tests/groupby_air_full.py b/perf_tests/groupby_air_full.py index e54ad139..2f2a6c1f 100755 --- a/perf_tests/groupby_air_full.py +++ b/perf_tests/groupby_air_full.py @@ -1,22 +1,21 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr -from dask_sql import Context - +import xarray_sql as xql +from datafusion import SessionContext if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - air = air.chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + air = air.chunk(chunks) - df = qr.read_xarray(air) + df = xql.read_xarray_table(air) - c = Context() - c.create_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = c.sql( - """ + query = ctx.sql( + """ SELECT "lat", "lon", SUM("air") as air_total FROM @@ -24,12 +23,12 @@ GROUP BY "lat", "lon" """ - ) + ) + + result = query.collect() - result = query.compute() + expected = air.sizes["lat"] * air.sizes["lon"] + actual = sum(len(batch) for batch in result) - expected = air.dims["lat"] * air.dims["lon"] - assert ( - len(result) == expected - ), f"Length must be {expected}, but was {len(result)}." - print(expected) + assert actual == expected, f"Length must be {expected}, but was {actual}." + print(expected) diff --git a/perf_tests/open_era5.py b/perf_tests/open_era5.py index 27e7edcc..e9f0b88c 100755 --- a/perf_tests/open_era5.py +++ b/perf_tests/open_era5.py @@ -1,15 +1,15 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql # Requires authenticating with GCP era5_ds = xr.open_zarr( "gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2", chunks={"time": 240, "level": 1}, ) -era5_wind_df = qr.read_xarray( +era5_wind_df = xql.read_xarray( era5_ds[["u_component_of_wind", "v_component_of_wind"]] ) -print(era5_wind_df.columns) +print(era5_wind_df.schema) diff --git a/perf_tests/sanity.py b/perf_tests/sanity.py index 2d51a684..9f8db6b6 100755 --- a/perf_tests/sanity.py +++ b/perf_tests/sanity.py @@ -1,16 +1,16 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240, "lat": 5, "lon": 7} + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240, "lat": 5, "lon": 7} - air_small = air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(chunks) + air_small = air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(chunks) - df = qr.read_xarray(air_small).compute() + df = xql.read_xarray(air_small).read_pandas() - print(len(df)) + print(len(df)) diff --git a/pyproject.toml b/pyproject.toml index ae1b3ffb..8a35be4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "maturin" [project] name = "xarray_sql" dynamic = ["version"] -description = "Querry Xarray with SQL." +description = "Query Xarray with SQL." readme = "README.md" requires-python = ">=3.10" # Python 3.9 EOL is October 31, 2025. license = {text = "Apache-2.0"} @@ -31,16 +31,40 @@ classifiers = [ ] dependencies = [ "dask>=2024.8.0", - "datafusion==51.0.0", # This needs to match the cargo datafusion version!! + "datafusion==54.0.0", # This needs to match the cargo datafusion version!! "xarray>=2024.7.0", ] [project.optional-dependencies] +duckdb = [ + "duckdb>=1.4.0", +] +polars = [ + # collect(engine="streaming") landed in polars 1.25; + # LazyFrame.collect_batches (streamed max_result_bytes enforcement) + # in 1.33. PolarsHandle degrades gracefully below both floors. + "polars>=1.33", +] +geo = [ + "pyproj", +] test = [ + "cftime", + "xarray-sql[duckdb,polars,geo]", "pytest", "xarray[io]", "gcsfs", ] +docs = [ + "zensical", + "mkdocstrings[python]", +] +dev = [ + "xarray_sql[docs]", + "pre-commit", + "pytest", + "watchfiles", +] [project.urls] Homepage = "https://github.com/alxmrs/xarray-sql" @@ -51,13 +75,15 @@ features = ["pyo3/extension-module"] module-name = "xarray_sql._native" [tool.setuptools.packages.find] -exclude = ["demo", "perf_tests"] +exclude = ["demo", "perf_tests", "tests", "tests.*"] -[tool.pyink] +[tool.ruff] line-length = 80 -preview = true -pyink-indentation = 2 -pyink-use-majority-quotes = true +indent-width = 4 + +[tool.ruff.format] +indent-style = "space" +quote-style = "double" [tool.mypy] python_version = "3.11" @@ -77,6 +103,7 @@ module = [ "pyarrow.*", "datafusion.*", "xarray.*", + "pandas.*", ] ignore_missing_imports = true @@ -85,11 +112,18 @@ ignore_missing_imports = true [dependency-groups] dev = [ "xarray_sql[test]", + "xarray_sql[docs]", "py-spy>=0.4.0", - "pyink>=24.10.1", + "ruff>=0.15.10", "maturin>=1.9.1", ] [tool.uv] # Rebuild package when any rust files change cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "**/*.rs"}] + +[tool.pytest.ini_options] +testpaths = ["tests"] +markers = [ + "integration: needs network (anonymous GCS); excluded from the CI unit run", +] diff --git a/src/lib.rs b/src/lib.rs index d61e9186..63a5a6bd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,55 +24,956 @@ //! incrementally. This enables processing of larger-than-memory datasets when combined //! with DataFusion's streaming execution. //! -//! ## Parallel Execution Note +//! ## Parallel Execution //! -//! When using DataFusion's parallel execution (multiple partitions), aggregation queries -//! without ORDER BY may return partial results due to how our stream interacts with -//! DataFusion's async runtime. To ensure complete results: -//! - Add ORDER BY to aggregation queries, or -//! - Use `SessionConfig().with_target_partitions(1)` for single-threaded execution -//! TODO(#106): Implement proper parallelism and partition handling. +//! Each xarray chunk becomes a separate partition, enabling parallel execution across +//! multiple cores. +//! +//! ## Filter Pushdown (Partition Pruning) +//! +//! When partition metadata is provided, SQL filters on dimension columns (time, lat, lon) +//! automatically prune partitions that can't contain matching rows. For example: +//! +//! ```sql +//! SELECT * FROM air WHERE time > '2020-02-01' +//! ``` +//! +//! Will skip loading partitions whose time ranges are entirely before 2020-02-01. +//! Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. +use std::collections::{HashMap, HashSet}; use std::ffi::CString; use std::fmt::Debug; use std::sync::Arc; use arrow::array::RecordBatch; -use arrow::datatypes::SchemaRef; -use arrow::pyarrow::FromPyArrow; +use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit}; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; use async_stream::try_stream; +use async_trait::async_trait; use datafusion::catalog::streaming::StreamingTable; -use datafusion::common::DataFusionError; +use datafusion::catalog::Session; +use datafusion::common::stats::Precision; +use datafusion::common::{ + ColumnStatistics, DataFusionError, Result as DFResult, ScalarValue, Statistics, +}; use datafusion::datasource::TableProvider; use datafusion::execution::TaskContext; +use datafusion::logical_expr::expr::InList; +use datafusion::logical_expr::{ + BinaryExpr, Expr, Operator, TableProviderFilterPushDown, TableType, +}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::streaming::PartitionStream; -use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; use pyo3::prelude::*; -use pyo3::types::PyCapsule; -use tokio::runtime::Handle; +use pyo3::types::{PyCapsule, PyList}; + +// ============================================================================ +// Partition Metadata Types for Filter Pushdown +// ============================================================================ + +// TODO(alxmrs, Claude): Support every valid xarray coordinate type. +/// Scalar value for dimension bounds, supporting common xarray coordinate types. +#[derive(Clone, Debug)] +pub enum ScalarBound { + /// 64-bit integer (for integer coordinates) + Int64(i64), + /// 64-bit float (for lat/lon coordinates) + Float64(f64), + /// Nanoseconds since Unix epoch (for datetime64[ns] coordinates) + TimestampNanos(i64), +} + +impl ScalarBound { + /// Compare this bound with a DataFusion ScalarValue. + /// Returns None if types are incompatible. + fn compare_to_scalar(&self, scalar: &ScalarValue) -> Option { + match (self, scalar) { + // Integer comparisons + (ScalarBound::Int64(a), ScalarValue::Int64(Some(b))) => Some(a.cmp(b)), + (ScalarBound::Int64(a), ScalarValue::Int32(Some(b))) => Some(a.cmp(&(*b as i64))), + + // Float comparisons + (ScalarBound::Float64(a), ScalarValue::Float64(Some(b))) => a.partial_cmp(b), + (ScalarBound::Float64(a), ScalarValue::Float32(Some(b))) => a.partial_cmp(&(*b as f64)), + + // Timestamp comparisons - convert to nanoseconds. + // Use checked_mul to avoid silent overflow in release builds; + // on overflow return None (conservative: include the partition). + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampNanosecond(Some(b), _)) => { + Some(a.cmp(b)) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMicrosecond(Some(b), _)) => { + b.checked_mul(1_000).map(|b_ns| a.cmp(&b_ns)) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMillisecond(Some(b), _)) => { + b.checked_mul(1_000_000).map(|b_ns| a.cmp(&b_ns)) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampSecond(Some(b), _)) => { + b.checked_mul(1_000_000_000).map(|b_ns| a.cmp(&b_ns)) + } + + // Incompatible types + _ => None, + } + } +} + +/// Range bounds for one dimension in a partition. +#[derive(Clone, Debug)] +pub struct DimensionRange { + /// The column name (dimension name from xarray) + pub column_name: String, + /// Minimum value (inclusive) - first coordinate value in this partition + pub min: ScalarBound, + /// Maximum value (inclusive) - last coordinate value in this partition + pub max: ScalarBound, +} + +/// Metadata for a single partition, used for filter-based pruning. +#[derive(Clone, Debug)] +pub struct PartitionMetadata { + /// Dimension ranges for this partition, keyed by column name + pub ranges: HashMap, + /// Exact number of rows in this partition (product of the chunk's + /// per-dimension sizes). Every partition supplies it, so the scan can + /// report exact `Statistics::num_rows` to the optimizer and cost-based + /// rules (join build-side selection, broadcast vs. shuffle) have real + /// cardinalities instead of guesses. xarray knows this exactly — it is + /// the product of the partition's dimension lengths — so unlike most + /// table providers these statistics are not estimates. + pub num_rows: usize, +} + +impl PartitionMetadata { + /// Get the range for a specific dimension column. + pub fn get_range(&self, column: &str) -> Option<&DimensionRange> { + self.ranges.get(column) + } +} + +// ============================================================================ +// Custom TableProvider with Filter Pushdown +// ============================================================================ + +/// A table provider that supports partition pruning via filter pushdown. +/// +/// This wraps partition streams with their metadata and implements +/// `TableProvider::supports_filters_pushdown` and partition pruning in `scan()`. +struct PrunableStreamingTable { + schema: SchemaRef, + /// Partition streams paired with their coordinate range metadata. + /// Stored behind the `ProjectableStream` trait so `PrunableStreamingTable` + /// is not coupled to `PyArrowStreamPartition`. + partitions: Vec<(Arc, PartitionMetadata)>, + /// Set of column names that are dimension columns (eligible for pruning) + dimension_columns: HashSet, +} + +impl PrunableStreamingTable { + fn new( + schema: SchemaRef, + partitions: Vec<(Arc, PartitionMetadata)>, + ) -> Self { + // Collect dimension column names from the first partition that has + // non-empty metadata. All partitions share the same dimension names, + // so we only need one representative. Using find_map keeps this O(D) + // rather than O(N × D) — important when N is in the hundreds of + // thousands (e.g. hourly chunks of a decades-long climate dataset). + let dimension_columns: HashSet = partitions + .iter() + .find_map(|(_, meta)| { + if meta.ranges.is_empty() { + None + } else { + Some(meta.ranges.keys().cloned().collect()) + } + }) + .unwrap_or_default(); + + Self { + schema, + partitions, + dimension_columns, + } + } + + /// Determine which partitions should be included based on filters. + /// Returns indices of partitions that may contain matching rows. + fn prune_partitions(&self, filters: &[Expr]) -> Vec { + self.partitions + .iter() + .enumerate() + .filter(|(_, (_, meta))| { + // Include partition unless a filter definitely excludes it + !filters + .iter() + .any(|f| self.filter_excludes_partition(f, meta)) + }) + .map(|(idx, _)| idx) + .collect() + } + + /// Returns true if this filter definitely excludes the partition. + /// Conservative: returns false (include) if uncertain. + fn filter_excludes_partition(&self, expr: &Expr, meta: &PartitionMetadata) -> bool { + match expr { + Expr::BinaryExpr(BinaryExpr { left, op, right }) => { + // Handle AND/OR logic + match op { + Operator::And => { + // For AND, exclude if either side excludes + self.filter_excludes_partition(left, meta) + || self.filter_excludes_partition(right, meta) + } + Operator::Or => { + // For OR, exclude only if both sides exclude + self.filter_excludes_partition(left, meta) + && self.filter_excludes_partition(right, meta) + } + // Handle comparison operators + _ => self.comparison_excludes(left, op, right, meta), + } + } + Expr::Not(_) => { + // NOT inverts the predicate. We cannot safely derive exclusion + // from the inner result: if inner returns false (uncertain), + // !false = true would incorrectly exclude the partition. + // Example: partition [1,10], NOT (col > 5) ≡ col <= 5 — + // inner returns false (max=10 > 5), but the partition contains + // values 1–5 which satisfy col <= 5 and must be included. + // Be conservative: never exclude on NOT. + false + } + Expr::Between(between) => self.between_excludes(between, meta), + Expr::InList(in_list) => self.in_list_excludes(in_list, meta), + // Unknown expression type - be conservative + _ => false, + } + } + + /// Check if a comparison expression excludes this partition. + fn comparison_excludes( + &self, + left: &Expr, + op: &Operator, + right: &Expr, + meta: &PartitionMetadata, + ) -> bool { + // Try to extract column and literal from either side + let (col_name, scalar, flipped) = match (left, right) { + (Expr::Column(c), Expr::Literal(s, _)) => (c.name.clone(), s, false), + (Expr::Literal(s, _), Expr::Column(c)) => (c.name.clone(), s, true), + _ => return false, // Not a simple column-literal comparison + }; + + // Get the dimension range for this column + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, // Not a dimension column, can't prune + }; + + // Flip operator if literal was on left side + let effective_op = if flipped { flip_operator(op) } else { *op }; + + // Determine if partition can be excluded based on operator + match effective_op { + // col > literal: exclude if max <= literal + Operator::Gt => matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) + ), + // col >= literal: exclude if max < literal + Operator::GtEq => { + matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less) + ) + } + // col < literal: exclude if min >= literal + Operator::Lt => matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater | std::cmp::Ordering::Equal) + ), + // col <= literal: exclude if min > literal + Operator::LtEq => { + matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater) + ) + } + // col = literal: exclude if literal outside [min, max] + Operator::Eq => { + let below_min = matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater) + ); + let above_max = matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less) + ); + below_min || above_max + } + // col != literal: can exclude only if range is a single point equal + // to the literal — every row has that value, so none satisfy !=. + Operator::NotEq => { + let min_eq = matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Equal) + ); + let max_eq = matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Equal) + ); + min_eq && max_eq + } + // Other operators: be conservative + _ => false, + } + } + + /// Check if a BETWEEN expression excludes this partition. + fn between_excludes( + &self, + between: &datafusion::logical_expr::Between, + meta: &PartitionMetadata, + ) -> bool { + if between.negated { + // NOT BETWEEN is complex, be conservative + return false; + } + + // Extract column name + let col_name = match between.expr.as_ref() { + Expr::Column(c) => c.name.clone(), + _ => return false, + }; + + // Get dimension range + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, + }; + + // Extract low and high bounds + let (low, high) = match (between.low.as_ref(), between.high.as_ref()) { + (Expr::Literal(l, _), Expr::Literal(h, _)) => (l, h), + _ => return false, + }; + + // Exclude if partition range doesn't overlap with [low, high] + // No overlap if: partition.max < low OR partition.min > high + let max_below_low = matches!( + range.max.compare_to_scalar(low), + Some(std::cmp::Ordering::Less) + ); + let min_above_high = matches!( + range.min.compare_to_scalar(high), + Some(std::cmp::Ordering::Greater) + ); + + max_below_low || min_above_high + } + + /// Check if an IN list expression excludes this partition. + fn in_list_excludes(&self, in_list: &InList, meta: &PartitionMetadata) -> bool { + if in_list.negated { + // NOT IN is complex, be conservative + return false; + } + + // Extract column name + let col_name = match in_list.expr.as_ref() { + Expr::Column(c) => c.name.clone(), + _ => return false, + }; + + // Get dimension range + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, + }; + + // Check if any value in the list could be in this partition's range + let any_in_range = in_list.list.iter().any(|expr| { + if let Expr::Literal(scalar, _) = expr { + // Value is in range if: min <= value <= max + let above_min = matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) + ); + let below_max = matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater | std::cmp::Ordering::Equal) + ); + above_min && below_max + } else { + // Non-literal in list, be conservative + true + } + }); + + // Exclude only if NO values could be in range + !any_in_range + } + + /// Check if an expression is a filter on a dimension column. + fn is_dimension_filter(&self, expr: &Expr) -> bool { + match expr { + Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op { + Operator::And | Operator::Or => { + self.is_dimension_filter(left) || self.is_dimension_filter(right) + } + _ => self.expr_references_dimension(left) || self.expr_references_dimension(right), + }, + Expr::Between(b) => self.expr_references_dimension(&b.expr), + Expr::InList(i) => self.expr_references_dimension(&i.expr), + Expr::Not(inner) => self.is_dimension_filter(inner), + _ => false, + } + } + + /// Check if an expression references a dimension column. + fn expr_references_dimension(&self, expr: &Expr) -> bool { + match expr { + Expr::Column(c) => self.dimension_columns.contains(&c.name), + _ => false, + } + } +} + +/// Extension trait for partition streams that support column projection. +/// +/// Implemented by `PyArrowStreamPartition` so that `PrunableStreamingTable` +/// can push projections to Python factories without coupling to the concrete type. +/// Any new stream implementation (e.g. for non-Python backends) can implement this +/// trait and be used with `PrunableStreamingTable` directly. +trait ProjectableStream: PartitionStream + Debug { + /// Return a new stream that emits only the specified columns. + fn clone_with_projection( + &self, + projection: Arc<[String]>, + projected_schema: SchemaRef, + ) -> Arc; + + /// Clone this stream as a generic `PartitionStream` Arc. + fn clone_as_stream(&self) -> Arc; +} + +/// Flip a comparison operator (for when literal is on left side). +fn flip_operator(op: &Operator) -> Operator { + match op { + Operator::Lt => Operator::Gt, + Operator::LtEq => Operator::GtEq, + Operator::Gt => Operator::Lt, + Operator::GtEq => Operator::LtEq, + other => *other, + } +} + +/// Convert a Python object to a ScalarBound using an explicit dtype tag. +fn python_to_scalar_bound(obj: &Bound<'_, PyAny>, dtype_tag: &str) -> PyResult { + match dtype_tag { + "timestamp_ns" => { + let val = obj.extract::()?; + Ok(ScalarBound::TimestampNanos(val)) + } + "float64" => { + let val = obj.extract::()?; + Ok(ScalarBound::Float64(val)) + } + "int64" => { + let val = obj.extract::()?; + Ok(ScalarBound::Int64(val)) + } + _ => Err(pyo3::exceptions::PyTypeError::new_err(format!( + "Unsupported dtype tag for partition bound: {dtype_tag}" + ))), + } +} + +/// Convert a bound Python metadata dict into per-dimension coordinate ranges. +/// +/// Operates on an already-bound reference so no additional GIL acquisition +/// is needed — this is called from within a `#[pymethods]` context where +/// the GIL is already held. The caller pairs the result with the partition's +/// row count to build a [`PartitionMetadata`]. +fn convert_python_ranges_from_bound( + meta_obj: &Bound<'_, PyAny>, +) -> PyResult> { + type MetaDict = HashMap, Py, String)>; + let meta_dict: MetaDict = meta_obj.extract()?; + let py = meta_obj.py(); + let mut ranges = HashMap::new(); + for (dim_name, (min_obj, max_obj, dtype_tag)) in meta_dict { + let min_bound = python_to_scalar_bound(min_obj.bind(py), &dtype_tag)?; + let max_bound = python_to_scalar_bound(max_obj.bind(py), &dtype_tag)?; + ranges.insert( + dim_name.clone(), + DimensionRange { + column_name: dim_name, + min: min_bound, + max: max_bound, + }, + ); + } + Ok(ranges) +} + +impl Debug for PrunableStreamingTable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrunableStreamingTable") + .field("schema", &self.schema) + .field("num_partitions", &self.partitions.len()) + .field("dimension_columns", &self.dimension_columns) + .finish() + } +} + +#[async_trait] +impl TableProvider for PrunableStreamingTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DFResult> { + // For dimension filters we can do exact pruning at partition level + // Return Inexact so DataFusion still applies row-level filtering + Ok(filters + .iter() + .map(|expr| { + if self.is_dimension_filter(expr) { + // We can prune partitions but not individual rows within + TableProviderFilterPushDown::Inexact + } else { + TableProviderFilterPushDown::Unsupported + } + }) + .collect()) + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DFResult> { + // Prune partitions based on filters + let included_indices = self.prune_partitions(filters); + + // Exact per-partition row counts for the partitions that survive + // pruning, in scan order. These feed `XarrayScanExec`'s statistics so + // the optimizer sees real cardinalities. + let included_metas: Vec<&PartitionMetadata> = included_indices + .iter() + .map(|&idx| &self.partitions[idx].1) + .collect(); + let partition_rows: Vec> = included_metas + .iter() + .map(|meta| Precision::Exact(meta.num_rows)) + .collect(); + + // Handle empty case — all partitions pruned, return empty plan + if included_indices.is_empty() { + let empty_table = StreamingTable::try_new(Arc::clone(&self.schema), vec![])?; + let inner = empty_table.scan(state, projection, filters, limit).await?; + let stats = build_scan_statistics(inner.schema().as_ref(), &included_metas); + return Ok(Arc::new(XarrayScanExec::new(inner, stats, partition_rows))); + } + + // Determine whether to push projection down to the Python factory. + // + // We push when the projection includes at least one data variable + // (non-dimension column), because xarray can selectively load only + // the requested data arrays while dimension coordinates are always + // available via xarray's coordinate system. + // + // We do NOT push when: + // - projection is None (load everything — factory receives None) + // - projection is Some([]) (COUNT(*) — let StreamingTable handle) + // - projection contains only dimension columns (ds.to_dataframe() + // needs at least one data variable; dimensions are always loaded) + let push_projection = match projection { + Some(indices) if !indices.is_empty() => indices + .iter() + .any(|&i| !self.dimension_columns.contains(self.schema.field(i).name())), + _ => false, + }; + + if push_projection { + let indices = projection.unwrap(); + + // Build the projected schema (only the requested fields) + let proj_fields: Vec<_> = indices + .iter() + .map(|&i| self.schema.field(i).clone()) + .collect(); + let projected_schema = Arc::new(Schema::new(proj_fields)); + + // Collect the requested column names to send to the factory. + // Stored in an Arc so each clone_with_projection call shares the + // same allocation via an atomic refcount increment (no N Vec copies). + let proj_col_names: Arc<[String]> = indices + .iter() + .map(|&i| self.schema.field(i).name().to_string()) + .collect::>() + .into(); + + // Clone each pruned partition with the projection baked in. + // The factory will receive proj_col_names and load only those vars. + let projected_partitions: Vec> = included_indices + .iter() + .map(|&idx| { + self.partitions[idx].0.clone_with_projection( + Arc::clone(&proj_col_names), + Arc::clone(&projected_schema), + ) + }) + .collect(); + + // StreamingTable already has the projected schema — pass None for + // projection so it doesn't wrap the stream in a redundant ProjectionExec. + let streaming = StreamingTable::try_new(projected_schema, projected_partitions)?; + let inner = streaming.scan(state, None, filters, limit).await?; + let stats = build_scan_statistics(inner.schema().as_ref(), &included_metas); + Ok(Arc::new(XarrayScanExec::new(inner, stats, partition_rows))) + } else { + // No projection pushdown — factory is called with None (loads all + // columns). StreamingTable applies projection via ProjectionExec. + let included_partitions: Vec> = included_indices + .iter() + .map(|&idx| self.partitions[idx].0.clone_as_stream()) + .collect(); + let streaming = StreamingTable::try_new(Arc::clone(&self.schema), included_partitions)?; + let inner = streaming.scan(state, projection, filters, limit).await?; + let stats = build_scan_statistics(inner.schema().as_ref(), &included_metas); + Ok(Arc::new(XarrayScanExec::new(inner, stats, partition_rows))) + } + } +} + +// ============================================================================ +// Exact Statistics + Scan Wrapper +// ============================================================================ + +/// Sum the exact per-partition row counts into a `Precision`. +/// +/// Every partition carries an exact count, so the total is exact too. The sum +/// is the table's row count and fits `usize`; `saturating_add` is defensive +/// against a pathological overflow rather than an expected case. +fn sum_row_counts<'a>(metas: impl Iterator) -> Precision { + Precision::Exact(metas.fold(0usize, |total, meta| total.saturating_add(meta.num_rows))) +} + +/// Fold two same-variant `ScalarBound`s, keeping the smaller (`keep_min`) or +/// larger one. Returns `None` if the variants differ (never expected within a +/// single dimension) so the caller can fall back to unknown. +fn fold_bound(a: &ScalarBound, b: &ScalarBound, keep_min: bool) -> Option { + let ord = match (a, b) { + (ScalarBound::Int64(x), ScalarBound::Int64(y)) => x.partial_cmp(y), + (ScalarBound::Float64(x), ScalarBound::Float64(y)) => x.partial_cmp(y), + (ScalarBound::TimestampNanos(x), ScalarBound::TimestampNanos(y)) => x.partial_cmp(y), + _ => return None, + }?; + let take_a = if keep_min { + ord != std::cmp::Ordering::Greater + } else { + ord != std::cmp::Ordering::Less + }; + Some(if take_a { a.clone() } else { b.clone() }) +} + +/// Convert a coordinate bound into a `ScalarValue` matching a column's Arrow +/// type, so the min/max we report line up with the column's own type. Returns +/// `None` for combinations we cannot convert without loss (e.g. a timestamp +/// unit we can't scale exactly), in which case the column is left without +/// min/max rather than risk a wrong value. +fn bound_to_scalar(bound: &ScalarBound, dtype: &DataType) -> Option { + match (bound, dtype) { + (ScalarBound::Int64(v), DataType::Int64) => Some(ScalarValue::Int64(Some(*v))), + (ScalarBound::Int64(v), DataType::Int32) => { + i32::try_from(*v).ok().map(|x| ScalarValue::Int32(Some(x))) + } + (ScalarBound::Float64(v), DataType::Float64) => Some(ScalarValue::Float64(Some(*v))), + (ScalarBound::Float64(v), DataType::Float32) => Some(ScalarValue::Float32(Some(*v as f32))), + // Datetime coordinates arrive as nanoseconds (see `cftime.partition_bounds` + // and the datetime64[ns] path in `_block_metadata`). Map them onto the + // column's own timestamp unit, but only when the scaling is exact so a + // reported bound is never a rounded value. + (ScalarBound::TimestampNanos(v), DataType::Timestamp(unit, tz)) => { + let scaled = match unit { + TimeUnit::Nanosecond => Some(*v), + TimeUnit::Microsecond if v % 1_000 == 0 => Some(v / 1_000), + TimeUnit::Millisecond if v % 1_000_000 == 0 => Some(v / 1_000_000), + TimeUnit::Second if v % 1_000_000_000 == 0 => Some(v / 1_000_000_000), + _ => None, + }?; + Some(match unit { + TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(scaled), tz.clone()), + TimeUnit::Microsecond => { + ScalarValue::TimestampMicrosecond(Some(scaled), tz.clone()) + } + TimeUnit::Millisecond => { + ScalarValue::TimestampMillisecond(Some(scaled), tz.clone()) + } + TimeUnit::Second => ScalarValue::TimestampSecond(Some(scaled), tz.clone()), + }) + } + _ => None, + } +} + +/// Exact in-memory byte size of `num_rows` rows of `schema`, or `Absent` if any +/// column is variable-width (e.g. Utf8) and cannot be sized from the row count +/// alone. Our data model is dense fixed-width grids, so this is normally exact. +fn total_byte_size(schema: &Schema, num_rows: &Precision) -> Precision { + let Precision::Exact(rows) = num_rows else { + return Precision::Absent; + }; + let mut row_width = 0usize; + for field in schema.fields() { + match field.data_type().primitive_width() { + Some(w) => row_width += w, + None => return Precision::Absent, + } + } + Precision::Exact(rows.saturating_mul(row_width)) +} + +/// Build `Statistics` for a scan over the given partitions. +/// +/// Every statistic here is derived from coordinate metadata xarray already +/// knows — none of it scans the data — and each is exact, not an estimate: +/// +/// * `num_rows`: summed product of each surviving chunk's dimension sizes. +/// Drives `JoinSelection`'s build-side choice and lets `COUNT(*)` skip the +/// scan entirely. +/// * `total_byte_size`: `num_rows × fixed row width`, for memory-cost rules. +/// * per dimension-coordinate column: exact `min`/`max` (folded coordinate +/// bounds — the join/filter keys) and `null_count = 0` (grid axes are always +/// fully populated). Data variables are left unknown; their bounds would need +/// a scan. +fn build_scan_statistics(output_schema: &Schema, metas: &[&PartitionMetadata]) -> Statistics { + let mut stats = Statistics::new_unknown(output_schema); + stats.num_rows = sum_row_counts(metas.iter().copied()); + stats.total_byte_size = total_byte_size(output_schema, &stats.num_rows); + + for (col_idx, field) in output_schema.fields().iter().enumerate() { + // Fold this column's min/max across every partition that carries a + // range for it. A column has a range iff it is a dimension coordinate + // with a representable bound; all such partitions share the same bound + // variant, so the fold is well-defined. + let mut folded: Option<(ScalarBound, ScalarBound)> = None; + for meta in metas { + if let Some(range) = meta.ranges.get(field.name()) { + folded = Some(match folded { + None => (range.min.clone(), range.max.clone()), + Some((lo, hi)) => ( + fold_bound(&lo, &range.min, true).unwrap_or(lo), + fold_bound(&hi, &range.max, false).unwrap_or(hi), + ), + }); + } + } + + let Some((lo, hi)) = folded else { continue }; + // This column is a coordinate axis: never null, so the null count is + // exactly zero regardless of whether the bound maps to a ScalarValue. + let dtype = field.data_type(); + stats.column_statistics[col_idx] = ColumnStatistics { + null_count: Precision::Exact(0), + min_value: bound_to_scalar(&lo, dtype) + .map(Precision::Exact) + .unwrap_or(Precision::Absent), + max_value: bound_to_scalar(&hi, dtype) + .map(Precision::Exact) + .unwrap_or(Precision::Absent), + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }; + } + + stats +} + +/// A thin scan operator that wraps an inner `StreamingTableExec` and reports +/// exact `Statistics` to the query optimizer. +/// +/// Execution, schema, ordering, and partitioning are delegated verbatim to the +/// inner plan (so projection mechanics are reused unchanged); the only thing +/// this node adds is real cardinality. `StreamingTableExec` reports unknown +/// statistics, and the physical `JoinSelection` rule reads statistics from the +/// `ExecutionPlan` (not from `TableProvider::statistics`) — even in DataFusion +/// 54, which forwards `ExecutionPlan` statistics across the FFI boundary — so +/// this wrapper is what carries the exact cardinality through to the optimizer. +#[derive(Debug)] +struct XarrayScanExec { + inner: Arc, + statistics: Statistics, + /// Exact row count per output partition (parallel to `inner` partitions), + /// so `partition_statistics(Some(i))` is exact too. + partition_rows: Vec>, +} + +impl XarrayScanExec { + fn new( + inner: Arc, + statistics: Statistics, + partition_rows: Vec>, + ) -> Self { + Self { + inner, + statistics, + partition_rows, + } + } +} + +impl DisplayAs for XarrayScanExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "XarrayScanExec: rows={:?}", self.statistics.num_rows) + } + DisplayFormatType::TreeRender => { + write!(f, "rows={:?}", self.statistics.num_rows) + } + } + } +} + +#[async_trait] +impl ExecutionPlan for XarrayScanExec { + fn name(&self) -> &str { + "XarrayScanExec" + } + + fn properties(&self) -> &Arc { + // Delegate partitioning + output ordering + boundedness to the inner + // StreamingTableExec. + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + // A scan is a leaf; the inner plan is an execution detail, not a child + // the optimizer should rewrite. + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DFResult> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> DFResult { + self.inner.execute(partition, ctx) + } + + fn partition_statistics(&self, partition: Option) -> DFResult> { + match partition { + None => Ok(Arc::new(self.statistics.clone())), + Some(i) => { + // Build a fresh, self-consistent per-partition summary rather + // than reusing the table-level one: the folded column min/max + // and `total_byte_size` describe the whole scan, and claiming + // them (as `Exact`) for a single partition could be wrong — + // partition `i` need not contain the table-wide min/max. We + // keep only what is exact per partition: its row count and the + // byte size derived from it. Column bounds are left unknown + // (we do not retain per-partition bounds here). + let num_rows = self + .partition_rows + .get(i) + .cloned() + .unwrap_or(Precision::Absent); + let schema = self.inner.schema(); + let mut s = Statistics::new_unknown(&schema); + s.num_rows = num_rows; + s.total_byte_size = total_byte_size(&schema, &s.num_rows); + Ok(Arc::new(s)) + } + } + } +} /// A partition stream that wraps a Python factory function that creates streams. /// /// The factory is called lazily on each `execute()` invocation, allowing /// the same table to be queried multiple times. +/// +/// When `projection` is set, the factory is called with that list of column +/// names so that xarray only loads the requested data variables rather than +/// materializing every variable in the dataset. struct PyArrowStreamPartition { schema: SchemaRef, - /// A Python callable (factory) that returns a fresh stream implementing `__arrow_c_stream__`. - /// Called on each execute() to create a new stream. - stream_factory: Py, + /// A Python callable (factory) that returns a fresh stream. + /// Signature: `make_stream(projection_names: Optional[List[str]]) -> RecordBatchReader` + /// + /// Wrapped in `Arc` so `ProjectableStream::clone_with_projection` can share + /// the same Python object across projected partitions without acquiring the + /// GIL — only an atomic reference-count increment is needed. + stream_factory: Arc>, + /// Column names to pass to the factory. `None` means load all columns. + /// Stored as `Arc<[String]>` so multiple projected clones share one allocation. + projection: Option>, } impl PyArrowStreamPartition { fn new(stream_factory: Py, schema: SchemaRef) -> Self { Self { schema, - stream_factory, + stream_factory: Arc::new(stream_factory), + projection: None, } } } +impl ProjectableStream for PyArrowStreamPartition { + /// Return a new partition that emits only the given columns. + /// + /// Clones the factory `Arc` (atomic refcount increment, no GIL) so the + /// same Python callable is shared across all projected partitions. + fn clone_with_projection( + &self, + projection: Arc<[String]>, + projected_schema: SchemaRef, + ) -> Arc { + Arc::new(Self { + schema: projected_schema, + stream_factory: Arc::clone(&self.stream_factory), + projection: Some(projection), + }) + } + + fn clone_as_stream(&self) -> Arc { + Arc::new(Self { + schema: Arc::clone(&self.schema), + stream_factory: Arc::clone(&self.stream_factory), + projection: self.projection.clone(), + }) + } +} + impl Debug for PyArrowStreamPartition { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PyArrowStreamPartition") @@ -89,21 +990,41 @@ impl PartitionStream for PyArrowStreamPartition { fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { let schema = Arc::clone(&self.schema); - // Clone the factory with the GIL held - let factory = Python::attach(|py| self.stream_factory.clone_ref(py)); + // Clone the factory Arc (no GIL needed) and the projection list. + let factory = Arc::clone(&self.stream_factory); + let projection = self.projection.clone(); // Create a lazy stream using try_stream! macro. - // This is cleaner than manual state management with unfold. - // Each iteration acquires the GIL and reads one batch. + // The GIL is acquired only for the duration of each Python call + // and released between batches. After each yielded batch we call + // yield_now() to explicitly suspend this task, giving the Tokio + // executor a chance to poll other partition streams (which can then + // acquire the GIL and make progress in parallel). let batch_stream = try_stream! { - // Call factory to get the PyArrow RecordBatchReader + // Call factory with the projection argument. + // `projection` is either a Python list of column names or None + // (load all columns). The factory always receives exactly one arg + // so it can distinguish "no projection" from "empty projection". let reader: Py = Python::attach(|py| { - factory.call0(py).map_err(|e| { + let proj_arg = match &projection { + Some(cols) => PyList::new(py, cols.iter().map(|s| s.as_str())) + .map_err(|e| { + DataFusionError::Execution(format!( + "Failed to build projection list: {e}" + )) + })? + .into_any(), + None => py.None().into_bound(py).into_any(), + }; + factory.call1(py, (proj_arg,)).map_err(|e| { DataFusionError::Execution(format!("Failed to call stream factory: {e}")) }) })?; - // Read batches until StopIteration + // Read batches until StopIteration. + // The GIL is released between iterations; yield_now() ensures + // other async tasks (i.e., other partitions) are scheduled + // before this stream is polled again. loop { let batch_result = Python::attach(|py| { let bound_reader = reader.bind(py); @@ -129,7 +1050,12 @@ impl PartitionStream for PyArrowStreamPartition { }); match batch_result { - Ok(Some(batch)) => yield batch, + Ok(Some(batch)) => { + yield batch; + // Yield to the executor so that other partition + // streams can acquire the GIL and make progress. + tokio::task::yield_now().await; + } Ok(None) => break, Err(e) => Err(e)?, } @@ -140,63 +1066,125 @@ impl PartitionStream for PyArrowStreamPartition { } } -/// A lazy table provider that wraps a Python stream factory. +// ============================================================================ +// FFI Helpers +// ============================================================================ + +/// Extract an `FFI_LogicalExtensionCodec` from a Python session object. +/// +/// DataFusion 52 passes the `SessionContext` to `__datafusion_table_provider__` +/// so that the provider can obtain the codec needed for physical-plan +/// serialisation across the FFI boundary. The session exposes this via +/// `__datafusion_logical_extension_codec__()`, which returns a PyCapsule +/// named `"datafusion_logical_extension_codec"`. +/// +/// Mirrors the helper in the official datafusion-python FFI example +/// (`examples/datafusion-ffi-example/src/utils.rs`). +fn ffi_logical_codec_from_pycapsule( + session: Bound<'_, PyAny>, +) -> PyResult { + let attr = "__datafusion_logical_extension_codec__"; + let capsule = if session.hasattr(attr)? { + session.getattr(attr)?.call0()? + } else { + session + }; + + let capsule = capsule.cast::().map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "session did not produce a PyCapsule for the logical extension codec: {e}" + )) + })?; + + // `pointer_checked` validates the capsule name matches before handing back + // the pointer, so an unexpectedly-named capsule is rejected here. + let expected = CString::new("datafusion_logical_extension_codec").unwrap(); + let ptr = capsule + .pointer_checked(Some(expected.as_c_str())) + .map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "capsule is not a datafusion_logical_extension_codec: {e}" + )) + })?; + + // SAFETY: The capsule was produced by datafusion-python and contains a + // valid FFI_LogicalExtensionCodec (#[repr(C)] StableAbi struct). + let codec = unsafe { &*(ptr.as_ptr() as *const FFI_LogicalExtensionCodec) }; + Ok(codec.clone()) +} + +// ============================================================================ +// Python-visible Table Class +// ============================================================================ + +/// A lazy table provider that wraps Python stream factory functions. /// /// This class implements the `__datafusion_table_provider__` protocol, allowing /// it to be registered with DataFusion's `SessionContext.register_table()`. /// /// Data is NOT read until query execution time - this enables true lazy evaluation. -/// The factory function is called on each query execution to create a fresh stream, -/// allowing the same table to be queried multiple times. +/// Each partition has its own factory function that is called on query execution +/// to create a fresh stream, enabling true parallelism in DataFusion. +/// +/// ## Filter Pushdown +/// +/// SQL filters on dimension columns (time, lat, lon, etc.) automatically prune +/// partitions that can't contain matching rows when metadata is supplied. /// /// # Example /// /// ```python /// from datafusion import SessionContext -/// from xarray_sql import LazyArrowStreamTable, XarrayRecordBatchReader +/// import pyarrow as pa /// -/// # Create a factory that produces lazy readers -/// def make_reader(): -/// return XarrayRecordBatchReader(ds, chunks={'time': 240}) +/// schema = pa.schema([("time", pa.int64()), ("air", pa.float32())]) /// -/// # Get schema from a sample reader -/// sample = make_reader() -/// schema = sample.schema +/// # Each element is a (factory_callable, metadata_dict, num_rows) tuple. +/// # metadata_dict maps dim name -> (min, max, dtype_str); use {} for no pruning. +/// # num_rows is the exact partition row count. +/// def make_partitions(): +/// yield (lambda: pa.RecordBatchReader.from_batches(schema, batches_0), +/// {"time": (0, 1_000_000_000, "int64")}, len(batches_0_rows)) +/// yield (lambda: pa.RecordBatchReader.from_batches(schema, batches_1), +/// {"time": (1_000_000_001, 2_000_000_000, "int64")}, len(batches_1_rows)) /// -/// # Wrap factory in lazy table - NO DATA LOADED -/// table = LazyArrowStreamTable(make_reader, schema) +/// table = LazyArrowStreamTable(make_partitions(), schema) /// -/// # Register with DataFusion - STILL NO DATA LOADED /// ctx = SessionContext() /// ctx.register_table("air", table) -/// -/// # Data only loaded HERE during collect() -/// # Each query creates a fresh stream via the factory -/// result = ctx.sql("SELECT AVG(air) FROM air").collect() -/// result2 = ctx.sql("SELECT * FROM air LIMIT 10").collect() # Works! +/// result = ctx.sql("SELECT AVG(air) FROM air WHERE time > 500000000").to_arrow_table() /// ``` + #[pyclass(name = "LazyArrowStreamTable")] struct LazyArrowStreamTable { - /// The underlying StreamingTable - table: Arc, + /// The underlying table provider with pruning support + table: Arc, } #[pymethods] impl LazyArrowStreamTable { - /// Create a new LazyArrowStreamTable from a stream factory function. + /// Create a new LazyArrowStreamTable from an iterable of partition pairs. /// /// Args: - /// stream_factory: A callable that returns a Python object implementing - /// the Arrow PyCapsule interface (`__arrow_c_stream__`). - /// Called on each query execution to create a fresh stream. - /// schema: A PyArrow Schema for the table. Required since the factory - /// hasn't been called yet. + /// partitions: Any Python iterable yielding + /// ``(factory, metadata_dict, num_rows)`` tuples, where: + /// - ``factory`` is a zero-argument callable returning a + /// ``pa.RecordBatchReader`` (called lazily at query time). + /// - ``metadata_dict`` is a ``dict[str, tuple[Any, Any, str]]`` + /// mapping dimension name to ``(min, max, dtype_str)``; pass + /// ``{}`` to skip pruning for a partition. + /// - ``num_rows`` is the exact row count for the partition, so + /// the scan reports exact ``Statistics`` to the optimizer. + /// Generators are accepted, so partition state can be produced + /// one item at a time and released after Rust stores it. + /// schema: A PyArrow Schema for the table. /// /// Raises: /// TypeError: If the schema is not a valid PyArrow Schema. + /// ValueError: If the partitions iterable is empty. #[new] - fn new(stream_factory: &Bound<'_, PyAny>, schema: &Bound<'_, PyAny>) -> PyResult { - // Convert the PyArrow schema to Arrow schema + #[pyo3(signature = (partitions, schema))] + fn new(partitions: &Bound<'_, PyAny>, schema: &Bound<'_, PyAny>) -> PyResult { use arrow::datatypes::Schema; use arrow::pyarrow::FromPyArrow; @@ -205,18 +1193,39 @@ impl LazyArrowStreamTable { })?; let schema_ref = Arc::new(arrow_schema); - // Create the partition stream with the factory - let partition = - PyArrowStreamPartition::new(stream_factory.clone().unbind(), schema_ref.clone()); + // Consume the Python iterable one item at a time. + // All GIL-bound work happens here, in a single GIL-held context, + // eliminating the per-partition Python::attach() calls of the old + // three-list approach. Python can release each block dict, factory + // closure, and metadata dict as soon as Rust has ingested them. + // Stored as Arc so PrunableStreamingTable + // is decoupled from PyArrowStreamPartition. + let mut partition_list: Vec<(Arc, PartitionMetadata)> = Vec::new(); + for item_result in partitions.try_iter()? { + let item = item_result?; + // Each partition is a ``(factory, metadata_dict, num_rows)`` tuple; + // `num_rows` is the exact per-partition row count that feeds the + // scan's statistics. + let (factory_obj, meta_obj, num_rows): (Py, Py, usize) = + item.extract().map_err(|e| { + pyo3::exceptions::PyTypeError::new_err(format!( + "each partition must be a (factory, metadata_dict, num_rows) tuple: {e}" + )) + })?; + let ranges = convert_python_ranges_from_bound(meta_obj.bind(partitions.py()))?; + let meta = PartitionMetadata { ranges, num_rows }; + let partition: Arc = + Arc::new(PyArrowStreamPartition::new(factory_obj, schema_ref.clone())); + partition_list.push((partition, meta)); + } - // Create the StreamingTable - let table = - StreamingTable::try_new(schema_ref, vec![Arc::new(partition)]).map_err(|e| { - pyo3::exceptions::PyRuntimeError::new_err(format!( - "Failed to create StreamingTable: {e}" - )) - })?; + if partition_list.is_empty() { + return Err(pyo3::exceptions::PyValueError::new_err( + "partitions iterable must not be empty", + )); + } + let table = PrunableStreamingTable::new(schema_ref, partition_list); Ok(Self { table: Arc::new(table), }) @@ -226,33 +1235,27 @@ impl LazyArrowStreamTable { /// /// This method is called by DataFusion's `register_table()` to get a /// foreign table provider that can be used in queries. + /// + /// In DataFusion 52+, the caller passes `session` (a `SessionContext`) + /// so that the provider can access task-context and codec information + /// needed for physical plan serialisation across the FFI boundary. fn __datafusion_table_provider__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { - // Create the FFI table provider - let provider: Arc = self.table.clone(); + let codec = ffi_logical_codec_from_pycapsule(session)?; - // Try to get the current tokio runtime handle (available when called from DataFusion context) - let runtime = Handle::try_current().ok(); + let provider: Arc = self.table.clone(); - // Create FFI wrapper (v49 API takes 3 arguments) - let ffi_provider = FFI_TableProvider::new( - provider, false, // can_support_pushdown_filters - runtime, - ); + let ffi_provider = FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec); - // Create the capsule name let name = CString::new("datafusion_table_provider").unwrap(); - - // Create the PyCapsule without a destructor closure - // The PyCapsule takes ownership of the FFI_TableProvider PyCapsule::new(py, ffi_provider, Some(name)) } /// Get the schema of the table as a PyArrow Schema. fn schema(&self, py: Python<'_>) -> PyResult> { - use arrow::pyarrow::ToPyArrow; self.table .schema() .to_pyarrow(py) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..add0b59e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,150 @@ +import pytest + +import numpy as np +import pandas as pd +import xarray as xr + + +def rand_wx(start: str, end: str) -> xr.Dataset: + np.random.seed(42) + lat = np.linspace(-90, 90, num=720) + lon = np.linspace(-180, 180, num=1440) + time = pd.date_range(start, end, freq="h") + level = np.array([1000, 500], dtype=np.int32) + reference_time = pd.Timestamp(start) + temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level)) + precipitation = 10 * np.random.rand(720, 1440, len(time), len(level)) + return xr.Dataset( + data_vars=dict( + temperature=(["lat", "lon", "time", "level"], temperature), + precipitation=(["lat", "lon", "time", "level"], precipitation), + ), + coords=dict( + lat=lat, + lon=lon, + time=time, + level=level, + reference_time=reference_time, + ), + attrs=dict(description="Random weather."), + ) + + +def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100): + """Create a large xarray dataset for memory testing.""" + np.random.seed(42) + + time = pd.date_range("2020-01-01", periods=time_steps, freq="h") + lat = np.linspace(-90, 90, lat_points) + lon = np.linspace(-180, 180, lon_points) + + temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10 + precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100 + + return xr.Dataset( + { + "temperature": (["time", "lat", "lon"], temp_data), + "precipitation": (["time", "lat", "lon"], precip_data), + }, + coords={"time": time, "lat": lat, "lon": lon}, + ) + + +@pytest.fixture +def air(): + ds = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + return ds.chunk(chunks) + + +@pytest.fixture +def air_small(air): + return air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk({"time": 240}) + + +@pytest.fixture +def randwx(): + return rand_wx("1995-01-13T00", "1995-01-13T01") + + +@pytest.fixture +def large_ds(): + return create_large_dataset().chunk({"time": 25}) + + +@pytest.fixture +def air_dataset_small(): + ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) + return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)) + + +@pytest.fixture +def air_dataset_large(): + return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) + + +@pytest.fixture +def rasm_ds(): + """rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time.""" + return xr.tutorial.open_dataset("rasm") + + +@pytest.fixture +def weather_dataset(): + ds = rand_wx("2023-01-01T00", "2023-01-01T12") + return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk( + {"time": 3} + ) + + +@pytest.fixture +def synthetic_dataset(): + return create_large_dataset( + time_steps=50, lat_points=20, lon_points=20 + ).chunk({"time": 25}) + + +@pytest.fixture +def station_dataset(): + return xr.Dataset( + { + "station_id": (["station"], [1, 2, 3, 4, 5]), + "elevation": (["station"], [100, 250, 500, 750, 1000]), + "name": ( + ["station"], + [ + "Station_A", + "Station_B", + "Station_C", + "Station_D", + "Station_E", + ], + ), + } + ).chunk({"station": 5}) + + +@pytest.fixture +def air_and_stations(): + air = ( + xr.tutorial.open_dataset("air_temperature") + .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8)) + .chunk({"time": 6}) + ) + stations = xr.Dataset( + { + "station_id": (["station"], [101, 102, 103]), + "lat": ( + ["station"], + [air.lat.values[0], air.lat.values[2], air.lat.values[4]], + ), + "lon": ( + ["station"], + [air.lon.values[1], air.lon.values[3], air.lon.values[5]], + ), + "elevation": (["station"], [100, 250, 500]), + } + ).chunk({"station": 3}) + return air, stations diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py new file mode 100644 index 00000000..34aebc72 --- /dev/null +++ b/tests/test_arrow_dataset.py @@ -0,0 +1,542 @@ +"""Contract tests for the engine-neutral pyarrow dataset view. + +``xql.arrow_dataset`` returns a real ``pyarrow.dataset.Dataset`` for +consumers of ``schema``, ``scanner``, ``get_fragments`` and the scan +conveniences — pyarrow itself and Polars are exercised here; DuckDB has +its own suite in ``test_duckdb_backend.py``. DataFusion's native Rust +``TableProvider`` (``XarrayContext``) fills the same role through +DataFusion's own extension trait and sits outside this contract. + +The contract, one section of this file per clause: + +1. Exactness — the pushed filter is applied row-exactly after pruning; + pruning never decides correctness. +2. Projection — exactly the requested columns come back; what gets read + is the projected columns plus the filter's, nothing else. +3. Pruning/counting — provably impossible regions are never read, + provable counts are pure arithmetic, and anything uncertain + (boundary chunks, NaN/NaT coordinates, opaque expressions) is + scanned conservatively. +4. Laziness — construction reads dimension coordinates only; scans + repeat, survive mid-scan abandonment, and run concurrently. +5. Tuning (``batch_size``, ``prefetch``, ``prefetch_bytes``, + ``coalesce_rows``) changes the shape of the work, never the result; + non-positive ``batch_size`` is rejected at construction; fragments + stay one per source chunk. +6. Schema stays on offset types (a view type disables DuckDB's + pushdown); the prefetch pool is fully started at construction and + shuts down when the dataset is collected. +""" + +import threading + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.compute as pc +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def ds() -> xr.Dataset: + np.random.seed(3) + return xr.Dataset( + { + "temperature": ( + ["time", "lat"], + 20 + 5 * np.random.randn(20, 6), + ), + "humidity": (["time", "lat"], np.random.rand(20, 6)), + }, + coords={ + "time": pd.date_range("2022-01-01", periods=20, freq="D"), + "lat": np.linspace(-25.0, 25.0, 6), + }, + ).chunk({"time": 5}) + + +class _ChunkCounter: + def __init__(self): + self.blocks = [] + self.column_sets = [] + + def __call__(self, block, names): + self.blocks.append(block) + self.column_sets.append(tuple(names)) + + +def _hourly_grid() -> xr.Dataset: + """100 hourly steps x 4 latitudes with sequential values.""" + return xr.Dataset( + {"t2m": (["time", "lat"], np.arange(100.0 * 4).reshape(100, 4))}, + coords={ + "time": pd.date_range("2020-01-01", periods=100, freq="h"), + "lat": np.linspace(-30.0, 30.0, 4), + }, + ) + + +@pytest.fixture +def counted(): + """A pushdown dataset over hourly data with a chunk-read counter.""" + source = _hourly_grid() + counter = _ChunkCounter() + dataset = XarrayPushdownDataset( + source, {"time": 10}, _iteration_callback=counter + ) + return dataset, counter + + +# -- Dataset protocol surface ------------------------------------------------ + + +def test_to_table_projects_and_filters(ds): + table = xql.arrow_dataset(ds).to_table( + columns=["time", "temperature"], + filter=pc.field("lat") > 0, + ) + assert table.column_names == ["time", "temperature"] + assert table.num_rows == 20 * 3 # lat > 0 keeps 3 of 6 latitudes + + +def test_count_rows_and_head(ds): + dataset = xql.arrow_dataset(ds) + assert dataset.count_rows() == 20 * 6 + assert dataset.head(7).num_rows == 7 + + +def test_get_fragments_prunes_and_scans(ds): + dataset = xql.arrow_dataset(ds) + assert len(dataset.get_fragments()) == 4 # time chunked by 5 + + # A time predicate covering the first chunk keeps one fragment. + early = pc.field("time") < pa.scalar( + pd.Timestamp("2022-01-06"), type=pa.timestamp("ns") + ) + kept = dataset.get_fragments(filter=early) + assert len(kept) == 1 + assert kept[0].to_table().num_rows == 5 * 6 + + # An unsatisfiable predicate prunes everything. + assert dataset.get_fragments(filter=pc.field("lat") > 100) == [] + + +def test_scanner_honors_batch_size(ds): + dataset = xql.arrow_dataset(ds) + batches = list(dataset.scanner(batch_size=7).to_batches()) + assert sum(b.num_rows for b in batches) == 20 * 6 + assert max(b.num_rows for b in batches) <= 7 + + # The kwarg travels through the inherited to_batches path, which is + # how Polars sizes its morsels. + sizes = [b.num_rows for b in dataset.to_batches(batch_size=11)] + assert sum(sizes) == 20 * 6 + assert max(sizes) <= 11 + + +def test_schema_never_uses_view_types(ds): + # A single view-typed column disables DuckDB's filter pushdown for + # the whole table; pin the schema to offset layouts so a pyarrow + # upgrade cannot regress this silently. + for field in xql.arrow_dataset(ds).schema: + assert field.type not in (pa.string_view(), pa.binary_view()) + + +# -- Consumer integrations --------------------------------------------------- + + +def test_datafusion_register_dataset_round_trips(ds): + from datafusion import SessionContext + + ctx = SessionContext() + ctx.register_dataset("t", xql.arrow_dataset(ds)) + out = ctx.sql( + "SELECT time, AVG(temperature) AS temperature FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ).to_pandas() + expected = ds.temperature.sel(lat=ds.lat[ds.lat > 0]).mean("lat").compute() + np.testing.assert_allclose(out["temperature"].values, expected.values) + + +def test_dask_from_map_over_fragments(ds): + dd = pytest.importorskip("dask.dataframe") + + frags = xql.arrow_dataset(ds).get_fragments() + ddf = dd.from_map(lambda f: f.to_table().to_pandas(), frags) + assert len(ddf.compute()) == 20 * 6 + + +def test_polars_scan_pushdown_round_trip(ds): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + out = ( + lf.filter(pl.col("lat") > 0) + .group_by("time") + .agg(pl.col("temperature").mean()) + .sort("time") + .collect() + ) + expected = ( + ds.temperature.sel(lat=ds.lat[ds.lat > 0]).mean("lat").compute().values + ) + np.testing.assert_allclose(out["temperature"].to_numpy(), expected) + + +def test_polars_result_round_trips_to_xarray(ds): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + frame = ( + lf.group_by("time") + .agg(pl.col("temperature").mean().alias("temperature")) + .sort("time") + .collect() + ) + # Polars DataFrames export Arrow via the PyCapsule protocol, so the + # engine-agnostic round-trip works unchanged. + out = xql.to_dataset(frame, template=ds) + assert list(out.dims) == ["time"] + assert out.sizes["time"] == 20 + + +# -- Projection: what is returned vs what is read ---------------------------- + + +def test_filter_only_columns_are_read_but_not_returned(): + src = xr.Dataset( + {"a": (["i"], np.arange(100.0)), "b": (["i"], np.arange(100.0) * 2)}, + coords={"i": np.arange(100.0)}, + ) + counter = _ChunkCounter() + dataset = XarrayPushdownDataset(src, {"i": 10}, _iteration_callback=counter) + table = dataset.to_table(columns=["a"], filter=pc.field("b") >= 100.0) + assert table.column_names == ["a"] + assert table.num_rows == 50 + # The filter column is read alongside the projected one, nothing else. + assert set(counter.column_sets) == {("a", "b")} + + +def test_empty_projection_is_a_real_projection(counted): + dataset, counter = counted + table = dataset.scanner(columns=[]).to_table() + assert table.num_columns == 0 + assert table.num_rows == 100 * 4 + + # With a filter, only the filter's column is read, still zero returned. + counter.column_sets.clear() + table = dataset.scanner( + columns=[], filter=pc.field("t2m") < 40.0 + ).to_table() + assert table.num_columns == 0 + assert table.num_rows == 40 # values 0..39: ten hours x four latitudes + assert set(counter.column_sets) == {("t2m",)} + + +# -- Counting and pruning ---------------------------------------------------- + +_T0 = pd.Timestamp("2020-01-01 03:00") +_T1 = pd.Timestamp("2020-01-02 03:00") + + +def _ts(value): + return pa.scalar(value, type=pa.timestamp("ns")) + + +@pytest.mark.parametrize( + "predicate, rows, reads", + [ + # No filter: pure arithmetic, no chunk is read. + (None, 100 * 4, 0), + # [03:00, 27:00): chunks 0 and 2 are boundary, chunk 1 is provably + # inside the range and must be counted arithmetically. + ( + (pc.field("time") >= _ts(_T0)) & (pc.field("time") < _ts(_T1)), + 24 * 4, + 2, + ), + # A data-variable filter carries no coordinate guarantee: every + # chunk is a boundary chunk, and the count must still be row-exact. + (pc.field("t2m") >= 200.0, 200, 10), + # Unsatisfiable: everything pruned, nothing read. + (pc.field("lat") > 100.0, 0, 0), + ], + ids=["unfiltered", "strict-chunks", "data-variable", "unsatisfiable"], +) +def test_count_rows_contract(counted, predicate, rows, reads): + dataset, counter = counted + assert dataset.count_rows(filter=predicate) == rows + assert len(counter.blocks) == reads + + +def test_count_rows_broad_filter_stays_arithmetic(): + # 100k single-element chunks with a filter keeping nearly all of + # them: the hierarchical strictness analysis must prove whole + # buckets at once instead of scanning every survivor. + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset( + {"v": (["step"], np.arange(100_000.0))}, + coords={"step": np.arange(100_000.0)}, + ), + {"step": 1}, + _iteration_callback=lambda b, n: reads.append(b), + ) + assert dataset.count_rows(filter=pc.field("step") >= 100.0) == 99_900 + assert len(reads) <= 2 # at most the bucket-edge chunk + + +def test_count_rows_cross_dimension_refinement(): + # Paired ranges across two chunked dims: per-dim pruning keeps the + # union of each dim's survivors (so the cross combinations too); + # the strictness pass must prune the crosses and count exactly. + t = np.arange(200.0) + lat = np.linspace(-45.0, 45.0, 20) + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset( + {"v": (["t", "lat"], np.arange(200.0 * 20).reshape(200, 20))}, + coords={"t": t, "lat": lat}, + ), + {"t": 10, "lat": 10}, + _iteration_callback=lambda b, n: reads.append(b), + ) + predicate = ((pc.field("t") < 5.0) & (pc.field("lat") < -40.0)) | ( + (pc.field("t") >= 190.0) & (pc.field("lat") > 40.0) + ) + n = dataset.count_rows(filter=predicate) + expected = int( + ( + ((t[:, None] < 5) & (lat[None, :] < -40)) + | ((t[:, None] >= 190) & (lat[None, :] > 40)) + ).sum() + ) + assert n == expected + # Per-dim pruning alone keeps 4 chunk combos (2 t-chunks x 2 + # lat-chunks); cross-dim refinement drops the 2 crosses. + assert len(reads) <= 2 + + +def test_poisoned_coordinates_prune_conservatively(): + # A NaN (or NaT) inside a coordinate chunk voids its range guarantee: + # that chunk must be scanned, never pruned or counted arithmetically, + # and the result must match the oracle (NaN compares False). + x = np.array([0.0, 1.0, np.nan, 3.0, 4.0, 5.0]) + reads: list = [] + dataset = XarrayPushdownDataset( + xr.Dataset({"v": (["x"], np.arange(6.0))}, coords={"x": x}), + {"x": 2}, + _iteration_callback=lambda b, n: reads.append(b), + ) + assert dataset.count_rows(filter=pc.field("x") > 0.5) == 4 + assert any(b["x"] == slice(2, 4) for b in reads) # the NaN chunk + + t = pd.to_datetime(["2020-01-01", "2020-01-02", "NaT", "2020-01-04"]) + nat = XarrayPushdownDataset( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": t}), + {"time": 2}, + ) + lo = _ts(pd.Timestamp("2020-01-02")) + assert nat.count_rows(filter=pc.field("time") >= lo) == 2 + + +# -- Scan scheduling knobs: shape of the work, never the result --------------- + + +@pytest.mark.parametrize("coalesce_rows", [None, 10 * 4, 30 * 4, 10_000]) +def test_coalesce_results_identical(coalesce_rows): + source = _hourly_grid() + dataset = XarrayPushdownDataset( + source, {"time": 10}, coalesce_rows=coalesce_rows + ) + predicate = (pc.field("time") >= _ts(_T0)) & ( + pc.field("time") < _ts(pd.Timestamp("2020-01-03 07:00")) + ) + table = dataset.to_table(filter=predicate) + assert table.num_rows == 52 * 4 + expected = source.t2m.isel(time=slice(3, 55)).values.ravel() + np.testing.assert_array_equal( + np.sort(table["t2m"].to_numpy()), np.sort(expected) + ) + + +def test_coalesce_merges_consecutive_chunk_runs(): + source = _hourly_grid() + reads: list[dict] = [] + dataset = XarrayPushdownDataset( + source, + {"time": 10}, + coalesce_rows=30 * 4, # up to 3 source chunks per read + _iteration_callback=lambda b, n: reads.append(b), + ) + # An unfiltered scan of 10 chunks arrives as ceil(10/3) = 4 reads. + assert dataset.to_table().num_rows == 400 + assert len(reads) == 4 + spans = sorted((b["time"].start, b["time"].stop) for b in reads) + assert spans == [(0, 30), (30, 60), (60, 90), (90, 100)] + + # Pruning still applies before merging: a filter keeping chunks + # 0-2 and 7-9 yields one merged read per consecutive run. + reads.clear() + keep = (pc.field("time") < _ts(pd.Timestamp("2020-01-02 06:00"))) | ( + pc.field("time") >= _ts(pd.Timestamp("2020-01-03 22:00")) + ) + table = dataset.to_table(filter=keep) + assert table.num_rows == (30 + 30) * 4 + spans = sorted((b["time"].start, b["time"].stop) for b in reads) + assert spans == [(0, 30), (70, 100)] + + +def test_coalesce_only_affects_scanner_not_fragments(): + dataset = XarrayPushdownDataset( + _hourly_grid(), {"time": 10}, coalesce_rows=10_000 + ) + # Fragment consumers (DataFusion, dask) keep one fragment per source + # chunk for their own parallelism. + assert len(dataset.get_fragments()) == 10 + + +def test_prefetch_bytes_scan_reads_every_block_once(): + source = xr.Dataset( + {"v": (["step"], np.arange(10_000.0))}, + coords={"step": np.arange(10_000.0)}, + ) + reads: list = [] + # 100 chunks of 100 rows x 16 bytes/row = 1600 bytes per block; a + # 4000-byte budget throttles admission well below the 8 threads + # prefetch allows. The scan must still visit every block exactly + # once and return the full table. + dataset = XarrayPushdownDataset( + source, + {"step": 100}, + prefetch=8, + prefetch_bytes=4_000, + _iteration_callback=lambda b, n: reads.append(b), + ) + table = dataset.to_table() + assert table.num_rows == 10_000 + assert len(reads) == 100 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"prefetch": 0}, # at most one worker: the pool-less path + {"prefetch": -1}, + {"coalesce_rows": 0}, + {"prefetch_bytes": 0}, + ], + ids=lambda kw: ( + next(iter(kw.items()))[0] + "=" + str(next(iter(kw.values()))) + ), +) +def test_degenerate_tuning_values_still_scan_exactly(kwargs): + # Degenerate knob values may degrade the schedule (prefetch <= 1 + # takes the pool-less path) but never the answer, and never hang. + src = xr.Dataset( + {"v": (["i"], np.arange(100.0))}, coords={"i": np.arange(100.0)} + ) + dataset = XarrayPushdownDataset(src, {"i": 10}, **kwargs) + assert dataset.to_table().num_rows == 100 + assert dataset.count_rows(filter=pc.field("i") >= 50.0) == 50 + + +@pytest.mark.parametrize("batch_size", [0, -5]) +def test_non_positive_batch_size_fails_at_construction(batch_size): + # batch_size cannot degrade gracefully: a zero size never advances + # the zero-column scan's row loop, so it is rejected eagerly. + src = xr.Dataset( + {"v": (["i"], np.arange(100.0))}, coords={"i": np.arange(100.0)} + ) + with pytest.raises(ValueError, match="batch_size"): + XarrayPushdownDataset(src, {"i": 10}, batch_size=batch_size) + + +# -- Laziness, re-scannability, concurrency ----------------------------------- + + +def test_abandoned_scanner_does_not_wedge_later_scans(counted): + dataset, counter = counted + batches = dataset.scanner().to_batches() + next(batches) + del batches # LIMIT-style early stop: consumer walks away mid-scan + counter.blocks.clear() + assert dataset.count_rows() == 400 + table = dataset.to_table(columns=["t2m"]) + assert table.num_rows == 400 + + +def test_concurrent_scans_are_isolated_and_exact(): + # Engines scan from their own worker threads; simultaneous filtered + # scans over one dataset must not cross-talk. + src = xr.Dataset( + {"v": (["i"], np.arange(20_000.0))}, coords={"i": np.arange(20_000.0)} + ) + dataset = XarrayPushdownDataset(src, {"i": 500}, prefetch=4) + results: list = [None] * 8 + errors: list = [] + + def worker(k): + try: + lo = k * 1000.0 + predicate = (pc.field("i") >= lo) & (pc.field("i") < lo + 3000.0) + results[k] = dataset.to_table(filter=predicate).num_rows + except Exception as exc: # noqa: BLE001 — reported by the assert + errors.append(f"{k}: {exc}") + + threads = [threading.Thread(target=worker, args=(k,)) for k in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(60) + assert not any(t.is_alive() for t in threads), "a scan wedged" + assert not errors + assert results == [3000] * 8 + + # Two batch iterators over the same dataset, consumed alternately, + # stay exact and independent. + a = dataset.scanner().to_batches() + b = dataset.scanner().to_batches() + rows_a = rows_b = 0 + exhausted_a = exhausted_b = False + while not (exhausted_a and exhausted_b): + batch = next(a, None) + if batch is None: + exhausted_a = True + else: + rows_a += batch.num_rows + batch = next(b, None) + if batch is None: + exhausted_b = True + else: + rows_b += batch.num_rows + assert rows_a == rows_b == 20_000 + + +# -- Pool lifecycle ------------------------------------------------------------ + + +def test_prefetch_pool_threads_all_started_at_construction(ds): + dataset = XarrayPushdownDataset(ds, {"time": 5}, prefetch=6) + # Every pool thread must exist before the first scan: a thread + # spawned later, from inside an engine's scan callback, is exactly + # the deadlock the pre-spawn exists to prevent. Thread accounting + # is only visible on the executor's private state. + assert len(dataset._pool._threads) == 6 + + +def test_prefetch_pool_shut_down_when_dataset_dies(ds): + import gc + + dataset = XarrayPushdownDataset(ds, {"time": 5}, prefetch=4) + pool = dataset._pool + del dataset + gc.collect() + # A shut-down executor refuses new work — the observable contract + # that its threads have been told to exit. + with pytest.raises(RuntimeError): + pool.submit(lambda: None) diff --git a/tests/test_arrow_dataset_integration.py b/tests/test_arrow_dataset_integration.py new file mode 100644 index 00000000..f7880e2b --- /dev/null +++ b/tests/test_arrow_dataset_integration.py @@ -0,0 +1,405 @@ +"""Integration tests: the arrow-dataset contract against real cloud stores. + +The same contract ``test_arrow_dataset.py`` pins on synthetic data, +exercised at scale: real Zarr stores read over the network, consumed +through the real engines (DuckDB, Polars, DataFusion). Assertions are +plan-shape — exactly which source chunks each query reads, exact row +counts, values matched against a direct xarray read of the same window — +so a pruning or fast-path regression fails long before it shows up in +wall-clock noise. + +Two axes, both extensible: + +* ``CASES`` — one :class:`StoreCase` per dataset. Expectations are + computed from the case's declared cadence and windows plus the store's + own coordinates. A new dataset is a config entry, provided its + temporal dimension is named ``time``, its cadence is regular, and its + grid is dense; anything else needs test changes, not just a case. +* ``ENGINES`` — engine name to runner function; every runner executes + the same windowed aggregation through its engine's idiomatic path + (DuckDB SQL, Polars lazy expressions, DataFusion SQL). + +Reads anonymously from public buckets. Excluded from the CI unit run +(``pytest -m "not integration"``); run deliberately with +``pytest -m integration tests/test_arrow_dataset_integration.py``. +""" + +import threading +from dataclasses import dataclass, field + +import pandas as pd +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + +pytestmark = pytest.mark.integration + + +@dataclass(frozen=True) +class StoreCase: + id: str + url: str + variable: str # the variable every scan queries + other_variable: str # registered alongside; must never be read + time_chunk: int # registration granularity, steps per chunk + steps_per_day: int # from the store's cadence + window_start: str # a day-window anchor, chunk-aligned + month: str # a chunk-aligned month for the arithmetic count + bbox: dict = field(default_factory=dict) # dim -> (lo, hi), inclusive + storage_options: dict = field(default_factory=dict) + + +CASES = [ + StoreCase( + id="arco-era5", + url="gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3", + variable="2m_temperature", + other_variable="10m_u_component_of_wind", + time_chunk=1, + steps_per_day=24, + window_start="2020-01-01", + month="2020-01", + bbox={"latitude": (36, 44), "longitude": (350, 360)}, + storage_options={"token": "anon"}, + ), +] + + +# -- Engine runners ------------------------------------------------------------ +# One function per engine: run count(*) + avg(variable) over a half-open +# time window (plus an optional bbox) through the engine's idiomatic +# path, returning (rows, mean). Adding an engine is one function and one +# registry entry; tests parametrize over the registry. + + +def _sql(case, t0, t1, bbox) -> str: + conds = [f"time >= TIMESTAMP '{t0}'", f"time < TIMESTAMP '{t1}'"] + for dim, (lo, hi) in bbox.items(): + conds.append(f'"{dim}" BETWEEN {lo} AND {hi}') + return ( + f'SELECT count(*), avg("{case.variable}") FROM t ' + f"WHERE {' AND '.join(conds)}" + ) + + +def _duckdb_scan(case, dataset, t0, t1, bbox): + duckdb = pytest.importorskip("duckdb") + con = duckdb.connect() + con.register("t", dataset) + n, mean = con.execute(_sql(case, t0, t1, bbox)).fetchone() + return int(n), float(mean) + + +def _datafusion_scan(case, dataset, t0, t1, bbox): + from datafusion import SessionContext + + ctx = SessionContext() + ctx.register_dataset("t", dataset) + row = ctx.sql(_sql(case, t0, t1, bbox)).to_pandas().iloc[0] + return int(row.iloc[0]), float(row.iloc[1]) + + +def _polars_scan(case, dataset, t0, t1, bbox): + pl = pytest.importorskip("polars") + lf = pl.scan_pyarrow_dataset(dataset).filter( + (pl.col("time") >= t0.to_pydatetime()) + & (pl.col("time") < t1.to_pydatetime()) + ) + for dim, (lo, hi) in bbox.items(): + lf = lf.filter((pl.col(dim) >= lo) & (pl.col(dim) <= hi)) + out = lf.select( + pl.len().alias("n"), pl.col(case.variable).mean().alias("mean") + ).collect() + return int(out["n"][0]), float(out["mean"][0]) + + +# DataFusion consumes the dataset through get_fragments() — one +# fragment per source chunk, which scanner-level coalescing +# deliberately leaves untouched (see the contract in +# test_arrow_dataset.py). The other engines scan the whole dataset. +_datafusion_scan.consumes_fragments = True # type: ignore[attr-defined] + +ENGINES = { + "duckdb": _duckdb_scan, + "polars": _polars_scan, + "datafusion": _datafusion_scan, +} + + +@pytest.fixture(params=sorted(ENGINES), ids=str) +def scan(request): + """The engine runner under test.""" + return ENGINES[request.param] + + +# -- Dataset cases ------------------------------------------------------------- + + +@pytest.fixture(scope="module", params=CASES, ids=lambda c: c.id) +def case(request) -> StoreCase: + c: StoreCase = request.param + # The expectations below assume windows land on chunk boundaries; + # reject a miswritten case loudly instead of failing tests obscurely. + assert c.steps_per_day % c.time_chunk == 0, ( + f"{c.id}: steps_per_day must be a multiple of time_chunk" + ) + steps_from_midnight = ( + pd.Timestamp(c.window_start) - pd.Timestamp(c.window_start).normalize() + ) / (pd.Timedelta(days=1) / c.steps_per_day) + assert steps_from_midnight % c.time_chunk == 0, ( + f"{c.id}: window_start is not chunk-aligned" + ) + return c + + +@pytest.fixture(scope="module") +def source(case) -> xr.Dataset: + ds = xr.open_zarr( + case.url, + chunks=None, + storage_options=case.storage_options, + consolidated=True, + )[[case.variable, case.other_variable]] + # Chunk boundaries are laid out from the store's own time origin; + # both declared anchors must land on one or the fast-path/read + # expectations below are silently wrong for this case. + origin = pd.Timestamp(ds.time.values[0]) + chunk_span = pd.Timedelta(days=1) / case.steps_per_day * case.time_chunk + for name in ("window_start", "month"): + anchor = pd.Timestamp(getattr(case, name)) + assert (anchor - origin) % chunk_span == pd.Timedelta(0), ( + f"{case.id}: {name} is not aligned to the store's chunk grid" + ) + return ds + + +def _tracked(case, source, variables=None, **kwargs): + """A pushdown dataset over ``variables`` recording every block read.""" + reads: list = [] + column_sets: list = [] + dataset = XarrayPushdownDataset( + source[variables or [case.variable]], + {"time": case.time_chunk}, + prefetch=16, + _iteration_callback=lambda b, names: ( + reads.append(b), + column_sets.append(tuple(names)), + ), + **kwargs, + ) + return dataset, reads, column_sets + + +def _grid_cells(case, source, use_bbox=True) -> int: + """Cells per time step, inside the case's bbox unless disabled.""" + cells = 1 + for dim in source[case.variable].dims: + if dim == "time": + continue + vals = source[dim].values + if use_bbox and dim in case.bbox: + lo, hi = case.bbox[dim] + cells *= int(((vals >= lo) & (vals <= hi)).sum()) + else: + cells *= len(vals) + return cells + + +def _day_window(case, day=0, days=1) -> tuple[pd.Timestamp, pd.Timestamp]: + start = pd.Timestamp(case.window_start) + pd.Timedelta(days=day) + return start, start + pd.Timedelta(days=days) + + +def _day_chunks(case: StoreCase) -> int: + return case.steps_per_day // case.time_chunk + + +# -- The contract, engine by engine --------------------------------------------- + + +def test_windowed_scan_prunes_and_is_exact(scan, case, source): + # One day + bbox: every engine must push the window down so only the + # day's chunks are read, and row count and mean must match a direct + # xarray read of the same window. + dataset, reads, _ = _tracked(case, source) + t0, t1 = _day_window(case) + n, mean = scan(case, dataset, t0, t1, case.bbox) + + assert len(reads) == _day_chunks(case), "the engine did not prune" + assert n == case.steps_per_day * _grid_cells(case, source) + + window = source[case.variable].sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + for dim, (lo, hi) in case.bbox.items(): + keep = window[dim][(window[dim] >= lo) & (window[dim] <= hi)] + window = window.sel({dim: keep}) + # rel=1e-5: the store is float32, so the two sides accumulate the + # mean in different orders and dtypes. + assert mean == pytest.approx(float(window.mean()), rel=1e-5) + + +def test_projection_reads_only_the_referenced_variable(scan, case, source): + # Two variables registered; a query touching one must never read the + # other, whichever engine decides the column set. + dataset, _, column_sets = _tracked( + case, source, variables=[case.variable, case.other_variable] + ) + t0, t1 = _day_window(case) + scan(case, dataset, t0, t1, {}) + read = {name for cols in column_sets for name in cols} + assert case.variable in read + assert case.other_variable not in read + + +def test_dataset_is_rescannable_across_queries(scan, case, source): + # One wrapper, two queries of different shapes: the second scan must + # see fresh state, not a consumed stream or stale pruning (a bbox + # left over) from the first. + dataset, reads, _ = _tracked(case, source) + t0, t1 = _day_window(case) + n, _ = scan(case, dataset, t0, t1, case.bbox) + assert n == case.steps_per_day * _grid_cells(case, source) + assert len(reads) == _day_chunks(case) + + reads.clear() + t0, t1 = _day_window(case, day=1) + n, _ = scan(case, dataset, t0, t1, {}) + assert n == case.steps_per_day * _grid_cells(case, source, use_bbox=False) + assert len(reads) == _day_chunks(case) + + +def test_coalescing_merges_consecutive_reads(scan, case, source): + # The same day window in a handful of merged reads instead of one + # per chunk; the answer must not change. + cells = _grid_cells(case, source, use_bbox=False) + merge_chunks = 6 # chunks per merged read + dataset, reads, _ = _tracked( + case, + source, + coalesce_rows=merge_chunks * case.time_chunk * cells, + ) + t0, t1 = _day_window(case) + n, _ = scan(case, dataset, t0, t1, {}) + assert n == case.steps_per_day * cells + if getattr(scan, "consumes_fragments", False): + # Fragment consumers read one source chunk per fragment. + assert len(reads) == _day_chunks(case) + else: + assert len(reads) == -(-_day_chunks(case) // merge_chunks) # ceil + + +def test_concurrent_queries_stay_exact(scan, case, source): + # Engines scan from worker threads; two simultaneous queries over + # disjoint days of one wrapper must both come back exact. + dataset, _, _ = _tracked(case, source) + cells = _grid_cells(case, source, use_bbox=False) + results: list = [None, None] + errors: list = [] + + def worker(day): + try: + t0, t1 = _day_window(case, day) + results[day] = scan(case, dataset, t0, t1, {})[0] + except Exception as exc: # noqa: BLE001 — reported by the assert + errors.append(str(exc)) + + # daemon: a genuinely wedged scan must fail the assert, not keep + # the interpreter alive after pytest reports it. + threads = [ + threading.Thread(target=worker, args=(d,), daemon=True) for d in (0, 1) + ] + for t in threads: + t.start() + for t in threads: + t.join(120) + assert not any(t.is_alive() for t in threads), "a scan wedged" + assert not errors + assert results == [case.steps_per_day * cells] * 2 + + +# -- Dataset-level fast paths (no engine in the loop) --------------------------- + + +def test_multiday_global_window_prunes_to_its_chunks(case, source): + # Windows wider than one day prune exactly (7 days of hourly chunks: + # 168 reads on arco-era5). A property of the dataset's own scanner, + # so it runs once here instead of once per engine; batches are + # streamed and dropped to keep the 100M-row scan out of memory. + import pyarrow as pa + import pyarrow.compute as pc + + dataset, reads, _ = _tracked(case, source) + days = 7 + t0, t1 = _day_window(case, days=days) + predicate = (pc.field("time") >= pa.scalar(t0, type=pa.timestamp("ns"))) & ( + pc.field("time") < pa.scalar(t1, type=pa.timestamp("ns")) + ) + scanner = dataset.scanner(columns=[case.variable], filter=predicate) + rows = sum(batch.num_rows for batch in scanner.to_batches()) + assert len(reads) == days * _day_chunks(case), "wide window mispruned" + assert rows == days * case.steps_per_day * _grid_cells( + case, source, use_bbox=False + ) + + +def test_count_rows_fast_path_reads_nothing(case, source): + # A chunk-aligned month: every surviving chunk is provably inside + # the range, so the count is pure arithmetic. + import pyarrow as pa + import pyarrow.compute as pc + + dataset, reads, _ = _tracked(case, source) + lo = pd.Timestamp(case.month) + hi = lo + pd.offsets.MonthBegin(1) + predicate = (pc.field("time") >= pa.scalar(lo, type=pa.timestamp("ns"))) & ( + pc.field("time") < pa.scalar(hi, type=pa.timestamp("ns")) + ) + steps = (hi - lo) / pd.Timedelta(days=1) * case.steps_per_day + count = dataset.count_rows(filter=predicate) + assert reads == [], "count_rows fast path must not read data" + assert count == int(steps) * _grid_cells(case, source, use_bbox=False) + + +def test_polars_lazy_roundtrip_window_reads_only_its_blocks(case, source): + # Lazy round-trip: construction reads nothing with template coords; + # a one-day window reads only its own coalesced blocks. Polars is + # the one engine whose results re-execute over this dataset — + # DuckDB relations cannot (see limitations.md) and DataFusion's + # chunked round-trip lives on its native path. + pl = pytest.importorskip("polars") + if tuple(int(p) for p in pl.__version__.split(".")[:2]) >= (1, 43): + # polars 1.43 regressed streaming re-execution over pyarrow + # datasets: this window takes >10 minutes against 7s on 1.42. + pytest.skip("polars >= 1.43 streaming re-execution regression") + cells = _grid_cells(case, source, use_bbox=False) + merge_chunks = 6 + dataset, reads, _ = _tracked( + case, + source, + variables=[case.variable], + coalesce_rows=merge_chunks * case.time_chunk * cells, + ) + lf = pl.scan_pyarrow_dataset(dataset) + reads.clear() + lazy = xql.to_dataset( + lf, + template=source[[case.variable]], + chunks={"time": case.steps_per_day}, + coords="template", + ) + assert reads == [], "lazy construction must not read the source" + t0, t1 = _day_window(case) + value = float( + lazy[case.variable] + .sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + .mean() + .compute() + ) + assert len(reads) == -(-_day_chunks(case) // merge_chunks) + oracle = float( + source[case.variable] + .sel(time=slice(t0, t1 - pd.Timedelta("1ns"))) + .mean() + ) + assert value == pytest.approx(oracle, rel=1e-6) diff --git a/tests/test_cft.py b/tests/test_cft.py new file mode 100644 index 00000000..24b362bb --- /dev/null +++ b/tests/test_cft.py @@ -0,0 +1,178 @@ +"""Unit tests for the cftime module (cftime ↔ Arrow bridge).""" + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +from xarray_sql import cftime as cft +from xarray_sql.df import _parse_schema + + +# -- Fixtures --------------------------------------------------------------- + + +@pytest.fixture +def ds_360day(): + """Synthetic 360-day calendar dataset.""" + import cftime + + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) + + +# -- Detection helpers ------------------------------------------------------ + + +class TestDetection: + def test_is_cftime_detects_cftime_array(self, rasm_ds): + assert cft.is_cftime(rasm_ds.coords["time"].values) + + def test_is_cftime_rejects_datetime64(self): + assert not cft.is_cftime(pd.date_range("2020-01-01", periods=10).values) + + def test_is_cftime_rejects_float(self): + assert not cft.is_cftime(np.array([1.0, 2.0, 3.0])) + + def test_is_cftime_index_detects_cftime(self, rasm_ds): + assert cft.is_cftime_index(rasm_ds, "time") + + def test_is_cftime_index_rejects_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert not cft.is_cftime_index(ds, "time") + + def test_is_cftime_index_rejects_nonexistent(self, rasm_ds): + assert not cft.is_cftime_index(rasm_ds, "nonexistent") + + +# -- Calendar classification ------------------------------------------------ + + +class TestCalendarClassification: + def test_calendar_returns_noleap(self, rasm_ds): + assert cft.calendar(rasm_ds, "time") == "noleap" + + def test_calendar_returns_360_day(self, ds_360day): + assert cft.calendar(ds_360day, "time") == "360_day" + + def test_calendar_returns_none_for_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert cft.calendar(ds, "time") is None + + def test_noleap_is_gregorian_like(self): + assert cft.is_gregorian_like("noleap") + assert cft.is_gregorian_like("standard") + assert cft.is_gregorian_like("proleptic_gregorian") + assert cft.is_gregorian_like("all_leap") + + def test_360_day_is_not_gregorian_like(self): + assert not cft.is_gregorian_like("360_day") + assert not cft.is_gregorian_like("julian") + + +# -- Numeric conversion ----------------------------------------------------- + + +class TestConversion: + def test_to_microseconds_returns_int64(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert us.dtype == np.int64 + + def test_to_microseconds_is_monotonic(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert np.all(np.diff(us) > 0) + + def test_to_microseconds_length_matches(self, rasm_ds): + values = rasm_ds.coords["time"].values + assert len(cft.to_microseconds(values)) == len(values) + + def test_to_offsets_returns_int64(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert offsets.dtype == np.int64 + + def test_to_offsets_is_monotonic(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert np.all(np.diff(offsets) > 0) + + def test_convert_for_field_gregorian_like(self, rasm_ds): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + result = cft.convert_for_field(rasm_ds.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) + + def test_convert_for_field_non_gregorian(self, ds_360day): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + result = cft.convert_for_field(ds_360day.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) + + +# -- Arrow schema helpers --------------------------------------------------- + + +class TestArrowField: + def test_gregorian_like_produces_timestamp_us(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + assert field.type == pa.timestamp("us") + assert field.metadata[b"xarray:calendar"] == b"noleap" + assert field.metadata[b"xarray:units"] == cft.DEFAULT_UNITS.encode() + + def test_non_gregorian_produces_int64(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + assert field.type == pa.int64() + assert field.metadata[b"xarray:calendar"] == b"360_day" + + +# -- Partition bounds ------------------------------------------------------- + + +class TestPartitionBounds: + def test_gregorian_like_returns_timestamp_ns_tag(self, rasm_ds): + values = rasm_ds.coords["time"].values[:10] + lo, hi, tag = cft.partition_bounds(values) + assert tag == "timestamp_ns" + assert lo < hi + + def test_non_gregorian_returns_int64_tag(self, ds_360day): + values = ds_360day.coords["time"].values + lo, hi, tag = cft.partition_bounds(values) + assert tag == "int64" + assert lo < hi + + def test_out_of_int64_range_returns_none(self): + # Year-1 gregorian dates exceed the int64 nanosecond range, so no + # pruning bound can be reported; the caller skips the dimension. + values = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ).values + assert cft.partition_bounds(values) is None + + +# -- Integration with _parse_schema ---------------------------------------- + + +class TestParseSchemaIntegration: + def test_noleap_produces_timestamp_us(self, rasm_ds): + schema = _parse_schema(rasm_ds[["Tair"]]) + time_field = schema.field("time") + assert time_field.type == pa.timestamp("us") + assert time_field.metadata[b"xarray:calendar"] == b"noleap" + + def test_360day_produces_int64(self, ds_360day): + schema = _parse_schema(ds_360day) + time_field = schema.field("time") + assert time_field.type == pa.int64() + assert time_field.metadata[b"xarray:calendar"] == b"360_day" + + def test_datetime64_unchanged(self): + ds = xr.tutorial.open_dataset("air_temperature") + schema = _parse_schema(ds) + time_field = schema.field("time") + assert pa.types.is_timestamp(time_field.type) + assert time_field.metadata is None # no xarray: metadata for native diff --git a/tests/test_df.py b/tests/test_df.py new file mode 100644 index 00000000..5185b573 --- /dev/null +++ b/tests/test_df.py @@ -0,0 +1,769 @@ +import tracemalloc + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +from xarray_sql.df import ( + DEFAULT_BATCH_SIZE, + _ensure_default_indexes, + _parse_schema, + block_slices, + compute_chunks, + dataset_to_record_batch, + explode, + from_map, + from_map_batched, + group_vars_by_dims, + iter_record_batches, + partition_metadata, + pivot, +) +from xarray_sql.reader import read_xarray, read_xarray_table + + +def test_explode_cardinality(air): + dss = explode(air) + assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()]) + + +def test_explode_dim_sizes_one(air): + chunks = {"time": 240} + ds = next(iter(explode(air))) + for k, v in chunks.items(): + assert k in ds.dims + assert v == ds.sizes[k] + + +def test_explode_data_equal_one_first(air): + ds = next(iter(explode(air))) + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + assert air.isel(iselection).equals(ds) + + +def test_explode_data_equal_one_last(air): + dss = list(explode(air)) + ds = dss[-1] + + # For the last chunk, we need to calculate where it actually starts + # The original logic slice(0, s) only works for the first chunk + iselection = {} + for dim in ds.dims: + # Get chunk boundaries + chunk_bounds = np.cumsum((0,) + air.chunks[dim]) + # Last chunk index + last_chunk_idx = len(air.chunks[dim]) - 1 + # Calculate actual start and end positions + start = chunk_bounds[last_chunk_idx] + end = chunk_bounds[last_chunk_idx + 1] + iselection[dim] = slice(start, end) + + assert air.isel(iselection).equals(ds) + + +def test_block_slices_scalar_dataset_yields_single_block(): + # A dimensionless dataset (e.g. scalar metadata variables) has exactly + # one block: the whole, empty selection. + ds = xr.Dataset({"projection": ((), 0)}) + assert list(block_slices(ds)) == [{}] + + +def test_block_slices_scalar_ignores_irrelevant_chunks(): + ds = xr.Dataset({"projection": ((), 0)}) + assert list(block_slices(ds, chunks={"time": 4})) == [{}] + + +def test_block_slices_filters_chunk_keys_to_dataset_dims(air_small): + # A chunk key for a dimension the dataset doesn't have is ignored, + # rather than raising. + base = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) + extra = list( + block_slices( + air_small, chunks={"time": 4, "lat": 3, "lon": 4, "absent": 2} + ) + ) + assert len(extra) == len(base) + + +def test_block_slices_dimensional_unchunked_raises(): + # A dataset with dimensions but no chunking is still a user error. + ds = xr.Dataset({"v": (["x"], np.arange(3))}, coords={"x": np.arange(3)}) + with pytest.raises(AssertionError): + list(block_slices(ds)) + + +def test_from_map_basic(): + def make_df(x): + return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]}) + + result = from_map(make_df, [1, 2, 3]) + assert isinstance(result, pa.Table) + assert len(result) == 6 + assert result.column_names == ["value", "index"] + + +def test_from_map_multiple_iterables(): + def add_values(x, y): + return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]}) + + result = from_map(add_values, [1, 2], [10, 20]) + assert isinstance(result, pa.Table) + assert len(result) == 2 + + df = result.to_pandas() + assert list(df["sum"]) == [11, 22] + + +def test_from_map_with_args(): + def multiply_and_add(x, multiplier, add_value): + return pd.DataFrame({"result": [x * multiplier + add_value]}) + + result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) + assert isinstance(result, pa.Table) + assert len(result) == 3 + + df = result.to_pandas() + assert list(df["result"]) == [12, 14, 16] + + +def test_from_map_with_pyarrow_tables(): + def make_arrow_table(x): + df = pd.DataFrame({"value": [x]}) + return pa.Table.from_pandas(df) + + result = from_map(make_arrow_table, [1, 2, 3]) + assert isinstance(result, pa.Table) + assert len(result) == 3 + + +def test_iter_record_batches_splits_into_multiple_batches(air_small): + """iter_record_batches should emit >1 batch when partition exceeds batch_size.""" + schema = _parse_schema(air_small) + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + ds_block = air_small.isel(block) + total_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes])) + + small_batch = 16 # force many small batches + batches = list( + iter_record_batches(ds_block, schema, batch_size=small_batch) + ) + + assert len(batches) == -(-total_rows // small_batch) # ceiling division + assert all(b.num_rows <= small_batch for b in batches) + assert sum(b.num_rows for b in batches) == total_rows + + +def test_iter_record_batches_matches_dataset_to_record_batch(air_small): + """Concatenating all iter_record_batches output must equal dataset_to_record_batch.""" + schema = _parse_schema(air_small) + dim_cols = [f.name for f in schema if f.name in air_small.dims] + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + ds_block = air_small.isel(block) + + batches = list(iter_record_batches(ds_block, schema, batch_size=16)) + actual_df = ( + pa.Table.from_batches(batches) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + expected_df = ( + dataset_to_record_batch(ds_block, schema) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + pd.testing.assert_frame_equal(actual_df, expected_df) + + +def test_iter_record_batches_projection_drops_cftime_dim(): + """A projection that drops a cftime dim (e.g. time under GROUP BY level) + must not call schema.field() for it. The dim is absent from the projected + schema, and cftime coords take the convert_for_field(schema.field(name)) + path, so an unguarded lookup raised KeyError during batch reading.""" + cftime = pytest.importorskip("cftime") + times = np.array( + [cftime.DatetimeGregorian(2020, m, 1) for m in (1, 2, 3)], dtype=object + ) + ds = xr.Dataset( + {"air": (["time", "lat"], np.arange(3 * 2, dtype=float).reshape(3, 2))}, + coords={"time": times, "lat": [0.0, 1.0]}, + ) + full = _parse_schema(ds) + projected = pa.schema( + [full.field("lat"), full.field("air")] + ) # time dropped + table = pa.Table.from_batches( + list(iter_record_batches(ds, projected, batch_size=16)), projected + ) + assert table.schema.names == ["lat", "air"] + assert table.num_rows == 6 + + +def test_iter_record_batches_default_batch_size(): + """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch.""" + ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2)) + schema = _parse_schema(ds) + total_rows = int(np.prod([ds.sizes[d] for d in ds.sizes])) + assert total_rows <= DEFAULT_BATCH_SIZE, "fixture too large — adjust isel" + batches = list(iter_record_batches(ds, schema)) + assert len(batches) == 1 + assert batches[0].num_rows == total_rows + + +def test_dataset_to_record_batch_matches_pivot(air_small): + """dataset_to_record_batch should contain the same rows as pivot. + + Row ordering may differ (pivot uses ds.dims key order; dataset_to_record_batch + uses the data variable's own dim order). Both orderings are valid for SQL, so + we sort by the coordinate columns before comparing. + """ + schema = _parse_schema(air_small) + dim_cols = [f.name for f in schema if f.name in air_small.dims] + blocks = list( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + + for block in blocks: + ds_block = air_small.isel(block) + actual_df = ( + dataset_to_record_batch(ds_block, schema) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + expected_df = ( + pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + + pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False) + + +def test_dataset_to_record_batch_column_order(air_small): + """Output column order must match schema (dims first, then data vars).""" + schema = _parse_schema(air_small) + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + batch = dataset_to_record_batch(air_small.isel(block), schema) + assert batch.schema.names == schema.names + + +def test_dataset_to_record_batch_row_count(air_small): + """Row count must equal the product of the block dimension sizes.""" + schema = _parse_schema(air_small) + chunks = {"time": 4, "lat": 3, "lon": 4} + for block in block_slices(air_small, chunks=chunks): + ds_block = air_small.isel(block) + expected_rows = int( + np.prod([ds_block.sizes[d] for d in ds_block.sizes]) + ) + batch = dataset_to_record_batch(ds_block, schema) + assert batch.num_rows == expected_rows + + +def test_from_map_batched_basic_functionality(air_small): + blocks = list( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + + first_block_df = pivot(air_small.isel(blocks[0])) + expected_schema = pa.Schema.from_pandas(first_block_df) + + reader = from_map_batched( + pivot, + [air_small.isel(block) for block in blocks], + schema=expected_schema, + ) + + assert isinstance(reader, pa.RecordBatchReader) + assert reader.schema == expected_schema + + batches = list(reader) + assert len(batches) > 0 + for batch in batches: + assert batch.schema == expected_schema + assert len(batch) > 0 + + +def adding_function(x, y): + """Simple function that adds two values and returns a DataFrame.""" + result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]}) + return result + + +def test_from_map_batched_multiple_iterables(): + x_values = [1, 2, 3, 4, 5] + y_values = [10, 20, 30, 40, 50] + + expected_schema = pa.schema( + [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())] + ) + + reader = from_map_batched( + adding_function, x_values, y_values, schema=expected_schema + ) + table = reader.read_all() + df = table.to_pandas() + + expected_df = pd.DataFrame( + { + "x": x_values, + "y": y_values, + "sum": [x + y for x, y in zip(x_values, y_values)], + } + ) + pd.testing.assert_frame_equal(df, expected_df) + + +def test_from_map_batched_with_args_and_kwargs(): + def multiply_and_add(x, multiplier, offset=0): + return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]}) + + values = [1, 2, 3] + expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())]) + + reader = from_map_batched( + multiply_and_add, values, args=(2,), offset=5, schema=expected_schema + ) + table = reader.read_all() + df = table.to_pandas() + + expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]}) + pd.testing.assert_frame_equal(df, expected_df) + + +def test_from_map_batched_empty_iterables(): + empty_schema = pa.schema([("value", pa.int64())]) + + reader = from_map_batched( + lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema + ) + batches = list(reader) + assert len(batches) == 0 + + +def test_from_map_batched_consistency_with_regular_map(air_small): + blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3})) + datasets = [air_small.isel(block) for block in blocks] + + first_df = pivot(datasets[0]) + schema = pa.Schema.from_pandas(first_df) + + reader = from_map_batched(pivot, datasets, schema=schema) + batched_table = reader.read_all() + + regular_dfs = [pivot(ds) for ds in datasets] + regular_table = pa.Table.from_pandas( + pd.concat(regular_dfs, ignore_index=True) + ) + + assert batched_table.schema == regular_table.schema + assert len(batched_table) == len(regular_table) + + batched_df = ( + batched_table.to_pandas() + .sort_values(["time", "lat", "lon"]) + .reset_index(drop=True) + ) + regular_df = ( + regular_table.to_pandas() + .sort_values(["time", "lat", "lon"]) + .reset_index(drop=True) + ) + + pd.testing.assert_frame_equal(batched_df, regular_df) + + +def test_from_map_batched_integration_with_datafusion_via_read_xarray(): + air = xr.tutorial.open_dataset("air_temperature") + air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15)) + air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8}) + + arrow_stream = read_xarray( + air_chunked, chunks={"time": 25, "lat": 5, "lon": 8} + ) + + assert hasattr(arrow_stream, "schema") + assert hasattr(arrow_stream, "__iter__") + + table = arrow_stream.read_all() + assert len(table) > 0 + + expected_columns = {"time", "lat", "lon", "air"} + actual_columns = set(table.column_names) + assert expected_columns.issubset(actual_columns) + + +def test_read_xarray_loads_one_chunk_at_a_time(large_ds): + tracemalloc.stop() # reset any state left by a previously-failed test + tracemalloc.start() + try: + iterable = read_xarray(large_ds) + first_size, first_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + + sizes, peaks = [], [] + + first_chunk = large_ds.isel(next(block_slices(large_ds))) + chunk_size = first_chunk.nbytes + + # Creating the iterator should be inexpensive -- less than one chunk. + # We multiply by constant factors because chunks have additional overhead + assert first_size < chunk_size * 3 + assert first_peak < chunk_size * 6 + + for it in iterable: + _ = it + cur_size, cur_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + sizes.append(cur_size) + peaks.append(cur_peak) + + for size in sizes: + # iter_record_batches' whole-partition fast path holds the + # data-variable arrays (≈1× chunk) plus repeat/tile-expanded + # coordinate columns (n_dims × 8 bytes × rows, ≈1.5× chunk + # for this 3-dim float64 dataset) for the partition being + # streamed; batches themselves are zero-copy slices. + assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" + assert chunk_size * 4.0 > size, f"size {size} unexpectedly high" + + for peak in peaks: + # Peak adds transient buffers on top of the steady state: + # np.repeat/np.tile intermediates for the coordinate columns + # and Arrow's from_pandas null scan; the first batch of each + # chunk is highest (Dask compute overhead). Observed ~5.04× + # on macOS. + assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" + assert chunk_size * 6.5 > peak, f"peak {peak} unexpectedly high" + + assert max(peaks) < large_ds.nbytes + finally: + tracemalloc.stop() + + +def test_read_xarray_table_memory_bounds(large_ds): + """read_xarray_table should not materialise data at registration time. + + Registration should only hold coordinate arrays and Rust partition metadata + (no data variables). Peak memory during a full-table query should be a + small fraction of the whole dataset — i.e. partitions are processed without + loading all of them simultaneously. + """ + from datafusion import SessionContext + + first_chunk = large_ds.isel(next(block_slices(large_ds))) + chunk_size = first_chunk.nbytes + + tracemalloc.stop() # reset any state left by a previously-failed test + # --- Registration phase --- + tracemalloc.start() + try: + table = read_xarray_table(large_ds) + reg_size, reg_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + + # The lazy generator only materialises coord arrays (~O(dim sizes)) and + # factory closure objects — no data arrays. Both metrics should be well + # below one chunk of data. + assert reg_size < chunk_size, ( + f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: " + "data may have been loaded eagerly" + ) + assert reg_peak < chunk_size * 2, ( + f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})" + ) + + # --- Query phase --- + ctx = SessionContext() + ctx.register_table("weather", table) + ctx.sql( + "SELECT AVG(temperature), AVG(precipitation) FROM weather" + ).collect() + _, query_peak = tracemalloc.get_traced_memory() + + # tracemalloc measures Python-heap allocations, which include Arrow + # buffer copies and object overhead on top of the raw data. The + # observed peak is typically 1.1–1.5× the raw dataset size; we use + # 2× as a generous bound that would still catch catastrophic regressions + # (e.g. loading all partitions twice simultaneously). + assert query_peak < large_ds.nbytes * 2, ( + f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: " + "may be holding excessive data in memory" + ) + finally: + tracemalloc.stop() + + +# --------------------------------------------------------------------------- +# compute_chunks: arithmetic replacement for ds.chunk(...).chunks. +# Dask serves as the source of truth. +# --------------------------------------------------------------------------- + + +def _dask_chunks(ds: xr.Dataset, chunks: dict) -> dict: + rechunked = ds.copy(data=None, deep=False).chunk(chunks) + return {str(k): tuple(v) for k, v in rechunked.chunks.items()} + + +def _normalise(result: dict) -> dict: + return {str(k): tuple(v) for k, v in result.items()} + + +def _simple_ds(shape: tuple[int, ...], dims: tuple[str, ...]) -> xr.Dataset: + return xr.Dataset( + {"v": (dims, np.zeros(shape))}, + coords={d: np.arange(s) for d, s in zip(dims, shape)}, + ) + + +@pytest.mark.parametrize( + "ds,chunks", + [ + # Even divide on a single dim. + (_simple_ds((10,), ("x",)), {"x": 5}), + # Uneven divide: trailing remainder chunk. + (_simple_ds((10,), ("x",)), {"x": 3}), + # Requested chunk size larger than the dim → single chunk. + (_simple_ds((5,), ("x",)), {"x": 100}), + # Multi-dim spec with a dim left unspecified (kept as one chunk). + (_simple_ds((4, 6), ("x", "y")), {"x": 2}), + # Multi-dim spec rechunking every dim. + (_simple_ds((7, 11, 13), ("a", "b", "c")), {"a": 3, "b": 4, "c": 5}), + ], +) +def test_compute_chunks_matches_dask(ds, chunks): + assert _normalise(compute_chunks(ds, chunks)) == _dask_chunks(ds, chunks) + + +def test_compute_chunks_preserves_existing_dask_chunking(): + # When the dataset is already dask-backed, rechunking one dim must + # leave other dims' existing chunk tuples alone. + ds = _simple_ds((4, 5), ("x", "y")).chunk({"x": 1, "y": 2}) + chunks = {"x": 2} + assert _normalise(compute_chunks(ds, chunks)) == _dask_chunks(ds, chunks) + + +def test_compute_chunks_tuples_sum_to_dim_size(): + # Dask-independent invariant: every per-dim chunk tuple must fully + # cover its dimension. + ds = _simple_ds((7, 11, 13), ("a", "b", "c")) + result = compute_chunks(ds, {"a": 3, "b": 4, "c": 5}) + for dim, tup in result.items(): + assert sum(tup) == ds.sizes[dim] + + +def test_iter_record_batches_large_string_dim_coord(): + """A string dim coord big enough that pa.array returns a ChunkedArray. + + Pivoting tiles a string dimension coordinate across every row of the + partition; for a few million rows pyarrow's numpy-unicode conversion + returns a ChunkedArray, which RecordBatch.from_arrays rejects. + Regression: found by the forecast-skill benchmark (a 2-model x 3.3M-row + window) streaming through the pyarrow dataset protocol into DuckDB. + """ + n_x = 1_700_000 # 2 * n_x rows: comfortably past the chunking threshold + ds = xr.Dataset( + {"value": (("model", "x"), np.zeros((2, n_x), dtype="float32"))}, + coords={"model": ["pangu", "graphcast"], "x": np.arange(n_x)}, + ) + schema = _parse_schema(ds) + got_rows = 0 + for batch in iter_record_batches(ds, schema, batch_size=DEFAULT_BATCH_SIZE): + got_rows += batch.num_rows + assert got_rows == 2 * n_x + + +# -- Object-dtype and out-of-ns-range coordinate support -------------------- + + +def _field_type(schema, name): + return schema.field(name).type + + +def test_parse_schema_maps_object_string_data_var_to_string(): + # A string variable arrives as numpy object dtype; _parse_schema must not + # hand it to pa.from_numpy_dtype (which raises "Unsupported numpy type 17"). + ds = xr.Dataset( + {"label": (["x"], np.array(["a", "b"], dtype=object))}, + coords={"x": [1, 2]}, + ) + schema = _parse_schema(_ensure_default_indexes(ds)) + assert _field_type(schema, "label") == pa.string() + + +def test_parse_schema_maps_object_string_coord_to_string(): + # A string dimension coordinate (e.g. station names) is object dtype too. + ds = xr.Dataset( + {"v": (["station"], [1.0, 2.0])}, + coords={"station": np.array(["A", "B"], dtype=object)}, + ) + schema = _parse_schema(_ensure_default_indexes(ds)) + assert _field_type(schema, "station") == pa.string() + + +def test_partition_metadata_skips_out_of_ns_datetime(): + # datetime64 coordinates outside the datetime64[ns] range (pre-1678 / + # post-2262) cannot be represented as int64 nanoseconds, so partition + # pruning must be skipped for that dimension rather than raising + # OverflowError. Registration must still succeed. + times = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ).to_datetimeindex(time_unit="us", unsafe=True) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) # must not raise + + assert len(meta) == len(blocks) + # "time" is unpruneable here, so it is omitted from every partition. + assert all("time" not in m for m in meta) + + +def test_parse_schema_all_null_object_var_stays_null(): + # An all-null object column has no data to infer a type from; let null be + # null rather than coercing it to a string column. + ds = _ensure_default_indexes( + xr.Dataset( + {"label": (["x"], np.array([None, None], dtype=object))}, + coords={"x": [1, 2]}, + ) + ) + schema = _parse_schema(ds) + assert pa.types.is_null(schema.field("label").type) + + +def test_partition_metadata_prunes_cftime_coord(): + # cftime dimension coordinates must produce pruning bounds; previously the + # object-dtype skip shadowed the cftime branch, silently disabling pruning. + times = xr.date_range( + "2000-01-01", periods=4, freq="1D", calendar="noleap", use_cftime=True + ) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) + + assert all("time" in m for m in meta) + for m in meta: + _, _, tag = m["time"] + assert tag == "timestamp_ns" + + +def test_partition_metadata_skips_ancient_cftime(): + # Ancient gregorian cftime dates overflow the int64 nanosecond range, so + # pruning must be skipped for that dim (no raise, dim omitted). + times = xr.date_range( + "0001-01-01", periods=3, freq="100YS", use_cftime=True + ) + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(3.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) # must not raise + + assert all("time" not in m for m in meta) + + +def test_string_dataset_round_trips_through_record_batch(): + # The schema fix must also flow through the batch builders: a string + # column has to materialize as an Arrow string array, not error out. + ds = _ensure_default_indexes( + xr.Dataset( + {"label": (["x"], np.array(["a", "b", "c", "d"], dtype=object))}, + coords={"x": [10, 20, 30, 40]}, + ) + ) + schema = _parse_schema(ds) + + batch = dataset_to_record_batch(ds, schema) + assert batch.schema.field("label").type == pa.string() + assert batch.column("label").to_pylist() == ["a", "b", "c", "d"] + + # The streaming path must agree with the one-shot path. + streamed = pa.Table.from_batches( + list(iter_record_batches(ds, schema, batch_size=2)), schema=schema + ) + assert streamed.column("label").to_pylist() == ["a", "b", "c", "d"] + + +def test_partition_metadata_in_range_datetime_still_pruned(): + # Regression guard: ordinary datetimes must keep producing timestamp_ns + # bounds so filter pushdown still works after the overflow fix. + times = pd.date_range("2000-01-01", periods=4, freq="D") + ds = _ensure_default_indexes( + xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times}) + ) + blocks = list(block_slices(ds, chunks={"time": 2})) + + meta = partition_metadata(ds, blocks) + + assert all("time" in m for m in meta) + for m in meta: + _, _, tag = m["time"] + assert tag == "timestamp_ns" + + +class TestGroupVarsByDims: + def test_single_dim_group(self): + ds = xr.Dataset( + { + "a": (["x", "y"], np.zeros((2, 3))), + "b": (["x", "y"], np.ones((2, 3))), + } + ) + groups = group_vars_by_dims(ds) + assert groups == {("x", "y"): ["a", "b"]} + + def test_multiple_dim_groups(self): + ds = xr.Dataset( + { + "surface": (["time", "lat", "lon"], np.zeros((2, 3, 4))), + "upper": ( + ["time", "lat", "lon", "level"], + np.zeros((2, 3, 4, 5)), + ), + } + ) + groups = group_vars_by_dims(ds) + assert set(groups.keys()) == { + ("time", "lat", "lon"), + ("time", "lat", "lon", "level"), + } + assert groups[("time", "lat", "lon")] == ["surface"] + assert groups[("time", "lat", "lon", "level")] == ["upper"] + + def test_empty_dataset(self): + assert group_vars_by_dims(xr.Dataset()) == {} + + def test_includes_scalar_group(self): + """Scalar (0-dim) variables group under the empty dims tuple.""" + ds = xr.Dataset( + {"band": (["y", "x"], np.zeros((2, 3))), "projection": ((), 0)} + ) + groups = group_vars_by_dims(ds) + assert groups == {("y", "x"): ["band"], (): ["projection"]} + + def test_ignores_coords(self): + """Coordinate variables shouldn't be returned as groups.""" + ds = xr.Dataset( + {"v": (["x"], np.arange(3))}, + coords={"x": np.arange(3), "label": ("x", ["a", "b", "c"])}, + ) + groups = group_vars_by_dims(ds) + assert groups == {("x",): ["v"]} diff --git a/tests/test_ds.py b/tests/test_ds.py new file mode 100644 index 00000000..3682dacf --- /dev/null +++ b/tests/test_ds.py @@ -0,0 +1,690 @@ +"""Tests for the SQL -> xarray reverse path. + +Covers the user-facing contract of ``ctx.sql(...).to_dataset(...)``: + +* Wrapper behavior on the object returned by ``ctx.sql`` and DataFusion + method passthrough. +* Round-trip identity across varied source Datasets (one parametrized + ``assert_identical`` test, not eight per-aspect checks). +* Aggregation, ``dimension_columns`` inference, and the template / + ``template`` resolution rules (name or Dataset) with their error paths. +* Sparsity handling and ``fill_value`` dtype behavior. +* The vectorized-indexer fallback through xarray's adapter. + +The tests favor checking the user-visible contract (values, dims, +attrs) over the implementation path (call counts, internal class +identity), so the suite stays useful as the lazy backend evolves. +""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from xarray_sql import XarrayContext +from xarray_sql.ds import XarrayDataFrame + + +# --------------------------------------------------------------------------- +# Wrapper: ctx.sql(...) returns XarrayDataFrame +# --------------------------------------------------------------------------- + + +def test_ctx_sql_returns_xarray_dataframe(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql("SELECT * FROM air LIMIT 5") + assert isinstance(result, XarrayDataFrame) + + +def test_to_pandas_unchanged_behavior(air_dataset_small): + """Wrapped ``.to_pandas()`` is bit-for-bit equal to the un-wrapped path.""" + from datafusion import SessionContext + + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + wrapped = ctx.sql("SELECT * FROM air LIMIT 7").to_pandas() + raw = SessionContext.sql(ctx, "SELECT * FROM air LIMIT 7").to_pandas() + pd.testing.assert_frame_equal(wrapped, raw) + + +def test_passthrough_methods(air_dataset_small): + """DataFusion methods we did not override forward via ``__getattr__``.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql("SELECT * FROM air LIMIT 5") + names = [f.name for f in result.schema()] + assert {"lat", "lon", "time", "air"}.issubset(set(names)) + + +# --------------------------------------------------------------------------- +# Round-trip identity (parametrized over local + tutorial datasets) +# --------------------------------------------------------------------------- + + +def _clear_encoding(ds: xr.Dataset) -> xr.Dataset: + """Strip ``encoding`` from a Dataset and all its variables. + + Round-trip identity tests should not be coupled to encoding choices, + since template-recovery deliberately drops dtype-bound keys. + """ + ds = ds.copy() + for v in ds.variables.values(): + v.encoding.clear() + ds.encoding.clear() + return ds + + +def _load_tutorial(name: str) -> xr.Dataset | None: + """Return a small xarray tutorial Dataset, or None when unavailable. + + Used to widen round-trip coverage beyond the conftest fixtures without + requiring network in CI. Pooch caches downloads locally on first run. + """ + try: + return xr.tutorial.open_dataset(name) + except (OSError, ValueError, ImportError): + return None + + +@pytest.mark.parametrize( + "fixture_name", + ["air_dataset_small", "weather_dataset", "synthetic_dataset", "eraint_uvz"], +) +def test_round_trip_identity(request, fixture_name): + """``SELECT *`` round-trips to a Dataset that is ``assert_identical`` + to the source: values, dims, coord values, dtypes, non-dim coords, + and attrs all match (modulo coord ordering, normalized on both + sides). One test covers what was previously a fan of narrow checks, + parametrized over local fixtures and one xarray tutorial dataset. + """ + if fixture_name == "eraint_uvz": + source = _load_tutorial("eraint_uvz") + if source is None: + pytest.skip("eraint_uvz tutorial dataset unavailable") + source = source.chunk() + else: + source = request.getfixturevalue(fixture_name).copy() + source.attrs["round_trip_marker"] = "yes" + first_var = next(iter(source.data_vars)) + source[first_var].attrs["units"] = "test_units" + + ctx = XarrayContext() + ctx.from_dataset("t", source) + out = ctx.sql("SELECT * FROM t").to_dataset().compute() + + sort_keys = list(out.dims) + actual = _clear_encoding(out.sortby(sort_keys)) + expected = _clear_encoding(source.compute().sortby(sort_keys)) + xr.testing.assert_identical(actual, expected) + + +def test_aggregation_drops_dim(air_dataset_small): + """``GROUP BY lat, lon`` over time -> 2D Dataset with the alias.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql( + "SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon" + ).to_dataset() + assert set(out.dims) == {"lat", "lon"} + assert "air_avg" in out.data_vars + assert "air" not in out.data_vars + expected = ( + air_dataset_small.compute() + .sortby(["lat", "lon"]) + .mean(dim="time")["air"] + .values + ) + actual = out.sortby(["lat", "lon"])["air_avg"].values + np.testing.assert_allclose(actual, expected) + + +def test_aggregation_infers_dims(air_dataset_small): + """to_dataset() infers the surviving GROUP BY dim when dims is omitted.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + + # Grouping by the time coordinate keeps time as the sole dimension; the + # ORDER BY makes the result order deterministic so no sort is needed below. + out = ctx.sql( + 'SELECT "time", AVG("air") AS air FROM "air" ' + 'GROUP BY "time" ORDER BY "time"' + ).to_dataset() + assert set(out.dims) == {"time"} + assert "air" in out.data_vars + expected = air_dataset_small.compute().mean(dim=["lat", "lon"])["air"] + np.testing.assert_allclose(out["air"].values, expected.values) + + +def test_barrier_query_scans_source_once(air_dataset_small): + """A barrier plan (aggregation) executes the source exactly once. + + The lazy scan path re-runs the whole upstream plan for every coordinate + discovery and every variable access; for an aggregation -- which cannot push + an indexer filter below the GROUP BY -- that is pure re-computation of an + expensive scan. ``to_dataset()`` on a barrier plan must instead make a + single streamed pass over the source, and ``.compute()`` must trigger no + further reads. + """ + from xarray_sql.df import block_slices + from xarray_sql.reader import read_xarray_table + + reads: list = [] + table = read_xarray_table( + air_dataset_small, + chunks={"time": 6}, + _iteration_callback=lambda block, proj: reads.append(block), + ) + n_partitions = len(list(block_slices(air_dataset_small, {"time": 6}))) + + ctx = XarrayContext() + ctx.register_table("air", table) + ctx._registered_datasets["air"] = air_dataset_small + + out = ctx.sql( + "SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon" + ).to_dataset() + reads_after_construct = len(reads) + out.compute() + reads_after_compute = len(reads) + + # Exactly one pass over the source (each partition read once) ... + assert reads_after_construct == n_partitions + # ... and computing the materialized result re-reads nothing. + assert reads_after_compute == reads_after_construct + + +def test_order_by_direction_sets_dim_order(air_dataset_small): + """A barrier query's ORDER BY direction carries through to the Dataset + dimension order, rather than being force-sorted ascending. + + ``ORDER BY lat DESC`` must yield a strictly descending ``lat`` dimension, + with data still correctly aligned to those (descending) coordinates. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql( + "SELECT lat, AVG(air) AS air_avg FROM air GROUP BY lat ORDER BY lat DESC" + ).to_dataset() + + lat = out["lat"].values + assert (np.diff(lat) < 0).all(), f"expected descending lat, got {lat}" + + # Values stay aligned to the descending coordinate (scatter handles order). + expected = ( + air_dataset_small.compute() + .mean(dim=["time", "lon"])["air"] + .sortby("lat", ascending=False) + ) + np.testing.assert_allclose(out["air_avg"].values, expected.values) + + +def test_unfiltered_scan_skips_dim_discovery(air_dataset_small): + """``SELECT * FROM `` does no per-dim discovery scans. + + The pre-fix :func:`_build_lazy_scan` ran one full source scan per dim + via ``inner_df.select(col(d)).distinct().sort(...)`` -- a single-column + projection per dim, multiplied across the dim count. With a registered + template and an unfiltered plan, :func:`_maybe_template_coords` returns + the template's coord arrays directly and the discovery loop is skipped + entirely. Verified by the iteration callback's projection list: no read + where the projection is a single dim column. + """ + from xarray_sql.reader import read_xarray_table + + reads: list = [] + table = read_xarray_table( + air_dataset_small, + chunks={"time": 6}, + _iteration_callback=lambda block, proj: reads.append(list(proj)), + ) + ctx = XarrayContext() + ctx.register_table("air", table) + ctx._registered_datasets["air"] = air_dataset_small + + reads.clear() + out = ctx.sql('SELECT * FROM "air"').to_dataset() + single_dim_reads = [ + proj for proj in reads if set(proj) <= {"time", "lat", "lon"} + ] + assert not single_dim_reads, ( + f"fast path should skip per-dim discovery scans, got " + f"{single_dim_reads!r}" + ) + # The fast path must not break correctness on first access. + out.compute() # safety: the Dataset still resolves + + +def test_filtered_scan_still_uses_dim_discovery(air_dataset_small): + """A ``WHERE`` clause forbids the template-coord fast path. + + With a filter the result's coord extent is a subset of the source, so + reusing template coords would silently widen the output. The dispatch + in :func:`_maybe_template_coords` falls back to per-dim discovery in + that case, which is observable as a positive read count on the source. + """ + from xarray_sql.reader import read_xarray_table + + reads: list = [] + table = read_xarray_table( + air_dataset_small, + chunks={"time": 6}, + _iteration_callback=lambda block, proj: reads.append(block), + ) + ctx = XarrayContext() + ctx.register_table("air", table) + ctx._registered_datasets["air"] = air_dataset_small + + reads.clear() + ctx.sql('SELECT * FROM "air" WHERE lat > 30').to_dataset() + assert reads, "filtered scan must hit the source to discover coord extent" + + +def test_fast_path_uses_scanned_tables_coords_not_user_template( + air_dataset_small, +): + """Fast path sources coord values from the scanned table, not ``template=``. + + With multiple registered Datasets, a user may pass ``template=other`` + for metadata recovery while the query scans a different registered + table. The coord values must come from the **scanned** table's + registered Dataset; using ``other``'s coords would silently widen the + output if their coord ranges differ. + """ + other = air_dataset_small.isel(lat=slice(0, 5)) + assert other.sizes["lat"] != air_dataset_small.sizes["lat"] + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small, chunks={"time": 24}) + ctx.from_dataset("other", other, chunks={"time": 24}) + out = ctx.sql('SELECT * FROM "air"').to_dataset( + dims=["time", "lat", "lon"], template=other + ) + # Scanned table is "air", so lat must match air's full lat axis. + np.testing.assert_array_equal( + out["lat"].values, air_dataset_small["lat"].values + ) + + +def test_round_trip_preserves_descending_lat_on_lazy_path(air_dataset_small): + """Lazy round-trip preserves source dim order. + + NCEP ``air_temperature`` ships descending lat (75.0 -> 15.0). The + discovery path's ``.distinct().sort()`` previously flipped lat to + ascending on the lazy result. The template-coord fast path returns the + source's coord arrays as-is, so the descending order survives. + """ + ds = air_dataset_small + assert (np.diff(ds["lat"].values) < 0).all(), ( + "test relies on a descending-lat fixture" + ) + ctx = XarrayContext() + ctx.from_dataset("air", ds, chunks={"time": 24}) + lazy = ctx.sql('SELECT * FROM "air"').to_dataset() + np.testing.assert_array_equal(lazy["lat"].values, ds["lat"].values) + + +def test_chunks_argument_controls_partitioning(synthetic_dataset): + """``chunks`` controls eager-vs-chunked and inherits the source grid. + + The default ``"inherit"`` reuses the source's genuinely multi-chunk + dimensions, so the output chunk grid maps onto the source partitions; + ``chunks=None`` forces an eager, in-memory result. Both reproduce the source. + """ + import dask.array as da + + ctx = XarrayContext() + ctx.from_dataset("t", synthetic_dataset) + var = next(iter(synthetic_dataset.data_vars)) + + inherited = ctx.sql("SELECT * FROM t").to_dataset() + assert isinstance(inherited[var].data, da.Array) + # Output time chunks align to the source's time partitions. + assert inherited.chunksizes["time"] == synthetic_dataset.chunksizes["time"] + + eager = ctx.sql("SELECT * FROM t").to_dataset(chunks=None) + assert not isinstance(eager[var].data, da.Array) + + xr.testing.assert_allclose( + inherited.compute().sortby(["time", "lat", "lon"]), + synthetic_dataset.compute().sortby(["time", "lat", "lon"]), + ) + + +def test_chunks_auto_snaps_to_source_partitions(): + """``chunks="auto"`` coarsens to the byte budget but snaps chunk boundaries + to whole source partitions (so no chunk splits a source partition).""" + import dask + + # 12 source partitions of size 2 along time. + ds = xr.Dataset( + { + "v": ( + ("time", "x"), + np.arange(24 * 4, dtype="float64").reshape(24, 4), + ) + }, + coords={"time": np.arange(24), "x": np.arange(4)}, + ).chunk({"time": 2}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + + # block bytes = 8 * 2(time) * 4(x) = 64; target 192 -> merge 3 partitions. + with dask.config.set({"array.chunk-size": "192B"}): + out = ctx.sql("SELECT * FROM t").to_dataset(chunks="auto") + + time_chunks = out.chunksizes["time"] + assert all(c % 2 == 0 for c in time_chunks) # aligned to source size 2 + assert time_chunks[0] > 2 # genuinely coarsened + assert len(time_chunks) < 12 # fewer chunks than source partitions + + xr.testing.assert_allclose( + out.compute().sortby(["time", "x"]), + ds.compute().sortby(["time", "x"]), + ) + + +# --------------------------------------------------------------------------- +# dimension_columns / template resolution rules +# --------------------------------------------------------------------------- + + +def test_to_dataset_multi_registered_requires_explicit_template( + air_dataset_small, +): + """With more than one registered Dataset, the caller disambiguates by + passing a registered table name as ``template=``.""" + ctx = XarrayContext() + ctx.from_dataset("air1", air_dataset_small) + ctx.from_dataset("air2", air_dataset_small) + out = ctx.sql("SELECT * FROM air1").to_dataset(template="air1") + assert set(out.dims) == {"time", "lat", "lon"} + + +def test_to_dataset_infer_fails_when_no_dim_survives(air_dataset_small): + """A global aggregation leaves no registered dim in the result -> clear error.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(ValueError, match="dims cannot be inferred"): + ctx.sql("SELECT AVG(air) AS air_avg FROM air").to_dataset() + + +def test_template_accepts_name_or_dataset(air_dataset_small): + """``template=`` accepts either a registered table name or a Dataset + object, with equivalent metadata recovery.""" + other = air_dataset_small.copy() + other.attrs = {"flag": "other"} + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + ctx.from_dataset("other", other) + + by_name = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template="other" + ) + by_object = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template=other + ) + assert by_name.attrs == {"flag": "other"} + assert by_object.attrs == {"flag": "other"} + + +def test_template_unknown_name_raises(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(ValueError, match="not a registered table"): + ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template="missing" + ) + + +def test_template_recovers_var_encoding_strips_dtype(air_dataset_small): + """``zlib`` survives; dtype-bound keys are stripped (SQL may have cast).""" + ds = air_dataset_small.copy() + ds["air"].encoding = { + "zlib": True, + "dtype": "int16", + "_FillValue": -999, + "missing_value": -999, + } + ctx = XarrayContext() + ctx.from_dataset("air", ds) + out = ctx.sql("SELECT * FROM air").to_dataset(dims=["time", "lat", "lon"]) + assert out["air"].encoding.get("zlib") is True + assert "dtype" not in out["air"].encoding + assert "_FillValue" not in out["air"].encoding + assert "missing_value" not in out["air"].encoding + + +def test_template_aggregation_alias_no_attrs(air_dataset_small): + """``air_avg`` from ``AVG(air)`` does NOT inherit attrs from ``air``.""" + ds = air_dataset_small.copy() + ds["air"].attrs = {"units": "K"} + ctx = XarrayContext() + ctx.from_dataset("air", ds) + out = ctx.sql( + "SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon" + ).to_dataset() + assert "air_avg" in out.data_vars + assert out["air_avg"].attrs == {} + + +def test_to_dataset_explicit_template_overrides_auto_resolve( + air_dataset_small, +): + """Explicit template= wins over the auto-resolved FROM-clause table.""" + other = air_dataset_small.copy() + other.attrs = {"flag": "explicit"} + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], template=other + ) + assert out.attrs == {"flag": "explicit"} + + +# --------------------------------------------------------------------------- +# Lazy backend: value-level contract (not call counts) +# --------------------------------------------------------------------------- + + +def test_lazy_isel_int_round_trip(air_dataset_small): + """``isel(time=0)`` on the lazy result matches the eager equivalent.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + lazy = ctx.sql("SELECT * FROM air").to_dataset() + eager = lazy.compute() + actual = lazy["air"].isel(time=0).sortby(["lat", "lon"]).values + expected = eager["air"].isel(time=0).sortby(["lat", "lon"]).values + np.testing.assert_array_equal(actual, expected) + + +def test_lazy_isel_slice_round_trip(air_dataset_small): + """isel(time=slice(0, 3)) round-trip matches the source.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql("SELECT * FROM air").to_dataset() + actual = out["air"].isel(time=slice(0, 3)).sortby(["lat", "lon"]).values + expected = ( + air_dataset_small["air"] + .compute() + .isel(time=slice(0, 3)) + .sortby(["lat", "lon"]) + .values + ) + np.testing.assert_array_equal(actual, expected) + + +def test_lazy_outer_indexer_array(air_dataset_small): + """Fancy index along one dim works (IN-equivalent pushdown).""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + lazy = ctx.sql("SELECT * FROM air").to_dataset() + eager = lazy.compute() + indices = [0, 3, 5] + np.testing.assert_array_equal( + lazy["air"].isel(lat=indices).values, + eager["air"].isel(lat=indices).values, + ) + + +def test_lazy_compute_returns_eager(air_dataset_small): + """``.compute()`` returns an in-memory Dataset matching the source.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + out = ctx.sql("SELECT * FROM air").to_dataset().compute() + np.testing.assert_array_equal( + out.sortby(["time", "lat", "lon"])["air"].values, + air_dataset_small.compute() + .sortby(["time", "lat", "lon"])["air"] + .values, + ) + + +def test_vectorized_indexer_falls_back_via_xarray_adapter( + air_dataset_small, +): + """VectorizedIndexer paths through xarray's adapter to outer + gather. + + Our SQLBackendArray declares ``IndexingSupport.OUTER``, so xarray's + ``explicit_indexing_adapter`` converts vectorized indexers into a + series of outer reads followed by an in-memory numpy gather. The + public contract: values match the eager-computed equivalent. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + lazy = ctx.sql("SELECT * FROM air").to_dataset() + eager = lazy.compute() + + points_t = xr.DataArray([0, 3, 1], dims="point") + points_lat = xr.DataArray([2, 0, 5], dims="point") + np.testing.assert_array_equal( + lazy["air"].isel(time=points_t, lat=points_lat).values, + eager["air"].isel(time=points_t, lat=points_lat).values, + ) + + +# --------------------------------------------------------------------------- +# Sparsity handling and fill_value +# --------------------------------------------------------------------------- + + +def test_sparsity_result_default_filters_lazy(air_dataset_small): + """Default sparsity='result' keeps only filtered coords (lazy path).""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + threshold = float(air_dataset_small["lat"].values[5]) + out = ctx.sql(f"SELECT * FROM air WHERE lat > {threshold}").to_dataset() + assert (out["lat"].values > threshold).all() + assert out.sizes["lat"] < air_dataset_small.sizes["lat"] + + +def test_sparsity_template_full_grid(air_dataset_small): + """sparsity='template' reindexes to the full grid with NaN fills.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + threshold = float(air_dataset_small["lat"].values[5]) + out = ctx.sql(f"SELECT * FROM air WHERE lat > {threshold}").to_dataset( + sparsity="template" + ) + assert out.sizes["lat"] == air_dataset_small.sizes["lat"] + lat_vals = out["lat"].values + below_mask = lat_vals <= threshold + above_mask = lat_vals > threshold + below = out["air"].isel(lat=np.where(below_mask)[0]) + above = out["air"].isel(lat=np.where(above_mask)[0]) + assert np.isnan(below.values).all() + assert not np.isnan(above.values).any() + + +def test_sparsity_template_requires_template(air_dataset_small): + """No resolvable template -> sparsity='template' raises.""" + other = air_dataset_small.copy() + ctx = XarrayContext() + ctx.from_dataset("a", air_dataset_small) + ctx.from_dataset("b", other) + with pytest.raises(ValueError, match="requires template= to be supplied"): + ctx.sql("SELECT * FROM a").to_dataset( + dims=["time", "lat", "lon"], + sparsity="template", + ) + + +def test_sparsity_invalid_value_raises(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(ValueError, match="sparsity must be"): + ctx.sql("SELECT * FROM air").to_dataset( + dims=["time", "lat", "lon"], + sparsity="bogus", # type: ignore[arg-type] + ) + + +def test_sparsity_template_with_aggregation(air_dataset_small): + """sparsity='template' on an aggregation respects dimension_columns subset.""" + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + threshold = float(air_dataset_small["lat"].values[5]) + out = ctx.sql( + f""" + SELECT lat, lon, AVG(air) AS air_avg + FROM air + WHERE lat > {threshold} + GROUP BY lat, lon + """ + ).to_dataset(dims=["lat", "lon"], sparsity="template") + assert out.sizes["lat"] == air_dataset_small.sizes["lat"] + assert "time" not in out.dims + below_mask = out["lat"].values <= threshold + below = out["air_avg"].isel(lat=np.where(below_mask)[0]) + assert np.isnan(below.values).all() + + +def test_fill_value_int_upcasts_to_float(): + """fill_value=NaN forces float upcast on int columns -- documented.""" + ds = xr.Dataset( + {"v": (("lat", "lon"), np.arange(6, dtype=np.int64).reshape(3, 2))}, + coords={"lat": [0, 1, 2], "lon": [10, 11]}, + ).chunk({"lat": 3}) + ctx = XarrayContext() + ctx.from_dataset("t", ds) + out = ctx.sql("SELECT * FROM t WHERE lat > 0").to_dataset( + sparsity="template" + ) + assert np.issubdtype(out["v"].dtype, np.floating) + assert np.isnan(out["v"].sel(lat=0).values).all() + + +def test_fill_value_custom_preserves_int(air_dataset_small): + """Passing a typed sentinel preserves the data var's int dtype.""" + source = xr.Dataset( + { + "v": ( + ("lat", "lon"), + np.arange(6, dtype=np.int64).reshape(3, 2) + 1, + ), + }, + coords={"lat": [0, 1, 2], "lon": [10, 11]}, + ).chunk({"lat": 3}) + ctx = XarrayContext() + ctx.from_dataset("t", source) + out = ctx.sql("SELECT * FROM t WHERE lat > 0").to_dataset( + sparsity="template", fill_value=-1 + ) + assert np.issubdtype(out["v"].dtype, np.integer) + assert (out["v"].sel(lat=0).values == -1).all() + assert out["v"].sel(lat=2, lon=11).item() == 6 + + +def test_sparsity_template_then_metadata(air_dataset_small): + """sparsity='template' composes with template metadata recovery.""" + ds = air_dataset_small.copy() + ds.attrs = {"src": "tmpl"} + ds["air"].attrs = {"units": "K"} + ctx = XarrayContext() + ctx.from_dataset("air", ds) + threshold = float(ds["lat"].values[5]) + out = ctx.sql(f"SELECT * FROM air WHERE lat > {threshold}").to_dataset( + sparsity="template" + ) + assert out.attrs == {"src": "tmpl"} + assert out["air"].attrs == {"units": "K"} + assert out.sizes["lat"] == ds.sizes["lat"] diff --git a/tests/test_duckdb_backend.py b/tests/test_duckdb_backend.py new file mode 100644 index 00000000..0a004502 --- /dev/null +++ b/tests/test_duckdb_backend.py @@ -0,0 +1,387 @@ +"""Tests for the DuckDB engine adapter and the engine-agnostic round-trip. + +Covers the two seams of the multi-engine design: ``xql.register`` puts a +lazy Dataset on a DuckDB connection, DuckDB executes its own SQL dialect +(including extensions), and ``xql.to_dataset`` rebuilds a labeled +Dataset from the Arrow result. +""" + +import duckdb +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.duckdb import ( + XarrayArrowStream, + XarrayPushdownDataset, +) + + +@pytest.fixture +def ds() -> xr.Dataset: + np.random.seed(7) + time = pd.date_range("2021-01-01", periods=8, freq="h") + lat = np.linspace(-10.0, 10.0, 5) + lon = np.linspace(0.0, 40.0, 6) + temperature = 15 + 8 * np.random.randn(8, 5, 6) + precipitation = 10 * np.random.rand(8, 5, 6) + return xr.Dataset( + data_vars=dict( + temperature=(["time", "lat", "lon"], temperature), + precipitation=(["time", "lat", "lon"], precipitation), + ), + coords=dict(time=time, lat=lat, lon=lon), + attrs=dict(description="Synthetic weather."), + ).chunk({"time": 4}) + + +@pytest.fixture +def con(ds) -> duckdb.DuckDBPyConnection: + connection = duckdb.connect() + xql.register(connection, "weather", ds) + return connection + + +def test_full_scan_round_trips(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature, precipitation FROM weather " + "ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + xr.testing.assert_allclose(out, ds.compute()) + assert out.attrs == ds.attrs + + +def test_aggregation_round_trips_on_surviving_dims(con, ds): + rel = con.sql( + "SELECT time, AVG(temperature) AS temperature FROM weather " + "GROUP BY time ORDER BY time" + ) + out = xql.to_dataset(rel, template=ds) + + expected = ds["temperature"].mean(["lat", "lon"]).compute() + assert list(out.dims) == ["time"] + np.testing.assert_allclose(out["temperature"].values, expected.values) + + +def test_registered_table_is_requeryable(con): + first = con.sql("SELECT COUNT(*) AS n FROM weather").fetchone()[0] + second = con.sql("SELECT COUNT(*) AS n FROM weather").fetchone()[0] + assert first == second == 8 * 5 * 6 + + +def test_registration_is_lazy(ds): + reads: list = [] + stream = XarrayArrowStream( + ds, _iteration_callback=lambda b, p: reads.append(b) + ) + + con = duckdb.connect() + con.register("weather", stream) + assert reads == [] # registration reads no data + + con.sql("SELECT AVG(temperature) FROM weather").fetchall() + assert len(reads) > 0 # data was read during query execution + + +def test_where_filter_yields_sparse_result(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE lat > 0 ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + expected = ds[["temperature"]].sel(lat=ds.lat[ds.lat > 0]).compute() + xr.testing.assert_allclose(out, expected) + + +def test_template_sparsity_reindexes_to_full_extent(con, ds): + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather WHERE lat > 0" + ) + out = xql.to_dataset(rel, template=ds, sparsity="template") + + assert out.sizes == {"time": 8, "lat": 5, "lon": 6} + assert out["temperature"].isnull().sum() == 8 * 3 * 6 # lat <= 0 cells + + +def test_duckdb_dialect_and_join(con, ds): + # Engine-native SQL: DuckDB's date_part plus a join against a local + # relation — nothing xarray-sql has to understand. + con.sql("CREATE TABLE labels AS SELECT 0 AS h, 'midnight' AS label") + rel = con.sql( + "SELECT w.time, AVG(w.temperature) AS temperature, ANY_VALUE(l.label) AS label " + "FROM weather w JOIN labels l ON date_part('hour', w.time) = l.h " + "GROUP BY w.time" + ) + out = xql.to_dataset(rel, dims=["time"]) + assert out.sizes == {"time": 1} + + +def test_to_dataset_accepts_plain_arrow_table(ds): + table = pa.table( + { + "time": pd.date_range("2021-01-01", periods=3, freq="h"), + "temperature": [1.0, 2.0, 3.0], + } + ) + out = xql.to_dataset(table, dims=["time"]) + np.testing.assert_allclose(out["temperature"].values, [1.0, 2.0, 3.0]) + + +def test_to_dataset_requires_dims_or_template(): + table = pa.table({"a": [1, 2], "b": [3.0, 4.0]}) + with pytest.raises(ValueError, match="dims cannot be inferred"): + xql.to_dataset(table) + + +def test_to_dataset_rejects_missing_dim_column(): + table = pa.table({"a": [1, 2], "b": [3.0, 4.0]}) + with pytest.raises(ValueError, match="not columns of the result"): + xql.to_dataset(table, dims=["z"]) + + +def test_register_splits_mixed_dimension_variables(ds): + mixed = ds.assign(surface=ds["temperature"].isel(time=0, drop=True)) + con = duckdb.connect() + xql.register(con, "weather", mixed) + + n_full = con.sql("SELECT COUNT(*) FROM weather_time_lat_lon").fetchone()[0] + n_surface = con.sql("SELECT COUNT(*) FROM weather_lat_lon").fetchone()[0] + assert n_full == 8 * 5 * 6 + assert n_surface == 5 * 6 + + +def test_pushdown_dataset_rejects_mixed_dimension_variables(ds): + mixed = ds.assign(surface=ds["temperature"].isel(time=0, drop=True)) + with pytest.raises(ValueError, match="dimensions must be equal"): + XarrayPushdownDataset(mixed) + + +def _tracked_connection(ds): + """Register ds with an iteration callback; returns (con, reads).""" + reads: list = [] + dataset = XarrayPushdownDataset( + ds, _iteration_callback=lambda block, cols: reads.append((block, cols)) + ) + con = duckdb.connect() + con.register("weather", dataset) + return con, reads + + +def test_projection_pushdown_skips_unrequested_variables(ds): + con, reads = _tracked_connection(ds) + con.sql("SELECT AVG(temperature) FROM weather").fetchall() + assert reads # data was read + for _, cols in reads: + assert "precipitation" not in cols + + +def test_filter_pushdown_prunes_chunks(ds): + # ds is chunked {"time": 4} -> 2 chunks; this predicate covers only + # the first chunk, so the second is never loaded. + con, reads = _tracked_connection(ds) + n = con.sql( + "SELECT COUNT(*) FROM weather WHERE time < '2021-01-01 04:00:00'" + ).fetchone()[0] + assert n == 4 * 5 * 6 + assert len(reads) == 1 + + +def test_pushed_filter_is_applied_exactly(ds): + # DuckDB trusts pushed comparison filters and does not re-apply + # them, so the scan itself must enforce the predicate row-exactly — + # including inside chunks that pruning keeps. + con, _ = _tracked_connection(ds) + out = con.sql( + "SELECT COUNT(*) FROM weather " + "WHERE time = '2021-01-01 02:00:00' AND lat > 0" + ).fetchone()[0] + expected = int( + (ds.time == np.datetime64("2021-01-01T02:00:00")).sum() + * (ds.lat > 0).sum() + * ds.sizes["lon"] + ) + assert out == expected + + +def test_filter_on_variable_outside_projection(ds): + # The filter references `temperature`, the projection only `lat`; + # the scan must widen its columns to evaluate the predicate. + con, _ = _tracked_connection(ds) + got = con.sql( + "SELECT COUNT(DISTINCT lat) FROM weather WHERE temperature > 20" + ).fetchone()[0] + expected = len( + np.unique( + ds.lat.values[np.where((ds.temperature > 20).any(["time", "lon"]))] + ) + ) + assert got == expected + + +def test_or_and_in_filters_round_trip(ds): + con, _ = _tracked_connection(ds) + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE lat < -5 OR lat > 5 ORDER BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + mask = (ds.lat < -5) | (ds.lat > 5) + expected = ds[["temperature"]].sel(lat=ds.lat[mask]).compute() + xr.testing.assert_allclose(out, expected) + + +def test_pushdown_dataset_rejects_unchunked_dataset(ds): + with pytest.raises(ValueError, match="must be chunked"): + XarrayPushdownDataset(ds.compute()) + + +def test_fully_pruned_scan_returns_empty(con, ds): + n = con.sql( + "SELECT COUNT(*) FROM weather WHERE time >= '2022-01-01'" + ).fetchone()[0] + assert n == 0 + rel = con.sql( + "SELECT time, lat, lon, temperature FROM weather " + "WHERE time >= '2022-01-01'" + ) + out = xql.to_dataset(rel, template=ds) + assert out.sizes.get("time", 0) == 0 + + +def test_limit_terminates_early(con): + rows = con.sql("SELECT time, temperature FROM weather LIMIT 5").fetchall() + assert len(rows) == 5 + + +def test_descending_coordinate_pruning_is_correct(ds): + # Latitude stored north→south, like most rasters and ERA5. + flipped = ds.isel(lat=slice(None, None, -1)).chunk({"lat": 2}) + con = duckdb.connect() + xql.register(con, "weather", flipped) + got = con.sql("SELECT COUNT(*) FROM weather WHERE lat > 4").fetchone()[0] + expected = int((ds.lat > 4).sum()) * ds.sizes["time"] * ds.sizes["lon"] + assert got == expected + + +def test_integer_and_string_variables_round_trip(): + ds = xr.Dataset( + { + "klass": (["y", "x"], np.arange(12, dtype=np.uint8).reshape(3, 4)), + "label": ( + ["y", "x"], + np.array([["a"] * 4, ["b"] * 4, ["c"] * 4]), + ), + }, + coords={"y": np.arange(3), "x": np.arange(4)}, + ).chunk({"y": 2}) + con = duckdb.connect() + xql.register(con, "grid", ds) + rows = con.sql( + "SELECT label, SUM(klass) AS total FROM grid " + "WHERE klass >= 4 GROUP BY label ORDER BY label" + ).fetchall() + assert rows == [("b", 22), ("c", 38)] + + +def test_finely_chunked_dimension_uses_bucketed_pruning(): + # 5000 single-step time chunks exceeds the shadow fanout (1024), so + # pruning goes through the coarse-then-refine path; an equality in + # the middle of the axis must load exactly one chunk. + n = 5000 + ds = xr.Dataset( + {"v": (["time", "x"], np.random.rand(n, 2))}, + coords={ + "time": pd.date_range("2000-01-01", periods=n, freq="h"), + "x": np.arange(2), + }, + ).chunk({"time": 1}) + reads: list = [] + dataset = XarrayPushdownDataset( + ds, _iteration_callback=lambda block, cols: reads.append(block) + ) + con = duckdb.connect() + con.register("t", dataset) + + got = con.sql( + "SELECT COUNT(*) FROM t WHERE time = '2000-03-15 07:00:00'" + ).fetchone()[0] + assert got == 2 + assert len(reads) == 1 + + # A range spanning most of the axis stays correct (refinement is + # skipped when it cannot pay for itself). + reads.clear() + got = con.sql( + "SELECT COUNT(*) FROM t WHERE time >= '2000-01-01 12:00:00'" + ).fetchone()[0] + assert got == (n - 12) * 2 + + +def test_register_kwargs_are_forwarded(ds): + con = duckdb.connect() + xql.register(con, "weather", ds, prefetch=1, batch_size=7) + n = con.sql("SELECT COUNT(*) FROM weather").fetchone()[0] + assert n == 8 * 5 * 6 + + +def test_register_dispatches_to_datafusion(): + # The same entry point serves the default engine. + ctx = xql.XarrayContext() + small = xr.Dataset( + {"v": (["x"], np.arange(4.0))}, coords={"x": np.arange(4)} + ).chunk({"x": 2}) + xql.register(ctx, "t", small) + out = ctx.sql("SELECT x, v FROM t ORDER BY x").to_dataset() + np.testing.assert_allclose(out["v"].values, np.arange(4.0)) + + +def test_register_rejects_unknown_connection(ds): + with pytest.raises(TypeError, match="No xarray-sql engine adapter"): + xql.register(object(), "weather", ds) + + +def test_nan_coordinate_chunk_is_not_pruned(): + # NaN in a chunk's coordinate must disable pruning for that span, + # never poison the range guarantee (which would silently drop rows). + ds = xr.Dataset( + {"v": (["lat"], np.arange(6.0))}, + coords={"lat": [np.nan, 5.0, 10.0, 20.0, 30.0, 40.0]}, + ).chunk({"lat": 2}) + con = duckdb.connect() + xql.register(con, "t", ds) + assert con.sql("SELECT v FROM t WHERE lat = 5.0").fetchall() == [(1.0,)] + + +def test_cftime_dataset_aggregates_under_projection(): + cftime = pytest.importorskip("cftime") + + times = xr.date_range( + "2000-01-01", periods=6, calendar="360_day", use_cftime=True + ) + ds = xr.Dataset( + {"v": (["time"], np.arange(6.0))}, coords={"time": times} + ).chunk({"time": 3}) + con = duckdb.connect() + xql.register(con, "t", ds) + # The scan projects only `v`; the cftime dim column is absent from + # the scan schema but still shapes the iteration. + assert con.sql("SELECT SUM(v) FROM t").fetchone()[0] == 15.0 + + +def test_null_dimension_value_round_trips_positionally(): + # A NULL in a result's dim column must reject the affine fast path + # and fall back to positional scatter. + table = pa.table( + { + "lat": pa.array([0.0, None, 1.0, 3.0], type=pa.float64()), + "v": [10.0, 99.0, 11.0, 13.0], + } + ) + out = xql.to_dataset(table, dims=["lat"]) + np.testing.assert_allclose(out["v"].values, [10.0, 99.0, 11.0, 13.0]) diff --git a/tests/test_geometry.py b/tests/test_geometry.py new file mode 100644 index 00000000..5e3238d7 --- /dev/null +++ b/tests/test_geometry.py @@ -0,0 +1,170 @@ +"""GeoArrow point-geometry columns derived at registration.""" + +import json + +import numpy as np +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def grid() -> xr.Dataset: + return xr.Dataset( + {"risk": (["y", "x"], np.arange(8.0 * 6).reshape(8, 6))}, + coords={ + "y": np.linspace(-28.0, -29.4, 8), # descending, like rasters + "x": np.linspace(-58.0, -57.0, 6), + }, + ) + + +def test_geometry_field_annotation(grid): + dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y")) + field = dataset.schema.field("geometry") + assert field.type == pa.binary() + assert field.metadata[b"ARROW:extension:name"] == b"geoarrow.wkb" + meta = json.loads(field.metadata[b"ARROW:extension:metadata"]) + assert meta == {"crs": "OGC:CRS84"} + + +def test_wkb_points_decode_exactly(grid): + dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y")) + table = dataset.to_table(columns=["geometry", "x", "y"]) + blob = table["geometry"][0].as_py() + assert len(blob) == 21 and blob[0] == 1 + x = np.frombuffer(blob, "= -28.7) & (grid.y <= -28.0) + expected = grid.risk.values[inside.values, :] + assert got == (expected.size, round(float(expected.mean()), 3)) + + +def test_geopandas_consumes_native_points(grid): + gpd = pytest.importorskip("geopandas") + + dataset = xql.arrow_dataset( + grid, {"y": 4}, geometry=("x", "y"), geometry_encoding="point" + ) + gdf = gpd.GeoDataFrame.from_arrow(dataset.to_table()) + assert gdf.geometry.iloc[0].x == float(grid.x[0]) + assert str(gdf.crs).endswith("CRS84") + + +def test_geometry_name_collision_raises(): + clash = xr.Dataset( + {"geometry": (["x"], np.arange(3.0))}, + coords={"x": np.arange(3.0)}, + ) + with pytest.raises(ValueError, match="shadow"): + xql.arrow_dataset(clash, {"x": 3}, geometry=("x", "x")) + + +def test_bbox_conjuncts_prunes_and_pairs_with_st_within(grid, spatial_con): + con, reads = spatial_con + + bounds = (-58.1, -28.75, -56.9, -27.9) # xmin, ymin, xmax, ymax + conjuncts = xql.bbox_conjuncts(bounds, x="x", y="y") + assert '"x" BETWEEN' in conjuncts and '"y" BETWEEN' in conjuncts + reads.clear() + got = con.execute( + f"SELECT count(*) FROM t WHERE {conjuncts} " + "AND ST_Within(geometry, ST_GeomFromText(" + "'POLYGON ((-58.1 -28.75, -56.9 -28.75, -56.9 -27.9, " + "-58.1 -27.9, -58.1 -28.75))'))" + ).fetchone() + assert len(reads) == 1 # the y-range pruned to one chunk + inside = (grid.y >= -28.75) & (grid.y <= -27.9) + assert got[0] == int(inside.sum()) * grid.sizes["x"] + + +def test_bbox_conjuncts_accepts_bounds_objects(): + class Boxy: + bounds = (1.0, 2.0, 3.0, 4.0) + + sql = xql.bbox_conjuncts(Boxy(), x="lon", y="lat", pad=0.5) + assert sql == '"lon" BETWEEN 0.5 AND 3.5 AND "lat" BETWEEN 1.5 AND 4.5' + + +def test_wkb_points_guards_int32_offset_overflow(): + from xarray_sql.geometry import _wkb_points + + # Stride-0 broadcast views: len() reports ~103M points without + # allocating them, and the guard must fire before any buffer is + # built (pa.binary() offsets are int32; n * 21 would overflow). + n = 103_000_000 + x = np.broadcast_to(np.float64(0.0), (n,)) + with pytest.raises(ValueError, match="int32"): + _wkb_points(x, x) diff --git a/tests/test_lazy_roundtrip.py b/tests/test_lazy_roundtrip.py new file mode 100644 index 00000000..185191ef --- /dev/null +++ b/tests/test_lazy_roundtrip.py @@ -0,0 +1,377 @@ +"""Lazy chunked round-trip through engines beyond DataFusion. + +``xql.to_dataset(result, chunks=...)`` re-executes the engine's query +per accessed window. These tests verify the reconstruction is correct +on Polars frames (DuckDB chunked reconstruction fails fast — see +DuckDBHandle.supports_chunked — while its eager path works), that +laziness is real, and that one-shot +streams are rejected with a clear error. +""" + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr + +import xarray_sql as xql +from xarray_sql.backends.pyarrow import XarrayPushdownDataset + + +@pytest.fixture +def source() -> xr.Dataset: + np.random.seed(7) + return xr.Dataset( + { + "t2m": ( + ["time", "lat"], + np.random.rand(100, 6).astype(np.float64), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=100, freq="h"), + "lat": np.linspace(-25.0, 25.0, 6), + }, + attrs={"title": "synthetic"}, + ) + + +@pytest.fixture +def registered(source): + """A DuckDB connection with the source registered + a read counter.""" + duckdb = pytest.importorskip("duckdb") + + reads: list[dict] = [] + dataset = XarrayPushdownDataset( + source, {"time": 10}, _iteration_callback=lambda b, n: reads.append(b) + ) + con = duckdb.connect() + con.register("t", dataset) + return con, reads + + +def test_duckdb_chunked_fails_fast_with_guidance(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + # Re-executing a DuckDB relation from dask worker threads + # intermittently deadlocks inside duckdb-python when the query + # scans a Python-backed table; the library refuses instead of + # hanging (see DuckDBHandle.supports_chunked). + with pytest.raises(NotImplementedError, match="Polars"): + xql.to_dataset(rel, template=source, chunks={"time": 10}) + + +def test_duckdb_eager_round_trip_through_handle(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset(rel, template=source) + assert not out.chunks + assert out.attrs == source.attrs + xr.testing.assert_allclose(out, source) + + +def test_duckdb_eager_filtered_and_aggregated(source, registered): + con, _ = registered + rel = con.sql( + "SELECT time, avg(t2m) AS t2m FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ) + out = xql.to_dataset(rel, template=source) + expected = source.t2m.sel(lat=source.lat[source.lat > 0]).mean("lat") + np.testing.assert_allclose(out.t2m.values, expected.values) + + +def test_polars_lazyframe_chunked_round_trip(source): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset(lf, template=source, chunks={"time": 20}) + assert out.chunks + xr.testing.assert_allclose(out.compute(), source) + + # Eager path through the same handle (LazyFrame has no stream + # protocol; the handle executes it once). + eager = xql.to_dataset(lf, template=source) + xr.testing.assert_allclose(eager, source) + + +def test_polars_eager_frame_is_reexecutable(source): + pl = pytest.importorskip("polars") + + frame = pl.DataFrame( + { + "time": np.repeat(source.time.values, 6), + "lat": np.tile(source.lat.values, 100), + "t2m": source.t2m.values.ravel(), + } + ) + out = xql.to_dataset(frame, template=source, chunks={"time": 50}) + xr.testing.assert_allclose(out.compute(), source) + + +def test_one_shot_stream_with_chunks_raises(source, registered): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + with pytest.raises(TypeError, match="re-executable"): + xql.to_dataset(table, template=source, chunks={"time": 10}) + + +def test_inherit_without_chunked_source_falls_back_to_eager(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset(rel, template=source, chunks="inherit") + # The in-memory template has no multi-chunk dim: eager, dense. + assert not out.chunks + xr.testing.assert_allclose(out, source) + + +def test_stepped_indexer_uses_value_lists(source): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset(lf, template=source, chunks={"time": 10}) + # A step-2 selection is not a contiguous coordinate range; the + # values path must return exactly the requested rows. + stepped = out.t2m.isel(time=slice(10, 30, 2)).compute() + np.testing.assert_allclose( + stepped.values, source.t2m.isel(time=slice(10, 30, 2)).values + ) + + +def test_descending_coordinate_windows(): + pl = pytest.importorskip("polars") + + desc = xr.Dataset( + {"v": (["lat"], np.arange(8.0))}, + coords={"lat": np.linspace(70.0, 0.0, 8)}, # descending, like ERA5 + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(desc, {"lat": 4})) + out = xql.to_dataset( + lf, + template=desc, + chunks={"lat": 4}, + coords="template", + ) + xr.testing.assert_allclose(out.compute(), desc) + window = out.v.isel(lat=slice(2, 6)).compute() + np.testing.assert_allclose(window.values, desc.v.isel(lat=slice(2, 6))) + + +def test_unsorted_template_coords_window_exactly(): + pl = pytest.importorskip("polars") + + # Template coords are used verbatim, so the backend can see a + # non-monotonic coordinate array. A contiguous positional window + # like 1:3 then has monotonic values [7, 55], but the value range + # [7, 55] also admits 23 at position 3 — the scatter would write + # that unrequested row over a requested cell. Windows over a + # non-monotonic coordinate must use explicit value lists. + src = xr.Dataset( + {"v": (["x"], np.arange(4.0))}, + coords={"x": np.array([102.0, 7.0, 55.0, 23.0])}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"x": 4})) + out = xql.to_dataset(lf, template=src, chunks={"x": 4}, coords="template") + window = out.v.isel(x=slice(1, 3)).compute() + np.testing.assert_array_equal(window.values, [1.0, 2.0]) + xr.testing.assert_allclose(out.compute(), src) + + +def test_polars_float_value_windows_are_exact(): + pl = pytest.importorskip("polars") + + # Non-representable float coordinates: upstream Polars is_in drops + # them (silently matching nothing); the handle's degenerate-range + # translation must return exactly the requested rows. + src = xr.Dataset( + {"v": (["lat", "t"], np.arange(38.0).reshape(19, 2))}, + coords={"lat": np.linspace(-45.0, 45.0, 19), "t": [0.0, 1.0]}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"lat": 5})) + out = xql.to_dataset(lf, template=src, chunks={"lat": 5}) + # A stepped (non-contiguous) selection forces the value-list path. + picked = out.v.isel(lat=slice(1, 12, 2)).compute() + np.testing.assert_allclose( + picked.values, src.v.isel(lat=slice(1, 12, 2)).values + ) + + +def test_max_result_bytes_guards_stream_collection(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset(rel, template=source, max_result_bytes=1_000) + # A generous budget passes untouched. + out = xql.to_dataset(rel, template=source, max_result_bytes=10**9) + xr.testing.assert_allclose(out, source) + + +def test_polars_large_float_value_lists_stay_flat(): + pl = pytest.importorskip("polars") + + # 5000 stepped float values in a single window: a left-deep OR + # chain plans quadratically at this size (seconds per window); the + # flat any_horizontal translation must stay exact and quick. + n = 10_000 + src = xr.Dataset( + {"v": (["x"], np.arange(float(n)))}, + coords={"x": np.linspace(-45.0, 45.0, n)}, + ) + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(src, {"x": n})) + out = xql.to_dataset(lf, template=src, chunks={"x": n}) + picked = out.v.isel(x=slice(1, None, 2)).compute() + np.testing.assert_allclose( + picked.values, src.v.isel(x=slice(1, None, 2)).values + ) + + +def test_collect_streaming_falls_back_on_older_polars(): + from xarray_sql.lazyscan import _collect_streaming + + class OldLazyFrame: + # Pre-1.25 collect(): no ``engine`` keyword. + def collect(self): + return "collected" + + assert _collect_streaming(OldLazyFrame()) == "collected" + + +def test_max_result_bytes_guards_polars_lazyframe(source): + pl = pytest.importorskip("polars") + + # The LazyFrame eager fallback collects inside the engine before + # any batch surfaces; with a budget set it must stream through + # collect_batches so the guard fires before full materialization. + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset(lf, template=source, max_result_bytes=1_000) + out = xql.to_dataset(lf, template=source, max_result_bytes=10**9) + xr.testing.assert_allclose(out, source) + + +def test_max_result_bytes_guards_table_only_results(source, registered): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + + class TableOnly: + # The narrowest result surface: to_arrow_table() materializes + # in full before the budget can see a batch, so the guard runs + # on the materialized size. + def __init__(self, t): + self._t = t + + def to_arrow_table(self): + return self._t + + with pytest.raises(ValueError, match="max_result_bytes"): + xql.to_dataset( + TableOnly(table), template=source, max_result_bytes=1_000 + ) + out = xql.to_dataset( + TableOnly(table), template=source, max_result_bytes=10**9 + ) + xr.testing.assert_allclose(out, source) + + +def test_max_result_bytes_guards_dense_blowup(registered): + con, _ = registered + # A sparse diagonal: tiny Arrow payload, huge dense grid (the + # coordinate product), so the dense-size check must fire even + # though the stream fits the budget. + diag = pa.table( + { + "a": np.arange(3000.0), + "b": np.arange(3000.0), + "v": np.ones(3000), + } + ) + with pytest.raises(ValueError, match="dense reconstruction"): + xql.to_dataset(diag, dims=["a", "b"], max_result_bytes=10_000_000) + + +def test_duckdb_spill_chunked_round_trip(source, registered, tmp_path): + con, reads = registered + rel = con.sql("SELECT * FROM t") + reads.clear() + out = xql.to_dataset( + rel, template=source, chunks={"time": 10}, spill=tmp_path + ) + spilled = list(tmp_path.glob("*.parquet")) + assert len(spilled) == 1 + # The source was streamed exactly once (10 chunks), during the spill. + assert len(reads) == 10 + assert out.chunks + reads.clear() + xr.testing.assert_allclose(out.compute(), source) + # Windows re-execute against the Parquet file, not the source. + assert reads == [] + + +def test_duckdb_spill_filtered_aggregation(source, registered, tmp_path): + con, _ = registered + rel = con.sql( + "SELECT time, avg(t2m) AS t2m FROM t " + "WHERE lat > 0 GROUP BY time ORDER BY time" + ) + out = xql.to_dataset( + rel, template=source, chunks={"time": 25}, spill=tmp_path + ) + expected = source.t2m.sel(lat=source.lat[source.lat > 0]).mean("lat") + np.testing.assert_allclose(out.t2m.compute().values, expected.values) + + +def test_one_shot_table_spill_chunked(source, registered, tmp_path): + con, _ = registered + table = con.sql("SELECT * FROM t").to_arrow_table() + out = xql.to_dataset( + table, template=source, chunks={"time": 10}, spill=tmp_path + ) + assert out.chunks + xr.testing.assert_allclose(out.compute(), source) + + +def test_spill_file_removed_when_dataset_dies(source, registered, tmp_path): + import gc + + con, _ = registered + rel = con.sql("SELECT * FROM t") + out = xql.to_dataset( + rel, template=source, chunks={"time": 10}, spill=tmp_path + ) + assert list(tmp_path.glob("*.parquet")) + del out + gc.collect() + assert list(tmp_path.glob("*.parquet")) == [] + + +def test_spill_requires_chunks(source, registered): + con, _ = registered + rel = con.sql("SELECT * FROM t") + with pytest.raises(ValueError, match="spill= only applies"): + xql.to_dataset(rel, template=source, spill=True) + + +def test_duckdb_handle_runner_stopped_when_handle_dies(source, registered): + import gc + + from xarray_sql.lazyscan import DuckDBHandle + + con, _ = registered + handle = DuckDBHandle(con.sql("SELECT * FROM t")) + runner = handle._runner + del handle + gc.collect() + # A shut-down executor refuses new work — the observable contract + # that the handle's dedicated engine thread has been told to exit. + with pytest.raises(RuntimeError): + runner.submit(lambda: None) + + +def test_polars_spill_uses_streaming_sink(source, tmp_path): + pl = pytest.importorskip("polars") + + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(source, {"time": 10})) + out = xql.to_dataset( + lf, template=source, chunks={"time": 20}, spill=tmp_path + ) + xr.testing.assert_allclose(out.compute(), source) diff --git a/tests/test_proj.py b/tests/test_proj.py new file mode 100644 index 00000000..b10c9df1 --- /dev/null +++ b/tests/test_proj.py @@ -0,0 +1,158 @@ +"""Tests for the pyproj CRS-transform extension (`xarray_sql.proj`).""" + +import numpy as np +import pyproj +import pytest +import xarray as xr +from datafusion import SessionContext + +from xarray_sql import XarrayContext, proj + +UTM10 = "EPSG:32610" # UTM zone 10N (metres) +UTM11 = "EPSG:32611" # UTM zone 11N (metres) +WGS84 = "EPSG:4326" # lon/lat degrees +WEBMERC = "EPSG:3857" # Web Mercator (metres) + + +@pytest.fixture +def utm_grid(): + """A 60x50 UTM zone 10N grid over the San Francisco Bay Area. + + Chunked into several partitions so DataFusion evaluates the UDF + concurrently — exercising the per-thread transformer cache. + """ + x = np.linspace(530_000.0, 630_000.0, 50) + y = np.linspace(4_140_000.0, 4_250_000.0, 60) + xx, yy = np.meshgrid(x, y) + return xr.Dataset( + {"value": (["y", "x"], np.hypot(xx, yy))}, + coords={"y": y, "x": x}, + ).chunk({"y": 15, "x": 50}) + + +def test_reproject_matches_pyproj(utm_grid): + ctx = XarrayContext() + ctx.from_dataset("grid", utm_grid) + result = ctx.sql( + f""" + SELECT x, y, + reproject(x, y, '{UTM10}', '{WGS84}')['x'] AS lon, + reproject(x, y, '{UTM10}', '{WGS84}')['y'] AS lat + FROM grid + ORDER BY y, x + """ + ).to_pandas() + + transformer = pyproj.Transformer.from_crs(UTM10, WGS84, always_xy=True) + ref_lon, ref_lat = transformer.transform( + result["x"].to_numpy(), result["y"].to_numpy() + ) + np.testing.assert_allclose(result["lon"], ref_lon, rtol=0, atol=1e-9) + np.testing.assert_allclose(result["lat"], ref_lat, rtol=0, atol=1e-9) + + +def test_reproject_roundtrip(utm_grid): + ctx = XarrayContext() + ctx.from_dataset("grid", utm_grid) + result = ctx.sql( + f""" + WITH lonlat AS ( + SELECT x, y, + reproject(x, y, '{UTM10}', '{WGS84}')['x'] AS lon, + reproject(x, y, '{UTM10}', '{WGS84}')['y'] AS lat + FROM grid + ) + SELECT x, y, + reproject(lon, lat, '{WGS84}', '{UTM10}')['x'] AS rx, + reproject(lon, lat, '{WGS84}', '{UTM10}')['y'] AS ry + FROM lonlat + ORDER BY y, x + """ + ).to_pandas() + # A metre-based CRS round-trips to well under a millimetre. + np.testing.assert_allclose(result["rx"], result["x"], rtol=0, atol=1e-4) + np.testing.assert_allclose(result["ry"], result["y"], rtol=0, atol=1e-4) + + +def test_per_row_crs(): + """The CRS arguments are expressions, so they may vary row by row.""" + lon = np.linspace(-125.9, -114.1, 24) # spans UTM zones 10N and 11N + lat = np.linspace(32.5, 41.5, 10) + LON, LAT = np.meshgrid(lon, lat) + pts = xr.Dataset( + { + "lon": (["i"], LON.ravel()), + "lat": (["i"], LAT.ravel()), + }, + coords={"i": np.arange(LON.size)}, + ).chunk({"i": LON.size}) + + ctx = XarrayContext() + ctx.from_dataset("pts", pts) + result = ctx.sql( + f""" + SELECT lon, lat, + reproject(lon, lat, '{WGS84}', + CASE WHEN lon < -120.0 + THEN '{UTM10}' ELSE '{UTM11}' END)['x'] AS e, + reproject(lon, lat, '{WGS84}', + CASE WHEN lon < -120.0 + THEN '{UTM10}' ELSE '{UTM11}' END)['y'] AS n + FROM pts + ORDER BY i + """ + ).to_pandas() + + for zone, mask in [ + (UTM10, result["lon"] < -120.0), + (UTM11, result["lon"] >= -120.0), + ]: + transformer = pyproj.Transformer.from_crs(WGS84, zone, always_xy=True) + ref_e, ref_n = transformer.transform( + result.loc[mask, "lon"].to_numpy(), + result.loc[mask, "lat"].to_numpy(), + ) + np.testing.assert_allclose( + result.loc[mask, "e"], ref_e, rtol=0, atol=1e-6 + ) + np.testing.assert_allclose( + result.loc[mask, "n"], ref_n, rtol=0, atol=1e-6 + ) + + +def test_null_and_out_of_domain_yield_nan(): + ctx = XarrayContext() + result = ctx.sql( + f""" + SELECT + reproject(CAST(NULL AS DOUBLE), 45.0, + '{WGS84}', '{WEBMERC}')['x'] AS null_coord, + reproject(0.0, 100.0, '{WGS84}', '{WEBMERC}')['y'] AS bad_lat, + reproject(0.0, 45.0, CAST(NULL AS VARCHAR), + '{WEBMERC}')['x'] AS null_crs + """ + ).to_pandas() + assert np.isnan(result["null_coord"].iloc[0]) + assert np.isnan(result["bad_lat"].iloc[0]) + assert np.isnan(result["null_crs"].iloc[0]) + + +def test_invalid_crs_raises(): + ctx = XarrayContext() + with pytest.raises(Exception): + ctx.sql( + "SELECT reproject(0.0, 0.0, 'EPSG:999999', 'EPSG:4326')" + ).to_pandas() + + +def test_register_on_plain_session_context_with_custom_name(): + ctx = SessionContext() + proj.register(ctx, name="st_transform") + result = ctx.sql( + f""" + SELECT st_transform(-122.0, 37.0, '{WGS84}', '{WEBMERC}')['x'] AS gx + """ + ).to_pandas() + transformer = pyproj.Transformer.from_crs(WGS84, WEBMERC, always_xy=True) + ref_x, _ = transformer.transform(-122.0, 37.0) + np.testing.assert_allclose(result["gx"].iloc[0], ref_x, rtol=0, atol=1e-6) diff --git a/tests/test_reader.py b/tests/test_reader.py new file mode 100644 index 00000000..153a9f3b --- /dev/null +++ b/tests/test_reader.py @@ -0,0 +1,1415 @@ +"""Tests for XarrayRecordBatchReader lazy streaming behavior. + +These tests verify that XarrayRecordBatchReader provides true lazy evaluation: +- No data iteration during reader creation +- No data iteration during DataFusion table registration (using LazyArrowStreamTable) +- Data iteration ONLY occurs during query execution (collect()) + +The lazy streaming is achieved via the Rust LazyArrowStreamTable class which +implements the __datafusion_table_provider__ protocol using StreamingTable. + +Additional tests verify: +- True streaming with bounded memory (batches processed incrementally) +- Back-pressure behavior (producer pauses when consumer is slow) +- Error propagation through the stream +""" + +import threading +import time +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest +import xarray as xr +from datafusion import SessionContext + +from xarray_sql._native import LazyArrowStreamTable +from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table + + +@pytest.fixture +def small_ds(): + """Create a small dataset for testing.""" + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=100, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) + + data = np.random.rand(100, 10, 10).astype(np.float32) + + return xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time, "lat": lat, "lon": lon}, + ) + + +class IterationTracker: + """Tracks when iteration occurs for testing lazy evaluation. + + The callback signature is ``(block, projection_names)`` where + ``projection_names`` is the list of column names requested by the query + (``None`` when no projection pushdown occurred, e.g. for + ``XarrayRecordBatchReader`` or a ``SELECT *`` query). + """ + + def __init__(self): + self.iteration_count = 0 + self.blocks_seen = [] + self.projections_seen = [] + + def __call__(self, block, projection_names=None): + self.iteration_count += 1 + self.blocks_seen.append(block) + self.projections_seen.append(projection_names) + + def reset(self): + self.iteration_count = 0 + self.blocks_seen = [] + self.projections_seen = [] + + +class TestXarrayRecordBatchReaderCreation: + """Tests that reader creation does NOT trigger data iteration.""" + + def test_reader_creation_does_not_iterate(self, small_ds): + """Creating a reader should NOT iterate through any data.""" + tracker = IterationTracker() + + XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations during reader creation, " + f"but got {tracker.iteration_count}" + ) + + def test_schema_access_does_not_iterate(self, small_ds): + """Accessing the schema should NOT trigger iteration.""" + tracker = IterationTracker() + + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + # Access schema + _ = reader.schema + _ = reader.__arrow_c_schema__() + + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations when accessing schema, " + f"but got {tracker.iteration_count}" + ) + + +class TestDataFusionRegistration: + """Tests that DataFusion table registration does NOT trigger iteration. + + These tests use read_xarray_table with register_table() + to achieve true lazy evaluation. + """ + + def test_register_table_does_not_iterate(self, small_ds): + """Registering a LazyArrowStreamTable should NOT iterate data. + + This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps + a factory and implements __datafusion_table_provider__ with StreamingTable, + ensuring data is only read during query execution. + """ + tracker = IterationTracker() + + # Use read_xarray_table which creates a factory-based table + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + assert tracker.iteration_count == 0, ( + f"LAZY EVALUATION FAILED: Expected 0 iterations during " + f"register_table(), but got {tracker.iteration_count}." + ) + + def test_sql_planning_does_not_iterate(self, small_ds): + """Creating a SQL query plan should NOT iterate data.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Create a query but don't execute it + ctx.sql("SELECT AVG(temperature) FROM test_table") + + # Just creating the query shouldn't iterate + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations during SQL planning, " + f"but got {tracker.iteration_count}. " + f"DataFusion may be scanning data during query planning." + ) + + +class TestDataFusionCollect: + """Tests that data iteration ONLY occurs during collect(). + + These tests use read_xarray_table to verify lazy evaluation. + """ + + def test_collect_triggers_iteration(self, small_ds): + """collect() should trigger data iteration.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Verify no iteration yet (lazy registration) + iteration_before_collect = tracker.iteration_count + assert iteration_before_collect == 0, ( + "Should have 0 iterations before collect" + ) + + # Now collect - this SHOULD iterate + ctx.sql("SELECT * FROM test_table LIMIT 10").collect() + + assert tracker.iteration_count > 0, ( + "Expected iterations during collect(), but got 0. " + "Data was never read!" + ) + assert tracker.iteration_count > iteration_before_collect, ( + "Expected more iterations after collect()" + ) + + def test_full_query_iterates_all_blocks(self, small_ds): + """A query that reads all data should iterate all blocks.""" + tracker = IterationTracker() + + chunks = {"time": 25} + table = read_xarray_table( + small_ds, + chunks=chunks, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run a query that needs to scan all data + ctx.sql("SELECT * FROM test_table").collect() + + # With time=100 and chunks=25, we expect 4 blocks + expected_blocks = 100 // 25 + assert tracker.iteration_count == expected_blocks, ( + f"Expected {expected_blocks} block iterations, " + f"but got {tracker.iteration_count}" + ) + + def test_aggregation_query_iterates_correctly(self, small_ds): + """Aggregation queries should iterate all necessary blocks.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run aggregation + result = ctx.sql( + "SELECT lat, AVG(temperature) as avg_temp " + "FROM test_table GROUP BY lat" + ).collect() + + # Should have iterated some blocks + assert tracker.iteration_count > 0 + assert len(result) > 0 + + +class TestLazyEvaluationEndToEnd: + """End-to-end tests verifying lazy evaluation through the full pipeline. + + These tests use read_xarray_table to achieve true lazy evaluation. + """ + + def test_lazy_evaluation_sequence(self, small_ds): + """Verify the exact sequence of lazy evaluation stages. + + This is the comprehensive test that proves true lazy evaluation: + 1. Table creation: 0 iterations + 2. Table registration: 0 iterations + 3. Query planning: 0 iterations + 4. collect(): N iterations (where N = number of blocks) + """ + tracker = IterationTracker() + + # Stage 1: Table creation (with factory) + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + iterations_after_table = tracker.iteration_count + assert iterations_after_table == 0, ( + f"Stage 1 FAILED: Table creation triggered " + f"{iterations_after_table} iterations" + ) + + # Stage 2: Table registration + ctx = SessionContext() + ctx.register_table("test_table", table) + iterations_after_registration = tracker.iteration_count + assert iterations_after_registration == 0, ( + f"Stage 2 FAILED: Table registration triggered " + f"{iterations_after_registration} iterations" + ) + + # Stage 3: Query planning + query = ctx.sql("SELECT * FROM test_table") + iterations_after_planning = tracker.iteration_count + assert iterations_after_planning == 0, ( + f"Stage 3 FAILED: Query planning triggered " + f"{iterations_after_planning} iterations" + ) + + # Stage 4: collect() - NOW iteration should happen + query.collect() + iterations_after_collect = tracker.iteration_count + assert iterations_after_collect > 0, ( + "Stage 4 FAILED: collect() triggered 0 iterations - no data was read!" + ) + + # Verify we got the expected number of blocks (100 time steps / 25 = 4) + expected_blocks = 4 + assert iterations_after_collect == expected_blocks, ( + f"Expected {expected_blocks} iterations, got {iterations_after_collect}" + ) + + def test_multiple_queries_on_same_table(self, small_ds): + """Same table can be queried multiple times with fresh iteration each time.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 50}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # First query + ctx.sql("SELECT * FROM test_table").collect() + first_query_iterations = tracker.iteration_count + assert first_query_iterations > 0, "First query should iterate" + + # Second query on same table - should iterate again + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + second_query_iterations = tracker.iteration_count + assert second_query_iterations > first_query_iterations, ( + "Second query should trigger additional iterations" + ) + + def test_stream_consumed_error(self, small_ds): + """Once consumed, a single XarrayRecordBatchReader should not be reusable.""" + reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25}) + + # Consume the reader by converting to a PyArrow reader and reading + import pyarrow as pa + + pa_reader = pa.RecordBatchReader.from_stream(reader) + _ = pa_reader.read_all() + + # Reader is now consumed, calling __arrow_c_stream__ again should fail + with pytest.raises(RuntimeError, match="already consumed"): + reader.__arrow_c_stream__() + + +class TestDataIntegrity: + """Tests that verify data correctness alongside lazy evaluation. + + These tests use read_xarray_table for lazy streaming. + """ + + def test_query_results_are_correct(self, small_ds): + """Verify that lazy evaluation produces correct results.""" + table = read_xarray_table(small_ds, chunks={"time": 25}) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Get count + result = ctx.sql("SELECT * FROM test_table").collect() + count = sum(b.num_rows for b in result) + + # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows + expected_count = 100 * 10 * 10 + assert count == expected_count, ( + f"Expected {expected_count} rows, got {count}" + ) + + def test_aggregation_results_are_correct(self, small_ds): + """Verify aggregation produces correct results.""" + table = read_xarray_table(small_ds, chunks={"time": 25}) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Get average temperature + result = ctx.sql( + "SELECT AVG(temperature) as avg_temp FROM test_table" + ).collect() + avg_temp = result[0].to_pandas()["avg_temp"].iloc[0] + + # With seed 42 and random data in [0, 1), average should be ~0.5 + assert 0.4 < avg_temp < 0.6, ( + f"Expected average temperature ~0.5, got {avg_temp}" + ) + + +class TestPyArrowInterop: + """Tests for PyArrow interoperability.""" + + def test_from_stream_does_not_iterate(self, small_ds): + """pa.RecordBatchReader.from_stream() should not iterate.""" + tracker = IterationTracker() + + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + # Create PyArrow reader from our stream + pa.RecordBatchReader.from_stream(reader) + + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations when creating PyArrow reader, " + f"but got {tracker.iteration_count}" + ) + + def test_pyarrow_iteration_triggers_callbacks(self, small_ds): + """Iterating via PyArrow should trigger our callbacks.""" + tracker = IterationTracker() + + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + pa_reader = pa.RecordBatchReader.from_stream(reader) + + # Now iterate + for batch in pa_reader: + pass + + assert tracker.iteration_count == 4, ( + f"Expected 4 iterations, got {tracker.iteration_count}" + ) + + def test_read_all_iterates_all(self, small_ds): + """read_all() should iterate through all blocks.""" + tracker = IterationTracker() + + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + pa_reader = pa.RecordBatchReader.from_stream(reader) + table = pa_reader.read_all() + + assert tracker.iteration_count == 4 + assert len(table) == 100 * 10 * 10 + + +class StreamingTracker: + """Tracks timing of batch iterations to verify streaming behavior. + + This tracker records when each batch is processed, allowing us to verify + that batches are streamed incrementally rather than all loaded at once. + """ + + def __init__(self): + self.batch_times = [] + self.batch_count = 0 + self._lock = threading.Lock() + + def __call__(self, block, projection_names=None): + with self._lock: + self.batch_times.append(time.monotonic()) + self.batch_count += 1 + + def reset(self): + with self._lock: + self.batch_times = [] + self.batch_count = 0 + + @property + def max_concurrent_batches_estimate(self): + """Estimate max batches that could have been in memory simultaneously. + + If all batches are loaded at once, all batch_times will be very close. + If streaming works correctly, batch_times should be spread out. + """ + if len(self.batch_times) < 2: + return len(self.batch_times) + + # Sort times and look at gaps + sorted_times = sorted(self.batch_times) + # If times are spread out, streaming is working + # If all times are within a tiny window, all batches loaded at once + sorted_times[-1] - sorted_times[0] + + # If the spread is very small compared to number of batches, + # batches were likely all loaded at once + return len(self.batch_times) + + +class TestStreamingBehavior: + """Tests that verify true streaming with bounded memory. + + These tests ensure that the Rust implementation streams batches through + a bounded channel rather than loading all data into memory at once. + """ + + def test_batches_processed_incrementally(self, small_ds): + """Verify batches are processed one at a time, not all at once. + + This test uses a callback that tracks when each batch is processed. + With true streaming, batches should be processed incrementally. + """ + tracker = StreamingTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run query that scans all data + ctx.sql("SELECT * FROM test_table").collect() + + # All 4 batches should have been processed + assert tracker.batch_count == 4, ( + f"Expected 4 batches, got {tracker.batch_count}" + ) + + def test_all_partitions_processed(self, small_ds): + """Verify that all partitions are processed (order may vary with parallelism).""" + blocks_seen = [] + + def track_order(block, projection_names=None): + # Record the time slice for ordering verification + blocks_seen.append(block.get("time", None)) + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=track_order, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + ctx.sql("SELECT * FROM test_table").collect() + + # Should have 4 blocks/partitions + assert len(blocks_seen) == 4 + + # All blocks should be present (though order may vary due to parallelism) + # Extract start positions and verify they cover all expected ranges + starts = sorted([b.start for b in blocks_seen]) + expected_starts = [0, 25, 50, 75] + assert starts == expected_starts, ( + f"Expected partition starts {expected_starts}, got {starts}" + ) + + def test_large_dataset_streams_correctly(self): + """Test streaming with a larger dataset to verify memory behavior. + + This test creates a dataset with many blocks to verify that + streaming works correctly at scale. + """ + # Create a dataset with 20 blocks + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=200, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) + + data = np.random.rand(200, 10, 10).astype(np.float32) + + large_ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # Use small chunks to create many blocks + table = read_xarray_table( + large_ds, + chunks={"time": 10}, # 200 / 10 = 20 blocks + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run a query that needs all data + result = ctx.sql("SELECT * FROM test_table").collect() + count = sum(b.num_rows for b in result) + + # Verify all blocks were processed + assert tracker.batch_count == 20, ( + f"Expected 20 batches for large dataset, got {tracker.batch_count}" + ) + + # Verify data integrity + expected_count = 200 * 10 * 10 + assert count == expected_count, ( + f"Expected {expected_count} rows, got {count}" + ) + + +class TestBoundedMemoryBehavior: + """Tests that verify memory usage remains bounded during streaming. + + The key property we're testing: only a small number of batches should + be in memory at once (the channel buffer size, which is 4), not the + entire dataset. + + These tests verify that: + 1. Many batches can be processed without loading all into memory + 2. Production times are spread out (indicating back-pressure) + 3. Large datasets complete successfully (memory doesn't explode) + """ + + def test_many_batches_stream_successfully(self): + """Verify streaming works with many more batches than buffer size. + + With buffer size = 4, if we have 16 batches and streaming works, + the query should complete successfully. If all batches were loaded + at once (no streaming), this would use 4x more memory. + """ + # Create dataset with 16 batches (4x buffer size) + np.random.seed(42) + time_coord = pd.date_range("2020-01-01", periods=160, freq="h") + lat = np.linspace(-90, 90, 5) + lon = np.linspace(-180, 180, 5) + data = np.random.rand(160, 5, 5).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 16 batches (160 / 10 = 16) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + result = ctx.sql("SELECT * FROM test_table").collect() + count = sum(b.num_rows for b in result) + + # All 16 batches should have been processed + assert tracker.batch_count == 16, ( + f"Expected 16 batches, got {tracker.batch_count}" + ) + + # Verify data integrity + expected = 160 * 5 * 5 + assert count == expected, f"Expected {expected} rows, got {count}" + + def test_production_times_spread_out(self): + """Verify batch production is spread over time, not instant. + + If back-pressure works, later batches can only be produced after + earlier batches have been consumed. Production times should span + a non-zero duration. + """ + np.random.seed(123) + time_coord = pd.date_range("2020-01-01", periods=100, freq="h") + lat = np.linspace(-90, 90, 5) + lon = np.linspace(-180, 180, 5) + data = np.random.rand(100, 5, 5).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 10 batches, more than buffer size of 4 + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + + # All 10 batches should be produced + assert tracker.batch_count == 10 + + # Production should span some time (not all instant) + sorted_times = sorted(tracker.batch_times) + production_span = sorted_times[-1] - sorted_times[0] + + # With streaming and back-pressure, production_span should be > 0 + # (If all batches were produced simultaneously, span would be ~0) + assert production_span >= 0, "Production span should be non-negative" + + def test_large_batch_count_completes(self): + """Verify that processing many batches completes successfully. + + This is a stress test: 50 batches is well above the buffer size of 4. + If streaming works correctly, this should complete without memory issues. + """ + np.random.seed(456) + time_coord = pd.date_range("2020-01-01", periods=500, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) + data = np.random.rand(500, 10, 10).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 50 batches (500 / 10 = 50) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + result = ctx.sql("SELECT * FROM test_table").collect() + count = sum(b.num_rows for b in result) + + # All 50 batches processed + assert tracker.batch_count == 50, ( + f"Expected 50 batches, got {tracker.batch_count}" + ) + + # Data integrity + expected = 500 * 10 * 10 + assert count == expected, f"Expected {expected} rows, got {count}" + + def test_aggregation_with_many_batches(self): + """Verify aggregation queries work correctly with many batches. + + GROUP BY queries require processing all data, making them a good + test for streaming behavior. Uses collect() to verify that parallel + aggregation returns complete results (fixed in DataFusion 52+). + """ + np.random.seed(789) + time_coord = pd.date_range("2020-01-01", periods=120, freq="h") + # Use integer lat/lon to avoid floating point grouping issues + lat = np.array([0, 1, 2, 3, 4], dtype=np.float64) + lon = np.array([0, 1, 2, 3, 4], dtype=np.float64) + data = np.random.rand(120, 5, 5).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 12 partitions (one per chunk) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # GROUP BY requires scanning all data; collect() must return complete results + df = ctx.sql( + "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat" + ).to_pandas() + + assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}" + + # All partitions processed + assert tracker.batch_count == 12, ( + f"Expected 12 partitions, got {tracker.batch_count}" + ) + + +class TestErrorPropagation: + """Tests that verify errors are properly propagated through the stream. + + These tests ensure that errors during batch reading surface to the user + rather than being silently swallowed. + """ + + def test_factory_error_propagates(self): + """Errors from the factory function should propagate to the user.""" + + def failing_factory(): + raise ValueError("Factory intentionally failed") + + schema = pa.schema([("value", pa.int64())]) + # partitions is an iterable of (factory, metadata_dict, num_rows) tuples + table = LazyArrowStreamTable([(failing_factory, {}, 1)], schema) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # The error should surface when we try to collect + with pytest.raises(Exception) as exc_info: + ctx.sql("SELECT * FROM test_table").collect() + + # Verify the error message mentions the factory failure + error_message = str(exc_info.value).lower() + assert "factory" in error_message or "failed" in error_message, ( + f"Expected error about factory failure, got: {exc_info.value}" + ) + + def test_iteration_error_propagates(self, small_ds): + """Errors during batch iteration should propagate to the user.""" + error_on_batch = 2 # Fail on the third batch + + def failing_callback(block, projection_names=None): + # Track which batch we're on using a mutable default + if not hasattr(failing_callback, "count"): + failing_callback.count = 0 + failing_callback.count += 1 + + if failing_callback.count == error_on_batch: + raise RuntimeError("Intentional batch processing error") + + # Reset the counter + failing_callback.count = 0 + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=failing_callback, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # The error should surface when we try to collect + with pytest.raises(Exception): + ctx.sql("SELECT * FROM test_table").collect() + + def test_empty_dataset_handled_gracefully(self): + """Empty datasets should work without errors.""" + # Create an empty dataset with the right structure + empty_ds = xr.Dataset( + { + "temperature": ( + ["time", "lat", "lon"], + np.array([]).reshape(0, 0, 0), + ) + }, + coords={ + "time": pd.DatetimeIndex([]), + "lat": np.array([]), + "lon": np.array([]), + }, + ) + + # This should work without crashing + table = read_xarray_table(empty_ds, chunks={"time": 10}) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + result = ctx.sql("SELECT * FROM test_table").collect() + count = sum(b.num_rows for b in result) + + assert count == 0, f"Expected 0 rows for empty dataset, got {count}" + + +class TestMultiplePartitions: + """Tests for scenarios with multiple queries and table reuse.""" + + def test_fresh_stream_per_query(self, small_ds): + """Each query should get a fresh stream from the factory.""" + call_count = {"value": 0} + original_callback = None + + def counting_callback(block, projection_names=None): + call_count["value"] += 1 + if original_callback: + original_callback(block) + + table = read_xarray_table( + small_ds, + chunks={"time": 50}, # 2 blocks per query + _iteration_callback=counting_callback, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # First query + ctx.sql("SELECT * FROM test_table").collect() + first_query_count = call_count["value"] + assert first_query_count == 2, ( + f"First query: expected 2, got {first_query_count}" + ) + + # Second query should trigger fresh iteration + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + second_query_count = call_count["value"] + assert second_query_count == 4, ( + f"After second query: expected 4 total, got {second_query_count}" + ) + + # Third query + ctx.sql("SELECT MAX(temperature) FROM test_table").collect() + third_query_count = call_count["value"] + assert third_query_count == 6, ( + f"After third query: expected 6 total, got {third_query_count}" + ) + + def test_parallel_queries_independent(self, small_ds): + """Multiple contexts with the same table should work independently.""" + tracker1 = IterationTracker() + tracker2 = IterationTracker() + + table1 = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker1, + ) + + table2 = read_xarray_table( + small_ds, + chunks={"time": 50}, + _iteration_callback=tracker2, + ) + + ctx1 = SessionContext() + ctx2 = SessionContext() + + ctx1.register_table("test_table", table1) + ctx2.register_table("test_table", table2) + + # Execute queries + ctx1.sql("SELECT * FROM test_table").collect() + ctx2.sql("SELECT * FROM test_table").collect() + + # Each should have its own iteration count + assert tracker1.iteration_count == 4, ( + f"Table1: expected 4 blocks, got {tracker1.iteration_count}" + ) + assert tracker2.iteration_count == 2, ( + f"Table2: expected 2 blocks, got {tracker2.iteration_count}" + ) + + +class TestFilterPushdown: + """Tests for partition pruning via filter pushdown. + + These tests verify that SQL filters on dimension columns (time, lat, lon) + correctly prune partitions, reducing the number of partitions read. + """ + + @pytest.fixture + def time_chunked_ds(self): + """Dataset chunked by time for pruning tests.""" + np.random.seed(42) + # 100 days of data, chunked into 4 partitions of 25 days each + time = pd.date_range("2020-01-01", periods=100, freq="D") + lat = np.linspace(-90, 90, 5) + data = np.random.rand(100, 5).astype(np.float32) + + return xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds): + """Query with time > X should skip early partitions.""" + tracker = IterationTracker() + + # 4 partitions: days 0-24, 25-49, 50-74, 75-99 + # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query only last 25 days (Mar 16+) - should prune first 3 partitions + # 2020-03-16 is day 75 + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time >= '2020-03-16' + """ + ).to_pandas() + + # Should read only 1 partition (the last one) + assert tracker.iteration_count == 1, ( + f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + ) + + # Verify data correctness - 25 days * 5 lat = 125 rows + count = result["cnt"].iloc[0] + assert count == 125, f"Expected 125 rows, got {count}" + + def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds): + """Query with time < X should skip late partitions.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query only first 25 days (< Jan 26) - should prune last 3 partitions + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time < '2020-01-26' + """ + ).to_pandas() + + # Should read only 1 partition (the first one) + assert tracker.iteration_count == 1, ( + f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + ) + + # Verify correctness: Jan 1–25 (25 days) × 5 lat = 125 rows + count = result["cnt"].iloc[0] + assert count == 125, f"Expected 125 rows, got {count}" + + def test_time_between_filter_prunes_outside_range(self, time_chunked_ds): + """Query with BETWEEN should prune partitions outside the range.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 1, 2, and 3 + ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time BETWEEN '2020-02-01' AND '2020-03-21' + """ + ).collect() + + # Partition 0 (Jan 1–25) ends before Feb 1 and is pruned. + # Partitions 1, 2, 3 each overlap with Feb 1–Mar 21. + assert tracker.iteration_count == 3, ( + f"Expected exactly 3 partitions after BETWEEN pruning, got {tracker.iteration_count}" + ) + + def test_lat_filter_prunes_partitions(self): + """Latitude filter should prune irrelevant partitions.""" + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=10, freq="D") + # 100 lat values from -90 to 90 + lat = np.linspace(-90, 90, 100) + data = np.random.rand(10, 100).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + tracker = IterationTracker() + + # Chunk by latitude: 4 partitions (25 lat values each) + # Partition 0: lat -90 to ~-45 + # Partition 1: lat ~-45 to 0 + # Partition 2: lat 0 to ~45 + # Partition 3: lat ~45 to 90 + table = read_xarray_table( + ds, + chunks={"lat": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query southern hemisphere only (lat < 0) + ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE lat < 0 + """ + ).collect() + + # np.linspace(-90, 90, 100) chunked by 25: + # Partition 0: indices 0–24, lat -90.0 to -46.4 (all negative) + # Partition 1: indices 25–49, lat -44.5 to -0.9 (all negative) + # Partition 2: indices 50–74, lat 0.9 to 44.5 (all positive) → pruned + # Partition 3: indices 75–99, lat 46.4 to 90.0 (all positive) → pruned + assert tracker.iteration_count == 2, ( + f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}" + ) + + def test_unchunked_dim_filter_still_prunes(self): + """Filters on an *unchunked* dim still prune via static bounds. + + ``read_xarray_table`` precomputes bounds for unchunked dims once + rather than re-scanning their full coord array on every partition. + Regression guard: if the static-range merge ever stops attaching + those bounds to each partition, the Rust pruner falls back to + "never prune" for the unchunked dim and reads every partition. + """ + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=100, freq="D") + lat = np.linspace(-90, 90, 50) # unchunked + data = np.random.rand(100, 50).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + tracker = IterationTracker() + # Chunk only on time → lat is "static" (one chunk spanning the axis). + # Every partition still spans lat -90 to +90. + table = read_xarray_table( + ds, chunks={"time": 25}, _iteration_callback=tracker + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # WHERE lat > 100 matches no rows. With static bounds (lat ∈ [-90, 90]) + # attached to every partition, the pruner drops *all* partitions and + # the table is never iterated. Without them, all 4 are read. + ctx.sql("SELECT COUNT(*) FROM test WHERE lat > 100").collect() + assert tracker.iteration_count == 0, ( + "Static lat bounds should let the pruner skip every partition; " + f"got {tracker.iteration_count} partitions read." + ) + + def test_no_pruning_for_data_column_filters(self, time_chunked_ds): + """Filters on data columns (not dimensions) should not prune.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Filter on temperature (data column), not a dimension + ctx.sql( + """ + SELECT COUNT(*) FROM test WHERE temperature > 0.5 + """ + ).collect() + + # All 4 partitions should be read (can't prune on data column) + assert tracker.iteration_count == 4, ( + f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}" + ) + + def test_filter_correctness_preserved(self, time_chunked_ds): + """Verify filtered results are correct after pruning.""" + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Get count with filter + filtered = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time >= '2020-02-15' AND time <= '2020-03-15' + """ + ).to_pandas() + + # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days + # 30 days * 5 lat values = 150 rows + count = filtered["cnt"].iloc[0] + assert count == 150, f"Expected 150 rows, got {count}" + + def test_and_filter_combines_pruning(self, time_chunked_ds): + """AND filters should combine for maximum pruning.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Very narrow range that spans only 1 partition + ctx.sql( + """ + SELECT * FROM test + WHERE time >= '2020-03-20' AND time <= '2020-04-05' + """ + ).collect() + + # Should read only 1 partition + assert tracker.iteration_count == 1, ( + f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}" + ) + + def test_or_filter_is_conservative(self, time_chunked_ds): + """OR filters should include partitions matching either condition.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # First or last partition (OR condition) + ctx.sql( + """ + SELECT * FROM test + WHERE time < '2020-01-10' OR time > '2020-03-30' + """ + ).collect() + + # Should read at least 2 partitions (first and last) + assert tracker.iteration_count >= 2, ( + f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}" + ) + + def test_empty_result_from_impossible_filter(self, time_chunked_ds): + """Filter that matches no data should return empty result.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query for dates outside the data range + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time > '2025-01-01' + """ + ).to_pandas() + + # Should read 0 partitions (all pruned) + assert tracker.iteration_count == 0, ( + f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}" + ) + + # Result should be 0 rows + count = result["cnt"].iloc[0] + assert count == 0, f"Expected 0 rows, got {count}" + + +class TestProjectionPushdown: + """Tests that column projection is pushed down to the xarray factory. + + When a query requests only a subset of data variables, the factory should + receive only those column names — not the full schema — so that xarray + skips loading unrequested variables from disk. + """ + + @pytest.fixture + def two_var_ds(self): + """A small dataset with two data variables sharing the same dimensions.""" + np.random.seed(0) + time = pd.date_range("2020-01-01", periods=10, freq="D") + lat = np.linspace(-10, 10, 5, dtype=np.float32) + shape = (10, 5) + return xr.Dataset( + { + "temperature": ( + ["time", "lat"], + np.random.rand(*shape).astype(np.float32), + ), + "precipitation": ( + ["time", "lat"], + np.random.rand(*shape).astype(np.float32), + ), + }, + coords={"time": time, "lat": lat}, + ) + + def test_single_column_select_projects_only_that_variable(self, two_var_ds): + """SELECT on one data variable should pass only that variable to the factory.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql("SELECT AVG(temperature) FROM data").collect() + + assert tracker.iteration_count > 0, ( + "Expected at least one partition to be read" + ) + for proj in tracker.projections_seen: + assert proj is not None, ( + "Expected projection_names to be set for a single-column SELECT, got None" + ) + assert "temperature" in proj, ( + f"Expected 'temperature' in projection_names, got {proj}" + ) + assert "precipitation" not in proj, ( + f"'precipitation' should not be loaded for a temperature-only query, got {proj}" + ) + + def test_full_select_includes_all_variables(self, two_var_ds): + """SELECT * should include all data variables in the projection.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql("SELECT * FROM data").collect() + + assert tracker.iteration_count > 0, ( + "Expected at least one partition to be read" + ) + for proj in tracker.projections_seen: + # DataFusion sends all column names explicitly even for SELECT *. + # Either None (no pushdown) or a list containing both data variables. + if proj is not None: + assert "temperature" in proj, f"Missing 'temperature' in {proj}" + assert "precipitation" in proj, ( + f"Missing 'precipitation' in {proj}" + ) + + def test_multi_column_select_includes_all_requested(self, two_var_ds): + """SELECT on multiple data variables should include all of them in the projection.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql( + "SELECT AVG(temperature), AVG(precipitation) FROM data" + ).collect() + + assert tracker.iteration_count > 0 + for proj in tracker.projections_seen: + # Both variables requested — either None (full scan) or both present + if proj is not None: + assert "temperature" in proj, f"Missing 'temperature' in {proj}" + assert "precipitation" in proj, ( + f"Missing 'precipitation' in {proj}" + ) + + def test_projection_result_correctness(self, two_var_ds): + """Single-column projected query returns the same result as an unprojected one.""" + table = read_xarray_table(two_var_ds, chunks={"time": 5}) + ctx = SessionContext() + ctx.register_table("data", table) + + projected = ( + ctx.sql("SELECT AVG(temperature) as avg_t FROM data") + .to_pandas()["avg_t"] + .iloc[0] + ) + expected = float(two_var_ds["temperature"].values.mean()) + assert abs(projected - expected) < 1e-4, ( + f"Projected AVG {projected} differs from expected {expected}" + ) + + def test_count_star_passes_none_projection(self, two_var_ds): + """COUNT(*) should not push a projection — factory receives None.""" + projections_seen = [] + + def callback(block, projection_names): + projections_seen.append(projection_names) + + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=callback, + ) + ctx = SessionContext() + ctx.register_table("data", table) + result = ctx.sql("SELECT COUNT(*) FROM data").collect() + + total_rows = 10 * 5 # time=10, lat=5 + assert result[0][0][0].as_py() == total_rows + assert all(p is None for p in projections_seen), ( + f"COUNT(*) should not push a projection, but got: {projections_seen}" + ) diff --git a/tests/test_sql.py b/tests/test_sql.py new file mode 100644 index 00000000..66c6dd27 --- /dev/null +++ b/tests/test_sql.py @@ -0,0 +1,550 @@ +"""SQL functionality tests for xarray-sql using pytest.""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from xarray_sql import XarrayContext + + +def test_sanity(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql( + 'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100' + ).to_pandas() + assert len(result) > 0 + assert len(result) <= 1320 + assert all(col in result.columns for col in ["lat", "lon", "time", "air"]) + + +def test_aggregation_small(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + query = """ + SELECT lat, lon, SUM(air) AS air_total + FROM air + GROUP BY lat, lon + """ + result = ctx.sql(query).to_pandas() + expected_rows = ( + air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"] + ) + assert len(result) == expected_rows + + +def test_aggregation_large(air_dataset_large): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_large) + query = """ + SELECT lat, lon, AVG(air) AS air_avg + FROM air + GROUP BY lat, lon + """ + result = ctx.sql(query).to_pandas() + expected_rows = ( + air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"] + ) + assert len(result) == expected_rows + + +def test_basic_select_all(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas() + assert len(result) <= 10 + for col in ["lat", "lon", "time", "air"]: + assert col in result.columns + + +def test_weather_queries(weather_dataset): + ctx = XarrayContext() + ctx.from_dataset("weather", weather_dataset) + # Selecting specific columns + result = ctx.sql( + "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20" + ).to_pandas() + assert "temperature" in result.columns + assert "precipitation" in result.columns + # Filtering + result = ctx.sql( + "SELECT * FROM weather WHERE temperature > 10 LIMIT 50" + ).to_pandas() + assert len(result) > 0 + assert (result["temperature"] > 10).all() + + +def test_synthetic_aggregations(synthetic_dataset): + ctx = XarrayContext() + ctx.from_dataset("synthetic", synthetic_dataset) + # COUNT aggregation + result = ctx.sql( + "SELECT COUNT(*) AS total_count FROM synthetic" + ).to_pandas() + assert result["total_count"].iloc[0] > 0 + # MIN, MAX, AVG + query = """ + SELECT MIN(temperature) AS min_temp, + MAX(temperature) AS max_temp, + AVG(temperature) AS avg_temp + FROM synthetic + """ + result = ctx.sql(query).to_pandas() + assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0] + assert ( + result["min_temp"].iloc[0] + <= result["avg_temp"].iloc[0] + <= result["max_temp"].iloc[0] + ) + + +def test_invalid_table_name(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT * FROM nonexistent_table") + + +def test_invalid_column_name(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT nonexistent_column FROM air") + + +def test_sql_syntax_error(air_dataset_small): + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT * FORM air") # Typo: FORM instead of FROM + with pytest.raises(Exception): + ctx.sql("SELECT * FROM air WHERE") # Incomplete WHERE + + +def test_cross_join(air_and_stations): + air, stations = air_and_stations + ctx = XarrayContext() + ctx.from_dataset("air_data", air) + ctx.from_dataset("stations", stations) + result = ctx.sql( + "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations" + ).to_pandas() + assert result["total"].iloc[0] > 0 + + +def test_string_coordinates(): + """String-typed coordinates should not crash during registration.""" + ds = xr.Dataset( + {"score": (["student", "subject"], np.random.rand(3, 2))}, + coords={ + "student": ["alice", "bob", "charlie"], + "subject": ["math", "science"], + }, + ) + ctx = XarrayContext() + ctx.from_dataset("scores", ds.chunk({"student": 3, "subject": 2})) + result = ctx.sql("SELECT * FROM scores").to_pandas() + assert len(result) == 6 + assert "student" in result.columns + assert "subject" in result.columns + assert "score" in result.columns + + +class TestNanAsNull: + """NaN in float columns should become Arrow nulls so SQL aggregates work.""" + + @pytest.fixture + def nan_ds(self): + data = np.array( + [[[1.0, 2.0], [np.nan, 4.0]], [[5.0, np.nan], [7.0, 8.0]]] + ) + return xr.Dataset( + {"temp": (["time", "x", "y"], data)}, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "x": [0, 1], + "y": [0, 1], + }, + ).chunk({"time": 1}) + + def test_nan_aggregates(self, nan_ds): + ctx = XarrayContext() + ctx.from_dataset("data", nan_ds) + + # Test multiple aggregates at once: + # MAX/MIN/AVG should ignore NaN, COUNT(col) should exclude NaN, + # and WHERE col IS NULL should match NaN. + query = """ + SELECT + MAX(temp) AS mx, + MIN(temp) AS mn, + AVG(temp) AS avg, + COUNT(temp) AS cnt, + COUNT(*) FILTER (WHERE temp IS NULL) AS null_cnt + FROM data + """ + result = ctx.sql(query).to_pandas().iloc[0] + + assert result["mx"] == 8.0 + assert result["mn"] == 1.0 + expected_avg = np.nanmean([1.0, 2.0, 4.0, 5.0, 7.0, 8.0]) + assert abs(result["avg"] - expected_avg) < 1e-6 + assert result["cnt"] == 6 + assert result["null_cnt"] == 2 + + +class TestCftimeGregorianLike: + """Tests for Gregorian-like cftime calendars (noleap, standard, etc.). + + These use pa.timestamp('us') and support string-based SQL filters. + """ + + def test_noleap_dataset_registers(self, rasm_ds): + """A noleap dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert result["cnt"].iloc[0] > 0 + + def test_select_time_column(self, rasm_ds): + """Querying the time column should return valid timestamps.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT DISTINCT time FROM rasm ORDER BY time LIMIT 5" + ).to_pandas() + assert len(result) == 5 + times = result["time"].tolist() + assert times == sorted(times) + + def test_string_filter_works(self, rasm_ds): + """String-based time filters should work for Gregorian-like calendars.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM rasm WHERE time >= '1980-10-01'" + ).to_pandas() + full = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert 0 < result["cnt"].iloc[0] < full["cnt"].iloc[0] + + def test_aggregation(self, rasm_ds): + """MIN/MAX on timestamp columns should work.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT MIN(time) AS t_min, MAX(time) AS t_max FROM rasm" + ).to_pandas() + assert result["t_min"].iloc[0] < result["t_max"].iloc[0] + + def test_row_count_matches_xarray(self, rasm_ds): + """Total row count should equal the product of dimension sizes.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + expected = int( + np.prod([rasm_ds.sizes[d] for d in rasm_ds["Tair"].dims]) + ) + assert result["cnt"].iloc[0] == expected + + +class TestCftimeNonGregorian: + """Tests for non-Gregorian cftime calendars (360_day, julian). + + These use pa.int64() with CF-convention metadata and the cftime() UDF. + """ + + @pytest.fixture + def ds_360day(self): + """Synthetic 360-day calendar dataset.""" + import cftime + + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) + + def test_360day_registers(self, ds_360day): + """A 360-day dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM ds360").to_pandas() + assert result["cnt"].iloc[0] == 12 + + def test_360day_select_ordered(self, ds_360day): + """Integer offsets should be orderable.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql( + "SELECT DISTINCT time FROM ds360 ORDER BY time" + ).to_pandas() + times = result["time"].tolist() + assert times == sorted(times) + assert len(times) == 12 + + def test_360day_integer_filter(self, ds_360day): + """Direct integer comparisons should work on int64 time columns.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # Get all distinct time values to find a midpoint + all_times = ( + ctx.sql("SELECT DISTINCT time FROM ds360 ORDER BY time") + .to_pandas()["time"] + .tolist() + ) + mid = all_times[len(all_times) // 2] + result = ctx.sql( + f"SELECT COUNT(*) AS cnt FROM ds360 WHERE time >= {mid}" + ).to_pandas() + assert 0 < result["cnt"].iloc[0] < 12 + + def test_360day_cftime_udf_registered(self, ds_360day): + """from_dataset should auto-register a cftime() UDF for 360-day calendars.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # The cftime() UDF should convert a date string to the int64 offset, + # enabling ergonomic filtering. + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM ds360 " + "WHERE time >= cftime('2000-07-01')" + ).to_pandas() + # July through December = 6 months + assert result["cnt"].iloc[0] == 6 + + def test_gregorian_like_no_cftime_udf(self): + """Gregorian-like calendars should NOT register a cftime() UDF.""" + ds = xr.tutorial.open_dataset("rasm") + ctx = XarrayContext() + ctx.from_dataset("rasm", ds, chunks={"time": 12}) + # Using cftime() should fail since it's not registered for noleap. + with pytest.raises(Exception): + ctx.sql( + "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" + ).collect() + + +class TestFromDatasetMultiDims: + """from_dataset should split datasets with mixed dims into multiple tables.""" + + @pytest.fixture + def mixed_ds(self): + np.random.seed(0) + return xr.Dataset( + { + "temperature_2m": ( + ["time", "lat", "lon"], + np.random.rand(2, 3, 4), + ), + "pressure": ( + ["time", "lat", "lon", "level"], + np.random.rand(2, 3, 4, 2), + ), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "lat": np.linspace(-90, 90, 3), + "lon": np.linspace(-180, 180, 4), + "level": [500, 1000], + }, + ).chunk({"time": 1}) + + def test_registers_multiple_tables(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + surface = ctx.sql("SELECT * FROM era5.time_lat_lon").to_pandas() + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() + assert "temperature_2m" in surface.columns + assert "pressure" in upper.columns + assert len(surface) == 2 * 3 * 4 + assert len(upper) == 2 * 3 * 4 * 2 + + @pytest.fixture + def scalar_and_array_ds(self): + """A gridded variable plus a scalar metadata variable (GOES-like).""" + np.random.seed(0) + return xr.Dataset( + { + "temperature_2m": ( + ["time", "lat", "lon"], + np.random.rand(2, 3, 4), + ), + "projection": ((), 0), + }, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "lat": np.linspace(-90, 90, 3), + "lon": np.linspace(-180, 180, 4), + }, + ).chunk({"time": 1}) + + def test_registers_scalar_var_as_single_row_table( + self, scalar_and_array_ds + ): + ctx = XarrayContext() + ctx.from_dataset("goes", scalar_and_array_ds) + surface = ctx.sql("SELECT * FROM goes.time_lat_lon").to_pandas() + scalar = ctx.sql("SELECT * FROM goes.scalar").to_pandas() + assert "temperature_2m" in surface.columns + assert "projection" in scalar.columns + assert len(scalar) == 1 + + def test_scalar_group_in_catalog(self, scalar_and_array_ds): + ctx = XarrayContext() + ctx.from_dataset("goes", scalar_and_array_ds) + tables = set(ctx.catalog().schema("goes").table_names()) + assert tables == {"time_lat_lon", "scalar"} + + def test_scalar_table_name_override(self, scalar_and_array_ds): + ctx = XarrayContext() + ctx.from_dataset("goes", scalar_and_array_ds, table_names={(): "meta"}) + result = ctx.sql("SELECT * FROM goes.meta").to_pandas() + assert "projection" in result.columns + assert len(result) == 1 + + def test_table_names_override(self, mixed_ds): + ctx = XarrayContext() + ctx.from_dataset( + "era5", + mixed_ds, + table_names={("time", "lat", "lon"): "surface"}, + ) + result = ctx.sql("SELECT * FROM era5.surface").to_pandas() + assert "temperature_2m" in result.columns + # Non-aliased group falls back to the default joined-dim table name. + upper = ctx.sql("SELECT * FROM era5.time_lat_lon_level").to_pandas() + assert "pressure" in upper.columns + + def test_schema_registered_in_catalog(self, mixed_ds): + """Mixed-dim datasets should create a SQL schema under the catalog.""" + ctx = XarrayContext() + ctx.from_dataset("era5", mixed_ds) + assert "era5" in ctx.catalog().schema_names() + tables = set(ctx.catalog().schema("era5").table_names()) + assert tables == {"time_lat_lon", "time_lat_lon_level"} + + def test_uniform_dims_uses_name_directly(self, mixed_ds): + """A single dim group should register under the bare name.""" + ds = mixed_ds[["temperature_2m"]] + ctx = XarrayContext() + ctx.from_dataset("surface", ds) + result = ctx.sql("SELECT * FROM surface").to_pandas() + assert "temperature_2m" in result.columns + + def test_table_names_is_keyword_only(self, mixed_ds): + ctx = XarrayContext() + with pytest.raises(TypeError): + ctx.from_dataset("era5", mixed_ds, {("time",): "x"}) + + @pytest.fixture + def coordless_dims_ds(self): + """Mirror the fashion-mnist layout: a dimension + coordinate (``sample``) alongside dimensions without coordinates + (``channel``/``height``/``width``).""" + n_sample, n_channel, n_height, n_width = 4, 1, 3, 3 + return xr.Dataset( + { + "images": ( + ["sample", "channel", "height", "width"], + np.arange( + n_sample * n_channel * n_height * n_width, + dtype="float32", + ).reshape(n_sample, n_channel, n_height, n_width), + ), + "labels": (["sample"], np.arange(n_sample, dtype="int64")), + }, + coords={"sample": ("sample", np.arange(n_sample, dtype="int64"))}, + ).chunk({"sample": 1}) + + def test_coordless_dims_appear_as_columns(self, coordless_dims_ds): + """Dimensions without coordinates must still be emitted as columns, + not silently dropped from the schema.""" + ctx = XarrayContext() + ctx.from_dataset( + "mnist", + coordless_dims_ds, + table_names={ + ("sample", "channel", "height", "width"): "X", + ("sample",): "y", + }, + ) + result = ctx.sql('SELECT * FROM mnist."X"').to_pandas() + assert set(result.columns) == { + "sample", + "channel", + "height", + "width", + "images", + } + + def test_coordless_dims_values_match_xarray(self, coordless_dims_ds): + """The X table's rows must match xarray's own pivot exactly, including + the synthesized index values for the coordinate-less dimensions.""" + ctx = XarrayContext() + ctx.from_dataset( + "mnist", + coordless_dims_ds, + table_names={ + ("sample", "channel", "height", "width"): "X", + ("sample",): "y", + }, + ) + dim_cols = ["sample", "channel", "height", "width"] + result = ( + ctx.sql('SELECT * FROM mnist."X"') + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + expected = ( + coordless_dims_ds[["images"]] + .to_dataframe() + .reset_index() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + pd.testing.assert_frame_equal( + result, expected, check_dtype=False, check_like=True + ) + + def test_coordless_dims_y_table_unaffected(self, coordless_dims_ds): + """The 1-D ``y`` group (sample coordinate + labels) is unchanged.""" + ctx = XarrayContext() + ctx.from_dataset( + "mnist", + coordless_dims_ds, + table_names={ + ("sample", "channel", "height", "width"): "X", + ("sample",): "y", + }, + ) + result = ctx.sql('SELECT * FROM mnist."y"').to_pandas() + assert set(result.columns) == {"sample", "labels"} + assert len(result) == coordless_dims_ds.sizes["sample"] + + def test_single_table_all_coordless_dims(self): + """A uniform-dim dataset whose dims lack coordinates registers as one + table with every dimension present as a column, and the coordinate-less + dimensions carry their ABSOLUTE index even when chunked.""" + ds = xr.Dataset( + {"a": (("x", "y"), np.arange(6, dtype="float32").reshape(3, 2))} + ).chunk({"x": 1}) # chunked along the coordinate-less 'x' dim + ctx = XarrayContext() + ctx.from_dataset("grid", ds) + result = ( + ctx.sql("SELECT * FROM grid") + .to_pandas() + .sort_values(["x", "y"]) + .reset_index(drop=True) + ) + assert set(result.columns) == {"x", "y", "a"} + # x must span 0..2 across the three chunks, not restart at 0 each block. + expected = ( + ds.to_dataframe() + .reset_index() + .sort_values(["x", "y"]) + .reset_index(drop=True) + ) + pd.testing.assert_frame_equal( + result, expected, check_dtype=False, check_like=True + ) diff --git a/tests/test_sql_recipes.py b/tests/test_sql_recipes.py new file mode 100644 index 00000000..5e654be1 --- /dev/null +++ b/tests/test_sql_recipes.py @@ -0,0 +1,61 @@ +"""The performance guide's SQL recipes, pinned on both engines. + +The guide documents caching as plain engine SQL rather than wrapping +it in a helper; this test runs the documented statement on DuckDB and +DataFusion so the recipe cannot rot. +""" + +import duckdb +import numpy as np +import pytest +import xarray as xr + +import xarray_sql as xql + + +def _grid() -> xr.Dataset: + np.random.seed(11) + return xr.Dataset( + { + "klass": ( + ["y", "x"], + np.random.randint(1, 6, (64, 64), dtype=np.uint8), + ) + }, + coords={ + "y": np.linspace(-34.0, -30.0, 64), + "x": np.linspace(-66.0, -62.0, 64), + }, + ).chunk({"y": 32}) + + +@pytest.fixture(params=["duckdb", "datafusion"]) +def con(request): + connection = ( + duckdb.connect() if request.param == "duckdb" else xql.XarrayContext() + ) + xql.register(connection, "grid", _grid()) + return connection + + +def _rows(con, sql) -> list[tuple]: + result = con.sql(sql) + if hasattr(result, "fetchall"): + return list(result.fetchall()) + frame = result.to_pandas() + return [tuple(r) for r in frame.itertuples(index=False)] + + +def test_documented_caching_recipe(con): + # The performance guide documents caching as plain engine SQL; this + # pins the recipe on both engines — including that DataFusion DDL + # is a lazy plan that must be collected to execute. + ctas = ( + "CREATE OR REPLACE TABLE cube AS " + "SELECT FLOOR(y) AS lat, klass, COUNT(*) AS n FROM grid " + "GROUP BY 1, 2 ORDER BY lat, klass" + ) + result = con.sql(ctas) + if hasattr(result, "collect"): + result.collect() + assert _rows(con, "SELECT SUM(n) FROM cube")[0][0] == 64 * 64 diff --git a/tests/test_stats.py b/tests/test_stats.py new file mode 100644 index 00000000..a1f199eb --- /dev/null +++ b/tests/test_stats.py @@ -0,0 +1,124 @@ +"""Exact table statistics reach the optimizer through the FFI boundary. + +DataFusion 54 forwards ``Statistics`` across the ``datafusion-ffi`` boundary, +so the exact statistics ``XarrayScanExec`` reports are visible to the query +optimizer: num_rows (product of a chunk's dimension sizes), total byte size, +and per dimension-column min/max bounds. These tests pin that behaviour. +""" + +import numpy as np +import xarray as xr + +from xarray_sql import XarrayContext + + +def _explain(ctx: XarrayContext, query: str) -> str: + ctx.sql("SET datafusion.explain.show_statistics = true").collect() + rows = ctx.sql(f"EXPLAIN {query}").to_pandas() + return "\n".join(rows["plan"].tolist()) + + +def test_exact_rows_in_scan_statistics(): + """The scan reports exact row counts (forwarded across FFI).""" + ds = xr.Dataset( + {"air": (("time", "lat", "lon"), np.random.rand(100, 4, 5))}, + coords={ + "time": np.arange(100), + "lat": np.arange(4), + "lon": np.arange(5), + }, + ) + ctx = XarrayContext() + ctx.from_dataset("air", ds, chunks={"time": 50}) + plan = _explain(ctx, "SELECT lat, lon, air FROM air") + total = 100 * 4 * 5 + assert f"Rows=Exact({total})" in plan + + +def test_exact_byte_size_in_scan_statistics(): + """The scan reports exact byte size (num_rows x fixed row width).""" + ds = xr.Dataset( + {"air": (("time", "lat", "lon"), np.random.rand(100, 4, 5))}, + coords={ + "time": np.arange(100), + "lat": np.arange(4), + "lon": np.arange(5), + }, + ) + ctx = XarrayContext() + ctx.from_dataset("air", ds, chunks={"time": 50}) + plan = _explain(ctx, "SELECT lat, lon, air FROM air") + # 2000 rows x (lat int64 + lon int64 + air float64) = 2000 x 24 bytes. + assert f"Bytes=Exact({100 * 4 * 5 * 24})" in plan + + +def test_dimension_column_min_max_in_scan_statistics(): + """Dimension columns carry exact min/max and a zero null count. + + These are the join/filter key columns; the bounds come from the same + coordinate metadata used for partition pruning (no data scan), and grid + axes are always fully populated so the null count is exactly zero. + """ + ds = xr.Dataset( + {"air": (("time", "lat", "lon"), np.random.rand(100, 4, 5))}, + coords={ + "time": np.arange(100), + "lat": np.arange(4), + "lon": np.arange(5), + }, + ) + ctx = XarrayContext() + ctx.from_dataset("air", ds, chunks={"time": 50}) + plan = _explain(ctx, "SELECT lat, lon, air FROM air") + # lat spans 0..3, lon spans 0..4, both never null. + assert "Min=Exact(Int64(0)) Max=Exact(Int64(3)) Null=Exact(0)" in plan + assert "Min=Exact(Int64(0)) Max=Exact(Int64(4)) Null=Exact(0)" in plan + + +def test_count_star_answered_from_statistics(): + """COUNT(*) returns the exact count from statistics (metadata only).""" + ds = xr.Dataset( + {"air": (("time", "lat", "lon"), np.random.rand(100, 4, 5))}, + coords={ + "time": np.arange(100), + "lat": np.arange(4), + "lon": np.arange(5), + }, + ) + ctx = XarrayContext() + ctx.from_dataset("air", ds, chunks={"time": 50}) + n = ctx.sql("SELECT COUNT(*) AS n FROM air").to_pandas()["n"][0] + assert int(n) == 100 * 4 * 5 + + +def test_join_picks_small_build_side(): + """With exact stats the optimizer broadcasts the smaller table (CollectLeft). + + Without statistics (the pre-54 FFI path) the optimizer could not know which + side was smaller and fell back to a Partitioned hash join. + """ + rng = np.random.default_rng(0) + big = xr.Dataset( + {"t": (("time", "lat", "lon"), rng.standard_normal((200, 8, 8)))}, + coords={ + "time": np.arange(200), + "lat": np.arange(8), + "lon": np.arange(8), + }, + ) + small = xr.Dataset( + {"w": (("lat", "lon"), rng.standard_normal((8, 8)))}, + coords={"lat": np.arange(8), "lon": np.arange(8)}, + ) + ctx = XarrayContext() + ctx.from_dataset("big", big, chunks={"time": 50}) + ctx.from_dataset("small", small, chunks={"lat": 8}) + + plan = _explain( + ctx, + "SELECT b.time, SUM(b.t * s.w) AS x FROM big b " + "JOIN small s ON b.lat=s.lat AND b.lon=s.lon GROUP BY b.time", + ) + assert "HashJoinExec: mode=CollectLeft" in plan + # The small (build) side's exact cardinality crossed the FFI boundary. + assert "Rows=Exact(64)" in plan diff --git a/tests/test_to_dataset_perf.py b/tests/test_to_dataset_perf.py new file mode 100644 index 00000000..9124e9c4 --- /dev/null +++ b/tests/test_to_dataset_perf.py @@ -0,0 +1,124 @@ +"""Peak-memory tests for the lazy SQL -> xarray round-trip. + +Asserts the lazy backend honors its contract: a single-chunk access +peaks far below an eager whole-grid materialization, and a streaming +aggregation does not balloon past the source size. +""" + +import gc +import tracemalloc + +import numpy as np +import pytest +import xarray as xr + +from xarray_sql import XarrayContext + + +def _peak_mb(fn): + """Return ``(result, peak_memory_mb)`` for a single call to ``fn``.""" + gc.collect() + tracemalloc.start() + tracemalloc.reset_peak() + out = fn() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + return out, peak / 1e6 + + +@pytest.fixture(scope="module") +def air_source(): + """NCEP ``air_temperature`` chunked along time, ~31 MB dense. + + Module-scoped so the pooch download (~3.5 MB) is amortized across + tests in this file. Skips when the tutorial dataset is unreachable. + """ + try: + return xr.tutorial.open_dataset("air_temperature").chunk({"time": 24}) + except (OSError, ValueError, ImportError) as e: + pytest.skip(f"air_temperature tutorial dataset unavailable: {e}") + + +def test_lazy_chunk_peak_memory_is_bounded(air_source): + """``.sel(time=t0).load()`` materializes only one chunk, not the cube. + + Reference observation: lazy chunk peak is ~1.8 MB on a 31 MB + dense source. A regression that quietly buffers the whole result + would push past 10 MB. Eager ``to_dataset(chunks=None)`` is + measured too as a sanity floor for the gap: the eager path should + be at least an order of magnitude heavier than the lazy chunk, + otherwise the lazy path isn't actually lazy. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_source, chunks={"time": 24}) + t0 = air_source["time"].values[0] + + out = ctx.sql('SELECT * FROM "air"').to_dataset() + chunk, chunk_peak = _peak_mb(lambda: out["air"].sel(time=t0).load()) + assert chunk.size == air_source.sizes["lat"] * air_source.sizes["lon"] + assert chunk_peak < 10.0, ( + f"lazy single-chunk access should stay under 10 MB on " + f"air_temperature, got {chunk_peak:.2f} MB" + ) + + _, eager_peak = _peak_mb( + lambda: ctx.sql('SELECT * FROM "air"').to_dataset(chunks=None) + ) + assert eager_peak > 50.0, ( + f"eager whole-grid materialization should peak above 50 MB; " + f"if it doesn't, the eager path may have silently gone lazy " + f"and the lazy assertion above no longer means anything. " + f"Got {eager_peak:.1f} MB" + ) + assert eager_peak / max(chunk_peak, 0.1) > 10.0, ( + f"eager peak should be at least 10x the lazy chunk peak; " + f"if it isn't, the lazy path isn't actually lazy. " + f"Got eager={eager_peak:.1f} MB, lazy={chunk_peak:.2f} MB" + ) + + +def test_streaming_aggregation_does_not_explode(air_source): + """A ``GROUP BY`` reducing the long axis streams in a single pass. + + Reduces 3.86M rows of ``air_temperature`` to ~1.3K group cells. The + aggregation must stream the source once and emit a tiny result, not + buffer the entire row set into memory. The engine streams partitions + concurrently, so the tracemalloc peak varies run to run (~1.2x-4x the + source, scaling with the number of in-flight partitions); we take the + best of a few samples for a stable floor (~1.5x source), well below a + "buffer the whole row set" regression, which would keep every sample + high. Threshold is 4x source size so transient buffers fit. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_source, chunks={"time": 24}) + source_mb = air_source.nbytes / 1e6 + + def run_agg(): + return ctx.sql( + 'SELECT lat, lon, AVG(air) AS air_avg FROM "air" GROUP BY lat, lon' + ).to_dataset(dims=["lat", "lon"]) + + peaks = [] + for _ in range(3): + agg, peak = _peak_mb(run_agg) + peaks.append(peak) + agg_peak = min(peaks) + + assert agg.sizes["lat"] * agg.sizes["lon"] == ( + air_source.sizes["lat"] * air_source.sizes["lon"] + ) + assert agg_peak < 4 * source_mb, ( + f"GROUP BY reduction should not balloon past 4x source size; got " + f"best-of-{len(peaks)} peak={agg_peak:.1f} MB on a " + f"{source_mb:.1f} MB source" + ) + + # Sanity: values agree with the xarray-native reduction. + ref = ( + air_source.compute() + .mean(dim="time")["air"] + .sortby(["lat", "lon"]) + .values + ) + got = agg.sortby(["lat", "lon"])["air_avg"].values + np.testing.assert_allclose(got, ref, rtol=1e-5) diff --git a/uv.lock b/uv.lock index 14d56c2f..cf72eb88 100644 --- a/uv.lock +++ b/uv.lock @@ -117,6 +117,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "anyio" +version = "4.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "idna" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, +] + [[package]] name = "asciitree" version = "0.3.3" @@ -154,40 +168,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/50/cd/30110dc0ffcf3b131156077b90e9f60ed75711223f306da4db08eff8403b/beautifulsoup4-4.13.4-py3-none-any.whl", hash = "sha256:9bbbb14bfde9d79f38b8cd5f8c7c85f4b8f2523190ebed90e950a8dea4cb1c4b", size = 187285, upload-time = "2025-04-15T17:05:12.221Z" }, ] -[[package]] -name = "black" -version = "24.10.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d8/0d/cc2fb42b8c50d80143221515dd7e4766995bd07c56c9a3ed30baf080b6dc/black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875", size = 645813, upload-time = "2024-10-07T19:20:50.361Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/f3/465c0eb5cddf7dbbfe1fecd9b875d1dcf51b88923cd2c1d7e9ab95c6336b/black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812", size = 1623211, upload-time = "2024-10-07T19:26:12.43Z" }, - { url = "https://files.pythonhosted.org/packages/df/57/b6d2da7d200773fdfcc224ffb87052cf283cec4d7102fab450b4a05996d8/black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea", size = 1457139, upload-time = "2024-10-07T19:25:06.453Z" }, - { url = "https://files.pythonhosted.org/packages/6e/c5/9023b7673904a5188f9be81f5e129fff69f51f5515655fbd1d5a4e80a47b/black-24.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:649fff99a20bd06c6f727d2a27f401331dc0cc861fb69cde910fe95b01b5928f", size = 1753774, upload-time = "2024-10-07T19:23:58.47Z" }, - { url = "https://files.pythonhosted.org/packages/e1/32/df7f18bd0e724e0d9748829765455d6643ec847b3f87e77456fc99d0edab/black-24.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:fe4d6476887de70546212c99ac9bd803d90b42fc4767f058a0baa895013fbb3e", size = 1414209, upload-time = "2024-10-07T19:24:42.54Z" }, - { url = "https://files.pythonhosted.org/packages/c2/cc/7496bb63a9b06a954d3d0ac9fe7a73f3bf1cd92d7a58877c27f4ad1e9d41/black-24.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5a2221696a8224e335c28816a9d331a6c2ae15a2ee34ec857dcf3e45dbfa99ad", size = 1607468, upload-time = "2024-10-07T19:26:14.966Z" }, - { url = "https://files.pythonhosted.org/packages/2b/e3/69a738fb5ba18b5422f50b4f143544c664d7da40f09c13969b2fd52900e0/black-24.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f9da3333530dbcecc1be13e69c250ed8dfa67f43c4005fb537bb426e19200d50", size = 1437270, upload-time = "2024-10-07T19:25:24.291Z" }, - { url = "https://files.pythonhosted.org/packages/c9/9b/2db8045b45844665c720dcfe292fdaf2e49825810c0103e1191515fc101a/black-24.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4007b1393d902b48b36958a216c20c4482f601569d19ed1df294a496eb366392", size = 1737061, upload-time = "2024-10-07T19:23:52.18Z" }, - { url = "https://files.pythonhosted.org/packages/a3/95/17d4a09a5be5f8c65aa4a361444d95edc45def0de887810f508d3f65db7a/black-24.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:394d4ddc64782e51153eadcaaca95144ac4c35e27ef9b0a42e121ae7e57a9175", size = 1423293, upload-time = "2024-10-07T19:24:41.7Z" }, - { url = "https://files.pythonhosted.org/packages/90/04/bf74c71f592bcd761610bbf67e23e6a3cff824780761f536512437f1e655/black-24.10.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e39e0fae001df40f95bd8cc36b9165c5e2ea88900167bddf258bacef9bbdc3", size = 1644256, upload-time = "2024-10-07T19:27:53.355Z" }, - { url = "https://files.pythonhosted.org/packages/4c/ea/a77bab4cf1887f4b2e0bce5516ea0b3ff7d04ba96af21d65024629afedb6/black-24.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37d422772111794b26757c5b55a3eade028aa3fde43121ab7b673d050949d65", size = 1448534, upload-time = "2024-10-07T19:26:44.953Z" }, - { url = "https://files.pythonhosted.org/packages/4e/3e/443ef8bc1fbda78e61f79157f303893f3fddf19ca3c8989b163eb3469a12/black-24.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14b3502784f09ce2443830e3133dacf2c0110d45191ed470ecb04d0f5f6fcb0f", size = 1761892, upload-time = "2024-10-07T19:24:10.264Z" }, - { url = "https://files.pythonhosted.org/packages/52/93/eac95ff229049a6901bc84fec6908a5124b8a0b7c26ea766b3b8a5debd22/black-24.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:30d2c30dc5139211dda799758559d1b049f7f14c580c409d6ad925b74a4208a8", size = 1434796, upload-time = "2024-10-07T19:25:06.239Z" }, - { url = "https://files.pythonhosted.org/packages/d0/a0/a993f58d4ecfba035e61fca4e9f64a2ecae838fc9f33ab798c62173ed75c/black-24.10.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cbacacb19e922a1d75ef2b6ccaefcd6e93a2c05ede32f06a21386a04cedb981", size = 1643986, upload-time = "2024-10-07T19:28:50.684Z" }, - { url = "https://files.pythonhosted.org/packages/37/d5/602d0ef5dfcace3fb4f79c436762f130abd9ee8d950fa2abdbf8bbc555e0/black-24.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f93102e0c5bb3907451063e08b9876dbeac810e7da5a8bfb7aeb5a9ef89066b", size = 1448085, upload-time = "2024-10-07T19:28:12.093Z" }, - { url = "https://files.pythonhosted.org/packages/47/6d/a3a239e938960df1a662b93d6230d4f3e9b4a22982d060fc38c42f45a56b/black-24.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddacb691cdcdf77b96f549cf9591701d8db36b2f19519373d60d31746068dbf2", size = 1760928, upload-time = "2024-10-07T19:24:15.233Z" }, - { url = "https://files.pythonhosted.org/packages/dd/cf/af018e13b0eddfb434df4d9cd1b2b7892bab119f7a20123e93f6910982e8/black-24.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:680359d932801c76d2e9c9068d05c6b107f2584b2a5b88831c83962eb9984c1b", size = 1436875, upload-time = "2024-10-07T19:24:42.762Z" }, - { url = "https://files.pythonhosted.org/packages/8d/a7/4b27c50537ebca8bec139b872861f9d2bf501c5ec51fcf897cb924d9e264/black-24.10.0-py3-none-any.whl", hash = "sha256:3bb2b7a1f7b685f85b11fed1ef10f8a9148bceb49853e47a294a3dd963c1dd7d", size = 206898, upload-time = "2024-10-07T19:20:48.317Z" }, -] - [[package]] name = "cachetools" version = "5.5.2" @@ -219,6 +199,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/84/ae/320161bd181fc06471eed047ecce67b693fd7515b16d495d8932db763426/certifi-2025.6.15-py3-none-any.whl", hash = "sha256:2e0c7ce7cb5d8f8634ca55d2ba7e6ec2689a2fd6537d8dec1296a477a4910057", size = 157650, upload-time = "2025-06-15T02:45:49.977Z" }, ] +[[package]] +name = "cfgv" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/b5/721b8799b04bf9afe054a3899c6cf4e880fcf8563cc71c15610242490a0c/cfgv-3.5.0.tar.gz", hash = "sha256:d5b1034354820651caa73ede66a6294d6e95c1b00acc5e9b098e917404669132", size = 7334, upload-time = "2025-11-19T20:55:51.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, +] + [[package]] name = "cftime" version = "1.6.4.post1" @@ -435,19 +424,25 @@ wheels = [ [[package]] name = "datafusion" -version = "51.0.0" +version = "54.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cloudpickle" }, { name = "pyarrow" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2c/6d/d0e2632c93bbcca0687eeda672af3f92042ecd349df7be55da86253594a9/datafusion-51.0.0.tar.gz", hash = "sha256:1887c7d5ed3ae5d9f389e62ba869864afad4006a3f7c99ef0ca4707782a7838f", size = 193751, upload-time = "2026-01-09T13:23:41.562Z" } +sdist = { url = "https://files.pythonhosted.org/packages/60/90/886f7e9cf827f07ebd60bd293e54e0a028a50dd49bbaef0ee42aae1981ea/datafusion-54.0.0.tar.gz", hash = "sha256:cfe7e8dfc026efc05824f49b53ad6a72caf5c2d6820759b6212a09e245a427ed", size = 276448, upload-time = "2026-06-29T11:19:34.816Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cf/a9/7717cec053a3309be3020fe3147e3f76e5bf21295fa8adf9b52dd44ea3ff/datafusion-51.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:0c0d265fe3ee0dcbfa7cc3c64c7cd94fc493f38418bd79debb7ec29f29b7176e", size = 30389413, upload-time = "2026-01-09T13:23:23.266Z" }, - { url = "https://files.pythonhosted.org/packages/55/45/72c9874fd3740a4cb9d55049fdbae0df512dc5433e9f1176f3cfd970f1a1/datafusion-51.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:43e6011db86e950bf9a21ed73cc089c2346b340a41a4f1044268af6c3a357acc", size = 26982206, upload-time = "2026-01-09T13:23:27.437Z" }, - { url = "https://files.pythonhosted.org/packages/21/ac/b32ba1f25d38fc16e7623cc4bfb7bd68db61be2ef27b2d9969ea5c865765/datafusion-51.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e76803907150159aa059d5cc9291645bbaac1b6a46d07e56035118d327b741ae", size = 33246117, upload-time = "2026-01-09T13:23:30.981Z" }, - { url = "https://files.pythonhosted.org/packages/0b/4e/437121422ef010690fc3cdd7f080203e986ba00e0e3c3b577e03f5b54ca2/datafusion-51.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9d0cfabfe1853994adc2e6e9da5f36c1eb061102e34a2f1101fa935c6991c9e1", size = 31421867, upload-time = "2026-01-09T13:23:34.436Z" }, - { url = "https://files.pythonhosted.org/packages/db/fc/58cf27fcb85b2fd2a698253ae46213b1cbda784407e205c148f4006c1429/datafusion-51.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:fd5f9abfd6669062debf0658d13e4583234c89d4df95faf381927b11cea411f5", size = 32517679, upload-time = "2026-01-09T13:23:39.615Z" }, + { url = "https://files.pythonhosted.org/packages/46/58/4c5b981e3d9ade32a906c15a4941eef50c9b862781cdc14bf4dff48d026a/datafusion-54.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:946f55e48b8d523d7b4ac106bdf588b4493c2c66f81877d6952aafeaf7c3ec73", size = 39810553, upload-time = "2026-06-29T11:19:02.1Z" }, + { url = "https://files.pythonhosted.org/packages/66/e5/5e4dbd42ce9a2affb3be90d9ab17cebde1a6f28b0d9fb4b83d612d5c8e42/datafusion-54.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a3bf43185c7e43e25242e5fb17b6a11b86bf976434c0bc493fdedbd9a080969", size = 37145255, upload-time = "2026-06-29T11:19:05.491Z" }, + { url = "https://files.pythonhosted.org/packages/c6/5e/dbb9e6e3e5006d34f295d7ac73f1302c8f2df140666402a06e6c55028edb/datafusion-54.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9432bf162381e9282cbc74915b8b773895de18be836f7e3f6d0de4d981f24630", size = 38853856, upload-time = "2026-06-29T11:19:08.732Z" }, + { url = "https://files.pythonhosted.org/packages/a8/81/e69008e3479f4d0134875bc4ae39503bedcd55ca2597e71392c963c651b4/datafusion-54.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3bcd4d213fa74710e75e6e182cc468c2bdbc5ffc74a08c8155d414fbbfa1b3f6", size = 41050149, upload-time = "2026-06-29T11:19:12.108Z" }, + { url = "https://files.pythonhosted.org/packages/61/d4/8ba6e3fe3291c9ccc94b5ca3ec3c1fbcbfbe5ece5ffb965e4550844e2c56/datafusion-54.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:b934e097e1bdca7d5768a81ac1bc4a1812cb459269f8b1a5d892a5d930f18376", size = 43444869, upload-time = "2026-06-29T11:19:15.963Z" }, + { url = "https://files.pythonhosted.org/packages/9d/41/5608323226f21a0fa180823c531dbc0ed270e9b694f299b7647505cb6a06/datafusion-54.0.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:c4e79048da82ad89b768bd0be7df39254cd2a0afe2b719d1f129e8a7229af683", size = 39796248, upload-time = "2026-06-29T11:19:19.208Z" }, + { url = "https://files.pythonhosted.org/packages/18/81/392ee323104ab14ca689384723b69e137064a828233c165574f97a74c0e9/datafusion-54.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:fe57038003b18e28b90752c1e32b44af74ec4f552a1904aee725e1129a00c447", size = 37153577, upload-time = "2026-06-29T11:19:22.397Z" }, + { url = "https://files.pythonhosted.org/packages/40/c4/ebd5ef5349ecbea7f5f9da76c213581c13e7bbe1b5735c9925b279eeb4eb/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:574f642832a106456cfc4f32aa82484c504fc32f4be2b510202bcb579de8e6d1", size = 38849839, upload-time = "2026-06-29T11:19:25.783Z" }, + { url = "https://files.pythonhosted.org/packages/5a/b9/2383d30d317bb913cab97dbf2e6e1d5f37f594860d5c5bc176e025cf7d4a/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:796fd5683927443c5bc61999d00b9007ef9b5ce107725ea8d241df718860985d", size = 41074623, upload-time = "2026-06-29T11:19:29.119Z" }, + { url = "https://files.pythonhosted.org/packages/35/5c/553fd1107dede0a56727fda7216a7198d41394f2d19697f4fb104cc695ea/datafusion-54.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:64973c63874ec31670dd97b32b18af7b07fad679cb20d58ed154038e3a5c204e", size = 43438801, upload-time = "2026-06-29T11:19:32.799Z" }, ] [[package]] @@ -459,6 +454,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "deepmerge" +version = "2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a8/3a/b0ba594708f1ad0bc735884b3ad854d3ca3bdc1d741e56e40bbda6263499/deepmerge-2.0.tar.gz", hash = "sha256:5c3d86081fbebd04dd5de03626a0607b809a98fb6ccba5770b62466fe940ff20", size = 19890, upload-time = "2024-08-30T05:31:50.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2d/82/e5d2c1c67d19841e9edc74954c827444ae826978499bde3dfc1d007c8c11/deepmerge-2.0-py3-none-any.whl", hash = "sha256:6de9ce507115cff0bed95ff0ce9ecc31088ef50cbdf09bc90a09349a318b3d00", size = 13475, upload-time = "2024-08-30T05:31:48.659Z" }, +] + +[[package]] +name = "distlib" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/96/8e/709914eb2b5749865801041647dc7f4e6d00b549cfe88b65ca192995f07c/distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d", size = 614605, upload-time = "2025-07-17T16:52:00.465Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" }, +] + [[package]] name = "donfig" version = "0.8.1.post1" @@ -471,6 +484,48 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/d5/c5db1ea3394c6e1732fb3286b3bd878b59507a8f77d32a2cebda7d7b7cd4/donfig-0.8.1.post1-py3-none-any.whl", hash = "sha256:2a3175ce74a06109ff9307d90a230f81215cbac9a751f4d1c6194644b8204f9d", size = 21592, upload-time = "2024-05-23T14:13:55.283Z" }, ] +[[package]] +name = "duckdb" +version = "1.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/19/e57151753576373c6696a12022648546cca6038e8833fda2908ee2342d9b/duckdb-1.5.5.tar.gz", hash = "sha256:72f33ee57ca7595b23957671a2cc7f7fe2be0ecc2d68f63abedcfcaa3a5c1238", size = 18066741, upload-time = "2026-07-22T10:55:17.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/d4/298acf9331a80b3ce6ac64dd940e7e13f4058fb69d18914445f02e3c7bfe/duckdb-1.5.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:3b805507f88171b428b21c966c30e9a3d54e30b24528918a44ed0032542bc26f", size = 32702934, upload-time = "2026-07-22T10:53:19.069Z" }, + { url = "https://files.pythonhosted.org/packages/d5/90/c489fb63d64b2e7ee109ce8460bdede003a0f256e5b41a03a2a1c4764058/duckdb-1.5.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b08e19cc856220d8a26fa62abc2264b349aff67255e9373c6a3f607addd56dc6", size = 17343604, upload-time = "2026-07-22T10:53:22.767Z" }, + { url = "https://files.pythonhosted.org/packages/0b/27/effa80a15b1f0c61c235622f797868485359e8c9ad6a8e358e7a0c479151/duckdb-1.5.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a17e6a922e42a5c06ed2353fe78c5dff2610f6632d603836f9606ad0bf754079", size = 15488179, upload-time = "2026-07-22T10:53:25.945Z" }, + { url = "https://files.pythonhosted.org/packages/5d/07/21212345c8d24ba62dceaa20be3b21f5c46f1510b1b42ce93bb058afe0c4/duckdb-1.5.5-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1bdc38922c365c37720149f90d90b1e9823eb82dad6830855b5f87537fa6fc0c", size = 19367323, upload-time = "2026-07-22T10:53:30.23Z" }, + { url = "https://files.pythonhosted.org/packages/3f/d0/10371ae875fb4b5ef61bb892743b4b2e90c512b371fdf29317deb744857d/duckdb-1.5.5-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e238060db5ca59879882a6e9b015e2c65d5c64ddf281ba1d7a9a2033764152cf", size = 21476568, upload-time = "2026-07-22T10:53:33.486Z" }, + { url = "https://files.pythonhosted.org/packages/b2/35/09568ce617dd7bc0757b3d7b6a981660b9e4f0b7594de8ed776755eae740/duckdb-1.5.5-cp310-cp310-win_amd64.whl", hash = "sha256:4acc72798ba1885a9c17d1242903d2cd502f13b1271c7677f7cab25d8578eceb", size = 13156129, upload-time = "2026-07-22T10:53:37.55Z" }, + { url = "https://files.pythonhosted.org/packages/9c/c2/b62ec24d57bb8df4e24b0b58f7f8facb32f5fdb9f1895aed9e9fcdded168/duckdb-1.5.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:1b543841b0ae18a9c982345cfa3987e9c065d3a4b0f067daa473d92d1e65f528", size = 32708371, upload-time = "2026-07-22T10:53:41.642Z" }, + { url = "https://files.pythonhosted.org/packages/8a/ce/769171ba45f0b73632dc3bc3108d891e81dd6c6bbfba630a34a75b4dcc0f/duckdb-1.5.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1a925d06c2a4c3b64553d6cc1aced5028d376d4479bed689a7d47e9b1dccd80a", size = 17343979, upload-time = "2026-07-22T10:53:44.951Z" }, + { url = "https://files.pythonhosted.org/packages/46/59/a8e3384ee916e00d5dcf985194c1511d61978540778a1e96fa47f9fb3e0d/duckdb-1.5.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0c42757cb34722144bd4dfb94b6f336339e7b2468f6813fa7fa9a319ba07bab4", size = 15493704, upload-time = "2026-07-22T10:53:47.912Z" }, + { url = "https://files.pythonhosted.org/packages/6f/1d/9840179c2607b90523a2884a129c4d4e6dbdc1178ba62a976c1043beba88/duckdb-1.5.5-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e72f9e1a4f90a5c8483ad4d540e495bf0834ba61c360b52499a573d7ed62a3f", size = 19366574, upload-time = "2026-07-22T10:53:51.876Z" }, + { url = "https://files.pythonhosted.org/packages/b5/55/f9641a4eebcc2f4df631287d6c3b9ed2eea3b92644f93acbad825e3972b6/duckdb-1.5.5-cp311-cp311-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9b6f86ed85d4ef5e0211eaebf75d057bd8bb520bba438a95dd0f4e42234bbfe", size = 21477952, upload-time = "2026-07-22T10:53:55.575Z" }, + { url = "https://files.pythonhosted.org/packages/3c/3a/07c3556e37a5c97b95917b029c8fdde4a25fbd76a660bacdac195cf20dcb/duckdb-1.5.5-cp311-cp311-win_amd64.whl", hash = "sha256:9f4287f97ccf0c1f3d471e7115be2b067cbf99627e2d34bffd462dd64703cddc", size = 13156986, upload-time = "2026-07-22T10:53:58.823Z" }, + { url = "https://files.pythonhosted.org/packages/4f/ff/07b48eef2078ca033847e9caa46cc7633b714c5f91ad1ce091c8ca89d792/duckdb-1.5.5-cp311-cp311-win_arm64.whl", hash = "sha256:179633a3fc6296c75d57c69c1e239fa9e5cdcb670fd1dbff88a02663f932905c", size = 14001317, upload-time = "2026-07-22T10:54:01.724Z" }, + { url = "https://files.pythonhosted.org/packages/d6/40/2e05d324400fdaa5656c9f48d6298da421cb034d85e509fa0e6e325cf04b/duckdb-1.5.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d4dd65f8941a604b947e0b9b4b4f7165988e29a23ec0b69b4038520956d9933e", size = 32753858, upload-time = "2026-07-22T10:54:05.514Z" }, + { url = "https://files.pythonhosted.org/packages/79/15/5ceb58ffb5bb8a62b3fd7abb39c41467cdf94850ece02e6d88664dfc75ce/duckdb-1.5.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33db46679b071f108d57139493dee2d37e1f5efcf5c5c039c2969eed11a6c8a7", size = 17368293, upload-time = "2026-07-22T10:54:09.139Z" }, + { url = "https://files.pythonhosted.org/packages/bf/5c/bf02da0b354fe83cca4f95a4fbf762181af466f7d551ab2a093f7698882a/duckdb-1.5.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f0b88535a5d86fdd63dba6ea02ab68c003dfb9e4892b11256ef24c4da208baae", size = 15509131, upload-time = "2026-07-22T10:54:12.228Z" }, + { url = "https://files.pythonhosted.org/packages/ea/a9/5f1f09da421d8e930e0b063d11c1b3f90363f40ede74438cd188afdd13a2/duckdb-1.5.5-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f316eae2323d9a851883fdf2dee91c1f9efe251ab33e14a2272f82a913422ed6", size = 19391959, upload-time = "2026-07-22T10:54:15.551Z" }, + { url = "https://files.pythonhosted.org/packages/4f/98/6549769f158126fa64fd6c1ac2eb59a18282146c939867a3eb31b7c1db07/duckdb-1.5.5-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7a6d2d11859d82a936ebdcb30ce3d8a1cbb3e990bff05c12abb9b54c44fa7bd1", size = 21510909, upload-time = "2026-07-22T10:54:19.681Z" }, + { url = "https://files.pythonhosted.org/packages/af/b7/5753b41d3124838f868f9f523362812d9fc45409e9e4dd70dcbb0a25826e/duckdb-1.5.5-cp312-cp312-win_amd64.whl", hash = "sha256:ddfbdb096c11d51ee22492397d342c90a82e62c5d09961477895934d0a25372f", size = 13168544, upload-time = "2026-07-22T10:54:22.789Z" }, + { url = "https://files.pythonhosted.org/packages/5c/28/44b679c7d46245f8398feae7edac959d1b83d4eb143e25b3fce0630b78bd/duckdb-1.5.5-cp312-cp312-win_arm64.whl", hash = "sha256:2725d2b9ace3a4e75d72fc5a239f6a44b502c580edadb8fb2676db772c5f9282", size = 13988684, upload-time = "2026-07-22T10:54:26.003Z" }, + { url = "https://files.pythonhosted.org/packages/47/37/4a38116e7700720fd152c666292214fd3abdf916496991296d8d1f66efbf/duckdb-1.5.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:cd98829b67788609017e65c761bd42a5dd0f9129441bed8bda4d6881ccf819f0", size = 32754294, upload-time = "2026-07-22T10:54:29.822Z" }, + { url = "https://files.pythonhosted.org/packages/66/42/7d392f1ba1eee0eaf4ab4c8c7a604bfe3536cd63f979cf5c98798664f807/duckdb-1.5.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:feead93c56679b79592d437c62975d39cb67adedffa7592c763baf8160ac7366", size = 17368211, upload-time = "2026-07-22T10:54:33.359Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/0a6f4fa60562faa615e55e15bd1953a2f2b17a8edd8105e5cda215e43457/duckdb-1.5.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:49c963d9469373d7aba8d750d9ea565ab823e94166efed953f184dd9b169b98c", size = 15509136, upload-time = "2026-07-22T10:54:36.369Z" }, + { url = "https://files.pythonhosted.org/packages/e4/cb/023c89f51978545b9fab318581bba0c457a58e7530d2d933e54ae7d8647c/duckdb-1.5.5-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a736217825461732b5442d05a220f3da2e23a0dae114efbf08c9bf171b53098a", size = 19392147, upload-time = "2026-07-22T10:54:39.551Z" }, + { url = "https://files.pythonhosted.org/packages/3e/c5/41bef391fb8b23dbc133c9f2ba016e7a7a8124513d2cc1b430f1897d87e4/duckdb-1.5.5-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:078e6a60dd8eedde5832f45422ca5c4a6b8c837aeabd8a56ca0b7d933f588053", size = 21511060, upload-time = "2026-07-22T10:54:42.788Z" }, + { url = "https://files.pythonhosted.org/packages/07/9f/c44dfc1f924ac29b3252dc1b91393c01d009dbfe9f8ed33f10b986151bd1/duckdb-1.5.5-cp313-cp313-win_amd64.whl", hash = "sha256:6826504277dba513c0c5d71d828456c94d729c9d2482f94b2e289f90a9167e28", size = 13168028, upload-time = "2026-07-22T10:54:46.127Z" }, + { url = "https://files.pythonhosted.org/packages/ca/88/591384b2cd59abddd6f5dc175e60374f9abae6064429f0c4402854c10f44/duckdb-1.5.5-cp313-cp313-win_arm64.whl", hash = "sha256:baa9c5702002fabb559ded2a39008f9f421fcbc7237d388b8213eff1e08858de", size = 13989955, upload-time = "2026-07-22T10:54:49.262Z" }, + { url = "https://files.pythonhosted.org/packages/3e/56/12c65bfa2d2605b81981b264788891bcf11ec72227889554cead5d8d13b9/duckdb-1.5.5-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8e6413dd40facb7b8ab21bd844450cd8f549b29e138635be9cf090ef4d2049e2", size = 32761946, upload-time = "2026-07-22T10:54:53.412Z" }, + { url = "https://files.pythonhosted.org/packages/b9/46/682ce155f17e0d2822d4f13ee3db9ca4b5b7c2da61b841b2629035e1f4bc/duckdb-1.5.5-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:64078acfd16541132ac6e191eb81b2845554444a0305cc1aa581ba107e514aa8", size = 17375069, upload-time = "2026-07-22T10:54:57.269Z" }, + { url = "https://files.pythonhosted.org/packages/39/ce/a24bcbd3289c8f305a430759c5fc12242740b4af3e17f7593f3a34e333d2/duckdb-1.5.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8c11775cc99a447618d5f1840126db17f2652f3eae05529df4f81f40e2df7151", size = 15519791, upload-time = "2026-07-22T10:55:00.681Z" }, + { url = "https://files.pythonhosted.org/packages/d9/76/3a01afbc615c1d418c0de58a6b68ac5ce2a8563232c0464bfbc2ce552398/duckdb-1.5.5-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77bbc1e6ba12e1e06f9020117bdf848627ecfdf36f907550e62e008e6109dece", size = 19398251, upload-time = "2026-07-22T10:55:04.168Z" }, + { url = "https://files.pythonhosted.org/packages/a1/43/3a5e81d1728f4d234c79bfe385808ee7c04834f7c37a4b5c257459c25614/duckdb-1.5.5-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fbf0f2d48b43c6c304d00463b463c27ead6c4b01c3c1816b750f728decf71afe", size = 21513851, upload-time = "2026-07-22T10:55:07.864Z" }, + { url = "https://files.pythonhosted.org/packages/91/41/fc7c829172c60ca22485251eab285f4f1a0d87b486a024c726f21471d86e/duckdb-1.5.5-cp314-cp314-win_amd64.whl", hash = "sha256:9dc826c4b50e64f6c4e4d07a3a9cb075ef70ba3899dc43ec5493dc3d7b04b353", size = 13691858, upload-time = "2026-07-22T10:55:11.181Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2c/95d9216b79e9273689d7ebce125a54503ed0c9bd7da931f0265888e99779/duckdb-1.5.5-cp314-cp314-win_arm64.whl", hash = "sha256:63e48d4b74b15aeacd688976432a7225163df8c226eddeb8536bba2d4d4ff433", size = 14470180, upload-time = "2026-07-22T10:55:14.445Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.0" @@ -492,6 +547,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/bf/fd60001b3abc5222d8eaa4a204cd8c0ae78e75adc688f33ce4bf25b7fafa/fasteners-0.19-py3-none-any.whl", hash = "sha256:758819cb5d94cdedf4e836988b74de396ceacb8e2794d21f82d131fd9ee77237", size = 18679, upload-time = "2023-09-19T17:11:18.725Z" }, ] +[[package]] +name = "filelock" +version = "3.25.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, +] + [[package]] name = "frozenlist" version = "1.7.0" @@ -613,6 +677,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ff/eb/9182e875592c48d282c5eab602000f0618817b9011b2b2925165e4b4b7f3/gcsfs-2025.5.1-py2.py3-none-any.whl", hash = "sha256:48712471ff71ac83d3e2152ba4dc232874698466e344d5e700feba06b0a0de7b", size = 36581, upload-time = "2025-05-24T12:12:57.011Z" }, ] +[[package]] +name = "ghp-import" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/29/d40217cbe2f6b1359e00c6c307bb3fc876ba74068cbab3dde77f03ca0dc4/ghp-import-2.1.0.tar.gz", hash = "sha256:9c535c4c61193c2df8871222567d7fd7e5014d835f97dc7b7439069e2413d343", size = 10943, upload-time = "2022-05-02T15:47:16.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/ec/67fbef5d497f86283db54c22eec6f6140243aae73265799baaaa19cd17fb/ghp_import-2.1.0-py3-none-any.whl", hash = "sha256:8337dd7b50877f163d4c0289bc1f1c7f127550241988d568c1db512c4324a619", size = 11034, upload-time = "2022-05-02T15:47:14.552Z" }, +] + [[package]] name = "google-api-core" version = "2.25.1" @@ -745,6 +821,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/86/f1/62a193f0227cf15a920390abe675f386dec35f7ae3ffe6da582d3ade42c7/googleapis_common_protos-1.70.0-py3-none-any.whl", hash = "sha256:b8bfcca8c25a2bb253e0e0b0adaf8c00773e5e6af6fd92397576680b807e0fd8", size = 294530, upload-time = "2025-04-14T10:17:01.271Z" }, ] +[[package]] +name = "griffelib" +version = "2.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/71/d7/2b805e89cdc609e5b304361d80586b272ef00f6287ee63de1e571b1f71ec/griffelib-2.0.1.tar.gz", hash = "sha256:59f39eabb4c777483a3823e39e8f9e03e69df271a7e49aee64e91a8cfa91bdf5", size = 166383, upload-time = "2026-03-23T21:05:25.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/4c/cc8c68196db727cfc1432f2ad5de50aa6707e630d44b2e6361dc06d8f134/griffelib-2.0.1-py3-none-any.whl", hash = "sha256:b769eed581c0e857d362fc8fcd8e57ecd2330c124b6104ac8b4c1c86d76970aa", size = 142377, upload-time = "2026-03-23T21:04:01.116Z" }, +] + [[package]] name = "h5netcdf" version = "1.6.3" @@ -790,6 +875,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/6d/0084ed0b78d4fd3e7530c32491f2884140d9b06365dac8a08de726421d4a/h5py-3.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:ae18e3de237a7a830adb76aaa68ad438d85fe6e19e0d99944a3ce46b772c69b3", size = 2852929, upload-time = "2025-06-06T14:05:47.659Z" }, ] +[[package]] +name = "identify" +version = "2.6.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/c4/7fb4db12296cdb11893d61c92048fe617ee853f8523b9b296ac03b43757e/identify-2.6.18.tar.gz", hash = "sha256:873ac56a5e3fd63e7438a7ecbc4d91aca692eb3fefa4534db2b7913f3fc352fd", size = 99580, upload-time = "2026-03-15T18:39:50.319Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl", hash = "sha256:8db9d3c8ea9079db92cafb0ebf97abdc09d52e97f4dcf773a2e694048b7cd737", size = 99394, upload-time = "2026-03-15T18:39:48.915Z" }, +] + [[package]] name = "idna" version = "3.10" @@ -820,6 +914,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/e1/e6716421ea10d38022b952c159d5161ca1193197fb744506875fbb87ea7b/iniconfig-2.1.0-py3-none-any.whl", hash = "sha256:9deba5723312380e77435581c6bf4935c94cbfab9b1ed33ef8d238ea168eb760", size = 6050, upload-time = "2025-03-19T20:10:01.071Z" }, ] +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + [[package]] name = "legacy-cgi" version = "2.6.3" @@ -912,6 +1018,100 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/a5/899a4719e02ff4383f3f96e5d1878f882f734377f10dfb69e73b5f223e44/lxml-6.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:c86df1c9af35d903d2b52d22ea3e66db8058d21dc0f59842ca5deb0595921141", size = 3517946, upload-time = "2025-06-26T16:28:07.665Z" }, ] +[[package]] +name = "markdown" +version = "3.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/4b/3541d44f3937ba468b75da9eebcae497dcf67adb65caa16760b0a6807ebb/markupsafe-3.0.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559", size = 11631, upload-time = "2025-09-27T18:36:05.558Z" }, + { url = "https://files.pythonhosted.org/packages/98/1b/fbd8eed11021cabd9226c37342fa6ca4e8a98d8188a8d9b66740494960e4/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419", size = 12057, upload-time = "2025-09-27T18:36:07.165Z" }, + { url = "https://files.pythonhosted.org/packages/40/01/e560d658dc0bb8ab762670ece35281dec7b6c1b33f5fbc09ebb57a185519/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ba88449deb3de88bd40044603fafffb7bc2b055d626a330323a9ed736661695", size = 22050, upload-time = "2025-09-27T18:36:08.005Z" }, + { url = "https://files.pythonhosted.org/packages/af/cd/ce6e848bbf2c32314c9b237839119c5a564a59725b53157c856e90937b7a/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f42d0984e947b8adf7dd6dde396e720934d12c506ce84eea8476409563607591", size = 20681, upload-time = "2025-09-27T18:36:08.881Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2a/b5c12c809f1c3045c4d580b035a743d12fcde53cf685dbc44660826308da/markupsafe-3.0.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c0c0b3ade1c0b13b936d7970b1d37a57acde9199dc2aecc4c336773e1d86049c", size = 20705, upload-time = "2025-09-27T18:36:10.131Z" }, + { url = "https://files.pythonhosted.org/packages/cf/e3/9427a68c82728d0a88c50f890d0fc072a1484de2f3ac1ad0bfc1a7214fd5/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:0303439a41979d9e74d18ff5e2dd8c43ed6c6001fd40e5bf2e43f7bd9bbc523f", size = 21524, upload-time = "2025-09-27T18:36:11.324Z" }, + { url = "https://files.pythonhosted.org/packages/bc/36/23578f29e9e582a4d0278e009b38081dbe363c5e7165113fad546918a232/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:d2ee202e79d8ed691ceebae8e0486bd9a2cd4794cec4824e1c99b6f5009502f6", size = 20282, upload-time = "2025-09-27T18:36:12.573Z" }, + { url = "https://files.pythonhosted.org/packages/56/21/dca11354e756ebd03e036bd8ad58d6d7168c80ce1fe5e75218e4945cbab7/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:177b5253b2834fe3678cb4a5f0059808258584c559193998be2601324fdeafb1", size = 20745, upload-time = "2025-09-27T18:36:13.504Z" }, + { url = "https://files.pythonhosted.org/packages/87/99/faba9369a7ad6e4d10b6a5fbf71fa2a188fe4a593b15f0963b73859a1bbd/markupsafe-3.0.3-cp310-cp310-win32.whl", hash = "sha256:2a15a08b17dd94c53a1da0438822d70ebcd13f8c3a95abe3a9ef9f11a94830aa", size = 14571, upload-time = "2025-09-27T18:36:14.779Z" }, + { url = "https://files.pythonhosted.org/packages/d6/25/55dc3ab959917602c96985cb1253efaa4ff42f71194bddeb61eb7278b8be/markupsafe-3.0.3-cp310-cp310-win_amd64.whl", hash = "sha256:c4ffb7ebf07cfe8931028e3e4c85f0357459a3f9f9490886198848f4fa002ec8", size = 15056, upload-time = "2025-09-27T18:36:16.125Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9e/0a02226640c255d1da0b8d12e24ac2aa6734da68bff14c05dd53b94a0fc3/markupsafe-3.0.3-cp310-cp310-win_arm64.whl", hash = "sha256:e2103a929dfa2fcaf9bb4e7c091983a49c9ac3b19c9061b6d5427dd7d14d81a1", size = 13932, upload-time = "2025-09-27T18:36:17.311Z" }, + { url = "https://files.pythonhosted.org/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad", size = 11631, upload-time = "2025-09-27T18:36:18.185Z" }, + { url = "https://files.pythonhosted.org/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a", size = 12058, upload-time = "2025-09-27T18:36:19.444Z" }, + { url = "https://files.pythonhosted.org/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50", size = 24287, upload-time = "2025-09-27T18:36:20.768Z" }, + { url = "https://files.pythonhosted.org/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf", size = 22940, upload-time = "2025-09-27T18:36:22.249Z" }, + { url = "https://files.pythonhosted.org/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f", size = 21887, upload-time = "2025-09-27T18:36:23.535Z" }, + { url = "https://files.pythonhosted.org/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a", size = 23692, upload-time = "2025-09-27T18:36:24.823Z" }, + { url = "https://files.pythonhosted.org/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115", size = 21471, upload-time = "2025-09-27T18:36:25.95Z" }, + { url = "https://files.pythonhosted.org/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a", size = 22923, upload-time = "2025-09-27T18:36:27.109Z" }, + { url = "https://files.pythonhosted.org/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19", size = 14572, upload-time = "2025-09-27T18:36:28.045Z" }, + { url = "https://files.pythonhosted.org/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01", size = 15077, upload-time = "2025-09-27T18:36:29.025Z" }, + { url = "https://files.pythonhosted.org/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c", size = 13876, upload-time = "2025-09-27T18:36:29.954Z" }, + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + [[package]] name = "maturin" version = "1.11.5" @@ -936,6 +1136,104 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/67/c94f8f5440bc42d54113a2d99de0d6107f06b5a33f31823e52b2715d856f/maturin-1.11.5-py3-none-win_arm64.whl", hash = "sha256:9348f7f0a346108e0c96e6719be91da4470bd43c15802435e9f4157f5cca43d4", size = 7624029, upload-time = "2026-01-09T11:06:08.728Z" }, ] +[[package]] +name = "mergedeep" +version = "1.3.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/41/580bb4006e3ed0361b8151a01d324fb03f420815446c7def45d02f74c270/mergedeep-1.3.4.tar.gz", hash = "sha256:0096d52e9dad9939c3d975a774666af186eda617e6ca84df4c94dec30004f2a8", size = 4661, upload-time = "2021-02-05T18:55:30.623Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/19/04f9b178c2d8a15b076c8b5140708fa6ffc5601fb6f1e975537072df5b2a/mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307", size = 6354, upload-time = "2021-02-05T18:55:29.583Z" }, +] + +[[package]] +name = "mkdocs" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "ghp-import" }, + { name = "jinja2" }, + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mergedeep" }, + { name = "mkdocs-get-deps" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "pyyaml" }, + { name = "pyyaml-env-tag" }, + { name = "watchdog" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/c6/bbd4f061bd16b378247f12953ffcb04786a618ce5e904b8c5a01a0309061/mkdocs-1.6.1.tar.gz", hash = "sha256:7b432f01d928c084353ab39c57282f29f92136665bdd6abf7c1ec8d822ef86f2", size = 3889159, upload-time = "2024-08-30T12:24:06.899Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/5b/dbc6a8cddc9cfa9c4971d59fb12bb8d42e161b7e7f8cc89e49137c5b279c/mkdocs-1.6.1-py3-none-any.whl", hash = "sha256:db91759624d1647f3f34aa0c3f327dd2601beae39a366d6e064c03468d35c20e", size = 3864451, upload-time = "2024-08-30T12:24:05.054Z" }, +] + +[[package]] +name = "mkdocs-autorefs" +version = "1.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mkdocs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/52/c0/f641843de3f612a6b48253f39244165acff36657a91cc903633d456ae1ac/mkdocs_autorefs-1.4.4.tar.gz", hash = "sha256:d54a284f27a7346b9c38f1f852177940c222da508e66edc816a0fa55fc6da197", size = 56588, upload-time = "2026-02-10T15:23:55.105Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/de/a3e710469772c6a89595fc52816da05c1e164b4c866a89e3cb82fb1b67c5/mkdocs_autorefs-1.4.4-py3-none-any.whl", hash = "sha256:834ef5408d827071ad1bc69e0f39704fa34c7fc05bc8e1c72b227dfdc5c76089", size = 25530, upload-time = "2026-02-10T15:23:53.817Z" }, +] + +[[package]] +name = "mkdocs-get-deps" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mergedeep" }, + { name = "platformdirs" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/25/b3cccb187655b9393572bde9b09261d267c3bf2f2cdabe347673be5976a6/mkdocs_get_deps-0.2.2.tar.gz", hash = "sha256:8ee8d5f316cdbbb2834bc1df6e69c08fe769a83e040060de26d3c19fad3599a1", size = 11047, upload-time = "2026-03-10T02:46:33.632Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/29/744136411e785c4b0b744d5413e56555265939ab3a104c6a4b719dad33fd/mkdocs_get_deps-0.2.2-py3-none-any.whl", hash = "sha256:e7878cbeac04860b8b5e0ca31d3abad3df9411a75a32cde82f8e44b6c16ff650", size = 9555, upload-time = "2026-03-10T02:46:32.256Z" }, +] + +[[package]] +name = "mkdocstrings" +version = "1.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mkdocs" }, + { name = "mkdocs-autorefs" }, + { name = "pymdown-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/46/62/0dfc5719514115bf1781f44b1d7f2a0923fcc01e9c5d7990e48a05c9ae5d/mkdocstrings-1.0.3.tar.gz", hash = "sha256:ab670f55040722b49bb45865b2e93b824450fb4aef638b00d7acb493a9020434", size = 100946, upload-time = "2026-02-07T14:31:40.973Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/41/1cf02e3df279d2dd846a1bf235a928254eba9006dd22b4a14caa71aed0f7/mkdocstrings-1.0.3-py3-none-any.whl", hash = "sha256:0d66d18430c2201dc7fe85134277382baaa15e6b30979f3f3bdbabd6dbdb6046", size = 35523, upload-time = "2026-02-07T14:31:39.27Z" }, +] + +[package.optional-dependencies] +python = [ + { name = "mkdocstrings-python" }, +] + +[[package]] +name = "mkdocstrings-python" +version = "2.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "griffelib" }, + { name = "mkdocs-autorefs" }, + { name = "mkdocstrings" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/29/33/c225eaf898634bdda489a6766fc35d1683c640bffe0e0acd10646b13536d/mkdocstrings_python-2.0.3.tar.gz", hash = "sha256:c518632751cc869439b31c9d3177678ad2bfa5c21b79b863956ad68fc92c13b8", size = 199083, upload-time = "2026-02-20T10:38:36.368Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/28/79f0f8de97cce916d5ae88a7bee1ad724855e83e6019c0b4d5b3fabc80f3/mkdocstrings_python-2.0.3-py3-none-any.whl", hash = "sha256:0b83513478bdfd803ff05aa43e9b1fca9dd22bcd9471f09ca6257f009bc5ee12", size = 104779, upload-time = "2026-02-20T10:38:34.517Z" }, +] + [[package]] name = "multidict" version = "6.6.3" @@ -1038,15 +1336,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/30/9aec301e9772b098c1f5c0ca0279237c9766d94b97802e9888010c64b0ed/multidict-6.6.3-py3-none-any.whl", hash = "sha256:8db10f29c7541fc5da4defd8cd697e1ca429db743fa716325f236079b96f775a", size = 12313, upload-time = "2025-06-30T15:53:45.437Z" }, ] -[[package]] -name = "mypy-extensions" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, -] - [[package]] name = "netcdf4" version = "1.7.2" @@ -1091,6 +1380,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cd/10/c52f12297965938d9b9be666ea1f9d8340c2aea31d6909d90aa650847248/netcdf4-1.7.2-cp311-abi3-win_amd64.whl", hash = "sha256:999bfc4acebf400ed724d5e7329e2e768accc7ee1fa1d82d505da782f730301b", size = 7148514, upload-time = "2025-10-13T18:32:33.121Z" }, ] +[[package]] +name = "nodeenv" +version = "1.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, +] + [[package]] name = "numcodecs" version = "0.13.1" @@ -1395,6 +1693,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "polars" +version = "1.43.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "polars-runtime-32" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/13/3873f213304bcbaaf39e63c8b905ceb460a0524448d57f86a829f6d4d0fd/polars-1.43.2.tar.gz", hash = "sha256:c699671b99eb71ff53334d237917aaa3db5ad4dda480abcb6c80e0eaee7b677b", size = 750312, upload-time = "2026-08-01T06:28:30.872Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/fe/0888040a24e4504098b85d8ad486b14cb01cf6b030bbe479dfc2dcffc2ac/polars-1.43.2-py3-none-any.whl", hash = "sha256:22aa0cb92a1ee2d60d6a15a638b2e8e0dd99aea21ac0cd8fb29da8e382e075a9", size = 847150, upload-time = "2026-08-01T06:27:15.543Z" }, +] + +[[package]] +name = "polars-runtime-32" +version = "1.43.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/06/11b578eeef05f867e3ee31b2a2fdd8e7684c2aa47822c49935d1be789c38/polars_runtime_32-1.43.2.tar.gz", hash = "sha256:d7b7c486bccee75a6af0158b87077da3d054657e3c60036b28644f4e1c7fdbf7", size = 3095669, upload-time = "2026-08-01T06:28:32.315Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/fc/12e6d4ca34d820297651134cfa35f86c33e898539fc6629cbb35d0089697/polars_runtime_32-1.43.2-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:91abf205d4ec93f92ba95386b7f8776559ae3dfce425ed2e527efa75d117d04a", size = 53088908, upload-time = "2026-08-01T06:27:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/32/81/833b0853551deb810854f96b43dea342b6e6c9b0ea1afcccf774157d519d/polars_runtime_32-1.43.2-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2cc3ff96fd44789b02eb5c15b98dfcb000101636b177d3034fae2feec19b118f", size = 47540529, upload-time = "2026-08-01T06:27:21.391Z" }, + { url = "https://files.pythonhosted.org/packages/83/55/7b2a75af14c9294d97f3bec132dd3018ddcd988bef32b5d28322150b8c11/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:10ed36e615ab362feb7406e6d084e124b445ad284caa73bd93ae7e65745ed894", size = 51366340, upload-time = "2026-08-01T06:27:24.776Z" }, + { url = "https://files.pythonhosted.org/packages/62/60/64deacb3abc70c52e2d88a808a052d1621c86a48fe9194f2c065579ab1cd/polars_runtime_32-1.43.2-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6d5a7ae004a2723ebf4427f6d6a639f30f86af4cf077075f6b35d04711154fc3", size = 57304599, upload-time = "2026-08-01T06:27:27.875Z" }, + { url = "https://files.pythonhosted.org/packages/52/95/d6e3a236d7630e17c40d0ddee839bf2be9acf548fdc0e5ad65ed9ff0cac6/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:09339eacc6d392206e78aabbaaa37d7276eb969b798f46cb1f367fd718798c60", size = 51520580, upload-time = "2026-08-01T06:27:30.913Z" }, + { url = "https://files.pythonhosted.org/packages/b6/5a/2deb8eac70e9a2ac26d88a66ae7cf52612865026f4f4a5e7ab11ad9d52bf/polars_runtime_32-1.43.2-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:452b400e59e7f56e4c6437f435e796903272a9388feee12de2bea049ae87025e", size = 55204471, upload-time = "2026-08-01T06:27:33.886Z" }, + { url = "https://files.pythonhosted.org/packages/29/9e/647401ae8a607bc0cc40ed7b8592d5b1be90ded0dc9b9d6d3aeb03f9524b/polars_runtime_32-1.43.2-cp310-abi3-win_amd64.whl", hash = "sha256:00e33c28e321410c8d66e814a90043101e3bdd9ed2c6dabda07565aa8adbbdf1", size = 52572176, upload-time = "2026-08-01T06:27:37.048Z" }, + { url = "https://files.pythonhosted.org/packages/96/8d/60a50c3f36c85218a7ffcb48c6fe2ce1f7bec799152d68b8658ebed2179c/polars_runtime_32-1.43.2-cp310-abi3-win_arm64.whl", hash = "sha256:350a4868cae85bf8b3f81b33ba47927c15256bd9264dfc8c0753f1b927eac9d3", size = 46582513, upload-time = "2026-08-01T06:27:40.025Z" }, +] + [[package]] name = "pooch" version = "1.8.2" @@ -1409,6 +1735,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/87/77cc11c7a9ea9fd05503def69e3d18605852cd0d4b0d3b8f15bbeb3ef1d1/pooch-1.8.2-py3-none-any.whl", hash = "sha256:3529a57096f7198778a5ceefd5ac3ef0e4d06a6ddaf9fc2d609b806f25302c47", size = 64574, upload-time = "2024-06-06T16:53:44.343Z" }, ] +[[package]] +name = "pre-commit" +version = "4.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cfgv" }, + { name = "identify" }, + { name = "nodeenv" }, + { name = "pyyaml" }, + { name = "virtualenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/40/f1/6d86a29246dfd2e9b6237f0b5823717f60cad94d47ddc26afa916d21f525/pre_commit-4.5.1.tar.gz", hash = "sha256:eb545fcff725875197837263e977ea257a402056661f09dae08e4b149b030a61", size = 198232, upload-time = "2025-12-16T21:14:33.552Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl", hash = "sha256:3b3afd891e97337708c1674210f8eba659b52a38ea5f822ff142d10786221f77", size = 226437, upload-time = "2025-12-16T21:14:32.409Z" }, +] + [[package]] name = "propcache" version = "0.3.2" @@ -1645,22 +1987,133 @@ wheels = [ ] [[package]] -name = "pyink" -version = "24.10.1" +name = "pymdown-extensions" +version = "10.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "black" }, - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "markdown" }, + { name = "pyyaml" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d1/a1/e5e28626fca4266a94c2e1c9264fbf915b9e83e94f52e965190e48fd0cbf/pyink-24.10.1.tar.gz", hash = "sha256:5ec4339aa4953f796e88d90bcac3e3472161e4c36dbde203d80f5f76721ac718", size = 267230, upload-time = "2025-01-10T11:28:09.907Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/63/06673d1eb6d8f83c0ea1f677d770e12565fb516928b4109c9e2055656a9e/pymdown_extensions-10.21.tar.gz", hash = "sha256:39f4a020f40773f6b2ff31d2cd2546c2c04d0a6498c31d9c688d2be07e1767d5", size = 853363, upload-time = "2026-02-15T20:44:06.748Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/12/2f271b3601ae25731879f160d6b3941d80eb6b4f3e24be90289e33fb1dc4/pyink-24.10.1-py3-none-any.whl", hash = "sha256:6349bf6ab75e2ea39a5f0bc3dee7ede7f4af8529291472638026de5fd4af80d2", size = 137118, upload-time = "2025-01-10T11:28:06.138Z" }, + { url = "https://files.pythonhosted.org/packages/6f/2c/5b079febdc65e1c3fb2729bf958d18b45be7113828528e8a0b5850dd819a/pymdown_extensions-10.21-py3-none-any.whl", hash = "sha256:91b879f9f864d49794c2d9534372b10150e6141096c3908a455e45ca72ad9d3f", size = 268877, upload-time = "2026-02-15T20:44:05.464Z" }, +] + +[[package]] +name = "pyproj" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "certifi", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/10/a8480ea27ea4bbe896c168808854d00f2a9b49f95c0319ddcbba693c8a90/pyproj-3.7.1.tar.gz", hash = "sha256:60d72facd7b6b79853f19744779abcd3f804c4e0d4fa8815469db20c9f640a47", size = 226339, upload-time = "2025-02-16T04:28:46.621Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/a3/c4cd4bba5b336075f145fe784fcaf4ef56ffbc979833303303e7a659dda2/pyproj-3.7.1-cp310-cp310-macosx_13_0_x86_64.whl", hash = "sha256:bf09dbeb333c34e9c546364e7df1ff40474f9fddf9e70657ecb0e4f670ff0b0e", size = 6262524, upload-time = "2025-02-16T04:27:19.725Z" }, + { url = "https://files.pythonhosted.org/packages/40/45/4fdf18f4cc1995f1992771d2a51cf186a9d7a8ec973c9693f8453850c707/pyproj-3.7.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:6575b2e53cc9e3e461ad6f0692a5564b96e7782c28631c7771c668770915e169", size = 4665102, upload-time = "2025-02-16T04:27:24.428Z" }, + { url = "https://files.pythonhosted.org/packages/0c/d2/360eb127380106cee83569954ae696b88a891c804d7a93abe3fbc15f5976/pyproj-3.7.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8cb516ee35ed57789b46b96080edf4e503fdb62dbb2e3c6581e0d6c83fca014b", size = 9432667, upload-time = "2025-02-16T04:27:27.04Z" }, + { url = "https://files.pythonhosted.org/packages/76/a5/c6e11b9a99ce146741fb4d184d5c468446c6d6015b183cae82ac822a6cfa/pyproj-3.7.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1e47c4e93b88d99dd118875ee3ca0171932444cdc0b52d493371b5d98d0f30ee", size = 9259185, upload-time = "2025-02-16T04:27:30.35Z" }, + { url = "https://files.pythonhosted.org/packages/41/56/a3c15c42145797a99363fa0fdb4e9805dccb8b4a76a6d7b2cdf36ebcc2a1/pyproj-3.7.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:3e8d276caeae34fcbe4813855d0d97b9b825bab8d7a8b86d859c24a6213a5a0d", size = 10469103, upload-time = "2025-02-16T04:27:33.542Z" }, + { url = "https://files.pythonhosted.org/packages/ef/73/c9194c2802fefe2a4fd4230bdd5ab083e7604e93c64d0356fa49c363bad6/pyproj-3.7.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f173f851ee75e54acdaa053382b6825b400cb2085663a9bb073728a59c60aebb", size = 10401391, upload-time = "2025-02-16T04:27:36.051Z" }, + { url = "https://files.pythonhosted.org/packages/c5/1d/ce8bb5b9251b04d7c22d63619bb3db3d2397f79000a9ae05b3fd86a5837e/pyproj-3.7.1-cp310-cp310-win32.whl", hash = "sha256:f550281ed6e5ea88fcf04a7c6154e246d5714be495c50c9e8e6b12d3fb63e158", size = 5869997, upload-time = "2025-02-16T04:27:38.302Z" }, + { url = "https://files.pythonhosted.org/packages/09/6a/ca145467fd2e5b21e3d5b8c2b9645dcfb3b68f08b62417699a1f5689008e/pyproj-3.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:3537668992a709a2e7f068069192138618c00d0ba113572fdd5ee5ffde8222f3", size = 6278581, upload-time = "2025-02-16T04:27:41.051Z" }, + { url = "https://files.pythonhosted.org/packages/ab/0d/63670fc527e664068b70b7cab599aa38b7420dd009bdc29ea257e7f3dfb3/pyproj-3.7.1-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:a94e26c1a4950cea40116775588a2ca7cf56f1f434ff54ee35a84718f3841a3d", size = 6264315, upload-time = "2025-02-16T04:27:44.539Z" }, + { url = "https://files.pythonhosted.org/packages/25/9d/cbaf82cfb290d1f1fa42feb9ba9464013bb3891e40c4199f8072112e4589/pyproj-3.7.1-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:263b54ba5004b6b957d55757d846fc5081bc02980caa0279c4fc95fa0fff6067", size = 4666267, upload-time = "2025-02-16T04:27:47.019Z" }, + { url = "https://files.pythonhosted.org/packages/79/53/24f9f9b8918c0550f3ff49ad5de4cf3f0688c9f91ff191476db8979146fe/pyproj-3.7.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f6d6a2ccd5607cd15ef990c51e6f2dd27ec0a741e72069c387088bba3aab60fa", size = 9680510, upload-time = "2025-02-16T04:27:49.239Z" }, + { url = "https://files.pythonhosted.org/packages/3c/ac/12fab74a908d40b63174dc704587febd0729414804bbfd873cabe504ff2d/pyproj-3.7.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8c5dcf24ede53d8abab7d8a77f69ff1936c6a8843ef4fcc574646e4be66e5739", size = 9493619, upload-time = "2025-02-16T04:27:52.65Z" }, + { url = "https://files.pythonhosted.org/packages/c4/45/26311d6437135da2153a178125db5dfb6abce831ce04d10ec207eabac70a/pyproj-3.7.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3c2e7449840a44ce860d8bea2c6c1c4bc63fa07cba801dcce581d14dcb031a02", size = 10709755, upload-time = "2025-02-16T04:27:55.239Z" }, + { url = "https://files.pythonhosted.org/packages/99/52/4ecd0986f27d0e6c8ee3a7bc5c63da15acd30ac23034f871325b297e61fd/pyproj-3.7.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0829865c1d3a3543f918b3919dc601eea572d6091c0dd175e1a054db9c109274", size = 10642970, upload-time = "2025-02-16T04:27:58.343Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a5/d3bfc018fc92195a000d1d28acc1f3f1df15ff9f09ece68f45a2636c0134/pyproj-3.7.1-cp311-cp311-win32.whl", hash = "sha256:6181960b4b812e82e588407fe5c9c68ada267c3b084db078f248db5d7f45d18a", size = 5868295, upload-time = "2025-02-16T04:28:01.712Z" }, + { url = "https://files.pythonhosted.org/packages/92/39/ef6f06a5b223dbea308cfcbb7a0f72e7b506aef1850e061b2c73b0818715/pyproj-3.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:5ad0ff443a785d84e2b380869fdd82e6bfc11eba6057d25b4409a9bbfa867970", size = 6279871, upload-time = "2025-02-16T04:28:04.988Z" }, + { url = "https://files.pythonhosted.org/packages/e6/c9/876d4345b8d17f37ac59ebd39f8fa52fc6a6a9891a420f72d050edb6b899/pyproj-3.7.1-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:2781029d90df7f8d431e29562a3f2d8eafdf233c4010d6fc0381858dc7373217", size = 6264087, upload-time = "2025-02-16T04:28:09.036Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e6/5f8691f8c90e7f402cc80a6276eb19d2ec1faa150d5ae2dd9c7b0a254da8/pyproj-3.7.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:d61bf8ab04c73c1da08eedaf21a103b72fa5b0a9b854762905f65ff8b375d394", size = 4669628, upload-time = "2025-02-16T04:28:10.944Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/16475bbb79c1c68845c0a0d9c60c4fb31e61b8a2a20bc18b1a81e81c7f68/pyproj-3.7.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:04abc517a8555d1b05fcee768db3280143fe42ec39fdd926a2feef31631a1f2f", size = 9721415, upload-time = "2025-02-16T04:28:13.342Z" }, + { url = "https://files.pythonhosted.org/packages/b3/a3/448f05b15e318bd6bea9a32cfaf11e886c4ae61fa3eee6e09ed5c3b74bb2/pyproj-3.7.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:084c0a475688f934d386c2ab3b6ce03398a473cd48adfda70d9ab8f87f2394a0", size = 9556447, upload-time = "2025-02-16T04:28:15.818Z" }, + { url = "https://files.pythonhosted.org/packages/6a/ae/bd15fe8d8bd914ead6d60bca7f895a4e6f8ef7e3928295134ff9a7dad14c/pyproj-3.7.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a20727a23b1e49c7dc7fe3c3df8e56a8a7acdade80ac2f5cca29d7ca5564c145", size = 10758317, upload-time = "2025-02-16T04:28:18.338Z" }, + { url = "https://files.pythonhosted.org/packages/9d/d9/5ccefb8bca925f44256b188a91c31238cae29ab6ee7f53661ecc04616146/pyproj-3.7.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bf84d766646f1ebd706d883755df4370aaf02b48187cedaa7e4239f16bc8213d", size = 10771259, upload-time = "2025-02-16T04:28:20.822Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7d/31dedff9c35fa703162f922eeb0baa6c44a3288469a5fd88d209e2892f9e/pyproj-3.7.1-cp312-cp312-win32.whl", hash = "sha256:5f0da2711364d7cb9f115b52289d4a9b61e8bca0da57f44a3a9d6fc9bdeb7274", size = 5859914, upload-time = "2025-02-16T04:28:23.303Z" }, + { url = "https://files.pythonhosted.org/packages/3e/47/c6ab03d6564a7c937590cff81a2742b5990f096cce7c1a622d325be340ee/pyproj-3.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:aee664a9d806612af30a19dba49e55a7a78ebfec3e9d198f6a6176e1d140ec98", size = 6273196, upload-time = "2025-02-16T04:28:25.227Z" }, + { url = "https://files.pythonhosted.org/packages/ef/01/984828464c9960036c602753fc0f21f24f0aa9043c18fa3f2f2b66a86340/pyproj-3.7.1-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:5f8d02ef4431dee414d1753d13fa82a21a2f61494737b5f642ea668d76164d6d", size = 6253062, upload-time = "2025-02-16T04:28:27.861Z" }, + { url = "https://files.pythonhosted.org/packages/68/65/6ecdcdc829811a2c160cdfe2f068a009fc572fd4349664f758ccb0853a7c/pyproj-3.7.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:0b853ae99bda66cbe24b4ccfe26d70601d84375940a47f553413d9df570065e0", size = 4660548, upload-time = "2025-02-16T04:28:29.526Z" }, + { url = "https://files.pythonhosted.org/packages/67/da/dda94c4490803679230ba4c17a12f151b307a0d58e8110820405ca2d98db/pyproj-3.7.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:83db380c52087f9e9bdd8a527943b2e7324f275881125e39475c4f9277bdeec4", size = 9662464, upload-time = "2025-02-16T04:28:31.437Z" }, + { url = "https://files.pythonhosted.org/packages/6f/57/f61b7d22c91ae1d12ee00ac4c0038714e774ebcd851b9133e5f4f930dd40/pyproj-3.7.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b35ed213892e211a3ce2bea002aa1183e1a2a9b79e51bb3c6b15549a831ae528", size = 9497461, upload-time = "2025-02-16T04:28:33.848Z" }, + { url = "https://files.pythonhosted.org/packages/b7/f6/932128236f79d2ac7d39fe1a19667fdf7155d9a81d31fb9472a7a497790f/pyproj-3.7.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a8b15b0463d1303bab113d1a6af2860a0d79013c3a66fcc5475ce26ef717fd4f", size = 10708869, upload-time = "2025-02-16T04:28:37.34Z" }, + { url = "https://files.pythonhosted.org/packages/1d/0d/07ac7712994454a254c383c0d08aff9916a2851e6512d59da8dc369b1b02/pyproj-3.7.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:87229e42b75e89f4dad6459200f92988c5998dfb093c7c631fb48524c86cd5dc", size = 10729260, upload-time = "2025-02-16T04:28:40.639Z" }, + { url = "https://files.pythonhosted.org/packages/b0/d0/9c604bc72c37ba69b867b6df724d6a5af6789e8c375022c952f65b2af558/pyproj-3.7.1-cp313-cp313-win32.whl", hash = "sha256:d666c3a3faaf3b1d7fc4a544059c4eab9d06f84a604b070b7aa2f318e227798e", size = 5855462, upload-time = "2025-02-16T04:28:42.827Z" }, + { url = "https://files.pythonhosted.org/packages/98/df/68a2b7f5fb6400c64aad82d72bcc4bc531775e62eedff993a77c780defd0/pyproj-3.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:d3caac7473be22b6d6e102dde6c46de73b96bc98334e577dfaee9886f102ea2e", size = 6266573, upload-time = "2025-02-16T04:28:44.727Z" }, +] + +[[package]] +name = "pyproj" +version = "3.7.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +dependencies = [ + { name = "certifi", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/90/67bd7260b4ea9b8b20b4f58afef6c223ecb3abf368eb4ec5bc2cdef81b49/pyproj-3.7.2.tar.gz", hash = "sha256:39a0cf1ecc7e282d1d30f36594ebd55c9fae1fda8a2622cee5d100430628f88c", size = 226279, upload-time = "2025-08-14T12:05:42.18Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/bd/f205552cd1713b08f93b09e39a3ec99edef0b3ebbbca67b486fdf1abe2de/pyproj-3.7.2-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:2514d61f24c4e0bb9913e2c51487ecdaeca5f8748d8313c933693416ca41d4d5", size = 6227022, upload-time = "2025-08-14T12:03:51.474Z" }, + { url = "https://files.pythonhosted.org/packages/75/4c/9a937e659b8b418ab573c6d340d27e68716928953273e0837e7922fcac34/pyproj-3.7.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:8693ca3892d82e70de077701ee76dd13d7bca4ae1c9d1e739d72004df015923a", size = 4625810, upload-time = "2025-08-14T12:03:53.808Z" }, + { url = "https://files.pythonhosted.org/packages/c0/7d/a9f41e814dc4d1dc54e95b2ccaf0b3ebe3eb18b1740df05fe334724c3d89/pyproj-3.7.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:5e26484d80fea56273ed1555abaea161e9661d81a6c07815d54b8e883d4ceb25", size = 9638694, upload-time = "2025-08-14T12:03:55.669Z" }, + { url = "https://files.pythonhosted.org/packages/ad/ab/9bdb4a6216b712a1f9aab1c0fcbee5d3726f34a366f29c3e8c08a78d6b70/pyproj-3.7.2-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:281cb92847814e8018010c48b4069ff858a30236638631c1a91dd7bfa68f8a8a", size = 9493977, upload-time = "2025-08-14T12:03:57.937Z" }, + { url = "https://files.pythonhosted.org/packages/c9/db/2db75b1b6190f1137b1c4e8ef6a22e1c338e46320f6329bfac819143e063/pyproj-3.7.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9c8577f0b7bb09118ec2e57e3babdc977127dd66326d6c5d755c76b063e6d9dc", size = 10841151, upload-time = "2025-08-14T12:04:00.271Z" }, + { url = "https://files.pythonhosted.org/packages/89/f7/989643394ba23a286e9b7b3f09981496172f9e0d4512457ffea7dc47ffc7/pyproj-3.7.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a23f59904fac3a5e7364b3aa44d288234af267ca041adb2c2b14a903cd5d3ac5", size = 10751585, upload-time = "2025-08-14T12:04:02.228Z" }, + { url = "https://files.pythonhosted.org/packages/53/6d/ad928fe975a6c14a093c92e6a319ca18f479f3336bb353a740bdba335681/pyproj-3.7.2-cp311-cp311-win32.whl", hash = "sha256:f2af4ed34b2cf3e031a2d85b067a3ecbd38df073c567e04b52fa7a0202afde8a", size = 5908533, upload-time = "2025-08-14T12:04:04.821Z" }, + { url = "https://files.pythonhosted.org/packages/79/e0/b95584605cec9ed50b7ebaf7975d1c4ddeec5a86b7a20554ed8b60042bd7/pyproj-3.7.2-cp311-cp311-win_amd64.whl", hash = "sha256:0b7cb633565129677b2a183c4d807c727d1c736fcb0568a12299383056e67433", size = 6320742, upload-time = "2025-08-14T12:04:06.357Z" }, + { url = "https://files.pythonhosted.org/packages/b7/4d/536e8f93bca808175c2d0a5ac9fdf69b960d8ab6b14f25030dccb07464d7/pyproj-3.7.2-cp311-cp311-win_arm64.whl", hash = "sha256:38b08d85e3a38e455625b80e9eb9f78027c8e2649a21dec4df1f9c3525460c71", size = 6245772, upload-time = "2025-08-14T12:04:08.365Z" }, + { url = "https://files.pythonhosted.org/packages/8d/ab/9893ea9fb066be70ed9074ae543914a618c131ed8dff2da1e08b3a4df4db/pyproj-3.7.2-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:0a9bb26a6356fb5b033433a6d1b4542158fb71e3c51de49b4c318a1dff3aeaab", size = 6219832, upload-time = "2025-08-14T12:04:10.264Z" }, + { url = "https://files.pythonhosted.org/packages/53/78/4c64199146eed7184eb0e85bedec60a4aa8853b6ffe1ab1f3a8b962e70a0/pyproj-3.7.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:567caa03021178861fad27fabde87500ec6d2ee173dd32f3e2d9871e40eebd68", size = 4620650, upload-time = "2025-08-14T12:04:11.978Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ac/14a78d17943898a93ef4f8c6a9d4169911c994e3161e54a7cedeba9d8dde/pyproj-3.7.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:c203101d1dc3c038a56cff0447acc515dd29d6e14811406ac539c21eed422b2a", size = 9667087, upload-time = "2025-08-14T12:04:13.964Z" }, + { url = "https://files.pythonhosted.org/packages/b8/be/212882c450bba74fc8d7d35cbd57e4af84792f0a56194819d98106b075af/pyproj-3.7.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1edc34266c0c23ced85f95a1ee8b47c9035eae6aca5b6b340327250e8e281630", size = 9552797, upload-time = "2025-08-14T12:04:16.624Z" }, + { url = "https://files.pythonhosted.org/packages/ba/c0/c0f25c87b5d2a8686341c53c1792a222a480d6c9caf60311fec12c99ec26/pyproj-3.7.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:aa9f26c21bc0e2dc3d224cb1eb4020cf23e76af179a7c66fea49b828611e4260", size = 10837036, upload-time = "2025-08-14T12:04:18.733Z" }, + { url = "https://files.pythonhosted.org/packages/5d/37/5cbd6772addde2090c91113332623a86e8c7d583eccb2ad02ea634c4a89f/pyproj-3.7.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f9428b318530625cb389b9ddc9c51251e172808a4af79b82809376daaeabe5e9", size = 10775952, upload-time = "2025-08-14T12:04:20.709Z" }, + { url = "https://files.pythonhosted.org/packages/69/a1/dc250e3cf83eb4b3b9a2cf86fdb5e25288bd40037ae449695550f9e96b2f/pyproj-3.7.2-cp312-cp312-win32.whl", hash = "sha256:b3d99ed57d319da042f175f4554fc7038aa4bcecc4ac89e217e350346b742c9d", size = 5898872, upload-time = "2025-08-14T12:04:22.485Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a6/6fe724b72b70f2b00152d77282e14964d60ab092ec225e67c196c9b463e5/pyproj-3.7.2-cp312-cp312-win_amd64.whl", hash = "sha256:11614a054cd86a2ed968a657d00987a86eeb91fdcbd9ad3310478685dc14a128", size = 6312176, upload-time = "2025-08-14T12:04:24.736Z" }, + { url = "https://files.pythonhosted.org/packages/5d/68/915cc32c02a91e76d02c8f55d5a138d6ef9e47a0d96d259df98f4842e558/pyproj-3.7.2-cp312-cp312-win_arm64.whl", hash = "sha256:509a146d1398bafe4f53273398c3bb0b4732535065fa995270e52a9d3676bca3", size = 6233452, upload-time = "2025-08-14T12:04:27.287Z" }, + { url = "https://files.pythonhosted.org/packages/be/14/faf1b90d267cea68d7e70662e7f88cefdb1bc890bd596c74b959e0517a72/pyproj-3.7.2-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:19466e529b1b15eeefdf8ff26b06fa745856c044f2f77bf0edbae94078c1dfa1", size = 6214580, upload-time = "2025-08-14T12:04:28.804Z" }, + { url = "https://files.pythonhosted.org/packages/35/48/da9a45b184d375f62667f62eba0ca68569b0bd980a0bb7ffcc1d50440520/pyproj-3.7.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c79b9b84c4a626c5dc324c0d666be0bfcebd99f7538d66e8898c2444221b3da7", size = 4615388, upload-time = "2025-08-14T12:04:30.553Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e7/d2b459a4a64bca328b712c1b544e109df88e5c800f7c143cfbc404d39bfb/pyproj-3.7.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:ceecf374cacca317bc09e165db38ac548ee3cad07c3609442bd70311c59c21aa", size = 9628455, upload-time = "2025-08-14T12:04:32.435Z" }, + { url = "https://files.pythonhosted.org/packages/f8/85/c2b1706e51942de19076eff082f8495e57d5151364e78b5bef4af4a1d94a/pyproj-3.7.2-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:5141a538ffdbe4bfd157421828bb2e07123a90a7a2d6f30fa1462abcfb5ce681", size = 9514269, upload-time = "2025-08-14T12:04:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/34/38/07a9b89ae7467872f9a476883a5bad9e4f4d1219d31060f0f2b282276cbe/pyproj-3.7.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f000841e98ea99acbb7b8ca168d67773b0191de95187228a16110245c5d954d5", size = 10808437, upload-time = "2025-08-14T12:04:36.485Z" }, + { url = "https://files.pythonhosted.org/packages/12/56/fda1daeabbd39dec5b07f67233d09f31facb762587b498e6fc4572be9837/pyproj-3.7.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8115faf2597f281a42ab608ceac346b4eb1383d3b45ab474fd37341c4bf82a67", size = 10745540, upload-time = "2025-08-14T12:04:38.568Z" }, + { url = "https://files.pythonhosted.org/packages/0d/90/c793182cbba65a39a11db2ac6b479fe76c59e6509ae75e5744c344a0da9d/pyproj-3.7.2-cp313-cp313-win32.whl", hash = "sha256:f18c0579dd6be00b970cb1a6719197fceecc407515bab37da0066f0184aafdf3", size = 5896506, upload-time = "2025-08-14T12:04:41.059Z" }, + { url = "https://files.pythonhosted.org/packages/be/0f/747974129cf0d800906f81cd25efd098c96509026e454d4b66868779ab04/pyproj-3.7.2-cp313-cp313-win_amd64.whl", hash = "sha256:bb41c29d5f60854b1075853fe80c58950b398d4ebb404eb532536ac8d2834ed7", size = 6310195, upload-time = "2025-08-14T12:04:42.974Z" }, + { url = "https://files.pythonhosted.org/packages/82/64/fc7598a53172c4931ec6edf5228280663063150625d3f6423b4c20f9daff/pyproj-3.7.2-cp313-cp313-win_arm64.whl", hash = "sha256:2b617d573be4118c11cd96b8891a0b7f65778fa7733ed8ecdb297a447d439100", size = 6230748, upload-time = "2025-08-14T12:04:44.491Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f0/611dd5cddb0d277f94b7af12981f56e1441bf8d22695065d4f0df5218498/pyproj-3.7.2-cp313-cp313t-macosx_13_0_x86_64.whl", hash = "sha256:d27b48f0e81beeaa2b4d60c516c3a1cfbb0c7ff6ef71256d8e9c07792f735279", size = 6241729, upload-time = "2025-08-14T12:04:46.274Z" }, + { url = "https://files.pythonhosted.org/packages/15/93/40bd4a6c523ff9965e480870611aed7eda5aa2c6128c6537345a2b77b542/pyproj-3.7.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:55a3610d75023c7b1c6e583e48ef8f62918e85a2ae81300569d9f104d6684bb6", size = 4652497, upload-time = "2025-08-14T12:04:48.203Z" }, + { url = "https://files.pythonhosted.org/packages/1b/ae/7150ead53c117880b35e0d37960d3138fe640a235feb9605cb9386f50bb0/pyproj-3.7.2-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:8d7349182fa622696787cc9e195508d2a41a64765da9b8a6bee846702b9e6220", size = 9942610, upload-time = "2025-08-14T12:04:49.652Z" }, + { url = "https://files.pythonhosted.org/packages/d8/17/7a4a7eafecf2b46ab64e5c08176c20ceb5844b503eaa551bf12ccac77322/pyproj-3.7.2-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:d230b186eb876ed4f29a7c5ee310144c3a0e44e89e55f65fb3607e13f6db337c", size = 9692390, upload-time = "2025-08-14T12:04:51.731Z" }, + { url = "https://files.pythonhosted.org/packages/c3/55/ae18f040f6410f0ea547a21ada7ef3e26e6c82befa125b303b02759c0e9d/pyproj-3.7.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:237499c7862c578d0369e2b8ac56eec550e391a025ff70e2af8417139dabb41c", size = 11047596, upload-time = "2025-08-14T12:04:53.748Z" }, + { url = "https://files.pythonhosted.org/packages/e6/2e/d3fff4d2909473f26ae799f9dda04caa322c417a51ff3b25763f7d03b233/pyproj-3.7.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8c225f5978abd506fd9a78eaaf794435e823c9156091cabaab5374efb29d7f69", size = 10896975, upload-time = "2025-08-14T12:04:55.875Z" }, + { url = "https://files.pythonhosted.org/packages/f2/bc/8fc7d3963d87057b7b51ebe68c1e7c51c23129eee5072ba6b86558544a46/pyproj-3.7.2-cp313-cp313t-win32.whl", hash = "sha256:2da731876d27639ff9d2d81c151f6ab90a1546455fabd93368e753047be344a2", size = 5953057, upload-time = "2025-08-14T12:04:58.466Z" }, + { url = "https://files.pythonhosted.org/packages/cc/27/ea9809966cc47d2d51e6d5ae631ea895f7c7c7b9b3c29718f900a8f7d197/pyproj-3.7.2-cp313-cp313t-win_amd64.whl", hash = "sha256:f54d91ae18dd23b6c0ab48126d446820e725419da10617d86a1b69ada6d881d3", size = 6375414, upload-time = "2025-08-14T12:04:59.861Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f8/1ef0129fba9a555c658e22af68989f35e7ba7b9136f25758809efec0cd6e/pyproj-3.7.2-cp313-cp313t-win_arm64.whl", hash = "sha256:fc52ba896cfc3214dc9f9ca3c0677a623e8fdd096b257c14a31e719d21ff3fdd", size = 6262501, upload-time = "2025-08-14T12:05:01.39Z" }, + { url = "https://files.pythonhosted.org/packages/42/17/c2b050d3f5b71b6edd0d96ae16c990fdc42a5f1366464a5c2772146de33a/pyproj-3.7.2-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:2aaa328605ace41db050d06bac1adc11f01b71fe95c18661497763116c3a0f02", size = 6214541, upload-time = "2025-08-14T12:05:03.166Z" }, + { url = "https://files.pythonhosted.org/packages/03/68/68ada9c8aea96ded09a66cfd9bf87aa6db8c2edebe93f5bf9b66b0143fbc/pyproj-3.7.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:35dccbce8201313c596a970fde90e33605248b66272595c061b511c8100ccc08", size = 4617456, upload-time = "2025-08-14T12:05:04.563Z" }, + { url = "https://files.pythonhosted.org/packages/81/e4/4c50ceca7d0e937977866b02cb64e6ccf4df979a5871e521f9e255df6073/pyproj-3.7.2-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:25b0b7cb0042444c29a164b993c45c1b8013d6c48baa61dc1160d834a277e83b", size = 9615590, upload-time = "2025-08-14T12:05:06.094Z" }, + { url = "https://files.pythonhosted.org/packages/05/1e/ada6fb15a1d75b5bd9b554355a69a798c55a7dcc93b8d41596265c1772e3/pyproj-3.7.2-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:85def3a6388e9ba51f964619aa002a9d2098e77c6454ff47773bb68871024281", size = 9474960, upload-time = "2025-08-14T12:05:07.973Z" }, + { url = "https://files.pythonhosted.org/packages/51/07/9d48ad0a8db36e16f842f2c8a694c1d9d7dcf9137264846bef77585a71f3/pyproj-3.7.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b1bccefec3875ab81eabf49059e2b2ea77362c178b66fd3528c3e4df242f1516", size = 10799478, upload-time = "2025-08-14T12:05:14.102Z" }, + { url = "https://files.pythonhosted.org/packages/85/cf/2f812b529079f72f51ff2d6456b7fef06c01735e5cfd62d54ffb2b548028/pyproj-3.7.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d5371ca114d6990b675247355a801925814eca53e6c4b2f1b5c0a956336ee36e", size = 10710030, upload-time = "2025-08-14T12:05:16.317Z" }, + { url = "https://files.pythonhosted.org/packages/99/9b/4626a19e1f03eba4c0e77b91a6cf0f73aa9cb5d51a22ee385c22812bcc2c/pyproj-3.7.2-cp314-cp314-win32.whl", hash = "sha256:77f066626030f41be543274f5ac79f2a511fe89860ecd0914f22131b40a0ec25", size = 5991181, upload-time = "2025-08-14T12:05:19.492Z" }, + { url = "https://files.pythonhosted.org/packages/04/b2/5a6610554306a83a563080c2cf2c57565563eadd280e15388efa00fb5b33/pyproj-3.7.2-cp314-cp314-win_amd64.whl", hash = "sha256:5a964da1696b8522806f4276ab04ccfff8f9eb95133a92a25900697609d40112", size = 6434721, upload-time = "2025-08-14T12:05:21.022Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ce/6c910ea2e1c74ef673c5d48c482564b8a7824a44c4e35cca2e765b68cfcc/pyproj-3.7.2-cp314-cp314-win_arm64.whl", hash = "sha256:e258ab4dbd3cf627809067c0ba8f9884ea76c8e5999d039fb37a1619c6c3e1f6", size = 6363821, upload-time = "2025-08-14T12:05:22.627Z" }, + { url = "https://files.pythonhosted.org/packages/e4/e4/5532f6f7491812ba782a2177fe9de73fd8e2912b59f46a1d056b84b9b8f2/pyproj-3.7.2-cp314-cp314t-macosx_13_0_x86_64.whl", hash = "sha256:bbbac2f930c6d266f70ec75df35ef851d96fdb3701c674f42fd23a9314573b37", size = 6241773, upload-time = "2025-08-14T12:05:24.577Z" }, + { url = "https://files.pythonhosted.org/packages/20/1f/0938c3f2bbbef1789132d1726d9b0e662f10cfc22522743937f421ad664e/pyproj-3.7.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:b7544e0a3d6339dc9151e9c8f3ea62a936ab7cc446a806ec448bbe86aebb979b", size = 4652537, upload-time = "2025-08-14T12:05:26.391Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a8/488b1ed47d25972f33874f91f09ca8f2227902f05f63a2b80dc73e7b1c97/pyproj-3.7.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:f7f5133dca4c703e8acadf6f30bc567d39a42c6af321e7f81975c2518f3ed357", size = 9940864, upload-time = "2025-08-14T12:05:27.985Z" }, + { url = "https://files.pythonhosted.org/packages/c7/cc/7f4c895d0cb98e47b6a85a6d79eaca03eb266129eed2f845125c09cf31ff/pyproj-3.7.2-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:5aff3343038d7426aa5076f07feb88065f50e0502d1b0d7c22ddfdd2c75a3f81", size = 9688868, upload-time = "2025-08-14T12:05:30.425Z" }, + { url = "https://files.pythonhosted.org/packages/b2/b7/c7e306b8bb0f071d9825b753ee4920f066c40fbfcce9372c4f3cfb2fc4ed/pyproj-3.7.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:b0552178c61f2ac1c820d087e8ba6e62b29442debddbb09d51c4bf8acc84d888", size = 11045910, upload-time = "2025-08-14T12:05:32.507Z" }, + { url = "https://files.pythonhosted.org/packages/42/fb/538a4d2df695980e2dde5c04d965fbdd1fe8c20a3194dc4aaa3952a4d1be/pyproj-3.7.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:47d87db2d2c436c5fd0409b34d70bb6cdb875cca2ebe7a9d1c442367b0ab8d59", size = 10895724, upload-time = "2025-08-14T12:05:35.465Z" }, + { url = "https://files.pythonhosted.org/packages/e8/8b/a3f0618b03957de9db5489a04558a8826f43906628bb0b766033aa3b5548/pyproj-3.7.2-cp314-cp314t-win32.whl", hash = "sha256:c9b6f1d8ad3e80a0ee0903a778b6ece7dca1d1d40f6d114ae01bc8ddbad971aa", size = 6056848, upload-time = "2025-08-14T12:05:37.553Z" }, + { url = "https://files.pythonhosted.org/packages/bc/56/413240dd5149dd3291eda55aa55a659da4431244a2fd1319d0ae89407cfb/pyproj-3.7.2-cp314-cp314t-win_amd64.whl", hash = "sha256:1914e29e27933ba6f9822663ee0600f169014a2859f851c054c88cf5ea8a333c", size = 6517676, upload-time = "2025-08-14T12:05:39.126Z" }, + { url = "https://files.pythonhosted.org/packages/15/73/a7141a1a0559bf1a7aa42a11c879ceb19f02f5c6c371c6d57fd86cefd4d1/pyproj-3.7.2-cp314-cp314t-win_arm64.whl", hash = "sha256:d9d25bae416a24397e0d85739f84d323b55f6511e45a522dd7d7eae70d10c7e4", size = 6391844, upload-time = "2025-08-14T12:05:40.745Z" }, ] [[package]] @@ -1693,6 +2146,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-discovery" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "platformdirs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9c/90/bcce6b46823c9bec1757c964dc37ed332579be512e17a30e9698095dcae4/python_discovery-1.2.0.tar.gz", hash = "sha256:7d33e350704818b09e3da2bd419d37e21e7c30db6e0977bb438916e06b41b5b1", size = 58055, upload-time = "2026-03-19T01:43:08.248Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/3c/2005227cb951df502412de2fa781f800663cccbef8d90ec6f1b371ac2c0d/python_discovery-1.2.0-py3-none-any.whl", hash = "sha256:1e108f1bbe2ed0ef089823d28805d5ad32be8e734b86a5f212bf89b71c266e4a", size = 31524, upload-time = "2026-03-19T01:43:07.045Z" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -1746,6 +2212,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446, upload-time = "2024-08-06T20:33:04.33Z" }, ] +[[package]] +name = "pyyaml-env-tag" +version = "1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/2e/79c822141bfd05a853236b504869ebc6b70159afc570e1d5a20641782eaa/pyyaml_env_tag-1.1.tar.gz", hash = "sha256:2eb38b75a2d21ee0475d6d97ec19c63287a7e140231e4214969d0eac923cd7ff", size = 5737, upload-time = "2025-05-13T15:24:01.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/11/432f32f8097b03e3cd5fe57e88efb685d964e2e5178a48ed61e841f7fdce/pyyaml_env_tag-1.1-py3-none-any.whl", hash = "sha256:17109e1a528561e32f026364712fee1264bc2ea6715120891174ed1b980d2e04", size = 4722, upload-time = "2025-05-13T15:23:59.629Z" }, +] + [[package]] name = "requests" version = "2.32.4" @@ -1803,6 +2281,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696, upload-time = "2025-04-16T09:51:17.142Z" }, ] +[[package]] +name = "ruff" +version = "0.15.17" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/a9/3abdf488f1bf3d24c699415e454ed554a6350d5d89ce183be1ee0a3361ac/ruff-0.15.17.tar.gz", hash = "sha256:2ec446937fd16c8c4de2674a209cc5af64d9c6f17d21fbf1151054fa0bcf5219", size = 4743346, upload-time = "2026-06-11T17:54:47.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/4d/e11259f5da07cb6afb2d074c31bf09da9671993f7329d4f15d2fdc458301/ruff-0.15.17-py3-none-linux_armv6l.whl", hash = "sha256:d9feddb927fc68bd295f5eebc587a7e42cfaf9b65f60ca4a2386febff575da8f", size = 10856677, upload-time = "2026-06-11T17:54:49.533Z" }, + { url = "https://files.pythonhosted.org/packages/29/3e/772d679e1a0dc058e58875bd2c0cb713a0530877b4a76fee3c7966df0d49/ruff-0.15.17-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:25805a226d741c47d274a35ad5c10a7dde175fcddfa511d7cf3da0a21eb3eab7", size = 11223443, upload-time = "2026-06-11T17:55:00.573Z" }, + { url = "https://files.pythonhosted.org/packages/68/58/bd41f7688b2fd5623012605130ed70e60aa7f2244baa3d5066bdd61530c8/ruff-0.15.17-py3-none-macosx_11_0_arm64.whl", hash = "sha256:f6ad73b14c2d18a3bf8ad7cb6974294d7f613a7898604826058e6ac64918ef4d", size = 10566458, upload-time = "2026-06-11T17:55:07.52Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5b/733371013fcf1ec339e477ece6ab42bfe10bdd9bba8ee88a9516aa56bfc0/ruff-0.15.17-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ba0c1e4f95bcb3869d0d30cbd5917071ef2e28665abfec970cdab0492c713ed", size = 10914483, upload-time = "2026-06-11T17:55:05.501Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cc/6f24251cc0252f7239391ccb85833f320efad14ebe5b443943f37ced6332/ruff-0.15.17-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:81647960f10bff57d2e51cadd0c3950fe598400c852863a038720ef5b8cca91e", size = 10647497, upload-time = "2026-06-11T17:54:57.733Z" }, + { url = "https://files.pythonhosted.org/packages/68/dd/0d10c17ce1a1624d6fc3156309c3f834fdb5dfaad026ec90c85684f3990e/ruff-0.15.17-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e01a84ddbc8c16c23055ba3924476850f1bbc1917cebbb9376665a63e74260d", size = 11416967, upload-time = "2026-06-11T17:54:51.461Z" }, + { url = "https://files.pythonhosted.org/packages/2f/91/556bfb156f6144f355e831c23db00b2fc4120f86b3ce81cc5f7fd2df51f3/ruff-0.15.17-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:84fe9f653152f8f294f9f7e03bf3a453d8b4a27f7a59c78c8666167f2b17b96c", size = 12335770, upload-time = "2026-06-11T17:54:45.793Z" }, + { url = "https://files.pythonhosted.org/packages/88/82/8b5999aa13355e926f06d9f42a32dcca862f623bf0363785ff89d607dffd/ruff-0.15.17-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c0fe88a7676e7a05b73174d4d4a59cb2ac21ff8263583f87a81a6018475a978", size = 11575441, upload-time = "2026-06-11T17:54:32.661Z" }, + { url = "https://files.pythonhosted.org/packages/11/93/f10377bb04109ca0e8cbc483ff1982c54b6d418210041776f93e8cdc7fa9/ruff-0.15.17-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ecfc3c7878fff94633ab0348524e093f9ce3243080416dd7d14f8ba400174719", size = 11557614, upload-time = "2026-06-11T17:54:34.698Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a6/eeeae7f7d5493df41649ab3db92f086b2d0a30199e4efdf8e3dd7a033f24/ruff-0.15.17-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:b8461180b22420b1bdc289909410930761629fddf2a5aaf60fae1ab26cedc4c4", size = 11544450, upload-time = "2026-06-11T17:54:39.042Z" }, + { url = "https://files.pythonhosted.org/packages/32/88/5991ce565129a24dd4a00db1254b3b5db2e53018cbe4018ea5a89738e727/ruff-0.15.17-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6eccbe50a038b503e7140b441aa9c7fc8c1f36edf23ebef9f4165c2f28f568b7", size = 10892524, upload-time = "2026-06-11T17:55:09.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/1d/0fdd248313425f55223968af04b0a42125466a8d88d21c1d99c6af0a51e8/ruff-0.15.17-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:382fc0521025f5a8ad447d8bdd523545d0d7646adb718eb1c2dac5065ec27c0f", size = 10659573, upload-time = "2026-06-11T17:54:36.824Z" }, + { url = "https://files.pythonhosted.org/packages/9e/0e/072e8260deb9461062ce9311ced27a8e541229a6ffd483013dd37661e43e/ruff-0.15.17-py3-none-musllinux_1_2_i686.whl", hash = "sha256:456d41fcd1b2777ad63f09a6e7121d43f7b688bbc76a800c10f7f8fb1f912c3f", size = 11127818, upload-time = "2026-06-11T17:55:03.124Z" }, + { url = "https://files.pythonhosted.org/packages/ab/b4/55060a34163121498014696b5f656db5b8c6963768f227dbf0d76b311073/ruff-0.15.17-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b1a04bcc94ae6194e9db05d16ad31f298a7194bfbcb08258bbe589cee1d587b8", size = 11655901, upload-time = "2026-06-11T17:54:53.562Z" }, + { url = "https://files.pythonhosted.org/packages/49/71/9b29d6b87cef468d697f43c6a91e3fae4a80185779d7d5a4ef27d173439f/ruff-0.15.17-py3-none-win32.whl", hash = "sha256:596065960ab1ff593f744220c9fe6580eda00a95003cffa9f4048bb5b1bf0392", size = 10925574, upload-time = "2026-06-11T17:54:55.723Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b2/8fc77f3723228836fa5d12497eb71c808f83782e10d058d2b15cfa14640b/ruff-0.15.17-py3-none-win_amd64.whl", hash = "sha256:6769e5fa1710b179b92e0bfa5a51735b35baea9013dadb06d5f44cbcf9547084", size = 12058788, upload-time = "2026-06-11T17:54:41.042Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c7/c53e8dbff9c9dc4b7928773421ae294a5d28fcb8dcda1a089579d3a7e510/ruff-0.15.17-py3-none-win_arm64.whl", hash = "sha256:f3be1fbb34bcdfd146240d8fb92a709d4c2c8191348580a3c044ec60fa0b4456", size = 11355275, upload-time = "2026-06-11T17:54:43.635Z" }, +] + [[package]] name = "scipy" version = "1.15.3" @@ -2020,6 +2523,157 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, ] +[[package]] +name = "virtualenv" +version = "21.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "distlib" }, + { name = "filelock" }, + { name = "platformdirs" }, + { name = "python-discovery" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/aa/92/58199fe10049f9703c2666e809c4f686c54ef0a68b0f6afccf518c0b1eb9/virtualenv-21.2.0.tar.gz", hash = "sha256:1720dc3a62ef5b443092e3f499228599045d7fea4c79199770499df8becf9098", size = 5840618, upload-time = "2026-03-09T17:24:38.013Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/59/7d02447a55b2e55755011a647479041bc92a82e143f96a8195cb33bd0a1c/virtualenv-21.2.0-py3-none-any.whl", hash = "sha256:1bd755b504931164a5a496d217c014d098426cddc79363ad66ac78125f9d908f", size = 5825084, upload-time = "2026-03-09T17:24:35.378Z" }, +] + +[[package]] +name = "watchdog" +version = "6.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/db/7d/7f3d619e951c88ed75c6037b246ddcf2d322812ee8ea189be89511721d54/watchdog-6.0.0.tar.gz", hash = "sha256:9ddf7c82fda3ae8e24decda1338ede66e1c99883db93711d8fb941eaa2d8c282", size = 131220, upload-time = "2024-11-01T14:07:13.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/56/90994d789c61df619bfc5ce2ecdabd5eeff564e1eb47512bd01b5e019569/watchdog-6.0.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d1cdb490583ebd691c012b3d6dae011000fe42edb7a82ece80965b42abd61f26", size = 96390, upload-time = "2024-11-01T14:06:24.793Z" }, + { url = "https://files.pythonhosted.org/packages/55/46/9a67ee697342ddf3c6daa97e3a587a56d6c4052f881ed926a849fcf7371c/watchdog-6.0.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:bc64ab3bdb6a04d69d4023b29422170b74681784ffb9463ed4870cf2f3e66112", size = 88389, upload-time = "2024-11-01T14:06:27.112Z" }, + { url = "https://files.pythonhosted.org/packages/44/65/91b0985747c52064d8701e1075eb96f8c40a79df889e59a399453adfb882/watchdog-6.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c897ac1b55c5a1461e16dae288d22bb2e412ba9807df8397a635d88f671d36c3", size = 89020, upload-time = "2024-11-01T14:06:29.876Z" }, + { url = "https://files.pythonhosted.org/packages/e0/24/d9be5cd6642a6aa68352ded4b4b10fb0d7889cb7f45814fb92cecd35f101/watchdog-6.0.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6eb11feb5a0d452ee41f824e271ca311a09e250441c262ca2fd7ebcf2461a06c", size = 96393, upload-time = "2024-11-01T14:06:31.756Z" }, + { url = "https://files.pythonhosted.org/packages/63/7a/6013b0d8dbc56adca7fdd4f0beed381c59f6752341b12fa0886fa7afc78b/watchdog-6.0.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ef810fbf7b781a5a593894e4f439773830bdecb885e6880d957d5b9382a960d2", size = 88392, upload-time = "2024-11-01T14:06:32.99Z" }, + { url = "https://files.pythonhosted.org/packages/d1/40/b75381494851556de56281e053700e46bff5b37bf4c7267e858640af5a7f/watchdog-6.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:afd0fe1b2270917c5e23c2a65ce50c2a4abb63daafb0d419fde368e272a76b7c", size = 89019, upload-time = "2024-11-01T14:06:34.963Z" }, + { url = "https://files.pythonhosted.org/packages/39/ea/3930d07dafc9e286ed356a679aa02d777c06e9bfd1164fa7c19c288a5483/watchdog-6.0.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:bdd4e6f14b8b18c334febb9c4425a878a2ac20efd1e0b231978e7b150f92a948", size = 96471, upload-time = "2024-11-01T14:06:37.745Z" }, + { url = "https://files.pythonhosted.org/packages/12/87/48361531f70b1f87928b045df868a9fd4e253d9ae087fa4cf3f7113be363/watchdog-6.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c7c15dda13c4eb00d6fb6fc508b3c0ed88b9d5d374056b239c4ad1611125c860", size = 88449, upload-time = "2024-11-01T14:06:39.748Z" }, + { url = "https://files.pythonhosted.org/packages/5b/7e/8f322f5e600812e6f9a31b75d242631068ca8f4ef0582dd3ae6e72daecc8/watchdog-6.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6f10cb2d5902447c7d0da897e2c6768bca89174d0c6e1e30abec5421af97a5b0", size = 89054, upload-time = "2024-11-01T14:06:41.009Z" }, + { url = "https://files.pythonhosted.org/packages/68/98/b0345cabdce2041a01293ba483333582891a3bd5769b08eceb0d406056ef/watchdog-6.0.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:490ab2ef84f11129844c23fb14ecf30ef3d8a6abafd3754a6f75ca1e6654136c", size = 96480, upload-time = "2024-11-01T14:06:42.952Z" }, + { url = "https://files.pythonhosted.org/packages/85/83/cdf13902c626b28eedef7ec4f10745c52aad8a8fe7eb04ed7b1f111ca20e/watchdog-6.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:76aae96b00ae814b181bb25b1b98076d5fc84e8a53cd8885a318b42b6d3a5134", size = 88451, upload-time = "2024-11-01T14:06:45.084Z" }, + { url = "https://files.pythonhosted.org/packages/fe/c4/225c87bae08c8b9ec99030cd48ae9c4eca050a59bf5c2255853e18c87b50/watchdog-6.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a175f755fc2279e0b7312c0035d52e27211a5bc39719dd529625b1930917345b", size = 89057, upload-time = "2024-11-01T14:06:47.324Z" }, + { url = "https://files.pythonhosted.org/packages/30/ad/d17b5d42e28a8b91f8ed01cb949da092827afb9995d4559fd448d0472763/watchdog-6.0.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:c7ac31a19f4545dd92fc25d200694098f42c9a8e391bc00bdd362c5736dbf881", size = 87902, upload-time = "2024-11-01T14:06:53.119Z" }, + { url = "https://files.pythonhosted.org/packages/5c/ca/c3649991d140ff6ab67bfc85ab42b165ead119c9e12211e08089d763ece5/watchdog-6.0.0-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:9513f27a1a582d9808cf21a07dae516f0fab1cf2d7683a742c498b93eedabb11", size = 88380, upload-time = "2024-11-01T14:06:55.19Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c7/ca4bf3e518cb57a686b2feb4f55a1892fd9a3dd13f470fca14e00f80ea36/watchdog-6.0.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7607498efa04a3542ae3e05e64da8202e58159aa1fa4acddf7678d34a35d4f13", size = 79079, upload-time = "2024-11-01T14:06:59.472Z" }, + { url = "https://files.pythonhosted.org/packages/5c/51/d46dc9332f9a647593c947b4b88e2381c8dfc0942d15b8edc0310fa4abb1/watchdog-6.0.0-py3-none-manylinux2014_armv7l.whl", hash = "sha256:9041567ee8953024c83343288ccc458fd0a2d811d6a0fd68c4c22609e3490379", size = 79078, upload-time = "2024-11-01T14:07:01.431Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/04edbf5e169cd318d5f07b4766fee38e825d64b6913ca157ca32d1a42267/watchdog-6.0.0-py3-none-manylinux2014_i686.whl", hash = "sha256:82dc3e3143c7e38ec49d61af98d6558288c415eac98486a5c581726e0737c00e", size = 79076, upload-time = "2024-11-01T14:07:02.568Z" }, + { url = "https://files.pythonhosted.org/packages/ab/cc/da8422b300e13cb187d2203f20b9253e91058aaf7db65b74142013478e66/watchdog-6.0.0-py3-none-manylinux2014_ppc64.whl", hash = "sha256:212ac9b8bf1161dc91bd09c048048a95ca3a4c4f5e5d4a7d1b1a7d5752a7f96f", size = 79077, upload-time = "2024-11-01T14:07:03.893Z" }, + { url = "https://files.pythonhosted.org/packages/2c/3b/b8964e04ae1a025c44ba8e4291f86e97fac443bca31de8bd98d3263d2fcf/watchdog-6.0.0-py3-none-manylinux2014_ppc64le.whl", hash = "sha256:e3df4cbb9a450c6d49318f6d14f4bbc80d763fa587ba46ec86f99f9e6876bb26", size = 79078, upload-time = "2024-11-01T14:07:05.189Z" }, + { url = "https://files.pythonhosted.org/packages/62/ae/a696eb424bedff7407801c257d4b1afda455fe40821a2be430e173660e81/watchdog-6.0.0-py3-none-manylinux2014_s390x.whl", hash = "sha256:2cce7cfc2008eb51feb6aab51251fd79b85d9894e98ba847408f662b3395ca3c", size = 79077, upload-time = "2024-11-01T14:07:06.376Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e8/dbf020b4d98251a9860752a094d09a65e1b436ad181faf929983f697048f/watchdog-6.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:20ffe5b202af80ab4266dcd3e91aae72bf2da48c0d33bdb15c66658e685e94e2", size = 79078, upload-time = "2024-11-01T14:07:07.547Z" }, + { url = "https://files.pythonhosted.org/packages/07/f6/d0e5b343768e8bcb4cda79f0f2f55051bf26177ecd5651f84c07567461cf/watchdog-6.0.0-py3-none-win32.whl", hash = "sha256:07df1fdd701c5d4c8e55ef6cf55b8f0120fe1aef7ef39a1c6fc6bc2e606d517a", size = 79065, upload-time = "2024-11-01T14:07:09.525Z" }, + { url = "https://files.pythonhosted.org/packages/db/d9/c495884c6e548fce18a8f40568ff120bc3a4b7b99813081c8ac0c936fa64/watchdog-6.0.0-py3-none-win_amd64.whl", hash = "sha256:cbafb470cf848d93b5d013e2ecb245d4aa1c8fd0504e863ccefa32445359d680", size = 79070, upload-time = "2024-11-01T14:07:10.686Z" }, + { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067, upload-time = "2024-11-01T14:07:11.845Z" }, +] + +[[package]] +name = "watchfiles" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c2/c9/8869df9b2a2d6c59d79220a4db37679e74f807c559ffe5265e08b227a210/watchfiles-1.1.1.tar.gz", hash = "sha256:a173cb5c16c4f40ab19cecf48a534c409f7ea983ab8fed0741304a1c0a31b3f2", size = 94440, upload-time = "2025-10-14T15:06:21.08Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/1a/206e8cf2dd86fddf939165a57b4df61607a1e0add2785f170a3f616b7d9f/watchfiles-1.1.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:eef58232d32daf2ac67f42dea51a2c80f0d03379075d44a587051e63cc2e368c", size = 407318, upload-time = "2025-10-14T15:04:18.753Z" }, + { url = "https://files.pythonhosted.org/packages/b3/0f/abaf5262b9c496b5dad4ed3c0e799cbecb1f8ea512ecb6ddd46646a9fca3/watchfiles-1.1.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:03fa0f5237118a0c5e496185cafa92878568b652a2e9a9382a5151b1a0380a43", size = 394478, upload-time = "2025-10-14T15:04:20.297Z" }, + { url = "https://files.pythonhosted.org/packages/b1/04/9cc0ba88697b34b755371f5ace8d3a4d9a15719c07bdc7bd13d7d8c6a341/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8ca65483439f9c791897f7db49202301deb6e15fe9f8fe2fed555bf986d10c31", size = 449894, upload-time = "2025-10-14T15:04:21.527Z" }, + { url = "https://files.pythonhosted.org/packages/d2/9c/eda4615863cd8621e89aed4df680d8c3ec3da6a4cf1da113c17decd87c7f/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f0ab1c1af0cb38e3f598244c17919fb1a84d1629cc08355b0074b6d7f53138ac", size = 459065, upload-time = "2025-10-14T15:04:22.795Z" }, + { url = "https://files.pythonhosted.org/packages/84/13/f28b3f340157d03cbc8197629bc109d1098764abe1e60874622a0be5c112/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3bc570d6c01c206c46deb6e935a260be44f186a2f05179f52f7fcd2be086a94d", size = 488377, upload-time = "2025-10-14T15:04:24.138Z" }, + { url = "https://files.pythonhosted.org/packages/86/93/cfa597fa9389e122488f7ffdbd6db505b3b915ca7435ecd7542e855898c2/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e84087b432b6ac94778de547e08611266f1f8ffad28c0ee4c82e028b0fc5966d", size = 595837, upload-time = "2025-10-14T15:04:25.057Z" }, + { url = "https://files.pythonhosted.org/packages/57/1e/68c1ed5652b48d89fc24d6af905d88ee4f82fa8bc491e2666004e307ded1/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:620bae625f4cb18427b1bb1a2d9426dc0dd5a5ba74c7c2cdb9de405f7b129863", size = 473456, upload-time = "2025-10-14T15:04:26.497Z" }, + { url = "https://files.pythonhosted.org/packages/d5/dc/1a680b7458ffa3b14bb64878112aefc8f2e4f73c5af763cbf0bd43100658/watchfiles-1.1.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:544364b2b51a9b0c7000a4b4b02f90e9423d97fbbf7e06689236443ebcad81ab", size = 455614, upload-time = "2025-10-14T15:04:27.539Z" }, + { url = "https://files.pythonhosted.org/packages/61/a5/3d782a666512e01eaa6541a72ebac1d3aae191ff4a31274a66b8dd85760c/watchfiles-1.1.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:bbe1ef33d45bc71cf21364df962af171f96ecaeca06bd9e3d0b583efb12aec82", size = 630690, upload-time = "2025-10-14T15:04:28.495Z" }, + { url = "https://files.pythonhosted.org/packages/9b/73/bb5f38590e34687b2a9c47a244aa4dd50c56a825969c92c9c5fc7387cea1/watchfiles-1.1.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:1a0bb430adb19ef49389e1ad368450193a90038b5b752f4ac089ec6942c4dff4", size = 622459, upload-time = "2025-10-14T15:04:29.491Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ac/c9bb0ec696e07a20bd58af5399aeadaef195fb2c73d26baf55180fe4a942/watchfiles-1.1.1-cp310-cp310-win32.whl", hash = "sha256:3f6d37644155fb5beca5378feb8c1708d5783145f2a0f1c4d5a061a210254844", size = 272663, upload-time = "2025-10-14T15:04:30.435Z" }, + { url = "https://files.pythonhosted.org/packages/11/a0/a60c5a7c2ec59fa062d9a9c61d02e3b6abd94d32aac2d8344c4bdd033326/watchfiles-1.1.1-cp310-cp310-win_amd64.whl", hash = "sha256:a36d8efe0f290835fd0f33da35042a1bb5dc0e83cbc092dcf69bce442579e88e", size = 287453, upload-time = "2025-10-14T15:04:31.53Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f8/2c5f479fb531ce2f0564eda479faecf253d886b1ab3630a39b7bf7362d46/watchfiles-1.1.1-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:f57b396167a2565a4e8b5e56a5a1c537571733992b226f4f1197d79e94cf0ae5", size = 406529, upload-time = "2025-10-14T15:04:32.899Z" }, + { url = "https://files.pythonhosted.org/packages/fe/cd/f515660b1f32f65df671ddf6f85bfaca621aee177712874dc30a97397977/watchfiles-1.1.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:421e29339983e1bebc281fab40d812742268ad057db4aee8c4d2bce0af43b741", size = 394384, upload-time = "2025-10-14T15:04:33.761Z" }, + { url = "https://files.pythonhosted.org/packages/7b/c3/28b7dc99733eab43fca2d10f55c86e03bd6ab11ca31b802abac26b23d161/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6e43d39a741e972bab5d8100b5cdacf69db64e34eb19b6e9af162bccf63c5cc6", size = 448789, upload-time = "2025-10-14T15:04:34.679Z" }, + { url = "https://files.pythonhosted.org/packages/4a/24/33e71113b320030011c8e4316ccca04194bf0cbbaeee207f00cbc7d6b9f5/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f537afb3276d12814082a2e9b242bdcf416c2e8fd9f799a737990a1dbe906e5b", size = 460521, upload-time = "2025-10-14T15:04:35.963Z" }, + { url = "https://files.pythonhosted.org/packages/f4/c3/3c9a55f255aa57b91579ae9e98c88704955fa9dac3e5614fb378291155df/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b2cd9e04277e756a2e2d2543d65d1e2166d6fd4c9b183f8808634fda23f17b14", size = 488722, upload-time = "2025-10-14T15:04:37.091Z" }, + { url = "https://files.pythonhosted.org/packages/49/36/506447b73eb46c120169dc1717fe2eff07c234bb3232a7200b5f5bd816e9/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5f3f58818dc0b07f7d9aa7fe9eb1037aecb9700e63e1f6acfed13e9fef648f5d", size = 596088, upload-time = "2025-10-14T15:04:38.39Z" }, + { url = "https://files.pythonhosted.org/packages/82/ab/5f39e752a9838ec4d52e9b87c1e80f1ee3ccdbe92e183c15b6577ab9de16/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9bb9f66367023ae783551042d31b1d7fd422e8289eedd91f26754a66f44d5cff", size = 472923, upload-time = "2025-10-14T15:04:39.666Z" }, + { url = "https://files.pythonhosted.org/packages/af/b9/a419292f05e302dea372fa7e6fda5178a92998411f8581b9830d28fb9edb/watchfiles-1.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aebfd0861a83e6c3d1110b78ad54704486555246e542be3e2bb94195eabb2606", size = 456080, upload-time = "2025-10-14T15:04:40.643Z" }, + { url = "https://files.pythonhosted.org/packages/b0/c3/d5932fd62bde1a30c36e10c409dc5d54506726f08cb3e1d8d0ba5e2bc8db/watchfiles-1.1.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:5fac835b4ab3c6487b5dbad78c4b3724e26bcc468e886f8ba8cc4306f68f6701", size = 629432, upload-time = "2025-10-14T15:04:41.789Z" }, + { url = "https://files.pythonhosted.org/packages/f7/77/16bddd9779fafb795f1a94319dc965209c5641db5bf1edbbccace6d1b3c0/watchfiles-1.1.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:399600947b170270e80134ac854e21b3ccdefa11a9529a3decc1327088180f10", size = 623046, upload-time = "2025-10-14T15:04:42.718Z" }, + { url = "https://files.pythonhosted.org/packages/46/ef/f2ecb9a0f342b4bfad13a2787155c6ee7ce792140eac63a34676a2feeef2/watchfiles-1.1.1-cp311-cp311-win32.whl", hash = "sha256:de6da501c883f58ad50db3a32ad397b09ad29865b5f26f64c24d3e3281685849", size = 271473, upload-time = "2025-10-14T15:04:43.624Z" }, + { url = "https://files.pythonhosted.org/packages/94/bc/f42d71125f19731ea435c3948cad148d31a64fccde3867e5ba4edee901f9/watchfiles-1.1.1-cp311-cp311-win_amd64.whl", hash = "sha256:35c53bd62a0b885bf653ebf6b700d1bf05debb78ad9292cf2a942b23513dc4c4", size = 287598, upload-time = "2025-10-14T15:04:44.516Z" }, + { url = "https://files.pythonhosted.org/packages/57/c9/a30f897351f95bbbfb6abcadafbaca711ce1162f4db95fc908c98a9165f3/watchfiles-1.1.1-cp311-cp311-win_arm64.whl", hash = "sha256:57ca5281a8b5e27593cb7d82c2ac927ad88a96ed406aa446f6344e4328208e9e", size = 277210, upload-time = "2025-10-14T15:04:45.883Z" }, + { url = "https://files.pythonhosted.org/packages/74/d5/f039e7e3c639d9b1d09b07ea412a6806d38123f0508e5f9b48a87b0a76cc/watchfiles-1.1.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:8c89f9f2f740a6b7dcc753140dd5e1ab9215966f7a3530d0c0705c83b401bd7d", size = 404745, upload-time = "2025-10-14T15:04:46.731Z" }, + { url = "https://files.pythonhosted.org/packages/a5/96/a881a13aa1349827490dab2d363c8039527060cfcc2c92cc6d13d1b1049e/watchfiles-1.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:bd404be08018c37350f0d6e34676bd1e2889990117a2b90070b3007f172d0610", size = 391769, upload-time = "2025-10-14T15:04:48.003Z" }, + { url = "https://files.pythonhosted.org/packages/4b/5b/d3b460364aeb8da471c1989238ea0e56bec24b6042a68046adf3d9ddb01c/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8526e8f916bb5b9a0a777c8317c23ce65de259422bba5b31325a6fa6029d33af", size = 449374, upload-time = "2025-10-14T15:04:49.179Z" }, + { url = "https://files.pythonhosted.org/packages/b9/44/5769cb62d4ed055cb17417c0a109a92f007114a4e07f30812a73a4efdb11/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2edc3553362b1c38d9f06242416a5d8e9fe235c204a4072e988ce2e5bb1f69f6", size = 459485, upload-time = "2025-10-14T15:04:50.155Z" }, + { url = "https://files.pythonhosted.org/packages/19/0c/286b6301ded2eccd4ffd0041a1b726afda999926cf720aab63adb68a1e36/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:30f7da3fb3f2844259cba4720c3fc7138eb0f7b659c38f3bfa65084c7fc7abce", size = 488813, upload-time = "2025-10-14T15:04:51.059Z" }, + { url = "https://files.pythonhosted.org/packages/c7/2b/8530ed41112dd4a22f4dcfdb5ccf6a1baad1ff6eed8dc5a5f09e7e8c41c7/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f8979280bdafff686ba5e4d8f97840f929a87ed9cdf133cbbd42f7766774d2aa", size = 594816, upload-time = "2025-10-14T15:04:52.031Z" }, + { url = "https://files.pythonhosted.org/packages/ce/d2/f5f9fb49489f184f18470d4f99f4e862a4b3e9ac2865688eb2099e3d837a/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dcc5c24523771db3a294c77d94771abcfcb82a0e0ee8efd910c37c59ec1b31bb", size = 475186, upload-time = "2025-10-14T15:04:53.064Z" }, + { url = "https://files.pythonhosted.org/packages/cf/68/5707da262a119fb06fbe214d82dd1fe4a6f4af32d2d14de368d0349eb52a/watchfiles-1.1.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1db5d7ae38ff20153d542460752ff397fcf5c96090c1230803713cf3147a6803", size = 456812, upload-time = "2025-10-14T15:04:55.174Z" }, + { url = "https://files.pythonhosted.org/packages/66/ab/3cbb8756323e8f9b6f9acb9ef4ec26d42b2109bce830cc1f3468df20511d/watchfiles-1.1.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:28475ddbde92df1874b6c5c8aaeb24ad5be47a11f87cde5a28ef3835932e3e94", size = 630196, upload-time = "2025-10-14T15:04:56.22Z" }, + { url = "https://files.pythonhosted.org/packages/78/46/7152ec29b8335f80167928944a94955015a345440f524d2dfe63fc2f437b/watchfiles-1.1.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:36193ed342f5b9842edd3532729a2ad55c4160ffcfa3700e0d54be496b70dd43", size = 622657, upload-time = "2025-10-14T15:04:57.521Z" }, + { url = "https://files.pythonhosted.org/packages/0a/bf/95895e78dd75efe9a7f31733607f384b42eb5feb54bd2eb6ed57cc2e94f4/watchfiles-1.1.1-cp312-cp312-win32.whl", hash = "sha256:859e43a1951717cc8de7f4c77674a6d389b106361585951d9e69572823f311d9", size = 272042, upload-time = "2025-10-14T15:04:59.046Z" }, + { url = "https://files.pythonhosted.org/packages/87/0a/90eb755f568de2688cb220171c4191df932232c20946966c27a59c400850/watchfiles-1.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:91d4c9a823a8c987cce8fa2690923b069966dabb196dd8d137ea2cede885fde9", size = 288410, upload-time = "2025-10-14T15:05:00.081Z" }, + { url = "https://files.pythonhosted.org/packages/36/76/f322701530586922fbd6723c4f91ace21364924822a8772c549483abed13/watchfiles-1.1.1-cp312-cp312-win_arm64.whl", hash = "sha256:a625815d4a2bdca61953dbba5a39d60164451ef34c88d751f6c368c3ea73d404", size = 278209, upload-time = "2025-10-14T15:05:01.168Z" }, + { url = "https://files.pythonhosted.org/packages/bb/f4/f750b29225fe77139f7ae5de89d4949f5a99f934c65a1f1c0b248f26f747/watchfiles-1.1.1-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:130e4876309e8686a5e37dba7d5e9bc77e6ed908266996ca26572437a5271e18", size = 404321, upload-time = "2025-10-14T15:05:02.063Z" }, + { url = "https://files.pythonhosted.org/packages/2b/f9/f07a295cde762644aa4c4bb0f88921d2d141af45e735b965fb2e87858328/watchfiles-1.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5f3bde70f157f84ece3765b42b4a52c6ac1a50334903c6eaf765362f6ccca88a", size = 391783, upload-time = "2025-10-14T15:05:03.052Z" }, + { url = "https://files.pythonhosted.org/packages/bc/11/fc2502457e0bea39a5c958d86d2cb69e407a4d00b85735ca724bfa6e0d1a/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:14e0b1fe858430fc0251737ef3824c54027bedb8c37c38114488b8e131cf8219", size = 449279, upload-time = "2025-10-14T15:05:04.004Z" }, + { url = "https://files.pythonhosted.org/packages/e3/1f/d66bc15ea0b728df3ed96a539c777acfcad0eb78555ad9efcaa1274688f0/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f27db948078f3823a6bb3b465180db8ebecf26dd5dae6f6180bd87383b6b4428", size = 459405, upload-time = "2025-10-14T15:05:04.942Z" }, + { url = "https://files.pythonhosted.org/packages/be/90/9f4a65c0aec3ccf032703e6db02d89a157462fbb2cf20dd415128251cac0/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:059098c3a429f62fc98e8ec62b982230ef2c8df68c79e826e37b895bc359a9c0", size = 488976, upload-time = "2025-10-14T15:05:05.905Z" }, + { url = "https://files.pythonhosted.org/packages/37/57/ee347af605d867f712be7029bb94c8c071732a4b44792e3176fa3c612d39/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bfb5862016acc9b869bb57284e6cb35fdf8e22fe59f7548858e2f971d045f150", size = 595506, upload-time = "2025-10-14T15:05:06.906Z" }, + { url = "https://files.pythonhosted.org/packages/a8/78/cc5ab0b86c122047f75e8fc471c67a04dee395daf847d3e59381996c8707/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:319b27255aacd9923b8a276bb14d21a5f7ff82564c744235fc5eae58d95422ae", size = 474936, upload-time = "2025-10-14T15:05:07.906Z" }, + { url = "https://files.pythonhosted.org/packages/62/da/def65b170a3815af7bd40a3e7010bf6ab53089ef1b75d05dd5385b87cf08/watchfiles-1.1.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c755367e51db90e75b19454b680903631d41f9e3607fbd941d296a020c2d752d", size = 456147, upload-time = "2025-10-14T15:05:09.138Z" }, + { url = "https://files.pythonhosted.org/packages/57/99/da6573ba71166e82d288d4df0839128004c67d2778d3b566c138695f5c0b/watchfiles-1.1.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:c22c776292a23bfc7237a98f791b9ad3144b02116ff10d820829ce62dff46d0b", size = 630007, upload-time = "2025-10-14T15:05:10.117Z" }, + { url = "https://files.pythonhosted.org/packages/a8/51/7439c4dd39511368849eb1e53279cd3454b4a4dbace80bab88feeb83c6b5/watchfiles-1.1.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:3a476189be23c3686bc2f4321dd501cb329c0a0469e77b7b534ee10129ae6374", size = 622280, upload-time = "2025-10-14T15:05:11.146Z" }, + { url = "https://files.pythonhosted.org/packages/95/9c/8ed97d4bba5db6fdcdb2b298d3898f2dd5c20f6b73aee04eabe56c59677e/watchfiles-1.1.1-cp313-cp313-win32.whl", hash = "sha256:bf0a91bfb5574a2f7fc223cf95eeea79abfefa404bf1ea5e339c0c1560ae99a0", size = 272056, upload-time = "2025-10-14T15:05:12.156Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f3/c14e28429f744a260d8ceae18bf58c1d5fa56b50d006a7a9f80e1882cb0d/watchfiles-1.1.1-cp313-cp313-win_amd64.whl", hash = "sha256:52e06553899e11e8074503c8e716d574adeeb7e68913115c4b3653c53f9bae42", size = 288162, upload-time = "2025-10-14T15:05:13.208Z" }, + { url = "https://files.pythonhosted.org/packages/dc/61/fe0e56c40d5cd29523e398d31153218718c5786b5e636d9ae8ae79453d27/watchfiles-1.1.1-cp313-cp313-win_arm64.whl", hash = "sha256:ac3cc5759570cd02662b15fbcd9d917f7ecd47efe0d6b40474eafd246f91ea18", size = 277909, upload-time = "2025-10-14T15:05:14.49Z" }, + { url = "https://files.pythonhosted.org/packages/79/42/e0a7d749626f1e28c7108a99fb9bf524b501bbbeb9b261ceecde644d5a07/watchfiles-1.1.1-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:563b116874a9a7ce6f96f87cd0b94f7faf92d08d0021e837796f0a14318ef8da", size = 403389, upload-time = "2025-10-14T15:05:15.777Z" }, + { url = "https://files.pythonhosted.org/packages/15/49/08732f90ce0fbbc13913f9f215c689cfc9ced345fb1bcd8829a50007cc8d/watchfiles-1.1.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3ad9fe1dae4ab4212d8c91e80b832425e24f421703b5a42ef2e4a1e215aff051", size = 389964, upload-time = "2025-10-14T15:05:16.85Z" }, + { url = "https://files.pythonhosted.org/packages/27/0d/7c315d4bd5f2538910491a0393c56bf70d333d51bc5b34bee8e68e8cea19/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce70f96a46b894b36eba678f153f052967a0d06d5b5a19b336ab0dbbd029f73e", size = 448114, upload-time = "2025-10-14T15:05:17.876Z" }, + { url = "https://files.pythonhosted.org/packages/c3/24/9e096de47a4d11bc4df41e9d1e61776393eac4cb6eb11b3e23315b78b2cc/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:cb467c999c2eff23a6417e58d75e5828716f42ed8289fe6b77a7e5a91036ca70", size = 460264, upload-time = "2025-10-14T15:05:18.962Z" }, + { url = "https://files.pythonhosted.org/packages/cc/0f/e8dea6375f1d3ba5fcb0b3583e2b493e77379834c74fd5a22d66d85d6540/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:836398932192dae4146c8f6f737d74baeac8b70ce14831a239bdb1ca882fc261", size = 487877, upload-time = "2025-10-14T15:05:20.094Z" }, + { url = "https://files.pythonhosted.org/packages/ac/5b/df24cfc6424a12deb41503b64d42fbea6b8cb357ec62ca84a5a3476f654a/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:743185e7372b7bc7c389e1badcc606931a827112fbbd37f14c537320fca08620", size = 595176, upload-time = "2025-10-14T15:05:21.134Z" }, + { url = "https://files.pythonhosted.org/packages/8f/b5/853b6757f7347de4e9b37e8cc3289283fb983cba1ab4d2d7144694871d9c/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:afaeff7696e0ad9f02cbb8f56365ff4686ab205fcf9c4c5b6fdfaaa16549dd04", size = 473577, upload-time = "2025-10-14T15:05:22.306Z" }, + { url = "https://files.pythonhosted.org/packages/e1/f7/0a4467be0a56e80447c8529c9fce5b38eab4f513cb3d9bf82e7392a5696b/watchfiles-1.1.1-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3f7eb7da0eb23aa2ba036d4f616d46906013a68caf61b7fdbe42fc8b25132e77", size = 455425, upload-time = "2025-10-14T15:05:23.348Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e0/82583485ea00137ddf69bc84a2db88bd92ab4a6e3c405e5fb878ead8d0e7/watchfiles-1.1.1-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:831a62658609f0e5c64178211c942ace999517f5770fe9436be4c2faeba0c0ef", size = 628826, upload-time = "2025-10-14T15:05:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/28/9a/a785356fccf9fae84c0cc90570f11702ae9571036fb25932f1242c82191c/watchfiles-1.1.1-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:f9a2ae5c91cecc9edd47e041a930490c31c3afb1f5e6d71de3dc671bfaca02bf", size = 622208, upload-time = "2025-10-14T15:05:25.45Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f4/0872229324ef69b2c3edec35e84bd57a1289e7d3fe74588048ed8947a323/watchfiles-1.1.1-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:d1715143123baeeaeadec0528bb7441103979a1d5f6fd0e1f915383fea7ea6d5", size = 404315, upload-time = "2025-10-14T15:05:26.501Z" }, + { url = "https://files.pythonhosted.org/packages/7b/22/16d5331eaed1cb107b873f6ae1b69e9ced582fcf0c59a50cd84f403b1c32/watchfiles-1.1.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:39574d6370c4579d7f5d0ad940ce5b20db0e4117444e39b6d8f99db5676c52fd", size = 390869, upload-time = "2025-10-14T15:05:27.649Z" }, + { url = "https://files.pythonhosted.org/packages/b2/7e/5643bfff5acb6539b18483128fdc0ef2cccc94a5b8fbda130c823e8ed636/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7365b92c2e69ee952902e8f70f3ba6360d0d596d9299d55d7d386df84b6941fb", size = 449919, upload-time = "2025-10-14T15:05:28.701Z" }, + { url = "https://files.pythonhosted.org/packages/51/2e/c410993ba5025a9f9357c376f48976ef0e1b1aefb73b97a5ae01a5972755/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bfff9740c69c0e4ed32416f013f3c45e2ae42ccedd1167ef2d805c000b6c71a5", size = 460845, upload-time = "2025-10-14T15:05:30.064Z" }, + { url = "https://files.pythonhosted.org/packages/8e/a4/2df3b404469122e8680f0fcd06079317e48db58a2da2950fb45020947734/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b27cf2eb1dda37b2089e3907d8ea92922b673c0c427886d4edc6b94d8dfe5db3", size = 489027, upload-time = "2025-10-14T15:05:31.064Z" }, + { url = "https://files.pythonhosted.org/packages/ea/84/4587ba5b1f267167ee715b7f66e6382cca6938e0a4b870adad93e44747e6/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:526e86aced14a65a5b0ec50827c745597c782ff46b571dbfe46192ab9e0b3c33", size = 595615, upload-time = "2025-10-14T15:05:32.074Z" }, + { url = "https://files.pythonhosted.org/packages/6a/0f/c6988c91d06e93cd0bb3d4a808bcf32375ca1904609835c3031799e3ecae/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:04e78dd0b6352db95507fd8cb46f39d185cf8c74e4cf1e4fbad1d3df96faf510", size = 474836, upload-time = "2025-10-14T15:05:33.209Z" }, + { url = "https://files.pythonhosted.org/packages/b4/36/ded8aebea91919485b7bbabbd14f5f359326cb5ec218cd67074d1e426d74/watchfiles-1.1.1-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5c85794a4cfa094714fb9c08d4a218375b2b95b8ed1666e8677c349906246c05", size = 455099, upload-time = "2025-10-14T15:05:34.189Z" }, + { url = "https://files.pythonhosted.org/packages/98/e0/8c9bdba88af756a2fce230dd365fab2baf927ba42cd47521ee7498fd5211/watchfiles-1.1.1-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:74d5012b7630714b66be7b7b7a78855ef7ad58e8650c73afc4c076a1f480a8d6", size = 630626, upload-time = "2025-10-14T15:05:35.216Z" }, + { url = "https://files.pythonhosted.org/packages/2a/84/a95db05354bf2d19e438520d92a8ca475e578c647f78f53197f5a2f17aaf/watchfiles-1.1.1-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:8fbe85cb3201c7d380d3d0b90e63d520f15d6afe217165d7f98c9c649654db81", size = 622519, upload-time = "2025-10-14T15:05:36.259Z" }, + { url = "https://files.pythonhosted.org/packages/1d/ce/d8acdc8de545de995c339be67711e474c77d643555a9bb74a9334252bd55/watchfiles-1.1.1-cp314-cp314-win32.whl", hash = "sha256:3fa0b59c92278b5a7800d3ee7733da9d096d4aabcfabb9a928918bd276ef9b9b", size = 272078, upload-time = "2025-10-14T15:05:37.63Z" }, + { url = "https://files.pythonhosted.org/packages/c4/c9/a74487f72d0451524be827e8edec251da0cc1fcf111646a511ae752e1a3d/watchfiles-1.1.1-cp314-cp314-win_amd64.whl", hash = "sha256:c2047d0b6cea13b3316bdbafbfa0c4228ae593d995030fda39089d36e64fc03a", size = 287664, upload-time = "2025-10-14T15:05:38.95Z" }, + { url = "https://files.pythonhosted.org/packages/df/b8/8ac000702cdd496cdce998c6f4ee0ca1f15977bba51bdf07d872ebdfc34c/watchfiles-1.1.1-cp314-cp314-win_arm64.whl", hash = "sha256:842178b126593addc05acf6fce960d28bc5fae7afbaa2c6c1b3a7b9460e5be02", size = 277154, upload-time = "2025-10-14T15:05:39.954Z" }, + { url = "https://files.pythonhosted.org/packages/47/a8/e3af2184707c29f0f14b1963c0aace6529f9d1b8582d5b99f31bbf42f59e/watchfiles-1.1.1-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:88863fbbc1a7312972f1c511f202eb30866370ebb8493aef2812b9ff28156a21", size = 403820, upload-time = "2025-10-14T15:05:40.932Z" }, + { url = "https://files.pythonhosted.org/packages/c0/ec/e47e307c2f4bd75f9f9e8afbe3876679b18e1bcec449beca132a1c5ffb2d/watchfiles-1.1.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:55c7475190662e202c08c6c0f4d9e345a29367438cf8e8037f3155e10a88d5a5", size = 390510, upload-time = "2025-10-14T15:05:41.945Z" }, + { url = "https://files.pythonhosted.org/packages/d5/a0/ad235642118090f66e7b2f18fd5c42082418404a79205cdfca50b6309c13/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3f53fa183d53a1d7a8852277c92b967ae99c2d4dcee2bfacff8868e6e30b15f7", size = 448408, upload-time = "2025-10-14T15:05:43.385Z" }, + { url = "https://files.pythonhosted.org/packages/df/85/97fa10fd5ff3332ae17e7e40e20784e419e28521549780869f1413742e9d/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6aae418a8b323732fa89721d86f39ec8f092fc2af67f4217a2b07fd3e93c6101", size = 458968, upload-time = "2025-10-14T15:05:44.404Z" }, + { url = "https://files.pythonhosted.org/packages/47/c2/9059c2e8966ea5ce678166617a7f75ecba6164375f3b288e50a40dc6d489/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f096076119da54a6080e8920cbdaac3dbee667eb91dcc5e5b78840b87415bd44", size = 488096, upload-time = "2025-10-14T15:05:45.398Z" }, + { url = "https://files.pythonhosted.org/packages/94/44/d90a9ec8ac309bc26db808a13e7bfc0e4e78b6fc051078a554e132e80160/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:00485f441d183717038ed2e887a7c868154f216877653121068107b227a2f64c", size = 596040, upload-time = "2025-10-14T15:05:46.502Z" }, + { url = "https://files.pythonhosted.org/packages/95/68/4e3479b20ca305cfc561db3ed207a8a1c745ee32bf24f2026a129d0ddb6e/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a55f3e9e493158d7bfdb60a1165035f1cf7d320914e7b7ea83fe22c6023b58fc", size = 473847, upload-time = "2025-10-14T15:05:47.484Z" }, + { url = "https://files.pythonhosted.org/packages/4f/55/2af26693fd15165c4ff7857e38330e1b61ab8c37d15dc79118cdba115b7a/watchfiles-1.1.1-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8c91ed27800188c2ae96d16e3149f199d62f86c7af5f5f4d2c61a3ed8cd3666c", size = 455072, upload-time = "2025-10-14T15:05:48.928Z" }, + { url = "https://files.pythonhosted.org/packages/66/1d/d0d200b10c9311ec25d2273f8aad8c3ef7cc7ea11808022501811208a750/watchfiles-1.1.1-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:311ff15a0bae3714ffb603e6ba6dbfba4065ab60865d15a6ec544133bdb21099", size = 629104, upload-time = "2025-10-14T15:05:49.908Z" }, + { url = "https://files.pythonhosted.org/packages/e3/bd/fa9bb053192491b3867ba07d2343d9f2252e00811567d30ae8d0f78136fe/watchfiles-1.1.1-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:a916a2932da8f8ab582f242c065f5c81bed3462849ca79ee357dd9551b0e9b01", size = 622112, upload-time = "2025-10-14T15:05:50.941Z" }, + { url = "https://files.pythonhosted.org/packages/ba/4c/a888c91e2e326872fa4705095d64acd8aa2fb9c1f7b9bd0588f33850516c/watchfiles-1.1.1-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:17ef139237dfced9da49fb7f2232c86ca9421f666d78c264c7ffca6601d154c3", size = 409611, upload-time = "2025-10-14T15:06:05.809Z" }, + { url = "https://files.pythonhosted.org/packages/1e/c7/5420d1943c8e3ce1a21c0a9330bcf7edafb6aa65d26b21dbb3267c9e8112/watchfiles-1.1.1-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:672b8adf25b1a0d35c96b5888b7b18699d27d4194bac8beeae75be4b7a3fc9b2", size = 396889, upload-time = "2025-10-14T15:06:07.035Z" }, + { url = "https://files.pythonhosted.org/packages/0c/e5/0072cef3804ce8d3aaddbfe7788aadff6b3d3f98a286fdbee9fd74ca59a7/watchfiles-1.1.1-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:77a13aea58bc2b90173bc69f2a90de8e282648939a00a602e1dc4ee23e26b66d", size = 451616, upload-time = "2025-10-14T15:06:08.072Z" }, + { url = "https://files.pythonhosted.org/packages/83/4e/b87b71cbdfad81ad7e83358b3e447fedd281b880a03d64a760fe0a11fc2e/watchfiles-1.1.1-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b495de0bb386df6a12b18335a0285dda90260f51bdb505503c02bcd1ce27a8b", size = 458413, upload-time = "2025-10-14T15:06:09.209Z" }, + { url = "https://files.pythonhosted.org/packages/d3/8e/e500f8b0b77be4ff753ac94dc06b33d8f0d839377fee1b78e8c8d8f031bf/watchfiles-1.1.1-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:db476ab59b6765134de1d4fe96a1a9c96ddf091683599be0f26147ea1b2e4b88", size = 408250, upload-time = "2025-10-14T15:06:10.264Z" }, + { url = "https://files.pythonhosted.org/packages/bd/95/615e72cd27b85b61eec764a5ca51bd94d40b5adea5ff47567d9ebc4d275a/watchfiles-1.1.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:89eef07eee5e9d1fda06e38822ad167a044153457e6fd997f8a858ab7564a336", size = 396117, upload-time = "2025-10-14T15:06:11.28Z" }, + { url = "https://files.pythonhosted.org/packages/c9/81/e7fe958ce8a7fb5c73cc9fb07f5aeaf755e6aa72498c57d760af760c91f8/watchfiles-1.1.1-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce19e06cbda693e9e7686358af9cd6f5d61312ab8b00488bc36f5aabbaf77e24", size = 450493, upload-time = "2025-10-14T15:06:12.321Z" }, + { url = "https://files.pythonhosted.org/packages/6e/d4/ed38dd3b1767193de971e694aa544356e63353c33a85d948166b5ff58b9e/watchfiles-1.1.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3e6f39af2eab0118338902798b5aa6664f46ff66bc0280de76fca67a7f262a49", size = 457546, upload-time = "2025-10-14T15:06:13.372Z" }, +] + [[package]] name = "webob" version = "1.8.9" @@ -2103,8 +2757,34 @@ dependencies = [ ] [package.optional-dependencies] +dev = [ + { name = "mkdocstrings", extra = ["python"] }, + { name = "pre-commit" }, + { name = "pytest" }, + { name = "watchfiles" }, + { name = "zensical" }, +] +docs = [ + { name = "mkdocstrings", extra = ["python"] }, + { name = "zensical" }, +] +duckdb = [ + { name = "duckdb" }, +] +geo = [ + { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +polars = [ + { name = "polars" }, +] test = [ + { name = "cftime" }, + { name = "duckdb" }, { name = "gcsfs" }, + { name = "polars" }, + { name = "pyproj", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pyproj", version = "3.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pytest" }, { name = "xarray", version = "2025.6.1", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version < '3.11'" }, { name = "xarray", version = "2025.7.0", source = { registry = "https://pypi.org/simple" }, extra = ["io"], marker = "python_full_version >= '3.11'" }, @@ -2114,26 +2794,38 @@ test = [ dev = [ { name = "maturin" }, { name = "py-spy" }, - { name = "pyink" }, - { name = "xarray-sql", extra = ["test"] }, + { name = "ruff" }, + { name = "xarray-sql", extra = ["docs", "test"] }, ] [package.metadata] requires-dist = [ + { name = "cftime", marker = "extra == 'test'" }, { name = "dask", specifier = ">=2024.8.0" }, - { name = "datafusion", specifier = "==51.0.0" }, + { name = "datafusion", specifier = "==54.0.0" }, + { name = "duckdb", marker = "extra == 'duckdb'", specifier = ">=1.4.0" }, { name = "gcsfs", marker = "extra == 'test'" }, + { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'" }, + { name = "polars", marker = "extra == 'polars'", specifier = ">=1.33" }, + { name = "pre-commit", marker = "extra == 'dev'" }, + { name = "pyproj", marker = "extra == 'geo'" }, + { name = "pytest", marker = "extra == 'dev'" }, { name = "pytest", marker = "extra == 'test'" }, + { name = "watchfiles", marker = "extra == 'dev'" }, { name = "xarray", specifier = ">=2024.7.0" }, { name = "xarray", extras = ["io"], marker = "extra == 'test'" }, + { name = "xarray-sql", extras = ["docs"], marker = "extra == 'dev'" }, + { name = "xarray-sql", extras = ["duckdb", "polars", "geo"], marker = "extra == 'test'" }, + { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["test"] +provides-extras = ["dev", "docs", "duckdb", "geo", "polars", "test"] [package.metadata.requires-dev] dev = [ { name = "maturin", specifier = ">=1.9.1" }, { name = "py-spy", specifier = ">=0.4.0" }, - { name = "pyink", specifier = ">=24.10.1" }, + { name = "ruff", specifier = ">=0.15.10" }, + { name = "xarray-sql", extras = ["docs"] }, { name = "xarray-sql", extras = ["test"] }, ] @@ -2276,6 +2968,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/57/3329346940f78de49047ddcb03fdbca9e16450c3a942688bf24201a322e5/zarr-3.0.10-py3-none-any.whl", hash = "sha256:110724c045fbe4ff5509a8a2a6b6098cb244a6af43da85eaeecef9821473163f", size = 209508, upload-time = "2025-07-03T17:29:25.6Z" }, ] +[[package]] +name = "zensical" +version = "0.0.29" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "deepmerge" }, + { name = "markdown" }, + { name = "pygments" }, + { name = "pymdown-extensions" }, + { name = "pyyaml" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/78/bd/5786ab618a60bd7469ab243a7fd2c9eecb0790c85c784abb8b97edb77a54/zensical-0.0.29.tar.gz", hash = "sha256:0d6282be7cb551e12d5806badf5e94c54a5e2f2cf07057a3e36d1eaf97c33ada", size = 3842641, upload-time = "2026-03-24T13:37:27.587Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/9c/8b681daa024abca9763017bec09ecee8008e110cae1254217c8dd22cc339/zensical-0.0.29-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:20ae0709ea14fce25ab33d0a82acdaf454a7a2e232a9ee20c019942205174476", size = 12311399, upload-time = "2026-03-24T13:36:53.809Z" }, + { url = "https://files.pythonhosted.org/packages/81/ae/4ebb4d8bb2ef0164d473698b92f11caf431fc436e1625524acd5641102ca/zensical-0.0.29-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:599af3ba66fcd0146d7019f3493ed3c316051fae6c4d5599bc59f3a8f4b8a6f0", size = 12191845, upload-time = "2026-03-24T13:36:56.909Z" }, + { url = "https://files.pythonhosted.org/packages/d5/35/67f89db06571a52283b3ecbe3bcf32fd3115ca50436b3ae177a948b83ea7/zensical-0.0.29-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:eea7e48a00a71c0586e875079b5f83a070c33a147e52ad4383e4b63ab524332b", size = 12554105, upload-time = "2026-03-24T13:36:59.945Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f6/ac79e5d9c18b28557c9ff1c7c23d695fbdd82645d69bfe02292f46d935e7/zensical-0.0.29-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:59a57db35542e98d2896b833de07d199320f8ada3b4e7ddccb7fe892292d8b74", size = 12498643, upload-time = "2026-03-24T13:37:02.376Z" }, + { url = "https://files.pythonhosted.org/packages/b1/70/5c22a96a69e0e91e569c26236918bb9bab1170f59b29ad04105ead64f199/zensical-0.0.29-cp310-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d42c2b2a96a80cf64c98ba7242f59ef95109914bd4c9499d7ebc12544663852c", size = 12854531, upload-time = "2026-03-24T13:37:04.962Z" }, + { url = "https://files.pythonhosted.org/packages/79/25/e32237a8fcb0ceae1ef8e192e7f8db53b38f1e48f1c7cdbacd0a7b713892/zensical-0.0.29-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b2fca39c5f6b1782c77cf6591cf346357cabee85ebdb956c5ddc0fd5169f3d9", size = 12596828, upload-time = "2026-03-24T13:37:07.817Z" }, + { url = "https://files.pythonhosted.org/packages/ff/74/89ac909cbb258903ea53802c184e4986c17ce0ba79b1c7f77b7e78a2dce3/zensical-0.0.29-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:dfc23a74ef672aa51088c080286319da1dc0b989cd5051e9e5e6d7d4abbc2fc1", size = 12732059, upload-time = "2026-03-24T13:37:11.651Z" }, + { url = "https://files.pythonhosted.org/packages/8c/31/2429de6a9328eed4acc7e9a3789f160294a15115be15f9870a0d02649302/zensical-0.0.29-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:c9336d4e4b232e3c9a70e30258e916dd7e60c0a2a08c8690065e60350c302028", size = 12768542, upload-time = "2026-03-24T13:37:14.39Z" }, + { url = "https://files.pythonhosted.org/packages/10/8a/55588b2a1dcbe86dad0404506c9ba367a06c663b1ff47147c84d26f7510e/zensical-0.0.29-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:30661148f0681199f3b598cbeb1d54f5cba773e54ae840bac639250d85907b84", size = 12917991, upload-time = "2026-03-24T13:37:16.795Z" }, + { url = "https://files.pythonhosted.org/packages/ec/5d/653901f0d3a3ca72daebc62746a148797f4e422cc3a2b66a4e6718e4398f/zensical-0.0.29-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6a566ac1fd4bfac5d711a7bd1ae06666712127c2718daa5083c7bf3f107e8578", size = 12868392, upload-time = "2026-03-24T13:37:19.42Z" }, + { url = "https://files.pythonhosted.org/packages/29/58/d7449bc88a174b98daa3f2fbdfbdac3493768a557d8987e88bdaa6c78b1a/zensical-0.0.29-cp310-abi3-win32.whl", hash = "sha256:a231a3a02a3851741dc4d2de8910b5c39fe81e55bf026d8edf4d803e91a922fb", size = 11905486, upload-time = "2026-03-24T13:37:22.154Z" }, + { url = "https://files.pythonhosted.org/packages/f5/09/3fd082d016497c4d26ff20f42a8be2cc91e27191c0c5f3cd6507827f666f/zensical-0.0.29-cp310-abi3-win_amd64.whl", hash = "sha256:7145c5504380a344b8cd4586da815cdde77ef4a42319fa4f35e78250f01985af", size = 12101510, upload-time = "2026-03-24T13:37:24.77Z" }, +] + [[package]] name = "zipp" version = "3.23.0" diff --git a/xarray_sql/__init__.py b/xarray_sql/__init__.py index 60508d59..0f98e1e0 100644 --- a/xarray_sql/__init__.py +++ b/xarray_sql/__init__.py @@ -1,10 +1,19 @@ +from . import cftime +from .backends import arrow_dataset, register +from .geometry import bbox_conjuncts +from .df import from_map from .reader import read_xarray, read_xarray_table +from .roundtrip import to_dataset from .sql import XarrayContext -from .df import from_map __all__ = [ + "cftime", "XarrayContext", "read_xarray_table", "read_xarray", + "arrow_dataset", + "bbox_conjuncts", + "register", + "to_dataset", "from_map", # deprecated ] diff --git a/xarray_sql/backends/__init__.py b/xarray_sql/backends/__init__.py new file mode 100644 index 00000000..82aad8eb --- /dev/null +++ b/xarray_sql/backends/__init__.py @@ -0,0 +1,34 @@ +"""Engine adapters — the *register* seam of xarray-sql. + +xarray-sql translates data, not queries, across two seams — the two +boundaries between xarray and a query engine that neither side builds +for itself: *register* (a lazy ``xarray.Dataset`` becomes a table on +the engine's own connection; this package) and *round-trip* (an Arrow +result becomes a labeled Dataset again; [xarray_sql.to_dataset][]). +SQL dialects, geometry, H3, and optimizers belong to each engine and +its extension ecosystem. + +Adapters register themselves on import via +[register_adapter][xarray_sql.backends.base.register_adapter]; +[register][xarray_sql.backends.base.register] dispatches on the connection +type. +""" + +from .base import EngineAdapter, get_adapter, register, register_adapter +from . import datafusion as _datafusion # noqa: F401 (self-registers) +from . import duckdb as _duckdb # noqa: F401 (self-registers) +from .pyarrow import ( + XarrayArrowStream, + XarrayPushdownDataset, + arrow_dataset, +) + +__all__ = [ + "EngineAdapter", + "XarrayArrowStream", + "XarrayPushdownDataset", + "arrow_dataset", + "get_adapter", + "register", + "register_adapter", +] diff --git a/xarray_sql/backends/base.py b/xarray_sql/backends/base.py new file mode 100644 index 00000000..59eb488c --- /dev/null +++ b/xarray_sql/backends/base.py @@ -0,0 +1,118 @@ +"""Engine-adapter dispatch for [xarray_sql.register][]. + +An *engine adapter* implements the register seam: given an engine's native +connection object and a lazy ``xarray.Dataset``, register the Dataset as +a queryable table on that connection. The Arrow C-stream protocol is the +common wire between xarray and every engine; adapters differ only in how +a stream is attached to the connection and in what pushdown the engine +can do against it. + +Adapters self-describe which connections they accept via ``matches``, +which must not require the engine's package to be importable (detection +is by type inspection), so optional engines stay optional. +""" + +from __future__ import annotations + +from typing import Any, Protocol, TypeGuard, TypeVar, cast + +import xarray as xr + +from ..df import Chunks + +ConT = TypeVar("ConT") +"""An engine's native connection type (e.g. ``duckdb.DuckDBPyConnection``).""" + + +class EngineAdapter(Protocol[ConT]): + """One engine's implementation of the register seam.""" + + @staticmethod + def matches(con: object) -> TypeGuard[ConT]: + """Whether *con* is a connection this adapter can register into.""" + ... + + @staticmethod + def register( + con: ConT, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> ConT: + """Register *ds* as table *name* on *con*; returns *con*.""" + ... + + +_ADAPTERS: list[type[EngineAdapter[Any]]] = [] + +_A = TypeVar("_A", bound=type[EngineAdapter[Any]]) + + +def register_adapter(cls: _A) -> _A: + """Class decorator adding an adapter to the dispatch list.""" + _ADAPTERS.append(cls) + return cls + + +def get_adapter(con: object) -> type[EngineAdapter[Any]]: + """Return the first adapter whose ``matches(con)`` is true.""" + for adapter in _ADAPTERS: + if adapter.matches(con): + return adapter + raise TypeError( + f"No xarray-sql engine adapter for connection of type " + f"{type(con).__module__}.{type(con).__qualname__}. " + f"Supported: DataFusion SessionContext and DuckDB connections." + ) + + +def register( + con: ConT, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, +) -> ConT: + """Register a lazy xarray Dataset as a table on an engine connection. + + The engine is inferred from the connection type. Data is not read at + registration time; the engine pulls Arrow record batches lazily during + query execution. Write your SQL in the engine's own dialect and use + the engine's extension ecosystem directly — xarray-sql translates the + data, not the queries. + + Example (DuckDB):: + + import duckdb + import xarray_sql as xql + + con = duckdb.connect() + xql.register(con, "era5", ds) + rel = con.sql("SELECT time, AVG(t2m) AS t2m FROM era5 GROUP BY time") + result = xql.to_dataset(rel, template=ds) + + Args: + con: An engine connection: a ``datafusion.SessionContext`` (or + [xarray_sql.XarrayContext][]) or a + ``duckdb.DuckDBPyConnection``. + name: The table name to register the Dataset under. Datasets + whose variables have differing dimensions are split into one + table per dimension group (a SQL schema ``name.group`` on + DataFusion; ``name_group`` tables on DuckDB). + ds: An xarray Dataset. + chunks: Xarray-like chunks specification controlling partition + granularity. Defaults to the Dataset's existing chunks. + **kwargs: Adapter-specific options, forwarded as-is — e.g. + ``table_names`` on DataFusion, ``batch_size`` / ``prefetch`` + on DuckDB. + + Returns: + The connection, to allow chaining. + """ + # The connection type is erased by the runtime dispatch; every adapter + # returns the connection it was given. + adapter: Any = get_adapter(con) + return cast(ConT, adapter.register(con, name, ds, chunks=chunks, **kwargs)) diff --git a/xarray_sql/backends/datafusion.py b/xarray_sql/backends/datafusion.py new file mode 100644 index 00000000..3d7c29e1 --- /dev/null +++ b/xarray_sql/backends/datafusion.py @@ -0,0 +1,46 @@ +"""DataFusion engine adapter. + +DataFusion is xarray-sql's default engine and the richest adapter: the +Rust ``LazyArrowStreamTable`` table provider gives partition pruning on +dimension predicates, projection pushdown, and exact per-partition +statistics for the optimizer. This module only routes the generic +[xarray_sql.register][] seam onto that existing machinery. +""" + +from __future__ import annotations + +from typing import Any, TypeGuard + +import xarray as xr +from datafusion import SessionContext + +from ..df import Chunks +from ..reader import read_xarray_table +from ..sql import XarrayContext +from .base import register_adapter + + +@register_adapter +class DataFusionAdapter: + """Registers Datasets on ``datafusion.SessionContext`` connections.""" + + @staticmethod + def matches(con: object) -> TypeGuard[SessionContext]: + return isinstance(con, SessionContext) + + @staticmethod + def register( + con: SessionContext, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> SessionContext: + # XarrayContext.from_dataset adds dim-group splitting, cftime UDF + # registration, and round-trip metadata tracking on top of the + # plain table registration; use it when available. + if isinstance(con, XarrayContext): + return con.from_dataset(name, ds, chunks=chunks, **kwargs) + con.register_table(name, read_xarray_table(ds, chunks, **kwargs)) + return con diff --git a/xarray_sql/backends/duckdb.py b/xarray_sql/backends/duckdb.py new file mode 100644 index 00000000..1651ac05 --- /dev/null +++ b/xarray_sql/backends/duckdb.py @@ -0,0 +1,89 @@ +"""DuckDB engine adapter. + +Registers a lazy ``xarray.Dataset`` on a ``duckdb.DuckDBPyConnection`` +as an [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]: +DuckDB classifies it with a real ``isinstance`` check against +``pyarrow.dataset.Dataset`` and calls ``scanner(columns=[...], +filter=)`` once per query, giving +projection pushdown, coordinate-range chunk pruning, and prefetched +parallel production (see [xarray_sql.backends.pyarrow][]). + +This adapter never imports the ``duckdb`` package at runtime — detection +is by connection type, and registration is a method call on the +connection — so DuckDB stays a purely optional dependency +(``pip install xarray-sql[duckdb]``). + +Zarr-native scanning inside DuckDB is what the [duckdb-zarr](https://github.com/xqlsystems/duckdb-zarr) extension provides; this +adapter instead covers everything xarray can open (NetCDF, GRIB, Xee, CF +decoding, in-memory) and pairs with [xarray_sql.to_dataset][] for +the labeled round-trip. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypeGuard + +import xarray as xr + +from ..df import Chunks, group_vars_by_dims +from .base import register_adapter +from .pyarrow import XarrayArrowStream, XarrayPushdownDataset + +if TYPE_CHECKING: + import duckdb + +__all__ = ["DuckDBAdapter", "XarrayArrowStream", "XarrayPushdownDataset"] + + +@register_adapter +class DuckDBAdapter: + """Registers Datasets on ``duckdb.DuckDBPyConnection`` connections.""" + + @staticmethod + def matches(con: object) -> TypeGuard[duckdb.DuckDBPyConnection]: + # The connection class lives in ``duckdb`` or, in newer releases, + # the ``_duckdb`` C-extension module. + root = type(con).__module__.split(".")[0] + return root in ("duckdb", "_duckdb") + + @staticmethod + def register( + con: duckdb.DuckDBPyConnection, + name: str, + ds: xr.Dataset, + *, + chunks: Chunks = None, + **kwargs: Any, + ) -> duckdb.DuckDBPyConnection: + """Register ``ds`` on a DuckDB connection. + + Datasets whose variables all share the same dimensions become a + single table named ``name``. Mixed-dimension datasets are split + into one table per dimension group, named + ``___...`` (DuckDB registration has no schema + namespace to mirror the DataFusion adapter's ``name.group`` + layout). Extra keyword arguments (``batch_size``, ``prefetch``) + are forwarded to [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + """ + groups = group_vars_by_dims(ds) + if len(groups) <= 1: + con.register(name, XarrayPushdownDataset(ds, chunks, **kwargs)) + return con + # Materialise dim coordinates once and share across sub-tables. + coord_arrays = { + str(dim): ds.coords[dim].values + for dim in ds.dims + if dim in ds.coords + } + for dims, var_names in groups.items(): + suffix = "_".join(dims) or "scalar" + con.register( + f"{name}_{suffix}", + XarrayPushdownDataset( + ds[var_names], + chunks, + coord_arrays=coord_arrays, + **kwargs, + ), + ) + return con diff --git a/xarray_sql/backends/pyarrow.py b/xarray_sql/backends/pyarrow.py new file mode 100644 index 00000000..dd0f2543 --- /dev/null +++ b/xarray_sql/backends/pyarrow.py @@ -0,0 +1,1143 @@ +"""Engine-neutral pyarrow views of lazy xarray Datasets. + +Two ways to hand a lazy ``xarray.Dataset`` to an Arrow-speaking query +engine: + +* [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset] — a real ``pyarrow.dataset.Dataset`` + subclass (the pattern Lance uses for ``LanceDataset``). Consumers of + the pyarrow dataset protocol — DuckDB via ``con.register``, Polars via + ``pl.scan_pyarrow_dataset``, or pyarrow itself — call + [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the columns a query needs + and the predicate it pushed down, so the scan loads only the needed + data variables from only the chunks whose coordinate ranges can + satisfy the predicate. Construct one with [arrow_dataset][xarray_sql.backends.pyarrow.arrow_dataset]. +* [XarrayArrowStream][xarray_sql.backends.pyarrow.XarrayArrowStream] — a re-scannable Arrow C-stream + (PyCapsule) view. No source-level pushdown, but works with any + PyCapsule consumer; the dependency-light fallback. + +Correctness contract shared by all consumers of the pushdown dataset: +engines may delete the filter conjuncts they push down (DuckDB does), +so the returned scanner applies the expression exactly via +``pyarrow.dataset.Scanner``; chunk pruning is only ever an optimization +on top. +""" + +from __future__ import annotations + +import itertools +import math +import re +import threading +import weakref +from collections import deque +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.dataset as pads +import pyarrow.fs as pafs +import xarray as xr + +from ..df import ( + Block, + Chunks, + DEFAULT_BATCH_SIZE, + _ensure_default_indexes, + _parse_schema, + iter_record_batches, + resolve_chunks, +) +from ..geometry import GEOMETRY_COLUMN, build_geometry, geometry_field +from ..reader import XarrayRecordBatchReader + +DEFAULT_PREFETCH = 4 +"""Chunk loads kept in flight ahead of the consumer during a scan.""" + +_SHADOW_FANOUT = 1024 +"""Maximum fragments per shadow level. + +A dimension with more chunks than this gets a two-level shadow: a coarse +level of at most this many buckets, refined per surviving bucket. This +bounds shadow construction cost for finely partitioned datasets (e.g. +hundreds of thousands of single-step time chunks) at registration and +query time alike. +""" + +_REFINE_MAX_FRACTION = 0.25 +"""Skip fine-level pruning when the coarse pass kept more buckets. + +Refinement builds one sub-shadow per surviving bucket; when a predicate +matches most of the axis that cost cannot pay for itself, so the scan +falls back to the (sound) coarse answer. +""" + +_STRICT_LEVEL_BUDGET = 4096 +"""Maximum guarantee fragments per level of the strictness analysis. + +Strictness classifies bucket-products of surviving chunks recursively: +whole buckets prove (or prune) at once, and only mixed cells refine. +Each level builds at most this many fragments, so grids of millions of +chunks resolve in a handful of vectorized passes. +""" + +_STRICT_MAX_DEPTH = 6 +"""Recursion bound for the strictness analysis (a backstop: realistic +grids terminate in 2-3 levels).""" + + +def _guarantee_shadow( + guarantees: list[tuple[str, pc.Expression]], schema: pa.Schema +) -> pads.FileSystemDataset: + """A path-only dataset whose fragments carry the given guarantees. + + Each ``(path, guarantee)`` pair becomes one fragment: the path is a + label (never opened) and the guarantee is its + ``partition_expression``, so ``get_fragments(filter=...)`` delegates + satisfiability to Arrow's guarantee simplification. + """ + fmt = pads.IpcFileFormat() + fs = pafs.LocalFileSystem() + fragments = [ + fmt.make_fragment(path, fs, partition_expression=guarantee) + for path, guarantee in guarantees + ] + return pads.FileSystemDataset(fragments, schema, fmt, fs) + + +class _DimShadow: + """Chunk-pruning index for one dimension of the source grid. + + Fragment ``i`` of a shadow ``FileSystemDataset`` carries the + guarantee ``dim ∈ [min, max]`` of chunk-span ``i`` as its + ``partition_expression``; ``get_fragments(filter=...)`` then lets + Arrow's guarantee simplification decide which spans can satisfy a + predicate — sound for every predicate shape, conservative on columns + the guarantee does not mention, and the fragments' paths are never + opened. + + Axes with more than ``_SHADOW_FANOUT`` chunks use two levels: a + coarse shadow over buckets of consecutive chunks, plus per-bucket + fine shadows built lazily for the buckets a query keeps. + """ + + def __init__( + self, + name: str, + schema: pa.Schema, + coord: np.ndarray, + bounds: np.ndarray, + ): + self._name = name + # The full table schema, not just this dimension's field: the + # pushed predicate may reference any column, and get_fragments + # must be able to bind all of them (guarantees stay per-dim; + # unmentioned columns are conservatively unconstrained). + self._schema = schema + self._field_type = schema.field(name).type + self._coord = coord + self._bounds = bounds + self._n = len(bounds) - 1 + self._step = max(1, math.ceil(self._n / _SHADOW_FANOUT)) + self._n_buckets = math.ceil(self._n / self._step) + self._coarse = self._build( + [ + (b * self._step, min((b + 1) * self._step, self._n)) + for b in range(self._n_buckets) + ] + ) + self._fine: dict[int, pads.FileSystemDataset] = {} + + def _build(self, spans: list[tuple[int, int]]) -> pads.FileSystemDataset: + """A shadow whose fragment ``i`` guarantees chunk-span ``spans[i]``.""" + guarantees: list[tuple[str, pc.Expression]] = [] + for i, (lo_chunk, hi_chunk) in enumerate(spans): + vals = self._coord[self._bounds[lo_chunk] : self._bounds[hi_chunk]] + if (vals.dtype.kind == "f" and np.isnan(vals).any()) or ( + vals.dtype.kind == "M" and np.isnat(vals).any() + ): + # NaN/NaT poisons min/max into a (dim >= NaN) guarantee + # that Arrow simplifies every predicate against as false, + # silently pruning rows. An always-true guarantee keeps + # the span unprunable instead. + guarantee = pc.scalar(True) + else: + # min/max (not first/last) so descending axes like + # latitude 90→-90 carry correct ranges. + lo = pa.scalar(vals.min(), type=self._field_type) + hi = pa.scalar(vals.max(), type=self._field_type) + guarantee = (pc.field(self._name) >= lo) & ( + pc.field(self._name) <= hi + ) + guarantees.append((str(i), guarantee)) + return _guarantee_shadow(guarantees, self._schema) + + @staticmethod + def _kept_indices( + shadow: pads.FileSystemDataset, filter: pc.Expression + ) -> list[int]: + return sorted( + int(frag.path) for frag in shadow.get_fragments(filter=filter) + ) + + def kept(self, filter: pc.Expression) -> list[int] | None: + """Chunk indices that can satisfy ``filter``; ``None`` means all.""" + try: + buckets = self._kept_indices(self._coarse, filter) + if self._step == 1: + return buckets if len(buckets) < self._n else None + if len(buckets) > _REFINE_MAX_FRACTION * self._n_buckets: + # Refining most of the axis costs more than it saves; + # answer with the coarse buckets, which is still sound. + return ( + None + if len(buckets) == self._n_buckets + else [ + i + for b in buckets + for i in range( + b * self._step, + min((b + 1) * self._step, self._n), + ) + ] + ) + kept: list[int] = [] + for b in buckets: + fine = self._fine.get(b) + if fine is None: + start = b * self._step + stop = min((b + 1) * self._step, self._n) + fine = self._build([(i, i + 1) for i in range(start, stop)]) + self._fine[b] = fine + start = b * self._step + kept.extend(start + i for i in self._kept_indices(fine, filter)) + return kept + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + return None # conservative: scan every chunk of this dim + + +class XarrayArrowStream: + """A re-scannable Arrow C-stream view over a lazy xarray Dataset. + + Arrow PyCapsule consumers (DuckDB among them) call + ``__arrow_c_stream__`` once per scan. Each call constructs a fresh + [XarrayRecordBatchReader][xarray_sql.reader.XarrayRecordBatchReader] over the same + lazy Dataset, so — unlike registering a ``pyarrow.RecordBatchReader`` + directly, which is exhausted after one query — the same registered + table supports any number of queries, and data is only read while a + query is executing. + + The PyCapsule scan path gets no source-level pushdown (the producer + never sees the query's columns or filters), so + [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset] is the default registration object; + this class remains as the dependency-light fallback. + """ + + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + # Validate eagerly (same checks XarrayRecordBatchReader runs) so + # registration fails fast instead of erroring mid-query. + probe = XarrayRecordBatchReader(ds, chunks, batch_size=batch_size) + self._ds = ds + self._chunks = chunks + self._batch_size = batch_size + self._schema = probe.schema + self._iteration_callback = _iteration_callback + + def __arrow_c_stream__( + self, requested_schema: object | None = None + ) -> object: + reader = XarrayRecordBatchReader( + self._ds, + self._chunks, + batch_size=self._batch_size, + _iteration_callback=self._iteration_callback, + ) + return reader.__arrow_c_stream__(requested_schema) + + def __arrow_c_schema__(self) -> object: + return self._schema.__arrow_c_schema__() + + +class XarrayPushdownDataset(pads.Dataset): + """A pushdown-capable ``pyarrow.dataset.Dataset`` view of a Dataset. + + Consumers that speak the pyarrow dataset protocol (DuckDB, Polars, + ...) call [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the columns a query needs and the + predicate it pushed down; the scan then loads only the needed data + variables from only the chunks whose coordinate ranges can satisfy + the predicate. + + The base class is never initialized (there is no C++ dataset behind + this object — the same construction Lance uses for ``LanceDataset``); + every entry point consumers touch is overridden in Python, and the + few inherited members that would read uninitialized native state are + stubbed out. + + References: + Lance's ``LanceDataset``, a ``pyarrow.dataset.Dataset`` subclass + built the same way: https://github.com/lancedb/lance + (``python/python/lance/dataset.py``). + """ + + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + prefetch: int = DEFAULT_PREFETCH, + prefetch_bytes: int | None = None, + coalesce_rows: int | None = None, + geometry: tuple[str, str] | None = None, + geometry_encoding: str = "wkb", + geometry_crs: str | None = "OGC:CRS84", + coord_arrays: dict[str, np.ndarray] | None = None, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + # Deliberately no super().__init__() — see class docstring. + ds = _ensure_default_indexes(ds) + if ds.data_vars: + fst = next(iter(ds.values())).dims + if not all(da.dims == fst for da in ds.values()): + raise ValueError( + "All dimensions must be equal. " + "Please filter data_vars in the Dataset." + ) + self._ds = ds + self._schema = _parse_schema(ds) + self._geometry = tuple(geometry) if geometry else None + self._geometry_encoding = geometry_encoding + if self._geometry: + if GEOMETRY_COLUMN in self._schema.names: + raise ValueError( + f"geometry= would shadow an existing column named " + f"{GEOMETRY_COLUMN!r}." + ) + missing = [d for d in self._geometry if d not in self._schema.names] + if missing: + raise ValueError( + f"geometry dims {missing} are not columns of the " + f"table; available: {self._schema.names}" + ) + self._schema = self._schema.append( + geometry_field(geometry_encoding, geometry_crs) + ) + self._resolved = resolve_chunks(ds, chunks) + if not self._resolved and ds.sizes: + raise ValueError( + "Dataset `ds` must be chunked or `chunks` must be provided." + ) + self._chunk_bounds = { + d: np.cumsum((0, *sizes)) for d, sizes in self._resolved.items() + } + # Reuse pre-materialised coordinate arrays where the caller has + # them (e.g. shared across the tables of a dim-group split); each + # missing dim costs one read, a network round-trip for Zarr. + self._coord_arrays = dict(coord_arrays or {}) + for d in ds.dims: + if str(d) not in self._coord_arrays: + self._coord_arrays[str(d)] = ds.coords[d].values + if batch_size <= 0: + # A zero size would never advance the zero-column scan's + # row loop; fail here rather than mid-scan. + raise ValueError(f"batch_size must be positive, got {batch_size}") + self._batch_size = batch_size + self._prefetch = prefetch + self._prefetch_bytes = prefetch_bytes + self._coalesce_rows = coalesce_rows + self._iteration_callback = _iteration_callback + self._shadows: dict[str, _DimShadow] | None = None + self._span_cache: dict[ + str, tuple[np.ndarray, np.ndarray, np.ndarray] + ] = {} + # One long-lived pool shared by every scan, its threads spawned + # NOW — never from inside an engine's scan callback. Creating a + # pool (and its OS threads) per scan deadlocks when the scan is + # driven from an engine executing under another thread pool + # (dask computing chunks of a lazy round-trip): thread startup + # and concurrent.futures' global shutdown lock interleave with + # the engine's callback needing the GIL. Scans that stop early + # cancel their queued loads instead of tearing the pool down. + # Like any thread state, the pool does not survive fork(); use + # from forked workers (e.g. PyTorch DataLoader) is tracked in + # https://github.com/xqlsystems/xarray-sql/issues/145. + self._pool: ThreadPoolExecutor | None = None + if self._prefetch > 1: + self._pool = ThreadPoolExecutor(max_workers=self._prefetch) + # Each pre-spawn task parks on the barrier, so no thread can + # take a second task and the executor is forced to start all + # ``prefetch`` OS threads before __init__ returns (submitting + # plain no-ops lets one idle thread absorb several of them, + # leaving the rest to spawn later inside an engine's scan + # callback — the deadlock this pre-spawn exists to prevent). + barrier = threading.Barrier(self._prefetch + 1) + spawn = [ + self._pool.submit(barrier.wait) for _ in range(self._prefetch) + ] + barrier.wait() + for f in spawn: + f.result() + # Stop the pool's threads when the dataset dies; live scans + # keep the dataset alive through their generator closures, + # so nothing in flight is cut short. The callback is bound + # to the executor, not the dataset, so the finalizer holds + # no reference cycle back to self. + weakref.finalize( + self, self._pool.shutdown, wait=False, cancel_futures=True + ) + + # ------------------------------------------------------------------ + # The consumer-facing surface + # ------------------------------------------------------------------ + + @property + def schema(self) -> pa.Schema: + return self._schema + + def scanner( + self, + columns: list[str] | None = None, + filter: pc.Expression | None = None, + batch_size: int | None = None, + **kwargs: Any, + ) -> pads.Scanner: + """Build a scanner for the requested columns and predicate. + + ``filter`` is applied exactly by the returned scanner (DuckDB + deletes the conjuncts it pushes down and trusts the source to + enforce them); chunk pruning and column selection only reduce + how much data is read to get there. ``batch_size`` caps rows per + emitted batch (Polars passes it through ``to_batches``). Extra + keyword arguments from other pyarrow-dataset consumers are + accepted and ignored. + """ + kept = None if filter is None else self._prune(filter) + blocks = ( + self._coalesced_blocks(kept) + if self._coalesce_rows + else self._blocks(kept) + ) + return self._scanner_for_blocks(blocks, columns, filter, batch_size) + + def _scanner_for_blocks( + self, + blocks: Iterator[Block] | list[Block], + columns: list[str] | None, + filter: pc.Expression | None, + batch_size: int | None = None, + ) -> pads.Scanner: + """A scanner over the given blocks; shared by dataset and fragments.""" + # ``None`` means every column (the pyarrow convention); an + # explicitly empty list is a real projection ("no payload"), not + # a request for the full schema. + proj = list(self._schema.names) if columns is None else list(columns) + scan_names = self._scan_columns(proj, filter) + scan_schema = pa.schema([self._schema.field(n) for n in scan_names]) + size = batch_size or self._batch_size + if self._geometry and GEOMETRY_COLUMN in scan_names: + batches = self._batches_with_geometry(scan_schema, blocks, size) + else: + batches = self._batch_generator(scan_schema, blocks, size) + return pads.Scanner.from_batches( + batches, schema=scan_schema, columns=proj, filter=filter + ) + + def _batches_with_geometry( + self, + scan_schema: pa.Schema, + blocks: Iterator[Block] | list[Block], + batch_size: int, + ) -> Iterator[pa.RecordBatch]: + """Emit ``scan_schema`` batches, synthesizing the geometry column. + + The pivot never materializes geometry: batches are produced with + the coordinate dims the geometry derives from, and the geometry + column is built per batch from those columns (for the native + encoding the point struct's children *are* the coordinate + arrays — a schema annotation, not a copy). + """ + assert self._geometry is not None + x_dim, y_dim = self._geometry + base_names = [n for n in scan_schema.names if n != GEOMETRY_COLUMN] + for d in (x_dim, y_dim): + if d not in base_names: + base_names.append(d) + base_schema = pa.schema([self._schema.field(n) for n in base_names]) + for batch in self._batch_generator(base_schema, blocks, batch_size): + geom = build_geometry( + self._geometry_encoding, + batch.column(base_names.index(x_dim)), + batch.column(base_names.index(y_dim)), + ) + arrays = [ + geom + if name == GEOMETRY_COLUMN + else batch.column(base_names.index(name)) + for name in scan_schema.names + ] + yield pa.RecordBatch.from_arrays(arrays, schema=scan_schema) + + def _batch_generator( + self, + scan_schema: pa.Schema, + blocks: Iterator[Block] | list[Block], + batch_size: int, + ) -> Iterator[pa.RecordBatch]: + names = list(scan_schema.names) + data_vars = [n for n in names if n in self._ds.data_vars] + # Select only the needed variables before slicing so unrequested + # variables are never loaded (dimension coords come via coords). + base = ( + self._ds[data_vars] + if data_vars + else self._ds.drop_vars(list(self._ds.data_vars)) + ) + + def load(block: Block) -> list[pa.RecordBatch]: + if self._iteration_callback is not None: + self._iteration_callback(block, names) + if not names: + # Zero-column projection: row counts are chunk + # arithmetic; no coordinate or variable data is read. + out = [] + rows = self._block_rows(block) + while rows > 0: + n = min(rows, batch_size) + out.append( + pa.table({"_": np.empty(n, np.int8)}) + .select([]) + .to_batches()[0] + ) + rows -= n + return out + return list( + iter_record_batches(base.isel(block), scan_schema, batch_size) + ) + + # Estimated pivoted bytes per row: gates admission when a + # byte budget is set, so peak memory tracks bytes in flight + # rather than block count (blocks vary in size under + # coalesce_rows). + row_width = 0 + for field in scan_schema: + try: + row_width += np.dtype(field.type.to_pandas_dtype()).itemsize + except (TypeError, NotImplementedError): + row_width += 8 + + def generate() -> Iterator[pa.RecordBatch]: + block_iter = iter(blocks) + first = next(block_iter, None) + if first is None: + return + second = next(block_iter, None) + if self._pool is None or second is None: + # Single-block scans (a lazy round-trip window that maps + # onto one source chunk) skip the pool entirely. + yield from load(first) + if second is not None: + yield from load(second) + for block in block_iter: + yield from load(block) + return + pool = self._pool + budget = self._prefetch_bytes + pending: deque = deque() + inflight = 0 + + def submit(block: Block) -> None: + nonlocal inflight + estimate = self._block_rows(block) * row_width + pending.append((pool.submit(load, block), estimate)) + inflight += estimate + + def drain_one() -> Iterator[pa.RecordBatch]: + nonlocal inflight + future, estimate = pending.popleft() + inflight -= estimate + yield from future.result() + + try: + submit(first) + submit(second) + for block in block_iter: + submit(block) + while len(pending) > 1 and ( + len(pending) >= self._prefetch + or (budget is not None and inflight > budget) + ): + yield from drain_one() + while pending: + yield from drain_one() + finally: + # Consumer may stop early (e.g. LIMIT): drop queued work + # without waiting for in-flight loads. The pool itself is + # shared across scans and stays up. + for future, _ in pending: + future.cancel() + + return generate() + + def get_fragments( + self, filter: pc.Expression | None = None + ) -> list["_XarrayFragment"]: + """One fragment per chunk of the source grid, pruned by ``filter``. + + This is how DataFusion consumes the dataset + (``SessionContext.register_dataset`` plans one partition per + fragment and scans them in parallel), and enables the Dask + pattern ``from_map(lambda f: f.to_table().to_pandas(), + ds.get_fragments())``. + """ + kept = None if filter is None else self._prune(filter) + return [_XarrayFragment(self, block) for block in self._blocks(kept)] + + def count_rows( + self, filter: pc.Expression | None = None, **kwargs: Any + ) -> int: + """Count rows, reading as little data as possible. + + Without a filter the count is pure chunk arithmetic — no I/O at + all. With a filter, chunks are split three ways: pruned chunks + contribute nothing, chunks whose coordinate ranges *prove* the + filter true contribute their exact size arithmetically, and only + the undecided boundary chunks are scanned (reading just the + columns the filter references). + """ + if not self._ds.sizes: + return int(self.scanner(columns=[], filter=filter).count_rows()) + if filter is None: + return int(np.prod([self._ds.sizes[d] for d in self._ds.dims])) + kept = self._prune(filter) + proven, boundary = self._strict_partition(kept, filter) + return proven + int( + self._scanner_for_blocks(boundary, [], filter).count_rows() + ) + + # Inherited convenience methods (to_table, head, to_batches, take) + # route through scanner() and keep working; the members below would + # touch the uninitialized native dataset. + + @property + def partition_expression(self) -> pc.Expression: + # The dataset-level guarantee: trivially true. The base class + # getter reads native state this object does not have. + return pc.scalar(True) + + def filter(self, expression: pc.Expression): + # A lazily-composed filter view is implementable; tracked in + # https://github.com/xqlsystems/xarray-sql/issues/239. + raise NotImplementedError( + "Use scanner(filter=...) or the engine's WHERE clause." + ) + + def replace_schema(self, schema: pa.Schema): + raise NotImplementedError + + # Guarded delegation for sort_by/join/join_asof (engine-level SQL + # joins already work through scanner()) is tracked in + # https://github.com/xqlsystems/xarray-sql/issues/240. + + def sort_by(self, sorting, **kwargs): + raise NotImplementedError + + def join(self, *args, **kwargs): + raise NotImplementedError + + def join_asof(self, *args, **kwargs): + raise NotImplementedError + + def __reduce__(self): + raise TypeError("XarrayPushdownDataset is not picklable.") + + # ------------------------------------------------------------------ + # Projection: which columns must be read + # ------------------------------------------------------------------ + + def _scan_columns( + self, proj: list[str], filter: pc.Expression | None + ) -> list[str]: + """Columns to read: the projection plus any the filter references. + + The consumer's column list need not include filter-only columns + (DuckDB drops pushed conjuncts from its plan and has no upstream + use for them). Rather than parsing the expression, probe it + against an empty table and grow the column set from the "no match + for field" errors until it evaluates; on anything unexpected fall + back to scanning every column, which is always correct. + """ + if filter is None: + return proj + wanted = set(proj) + for _ in range(len(self._schema.names) + 1): + probe = pa.table( + { + n: pa.array([], type=self._schema.field(n).type) + for n in self._schema.names + if n in wanted + } + ) + try: + probe.filter(filter) + except pa.lib.ArrowInvalid as exc: + match = re.search(r"FieldRef\.Name\((.*?)\)", str(exc)) + name = match.group(1) if match else None + if ( + name is not None + and name in set(self._schema.names) - wanted + ): + wanted.add(name) + continue + return list(self._schema.names) + else: + return [n for n in self._schema.names if n in wanted] + return list(self._schema.names) + + # ------------------------------------------------------------------ + # Pruning: which chunks can satisfy the predicate + # ------------------------------------------------------------------ + + def _dim_shadows(self) -> dict[str, _DimShadow]: + """One pruning index per prunable dimension, built lazily. + + Keeping one shadow per dimension (Σ n_d fragments) instead of one + per chunk (Π n_d) is what keeps this cheap for finely partitioned + datasets, and is sound: a chunk is dropped only when the full + predicate is provably false given that single dimension's range. + """ + if self._shadows is not None: + return self._shadows + shadows: dict[str, _DimShadow] = {} + for dim in self._resolved: + name = str(dim) + if name not in self._schema.names: + continue + coord = self._coord_arrays[name] + if coord.dtype.kind not in ("i", "u", "f", "M"): + continue # strings/objects/cftime: never prune this dim + try: + shadows[name] = _DimShadow( + name, self._schema, coord, self._chunk_bounds[dim] + ) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + continue # conservative: no pruning on this dim + self._shadows = shadows + return shadows + + def _prune(self, filter: pc.Expression) -> dict[str, list[int]]: + """Per-dimension chunk indices that can satisfy ``filter``. + + Satisfiability is delegated to Arrow's guarantee simplification + (see ``_DimShadow``) — no expression decoding here, and + predicates on columns a shadow knows nothing about are + conservatively kept. Dimensions without a shadow, or where every + chunk survives, are absent from the result. + """ + kept: dict[str, list[int]] = {} + for name, shadow in self._dim_shadows().items(): + indices = shadow.kept(filter) + if indices is not None: + kept[name] = indices + return kept + + # ------------------------------------------------------------------ + # Strictness: which surviving chunks satisfy the filter entirely + # ------------------------------------------------------------------ + + def _chunk_spans( + self, name: str + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Cached vectorized per-chunk ``(lo, hi, poisoned)`` for a dim.""" + cached = self._span_cache.get(name) + if cached is not None: + return cached + coord = self._coord_arrays[name] + starts = self._chunk_bounds[name][:-1] + lo = np.minimum.reduceat(coord, starts) + hi = np.maximum.reduceat(coord, starts) + if coord.dtype.kind == "f": + bad_values = np.isnan(coord) + elif coord.dtype.kind == "M": + bad_values = np.isnat(coord) + else: + bad_values = np.zeros(len(coord), dtype=bool) + bad = np.bitwise_or.reduceat(bad_values, starts) + self._span_cache[name] = (lo, hi, bad) + return self._span_cache[name] + + def _strict_partition( + self, + kept: dict[str, list[int]] | None, + filter: pc.Expression, + ) -> tuple[int, list[Block] | Iterator[Block]]: + """``(rows proven to satisfy filter, boundary blocks to scan)``. + + A cell of the surviving chunk grid with conjunctive coordinate + guarantee ``G`` satisfies ``filter`` everywhere iff + ``G ∧ ¬filter`` is unsatisfiable, and can be *dropped* entirely + iff ``G ∧ filter`` is — both decided by Arrow's guarantee + simplification. Cells are classified hierarchically: each level + buckets the surviving indices into at most + ``_STRICT_LEVEL_BUDGET`` products, proves/prunes whole buckets + at once (the prune side refines the per-dimension pruning with + cross-dimension information), and recurses only into mixed + cells — so million-chunk axes resolve in two or three levels. + Anything undecidable (NaN spans, non-numeric dims, expression + shapes the simplifier rejects) conservatively lands in the + boundary set, which the caller scans exactly. + """ + dims = list(self._resolved.keys()) + lists = {d: list(self._surviving(kept, d)) for d in dims} + if not dims or any(not v for v in lists.values()): + return 0, [] + lens = {d: np.diff(self._chunk_bounds[d]) for d in dims} + outer = self._outer_rows() + usable = { + d: str(d) in self._schema.names + and self._coord_arrays[str(d)].dtype.kind in ("i", "u", "f", "M") + for d in dims + } + proven = 0 + boundary: list[Block] = [] + + def bucket_guarantee( + d: Any, indices: np.ndarray + ) -> pc.Expression | None: + """Conjunctive [min, max] guarantee over one dimension's index + bucket, or None when unprovable (NaN spans, non-numeric dims).""" + if not usable[d]: + return None + lo, hi, bad = self._chunk_spans(str(d)) + if bad[indices].any(): + return None + field_type = self._schema.field(str(d)).type + return ( + pc.field(str(d)) >= pa.scalar(lo[indices].min(), field_type) + ) & (pc.field(str(d)) <= pa.scalar(hi[indices].max(), field_type)) + + def rows_of(cell: dict) -> int: + """Rows spanned by ``cell``: the product of its chunks' lengths + per dimension, times the rows of dimensions outside the grid.""" + rows = outer + for d in dims: + rows *= int(lens[d][np.asarray(cell[d])].sum()) + return rows + + def classify(cell: dict, depth: int) -> None: + """Prove, prune, or split one cell of the surviving chunk grid. + + A cell is a hyper-rectangle of chunk indices per dimension. + Its indices are bucketed into at most + ``_STRICT_LEVEL_BUDGET`` products; Arrow's guarantee + simplification then decides each bucket-product wholesale: + proven (every row satisfies ``filter``, counted into + ``proven`` without reading), pruned (provably empty, + dropped), or mixed (recurse). Undecidable single-chunk + cells land in ``boundary`` for exact scanning. + """ + nonlocal proven + ks = {d: len(cell[d]) for d in dims} + while int(np.prod(list(ks.values()))) > _STRICT_LEVEL_BUDGET: + widest = max(ks, key=lambda d: ks[d]) + if ks[widest] == 1: + break + ks[widest] = max(1, ks[widest] // 2) + buckets = { + d: np.array_split(np.asarray(cell[d]), ks[d]) for d in dims + } + combos = list(itertools.product(*(range(ks[d]) for d in dims))) + guarantees: list[pc.Expression | None] = [] + for combo in combos: + g: pc.Expression | None = None + complete = True + for d, b in zip(dims, combo): + gd = bucket_guarantee(d, buckets[d][b]) + if gd is None: + if usable[d]: + complete = False # NaN span: never provable + continue + g = gd if g is None else g & gd + guarantees.append(g if (g is not None and complete) else None) + decidable = [i for i, g in enumerate(guarantees) if g is not None] + satisfiable = set(decidable) + unstrict = set(decidable) + if decidable: + shadow = _guarantee_shadow( + [(str(i), guarantees[i]) for i in decidable], self._schema + ) + satisfiable = { + int(f.path) for f in shadow.get_fragments(filter=filter) + } + unstrict = { + int(f.path) for f in shadow.get_fragments(filter=~filter) + } + for i, combo in enumerate(combos): + subcell = {d: list(buckets[d][b]) for d, b in zip(dims, combo)} + if guarantees[i] is not None: + if i not in satisfiable: + continue # provably empty: cross-dim refinement + if i not in unstrict: + proven += rows_of(subcell) + continue + if all(len(v) == 1 for v in subcell.values()): + boundary.append( + self._block_for_combo( + tuple(v[0] for v in subcell.values()) + ) + ) + elif depth < _STRICT_MAX_DEPTH: + classify(subcell, depth + 1) + else: + # Depth backstop; unreachable for realistic grids. + for c in itertools.product(*subcell.values()): + boundary.append(self._block_for_combo(c)) + + try: + classify(lists, 0) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError): + return 0, self._blocks(kept) # scan every survivor, exactly + return proven, boundary + + def _block_rows(self, block: Block) -> int: + rows = 1 + for d, sl in block.items(): + size = self._ds.sizes[d] + start, stop, _ = sl.indices(size) + rows *= stop - start + return rows + + # ------------------------------------------------------------------ + # Scan: load surviving chunks, prefetching ahead of the consumer + # ------------------------------------------------------------------ + + def _surviving( + self, kept: dict[str, list[int]] | None, dim: Any + ) -> list[int] | range: + """One dim's surviving chunk indices; every chunk when unpruned.""" + return (kept or {}).get(str(dim), range(len(self._resolved[dim]))) + + def _outer_rows(self) -> int: + """Rows contributed per grid cell by dims outside the chunk grid. + + Unresolved dims span their full extent in every block, so they + multiply every block's row count uniformly. + """ + rows = 1 + for d in self._ds.dims: + if d not in self._resolved: + rows *= self._ds.sizes[d] + return rows + + def _combos( + self, kept: dict[str, list[int]] | None + ) -> Iterator[tuple[int, ...]]: + """Surviving chunk-index combinations, in grid order.""" + dims = list(self._resolved.keys()) + if not dims: + return + yield from itertools.product(*(self._surviving(kept, d) for d in dims)) + + def _block_for_combo(self, combo: tuple[int, ...]) -> Block: + block: Block = {d: slice(None) for d in self._ds.dims} + for d, i in zip(self._resolved.keys(), combo): + bounds = self._chunk_bounds[d] + block[d] = slice(int(bounds[i]), int(bounds[i + 1])) + return block + + def _blocks(self, kept: dict[str, list[int]] | None) -> Iterator[Block]: + """Yield isel-able block slices for the surviving chunk grid.""" + if not self._resolved: + yield {} + return + for combo in self._combos(kept): + yield self._block_for_combo(combo) + + def _coalesced_blocks( + self, kept: dict[str, list[int]] | None + ) -> Iterator[Block]: + """Blocks with runs of consecutive chunks merged along one dim. + + Runs are merged along the most finely chunked dimension while + the merged block stays under ``coalesce_rows`` rows. One merged + block is one ``isel`` — on Zarr sources its member chunks are + fetched by the store's own concurrent batch read instead of one + request per chunk through the prefetch pool. + """ + if not self._resolved: + yield {} + return + dims = list(self._resolved.keys()) + merge_dim = max(dims, key=lambda d: len(self._resolved[d])) + others = [d for d in dims if d != merge_dim] + ranges = {d: list(self._surviving(kept, d)) for d in dims} + merge_bounds = self._chunk_bounds[merge_dim] + outer_rows = self._outer_rows() + + def flush(prefix: tuple[int, ...], run: list[int]) -> Block: + block: Block = {d: slice(None) for d in self._ds.dims} + for d, i in zip(others, prefix): + bounds = self._chunk_bounds[d] + block[d] = slice(int(bounds[i]), int(bounds[i + 1])) + block[merge_dim] = slice( + int(merge_bounds[run[0]]), int(merge_bounds[run[-1] + 1]) + ) + return block + + coalesce_rows = self._coalesce_rows + assert coalesce_rows is not None # only reached with coalescing on + for prefix in itertools.product(*(ranges[d] for d in others)): + per_row = outer_rows + for d, i in zip(others, prefix): + bounds = self._chunk_bounds[d] + per_row *= int(bounds[i + 1] - bounds[i]) + run: list[int] = [] + run_rows = 0 + for i in ranges[merge_dim]: + rows = int(merge_bounds[i + 1] - merge_bounds[i]) * per_row + if run and ( + i != run[-1] + 1 or run_rows + rows > coalesce_rows + ): + yield flush(prefix, run) + run, run_rows = [], 0 + run.append(i) + run_rows += rows + if run: + yield flush(prefix, run) + + +class _XarrayFragment: + """One chunk of the source grid, presented as a dataset fragment. + + Fragment consumers (DataFusion's ``DatasetExec`` plans one partition + per fragment; Dask maps over them) call [scanner][xarray_sql.backends.pyarrow.XarrayPushdownDataset.scanner] with the + columns and predicate for this piece; the pushed filter is applied + row-exactly, same as the parent dataset's scanner. + """ + + def __init__(self, dataset: XarrayPushdownDataset, block: Block): + self._dataset = dataset + self._block = block + + @property + def physical_schema(self) -> pa.Schema: + return self._dataset.schema + + def scanner( + self, + schema: pa.Schema | None = None, + columns: list[str] | None = None, + filter: pc.Expression | None = None, + batch_size: int | None = None, + **kwargs: Any, + ) -> pads.Scanner: + return self._dataset._scanner_for_blocks( + [self._block], columns, filter, batch_size + ) + + def to_batches(self, **kwargs: Any) -> Iterator[pa.RecordBatch]: + return iter(self.scanner(**kwargs).to_batches()) + + def to_table(self, **kwargs: Any) -> pa.Table: + return self.scanner(**kwargs).to_table() + + def count_rows(self, **kwargs: Any) -> int: + return int(self.scanner(**kwargs).count_rows()) + + def __dask_tokenize__(self) -> tuple: + # Dask hashes from_map inputs; the parent dataset is unpicklable, + # so provide a deterministic token from the fragment's identity. + return ( + "xarray_sql._XarrayFragment", + repr(self._block), + self._dataset.schema.to_string(), + ) + + def __repr__(self) -> str: + return f"_XarrayFragment({self._block!r})" + + +def arrow_dataset( + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + prefetch: int = DEFAULT_PREFETCH, + prefetch_bytes: int | None = None, + coalesce_rows: int | None = None, + geometry: tuple[str, str] | None = None, + geometry_encoding: str = "wkb", + geometry_crs: str | None = "OGC:CRS84", +) -> XarrayPushdownDataset: + """A pushdown-capable ``pyarrow.dataset.Dataset`` view of ``ds``. + + The returned object works anywhere a pyarrow dataset does, keeping + projection pushdown and coordinate-range chunk pruning:: + + import polars as pl + lf = pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) + + import duckdb + duckdb.connect().register("t", xql.arrow_dataset(ds)) + + xql.arrow_dataset(ds).to_table(columns=["t2m"], filter=...) + + Args: + ds: An xarray Dataset. All data variables must share the same + dimensions (select a variable subset first otherwise). + chunks: Xarray-like chunks specification controlling partition + granularity. Defaults to the Dataset's existing chunks. + batch_size: Maximum rows per emitted Arrow RecordBatch. + prefetch: Chunk loads kept in flight ahead of the consumer + (memory scales with ``prefetch`` x pivoted chunk size). + prefetch_bytes: Optional cap on *estimated pivoted bytes* in + flight; admission then tracks bytes rather than block count, + which keeps peak memory steady when ``coalesce_rows`` makes + blocks large or uneven. ``prefetch`` still bounds + concurrency (thread count). + coalesce_rows: When set, merge runs of consecutive surviving + chunks along the most finely chunked dimension into single + reads of at most this many rows. Fewer, larger source + requests — the win on remote stores, where each merged read + fetches its member chunks through the store's own concurrent + batching. Memory scales with ``prefetch`` x the *merged* + block size, so size accordingly (e.g. ``8_000_000``). + geometry: ``(x_dim, y_dim)`` coordinate dims to derive a + ``geometry`` point column from (see + [xarray_sql.geometry][]). With the default ``"wkb"`` + encoding, DuckDB (spatial loaded) sees a native ``GEOMETRY`` + column, so ``ST_Within(geometry, ...)`` works directly. + geometry_encoding: ``"wkb"`` (default; DuckDB-consumable) or + ``"point"`` (GeoArrow native separated coordinates — the + struct children are the coordinate arrays; for GeoPandas, + lonboard, geoarrow-rs consumers). + geometry_crs: CRS tag carried in the extension metadata. + Defaults to ``OGC:CRS84`` (plain longitude/latitude); pass + ``None`` to omit, or an authority code / PROJJSON string. + + Returns: + An [XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset]. + """ + return XarrayPushdownDataset( + ds, + chunks, + batch_size=batch_size, + prefetch=prefetch, + prefetch_bytes=prefetch_bytes, + coalesce_rows=coalesce_rows, + geometry=geometry, + geometry_encoding=geometry_encoding, + geometry_crs=geometry_crs, + ) diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py new file mode 100644 index 00000000..1f7c49ce --- /dev/null +++ b/xarray_sql/cftime.py @@ -0,0 +1,259 @@ +"""Bridge between cftime calendars and Arrow/DataFusion types. + +cftime (https://unidata.github.io/cftime/) provides datetime objects for +calendars used in climate science — noleap, 360-day, all-leap, julian, etc. +Arrow and DataFusion have no native concept of non-Gregorian calendars, so +this module handles the conversion in two tiers: + +* **Gregorian-like calendars** (standard, gregorian, proleptic_gregorian, + noleap/365_day, all_leap/366_day): mapped to ``pa.timestamp('us')`` so + that string-based SQL filters like ``WHERE time > '1980-01-01'`` work + naturally. Microsecond resolution avoids the 1678–2262 overflow of + nanoseconds while preserving sub-second precision. + +* **Non-Gregorian calendars** (360_day, julian): mapped to ``pa.int64()`` + with ``xarray:units`` and ``xarray:calendar`` metadata on the Arrow field. + This preserves the original CF-convention encoding losslessly. A + ``cftime()`` DataFusion UDF (registered automatically by + ``XarrayContext.from_dataset``) provides ergonomic SQL filtering. +""" + +from __future__ import annotations + +import numpy as np +import pyarrow as pa +import xarray as xr + + +# --------------------------------------------------------------------------- +# Calendar classification +# --------------------------------------------------------------------------- + +GREGORIAN_LIKE_CALENDARS: frozenset[str] = frozenset( + { + "standard", + "gregorian", + "proleptic_gregorian", + "noleap", + "365_day", + "all_leap", + "366_day", + } +) +"""Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``.""" + +DEFAULT_UNITS: str = "microseconds since 1970-01-01T00:00:00" +"""Default CF-convention units when no encoding is available on the coordinate. + +Microseconds give sub-second precision and fit int64 for ±292 k years. +""" + + +def is_gregorian_like(calendar: str) -> bool: + """Return True if *calendar* is close enough to Gregorian for ``pa.timestamp``.""" + return calendar in GREGORIAN_LIKE_CALENDARS + + +# --------------------------------------------------------------------------- +# Detection helpers (avoid materializing Dask/Zarr data where possible) +# --------------------------------------------------------------------------- + + +def is_cftime(values: np.ndarray) -> bool: + """Check if a numpy array contains cftime datetime objects.""" + try: + import cftime + + if values.dtype == np.dtype("O") and len(values) > 0: + sample = values.ravel()[0] + return isinstance(sample, cftime.datetime) + except ImportError: + pass + return False + + +def is_cftime_index(ds: xr.Dataset, coord_name: str) -> bool: + """Check if a coordinate uses a ``CFTimeIndex`` without materializing data.""" + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex + + return isinstance(idx, CFTimeIndex) + except (ImportError, AttributeError): + pass + return False + + +def calendar(ds: xr.Dataset, coord_name: str) -> str | None: + """Return the calendar name for a cftime coordinate, or ``None``. + + Checks the xarray index first (no data materialization), then falls + back to inspecting element 0 of the coordinate values. + """ + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex + + if isinstance(idx, CFTimeIndex): + return str(idx.calendar) # type: ignore[attr-defined] + except (ImportError, AttributeError): + pass + try: + values = ds.coords[coord_name].values + if is_cftime(values): + return str(values.ravel()[0].calendar) + except (AttributeError, KeyError): + pass + return None + + +def encoding(ds: xr.Dataset, coord_name: str) -> tuple[str, str]: + """Return ``(units, calendar)`` for a cftime coordinate. + + Reads xarray ``.encoding`` metadata (from the originating NetCDF file) + first, falling back to [DEFAULT_UNITS][xarray_sql.cftime.DEFAULT_UNITS]. + """ + cal = calendar(ds, coord_name) or "standard" + enc = ds.coords[coord_name].encoding + units = enc.get("units", DEFAULT_UNITS) + return units, cal + + +# --------------------------------------------------------------------------- +# Numeric conversion +# --------------------------------------------------------------------------- + + +def to_microseconds(values) -> np.ndarray: + """Convert cftime objects to int64 microseconds since Unix epoch. + + Used for Gregorian-like calendars. Vectorised via ``cftime.date2num`` + (implemented in C). + """ + import cftime as _cftime + + us = _cftime.date2num( + values.ravel(), + units=DEFAULT_UNITS, + calendar=values.ravel()[0].calendar, + ) + return np.asarray(us, dtype=np.float64).astype(np.int64) + + +def to_offsets(values, units: str, cal: str) -> np.ndarray: + """Convert cftime objects to int64 offsets in the given *units*/*calendar*. + + Used for non-Gregorian calendars where data is stored as ``pa.int64()``. + """ + import cftime as _cftime + + raw = _cftime.date2num(values.ravel(), units=units, calendar=cal) + return np.asarray(raw, dtype=np.float64).astype(np.int64) + + +def convert_for_field(values, field: pa.Field) -> np.ndarray: + """Convert cftime values to the numeric type dictated by *field*. + + Reads ``xarray:calendar`` and ``xarray:units`` from the field's Arrow + metadata to choose between the timestamp path and the integer-offset path. + """ + meta = field.metadata or {} + cal = meta.get(b"xarray:calendar", b"standard").decode() + units = meta.get(b"xarray:units", DEFAULT_UNITS.encode()).decode() + if is_gregorian_like(cal): + return to_microseconds(values) + return to_offsets(values, units, cal) + + +# --------------------------------------------------------------------------- +# Partition pruning helpers +# --------------------------------------------------------------------------- + + +def partition_bounds( + values, +) -> tuple[int, int, str] | None: + """Return ``(min, max, dtype_tag)`` for a cftime coordinate slice. + + Gregorian-like calendars return nanosecond bounds tagged + ``"timestamp_ns"`` (compatible with ``ScalarBound::TimestampNanos`` + in the Rust pruning layer). Non-Gregorian calendars return int64 + offsets tagged ``"int64"``. + + Returns ``None`` when the nanosecond bound falls outside the int64 range + (e.g. paleoclimate dates before ~1678), signalling the caller to skip + pruning for that dimension rather than emit a bound the Rust layer would + reject. + """ + cal = values.ravel()[0].calendar + if is_gregorian_like(cal): + us = to_microseconds(values) + lo, hi = int(us.min()) * 1_000, int(us.max()) * 1_000 + int64 = np.iinfo(np.int64) + if lo < int64.min or hi > int64.max: + return None + return lo, hi, "timestamp_ns" + offsets = to_offsets(values, DEFAULT_UNITS, cal) + return int(offsets.min()), int(offsets.max()), "int64" + + +# --------------------------------------------------------------------------- +# Arrow schema helpers +# --------------------------------------------------------------------------- + + +def arrow_field(name: str, units: str, cal: str) -> pa.Field: + """Build a ``pa.Field`` for a cftime coordinate. + + Gregorian-like → ``pa.timestamp('us')``; non-Gregorian → ``pa.int64()``. + Both carry ``xarray:calendar`` and ``xarray:units`` metadata for + round-trip fidelity. + """ + meta = { + b"xarray:calendar": cal.encode(), + b"xarray:units": units.encode(), + } + if is_gregorian_like(cal): + return pa.field(name, pa.timestamp("us"), metadata=meta) + return pa.field(name, pa.int64(), metadata=meta) + + +# --------------------------------------------------------------------------- +# DataFusion UDF +# --------------------------------------------------------------------------- + + +def make_cftime_udf(units: str, calendar: str): + """Create a DataFusion scalar UDF that converts date strings to int64 offsets. + + This enables ergonomic SQL filtering on non-Gregorian cftime columns:: + + SELECT * FROM ds360 WHERE time > cftime('0500-01-01') + + The UDF parses the input string as a cftime datetime in the given + calendar system and returns the corresponding int64 offset in the + specified units. + """ + import cftime as _cftime + from datafusion import udf + + def _cftime_scalar(date_strings: pa.Array) -> pa.Array: + results: list[int | None] = [] + for s in date_strings.to_pylist(): + if s is None: + results.append(None) + continue + dt = _cftime.datetime.strptime(s, "%Y-%m-%d", calendar=calendar) + val = _cftime.date2num(dt, units=units, calendar=calendar) + results.append(int(val)) + return pa.array(results, type=pa.int64()) + + return udf( + _cftime_scalar, + [pa.utf8()], + pa.int64(), + "immutable", + "cftime", + ) diff --git a/xarray_sql/core.py b/xarray_sql/core.py index afbd2eca..6ff06b6a 100644 --- a/xarray_sql/core.py +++ b/xarray_sql/core.py @@ -1,48 +1,49 @@ import itertools -import typing as t +from collections.abc import Iterator +from typing import Any import numpy as np import xarray as xr -Row = t.List[t.Any] +Row = list[Any] # deprecated -def get_columns(ds: xr.Dataset) -> t.List[str]: - return list(ds.sizes.keys()) + list(ds.data_vars.keys()) +def get_columns(ds: xr.Dataset) -> list[str]: + return list(ds.sizes.keys()) + list(ds.data_vars.keys()) # Deprecated -def unravel(ds: xr.Dataset) -> t.Iterator[Row]: - dim_keys, dim_vals = zip(*ds.sizes.items()) +def unravel(ds: xr.Dataset) -> Iterator[Row]: + dim_keys, dim_vals = zip(*ds.sizes.items()) - for idx in itertools.product(*(range(d) for d in dim_vals)): - coord_idx = dict(zip(dim_keys, idx)) - data = ds.isel(coord_idx) - coord_data = [ds.coords[v][coord_idx[v]] for v in dim_keys] - row = [v.values for v in coord_data + list(data.data_vars.values())] - yield row + for idx in itertools.product(*(range(d) for d in dim_vals)): + coord_idx = dict(zip(dim_keys, idx)) + data = ds.isel(coord_idx) + coord_data = [ds.coords[v][coord_idx[v]] for v in dim_keys] + row = [v.values for v in coord_data + list(data.data_vars.values())] + yield row # Deprecated def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: - """Unravel with unbounded memory (as a NumPy Array).""" - dim_keys, dim_vals = zip(*ds.sizes.items()) - columns = get_columns(ds) + """Unravel with unbounded memory (as a NumPy Array).""" + dim_keys, dim_vals = zip(*ds.sizes.items()) + columns = get_columns(ds) - N = np.prod([d for d in dim_vals]) + N = np.prod([d for d in dim_vals]) - out = np.recarray((N,), dtype=[(c, ds[c].dtype) for c in columns]) + out = np.recarray((N,), dtype=[(c, ds[c].dtype) for c in columns]) - for name, da in ds.items(): - out[name] = da.values.ravel() + for name, da in ds.items(): + out[name] = da.values.ravel() - prod_vals = (ds.coords[k].values for k in dim_keys) - coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( - -1, len(dim_keys) - ) + prod_vals = (ds.coords[k].values for k in dim_keys) + coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( + -1, len(dim_keys) + ) - for i, d in enumerate(dim_keys): - out[d] = coords[:, i] + for i, d in enumerate(dim_keys): + out[d] = coords[:, i] - return out + return out diff --git a/xarray_sql/df.py b/xarray_sql/df.py index c799626f..83879f77 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,172 +1,656 @@ import itertools -import typing as t -import warnings +from collections import defaultdict +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping +from typing import Any import numpy as np import pandas as pd import pyarrow as pa import xarray as xr -from datafusion.context import ArrowStreamExportable -Block = t.Dict[t.Hashable, slice] -Chunks = t.Optional[t.Dict[str, int]] +from . import cftime as cft + + +Block = dict[Hashable, slice] +Chunks = dict[str, int] | None # Borrowed from Xarray def _get_chunk_slicer( - dim: t.Hashable, chunk_index: t.Mapping, chunk_bounds: t.Mapping + dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping ): - if dim in chunk_index: - which_chunk = chunk_index[dim] - return slice( - chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] + if dim in chunk_index: + which_chunk = chunk_index[dim] + return slice( + chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] + ) + return slice(None) + + +def compute_chunks( + ds: xr.Dataset, chunks: dict[str, int] +) -> dict[Hashable, tuple[int, ...]]: + """Per-dim chunk-size tuples matching ``ds.chunk(chunks).chunks``. + + Pure arithmetic replacement for the dask rechunk round-trip; dask's + ``.chunk()`` eagerly builds a task graph, which dominates + ``block_slices()`` cost on large datasets. + """ + existing = dict(ds.chunks) if ds.chunks else {} + result: dict[Hashable, tuple[int, ...]] = {} + for dim in ds.dims: + size = ds.sizes[dim] + if dim in chunks: + cs = chunks[dim] + if cs <= 0 or cs >= size: + result[dim] = (size,) + else: + n_full, rem = divmod(size, cs) + result[dim] = (cs,) * n_full + ((rem,) if rem else ()) + elif dim in existing: + result[dim] = tuple(existing[dim]) + else: + result[dim] = (size,) + return result + + +def resolve_chunks( + ds: xr.Dataset, chunks: Chunks +) -> Mapping[Hashable, tuple[int, ...]]: + """Normalise the user's ``chunks`` argument to per-dim size tuples. + + Filters out keys for dims this dataset doesn't have (sub-datasets in a + heterogeneous group need not contain every dimension named in the + spec), then either rechunks arithmetically via ``compute_chunks`` or + falls back to the dataset's existing dask chunks. + + Returns an empty mapping for scalar datasets; callers should treat that + as "one block covering everything". + """ + if chunks is not None: + chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes} + if chunks: + return compute_chunks(ds, chunks) + return {d: tuple(c) for d, c in ds.chunks.items()} + + +def _ensure_default_indexes(ds: xr.Dataset) -> xr.Dataset: + """Attach a default integer index coordinate to every dimension lacking one. + + xarray allows "dimensions without coordinates"; these are absent from + ``ds.coords``, so they are dropped from the SQL schema and, once a block is + sliced out with ``isel``, their position is synthesized *relative to the + block* (restarting at 0 in every partition). Materialising an explicit + ``arange`` index up front turns them into ordinary dimension coordinates, so + they appear as columns and carry their absolute position through chunked + reads. Datasets whose dimensions already have coordinates are returned + unchanged. + """ + missing = { + dim: np.arange(ds.sizes[dim]) for dim in ds.dims if dim not in ds.coords + } + return ds.assign_coords(missing) if missing else ds + + +def _block_slices_from_resolved( + ds: xr.Dataset, resolved: Mapping[Hashable, tuple[int, ...]] +) -> Iterator[Block]: + """Emit blocks given pre-resolved per-dim chunk tuples.""" + if not resolved: + # No chunkable dimensions. A dimensionless dataset (e.g. scalar + # metadata variables) is a single block; a dataset that has + # dimensions but no chunking is a user error. + assert not ds.sizes, ( + "Dataset `ds` must be chunked or `chunks` must be provided." + ) + yield {} + return + + chunk_bounds = { + dim: np.cumsum((0,) + tuple(c)) for dim, c in resolved.items() + } + ichunk = {dim: range(len(tuple(c))) for dim, c in resolved.items()} + ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. + chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) + yield from ( + { + dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) + for dim in ds.dims + } + for chunk_index in chunk_idxs ) - return slice(None) # Adapted from Xarray `map_blocks` implementation. -def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]: - """Compute block slices for a chunked Dataset.""" - if chunks is not None: - for_chunking = ds.copy(data=None, deep=False).chunk(chunks) - chunks = for_chunking.chunks - del for_chunking - else: - chunks = ds.chunks - - assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided." - - # chunks is Dict[str, Tuple[int, ...]] from xarray - chunk_bounds = { - dim: np.cumsum((0,) + tuple(c)) # type: ignore[arg-type] - for dim, c in chunks.items() - } - ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()} # type: ignore[arg-type] - ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. - chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) - blocks = ( - { - dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) - for dim in ds.dims - } - for chunk_index in chunk_idxs - ) - yield from blocks - - -def explode(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[xr.Dataset]: - """Explodes a dataset into its chunks.""" - yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks)) +def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: + """Compute block slices for a chunked Dataset.""" + yield from _block_slices_from_resolved(ds, resolve_chunks(ds, chunks)) + + +def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]: + """Explodes a dataset into its chunks.""" + yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks)) + + +def group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]: + """Group a Dataset's data variables by their exact dimension tuple. + + Variables that share dimensions can share a table; each distinct + dimension tuple becomes its own table when a mixed-dimension Dataset + is registered:: + + ("time", "lat", "lon"): ["temperature_2m", "wind_speed"], + ("time", "lat", "lon", "level"): ["pressure", "humidity"] + """ + groups = defaultdict(list) + for var_name, var in ds.data_vars.items(): + dims = var.dims + groups[dims].append(var_name) + return groups def _block_len(block: Block) -> int: - return int(np.prod([v.stop - v.start for v in block.values()])) + return int(np.prod([v.stop - v.start for v in block.values()])) def from_map_batched( - func: t.Callable[..., pd.DataFrame], - *iterables, - args: t.Optional[t.Tuple] = None, + func: Callable[..., pd.DataFrame], + *iterables: tuple[Any, ...], + args: tuple | None = None, schema: pa.Schema = None, - **kwargs, + **kwargs: dict[str, Any], ) -> pa.RecordBatchReader: - """Create a PyArrow RecordBatchReader by mapping a function over iterables. + """Create a PyArrow RecordBatchReader by mapping a function over iterables. - This is equivalent to dask's from_map but returns a PyArrow - RecordBatchReader that can be used with DataFusion. It iterates over - RecordBatches which are created via the `func` one-at-a-time. + This is equivalent to dask's from_map but returns a PyArrow + RecordBatchReader that can be used with DataFusion. It iterates over + RecordBatches which are created via the `func` one-at-a-time. - Args: - func: Function to apply to each element of the iterables. Currently, the function - must return a Pandas DataFrame. - *iterables: Iterable objects to map the function over. - schema: Optional schema needed for the RecordBatchReader. - args: Additional positional arguments to pass to func. - **kwargs: Additional keyword arguments to pass to func. + Args: + func: Function to apply to each element of the iterables. Currently, the + function must return a Pandas DataFrame. + *iterables: Iterable objects to map the function over. + schema: Optional schema needed for the RecordBatchReader. + args: Additional positional arguments to pass to func. + **kwargs: Additional keyword arguments to pass to func. - Returns: - A PyArrow RecordBatchReader containing the stream of RecordBatches. - """ - if args is None: - args = () + Returns: + A PyArrow RecordBatchReader containing the stream of RecordBatches. + """ + if args is None: + args = () - def map_batches(): - for items in zip(*iterables): - df = func(*items, *args, **kwargs) - yield pa.RecordBatch.from_pandas(df, schema=schema) + def map_batches(): + for items in zip(*iterables): + df = func(*items, *args, **kwargs) + yield pa.RecordBatch.from_pandas(df, schema=schema) - return pa.RecordBatchReader.from_batches(schema, map_batches()) + return pa.RecordBatchReader.from_batches(schema, map_batches()) def from_map( - func: t.Callable, *iterables, args: t.Optional[t.Tuple] = None, **kwargs + func: Callable, + *iterables: tuple[Any, ...], + args: tuple | None = None, + **kwargs: dict[str, Any], ) -> pa.Table: - """Create a PyArrow Table by mapping a function over iterables. - - This is equivalent to dask's from_map but returns a PyArrow Table - that can be used with DataFusion instead of a Dask DataFrame. - - Args: - func: Function to apply to each element of the iterables. - *iterables: Iterable objects to map the function over. - args: Additional positional arguments to pass to func. - **kwargs: Additional keyword arguments to pass to func. - - Returns: - A PyArrow Table containing the concatenated results. - """ - if args is None: - args = () - - # Apply the function to each combination of iterable elements - results = [] - for items in zip(*iterables) if len(iterables) > 1 else iterables[0]: - if isinstance(items, tuple): - result = func(*items, *args, **kwargs) - else: - result = func(items, *args, **kwargs) - - # Convert result to PyArrow Table - if isinstance(result, pd.DataFrame): - pa_table = pa.Table.from_pandas(result) - elif isinstance(result, pa.Table): - pa_table = result - else: - # Try to convert to pandas first, then to PyArrow - try: - df = pd.DataFrame(result) - pa_table = pa.Table.from_pandas(df) - except Exception as e: - raise ValueError( - f"Cannot convert function result to PyArrow Table: {e}" - ) - - results.append(pa_table) - - # Concatenate all results - if not results: - raise ValueError("No results to concatenate") - - return pa.concat_tables(results) + """Create a PyArrow Table by mapping a function over iterables. + + This is equivalent to dask's from_map but returns a PyArrow Table + that can be used with DataFusion instead of a Dask DataFrame. + + Args: + func: Function to apply to each element of the iterables. + *iterables: Iterable objects to map the function over. + args: Additional positional arguments to pass to func. + **kwargs: Additional keyword arguments to pass to func. + + Returns: + A PyArrow Table containing the concatenated results. + """ + if args is None: + args = () + + # Apply the function to each combination of iterable elements + results = [] + for items in zip(*iterables) if len(iterables) > 1 else iterables[0]: + if isinstance(items, tuple): + result = func(*items, *args, **kwargs) + else: + result = func(items, *args, **kwargs) + + # Convert result to PyArrow Table + if isinstance(result, pd.DataFrame): + pa_table = pa.Table.from_pandas(result) + elif isinstance(result, pa.Table): + pa_table = result + else: + # Try to convert to pandas first, then to PyArrow + try: + df = pd.DataFrame(result) + pa_table = pa.Table.from_pandas(df) + except Exception as e: + raise ValueError( + f"Cannot convert function result to PyArrow Table: {e}" + ) + + results.append(pa_table) + + # Concatenate all results + if not results: + raise ValueError("No results to concatenate") + + return pa.concat_tables(results) def pivot(ds: xr.Dataset) -> pd.DataFrame: - """Converts an xarray Dataset to a pandas DataFrame.""" - return ds.to_dataframe().reset_index() # type: ignore[no-any-return] - - -def _parse_schema(ds) -> pa.Schema: - """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.""" - columns = [] - - for coord_name, coord_var in ds.coords.items(): - # Only include dimension coordinates - if coord_name in ds.dims: - pa_type = pa.from_numpy_dtype(coord_var.dtype) - columns.append(pa.field(coord_name, pa_type)) - - for var_name, var in ds.data_vars.items(): - pa_type = pa.from_numpy_dtype(var.dtype) - columns.append(pa.field(var_name, pa_type)) - - return pa.schema(columns) + """Converts an xarray Dataset to a pandas DataFrame.""" + return ds.to_dataframe().reset_index() # type: ignore[no-any-return] + + +def dataset_to_record_batch( + ds: xr.Dataset, schema: pa.Schema +) -> pa.RecordBatch: + """Convert an xarray Dataset partition to an Arrow RecordBatch. + + Builds the RecordBatch directly from numpy arrays, bypassing the pandas + round-trip (to_dataframe → reset_index → from_pandas) used by pivot(). + For large partitions this reduces peak memory from ~5× to ~2× the + partition size. + + Dimension coordinates are broadcast to the full partition shape and + ravelled. np.broadcast_to() is zero-copy; the ravel() forces one copy + per coordinate (unavoidable, since broadcast arrays are non-contiguous). + Data variable arrays are ravelled in-place — a zero-copy view when the + underlying array is already C-contiguous (the common case for numpy-backed + xarray datasets). + + Args: + ds: A partition-sized xarray Dataset (already sliced via isel). + schema: The Arrow schema for the output, as produced by _parse_schema. + Column order in the output matches schema field order. + + Returns: + A RecordBatch with one column per dimension coordinate and data + variable, in schema order. + """ + # Use the data variable's dimension order as canonical so coordinate + # broadcasts and data variable ravels use the same layout. All data + # variables are validated to share the same dims tuple. + if ds.data_vars: + first_var = next(iter(ds.data_vars.values())) + dim_names = list(first_var.dims) + shape = first_var.shape + else: + dim_names = list(ds.sizes.keys()) + shape = tuple(ds.sizes[d] for d in dim_names) + + arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + # Broadcast 1-D coordinate to the full N-D partition shape, then ravel. + axis = dim_names.index(name) + coord = ds.coords[name].values + if cft.is_cftime(coord): + coord = cft.convert_for_field(coord, field) + reshape = [1] * len(shape) + reshape[axis] = coord.shape[0] + arr = np.broadcast_to(coord.reshape(reshape), shape).ravel() + arrays.append(pa.array(arr, type=field.type)) + else: + # Data variable: ravel to 1-D (zero-copy for C-contiguous arrays). + raw = ds[name].values.ravel() + if cft.is_cftime(ds[name].values): + raw = cft.convert_for_field(ds[name].values, field) + + # from_pandas=True maps NaN → Arrow null inside the C++ copy kernel, + # so SQL aggregates (MAX, MIN, AVG) skip missing values correctly. + arrays.append(pa.array(raw, type=field.type, from_pandas=True)) + + return pa.RecordBatch.from_arrays(arrays, schema=schema) + + +DEFAULT_BATCH_SIZE: int = 65_536 +"""Default number of rows per emitted Arrow RecordBatch. + +64 K rows balances DataFusion pipeline depth against per-batch overhead. +""" + +_FULL_PIVOT_MAX_ROWS: int = 8_388_608 +"""Row cap for the whole-partition coordinate fast path in +iter_record_batches. + +Below this, coordinate columns are materialised for the full partition +with repeat/tile (sequential writes, ~3x faster than per-batch index +arithmetic) and batches are zero-copy slices; the cost is holding every +coordinate column of the partition in memory at once (rows x 8 bytes x +n_dims). Above it — e.g. single-time-step reanalysis partitions with +tens of millions of rows — the per-batch path keeps peak memory at +O(batch_size) per coordinate instead. +""" + + +def _as_single_array(values, type: pa.DataType, *, from_pandas: bool = False): + """``pa.array`` that always returns a contiguous ``pa.Array``. + + ``pa.array`` may return a ``ChunkedArray`` instead of an ``Array`` for + large inputs (observed for numpy fixed-width unicode columns of a few + million rows — e.g. a string dimension coordinate tiled across a full + partition). ``RecordBatch.from_arrays`` rejects chunked input, so + flatten it back to one contiguous array. + """ + arr = pa.array(values, type=type, from_pandas=from_pandas) + if isinstance(arr, pa.ChunkedArray): + arr = arr.combine_chunks() + return arr + + +def iter_record_batches( + ds: xr.Dataset, + schema: pa.Schema, + batch_size: int = DEFAULT_BATCH_SIZE, +) -> Iterator[pa.RecordBatch]: + """Yield RecordBatches of at most *batch_size* rows from a partition Dataset. + + Unlike `dataset_to_record_batch`, which materialises the entire + partition as one batch, this generator emits smaller batches so that + DataFusion can begin filtering and aggregating before the full partition + is loaded. Peak memory per batch is O(batch_size) for coordinate columns + and O(partition_size) for data-variable columns (which must be loaded in + full from storage). + + Coordinate values are computed per batch via strided index arithmetic — + no broadcast array spanning the whole partition is ever allocated. Data + variable flat arrays are loaded once (triggering any remote I/O) and then + sliced as zero-copy views for each batch. + + Args: + ds: A partition-sized xarray Dataset (already sliced via isel). + schema: The Arrow schema for the output, as produced by _parse_schema. + batch_size: Maximum number of rows per yielded RecordBatch. + + Yields: + RecordBatches in schema column order, covering all rows of the + partition exactly once. + """ + if ds.data_vars: + first_var = next(iter(ds.data_vars.values())) + dim_names = list(first_var.dims) + shape = first_var.shape + else: + dim_names = list(ds.sizes.keys()) + shape = tuple(ds.sizes[d] for d in dim_names) + + total_rows = int(np.prod(shape)) + + # Preload small 1-D coordinate arrays (negligible memory). + # Convert cftime objects to numeric values matching the schema type. + # Projected scans may omit dimension columns from the schema; those + # dims still shape the iteration but never emit a column. + coord_values = {} + schema_names = set(schema.names) + for name in dim_names: + # A dim the projection dropped (e.g. time under GROUP BY level) is never + # read in the batch loop below, which only iterates the schema's fields. + # Skip it so schema.field(name) is not called for a projected-away name + # (it raises for cftime coords, which take the convert_for_field path). + if name not in schema_names: + continue + vals = ds.coords[name].values + if cft.is_cftime(vals): + coord_values[name] = cft.convert_for_field(vals, schema.field(name)) + else: + coord_values[name] = vals + + # C-order stride for each dimension: stride[k] = prod(shape[k+1:]). + # Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k]. + strides = [int(np.prod(shape[k + 1 :])) for k in range(len(shape))] + + # Load data-variable arrays fully (triggers Dask/Zarr compute once). + # ravel() is a zero-copy view for C-contiguous arrays. + data_arrays = {} + for field in schema: + if field.name not in ds.dims: + raw = ds[field.name].values + if cft.is_cftime(raw): + data_arrays[field.name] = cft.convert_for_field(raw, field) + else: + data_arrays[field.name] = raw.ravel() + + if 0 < total_rows <= _FULL_PIVOT_MAX_ROWS: + # Fast path: build each coordinate column once for the whole + # partition. In C order, dim k's flat column is its coord values + # each repeated prod(shape[k+1:]) times, with that pattern tiled + # prod(shape[:k]) times — two sequential-write kernels, much + # faster than per-batch division/modulo plus gather. Batches are + # then zero-copy slices of the full-partition Arrow arrays. + full_arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + k = dim_names.index(name) + outer = int(np.prod(shape[:k])) + col = np.repeat(coord_values[name], strides[k]) + if outer > 1: + col = np.tile(col, outer) + full_arrays.append(_as_single_array(col, field.type)) + else: + full_arrays.append( + _as_single_array( + data_arrays[name], field.type, from_pandas=True + ) + ) + for row_start in range(0, total_rows, batch_size): + yield pa.RecordBatch.from_arrays( + [a.slice(row_start, batch_size) for a in full_arrays], + schema=schema, + ) + return + + for row_start in range(0, total_rows, batch_size): + row_end = min(row_start + batch_size, total_rows) + row_idx = np.arange(row_start, row_end) + + arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + k = dim_names.index(name) + coord_idx = (row_idx // strides[k]) % shape[k] + arrays.append( + _as_single_array(coord_values[name][coord_idx], field.type) + ) + else: + arrays.append( + _as_single_array( + data_arrays[name][row_start:row_end], + field.type, + from_pandas=True, + ) + ) + + yield pa.RecordBatch.from_arrays(arrays, schema=schema) + + +def _arrow_type_for_object(values: np.ndarray) -> pa.DataType: + """Infer an Arrow type for a non-cftime object-dtype array. + + ``pa.from_numpy_dtype`` cannot map numpy object dtype, so let pyarrow infer + the type from the data instead: strings become ``pa.string()``, bytes + ``pa.binary()``, and other representable Python scalars their Arrow + equivalent. An all-null array stays ``pa.null()``, and a column mixing + incompatible types (e.g. str and int) raises, surfacing a clear error + rather than a silent coercion. Object-dtype arrays are never Dask/Zarr + backed, so this triggers no remote I/O. + """ + return pa.array(np.asarray(values).ravel()).type + + +def _parse_schema(ds: xr.Dataset) -> pa.Schema: + """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. + + Only *dimension coordinates* become dimension columns, so a dimension + without a coordinate would be dropped. Callers must run the Dataset through + ``_ensure_default_indexes`` first (the readers do) so every dimension + has a coordinate and appears as a column. + + Uses the xarray index type to detect cftime coordinates without + materializing their data — important for Dask/Zarr-backed datasets + where .values would trigger eager computation. + + cftime coordinates are mapped to one of two Arrow types: + + * **Gregorian-like calendars** (standard, noleap, all_leap, etc.): + ``pa.timestamp('us')`` so string-based SQL filters work naturally. + * **Non-Gregorian calendars** (360_day, julian): + ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata + on the field, preserving lossless CF-convention encoding. + """ + columns = [] + + for coord_name, coord_var in ds.coords.items(): + # Only include dimension coordinates + if coord_name in ds.dims: + if cft.is_cftime_index(ds, coord_name): + units, calendar = cft.encoding(ds, coord_name) + columns.append(cft.arrow_field(coord_name, units, calendar)) + elif coord_var.dtype == np.dtype("O"): + # Object dtype that isn't cftime (e.g. string station names). + arrow_type = _arrow_type_for_object(coord_var.values) + columns.append(pa.field(coord_name, arrow_type)) + else: + pa_type = pa.from_numpy_dtype(coord_var.dtype) + columns.append(pa.field(coord_name, pa_type)) + + for var_name, var in ds.data_vars.items(): + # An object-dtype data variable may hold cftime objects (encode it like + # a cftime coordinate) or strings/other Python scalars (infer the Arrow + # type from the data). The dtype check keeps the common numeric path off + # the object branch. + if var.dtype == np.dtype("O"): + if cft.is_cftime(var.values): + # Encode with the same units/calendar as a cftime coordinate. + cal = var.values.ravel()[0].calendar + columns.append( + cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal) + ) + else: + arrow_type = _arrow_type_for_object(var.values) + columns.append(pa.field(var_name, arrow_type)) + else: + pa_type = pa.from_numpy_dtype(var.dtype) + columns.append(pa.field(var_name, pa_type)) + + return pa.schema(columns) + + +# Type alias for partition metadata: maps dimension name to (min, max, dtype_str) values +PartitionBounds = dict[str, tuple[Any, Any, str]] + + +def _block_metadata( + coord_arrays: dict, + block: Block, + dims: Iterable[Hashable] | None = None, +) -> PartitionBounds: + """Compute min/max coordinate values for a single partition block. + + Args: + coord_arrays: Pre-materialised coordinate arrays keyed by dimension name + string. Hoist this outside any loop to avoid repeated remote I/O + for Zarr-backed datasets. + block: A single block slice dict from block_slices(). + dims: Optional restriction to a subset of dims to compute. Used by + ``read_xarray_table`` to skip unchunked dims whose bounds are + constant across all partitions and have been precomputed once. + Defaults to all dims present in ``block``. + + Returns: + Dict mapping dimension name to (min_value, max_value, dtype_str). + Dimensions with an empty slice are omitted; the Rust pruning logic + treats missing dimensions conservatively (never prunes on them). + """ + items = ((d, block[d]) for d in dims) if dims is not None else block.items() + ranges: PartitionBounds = {} + for dim, slc in items: + coord_values = coord_arrays[str(dim)][slc] + if len(coord_values) == 0: + continue + # cftime coordinates are object dtype but carry their own bound + # encoding, so they must be handled before the string/object skip + # below (otherwise pruning is silently disabled for them). + # partition_bounds returns None when the bound overflows int64. + if cft.is_cftime(coord_values): + bounds = cft.partition_bounds(coord_values) + if bounds is not None: + ranges[str(dim)] = bounds + continue + # String/object dtypes are not representable as ScalarBound + # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not + # support them. Skip so pruning treats the dimension conservatively. + if coord_values.dtype.kind in ("U", "S", "O"): + continue + + # Use actual min/max rather than first/last so that non-monotonic + # coordinate axes (e.g. descending latitude 90→-90) are handled + # correctly. np.min/max work for both numeric and datetime64 arrays. + min_val = coord_values.min() + max_val = coord_values.max() + + if isinstance(min_val, (np.datetime64, pd.Timestamp)): + # The Rust pruning layer only accepts int64 nanosecond bounds + # (ScalarBound::TimestampNanos). Dates outside the + # datetime64[ns] range (pre-1678 / post-2262) cannot be + # represented, so skip pruning for this dimension rather than + # raising -- registration still succeeds and the Rust pruner + # treats a missing dimension conservatively (never prunes on it). + try: + min_ns = int(pd.Timestamp(min_val).value) + max_ns = int(pd.Timestamp(max_val).value) + except (OverflowError, pd.errors.OutOfBoundsDatetime): + continue + ranges[str(dim)] = (min_ns, max_ns, "timestamp_ns") + elif hasattr(min_val, "item"): + min_val = min_val.item() + max_val = max_val.item() + dtype = "float64" if isinstance(min_val, float) else "int64" + ranges[str(dim)] = (min_val, max_val, dtype) + else: + dtype = "float64" if isinstance(min_val, float) else "int64" + ranges[str(dim)] = (min_val, max_val, dtype) + return ranges + + +def partition_metadata( + ds: xr.Dataset, blocks: list[Block] +) -> list[PartitionBounds]: + """Compute min/max coordinate values for each partition. + + This metadata enables filter pushdown: SQL queries with WHERE clauses + on dimension columns can prune partitions that can't contain matching rows. + + Args: + ds: The xarray Dataset containing coordinate values. + blocks: List of block slices from block_slices(). + + Returns: + List of dicts mapping dimension name to + (min_value, max_value, dtype_str) tuples. + + - For datetime64, values are nanoseconds since Unix epoch + (int64), dtype_str is "timestamp_ns" + - For numeric types, values are Python int or float, + dtype_str is "int64" or "float64" + + Note: + If a partition has an empty slice for a dimension, that dimension is + omitted from the partition's metadata. The Rust pruning logic treats + missing dimensions conservatively (never prunes on them). + """ + # Hoist coordinate array reads outside the partition loop. + # ds.coords[dim].values materializes the full array on every call; doing it + # N_partitions × N_dims times is wasteful and, for remote Zarr-backed datasets + # (e.g. ARCO-ERA5 on GCS), may trigger repeated network I/O. + coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} + return [_block_metadata(coord_arrays, block) for block in blocks] diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py deleted file mode 100644 index 8898ef4f..00000000 --- a/xarray_sql/df_test.py +++ /dev/null @@ -1,341 +0,0 @@ -import itertools -import tracemalloc - -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -import xarray as xr - -from .reader import read_xarray -from .df import explode, block_slices, from_map, pivot, from_map_batched - - -def rand_wx(start: str, end: str) -> xr.Dataset: - np.random.seed(42) - lat = np.linspace(-90, 90, num=720) - lon = np.linspace(-180, 180, num=1440) - time = pd.date_range(start, end, freq="h") - level = np.array([1000, 500], dtype=np.int32) - reference_time = pd.Timestamp(start) - temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level)) - precipitation = 10 * np.random.rand(720, 1440, len(time), len(level)) - return xr.Dataset( - data_vars=dict( - temperature=(["lat", "lon", "time", "level"], temperature), - precipitation=(["lat", "lon", "time", "level"], precipitation), - ), - coords=dict( - lat=lat, - lon=lon, - time=time, - level=level, - reference_time=reference_time, - ), - attrs=dict(description="Random weather."), - ) - - -def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100): - """Create a large xarray dataset for memory testing.""" - np.random.seed(42) - - time = pd.date_range("2020-01-01", periods=time_steps, freq="h") - lat = np.linspace(-90, 90, lat_points) - lon = np.linspace(-180, 180, lon_points) - - temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10 - precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100 - - return xr.Dataset( - { - "temperature": (["time", "lat", "lon"], temp_data), - "precipitation": (["time", "lat", "lon"], precip_data), - }, - coords={"time": time, "lat": lat, "lon": lon}, - ) - - -def adding_function(x, y): - """Simple function that adds two values and returns a DataFrame.""" - result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]}) - return result - - -@pytest.fixture -def air(): - ds = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - return ds.chunk(chunks) - - -@pytest.fixture -def air_small(air): - return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk( - {"time": 240} - ) - - -@pytest.fixture -def randwx(): - return rand_wx("1995-01-13T00", "1995-01-13T01") - - -@pytest.fixture -def large_ds(): - return create_large_dataset().chunk({"time": 25}) - - -def test_explode_cardinality(air): - dss = explode(air) - assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()]) - - -def test_explode_dim_sizes_one(air): - chunks = {"time": 240} - ds = next(iter(explode(air))) - for k, v in chunks.items(): - assert k in ds.dims - assert v == ds.sizes[k] - - -@pytest.mark.skip(reason="TODO(alxmrs): Why is this test slow?") -def test_explode_dim_sizes_all(air): - dss = explode(air) - assert [tuple(ds.dims.values()) for ds in dss] == list( - itertools.product(*air.chunksizes.values()) - ) - - -def test_explode_data_equal_one_first(air): - ds = next(iter(explode(air))) - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - assert air.isel(iselection).equals(ds) - - -def test_explode_data_equal_one_last(air): - dss = list(explode(air)) - ds = dss[-1] - - # For the last chunk, we need to calculate where it actually starts - # The original logic slice(0, s) only works for the first chunk - iselection = {} - for dim in ds.dims: - # Get chunk boundaries - chunk_bounds = np.cumsum((0,) + air.chunks[dim]) - # Last chunk index - last_chunk_idx = len(air.chunks[dim]) - 1 - # Calculate actual start and end positions - start = chunk_bounds[last_chunk_idx] - end = chunk_bounds[last_chunk_idx + 1] - iselection[dim] = slice(start, end) - - assert air.isel(iselection).equals(ds) - - -def test_from_map_basic(): - def make_df(x): - return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]}) - - result = from_map(make_df, [1, 2, 3]) - assert isinstance(result, pa.Table) - assert len(result) == 6 - assert result.column_names == ["value", "index"] - - -def test_from_map_multiple_iterables(): - def add_values(x, y): - return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]}) - - result = from_map(add_values, [1, 2], [10, 20]) - assert isinstance(result, pa.Table) - assert len(result) == 2 - - df = result.to_pandas() - assert list(df["sum"]) == [11, 22] - - -def test_from_map_with_args(): - def multiply_and_add(x, multiplier, add_value): - return pd.DataFrame({"result": [x * multiplier + add_value]}) - - result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) - assert isinstance(result, pa.Table) - assert len(result) == 3 - - df = result.to_pandas() - assert list(df["result"]) == [12, 14, 16] - - -def test_from_map_with_pyarrow_tables(): - def make_arrow_table(x): - df = pd.DataFrame({"value": [x]}) - return pa.Table.from_pandas(df) - - result = from_map(make_arrow_table, [1, 2, 3]) - assert isinstance(result, pa.Table) - assert len(result) == 3 - - -def test_from_map_batched_basic_functionality(air_small): - blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) - - first_block_df = pivot(air_small.isel(blocks[0])) - expected_schema = pa.Schema.from_pandas(first_block_df) - - reader = from_map_batched( - pivot, [air_small.isel(block) for block in blocks], schema=expected_schema - ) - - assert isinstance(reader, pa.RecordBatchReader) - assert reader.schema == expected_schema - - batches = list(reader) - assert len(batches) > 0 - for batch in batches: - assert batch.schema == expected_schema - assert len(batch) > 0 - - -def test_from_map_batched_multiple_iterables(): - x_values = [1, 2, 3, 4, 5] - y_values = [10, 20, 30, 40, 50] - - expected_schema = pa.schema( - [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())] - ) - - reader = from_map_batched( - adding_function, x_values, y_values, schema=expected_schema - ) - table = reader.read_all() - df = table.to_pandas() - - expected_df = pd.DataFrame( - { - "x": x_values, - "y": y_values, - "sum": [x + y for x, y in zip(x_values, y_values)], - } - ) - pd.testing.assert_frame_equal(df, expected_df) - - -def test_from_map_batched_with_args_and_kwargs(): - def multiply_and_add(x, multiplier, offset=0): - return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]}) - - values = [1, 2, 3] - expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())]) - - reader = from_map_batched( - multiply_and_add, values, args=(2,), offset=5, schema=expected_schema - ) - table = reader.read_all() - df = table.to_pandas() - - expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]}) - pd.testing.assert_frame_equal(df, expected_df) - - -def test_from_map_batched_empty_iterables(): - empty_schema = pa.schema([("value", pa.int64())]) - - reader = from_map_batched( - lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema - ) - batches = list(reader) - assert len(batches) == 0 - - -def test_from_map_batched_consistency_with_regular_map(air_small): - blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3})) - datasets = [air_small.isel(block) for block in blocks] - - first_df = pivot(datasets[0]) - schema = pa.Schema.from_pandas(first_df) - - reader = from_map_batched(pivot, datasets, schema=schema) - batched_table = reader.read_all() - - regular_dfs = [pivot(ds) for ds in datasets] - regular_table = pa.Table.from_pandas( - pd.concat(regular_dfs, ignore_index=True) - ) - - assert batched_table.schema == regular_table.schema - assert len(batched_table) == len(regular_table) - - batched_df = ( - batched_table.to_pandas() - .sort_values(["time", "lat", "lon"]) - .reset_index(drop=True) - ) - regular_df = ( - regular_table.to_pandas() - .sort_values(["time", "lat", "lon"]) - .reset_index(drop=True) - ) - - pd.testing.assert_frame_equal(batched_df, regular_df) - - -def test_from_map_batched_integration_with_datafusion_via_read_xarray(): - air = xr.tutorial.open_dataset("air_temperature") - air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15)) - air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8}) - - arrow_stream = read_xarray( - air_chunked, chunks={"time": 25, "lat": 5, "lon": 8} - ) - - assert hasattr(arrow_stream, "schema") - assert hasattr(arrow_stream, "__iter__") - - table = arrow_stream.read_all() - assert len(table) > 0 - - expected_columns = {"time", "lat", "lon", "air"} - actual_columns = set(table.column_names) - assert expected_columns.issubset(actual_columns) - - -def test_read_xarray_loads_one_chunk_at_a_time(large_ds): - tracemalloc.start() - iterable = read_xarray(large_ds) - first_size, first_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - - sizes, peaks = [], [] - - first_chunk = large_ds.isel(next(block_slices(large_ds))) - chunk_size = first_chunk.nbytes - - # Creating the iterator should be inexpensive -- less than one chunk. - # We multiply by constant factors because chunks have additional overhead - assert first_size < chunk_size * 3 - assert first_peak < chunk_size * 6 - - for it in iterable: - _ = it - cur_size, cur_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - sizes.append(cur_size) - peaks.append(cur_peak) - - mean_size = np.mean(sizes) - mean_peak = np.mean(peaks) - - for size in sizes: - assert mean_size * 1.1 > size - assert chunk_size * 3 > size - assert chunk_size * 2 < size - - for peak in peaks: - assert mean_peak * 1.1 > peak - assert chunk_size * 7 > peak - assert chunk_size * 4 < peak - - assert max(peaks) < large_ds.nbytes - - tracemalloc.stop() diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py new file mode 100644 index 00000000..99ef7b31 --- /dev/null +++ b/xarray_sql/ds.py @@ -0,0 +1,1085 @@ +"""Reconstruct xarray Datasets from SQL query results. + +The inverse of the forward Dataset-to-table pivot done by +[xarray_sql.df.pivot][]. Internally defines an [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] +wrapper around the DataFusion ``DataFrame`` returned by +[XarrayContext.sql][xarray_sql.sql.XarrayContext.sql], with a [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] +method that round-trips a query result back to ``xr.Dataset``. + +Reconstruction is controlled by the ``chunks`` argument to +[XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] -- the xarray idiom for tuning how a +result is partitioned -- rather than by reflecting on the query plan: + +* **Eager** (``chunks=None``, or the default ``"inherit"`` when the + result keeps no multi-chunk source dimension): the plan executes + exactly once via ``execute_stream`` and the result is scattered into a + dense in-memory Dataset. This is the right default for reductions + (aggregations), whose results are small, and it never re-executes. +* **Lazy / chunked** (``chunks`` is a mapping, ``"auto"``, or + ``"inherit"`` over a multi-chunk source dimension): data variables are + backed by [SQLBackendArray][xarray_sql.ds.SQLBackendArray] wrapped in + ``xarray.core.indexing.LazilyIndexedArray`` and chunked via xarray's + configured chunk manager (dask, cubed, ...). Each chunk maps onto the + source partitions and reads its coordinate range on access by + translating the indexer into a DataFusion ``filter`` expression, so only + the requested partitions are materialized as Arrow ``RecordBatch`` es + and scattered into numpy. + +``.compute()`` materializes the whole Dataset in memory. +""" + +from __future__ import annotations + +import warnings +from collections.abc import Mapping +from typing import Any, Literal, cast + +import numpy as np +import pandas as pd +import pyarrow as pa +import xarray as xr + +from .lazyscan import DataFusionHandle, DimSpec, LazyResultHandle + +Sparsity = Literal["result", "template"] +"""Output coordinate extent for a filtered round-trip. + +* ``"result"`` keeps only the dim values present in the query result, so + the output is sparse and equal to whatever rows came back. +* ``"template"`` reindexes to the registered Dataset's full coord ranges + and fills absent cells with ``fill_value``. +""" + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + + +def _ds_var_dims(ds: xr.Dataset) -> list[str]: + """Return a Dataset's data-variable dim order. + + The forward path validates that all data variables share the same dims + tuple, so the first var's dim order is canonical. Falls back to + ``ds.dims`` keys for empty Datasets. Always use this rather than + ``list(ds.dims)`` when round-tripping, since the latter is in + canonical name order and may not match the variable's axis order. + """ + if ds.data_vars: + return list(next(iter(ds.data_vars.values())).dims) + return list(ds.dims) + + +def _apply_template(ds: xr.Dataset, template: xr.Dataset) -> xr.Dataset: + """Recover metadata that the forward SQL pivot strips. + + Adds back, where unambiguous: + + * Data-variable ``attrs`` and ``encoding`` for vars present in + ``template`` (aggregation aliases like ``air_avg`` get nothing). + Dtype-bound encoding keys (``dtype``, ``_FillValue``, + ``missing_value``) are intentionally dropped: SQL may have + changed the column's dtype (e.g. ``int16`` -> ``float64`` after + ``AVG`` or a null-introducing filter), and reattaching the + source's packing would make a later ``ds.to_netcdf()`` write + corrupt values. + * Dim-coordinate dtype, where SQL upcasted (datetime is the + canonical case). + * Non-dim coordinates whose dims are all present in ``ds`` (scalar + coords attach as-is; vector coords use ``.sel``). + * Dataset-level ``attrs``. + + Skipped coords are warned about once per call. + """ + out = ds.copy() + + # 1. Data-var attrs / encoding for vars present in the template. + # Aggregation aliases absent from template intentionally inherit nothing. + for name in list(out.data_vars): + if name in template.data_vars: + out[name].attrs = dict(template[name].attrs) + # Drop dtype-bound encoding keys; SQL may have changed dtype. + enc = { + k: v + for k, v in template[name].encoding.items() + if k not in {"dtype", "_FillValue", "missing_value"} + } + out[name].encoding = enc + + # 2. Restore dim-coordinate dtype when SQL changed it (e.g. datetime + # upcast through pyarrow / pandas) and copy the source's dim-coord + # attrs (``standard_name``, ``long_name``, ``units``, etc.). + for d in list(out.dims): + if d in template.coords: + tdt = template.coords[d].dtype + if out.coords[d].dtype != tdt: + try: + out = out.assign_coords({d: out.coords[d].astype(tdt)}) + except (ValueError, TypeError): + pass # incompatible cast; leave as-is + out[d].attrs = dict(template.coords[d].attrs) + + # 3. Non-dim coordinates whose dims are all present in the result. + out_dims = set(out.dims) + skipped: list[str] = [] + for cname, coord in template.coords.items(): + if cname in template.dims: + continue # dim coord; already in out + if not set(coord.dims) <= out_dims: + continue # spans dims the result lacks + try: + if not coord.dims: + # Scalar coord (e.g. weather_dataset.reference_time). + out = out.assign_coords({cname: coord}) + else: + sel = {d: out.coords[d] for d in coord.dims} + out = out.assign_coords({cname: coord.sel(sel)}) + except (KeyError, ValueError, TypeError): + skipped.append(cname) + + # 4. Dataset-level attrs. + out.attrs = dict(template.attrs) + + if skipped: + warnings.warn( + f"Could not re-attach non-dim coordinates from template: {skipped}", + stacklevel=3, + ) + return out + + +def _axis_numeric(values: np.ndarray) -> np.ndarray: + """View an axis as float64 for affine position arithmetic.""" + if values.dtype.kind == "M": + return values.astype("datetime64[ns]").view("int64").astype("float64") + return values.astype("float64", copy=False) + + +def _affine_axis(requested: np.ndarray) -> tuple[float, float] | None: + """``(origin, step)`` when *requested* is uniformly spaced, else None. + + Uniform spacing must hold exactly enough that ``rint((v - origin) / + step)`` recovers every index: the deviation of each element from its + affine prediction is checked against a quarter step. Non-numeric + axes (strings, cftime objects) never qualify. + """ + if requested.dtype.kind not in ("i", "u", "f", "M") or len(requested) < 2: + return None + numeric = _axis_numeric(requested) + step = (numeric[-1] - numeric[0]) / (len(numeric) - 1) + if step == 0 or not np.isfinite(step): + return None + predicted = numeric[0] + step * np.arange(len(numeric)) + # Written as a <= comparison so a NaN anywhere in the axis (e.g. a + # NULL dim value in the result) fails the check and falls back to + # the searchsorted path, which handles it positionally. + if not (np.abs(numeric - predicted) <= 0.25 * abs(step)).all(): + return None + return float(numeric[0]), float(step) + + +def _scatter_batches_to_ndarray( + batches: list[pa.RecordBatch], + dimension_columns: list[str], + requested: dict[str, np.ndarray], + var_name: str, + out_shape: tuple[int, ...], + dtype: np.dtype, + drop_axes: list[int], +) -> np.ndarray: + """Convert filtered Arrow ``RecordBatch`` rows into a dense N-D numpy array. + + SQL query results arrive as flat rows; xarray expects N-D arrays. + This bridges the two: each row carries the dim-coord values that + identify its cell in the output cube plus the value to write there. + We look up the row's N-D position by binary-searching its coord + values within the caller's requested coord arrays + (``np.searchsorted``), then scatter-write the value at that index. + + Missing combinations (sparse results from filtered queries) stay as + ``NaN`` for floating-point outputs by pre-filling the buffer; integer + outputs leave them as ``np.empty``-style undefined values. + """ + # NaN fill for float outputs; default for int/datetime falls through + # to ``np.empty``-style undefined values (but every output cell is + # written below for non-sparse cases). + out = ( + np.full(out_shape, np.nan, dtype=dtype) + if np.issubdtype(dtype, np.floating) + else np.empty(out_shape, dtype=dtype) + ) + + # ``requested[d]`` may be in any order (callers can iselect arbitrary + # positions, and template coords like air_temperature.lat are descending). + # ``np.searchsorted`` requires ascending input, so we sort each requested + # array once, search there, and remap back to the original positions. + # Uniformly spaced axes (the norm for rasters and regular time steps, + # ascending or descending) skip the search entirely: the position is + # ``rint((value - origin) / step)``, a fused vector op several times + # faster than a per-row binary search. + affine = {d: _affine_axis(requested[d]) for d in dimension_columns} + sorted_idx = { + d: np.argsort(requested[d]) + for d in dimension_columns + if affine[d] is None + } + sorted_req = {d: requested[d][sorted_idx[d]] for d in sorted_idx} + + for batch in batches: + if batch.num_rows == 0: + continue + schema_names = batch.schema.names + # Build per-dim position arrays for this batch (positions within + # the caller's requested coord order). + positions = [] + for d in dimension_columns: + col_arr = batch.column(schema_names.index(d)) + vals = col_arr.to_numpy(zero_copy_only=False) + pair = affine[d] + if pair is not None: + origin, step = pair + pos = np.rint((_axis_numeric(vals) - origin) / step).astype( + np.intp + ) + positions.append(pos) + else: + pos_in_sorted = np.searchsorted(sorted_req[d], vals) + positions.append(sorted_idx[d][pos_in_sorted]) + value_arr = batch.column(schema_names.index(var_name)).to_numpy( + zero_copy_only=False + ) + out[tuple(positions)] = value_arr.astype(dtype, copy=False) + + if drop_axes: + out = np.squeeze(out, axis=tuple(drop_axes)) + return cast(np.ndarray, out) + + +class SQLBackendArray(xr.backends.BackendArray): + """Read-only lazy N-D array view over a re-executable SQL result. + + Bridges xarray's lazy-indexing interface + (``xarray.backends.BackendArray``) to an engine query result, + so an xarray ``Dataset`` can present a SQL query as if it were a + materialized N-D array without actually loading any data until the + caller asks for it. This is the workhorse that lets + [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset] (and the engine-agnostic + ``xql.to_dataset(chunks=...)``) return a Dataset cheaply. + + On each ``__getitem__`` call, the requested xarray indexer is + translated into per-dimension coordinate windows and a column + projection, executed through a + [LazyResultHandle][xarray_sql.lazyscan.LazyResultHandle] (DataFusion, DuckDB, + or Polars — each renders the windows with its own typed expression + API). The resulting Arrow ``RecordBatch`` es are scattered into a + preallocated numpy buffer, so only the requested data is + materialized. + + Constraints and caveats: + + - Read-only: there is no write path; the backend exists to surface + query results, not to round-trip writes into a SQL store. + - The underlying engine object may hold non-picklable references + (DataFusion's ``SessionContext``, a DuckDB connection). The class + therefore overrides ``__copy__`` and ``__deepcopy__`` to return + ``self`` -- this is safe because the backend is read-only. + - ``IndexingSupport.OUTER``: ``BasicIndexer`` and ``OuterIndexer`` + are translated to filter predicates directly; ``VectorizedIndexer`` + paths through xarray's adapter to outer-then-gather and so still + works, just less efficiently. + + Raises: + ValueError, engine exceptions: propagated from the underlying + filter/project/execute chain if a predicate refers to a + missing column, the dtype of a literal is incompatible, or + the execution itself fails. + AssertionError: from ``np.searchsorted`` mis-alignment, which + indicates the result contains coordinate values not present + in the wrapper's pre-computed coord arrays -- usually a + symptom of a filtered query whose coord discovery missed a + value. + + Constructed by ``_build_lazy_scan``; users should not instantiate + this class directly. + """ + + def __init__( + self, + handle: LazyResultHandle, + var_name: str, + dimension_columns: list[str], + coord_arrays: dict[str, np.ndarray], + shape: tuple[int, ...], + dtype: np.dtype, + ) -> None: + self._handle = handle + self._var_name = var_name + self._dimension_columns = list(dimension_columns) + self._coord_arrays = coord_arrays + # Computed once per dim: whether the whole coordinate array is + # strictly monotonic, the precondition for translating contiguous + # positional windows into value ranges (see _dim_spec). + self._monotonic = { + d: _strictly_monotonic(coord_arrays[d]) for d in dimension_columns + } + self.shape = tuple(shape) + self.dtype = np.dtype(dtype) + + def __getitem__(self, key: Any) -> np.ndarray: + return cast( + np.ndarray, + xr.core.indexing.explicit_indexing_adapter( + key, + self.shape, + xr.core.indexing.IndexingSupport.OUTER, + self._raw_getitem, + ), + ) + + def __copy__(self) -> "SQLBackendArray": + # The backend is read-only; the underlying DataFusion DataFrame + # holds a non-picklable SessionContext reference, so sharing the + # same backend across a copy is both safe and necessary. + return self + + def __deepcopy__(self, memo: dict) -> "SQLBackendArray": + return self + + # ------------------------------------------------------------------ + + def _raw_getitem(self, key: tuple) -> np.ndarray: + """Materialize the indexed region described by *key* via the engine. + + ``key`` is a tuple of ``int``/``slice``/1-D integer-array, one per + dim, in ``_dimension_columns`` order. + """ + requested: dict[str, np.ndarray] = {} + # Per-dim windows for the engine. Dims whose indexer covers the + # full extent are omitted entirely so the engine doesn't have to + # evaluate a tautology. + specs: dict[str, DimSpec] = {} + drop_axes: list[int] = [] + for axis, (dim, k) in enumerate( + zip(self._dimension_columns, key, strict=True) + ): + coord = self._coord_arrays[dim] + contiguous = False + if isinstance(k, slice): + start = 0 if k.start is None else k.start + stop = len(coord) if k.stop is None else k.stop + step = 1 if k.step is None else k.step + requested[dim] = np.asarray(coord[start:stop:step]) + contiguous = step == 1 + if start == 0 and stop >= len(coord) and step == 1: + continue + elif isinstance(k, (int, np.integer)): + requested[dim] = np.asarray([coord[int(k)]]) + drop_axes.append(axis) + else: + arr = np.asarray(k) + requested[dim] = np.asarray(coord[arr]) + if ( + len(arr) == len(coord) + and (arr == np.arange(len(coord))).all() + ): + continue + contiguous = len(arr) > 1 and bool((np.diff(arr) == 1).all()) + specs[dim] = _dim_spec( + requested[dim], contiguous, self._monotonic[dim] + ) + + out_shape = tuple(len(requested[d]) for d in self._dimension_columns) + if any(n == 0 for n in out_shape): + empty = np.empty(out_shape, dtype=self.dtype) + squeezed = ( + np.squeeze(empty, axis=tuple(drop_axes)) if drop_axes else empty + ) + return cast(np.ndarray, squeezed) + + batches = self._handle.fetch( + specs, self._dimension_columns + [self._var_name] + ) + return _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=self._dimension_columns, + requested=requested, + var_name=self._var_name, + out_shape=out_shape, + dtype=self.dtype, + drop_axes=drop_axes, + ) + + +def _strictly_monotonic(coord: np.ndarray) -> bool: + """Whether ``coord`` is strictly increasing or strictly decreasing. + + Strict monotonicity of the whole coordinate array is the + precondition for translating a contiguous positional window into a + value range: with duplicated or unsorted values, ``[min, max]`` of a + window admits coordinate values at positions outside the window. + NaN/NaT (whose comparisons are all false) and non-comparable object + arrays report ``False``, which safely falls back to value lists. + """ + if len(coord) < 2: + return True + head, tail = coord[:-1], coord[1:] + try: + return bool((tail > head).all() or (tail < head).all()) + except TypeError: + return False + + +def _dim_spec( + vals: np.ndarray, contiguous: bool, coord_monotonic: bool +) -> DimSpec: + """The engine window for one dim's requested coordinate values. + + A contiguous run of positions over a strictly monotonic coordinate + array is exactly the value range ``[min, max]`` — a two-literal + predicate engines can push into range pruning. Monotonicity must + hold for the *entire* coordinate array (``coord_monotonic``), not + just the requested window: template coords are used verbatim, and + over a non-monotonic array a window's ``[min, max]`` admits values + at positions outside the window, which the scatter would then write + to wrong cells. Anything else (stepped slices, fancy indexers, + non-monotonic or duplicated coords) must be an explicit value list: + a range would admit rows the scatter did not request. + """ + if contiguous and coord_monotonic and len(vals) > 1: + return ("range", vals.min(), vals.max()) + return ("values", vals, None) + + +def _c_order_grid( + dim_cols: dict[str, np.ndarray], + coord_arrays: dict[str, np.ndarray], + dimension_columns: list[str], + total_rows: int, +) -> bool: + """Whether the result rows form the complete grid in C order. + + True iff the row count is exactly the coordinate product and every + dimension column is its coordinates repeated/tiled in C order — the + shape any unfiltered or bbox-windowed scan produces. When it holds, + data variables are dense row-major arrays already and can be + reshaped instead of scatter-written (one memcpy versus a + ``searchsorted`` per dimension per row). + """ + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + if total_rows != int(np.prod(shape)) or total_rows == 0: + return False + for k, d in enumerate(dimension_columns): + inner = int(np.prod(shape[k + 1 :])) + outer = int(np.prod(shape[:k])) + view = dim_cols[d].reshape(outer, shape[k], inner) + if not (view == coord_arrays[d][None, :, None]).all(): + return False + return True + + +def _dataset_from_batches( + batches: list[pa.RecordBatch], + dimension_columns: list[str], + field_names: list[str], + field_types: dict[str, Any], +) -> xr.Dataset: + """Build a dense in-memory Dataset from Arrow ``RecordBatch`` es. + + The engine-agnostic core of the eager round-trip: derives the + coordinates and every data variable from a single already-executed + result, whichever engine produced it. ``field_types`` values only + need a ``to_pandas_dtype()`` method (both ``pyarrow.DataType`` and + DataFusion's Arrow type wrappers qualify). + + Complete grid-ordered results (unfiltered scans, bbox windows) are + reshaped directly; anything else — sparse results from filtered + queries, engine-reordered rows — falls back to the positional + scatter, which handles arbitrary row order. + """ + dim_cols: dict[str, np.ndarray] = {} + coord_arrays: dict[str, np.ndarray] = {} + for d in dimension_columns: + if not batches: + dim_cols[d] = np.asarray([]) + coord_arrays[d] = np.asarray([]) + continue + vals = np.concatenate( + [ + b.column(b.schema.names.index(d)).to_numpy(zero_copy_only=False) + for b in batches + ] + ) + dim_cols[d] = vals + # Preserve the order coordinate values first appear in the result so an + # ORDER BY direction (e.g. ``ORDER BY level DESC``) carries through to + # the Dataset dimension instead of being force-sorted ascending. + # pd.unique keeps first-appearance order; the scatter below argsorts + # internally, so arbitrarily-ordered coordinates are placed correctly. + coord_arrays[d] = np.asarray(pd.unique(vals)) + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + total_rows = sum(b.num_rows for b in batches) + + grid_ordered = _c_order_grid( + dim_cols, coord_arrays, dimension_columns, total_rows + ) + + data_vars: dict[str, xr.Variable] = {} + for name in field_names: + if name in dimension_columns: + continue + np_dtype = np.dtype(field_types[name].to_pandas_dtype()) + if grid_ordered: + flat = np.concatenate( + [ + b.column(b.schema.names.index(name)).to_numpy( + zero_copy_only=False + ) + for b in batches + ] + ) + dense = flat.astype(np_dtype, copy=False).reshape(shape) + else: + dense = _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=dimension_columns, + requested=coord_arrays, + var_name=name, + out_shape=shape, + dtype=np_dtype, + drop_axes=[], + ) + data_vars[name] = xr.Variable(dimension_columns, dense) + + coords_arg = {d: coord_arrays[d] for d in dimension_columns} + return xr.Dataset(data_vars=data_vars, coords=coords_arg) + + +def _materialize( + inner_df: Any, + dimension_columns: list[str], + field_names: list[str], + field_types: dict[str, Any], +) -> xr.Dataset: + """Execute the query once and build a dense in-memory Dataset. + + Runs the plan exactly once via ``execute_stream()`` -- streaming the result + as Arrow ``RecordBatch`` es (``datafusion.RecordBatch.to_pyarrow()``) -- then + derives both the coordinates and every data variable from that single pass. + This is the eager path, used when no output chunking is requested. It never + re-executes, so an aggregation over a remote Zarr scan costs exactly one + scan, regardless of how many dimensions or variables the result has. + """ + batches = [b.to_pyarrow() for b in inner_df.execute_stream()] + return _dataset_from_batches( + batches, dimension_columns, field_names, field_types + ) + + +_PURE_SCAN_NODES = {"Projection", "Sort", "TableScan", "SubqueryAlias"} + + +def _unfiltered_scan_table(inner_df: Any) -> str | None: + """Return the scanned table name iff the query is a pure unfiltered scan. + + A pure scan only contains ``Projection``, ``Sort``, ``TableScan``, + ``SubqueryAlias`` nodes and exactly one ``TableScan``. Anything else + (``Filter``, ``Aggregate``, ``Join``, ``Union``, ``Limit``, multi-table + joins, ...) returns ``None`` so the caller falls back to per-dim + discovery. The returned name is the registered table the caller can + look up to source coord arrays from. + """ + try: + lp = inner_df.logical_plan() + except Exception: + return None + table_name: str | None = None + stack = [lp] + while stack: + node = stack.pop() + try: + variant = node.to_variant() + except Exception: + return None + cls = type(variant).__name__ + if cls not in _PURE_SCAN_NODES: + return None + if cls == "TableScan": + try: + this = variant.table_name() + except (AttributeError, TypeError): + return None + if not isinstance(this, str): + return None + if table_name is not None and table_name != this: + return None # multi-table scan; not a single source + table_name = this + stack.extend(node.inputs()) + return table_name + + +def _maybe_template_coords( + templates: dict[str, xr.Dataset] | None, + dimension_columns: list[str], + inner_df: Any, +) -> dict[str, np.ndarray] | None: + """Use the scanned table's registered coord arrays directly when safe. + + Returns coord arrays sourced from the registered Dataset for the + scanned table iff the query is an unfiltered scan over that single + table and the registered Dataset carries all requested dims. Returns + ``None`` otherwise so the caller falls back to per-dim discovery. + Skipping discovery avoids one full plan execution per dim and + preserves the source's coordinate order. + + Coord values come from the **scanned** registered Dataset, not from + any user-supplied ``template=`` (which is for metadata recovery + only). That keeps the fast path correct when a user with multiple + registered Datasets passes a metadata template that differs from + the query's source. + """ + if not templates: + return None + table = _unfiltered_scan_table(inner_df) + if table is None or table not in templates: + return None + source = templates[table] + if not all(d in source.coords for d in dimension_columns): + return None + return {d: np.asarray(source.coords[d].values) for d in dimension_columns} + + +def _build_lazy_scan( + handle: LazyResultHandle, + dimension_columns: list[str], + field_names: list[str], + field_types: dict[str, Any], + coord_arrays: dict[str, np.ndarray] | None = None, +) -> xr.Dataset: + """Build a lazy Dataset whose data vars are [SQLBackendArray][xarray_sql.ds.SQLBackendArray]. + + Used when output chunking is requested: each data variable stays lazy and, + once wrapped by ``Dataset.chunk``, every chunk reads its coordinate range + via a pushdown filter on first access. Coordinates come either from the + caller (the scanned table's registered Dataset for unfiltered DataFusion + scans -- see ``_maybe_template_coords`` -- or an explicitly trusted + template) or from per-dim distinct queries through the handle; over a + registered pushdown table the engine projects to that single coordinate + column, so discovery reads coordinate values only (no data-variable I/O). + """ + if coord_arrays is None: + coord_arrays = {} + for d in dimension_columns: + # ``distinct`` returns engine order; sort ascending so + # positional slices map onto contiguous value ranges. + coord_arrays[d] = np.sort(handle.distinct(d)) + shape = tuple(len(coord_arrays[d]) for d in dimension_columns) + + data_vars: dict[str, xr.Variable] = {} + for name in field_names: + if name in dimension_columns: + continue + np_dtype = field_types[name].to_pandas_dtype() + backend = SQLBackendArray( + handle=handle, + var_name=name, + dimension_columns=dimension_columns, + coord_arrays=coord_arrays, + shape=shape, + dtype=np_dtype, + ) + lazy = xr.core.indexing.LazilyIndexedArray(backend) + data_vars[name] = xr.Variable(dimension_columns, lazy) + + coords_arg = {d: coord_arrays[d] for d in dimension_columns} + return xr.Dataset(data_vars=data_vars, coords=coords_arg) + + +def _auto_chunk_target_bytes() -> int: + """Byte target for ``chunks="auto"`` (the chunk manager's, else 128 MiB).""" + try: + import dask + from dask.utils import parse_bytes + + return int(parse_bytes(dask.config.get("array.chunk-size"))) + except Exception: + return 128 * 1024 * 1024 + + +def _auto_chunks( + template: xr.Dataset | None, + dimension_columns: list[str], + field_types: dict[str, Any], +) -> dict[str, int] | None: + """Resolve ``chunks="auto"`` to a source-partition-aligned chunk spec. + + Sizes chunks to roughly the chunk manager's byte target (dask's + ``array.chunk-size``, default 128 MiB) but snaps boundaries to whole source + partitions, so every chunk is a union of source partitions -- no chunk splits + a partition (which would make adjacent chunks re-read it). This is what makes + ``"auto"`` useful for finely partitioned sources (e.g. ERA5 + ``chunks={"time": 1}``): it coarsens many tiny partitions into memory-sized, + aligned chunks. Returns ``None`` when there is no resolvable source grid to + align to, so the caller falls back to the chunk manager's own ``"auto"``. + """ + if template is None: + return None + part = template.chunksizes # dim -> tuple of source chunk lengths + chunked_dims = [ + d for d in dimension_columns if d in part and len(part[d]) > 1 + ] + if not chunked_dims: + return None + + itemsizes = [ + np.dtype(t.to_pandas_dtype()).itemsize + for name, t in field_types.items() + if name not in dimension_columns + ] + itemsize = max(itemsizes) if itemsizes else 8 + + # Bytes in one source-partition block: the nominal source chunk length per + # dimension (``part[d][0]``) multiplied across all dims, times itemsize. + block_bytes = itemsize + for d in dimension_columns: + if d in part: + block_bytes *= int(part[d][0]) + # Number of source partitions to merge per chunk to approach the target. + merge = max(1, _auto_chunk_target_bytes() // max(block_bytes, 1)) + + # Absorb the coarsening into the most finely partitioned dimension; the rest + # keep their source chunk length. xarray caps an oversize chunk at the dim + # length, so an over-large merge simply yields a single chunk on that dim. + primary = max(chunked_dims, key=lambda d: len(part[d])) + return { + d: int(part[d][0]) * (merge if d == primary else 1) + for d in chunked_dims + } + + +def _result_to_xarray( + inner_df: Any, + dimension_columns: list[str], + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str | None, + templates: dict[str, xr.Dataset] | None = None, +) -> xr.Dataset: + """Reconstruct an ``xr.Dataset`` from a SQL result. + + ``chunks`` (already resolved by ``XarrayDataFrame._resolve_chunks``) + selects the execution strategy: + + * ``None`` -> eager: execute once and materialize a dense Dataset + (``_materialize``). Correct for any query and the right default for + reductions, whose results are small. + * a mapping (or ``"auto"``) -> lazy/chunked: build [SQLBackendArray][xarray_sql.ds.SQLBackendArray] + data variables (``_build_lazy_scan``) and wrap them with + ``Dataset.chunk`` so each chunk reads its coordinate range via filter + pushdown. The chunk grid maps onto the source partitions. Chunking goes + through xarray's configured chunk manager (dask, cubed, ...), so no + chunked-array backend is imported directly here. + """ + if sparsity not in ("result", "template"): + raise ValueError( + f"sparsity must be 'result' or 'template', got {sparsity!r}" + ) + if sparsity == "template" and template is None: + raise ValueError( + "sparsity='template' requires template= to be supplied" + ) + + schema = inner_df.schema() + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + + if chunks is None: + ds = _materialize(inner_df, dimension_columns, field_names, field_types) + else: + ds = _build_lazy_scan( + DataFusionHandle(inner_df), + dimension_columns, + field_names, + field_types, + coord_arrays=_maybe_template_coords( + templates, dimension_columns, inner_df + ), + ) + return _finish_dataset( + ds, + dimension_columns, + template, + sparsity, + fill_value, + chunks, + field_types, + ) + + +def _finish_dataset( + ds: xr.Dataset, + dimension_columns: list[str], + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str | None, + field_types: dict[str, Any], +) -> xr.Dataset: + """Shared reconstruction tail: sparsity, template metadata, chunking.""" + if sparsity == "template": + assert template is not None + indexers = { + d: template.coords[d].values + for d in dimension_columns + if d in template.coords and d in template.dims + } + if indexers: + ds = ds.reindex(indexers, fill_value=fill_value) + + if template is not None: + ds = _apply_template(ds, template) + + if chunks is not None: + if chunks == "auto": + # Snap the byte-budgeted "auto" sizing to source partition + # boundaries; fall back to the chunk manager's own "auto" when there + # is no source grid to align to. + chunks = ( + _auto_chunks(template, dimension_columns, field_types) or "auto" + ) + # Wrap the lazy data variables in the configured chunk manager (dask by + # default). Each chunk reads its coordinate range via pushdown on access. + ds = ds.chunk(chunks) + return ds + + +# --------------------------------------------------------------------------- +# Public wrapper +# --------------------------------------------------------------------------- + + +class XarrayDataFrame: + """Wrapper around a DataFusion ``DataFrame`` with xarray-aware helpers. + + Returned by [xarray_sql.XarrayContext.sql][]. Forwards every + attribute it does not define itself to the wrapped DataFrame, so + ``.collect()``, ``.schema()``, ``.show()``, ``.count()`` all work + unchanged. + + Carries a private snapshot of the context's registered Datasets so + ``to_dataset`` can default ``dims`` and recover metadata + dropped by the forward pivot. + + Users should not construct this class directly; let + [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql] produce it. + """ + + def __init__( + self, + inner: Any, + templates: dict[str, xr.Dataset] | None = None, + ) -> None: + """Construct a wrapper. + + Args: + inner: The underlying ``datafusion.DataFrame`` returned by + [XarrayContext.sql][xarray_sql.sql.XarrayContext.sql]. + templates: Snapshot of the registered Datasets on the producing + context, keyed by the SQL identifier each was registered + under. Used by ``to_dataset`` to recover metadata that + the forward pivot strips. ``None`` means no metadata + recovery is possible from registrations alone; callers may + still pass ``template=`` to ``to_dataset`` explicitly. + """ + object.__setattr__(self, "_inner", inner) + object.__setattr__(self, "_templates", dict(templates or {})) + + def to_pandas(self) -> pd.DataFrame: + """Materialize the result as a ``pd.DataFrame`` (DataFusion API).""" + return self._inner.to_pandas() + + def to_dataset( + self, + dims: list[str] | None = None, + template: xr.Dataset | str | None = None, + sparsity: Sparsity = "result", + fill_value: Any = np.nan, + chunks: Mapping[str, int] | str | None = "inherit", + ) -> xr.Dataset: + """Convert the result to an ``xr.Dataset``. + + Args: + dims: Result columns to use as Dataset dimensions. When + ``None``, defaults to a registered Dataset's dimensions that + survive into the result columns, so an aggregation that drops + dims (e.g. ``GROUP BY time`` over a ``(time, lat, lon)`` grid) + round-trips on the remaining dim. Raises when no dimension + survives, or when several registered Datasets imply different + dims (pass ``dims`` explicitly then). + template: Source to recover metadata (attrs, encoding, non-dim + coordinates, dim-coord dtype) from. Either an ``xr.Dataset`` + used directly, or the name of a registered table (e.g. + ``"era5.surface"``) whose Dataset is looked up. When ``None`` + and exactly one Dataset is registered, that one is used. + sparsity: ``"result"`` (default) keeps only dim values + present in the result. ``"template"`` reindexes to the + template's full coord ranges, filling absent cells with + ``fill_value``; requires a template. + fill_value: Used when ``sparsity="template"`` reindexes + to a wider extent. Defaults to ``np.nan``. + chunks: Output chunking, controlling laziness (an xarray idiom). + + * ``"inherit"`` (default): reuse the source Dataset's chunk + sizes, but only for dimensions that were genuinely split into + multiple chunks in the input -- so the output chunk grid maps + onto the source partitions. A reduction that drops the chunked + dimension (e.g. a global aggregation) inherits nothing and so + is materialized eagerly. Falls back to eager when no source + Dataset is resolvable. + * ``None``: eager. Execute the query once and return a dense + in-memory Dataset. Best for reductions (small results). + * a mapping (e.g. ``{"time": 100}``): chunk explicitly. Each + chunk reads its coordinate range lazily via filter pushdown on + access, through xarray's configured chunk manager (dask, + cubed, ...). + * ``"auto"``: size chunks to the chunk manager's byte target but + snap boundaries to whole source partitions, so each chunk is a + union of source partitions. Useful for finely partitioned + sources (e.g. ERA5 ``chunks={"time": 1}``), coarsening many + tiny partitions into memory-sized, aligned chunks. + + Returns: + An ``xr.Dataset`` with ``dims`` as dimensions and the + remaining result columns as data variables. + + Raises: + ValueError: ``dims`` cannot be inferred, names a missing + column, or the result has duplicate dim tuples; + ``template`` names an unknown registered table; or + ``sparsity="template"`` is requested without a + resolvable template. + """ + if not isinstance(template, xr.Dataset): + # ``template`` is a registered-table name or None; look it up. + template = self._resolve_template(template) + if dims is None: + dims = self._infer_dimension_columns(preferred_template=template) + resolved_chunks = self._resolve_chunks(chunks, template, dims) + return _result_to_xarray( + inner_df=self._inner, + dimension_columns=dims, + template=template, + sparsity=sparsity, + fill_value=fill_value, + chunks=resolved_chunks, + templates=self._templates, + ) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + @staticmethod + def _resolve_chunks( + chunks: Mapping[str, int] | str | None, + template: xr.Dataset | None, + dimension_columns: list[str], + ) -> Mapping[str, int] | str | None: + """Resolve the ``chunks`` argument to a concrete spec or ``None``. + + ``None`` selects the eager path; anything else selects the lazy/chunked + path. ``"inherit"`` reuses the source Dataset's chunk sizes -- but only + for dimensions actually split into more than one chunk in the input + (a single full chunk is not "chunked"), so reductions that drop the + chunked dimension resolve to ``None`` (eager) automatically. Mappings + pass through unchanged; ``"auto"`` passes through here and is snapped to + source partition boundaries later (see ``_auto_chunks``). + """ + if chunks is None: + return None + if chunks == "inherit": + if template is None: + return None + sizes = template.chunksizes # dim -> tuple of chunk lengths + inherited = { + d: sizes[d][0] + for d in dimension_columns + if d in sizes and len(sizes[d]) > 1 + } + return inherited or None + return chunks + + def _resolve_template(self, name: str | None) -> xr.Dataset | None: + """Pick a template Dataset for metadata recovery by registered name. + + Priority: + 1. The named registered table (``name``). + 2. If exactly one Dataset is registered on the context, use it. + 3. None. + """ + templates = self._templates + if name is not None: + if name not in templates: + raise ValueError( + f"template={name!r} is not a registered table on this " + f"context. Registered: {list(templates)}" + ) + return templates[name] + if len(templates) == 1: + return next(iter(templates.values())) + return None + + def _infer_dimension_columns( + self, preferred_template: xr.Dataset | None = None + ) -> list[str]: + """Pick a default ``dimension_columns`` from the registry, or raise. + + A registered Dataset's dims that survive into the result columns + become the dimensions, so aggregations that drop dims (e.g. + ``GROUP BY time`` over a ``(time, lat, lon)`` grid) round-trip on the + surviving dim(s). Uses the data variable's dim order (via + ``_ds_var_dims``) so the original axis order is preserved. + """ + result_cols = set(self._result_columns()) + + def surviving(template: xr.Dataset) -> list[str]: + # Template dims still present in the result, in var axis order. + return [d for d in _ds_var_dims(template) if d in result_cols] + + if preferred_template is not None: + preferred = surviving(preferred_template) + if preferred: + return preferred + if not self._templates: + raise ValueError( + "dims cannot be inferred (no registered " + "Dataset on this result); pass dims=[...] " + "explicitly." + ) + candidates = {tuple(surviving(t)) for t in self._templates.values()} + candidates.discard(()) # templates with no surviving dim + if len(candidates) == 1: + return list(next(iter(candidates))) + if not candidates: + raise ValueError( + "dims cannot be inferred: no registered Dataset " + "dimension survives in the result columns. Pass " + "dims=[...] explicitly." + ) + raise ValueError( + "dims cannot be inferred unambiguously: multiple " + "registered Datasets are compatible with the result. Pass " + "dims=[...] explicitly." + ) + + def _result_columns(self) -> list[str]: + """Return the result's column names without materializing rows.""" + return [field.name for field in self._inner.schema()] + + def __getattr__(self, name: str) -> Any: + # Runs only when ``name`` is not found via normal lookup, so this + # safely forwards anything we have not overridden. + return getattr(self._inner, name) + + def __repr__(self) -> str: + return repr(self._inner) diff --git a/xarray_sql/geometry.py b/xarray_sql/geometry.py new file mode 100644 index 00000000..6e3e6406 --- /dev/null +++ b/xarray_sql/geometry.py @@ -0,0 +1,131 @@ +"""GeoArrow point-geometry columns derived from coordinate dimensions. + +A regular grid's pivot already materializes per-row x/y coordinate +columns; a point-geometry column is those same values under a GeoArrow +extension annotation. Two encodings: + +* ``"wkb"`` (default) — 21-byte WKB points under the ``geoarrow.wkb`` + extension name. DuckDB (>= 1.2, spatial loaded) ingests the column as + a native ``GEOMETRY`` with the CRS attached, so ``ST_Within(geometry, + ...)`` works with no ``ST_Point(x, y)`` construction in user SQL. +* ``"point"`` — GeoArrow native points with *separated* coordinates + (``struct`` under ``geoarrow.point``): the + child arrays are the coordinate columns themselves, no per-row + parsing for consumers that execute on native layouts (GeoPandas 1.x, + geoarrow-rs, lonboard, SedonaDB). DuckDB does not consume this + encoding. + +The CRS rides in the extension metadata (GeoArrow 0.2 allows +authority:code strings alongside PROJJSON). ``OGC:CRS84`` is the +correct tag for plain longitude/latitude grids. +""" + +from __future__ import annotations + +import json +from typing import Any + +import numpy as np +import pyarrow as pa + +GEOMETRY_COLUMN = "geometry" + +_ENCODINGS = ("wkb", "point") + + +def geometry_field(encoding: str, crs: str | None) -> pa.Field: + """The schema field for the derived geometry column.""" + if encoding not in _ENCODINGS: + raise ValueError( + f"geometry_encoding must be one of {_ENCODINGS}, got {encoding!r}" + ) + metadata = { + b"ARROW:extension:name": f"geoarrow.{encoding}".encode(), + } + if crs is not None: + metadata[b"ARROW:extension:metadata"] = json.dumps( + {"crs": crs} + ).encode() + storage = ( + pa.binary() + if encoding == "wkb" + else pa.struct([("x", pa.float64()), ("y", pa.float64())]) + ) + return pa.field(GEOMETRY_COLUMN, storage, metadata=metadata) + + +def build_geometry(encoding: str, x: pa.Array, y: pa.Array) -> pa.Array: + """Point geometries for one batch's x/y coordinate columns.""" + if encoding == "point": + return pa.StructArray.from_arrays( + [x.cast(pa.float64()), y.cast(pa.float64())], ["x", "y"] + ) + return _wkb_points( + np.ascontiguousarray(x.to_numpy(zero_copy_only=False), " pa.Array: + """Vectorized 21-byte little-endian WKB point encoding.""" + n = len(x) + # ``pa.binary()`` carries int32 offsets, which the final offset + # (n * 21) overflows past ~102M points; the buffers would build + # silently corrupt. Unreachable through the pivot (batch_size caps + # rows per batch well below this), so guard rather than widen the + # storage to large_binary, which DuckDB's ingestion expects not to + # see. + if n * 21 > np.iinfo(np.int32).max: + raise ValueError( + f"cannot WKB-encode {n:,} points in a single batch: " + "pa.binary() offsets are int32 and n * 21 bytes would " + "overflow them. Use a smaller batch_size." + ) + buf = np.empty((n, 21), dtype=np.uint8) + buf[:, 0] = 1 # little-endian byte order mark + buf[:, 1:5] = np.array([1, 0, 0, 0], dtype=np.uint8) # WKB type 1: Point + buf[:, 5:13] = x.view(np.uint8).reshape(n, 8) + buf[:, 13:21] = y.view(np.uint8).reshape(n, 8) + offsets = pa.py_buffer( + np.arange(0, (n + 1) * 21, 21, dtype=np.int32).tobytes() + ) + return pa.Array.from_buffers( + pa.binary(), n, [None, offsets, pa.py_buffer(buf.tobytes())] + ) + + +def bbox_conjuncts( + bounds: Any, x: str = "x", y: str = "y", pad: float = 0.0 +) -> str: + """SQL bbox conjuncts for a geometry's envelope — the pruning half. + + Engines do not push ``ST_*`` functions into the scan, so a + geometry-only predicate reads every chunk; pairing it with range + conjuncts on the coordinate columns restores pruning. This helper + renders those conjuncts from a geometry's envelope:: + + poly = shapely.from_wkt("POLYGON (...)") + sql = ( + f"SELECT avg(risk) FROM eri " + f"WHERE {xql.bbox_conjuncts(poly, x='x', y='y')} " + f"AND ST_Within(geometry, ST_GeomFromText('{poly.wkt}'))" + ) + + Args: + bounds: ``(xmin, ymin, xmax, ymax)``, or any object with a + ``.bounds`` attribute in that convention (shapely + geometries qualify). + x: The x/longitude column name. + y: The y/latitude column name. + pad: Optional margin added on every side (e.g. to be safe + around ``ST_DWithin``-style predicates). + + Returns: + A SQL snippet ``"x" BETWEEN a AND b AND "y" BETWEEN c AND d``. + """ + values = getattr(bounds, "bounds", bounds) + xmin, ymin, xmax, ymax = (float(v) for v in values) + return ( + f'"{x}" BETWEEN {xmin - pad!r} AND {xmax + pad!r} ' + f'AND "{y}" BETWEEN {ymin - pad!r} AND {ymax + pad!r}' + ) diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py new file mode 100644 index 00000000..cbf0551d --- /dev/null +++ b/xarray_sql/lazyscan.py @@ -0,0 +1,369 @@ +"""Re-executable engine handles behind the lazy chunked round-trip. + +The lazy path of ``to_dataset(chunks=...)`` re-executes the engine's +query per accessed chunk, narrowed to that chunk's coordinate window and +columns. That requires the engine result to be *re-executable* — a +handle onto the query, not a one-shot stream of its rows. Each handle +here adapts one engine's native lazy surface to the three operations the +reconstruction needs: + +* [schema][xarray_sql.lazyscan.LazyResultHandle.schema] — result column names/types, without + executing the query; +* [distinct][xarray_sql.lazyscan.LazyResultHandle.distinct] — one column's distinct values + (coordinate discovery; the caller sorts); +* [fetch][xarray_sql.lazyscan.LazyResultHandle.fetch] — the result narrowed by per-dimension + windows and projected to the requested columns, as Arrow batches. + +Windows are passed as [DimSpec][xarray_sql.lazyscan.DimSpec] values instead of rendered SQL so +each engine can express them with its own *typed* expression API — +strings would re-open every literal-formatting pitfall (timestamps, +floats, quoting) per dialect. + +Handles compose with the registration seam: when the wrapped query +scans a Dataset registered through xarray-sql's pushdown machinery, the +per-chunk range filter flows back through the engine into +[XarrayPushdownDataset][xarray_sql.backends.pyarrow.XarrayPushdownDataset], so each +output chunk's access reads only the source chunks it maps onto. +""" + +from __future__ import annotations + +import weakref +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Literal, Protocol, cast + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +from datafusion import col, literal + +DimSpec = tuple[Literal["range", "values"], Any, Any] +"""One dimension's window: ``("range", lo, hi)`` (inclusive bounds; the +requested coordinate positions are contiguous) or ``("values", array, +None)`` (explicit value list, for stepped/fancy indexers). + +The payload stays ``Any``: the values are coordinate scalars handed to +the engine's typed expression API, which does the comparing — Python +never orders them, and a concrete union over coordinate dtypes +(timestamps, cftime, numerics, strings) would stay incomplete.""" + + +def _collect_streaming(lf: Any) -> Any: + """Collect a Polars LazyFrame on the streaming engine, if available. + + ``collect(engine="streaming")`` needs polars >= 1.25 (the test + extra pins higher); an older installed polars raises TypeError on + the unknown keyword, and the plain in-memory collect is a correct, + if less memory-frugal, fallback. + """ + try: + return lf.collect(engine="streaming") + except TypeError: + return lf.collect() + + +def _plain(value: Any) -> Any: + """A plain-Python literal (numpy scalars don't travel to engines).""" + if isinstance(value, np.datetime64): + return pd.Timestamp(value) + if isinstance(value, np.timedelta64): + return pd.Timedelta(value) + if isinstance(value, np.generic): + return value.item() + return value + + +class LazyResultHandle(Protocol): + """A re-executable query result (see module docstring).""" + + supports_chunked: bool = True + """Whether fetch() may be driven from consumer worker threads (the + chunked reconstruction). Handles for engines that cannot safely + re-execute under foreign threads set this False; the eager path + remains available.""" + + def schema(self) -> pa.Schema: ... + + def distinct(self, column: str) -> np.ndarray: ... + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: ... + + def spill_parquet(self, path: str) -> None: ... + + # Handles may additionally offer ``stream(columns)``, yielding the + # unfiltered result's Arrow batches incrementally instead of + # materializing it first. The eager round-trip uses it to enforce + # ``max_result_bytes`` while collecting; handles without it refuse + # the budget rather than blow past it after collecting. + + +class DataFusionHandle: + """Handle over a ``datafusion.DataFrame``.""" + + supports_chunked = True + + def __init__(self, df: Any) -> None: + self._df = df + + def schema(self) -> pa.Schema: + return self._df.schema() + + def distinct(self, column: str) -> np.ndarray: + dim_only = self._df.select(col(f'"{column}"')).distinct() + batches = [b.to_pyarrow() for b in dim_only.execute_stream()] + if not batches: + return np.asarray([]) + return np.concatenate( + [b.column(0).to_numpy(zero_copy_only=False) for b in batches] + ) + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + predicate = None + for dim, (kind, a, b) in specs.items(): + c = col(f'"{dim}"') + if kind == "range": + p = (c >= literal(a)) & (c <= literal(b)) + else: + # DataFusion 52.0.0 exposes no clean ``Expr.in_list`` + # from Python; OR-chained equalities constant-fold + # equivalently and stay typed. + p = c == literal(a[0]) + for v in a[1:]: + p = p | (c == literal(v)) + predicate = p if predicate is None else predicate & p + out = self._df if predicate is None else self._df.filter(predicate) + out = out.select(*(col(f'"{n}"') for n in columns)) + return [b.to_pyarrow() for b in out.execute_stream()] + + def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: + """Execute once, yielding Arrow batches as the plan produces them.""" + out = self._df.select(*(col(f'"{n}"') for n in columns)) + return (b.to_pyarrow() for b in out.execute_stream()) + + def spill_parquet(self, path: str) -> None: + with pq.ParquetWriter(path, self.schema()) as writer: + for batch in self._df.execute_stream(): + writer.write_batch(batch.to_pyarrow()) + + +class DuckDBHandle: + """Handle over a ``duckdb.DuckDBPyRelation``. + + Relations are lazy relational algebra: ``filter``/``project`` derive + new relations and every materialization runs the query again, so a + single relation can serve any number of per-chunk fetches, each + narrowed to its own window — the *re-executable* property the + module docstring requires of every handle. Predicates are built + with DuckDB's typed expression API, never rendered SQL text. + + Every engine call runs on one dedicated thread owned by the handle. + A relation is bound to one connection, and a query over a table + registered through xarray-sql re-enters Python from DuckDB's + execution threads (the Arrow scan callback); driving such queries + directly from several consumer threads at once (dask computing + output chunks of a lazy round-trip) deadlocks between the + connection's serialization, the callback's need for the GIL, and + the consumer pool's own thread management. Funnelling execution + through a single pre-started thread reproduces the topology that is + known safe — one thread inside the engine, every other thread + parked on a GIL-releasing wait. + """ + + supports_chunked = False + """Chunked (lazy) reconstruction is disabled for DuckDB relations. + + Windows of a chunked round-trip re-execute the relation from the + consumer's worker threads (dask). A DuckDB query whose source is a + Python-callback Arrow scan (any table registered through xarray-sql) + intermittently deadlocks inside duckdb-python/CPython when other + Python threads start or stop during execution — reproduced on + duckdb 1.4-1.5 / CPython 3.12 / macOS at ~50% of runs, regardless + of ``SET threads=1``, connection serialization, or pool pre-warming. + Until that upstream race is fixed, chunked DuckDB round-trips fail + fast instead of hanging; the eager path (and every other handle + operation) runs on one dedicated thread and is unaffected. + """ + + def __init__(self, rel: Any) -> None: + self._rel = rel + self._runner = ThreadPoolExecutor(max_workers=1) + self._runner.submit(lambda: None).result() # start the thread now + # Stop the dedicated engine thread when the handle dies; it + # would otherwise linger for the life of the process, one per + # discarded handle. The callback is bound to the executor, not + # the handle, so the finalizer holds no reference back to self. + weakref.finalize( + self, self._runner.shutdown, wait=False, cancel_futures=True + ) + + def _run(self, fn: Any) -> Any: + return self._runner.submit(fn).result() + + @staticmethod + def _to_arrow_table(rel: Any) -> pa.Table: + if hasattr(rel, "to_arrow_table"): + return rel.to_arrow_table() + return rel.fetch_arrow_table() # duckdb < 1.5 + + @staticmethod + def _to_arrow_reader(rel: Any) -> pa.RecordBatchReader: + if hasattr(rel, "to_arrow_reader"): + return rel.to_arrow_reader() + return rel.fetch_record_batch() # duckdb < 1.5 + + def schema(self) -> pa.Schema: + return self._run( + lambda: self._to_arrow_table(self._rel.limit(0)).schema + ) + + def distinct(self, column: str) -> np.ndarray: + import duckdb + + table = self._run( + lambda: self._to_arrow_table( + self._rel.project(duckdb.ColumnExpression(column)).distinct() + ) + ) + return np.asarray(table.column(0).to_numpy(zero_copy_only=False)) + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + import duckdb + + predicate = None + for dim, (kind, a, b) in specs.items(): + c = duckdb.ColumnExpression(dim) + if kind == "range": + p = (c >= duckdb.ConstantExpression(_plain(a))) & ( + c <= duckdb.ConstantExpression(_plain(b)) + ) + else: + p = c.isin(*(duckdb.ConstantExpression(_plain(v)) for v in a)) + predicate = p if predicate is None else predicate & p + rel = self._rel if predicate is None else self._rel.filter(predicate) + rel = rel.project(*(duckdb.ColumnExpression(n) for n in columns)) + + return cast( + list[pa.RecordBatch], + self._run(lambda: list(self._to_arrow_reader(rel))), + ) + + def spill_parquet(self, path: str) -> None: + def run() -> None: + reader = self._to_arrow_reader(self._rel) + with pq.ParquetWriter(path, reader.schema) as writer: + for batch in reader: + writer.write_batch(batch) + + self._run(run) + + +class PolarsHandle: + """Handle over a ``polars.LazyFrame``. + + Per-window fetches run on the streaming engine, so a window read + never materializes more than the window even when the frame scans + an out-of-core source. + """ + + supports_chunked = True + + def __init__(self, lf: Any) -> None: + self._lf = lf + + def schema(self) -> pa.Schema: + import polars as pl + + return pl.DataFrame(schema=self._lf.collect_schema()).to_arrow().schema + + def distinct(self, column: str) -> np.ndarray: + import polars as pl + + out = _collect_streaming(self._lf.select(pl.col(column).unique())) + return out.to_series().to_numpy() + + def fetch( + self, specs: dict[str, DimSpec], columns: list[str] + ) -> list[pa.RecordBatch]: + import polars as pl + + exprs = [] + for dim, (kind, a, b) in specs.items(): + if kind == "range": + exprs.append(pl.col(dim).is_between(_plain(a), _plain(b))) + elif getattr(a, "dtype", None) is not None and a.dtype.kind == "f": + # Upstream Polars translates float ``is_in`` literals + # imprecisely (silently matching nothing); degenerate + # ranges compare exactly. Reproduced on polars 1.42. + # ``any_horizontal`` keeps the disjunction flat — a + # left-deep OR chain plans quadratically in the number + # of values. + exprs.append( + pl.any_horizontal( + [pl.col(dim).is_between(*(_plain(v),) * 2) for v in a] + ) + ) + else: + exprs.append(pl.col(dim).is_in([_plain(v) for v in a])) + lf = self._lf.filter(*exprs) if exprs else self._lf + out = _collect_streaming(lf.select([pl.col(n) for n in columns])) + return cast(list[pa.RecordBatch], out.to_arrow().to_batches()) + + def stream(self, columns: list[str]) -> Iterator[pa.RecordBatch]: + """Execute once, yielding Arrow batches incrementally. + + Unlike [fetch][xarray_sql.lazyscan.PolarsHandle.fetch], whose ``collect`` materializes the whole + result inside the engine before any batch surfaces, this yields + batches as the streaming engine produces them, so a byte budget + can fire before the result is fully in memory. Requires + ``LazyFrame.collect_batches`` (polars >= 1.33); older Polars + raises here, before anything is collected. + """ + import polars as pl + + lf = self._lf.select([pl.col(n) for n in columns]) + if not hasattr(lf, "collect_batches"): + raise ValueError( + "streaming collection of a Polars LazyFrame requires " + "polars >= 1.33 (LazyFrame.collect_batches)." + ) + + def generate() -> Iterator[pa.RecordBatch]: + for frame in lf.collect_batches(engine="streaming"): + yield from frame.to_arrow().to_batches() + + return generate() + + def spill_parquet(self, path: str) -> None: + self._lf.sink_parquet(path) + + +def resolve_lazy_handle(result: Any) -> LazyResultHandle | None: + """Adapt an engine result to a handle, or ``None`` if it is one-shot. + + Recognizes DuckDB relations, Polars lazy *and* eager frames (an + eager frame re-executes trivially over its in-memory data), and + DataFusion DataFrames. ``pyarrow`` tables/readers and bare + ``__arrow_c_stream__`` objects are one-shot streams: there is no + query to re-execute, so the lazy path cannot serve them. + """ + root = type(result).__module__.split(".")[0] + if root in ("duckdb", "_duckdb") and hasattr(result, "filter"): + return DuckDBHandle(result) + if root == "polars": + import polars as pl + + if isinstance(result, pl.LazyFrame): + return PolarsHandle(result) + if isinstance(result, pl.DataFrame): + return PolarsHandle(result.lazy()) + if hasattr(result, "execute_stream") and hasattr(result, "logical_plan"): + return DataFusionHandle(result) + return None diff --git a/xarray_sql/proj.py b/xarray_sql/proj.py new file mode 100644 index 00000000..2470d1b8 --- /dev/null +++ b/xarray_sql/proj.py @@ -0,0 +1,235 @@ +"""PROJ-backed CRS transforms for SQL — the optional geo extension. + +Geospatial SQL dialects expose coordinate reference system (CRS) +transforms as a scalar function — PostGIS and DuckDB-spatial both call it +``ST_Transform`` — because a CRS transform is row-independent: each +point's new coordinate depends only on its own old coordinate. This +module brings the same capability to xarray-sql as a vectorized scalar +UDF over Arrow arrays:: + + SELECT x, y, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon, + reproject(x, y, 'EPSG:32610', 'EPSG:4326')['y'] AS lat + FROM grid + +The CRS pair is part of the *query*, not baked in at registration time, +so one registered UDF serves any transform — and, because the arguments +are ordinary SQL expressions, the CRS may even vary per row (e.g. a +``CASE`` expression selecting the UTM zone from the longitude). + +Design notes: + +* **Both output coordinates come from one call**, returned as an Arrow + struct ``{x, y}`` (in ``always_xy`` order: easting/longitude first). + Splitting the transform into two scalar UDFs would run PROJ twice per + row and, worse, evaluate the two projections concurrently on separate + expression trees. +* **All pyproj work runs on a dedicated pool of Python threads.** + DataFusion's runtime workers are not Python-created threads, and + pyproj (< 3.8, see pyproj#1541) leaves a dangling ``PJ_CONTEXT`` + behind when their ephemeral Python thread states are torn down, so + calling pyproj in place segfaults — the UDF hands each batch to the + pool instead. Pool threads are long-lived, so each caches one + transformer per CRS pair (transformers must not be shared across + threads), amortizing the expensive construction — PROJ database + lookups and candidate-operation selection — across record batches. + Concurrent partitions still transform in parallel across the pool. +* Any CRS spelling ``pyproj.CRS`` accepts works: authority codes + (``EPSG:4326``), WKT, PROJ strings (``+proj=utm +zone=10``), etc. + An unknown CRS raises ``pyproj.exceptions.CRSError`` and fails the + query loudly rather than returning wrong coordinates. +* Non-finite or NULL input coordinates yield NaN output (PROJ itself + would return ``inf``); NULL CRS arguments yield NaN as well. + +Requires ``pyproj`` (``pip install xarray-sql[geo]``). When pyproj is +installed, [xarray_sql.XarrayContext][] registers ``reproject()`` +automatically; [register][xarray_sql.proj.register] is the explicit hook for plain +DataFusion ``SessionContext`` objects or custom UDF names. +""" + +from __future__ import annotations + +import os +import threading +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyproj +from datafusion import udf + +__all__ = ["register"] + +RETURN_TYPE = pa.struct([("x", pa.float64()), ("y", pa.float64())]) +"""Arrow type returned by ``reproject()``: destination coordinates in +``always_xy`` order — ``x`` is easting/longitude, ``y`` is +northing/latitude.""" + + +# --------------------------------------------------------------------------- +# The PROJ worker pool +# --------------------------------------------------------------------------- +# +# DataFusion evaluates UDFs on its tokio runtime's worker threads, which +# are not created by Python: a Python thread state is created and +# destroyed around every UDF call. pyproj keeps its per-thread PJ_CONTEXT +# in CPython thread-specific storage but does not clear that pointer when +# the context dies with the ephemeral thread state (fixed by pyproj#1541, +# unreleased as of 3.7.2), so the next call on the same OS thread +# dereferences a dangling context and segfaults inside ``proj_create``. +# Python-owned threads keep their thread state — and thus their contexts — +# alive for the thread's lifetime, so the UDF never calls pyproj in place: +# every batch is handed to a small pool of Python-owned worker threads. +# The pool stays worthwhile on fixed pyproj too: ephemeral thread states +# would rebuild context and transformer per batch (0.07–12 ms measured) +# versus ~10 µs for the pool round-trip. pyproj releases the GIL during +# the transform loop, so concurrent partitions still run in parallel +# across the pool. + +_local = threading.local() +_pool_lock = threading.Lock() +_pool: ThreadPoolExecutor | None = None + + +def _proj_pool() -> ThreadPoolExecutor: + """Return the process-wide pool that runs all pyproj work.""" + global _pool + if _pool is None: + with _pool_lock: + if _pool is None: + _pool = ThreadPoolExecutor( + max_workers=os.cpu_count() or 4, + thread_name_prefix="xarray-sql-proj", + ) + return _pool + + +def _transformer(src_crs: str, dst_crs: str) -> pyproj.Transformer: + """Return a cached ``Transformer`` owned by the calling pool thread. + + PROJ transformers are not safe to share across threads, so each pool + thread keeps its own transformer per ``(src, dst)`` pair; the cache + also amortizes construction (expensive PROJ database lookups) across + record batches. ``always_xy=True`` fixes the argument order to + (easting/longitude, northing/latitude) regardless of the CRS's + declared axis order. + """ + cache = getattr(_local, "transformers", None) + if cache is None: + cache = _local.transformers = {} + key = (src_crs, dst_crs) + transformer = cache.get(key) + if transformer is None: + transformer = cache[key] = pyproj.Transformer.from_crs( + src_crs, dst_crs, always_xy=True + ) + return transformer + + +def _transform_chunk( + src_crs: str, dst_crs: str, xs: np.ndarray, ys: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Transform one coordinate chunk; runs on a PROJ pool thread.""" + tx, ty = _transformer(src_crs, dst_crs).transform(xs, ys) + return tx, ty + + +# --------------------------------------------------------------------------- +# The UDF +# --------------------------------------------------------------------------- + + +def _reproject( + x: pa.Array, y: pa.Array, src_crs: pa.Array, dst_crs: pa.Array +) -> pa.Array: + """Vectorized ``reproject`` kernel over one Arrow record batch. + + DataFusion broadcasts scalar arguments (the usual literal CRS + strings) to full-length arrays before calling in, so all four + arguments arrive with one value per row. The common case — one CRS + pair for the whole batch — never touches the strings row by row: + uniqueness is established with a vectorized Arrow kernel and the + batch becomes a single PROJ call. (Materializing the CRS columns + as Python strings costs two object allocations per row, which at + billions of rows dwarfs the transform itself.) Only when the CRS + genuinely varies within the batch are rows grouped by pair and + transformed per group. + """ + # Zero-copy read-only views when the batch has no nulls; with nulls, + # pyarrow must materialize the validity bitmap as NaN (the NULL -> NaN + # contract). pyproj copies into its own writable buffer either way -- + # PROJ mutates buffers in place -- so this is the minimal-copy path. + xs = np.asarray(x.to_numpy(zero_copy_only=False), dtype="float64") + ys = np.asarray(y.to_numpy(zero_copy_only=False), dtype="float64") + + out_x = np.full(xs.shape, np.nan) + out_y = np.full(ys.shape, np.nan) + valid = np.isfinite(xs) & np.isfinite(ys) + src_unique = pc.unique(src_crs) + dst_unique = pc.unique(dst_crs) + + if len(src_unique) == 1 and len(dst_unique) == 1: + groups = [(src_unique[0].as_py(), dst_unique[0].as_py(), valid)] + else: + pairs = list(zip(src_crs.to_pylist(), dst_crs.to_pylist())) + groups = [ + ( + src, + dst, + valid + & np.fromiter( + (p == (src, dst) for p in pairs), + dtype=bool, + count=len(pairs), + ), + ) + for src, dst in set(pairs) + ] + + for src, dst, mask in groups: + if src is None or dst is None or not mask.any(): + continue + tx, ty = ( + _proj_pool() + .submit(_transform_chunk, src, dst, xs[mask], ys[mask]) + .result() + ) + out_x[mask] = tx + out_y[mask] = ty + + # PROJ signals out-of-domain points with inf; normalize to NaN so + # the result round-trips to xarray like any other missing value. + invalid = ~(np.isfinite(out_x) & np.isfinite(out_y)) + out_x[invalid] = np.nan + out_y[invalid] = np.nan + + return pa.StructArray.from_arrays( + [pa.array(out_x), pa.array(out_y)], names=["x", "y"] + ) + + +def register(ctx, name: str = "reproject") -> None: + """Register the ``reproject(x, y, src_crs, dst_crs)`` scalar UDF. + + Works on any DataFusion ``SessionContext`` (``XarrayContext`` + registers it automatically when pyproj is installed). The UDF + returns a ``{x, y}`` struct of destination coordinates, so a query + selects components with subscripts:: + + SELECT reproject(x, y, 'EPSG:32610', 'EPSG:4326')['x'] AS lon + FROM grid + + Args: + ctx: The DataFusion session context to register the UDF on. + name: SQL name for the function (default ``"reproject"``). + """ + ctx.register_udf( + udf( + _reproject, + [pa.float64(), pa.float64(), pa.utf8(), pa.utf8()], + RETURN_TYPE, + "immutable", + name, + ) + ) diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index f89cd276..153ed8c3 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -10,211 +10,331 @@ from __future__ import annotations -import typing as t +from collections.abc import Callable, Iterator +from typing import TYPE_CHECKING +import numpy as np import pyarrow as pa import xarray as xr -from .df import Block, Chunks, block_slices, pivot, _parse_schema +from .df import ( + Block, + Chunks, + DEFAULT_BATCH_SIZE, + _block_len, + _block_metadata, + _block_slices_from_resolved, + _ensure_default_indexes, + _parse_schema, + block_slices, + iter_record_batches, + resolve_chunks, +) -if t.TYPE_CHECKING: - from ._native import LazyArrowStreamTable +if TYPE_CHECKING: + from ._native import LazyArrowStreamTable class XarrayRecordBatchReader: - """A lazy Arrow stream reader for xarray Datasets. - - Implements the Arrow PyCapsule Interface (__arrow_c_stream__) to enable - zero-copy, lazy streaming of xarray data to DataFusion and other Arrow - consumers. - - The key property is that xarray blocks are only converted to Arrow - RecordBatches when the consumer calls get_next (e.g., during DataFusion's - collect()), NOT when the reader is created or registered. - - Attributes: - schema: The Arrow schema for the stream. - - Example: - >>> import xarray as xr - >>> from xarray_sql import XarrayRecordBatchReader - >>> ds = xr.tutorial.open_dataset('air_temperature') - >>> reader = XarrayRecordBatchReader(ds, chunks={'time': 240}) - >>> # At this point, NO data has been read from xarray - >>> # Data is only read when consumed: - >>> import pyarrow as pa - >>> pa_reader = pa.RecordBatchReader.from_stream(reader) - >>> for batch in pa_reader: - ... print(batch.num_rows) # Data read here - """ - - def __init__( - self, - ds: xr.Dataset, - chunks: Chunks = None, - *, - _iteration_callback: t.Optional[t.Callable[[Block], None]] = None, - ): - """Initialize the lazy reader. - - Args: - ds: An xarray Dataset. All data_vars must share the same dimensions. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - _iteration_callback: Internal callback for testing. Called with - each block dict just before it's converted to Arrow. This - allows tests to track when iteration actually occurs. - """ - self._ds = ds - self._chunks = chunks - self._schema = _parse_schema(ds) - self._iteration_callback = _iteration_callback - self._consumed = False - - # Validate dimensions - fst = next(iter(ds.values())).dims - if not all(da.dims == fst for da in ds.values()): - raise ValueError( - "All dimensions must be equal. Please filter data_vars in the Dataset." - ) - - @property - def schema(self) -> pa.Schema: - """The Arrow schema for this stream.""" - return self._schema - - def _generate_batches(self) -> t.Iterator[pa.RecordBatch]: - """Generate RecordBatches lazily from xarray blocks. - - This generator is only consumed when the Arrow stream's get_next - is called, ensuring true lazy evaluation. + """A lazy Arrow stream reader for xarray Datasets. + + Implements the Arrow PyCapsule Interface (__arrow_c_stream__) to enable + zero-copy, lazy streaming of xarray data to DataFusion and other Arrow + consumers. + + The key property is that xarray blocks are only converted to Arrow + RecordBatches when the consumer calls get_next (e.g., during DataFusion's + collect()), NOT when the reader is created or registered. + + Attributes: + schema: The Arrow schema for the stream. + + Example: + >>> import xarray as xr + >>> from xarray_sql import XarrayRecordBatchReader + >>> ds = xr.tutorial.open_dataset('air_temperature') + >>> reader = XarrayRecordBatchReader(ds, chunks={'time': 240}) + >>> # At this point, NO data has been read from xarray + >>> # Data is only read when consumed: + >>> import pyarrow as pa + >>> pa_reader = pa.RecordBatchReader.from_stream(reader) + >>> for batch in pa_reader: + ... print(batch.num_rows) # Data read here """ - for block in block_slices(self._ds, self._chunks): - # Call the iteration callback if provided (for testing) - if self._iteration_callback is not None: - self._iteration_callback(block) - # Convert this block to a RecordBatch - df = pivot(self._ds.isel(block)) - yield pa.RecordBatch.from_pandas(df, schema=self._schema) + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + """Initialize the lazy reader. + + Args: + ds: An xarray Dataset. All data_vars must share the same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + batch_size: Maximum rows per emitted Arrow RecordBatch. Smaller + values let DataFusion start processing earlier at the cost of + more Python→Arrow conversion calls. + _iteration_callback: Internal callback for testing. Called with + each block dict just before it's converted to Arrow. This + allows tests to track when iteration actually occurs. + """ + self._ds = ds + self._chunks = chunks + self._batch_size = batch_size + self._schema = _parse_schema(ds) + self._iteration_callback = _iteration_callback + self._consumed = False + + # Validate dimensions + fst = next(iter(ds.values())).dims + if not all(da.dims == fst for da in ds.values()): + raise ValueError( + "All dimensions must be equal. Please filter data_vars in the Dataset." + ) + + @property + def schema(self) -> pa.Schema: + """The Arrow schema for this stream.""" + return self._schema + + def _generate_batches(self) -> Iterator[pa.RecordBatch]: + """Generate RecordBatches lazily from xarray blocks. + + This generator is only consumed when the Arrow stream's get_next + is called, ensuring true lazy evaluation. Each xarray block is + emitted as one or more RecordBatches of at most self._batch_size rows. + """ + for block in block_slices(self._ds, self._chunks): + # Call the iteration callback if provided (for testing). + # XarrayRecordBatchReader has no projection concept, so always passes None. + if self._iteration_callback is not None: + self._iteration_callback(block, None) + + yield from iter_record_batches( + self._ds.isel(block), self._schema, self._batch_size + ) + + def __arrow_c_stream__( + self, requested_schema: object | None = None + ) -> object: + """Export as Arrow C Stream via PyCapsule. + + This method is called by Arrow consumers (like DataFusion) to get + a C-level stream interface. The actual data iteration only begins + when the consumer calls get_next on the stream. + + Args: + requested_schema: Optional schema for type casting. Currently + passed through to PyArrow's implementation. + + Returns: + PyCapsule containing ArrowArrayStream pointer with name + "arrow_array_stream". + + Raises: + RuntimeError: If the stream has already been consumed. + """ + if self._consumed: + raise RuntimeError( + "Stream already consumed. XarrayRecordBatchReader can only " + "be iterated once. Create a new reader for additional iterations." + ) + self._consumed = True + + # Create a PyArrow RecordBatchReader from our generator + # The generator is NOT consumed here - only when get_next is called + reader = pa.RecordBatchReader.from_batches( + self._schema, self._generate_batches() + ) + + # Delegate to PyArrow's __arrow_c_stream__ implementation + return reader.__arrow_c_stream__(requested_schema) + + def __arrow_c_schema__( + self, requested_schema: object | None = None + ) -> object: + """Export the schema as Arrow C Schema via PyCapsule. + + This allows consumers to inspect the schema without consuming the stream. + + Args: + requested_schema: Optional schema for negotiation (unused). + + Returns: + PyCapsule containing ArrowSchema pointer. + """ + return self._schema.__arrow_c_schema__() - def __arrow_c_stream__( - self, requested_schema: t.Optional[object] = None - ) -> object: - """Export as Arrow C Stream via PyCapsule. - This method is called by Arrow consumers (like DataFusion) to get - a C-level stream interface. The actual data iteration only begins - when the consumer calls get_next on the stream. +def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: + """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. Args: - requested_schema: Optional schema for type casting. Currently - passed through to PyArrow's implementation. + ds: An Xarray Dataset. All `data_vars` must share the same dimensions. + chunks: Xarray-like chunks. If not provided, will default to the + Dataset's chunks. The product of the chunk sizes becomes the + standard length of each dataframe partition. Returns: - PyCapsule containing ArrowArrayStream pointer with name - "arrow_array_stream". - - Raises: - RuntimeError: If the stream has already been consumed. + A PyArrow RecordBatchReader, which is a table representation of the input + Dataset. """ - if self._consumed: - raise RuntimeError( - "Stream already consumed. XarrayRecordBatchReader can only " - "be iterated once. Create a new reader for additional iterations." - ) - self._consumed = True - - # Create a PyArrow RecordBatchReader from our generator - # The generator is NOT consumed here - only when get_next is called - reader = pa.RecordBatchReader.from_batches( - self._schema, self._generate_batches() - ) - - # Delegate to PyArrow's __arrow_c_stream__ implementation - return reader.__arrow_c_stream__(requested_schema) + ds = _ensure_default_indexes(ds) + reader = XarrayRecordBatchReader(ds, chunks=chunks) + return pa.RecordBatchReader.from_stream(reader) - def __arrow_c_schema__( - self, requested_schema: t.Optional[object] = None - ) -> object: - """Export the schema as Arrow C Schema via PyCapsule. - This allows consumers to inspect the schema without consuming the stream. - - Args: - requested_schema: Optional schema for negotiation (unused). +def read_xarray_table( + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + coord_arrays: dict[str, np.ndarray] | None = None, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, +) -> "LazyArrowStreamTable": + """Create a lazy DataFusion table from an xarray Dataset. - Returns: - PyCapsule containing ArrowSchema pointer. - """ - return self._schema.__arrow_c_schema__() + This is the simplest way to register xarray data with DataFusion. + Data is only read when queries are executed, not during registration. + The table can be queried multiple times. + Each chunk becomes a separate partition, enabling DataFusion's parallel + execution across multiple cores. -def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: - """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. + Note: + SQL queries with WHERE clauses on dimension columns (time, lat, lon, etc.) + automatically prune partitions that can't contain matching rows — this is + called *filter pushdown*. For example: - Args: - ds: An Xarray Dataset. All `data_vars` mush share the same dimensions. - chunks: Xarray-like chunks. If not provided, will default to the Dataset's - chunks. The product of the chunk sizes becomes the standard length of each - dataframe partition. + # This query will skip loading partitions with time < '2020-02-01' + result = ctx.sql('SELECT * FROM air WHERE time > "2020-02-01"').collect() - Returns: - A PyArrow RecordBatchReader, which is a table representation of the input - Dataset. - """ - reader = XarrayRecordBatchReader(ds, chunks=chunks) - return pa.RecordBatchReader.from_stream(reader) + Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. + Args: + ds: An xarray Dataset. All data_vars must share the same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + batch_size: Maximum rows per Arrow RecordBatch emitted per partition. + Smaller values let DataFusion start processing earlier; the default + (65 536) works well for most datasets. + coord_arrays: Pre-materialised coordinate arrays keyed by dim-name + string. Hand in to share a single read across multiple tables + built from the same parent Dataset (e.g. surface + atmosphere + from ARCO-ERA5); the dim coords are otherwise read once per + ``read_xarray_table`` call, which is a network round-trip for + Zarr-backed datasets. + _iteration_callback: Internal callback for testing. Called with + each block dict just before it's converted to Arrow. -def read_xarray_table( - ds: xr.Dataset, - chunks: Chunks = None, - *, - _iteration_callback: t.Optional[t.Callable[[Block], None]] = None, -) -> "LazyArrowStreamTable": - """Create a lazy DataFusion table from an xarray Dataset. - - This is the simplest way to register xarray data with DataFusion. - Data is only read when queries are executed (during collect()), - not during registration. The table can be queried multiple times. - - Args: - ds: An xarray Dataset. All data_vars must share the same dimensions. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - _iteration_callback: Internal callback for testing. Called with - each block dict just before it's converted to Arrow. - - Returns: - A LazyArrowStreamTable ready for registration with DataFusion. - - Example: - >>> from datafusion import SessionContext - >>> import xarray as xr - >>> from xarray_sql import read_xarray_table - >>> - >>> ds = xr.tutorial.open_dataset('air_temperature') - >>> table = read_xarray_table(ds, chunks={'time': 240}) - >>> - >>> ctx = SessionContext() - >>> ctx.register_table('air', table) - >>> - >>> # Data is only read here, during collect() - >>> result = ctx.sql('SELECT AVG(air) FROM air').collect() - >>> # Can query again - each query creates a fresh stream - >>> result2 = ctx.sql('SELECT * FROM air LIMIT 10').collect() - """ - from ._native import LazyArrowStreamTable - - # Get schema from dataset without creating a stream - schema = _parse_schema(ds) - - # Create a factory function that produces fresh RecordBatchReaders on each call - def make_stream() -> pa.RecordBatchReader: - stream = XarrayRecordBatchReader( - ds, chunks, _iteration_callback=_iteration_callback + Returns: + A LazyArrowStreamTable ready for registration with DataFusion. + + Example: + >>> from datafusion import SessionContext + >>> import xarray as xr + >>> from xarray_sql import read_xarray_table + >>> + >>> ds = xr.tutorial.open_dataset('air_temperature') + >>> table = read_xarray_table(ds, chunks={'time': 240}) + >>> + >>> ctx = SessionContext() + >>> ctx.register_table('air', table) + >>> + >>> # Data is only read here, during query execution + >>> # Filters on 'time' will prune partitions automatically! + >>> result = ctx.sql('SELECT AVG(air) FROM air').collect() + """ + from ._native import LazyArrowStreamTable + + ds = _ensure_default_indexes(ds) + schema = _parse_schema(ds) + + # Hoist coordinate reads once; avoids N_partitions remote I/O calls for + # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS). When the caller supplies + # pre-materialised arrays (e.g. shared across surface + atmosphere + # tables), reuse them and skip the extra read. + if coord_arrays is None: + coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} + + # Determine which column names are data variables (not dimension coordinates). + # Used by the factory to skip loading unrequested variables. + data_var_names = set(ds.data_vars.keys()) + + def make_partition_factory( + block: Block, + ) -> Callable[[list[str] | None], pa.RecordBatchReader]: + def make_stream( + projection_names: list[str] | None, + ) -> pa.RecordBatchReader: + if _iteration_callback is not None: + _iteration_callback(block, projection_names) + + if projection_names is not None: + # Restrict to the data variables mentioned in the projection. + # Dimension coordinates come along automatically via coords. + data_vars_needed = [ + c for c in projection_names if c in data_var_names + ] + if data_vars_needed: + ds_block = ds[data_vars_needed].isel(block) + else: + # Only dimension coords requested — drop all data vars to avoid + # loading them unnecessarily (e.g. for queries like SELECT lat, lon). + ds_block = ds.drop_vars(list(ds.data_vars)).isel(block) + batch_schema = pa.schema( + [schema.field(name) for name in projection_names] + ) + else: + ds_block = ds.isel(block) + batch_schema = schema + + return pa.RecordBatchReader.from_batches( + batch_schema, + iter_record_batches(ds_block, batch_schema, batch_size), + ) + + return make_stream + + # Separate dims whose chunk bounds vary across partitions from those + # whose bounds are constant (one chunk spanning the whole axis). For the + # latter we compute min/max once instead of re-scanning the full coord + # array on every partition — dominant cost when registering hundreds of + # thousands of single-time-step partitions on a 4-D dataset like ERA5. + resolved = resolve_chunks(ds, chunks) + varying_dims = [d for d, tup in resolved.items() if len(tup) > 1] + static_dims = [d for d in ds.dims if d not in varying_dims] + static_block: Block = {d: slice(None) for d in static_dims} + static_ranges = _block_metadata( + coord_arrays, static_block, dims=static_dims ) - return pa.RecordBatchReader.from_stream(stream) - return LazyArrowStreamTable(make_stream, schema) + def partition_pairs(): + """Lazily yield (factory, metadata, num_rows) for each partition. + + Consuming this generator one item at a time means Python never holds + all N block dicts, metadata dicts, and factory closures simultaneously. + Peak Python memory during registration is O(1) per partition instead + of O(N_partitions). + """ + for block in _block_slices_from_resolved(ds, resolved): + dynamic = _block_metadata(coord_arrays, block, dims=varying_dims) + yield ( + make_partition_factory(block), + {**static_ranges, **dynamic}, + # Exact row count for this partition (product of the chunk's + # per-dimension sizes), so the scan can report exact + # Statistics::num_rows to the optimizer. + _block_len(block), + ) + + return LazyArrowStreamTable(partition_pairs(), schema) diff --git a/xarray_sql/reader_test.py b/xarray_sql/reader_test.py deleted file mode 100644 index 1527ba61..00000000 --- a/xarray_sql/reader_test.py +++ /dev/null @@ -1,944 +0,0 @@ -"""Tests for XarrayRecordBatchReader lazy streaming behavior. - -These tests verify that XarrayRecordBatchReader provides true lazy evaluation: -- No data iteration during reader creation -- No data iteration during DataFusion table registration (using LazyArrowStreamTable) -- Data iteration ONLY occurs during query execution (collect()) - -The lazy streaming is achieved via the Rust LazyArrowStreamTable class which -implements the __datafusion_table_provider__ protocol using StreamingTable. - -Additional tests verify: -- True streaming with bounded memory (batches processed incrementally) -- Back-pressure behavior (producer pauses when consumer is slow) -- Error propagation through the stream -""" - -import threading -import time -import numpy as np -import pandas as pd -import pyarrow as pa -import pytest -import xarray as xr -from datafusion import SessionContext - -from ._native import LazyArrowStreamTable -from .reader import XarrayRecordBatchReader, read_xarray_table -from .df import _parse_schema - - -@pytest.fixture -def small_ds(): - """Create a small dataset for testing.""" - np.random.seed(42) - time = pd.date_range("2020-01-01", periods=100, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) - - data = np.random.rand(100, 10, 10).astype(np.float32) - - return xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time, "lat": lat, "lon": lon}, - ) - - -class IterationTracker: - """Tracks when iteration occurs for testing lazy evaluation.""" - - def __init__(self): - self.iteration_count = 0 - self.blocks_seen = [] - - def __call__(self, block): - self.iteration_count += 1 - self.blocks_seen.append(block) - - def reset(self): - self.iteration_count = 0 - self.blocks_seen = [] - - -class TestXarrayRecordBatchReaderCreation: - """Tests that reader creation does NOT trigger data iteration.""" - - def test_reader_creation_does_not_iterate(self, small_ds): - """Creating a reader should NOT iterate through any data.""" - tracker = IterationTracker() - - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations during reader creation, " - f"but got {tracker.iteration_count}" - ) - - def test_schema_access_does_not_iterate(self, small_ds): - """Accessing the schema should NOT trigger iteration.""" - tracker = IterationTracker() - - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - # Access schema - _ = reader.schema - _ = reader.__arrow_c_schema__() - - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations when accessing schema, " - f"but got {tracker.iteration_count}" - ) - - -class TestDataFusionRegistration: - """Tests that DataFusion table registration does NOT trigger iteration. - - These tests use read_xarray_table with register_table() - to achieve true lazy evaluation. - """ - - def test_register_table_does_not_iterate(self, small_ds): - """Registering a LazyArrowStreamTable should NOT iterate data. - - This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps - a factory and implements __datafusion_table_provider__ with StreamingTable, - ensuring data is only read during query execution. - """ - tracker = IterationTracker() - - # Use read_xarray_table which creates a factory-based table - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - assert tracker.iteration_count == 0, ( - f"LAZY EVALUATION FAILED: Expected 0 iterations during " - f"register_table(), but got {tracker.iteration_count}." - ) - - def test_sql_planning_does_not_iterate(self, small_ds): - """Creating a SQL query plan should NOT iterate data.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Create a query but don't execute it - query = ctx.sql("SELECT AVG(temperature) FROM test_table") - - # Just creating the query shouldn't iterate - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations during SQL planning, " - f"but got {tracker.iteration_count}. " - f"DataFusion may be scanning data during query planning." - ) - - -class TestDataFusionCollect: - """Tests that data iteration ONLY occurs during collect(). - - These tests use read_xarray_table to verify lazy evaluation. - """ - - def test_collect_triggers_iteration(self, small_ds): - """collect() should trigger data iteration.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Verify no iteration yet (lazy registration) - iteration_before_collect = tracker.iteration_count - assert ( - iteration_before_collect == 0 - ), "Should have 0 iterations before collect" - - # Now collect - this SHOULD iterate - result = ctx.sql("SELECT * FROM test_table LIMIT 10").collect() - - assert tracker.iteration_count > 0, ( - f"Expected iterations during collect(), but got 0. " - f"Data was never read!" - ) - assert ( - tracker.iteration_count > iteration_before_collect - ), f"Expected more iterations after collect()" - - def test_full_query_iterates_all_blocks(self, small_ds): - """A query that reads all data should iterate all blocks.""" - tracker = IterationTracker() - - chunks = {"time": 25} - table = read_xarray_table( - small_ds, - chunks=chunks, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Run a query that needs to scan all data - result = ctx.sql("SELECT COUNT(*) FROM test_table").collect() - - # With time=100 and chunks=25, we expect 4 blocks - expected_blocks = 100 // 25 - assert tracker.iteration_count == expected_blocks, ( - f"Expected {expected_blocks} block iterations, " - f"but got {tracker.iteration_count}" - ) - - def test_aggregation_query_iterates_correctly(self, small_ds): - """Aggregation queries should iterate all necessary blocks.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Run aggregation - result = ctx.sql( - "SELECT lat, AVG(temperature) as avg_temp " - "FROM test_table GROUP BY lat" - ).collect() - - # Should have iterated some blocks - assert tracker.iteration_count > 0 - assert len(result) > 0 - - -class TestLazyEvaluationEndToEnd: - """End-to-end tests verifying lazy evaluation through the full pipeline. - - These tests use read_xarray_table to achieve true lazy evaluation. - """ - - def test_lazy_evaluation_sequence(self, small_ds): - """Verify the exact sequence of lazy evaluation stages. - - This is the comprehensive test that proves true lazy evaluation: - 1. Table creation: 0 iterations - 2. Table registration: 0 iterations - 3. Query planning: 0 iterations - 4. collect(): N iterations (where N = number of blocks) - """ - tracker = IterationTracker() - - # Stage 1: Table creation (with factory) - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - iterations_after_table = tracker.iteration_count - assert iterations_after_table == 0, ( - f"Stage 1 FAILED: Table creation triggered " - f"{iterations_after_table} iterations" - ) - - # Stage 2: Table registration - ctx = SessionContext() - ctx.register_table("test_table", table) - iterations_after_registration = tracker.iteration_count - assert iterations_after_registration == 0, ( - f"Stage 2 FAILED: Table registration triggered " - f"{iterations_after_registration} iterations" - ) - - # Stage 3: Query planning - query = ctx.sql("SELECT * FROM test_table") - iterations_after_planning = tracker.iteration_count - assert iterations_after_planning == 0, ( - f"Stage 3 FAILED: Query planning triggered " - f"{iterations_after_planning} iterations" - ) - - # Stage 4: collect() - NOW iteration should happen - result = query.collect() - iterations_after_collect = tracker.iteration_count - assert ( - iterations_after_collect > 0 - ), f"Stage 4 FAILED: collect() triggered 0 iterations - no data was read!" - - # Verify we got the expected number of blocks (100 time steps / 25 = 4) - expected_blocks = 4 - assert ( - iterations_after_collect == expected_blocks - ), f"Expected {expected_blocks} iterations, got {iterations_after_collect}" - - def test_multiple_queries_on_same_table(self, small_ds): - """Same table can be queried multiple times with fresh iteration each time.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 50}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # First query - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - first_query_iterations = tracker.iteration_count - assert first_query_iterations > 0, "First query should iterate" - - # Second query on same table - should iterate again - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - second_query_iterations = tracker.iteration_count - assert ( - second_query_iterations > first_query_iterations - ), "Second query should trigger additional iterations" - - def test_stream_consumed_error(self, small_ds): - """Once consumed, a single XarrayRecordBatchReader should not be reusable.""" - reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25}) - - # Consume the reader by converting to a PyArrow reader and reading - import pyarrow as pa - - pa_reader = pa.RecordBatchReader.from_stream(reader) - _ = pa_reader.read_all() - - # Reader is now consumed, calling __arrow_c_stream__ again should fail - with pytest.raises(RuntimeError, match="already consumed"): - reader.__arrow_c_stream__() - - -class TestDataIntegrity: - """Tests that verify data correctness alongside lazy evaluation. - - These tests use read_xarray_table for lazy streaming. - """ - - def test_query_results_are_correct(self, small_ds): - """Verify that lazy evaluation produces correct results.""" - table = read_xarray_table(small_ds, chunks={"time": 25}) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Get count - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows - expected_count = 100 * 10 * 10 - assert ( - count == expected_count - ), f"Expected {expected_count} rows, got {count}" - - def test_aggregation_results_are_correct(self, small_ds): - """Verify aggregation produces correct results.""" - table = read_xarray_table(small_ds, chunks={"time": 25}) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Get average temperature - result = ctx.sql( - "SELECT AVG(temperature) as avg_temp FROM test_table" - ).collect() - avg_temp = result[0].to_pandas()["avg_temp"].iloc[0] - - # With seed 42 and random data in [0, 1), average should be ~0.5 - assert ( - 0.4 < avg_temp < 0.6 - ), f"Expected average temperature ~0.5, got {avg_temp}" - - -class TestPyArrowInterop: - """Tests for PyArrow interoperability.""" - - def test_from_stream_does_not_iterate(self, small_ds): - """pa.RecordBatchReader.from_stream() should not iterate.""" - tracker = IterationTracker() - - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - # Create PyArrow reader from our stream - pa_reader = pa.RecordBatchReader.from_stream(reader) - - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations when creating PyArrow reader, " - f"but got {tracker.iteration_count}" - ) - - def test_pyarrow_iteration_triggers_callbacks(self, small_ds): - """Iterating via PyArrow should trigger our callbacks.""" - tracker = IterationTracker() - - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - pa_reader = pa.RecordBatchReader.from_stream(reader) - - # Now iterate - for batch in pa_reader: - pass - - assert ( - tracker.iteration_count == 4 - ), f"Expected 4 iterations, got {tracker.iteration_count}" - - def test_read_all_iterates_all(self, small_ds): - """read_all() should iterate through all blocks.""" - tracker = IterationTracker() - - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - pa_reader = pa.RecordBatchReader.from_stream(reader) - table = pa_reader.read_all() - - assert tracker.iteration_count == 4 - assert len(table) == 100 * 10 * 10 - - -class StreamingTracker: - """Tracks timing of batch iterations to verify streaming behavior. - - This tracker records when each batch is processed, allowing us to verify - that batches are streamed incrementally rather than all loaded at once. - """ - - def __init__(self): - self.batch_times = [] - self.batch_count = 0 - self._lock = threading.Lock() - - def __call__(self, block): - with self._lock: - self.batch_times.append(time.monotonic()) - self.batch_count += 1 - - def reset(self): - with self._lock: - self.batch_times = [] - self.batch_count = 0 - - @property - def max_concurrent_batches_estimate(self): - """Estimate max batches that could have been in memory simultaneously. - - If all batches are loaded at once, all batch_times will be very close. - If streaming works correctly, batch_times should be spread out. - """ - if len(self.batch_times) < 2: - return len(self.batch_times) - - # Sort times and look at gaps - sorted_times = sorted(self.batch_times) - # If times are spread out, streaming is working - # If all times are within a tiny window, all batches loaded at once - total_duration = sorted_times[-1] - sorted_times[0] - - # If the spread is very small compared to number of batches, - # batches were likely all loaded at once - return len(self.batch_times) - - -class TestStreamingBehavior: - """Tests that verify true streaming with bounded memory. - - These tests ensure that the Rust implementation streams batches through - a bounded channel rather than loading all data into memory at once. - """ - - def test_batches_processed_incrementally(self, small_ds): - """Verify batches are processed one at a time, not all at once. - - This test uses a callback that tracks when each batch is processed. - With true streaming, batches should be processed incrementally. - """ - tracker = StreamingTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Run query that scans all data - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - - # All 4 batches should have been processed - assert ( - tracker.batch_count == 4 - ), f"Expected 4 batches, got {tracker.batch_count}" - - def test_streaming_preserves_order(self, small_ds): - """Verify that streaming preserves the order of batches.""" - blocks_seen = [] - - def track_order(block): - # Record the time slice for ordering verification - blocks_seen.append(block.get("time", None)) - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=track_order, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - ctx.sql("SELECT * FROM test_table").collect() - - # Should have 4 blocks - assert len(blocks_seen) == 4 - - # Blocks should be in order (each slice should start after previous) - for i in range(1, len(blocks_seen)): - prev_end = blocks_seen[i - 1].stop - curr_start = blocks_seen[i].start - assert curr_start == prev_end, ( - f"Block {i} starts at {curr_start}, expected {prev_end}. " - f"Blocks are out of order!" - ) - - def test_large_dataset_streams_correctly(self): - """Test streaming with a larger dataset to verify memory behavior. - - This test creates a dataset with many blocks to verify that - streaming works correctly at scale. - """ - # Create a dataset with 20 blocks - np.random.seed(42) - time = pd.date_range("2020-01-01", periods=200, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) - - data = np.random.rand(200, 10, 10).astype(np.float32) - - large_ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # Use small chunks to create many blocks - table = read_xarray_table( - large_ds, - chunks={"time": 10}, # 200 / 10 = 20 blocks - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Run a query that needs all data - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - # Verify all blocks were processed - assert ( - tracker.batch_count == 20 - ), f"Expected 20 batches for large dataset, got {tracker.batch_count}" - - # Verify data integrity - expected_count = 200 * 10 * 10 - assert ( - count == expected_count - ), f"Expected {expected_count} rows, got {count}" - - -class TestBoundedMemoryBehavior: - """Tests that verify memory usage remains bounded during streaming. - - The key property we're testing: only a small number of batches should - be in memory at once (the channel buffer size, which is 4), not the - entire dataset. - - These tests verify that: - 1. Many batches can be processed without loading all into memory - 2. Production times are spread out (indicating back-pressure) - 3. Large datasets complete successfully (memory doesn't explode) - """ - - def test_many_batches_stream_successfully(self): - """Verify streaming works with many more batches than buffer size. - - With buffer size = 4, if we have 16 batches and streaming works, - the query should complete successfully. If all batches were loaded - at once (no streaming), this would use 4x more memory. - """ - # Create dataset with 16 batches (4x buffer size) - np.random.seed(42) - time_coord = pd.date_range("2020-01-01", periods=160, freq="h") - lat = np.linspace(-90, 90, 5) - lon = np.linspace(-180, 180, 5) - data = np.random.rand(160, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 16 batches (160 / 10 = 16) - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - # All 16 batches should have been processed - assert ( - tracker.batch_count == 16 - ), f"Expected 16 batches, got {tracker.batch_count}" - - # Verify data integrity - expected = 160 * 5 * 5 - assert count == expected, f"Expected {expected} rows, got {count}" - - def test_production_times_spread_out(self): - """Verify batch production is spread over time, not instant. - - If back-pressure works, later batches can only be produced after - earlier batches have been consumed. Production times should span - a non-zero duration. - """ - np.random.seed(123) - time_coord = pd.date_range("2020-01-01", periods=100, freq="h") - lat = np.linspace(-90, 90, 5) - lon = np.linspace(-180, 180, 5) - data = np.random.rand(100, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 10 batches, more than buffer size of 4 - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - - # All 10 batches should be produced - assert tracker.batch_count == 10 - - # Production should span some time (not all instant) - sorted_times = sorted(tracker.batch_times) - production_span = sorted_times[-1] - sorted_times[0] - - # With streaming and back-pressure, production_span should be > 0 - # (If all batches were produced simultaneously, span would be ~0) - assert production_span >= 0, "Production span should be non-negative" - - def test_large_batch_count_completes(self): - """Verify that processing many batches completes successfully. - - This is a stress test: 50 batches is well above the buffer size of 4. - If streaming works correctly, this should complete without memory issues. - """ - np.random.seed(456) - time_coord = pd.date_range("2020-01-01", periods=500, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) - data = np.random.rand(500, 10, 10).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 50 batches (500 / 10 = 50) - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - # All 50 batches processed - assert ( - tracker.batch_count == 50 - ), f"Expected 50 batches, got {tracker.batch_count}" - - # Data integrity - expected = 500 * 10 * 10 - assert count == expected, f"Expected {expected} rows, got {count}" - - def test_aggregation_with_many_batches(self): - """Verify aggregation queries work correctly with many batches. - - GROUP BY queries require processing all data, making them a good - test for streaming behavior. - - Note: ORDER BY is used to ensure deterministic results. Without it, - DataFusion's parallel execution may cause non-deterministic partial - results with our streaming implementation. - """ - np.random.seed(789) - time_coord = pd.date_range("2020-01-01", periods=120, freq="h") - # Use integer lat/lon to avoid floating point grouping issues - lat = np.array([0, 1, 2, 3, 4], dtype=np.float64) - lon = np.array([0, 1, 2, 3, 4], dtype=np.float64) - data = np.random.rand(120, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 12 batches - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # GROUP BY requires scanning all data - # ORDER BY ensures all partial aggregates are collected before returning - # TODO(#106): Fix the underlying partitioning issue. - result = ctx.sql( - "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat ORDER BY lat" - ).collect() - - # Should have result for each lat value - df = result[0].to_pandas() - assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}" - - # All batches processed - assert ( - tracker.batch_count == 12 - ), f"Expected 12 batches, got {tracker.batch_count}" - - -class TestErrorPropagation: - """Tests that verify errors are properly propagated through the stream. - - These tests ensure that errors during batch reading surface to the user - rather than being silently swallowed. - """ - - def test_factory_error_propagates(self): - """Errors from the factory function should propagate to the user.""" - - def failing_factory(): - raise ValueError("Factory intentionally failed") - - schema = pa.schema([("value", pa.int64())]) - table = LazyArrowStreamTable(failing_factory, schema) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # The error should surface when we try to collect - with pytest.raises(Exception) as exc_info: - ctx.sql("SELECT * FROM test_table").collect() - - # Verify the error message mentions the factory failure - error_message = str(exc_info.value).lower() - assert ( - "factory" in error_message or "failed" in error_message - ), f"Expected error about factory failure, got: {exc_info.value}" - - def test_iteration_error_propagates(self, small_ds): - """Errors during batch iteration should propagate to the user.""" - error_on_batch = 2 # Fail on the third batch - - def failing_callback(block): - # Track which batch we're on using a mutable default - if not hasattr(failing_callback, "count"): - failing_callback.count = 0 - failing_callback.count += 1 - - if failing_callback.count == error_on_batch: - raise RuntimeError("Intentional batch processing error") - - # Reset the counter - failing_callback.count = 0 - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=failing_callback, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # The error should surface when we try to collect - with pytest.raises(Exception): - ctx.sql("SELECT * FROM test_table").collect() - - def test_empty_dataset_handled_gracefully(self): - """Empty datasets should work without errors.""" - # Create an empty dataset with the right structure - empty_ds = xr.Dataset( - { - "temperature": ( - ["time", "lat", "lon"], - np.array([]).reshape(0, 0, 0), - ) - }, - coords={ - "time": pd.DatetimeIndex([]), - "lat": np.array([]), - "lon": np.array([]), - }, - ) - - # This should work without crashing - table = read_xarray_table(empty_ds, chunks={"time": 10}) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - assert count == 0, f"Expected 0 rows for empty dataset, got {count}" - - -class TestMultiplePartitions: - """Tests for scenarios with multiple queries and table reuse.""" - - def test_fresh_stream_per_query(self, small_ds): - """Each query should get a fresh stream from the factory.""" - call_count = {"value": 0} - original_callback = None - - def counting_callback(block): - call_count["value"] += 1 - if original_callback: - original_callback(block) - - table = read_xarray_table( - small_ds, - chunks={"time": 50}, # 2 blocks per query - _iteration_callback=counting_callback, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # First query - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - first_query_count = call_count["value"] - assert ( - first_query_count == 2 - ), f"First query: expected 2, got {first_query_count}" - - # Second query should trigger fresh iteration - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - second_query_count = call_count["value"] - assert ( - second_query_count == 4 - ), f"After second query: expected 4 total, got {second_query_count}" - - # Third query - ctx.sql("SELECT MAX(temperature) FROM test_table").collect() - third_query_count = call_count["value"] - assert ( - third_query_count == 6 - ), f"After third query: expected 6 total, got {third_query_count}" - - def test_parallel_queries_independent(self, small_ds): - """Multiple contexts with the same table should work independently.""" - tracker1 = IterationTracker() - tracker2 = IterationTracker() - - table1 = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker1, - ) - - table2 = read_xarray_table( - small_ds, - chunks={"time": 50}, - _iteration_callback=tracker2, - ) - - ctx1 = SessionContext() - ctx2 = SessionContext() - - ctx1.register_table("test_table", table1) - ctx2.register_table("test_table", table2) - - # Execute queries - ctx1.sql("SELECT COUNT(*) FROM test_table").collect() - ctx2.sql("SELECT COUNT(*) FROM test_table").collect() - - # Each should have its own iteration count - assert ( - tracker1.iteration_count == 4 - ), f"Table1: expected 4 blocks, got {tracker1.iteration_count}" - assert ( - tracker2.iteration_count == 2 - ), f"Table2: expected 2 blocks, got {tracker2.iteration_count}" diff --git a/xarray_sql/roundtrip.py b/xarray_sql/roundtrip.py new file mode 100644 index 00000000..e342828f --- /dev/null +++ b/xarray_sql/roundtrip.py @@ -0,0 +1,495 @@ +"""Engine-agnostic round-trip: Arrow query results → labeled ``xr.Dataset``. + +The second seam of xarray-sql. Any engine's result — a DuckDB relation, +a ``pyarrow.Table``, a ``pyarrow.RecordBatchReader``, or any object +implementing the Arrow PyCapsule stream protocol — plus the registered +Dataset as a *template* is enough to rebuild a labeled, metadata-carrying +Dataset. Nothing here is engine-specific: results arrive as Arrow record +batches regardless of which engine executed the SQL. + +Reconstruction is eager by default (the result is materialized once +into a dense in-memory Dataset). Passing ``chunks=`` selects the +lazy/chunked path instead: data variables are reconstructed on access, +window by window, by re-executing the engine's query narrowed to each +chunk's coordinate range. That requires the result to be +*re-executable* — a Polars LazyFrame (or eager DataFrame) or a +DataFusion DataFrame — not a one-shot Arrow stream; see +[xarray_sql.lazyscan][]. DuckDB relations are re-executable but +refuse the chunked path (a thread-safety limitation noted on +[DuckDBHandle][xarray_sql.lazyscan.DuckDBHandle]); pair them with +``spill=True`` instead. +""" + +from __future__ import annotations + +import os +import tempfile +import weakref +from collections.abc import Mapping +from typing import Any, Literal + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq +import xarray as xr + +from .ds import ( + Sparsity, + XarrayDataFrame, + _build_lazy_scan, + _dataset_from_batches, + _ds_var_dims, + _finish_dataset, +) +from .lazyscan import LazyResultHandle, PolarsHandle, resolve_lazy_handle + + +def _guarded(batches: Any, max_bytes: int | None) -> list[pa.RecordBatch]: + """Collect a batch iterable, erroring cleanly past ``max_bytes``. + + A result that would blow past the budget raises with the running + size instead of exhausting memory, before the (larger) dense + reconstruction is even attempted. + """ + if max_bytes is None: + return list(batches) + out: list[pa.RecordBatch] = [] + total = 0 + for batch in batches: + total += batch.nbytes + if total > max_bytes: + raise ValueError( + f"result exceeded max_result_bytes={max_bytes:,} while " + f"materializing (>= {total:,} bytes after " + f"{sum(b.num_rows for b in out) + batch.num_rows:,} rows). " + "Aggregate further, or reconstruct lazily with chunks=." + ) + out.append(batch) + return out + + +def _open_stream(result: Any) -> tuple[pa.Schema, Any] | None: + """The result's Arrow batches as ``(schema, iterable)``, or ``None``. + + Probes, in order: ``pyarrow.Table`` / ``pyarrow.RecordBatch``, + ``pyarrow.RecordBatchReader``, ``__arrow_c_stream__`` (the Arrow + PyCapsule protocol — DuckDB relations qualify on duckdb >= 1.1), and + a ``fetch_record_batch()`` method (DuckDB relations on older + versions). + """ + if isinstance(result, pa.RecordBatch): + return result.schema, [result] + if isinstance(result, pa.Table): + return result.schema, result.to_batches() + if isinstance(result, pa.RecordBatchReader): + return result.schema, result + if hasattr(result, "__arrow_c_stream__"): + reader = pa.RecordBatchReader.from_stream(result) + return reader.schema, reader + if hasattr(result, "fetch_record_batch"): + reader = result.fetch_record_batch() + return reader.schema, reader + return None + + +def _result_to_batches( + result: Any, max_bytes: int | None = None +) -> tuple[pa.Schema, list[pa.RecordBatch]]: + """Normalize an engine result into ``(schema, record batches)``. + + Accepts everything ``_open_stream`` recognizes, then objects + with a ``to_arrow_table()`` method (DataFusion DataFrames and the + [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper), then re-executable + results without a stream protocol (a Polars LazyFrame), executed + once through their lazy handle. + """ + opened = _open_stream(result) + if opened is not None: + schema, batches = opened + if isinstance(result, (pa.RecordBatch, pa.Table)): + # Already in memory: nothing left for the budget to bound + # (the dense-size check still applies downstream). + return schema, list(batches) + return schema, _guarded(batches, max_bytes) + handle = resolve_lazy_handle(result) + if hasattr(result, "to_arrow_table") and ( + max_bytes is None or handle is None + ): + # This branch materializes the whole result in one call before + # the budget can observe a single batch, so with a budget set a + # re-executable result streams through its handle below instead; + # the post-materialization nbytes check is the fallback guard + # for one-shot results whose only surface is to_arrow_table(). + table = result.to_arrow_table() + if max_bytes is not None and table.nbytes > max_bytes: + raise ValueError( + f"result materialized to {table.nbytes:,} bytes, over " + f"max_result_bytes={max_bytes:,}. Aggregate further, or " + "reconstruct lazily with chunks=." + ) + return table.schema, table.to_batches() + if handle is not None: + schema = handle.schema() + names = list(schema.names) + if max_bytes is None: + return schema, handle.fetch({}, names) + # fetch() may materialize the whole result inside the engine + # before any batch surfaces (Polars collect()), which would + # defeat the budget; enforce it on a true batch stream, or + # refuse up front instead of erroring after the memory is spent. + stream = getattr(handle, "stream", None) + if stream is None: + raise ValueError( + "max_result_bytes cannot be enforced for " + f"{type(result).__qualname__}: the result materializes " + "fully before batches surface. Drop max_result_bytes=, " + "or reconstruct lazily with chunks=." + ) + return schema, _guarded(stream(names), max_bytes) + raise TypeError( + f"Cannot read an Arrow stream from {type(result).__qualname__}; " + "expected a pyarrow Table/RecordBatch/RecordBatchReader, an object " + "implementing __arrow_c_stream__, or an engine result exposing " + "fetch_record_batch()/to_arrow_table()." + ) + + +def to_dataset( + result: Any, + dims: list[str] | None = None, + template: xr.Dataset | None = None, + sparsity: Sparsity = "result", + fill_value: Any = np.nan, + chunks: Mapping[str, int] | str | None = None, + coords: Literal["discover", "template"] = "discover", + max_result_bytes: int | None = None, + spill: bool | str | os.PathLike = False, +) -> xr.Dataset: + """Convert an engine's Arrow result into a labeled ``xr.Dataset``. + + The engine-agnostic counterpart of + [XarrayDataFrame.to_dataset][xarray_sql.ds.XarrayDataFrame.to_dataset]: SQL in, array out, for engines + xarray-sql does not wrap in a session of its own. + + Example (DuckDB):: + + con = duckdb.connect() + xql.register(con, "era5", ds) + rel = con.sql( + "SELECT time, lat, lon, AVG(t2m) AS t2m FROM era5 " + "GROUP BY time, lat, lon" + ) + out = xql.to_dataset(rel, template=ds) + + Args: + result: The engine's query result: a ``pyarrow.Table``, + ``RecordBatch`` or ``RecordBatchReader``, any object + implementing ``__arrow_c_stream__`` (DuckDB relations), or an + object with ``fetch_record_batch()`` / ``to_arrow_table()``. + The result is consumed once. + dims: Result columns to use as Dataset dimensions. When ``None``, + defaults to the ``template``'s dimensions that survive into + the result columns (so aggregations that drop dims round-trip + on the remaining ones). Either ``dims`` or ``template`` must + be given. + template: The source Dataset registered with the engine. Recovers + metadata the tabular pivot strips (attrs, encoding, non-dim + coordinates, dim-coord dtype) and provides the ``dims`` + default. + sparsity: ``"result"`` (default) keeps only dim values present in + the result. ``"template"`` reindexes to the template's full + coord ranges, filling absent cells with ``fill_value``. + fill_value: Fill for ``sparsity="template"``. Defaults to NaN. + chunks: ``None`` (default) materializes eagerly. A mapping + (e.g. ``{"time": 100}``), ``"auto"``, or ``"inherit"`` + selects the lazy/chunked path: data variables are + reconstructed window by window on access, each window + re-executing the engine's query narrowed to its coordinate + range (over a table registered through xarray-sql, that + filter flows back into chunk pruning at the source). + Requires a re-executable ``result`` — a Polars + LazyFrame/DataFrame or a DataFusion DataFrame; DuckDB + relations refuse the chunked path (add ``spill=True``). + coords: How the lazy path learns each dimension's coordinate + values. ``"discover"`` (default) runs one ``DISTINCT`` query + per dim — correct for any query. ``"template"`` trusts the + template's coord arrays instead, skipping discovery; only + valid when the result spans the template's full extent (an + unfiltered scan), and requires ``template=``. + max_result_bytes: Optional budget for the eager path. Raises a + clean ``ValueError`` (with the running size) as soon as the + materializing result exceeds it — both while collecting the + Arrow stream and before allocating the dense arrays — + instead of exhausting memory. ``None`` (default) means + unlimited. Results whose only surface is + ``to_arrow_table()`` necessarily materialize in full before + the budget can be checked (the check then runs on the + materialized size); re-executable results stream instead, + so the budget fires before full materialization. + spill: Chunked reconstruction from a one-pass on-disk spill + instead of per-window re-execution: the result is streamed + *once* (bounded memory) into a temporary Parquet file, and + windows re-execute against that file. This serves the two + results the re-execution path cannot — DuckDB relations and + one-shot Arrow streams — and trades per-window narrowness + for a single full pass plus temporary disk. ``True`` spills + to the system temp dir; a path spills into that directory. + The file is removed when the returned Dataset is garbage + collected. Requires Polars; only valid with ``chunks=``. + + Returns: + An ``xr.Dataset`` with ``dims`` as dimensions and the remaining + result columns as data variables — dense and in-memory by + default, lazily chunked when ``chunks`` is given. + + Raises: + ValueError: When neither ``dims`` nor ``template`` resolves the + dimension columns, a requested dim is missing from the result, + or ``sparsity="template"`` is used without a template. + TypeError: When ``result`` exposes no readable Arrow stream, or + ``chunks`` is requested for a one-shot stream that cannot be + re-executed. + """ + if sparsity not in ("result", "template"): + raise ValueError( + f"sparsity must be 'result' or 'template', got {sparsity!r}" + ) + if sparsity == "template" and template is None: + raise ValueError("sparsity='template' requires template= to be given") + if coords not in ("discover", "template"): + raise ValueError( + f"coords must be 'discover' or 'template', got {coords!r}" + ) + if coords == "template" and template is None: + raise ValueError("coords='template' requires template= to be given") + if spill and chunks is None: + raise ValueError( + "spill= only applies to chunked reconstruction; pass chunks=." + ) + + if chunks is not None: + if spill: + return _to_dataset_spilled( + result, + dims, + template, + sparsity, + fill_value, + chunks, + coords, + spill, + ) + return _to_dataset_lazy( + result, dims, template, sparsity, fill_value, chunks, coords + ) + + schema, batches = _result_to_batches(result, max_result_bytes) + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + + dims = _resolve_dims(dims, template, field_names) + + if max_result_bytes is not None: + _check_dense_size( + batches, dims, field_names, field_types, max_result_bytes + ) + ds = _dataset_from_batches(batches, dims, field_names, field_types) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, None, field_types + ) + + +def _check_dense_size( + batches: list[pa.RecordBatch], + dims: list[str], + field_names: list[str], + field_types: dict[str, Any], + max_bytes: int, +) -> None: + """Error before allocating dense arrays larger than the budget. + + The dense grid is the coordinate product, which for sparse results + can dwarf the Arrow input; check it against the same budget before + a single output array is allocated. + """ + sizes = [] + for d in dims: + # Vectorized distinct count: a per-row Python set (to_pylist) + # costs orders of magnitude more on wide results. + arrays = [b.column(b.schema.names.index(d)) for b in batches] + sizes.append(len(pc.unique(pa.chunked_array(arrays))) if arrays else 0) + cells = int(np.prod(sizes)) if sizes else 0 + total = sum( + cells * np.dtype(field_types[n].to_pandas_dtype()).itemsize + for n in field_names + if n not in dims + ) + if total > max_bytes: + raise ValueError( + f"dense reconstruction needs {total:,} bytes " + f"({cells:,} grid cells), over max_result_bytes=" + f"{max_bytes:,}. Aggregate further, or reconstruct lazily " + "with chunks=." + ) + + +def _resolve_dims( + dims: list[str] | None, + template: xr.Dataset | None, + field_names: list[str], +) -> list[str]: + """Dimension columns, inferred from the template when not given.""" + if dims is None: + if template is None: + raise ValueError( + "dims cannot be inferred without a template; pass " + "dims=[...] or template=." + ) + dims = [d for d in _ds_var_dims(template) if d in field_names] + if not dims: + raise ValueError( + "dims cannot be inferred: no template dimension survives " + "in the result columns. Pass dims=[...] explicitly." + ) + missing = [d for d in dims if d not in field_names] + if missing: + raise ValueError( + f"dims {missing} are not columns of the result {field_names}." + ) + return dims + + +def _to_dataset_lazy( + result: Any, + dims: list[str] | None, + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str, + coords: Literal["discover", "template"], + _handle: LazyResultHandle | None = None, +) -> xr.Dataset: + """The chunked reconstruction behind ``to_dataset(chunks=...)``.""" + handle = _handle if _handle is not None else resolve_lazy_handle(result) + if handle is None: + raise TypeError( + "chunks= requires a re-executable engine result (a Polars " + "LazyFrame/DataFrame or a DataFusion DataFrame); got " + f"{type(result).__qualname__}, which is a one-shot stream. " + "Pass the engine's lazy handle instead of a materialized " + "result, add spill=True to reconstruct from a one-pass " + "on-disk spill, or use chunks=None." + ) + schema = handle.schema() + field_names = [f.name for f in schema] + field_types = {f.name: f.type for f in schema} + dims = _resolve_dims(dims, template, field_names) + + coord_arrays = None + if coords == "template": + assert template is not None + missing = [d for d in dims if d not in template.coords] + if missing: + raise ValueError( + f"coords='template' requires the template to carry coords " + f"for every dim; missing {missing}." + ) + coord_arrays = {d: np.asarray(template.coords[d].values) for d in dims} + + resolved = XarrayDataFrame._resolve_chunks(chunks, template, dims) + if resolved is None: + # "inherit" with no chunked source dimension to inherit from: + # eager is the right execution, exactly as on the wrapper path. + batches = handle.fetch({}, field_names) + ds = _dataset_from_batches(batches, dims, field_names, field_types) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, None, field_types + ) + if not getattr(handle, "supports_chunked", True): + raise NotImplementedError( + "Chunked reconstruction is not supported for " + f"{type(result).__qualname__}: re-executing a DuckDB " + "relation from worker threads intermittently deadlocks in " + "duckdb-python when the query scans a Python-backed table " + "(see xarray_sql.lazyscan.DuckDBHandle.supports_chunked). " + "Add spill=True to reconstruct from a one-pass on-disk " + "spill, use chunks=None (eager), or run the query through " + "Polars (pl.scan_pyarrow_dataset(xql.arrow_dataset(ds))) " + "or a DataFusion context." + ) + ds = _build_lazy_scan( + handle, dims, field_names, field_types, coord_arrays=coord_arrays + ) + return _finish_dataset( + ds, dims, template, sparsity, fill_value, resolved, field_types + ) + + +def _to_dataset_spilled( + result: Any, + dims: list[str] | None, + template: xr.Dataset | None, + sparsity: Sparsity, + fill_value: Any, + chunks: Mapping[str, int] | str, + coords: Literal["discover", "template"], + spill: bool | str | os.PathLike, +) -> xr.Dataset: + """Chunked reconstruction from a one-pass temporary Parquet spill. + + The result is streamed exactly once with bounded memory — through + the engine handle where one exists (DuckDB spills on its dedicated + engine thread; Polars uses its streaming sink), or straight from + the Arrow stream for one-shot results — and the ordinary lazy + reconstruction then runs against a Polars scan of the file, whose + per-window predicates enjoy Parquet row-group pruning. The file is + removed when the reconstruction handle is garbage collected. + """ + import polars as pl + + directory = os.fspath(spill) if not isinstance(spill, bool) else None + fd, path = tempfile.mkstemp(suffix=".parquet", dir=directory) + os.close(fd) + try: + handle = resolve_lazy_handle(result) + if handle is not None: + handle.spill_parquet(path) + else: + _stream_to_parquet(result, path) + except BaseException: + os.unlink(path) + raise + spilled = PolarsHandle(pl.scan_parquet(path)) + weakref.finalize(spilled, _unlink_quietly, path) + return _to_dataset_lazy( + result, + dims, + template, + sparsity, + fill_value, + chunks, + coords, + _handle=spilled, + ) + + +def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except OSError: + pass + + +def _stream_to_parquet(result: Any, path: str) -> None: + """Write a one-shot Arrow result to Parquet, batch by batch.""" + opened = _open_stream(result) + if opened is None: + raise TypeError( + f"cannot spill {type(result).__qualname__}: no readable " + "Arrow stream." + ) + schema, batches = opened + with pq.ParquetWriter(path, schema) as writer: + for batch in batches: + writer.write_batch(batch) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 4bdb7058..a5df4fb9 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -1,18 +1,190 @@ import xarray as xr from datafusion import SessionContext +from datafusion.catalog import Schema +from types import ModuleType -from .df import Chunks +from . import cftime as cft +from .df import Chunks, group_vars_by_dims +from .ds import XarrayDataFrame from .reader import read_xarray_table +_proj: ModuleType | None +try: # pyproj is an optional dependency (`pip install xarray-sql[geo]`). + from . import proj + + _proj = proj +except ImportError: # pragma: no cover - depends on the environment + _proj = None + class XarrayContext(SessionContext): - """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" - - def from_dataset( - self, - table_name: str, - input_table: xr.Dataset, - chunks: Chunks = None, - ): - table = read_xarray_table(input_table, chunks) - self.register_table(table_name, table) + """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Track registered xarray Datasets so XarrayDataFrame can recover + # defaults (dimension_columns) and metadata (var/dataset attrs, + # non-dim coords, dim-coord dtype) that the forward pivot drops. + # Keys are the fully-qualified table names users will reference + # in SQL (e.g. ``"air"`` for a uniform-dim Dataset, or + # ``"era5.surface"`` for one entry from a multi-dim-group split). + self._registered_datasets: dict[str, xr.Dataset] = {} + # With pyproj installed, every context speaks CRS out of the box: + # reproject(x, y, src_crs, dst_crs), à la PostGIS ST_Transform. + if _proj is not None: + _proj.register(self) + + def from_dataset( + self, + name: str, + input_table: xr.Dataset, + *, + table_names: dict[tuple[str, ...], str] | None = None, + chunks: Chunks = None, + ): + """Register an xarray Dataset as one or more queryable SQL tables. + + When all data variables share the same dimensions, the dataset is + registered as a single table named ``name``. When variables have + differing dimensions (e.g. some on a 3D grid and others on a 4D + grid), the dataset is split into one table per dimension group. + The tables are registered under a SQL schema (namespace) named + ``name`` and named ``__...`` by default:: + + ctx.from_dataset('era5', ds, chunks={'time': 24}) + # registers tables: 'era5.time_lat_lon' and + # 'era5.time_lat_lon_level' + ctx.sql('SELECT AVG(temperature_2m) FROM era5.time_lat_lon') + + Use ``table_names`` to override the name for specific dimension + tuples:: + + ctx.from_dataset( + 'era5', ds, + table_names={('time', 'lat', 'lon'): 'surface'}, + ) + ctx.sql('SELECT * FROM era5.surface') + + For datasets with non-Gregorian cftime coordinates (e.g. 360_day, + julian), a ``cftime()`` scalar UDF is automatically registered so + you can write ergonomic SQL filters:: + + ctx.from_dataset("ds360", ds, chunks={"time": 6}) + ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") + + .. note:: + + Only one ``cftime()`` UDF is registered per context, using the + units and calendar of the *first* non-Gregorian coordinate + encountered. If you register multiple datasets with *different* + non-Gregorian calendars (e.g. one 360_day and one julian), the + UDF from the first registration will be used for all subsequent + ``cftime()`` calls and may produce incorrect offsets for the + other dataset. In that case, create a separate ``XarrayContext`` + for each calendar. + + Args: + name: The SQL identifier under which the dataset is registered. + For datasets with uniform dimensions, this is the table + name. For datasets with mixed dimensions, this is the name + of a SQL schema (namespace) containing one table per + dimension group. + input_table: An xarray Dataset. + table_names: Optional mapping from dimension tuples to custom + table names within the schema, used when the dataset has + variables with differing dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + + Returns: + self, to allow chaining. + """ + groups = group_vars_by_dims(input_table) + + # Materialise dim coordinates once and share across every sub-table. + # For Zarr-backed parents (e.g. ARCO-ERA5 on GCS) this saves one + # network round-trip per dim per dim-group. + coord_arrays = { + str(dim): input_table.coords[dim].values for dim in input_table.dims + } + + if len(groups) <= 1: + self._registered_datasets[name] = input_table + return self._from_dataset( + name, input_table, chunks, coord_arrays=coord_arrays + ) + + table_names = table_names or {} + schema = Schema.memory_schema(self) + self.catalog().register_schema(name, schema) + + for dims, var_names in groups.items(): + # Scalar variables group under empty dims, where "_".join(()) is + # the empty string; fall back to a valid default table name. + sub_name = table_names.get(dims, "_".join(dims) or "scalar") + sub_ds = input_table[var_names] + self._from_dataset( + sub_name, + sub_ds, + chunks, + schema=schema, + coord_arrays=coord_arrays, + ) + # Track the fully-qualified name so XarrayDataFrame metadata + # recovery can find this Dataset on round-trip. + self._registered_datasets[f"{name}.{sub_name}"] = sub_ds + + return self + + def _from_dataset( + self, + table_name: str, + input_table: xr.Dataset, + chunks: Chunks = None, + schema: Schema | None = None, + coord_arrays: dict | None = None, + ): + """Register a Dataset as a single SQL table. + + Registers a top-level table by default, or a table inside ``schema`` + (a SQL namespace) when one is given. + """ + register = ( + self.register_table if schema is None else schema.register_table + ) + register( + table_name, + read_xarray_table(input_table, chunks, coord_arrays=coord_arrays), + ) + self._maybe_register_cftime_udf(input_table) + return self + + def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None: + """Auto-register a cftime() UDF for non-Gregorian cftime coordinates.""" + for coord_name in ds.dims: + if cft.is_cftime_index(ds, coord_name): + units, cal = cft.encoding(ds, coord_name) + if not cft.is_gregorian_like(cal): + self.register_udf(cft.make_cftime_udf(units, cal)) + break # One UDF per context is enough. + + def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame: + """Run a SQL query, returning an [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapper. + + Identical to ``datafusion.SessionContext.sql`` except the returned + object wraps the DataFusion DataFrame. The wrapper exposes + ``.to_pandas()`` (unchanged), forwards every other DataFusion + method via ``__getattr__``, and adds + ``.to_dataset(dimension_columns=[...])`` for round-tripping the + result back to an ``xr.Dataset``. + + Args: + query: A SQL query string. + *args: Forwarded to ``SessionContext.sql``. + **kwargs: Forwarded to ``SessionContext.sql``. + + Returns: + An [XarrayDataFrame][xarray_sql.ds.XarrayDataFrame] wrapping the DataFusion DataFrame. + """ + inner = super().sql(query, *args, **kwargs) + return XarrayDataFrame(inner, templates=self._registered_datasets) diff --git a/xarray_sql/sql_test.py b/xarray_sql/sql_test.py deleted file mode 100644 index 9374cb64..00000000 --- a/xarray_sql/sql_test.py +++ /dev/null @@ -1,196 +0,0 @@ -"""SQL functionality tests for xarray-sql using pytest.""" - -import numpy as np -import pandas as pd -import pytest -import xarray as xr - -from . import XarrayContext -from .df_test import create_large_dataset, rand_wx - - -@pytest.fixture -def air_dataset_small(): - ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) - return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)) - - -@pytest.fixture -def air_dataset_large(): - return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) - - -@pytest.fixture -def weather_dataset(): - ds = rand_wx("2023-01-01T00", "2023-01-01T12") - return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk( - {"time": 3} - ) - - -@pytest.fixture -def synthetic_dataset(): - return create_large_dataset( - time_steps=50, lat_points=20, lon_points=20 - ).chunk({"time": 25}) - - -@pytest.fixture -def station_dataset(): - return xr.Dataset( - { - "station_id": (["station"], [1, 2, 3, 4, 5]), - "elevation": (["station"], [100, 250, 500, 750, 1000]), - "name": ( - ["station"], - ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"], - ), - } - ).chunk({"station": 5}) - - -@pytest.fixture -def air_and_stations(): - air = ( - xr.tutorial.open_dataset("air_temperature") - .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8)) - .chunk({"time": 6}) - ) - stations = xr.Dataset( - { - "station_id": (["station"], [101, 102, 103]), - "lat": ( - ["station"], - [air.lat.values[0], air.lat.values[2], air.lat.values[4]], - ), - "lon": ( - ["station"], - [air.lon.values[1], air.lon.values[3], air.lon.values[5]], - ), - "elevation": (["station"], [100, 250, 500]), - } - ).chunk({"station": 3}) - return air, stations - - -def test_sanity(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - result = ctx.sql( - 'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100' - ).to_pandas() - assert len(result) > 0 - assert len(result) <= 1320 - assert all(col in result.columns for col in ["lat", "lon", "time", "air"]) - - -def test_aggregation_small(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - query = """ - SELECT lat, lon, SUM(air) AS air_total - FROM air - GROUP BY lat, lon - """ - result = ctx.sql(query).to_pandas() - expected_rows = ( - air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"] - ) - assert len(result) == expected_rows - - -def test_aggregation_large(air_dataset_large): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_large) - query = """ - SELECT lat, lon, AVG(air) AS air_avg - FROM air - GROUP BY lat, lon - """ - result = ctx.sql(query).to_pandas() - expected_rows = ( - air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"] - ) - assert len(result) == expected_rows - - -def test_basic_select_all(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas() - assert len(result) <= 10 - for col in ["lat", "lon", "time", "air"]: - assert col in result.columns - - -def test_weather_queries(weather_dataset): - ctx = XarrayContext() - ctx.from_dataset("weather", weather_dataset) - # Selecting specific columns - result = ctx.sql( - "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20" - ).to_pandas() - assert "temperature" in result.columns - assert "precipitation" in result.columns - # Filtering - result = ctx.sql( - "SELECT * FROM weather WHERE temperature > 10 LIMIT 50" - ).to_pandas() - assert len(result) > 0 - assert (result["temperature"] > 10).all() - - -def test_synthetic_aggregations(synthetic_dataset): - ctx = XarrayContext() - ctx.from_dataset("synthetic", synthetic_dataset) - # COUNT aggregation - result = ctx.sql("SELECT COUNT(*) AS total_count FROM synthetic").to_pandas() - assert result["total_count"].iloc[0] > 0 - # MIN, MAX, AVG - query = """ - SELECT MIN(temperature) AS min_temp, - MAX(temperature) AS max_temp, - AVG(temperature) AS avg_temp - FROM synthetic - """ - result = ctx.sql(query).to_pandas() - assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0] - assert ( - result["min_temp"].iloc[0] - <= result["avg_temp"].iloc[0] - <= result["max_temp"].iloc[0] - ) - - -def test_invalid_table_name(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT * FROM nonexistent_table") - - -def test_invalid_column_name(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT nonexistent_column FROM air") - - -def test_sql_syntax_error(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT * FORM air") # Typo: FORM instead of FROM - with pytest.raises(Exception): - ctx.sql("SELECT * FROM air WHERE") # Incomplete WHERE - - -def test_cross_join(air_and_stations): - air, stations = air_and_stations - ctx = XarrayContext() - ctx.from_dataset("air_data", air) - ctx.from_dataset("stations", stations) - result = ctx.sql( - "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations" - ).to_pandas() - assert result["total"].iloc[0] > 0 diff --git a/zensical.toml b/zensical.toml new file mode 100644 index 00000000..b80208b5 --- /dev/null +++ b/zensical.toml @@ -0,0 +1,138 @@ +[project] +site_name = "xarray-sql" +site_description = "Query Xarray with SQL" +site_author = "Alexander Merose" +site_url = "https://xqlsystems.github.io/xarray-sql" +repo_url = "https://github.com/xqlsystems/xarray-sql" +repo_name = "xqlsystems/xarray-sql" +edit_uri = "edit/main/docs/" +nav = [ + {"Home" = "index.md"}, + {"Getting started" = "examples.md"}, + {"Concepts" = [ + {"Engines" = "engines.md"}, + ]}, + {"Guides" = [ + {"Geospatial in SQL" = "geospatial.md"}, + {"Performance" = "performance.md"}, + ]}, + {"Reference" = [ + {"API" = "reference/xarray_sql.md"}, + {"Known issues" = "limitations.md"}, + {"Contributing" = "contributing.md"}, + ]} +] + +# Theme configuration +[project.theme] +variant = "modern" +logo = "assets/logo.svg" +features = [ + "announce.dismiss", + "content.action.edit", + "content.action.view", + "content.code.annotate", + "content.code.copy", + "content.tooltips", + "navigation.footer", + "navigation.indexes", + "navigation.instant", + "navigation.instant.prefetch", + "navigation.instant.progress", + "navigation.sections", + "navigation.top", + "navigation.tracking", + "search.highlight", + "search.share", + "search.suggest", + "toc.follow" +] + +# Color palette - light mode +[[project.theme.palette]] +media = "(prefers-color-scheme: light)" +scheme = "default" +primary = "custom" +accent = "blue" + +[project.theme.palette.toggle] +icon = "material/brightness-7" +name = "Switch to dark mode" + +# Color palette - dark mode +[[project.theme.palette]] +media = "(prefers-color-scheme: dark)" +scheme = "slate" +primary = "custom" +accent = "black" + +[project.theme.palette.toggle] +icon = "material/brightness-4" +name = "Switch to light mode" + +# Markdown extensions +[project.markdown_extensions.toc] +permalink = true + +[project.markdown_extensions.pymdownx.highlight] +anchor_linenums = true +line_spans = "__span" +pygments_lang_class = true + +[project.markdown_extensions.pymdownx.superfences] +[[project.markdown_extensions.pymdownx.superfences.custom_fences]] +name = "mermaid" +class = "mermaid" +format = "pymdownx.superfences.fence_code_format" + +[project.markdown_extensions.pymdownx.tasklist] +custom_checkbox = true + +[project.markdown_extensions.pymdownx.inlinehilite] + +[project.markdown_extensions.admonition] + +[project.markdown_extensions.footnotes] + +[project.markdown_extensions.pymdownx.snippets] +url_download = true +base_path = ["."] + +[project.markdown_extensions.pymdownx.tabbed] +alternate_style = true + +[project.markdown_extensions.pymdownx.emoji] +emoji_index = "zensical.extensions.emoji.twemoji" +emoji_generator = "zensical.extensions.emoji.to_svg" + +[project.markdown_extensions.attr_list] + +[project.markdown_extensions.md_in_html] + +[project.markdown_extensions.abbr] + +[project.markdown_extensions.def_list] + +# Plugins +[project.plugins] +search = {} +typeset = {} + +[project.plugins.mkdocstrings.handlers.python] +paths = ["."] + +[project.plugins.mkdocstrings.handlers.python.options] +docstring_style = "google" +show_if_no_docstring = true +filters = ["!^_"] + +# Extra configuration +[project.extra] + +[[project.extra.social]] +icon = "fontawesome/brands/github" +link = "https://github.com/xqlsystems/xarray-sql" + +[[project.extra.social]] +icon = "fontawesome/brands/python" +link = "https://pypi.org/project/xarray-sql"