diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 55f37fe21..7387de7a3 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -40,7 +40,7 @@ jobs: - run: | sudo apt remove -y postgres* - sudo apt -y install curl ca-certificates build-essential pkg-config libssl-dev + sudo apt -y install curl ca-certificates build-essential cmake pkg-config libssl-dev sudo install -d /usr/share/postgresql-common/pgdg sudo curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail https://www.postgresql.org/media/keys/ACCC4CF8.asc . /etc/os-release @@ -53,9 +53,11 @@ jobs: - run: cargo install --locked cargo-pgrx --version 0.19.2 - run: cargo pgrx init --pg15 /usr/lib/postgresql/15/bin/pg_config - - name: Build docker images + - name: Start test services and wait for fixtures run: | docker compose -f wrappers/.ci/docker-compose-native.yaml up -d + docker compose -f wrappers/.ci/docker-compose-native.yaml wait s3-init + test "$(docker inspect -f '{{.State.ExitCode}}' s3-init)" = "0" - name: Generate code coverage id: coverage diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index afa07af92..9fc3284a7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -163,7 +163,7 @@ jobs: # Add postgres package repo and install requested postgres version sudo apt update sudo apt remove -y postgres* - sudo apt -y install curl ca-certificates pkg-config libssl-dev + sudo apt -y install curl ca-certificates cmake pkg-config libssl-dev sudo install -d /usr/share/postgresql-common/pgdg sudo curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail https://www.postgresql.org/media/keys/ACCC4CF8.asc . /etc/os-release diff --git a/.github/workflows/test_wrappers.yml b/.github/workflows/test_wrappers.yml index 8cb000278..47c62cf97 100644 --- a/.github/workflows/test_wrappers.yml +++ b/.github/workflows/test_wrappers.yml @@ -54,9 +54,11 @@ jobs: - name: Checkout code uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - - name: Build docker images + - name: Start test services and wait for fixtures run: | docker compose -f wrappers/.ci/docker-compose-native.yaml up -d + docker compose -f wrappers/.ci/docker-compose-native.yaml wait s3-init + test "$(docker inspect -f '{{.State.ExitCode}}' s3-init)" = "0" - uses: actions-rust-lang/setup-rust-toolchain@150fca883cd4034361b621bd4e6a9d34e5143606 # v1.15.4 with: @@ -65,13 +67,13 @@ jobs: - run: | sudo apt remove -y postgres* - sudo apt -y install curl ca-certificates build-essential pkg-config libssl-dev + sudo apt -y install curl ca-certificates build-essential cmake pkg-config libssl-dev sudo install -d /usr/share/postgresql-common/pgdg sudo curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail https://www.postgresql.org/media/keys/ACCC4CF8.asc . /etc/os-release sudo sh -c "echo 'deb [signed-by=/usr/share/postgresql-common/pgdg/apt.postgresql.org.asc] https://apt.postgresql.org/pub/repos/apt $VERSION_CODENAME-pgdg main' > /etc/apt/sources.list.d/pgdg.list" sudo apt update -y -qq --fix-missing - sudo apt -y install postgresql-client-15 postgresql-15 postgresql-server-dev-15 + sudo apt -y install postgresql-client-15 postgresql-15 postgresql-15-postgis-3 postgresql-15-postgis-3-scripts postgresql-server-dev-15 sudo apt -y autoremove && sudo apt -y clean sudo chmod a+rwx `/usr/lib/postgresql/15/bin/pg_config --pkglibdir` `/usr/lib/postgresql/15/bin/pg_config --sharedir`/extension /var/run/postgresql/ @@ -183,7 +185,7 @@ jobs: - name: Install PostgreSQL 15 run: | sudo apt remove -y postgres* - sudo apt -y install curl ca-certificates build-essential pkg-config libssl-dev + sudo apt -y install curl ca-certificates build-essential cmake pkg-config libssl-dev sudo install -d /usr/share/postgresql-common/pgdg sudo curl -o /usr/share/postgresql-common/pgdg/apt.postgresql.org.asc --fail https://www.postgresql.org/media/keys/ACCC4CF8.asc . /etc/os-release diff --git a/.gitignore b/.gitignore index b5a69196d..a47b68ccd 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,11 @@ results/ wrappers/results/ wrappers/regression.* .venv/ + +# Local research and generated non-code artifacts +research/** +outputs/ +.tickets/ +.aura/ +graphify-out/ +.sharding-build-*/ diff --git a/Cargo.lock b/Cargo.lock index b3d9e4c84..689b72c65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -83,6 +83,12 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "ambient-authority" +version = "0.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9d4ee0d472d1cd2e28c97dfa124b3d8d992e10eb0a035f33f5d12e3a177ba3b" + [[package]] name = "android_system_properties" version = "0.1.5" @@ -1478,6 +1484,8 @@ dependencies = [ "cexpr", "clang-sys", "itertools 0.13.0", + "log", + "prettyplease", "proc-macro2", "quote", "regex", @@ -1576,6 +1584,25 @@ dependencies = [ "piper", ] +[[package]] +name = "blosc-rs" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d194effa695b1fdfb156c517b4fb66a6fa1adabf33d19749bd34a483d564f9" +dependencies = [ + "blosc-rs-sys", +] + +[[package]] +name = "blosc-rs-sys" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f2fa86b14bf873f98a075591d212309121642929d812aa759fec348758c8c42" +dependencies = [ + "bindgen", + "cmake", +] + [[package]] name = "bon" version = "3.9.1" @@ -1784,6 +1811,36 @@ dependencies = [ "serde_core", ] +[[package]] +name = "cap-primitives" +version = "3.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e0bf07d379916947be6c4a07f43684153d710a2896c31f9e97781362895596c" +dependencies = [ + "ambient-authority", + "fs-set-times", + "io-extras", + "io-lifetimes 2.0.4", + "ipnet", + "maybe-owned", + "rustix 1.1.4", + "rustix-linux-procfs", + "windows-sys 0.59.0", + "winx", +] + +[[package]] +name = "cap-std" +version = "3.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a59e59fa26472d29680ece6a9f8ee8b0551a719a33df2f5240bde065ecbddfd7" +dependencies = [ + "cap-primitives", + "io-extras", + "io-lifetimes 2.0.4", + "rustix 1.1.4", +] + [[package]] name = "cargo-platform" version = "0.3.3" @@ -3365,6 +3422,17 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs-set-times" +version = "0.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" +dependencies = [ + "io-lifetimes 2.0.4", + "rustix 1.1.4", + "windows-sys 0.59.0", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -4479,6 +4547,16 @@ version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" +[[package]] +name = "io-extras" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2285ddfe3054097ef4b2fe909ef8c3bcd1ea52a8f0d274416caebeef39f04a65" +dependencies = [ + "io-lifetimes 2.0.4", + "windows-sys 0.59.0", +] + [[package]] name = "io-lifetimes" version = "1.0.11" @@ -4490,6 +4568,12 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "io-lifetimes" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06432fb54d3be7964ecd3649233cddf80db2832f47fec34c01f65b3d9d774983" + [[package]] name = "ipconfig" version = "0.3.4" @@ -5131,6 +5215,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4" +[[package]] +name = "maybe-owned" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4facc753ae494aeb6e3c22f839b158aebd4f9270f55cd3c79906c45476c47ab4" + [[package]] name = "md-5" version = "0.10.6" @@ -7489,7 +7579,7 @@ checksum = "519165d378b97752ca44bbe15047d5d3409e875f39327546b42ac81d7e18c1b6" dependencies = [ "bitflags 1.3.2", "errno", - "io-lifetimes", + "io-lifetimes 1.0.11", "libc", "linux-raw-sys 0.3.8", "windows-sys 0.48.0", @@ -7521,6 +7611,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustix-linux-procfs" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fc84bf7e9aa16c4f2c758f27412dc9841341e16aa682d9c7ac308fe3ee12056" +dependencies = [ + "once_cell", + "rustix 1.1.4", +] + [[package]] name = "rustls" version = "0.21.12" @@ -10474,6 +10574,16 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +[[package]] +name = "winx" +version = "0.36.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" +dependencies = [ + "bitflags 2.13.1", + "windows-sys 0.59.0", +] + [[package]] name = "wiremock" version = "0.5.22" @@ -10633,11 +10743,14 @@ dependencies = [ "aws-sdk-s3vectors", "aws-smithy-async", "aws-smithy-types", + "blosc-rs", "bson", "bytes", + "cap-std", "chrono", "chrono-tz 0.6.3", "clickhouse-rs", + "crc32c", "crossbeam", "csv", "dirs", @@ -10652,6 +10765,8 @@ dependencies = [ "iceberg-catalog-rest", "iceberg-catalog-s3tables", "jwt-simple", + "libc", + "lru 0.12.5", "mongodb", "mysql_async", "num-traits 0.2.19", @@ -10681,6 +10796,7 @@ dependencies = [ "wasmtime", "wiremock", "yup-oauth2", + "zstd", ] [[package]] diff --git a/README.md b/README.md index d0637b6d8..9e93bb5be 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ | [Redis](./wrappers/src/fdw/redis_fdw) | A FDW for [Redis](https://redis.io/) | ✅ | ❌ | | [S3](./wrappers/src/fdw/s3_fdw) | A FDW for [AWS S3](https://aws.amazon.com/s3/) | ✅ | ❌ | | [S3 Vectors](./wrappers/src/fdw/s3vectors_fdw) | A FDW for [AWS S3 Vectors](https://aws.amazon.com/s3/features/vectors/) | ✅ | ✅ | +| [Zarr](./wrappers/src/fdw/zarr_fdw) | A read-only FDW for Zarr v2/v3 scientific arrays | ✅ | ❌ | | [SQL Server](./wrappers/src/fdw/mssql_fdw) | A FDW for [Microsoft SQL Server](https://www.microsoft.com/en-au/sql-server/) | ✅ | ❌ | | [Slack](./wasm-wrappers/fdw/slack_fdw) | A Wasm FDW for [Slack](https://www.slack.com/) | ✅ | ❌ | | [Snowflake](./wasm-wrappers/fdw/snowflake_fdw) | A Wasm FDW for [Snowflake](https://www.snowflake.com/) | ✅ | ✅ | diff --git a/docs/catalog/index.md b/docs/catalog/index.md index e55745005..f0fdaa61f 100644 --- a/docs/catalog/index.md +++ b/docs/catalog/index.md @@ -38,6 +38,7 @@ Each FDW documentation includes a detailed "Limitations" section that describes | Redis | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | S3 | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | S3 Vectors | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | +| Zarr | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | | Shopify | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | Snowflake | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | | Stripe | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | diff --git a/docs/catalog/zarr.md b/docs/catalog/zarr.md new file mode 100644 index 000000000..027c09f8a --- /dev/null +++ b/docs/catalog/zarr.md @@ -0,0 +1,952 @@ +--- +source: +documentation: +author: HamzaMPSY(https://github.com/HamzaMPSY) +tags: + - native + - community +--- + +# Zarr + +The Zarr Wrapper provides read-only access to Zarr v2 arrays and a core subset +of Zarr v3 arrays in S3-compatible object storage, trusted anonymous HTTP(S) +object stores, or a secured local filesystem directory. A scan reads one +rank-1 through rank-64 value array whose named dimensions resolve to sibling +coordinate arrays. + +## Enable the wrapper + +```sql +create extension if not exists wrappers with schema extensions; + +create foreign data wrapper zarr_wrapper + handler zarr_fdw_handler + validator zarr_fdw_validator; +``` + +Create a server for the root of one Zarr store: + +```sql +create server public_zarr_server + foreign data wrapper zarr_wrapper + options ( + store_url 's3://example-bucket/datasets/climate.zarr', + aws_region 'us-east-1', + anonymous 'true' + ); +``` + +For S3-compatible services such as MinIO, also set `endpoint_url` and, when +required, `path_style_url 'true'`. Authentication can use the AWS provider +chain, `anonymous 'true'`, a complete direct key pair, or a complete Vault key +pair. Authentication modes cannot be combined. + +An absolute local directory can be used without copying the dataset into +object storage: + +```sql +create server local_zarr_server + foreign data wrapper zarr_wrapper + options ( + store_url 'file:///srv/zarr/climate.zarr' + ); +``` + +Local URLs must have the exact `file:///absolute/path` form, without a host, +userinfo, query, or fragment. S3 authentication, endpoint, region, and +path-style options cannot be used with a local store. Creating or altering a +local Zarr server requires a PostgreSQL superuser, and the foreign server must +remain owned by a superuser at execution time. A superuser can grant `USAGE` +on a fixed server and `SELECT` on its foreign tables to other roles. + +The configured local root and its contents must be administered outside +PostgreSQL and must not be writable by untrusted operating-system users. Reads +reject traversal and symbolic links that escape the configured root; +directory discovery does not follow symbolic-link entries, and final Zarr +objects must be regular files. Missing regular object paths retain normal Zarr +fill-value behavior; permission, file-type, containment, and mutation failures +are errors rather than missing chunks. Errors identify store-relative object +keys without exposing the ambient filesystem root. + +A trusted server that exposes Zarr objects as ordinary anonymous HTTPS `GET` +requests can be used directly: + +```sql +create server https_zarr_server + foreign data wrapper zarr_wrapper + options ( + store_url 'https://datasets.example.org/climate.zarr' + ); +``` + +HTTP(S) store URLs require a host and may contain a port and path, but not +userinfo, credentials, a query, or a fragment. S3 authentication, endpoint, +region, and path-style options cannot be used. The backend does not send +authorization headers or cookies, follow redirects, use ambient HTTP proxies, +retry requests, or transparently decode HTTP content encodings. Requests use +`Accept-Encoding: identity`, and any non-identity response encoding is an +error. Creating or altering an HTTP(S) Zarr server requires a PostgreSQL +superuser, and its foreign-server owner must remain a superuser. Grant fixed +servers to readers with PostgreSQL `USAGE` and `SELECT` privileges. + +HTTPS is required by default and uses normal certificate and hostname +verification. Plain HTTP is unencrypted and must be enabled explicitly only +for a trusted network or local test server: + +```sql +create server insecure_test_zarr_server + foreign data wrapper zarr_wrapper + options ( + store_url 'http://127.0.0.1:8787/climate.zarr', + allow_insecure_http 'true' + ); +``` + +An HTTP object server must return `200` for a complete object and `404` only +when that object is absent. Missing chunks then retain normal Zarr fill-value +behavior; redirects and every other status are errors. Indexed Zarr v3 shards +additionally require single byte-range `GET` support with exact `206`, +`Content-Range`, an exact `Content-Length` when that header is present, and a +quoted strong `ETag`. Payload ranges use `If-Match`; a missing, changed, or +unsatisfiable conditioned object fails instead of combining two shard +generations. Responses remain bounded when `Content-Length` is absent, and +PostgreSQL cancellation is polled while awaiting headers and body chunks. + +Chunk execution is bounded by three optional server settings: + +| Option | Default | Range | +| --- | ---: | ---: | +| `max_concurrent_reads` | `4` | `1`–`32` | +| `max_inflight_bytes` | `269484036` | 1 MiB–1 GiB | +| `compressed_cache_bytes` | `67108864` | `0`–1 GiB; `0` disables caching | + +Reads are prefetched in deterministic chunk order without background tasks. +S3 and HTTP(S) use the configured read concurrency; local stores use one +effective read at a time while retaining the same byte and cache limits. +The compressed cache belongs to one query execution, so bytes never cross +roles, credentials, server changes, or queries. A rescan within the same query +can reuse cached chunks. +For sharded arrays, that same byte budget is divided between decoded shard +indexes and encoded inner-payload ranges; sharding does not add a second +unbounded cache. + +S3 inspection requires `s3:ListBucket` and `s3:GetObject`, while scanning +requires `s3:GetObject`. Local inspection and scans require the PostgreSQL +operating-system account to traverse the configured directory and read the +required metadata and chunk files. HTTP(S) supports exact-object scans, +including explicit OME multiscale selection, but not hierarchy listing; +`zarr_inspect` and `zarr_multiscales` therefore reject HTTP(S) servers before +making an object request. + +## Inspect a dataset + +Use `zarr_inspect` before defining a foreign table: + +```sql +select path, + kind, + variable, + dimensions, + shape, + chunks, + dtype, + units, + calendar +from zarr_inspect('public_zarr_server') +order by path; +``` + +The caller must have `USAGE` on the foreign server. Inspection traverses the +group hierarchy and reads only v2 `.zgroup`, `.zarray`, and `.zattrs` objects or +v3 `zarr.json` objects; it does not read or decode chunk objects. Zarr v3 groups +must be explicit. A node that contains both v2 and v3 metadata is rejected +rather than interpreted using an arbitrary precedence rule. + +The function returns these fields: + +| Field | Meaning | +| --- | --- | +| `path` | Path relative to the configured Zarr root; `/` is the root node | +| `kind` | `group` or `array` | +| `group_path` | Parent group for a non-root node | +| `variable` | Array name; `NULL` for groups | +| `zarr_format` | Format version recorded by `.zgroup`, `.zarray`, or `zarr.json` | +| `shape`, `chunks` | Raw JSON arrays, preserving the metadata integer range | +| `dimensions` | Named dimensions from v2 xarray `_ARRAY_DIMENSIONS` or v3 `dimension_names` | +| `dtype` | Native v2 NumPy dtype string or v3 data-type identifier | +| `codecs` | Native v2 `{filters, compressor}` object or v3 ordered codec array | +| `fill_value` | Raw v2 or v3 fill value | +| `units`, `calendar` | Common scientific attributes when they are strings | +| `scale_factor`, `add_offset` | Finite numeric scientific attributes | +| `crs` | Best-effort CRS metadata from direct `crs`, `spatial_ref`, or `crs_wkt`, or from a resolved sibling `grid_mapping` reference | +| `attributes` | Complete v2 `.zattrs` object or v3 `attributes` object | +| `warnings` | Non-fatal metadata issues, such as malformed named dimensions | + +The inspection surface exposes scientific metadata. Scans can opt into the +CF-style value and time-coordinate decoding described below; physical-unit and +CRS transformations are not applied yet. The complete `attributes` value remains +authoritative because scientific metadata conventions vary between datasets. + +For CRS metadata, `zarr_inspect` keeps the raw node attributes in `attributes` +and fills `crs` as a convenience projection. Direct CRS metadata on +the current node wins in this order: `crs`, `spatial_ref`, then `crs_wkt`. If an +array has no direct CRS metadata but has a simple `grid_mapping` string, the +inspector attempts to resolve that name to a sibling array in the same group and +uses the sibling's direct CRS value. If that reference cannot be resolved, the +raw `grid_mapping` value remains visible in `crs` and a warning is emitted. +CRS strings, WKT, EPSG labels, and GeoTransform attributes remain visible as +metadata. Ordinary foreign-table scans do not assign SRIDs, transform +coordinates, or accept PostGIS predicates. The point-sampling function below +uses supported EPSG metadata from the selected array's resolved grid mapping. + +## Query an array + +After inspection, define a foreign table using PostgreSQL types that match the +array dtype: + +```sql +create foreign table climate_temperature ( + time timestamptz, + y double precision, + x double precision, + temperature real +) +server public_zarr_server +options ( + array_group 'climate/temperature', + time_unit 'seconds', + time_origin 'unix' +); + +select time, y, x, temperature +from climate_temperature +where time >= timestamptz '2025-01-01 00:00:00+00' + and time < timestamptz '2025-01-02 00:00:00+00' + and y between 30 and 31 + and x between -8 and -7; +``` + +Chunk indexes are generated lazily in C order. Their memory is proportional to +array rank rather than the number of selected chunks. PostgreSQL `LIMIT` can +therefore stop later chunk requests naturally, but it does not bypass metadata +reads or the bounded coordinate vectors required for projection and pruning. + +`EXPLAIN ANALYZE` reports actual shape/chunk selection, request and byte counts, +cache activity, synthesized fill bytes, decoded cells, tuple counts, timings, +and aggregate mode. Plain `EXPLAIN` remains network-free, so runtime metadata is +not fabricated or fetched during planning. Chunk-statistic pruning is reported +as disabled until the separately validated statistics catalog is implemented. +The query-local EXPLAIN counters include work initiated before an early `LIMIT`. +Errors and cancellations clean up queued reads safely but do not return an +`EXPLAIN ANALYZE` plan. The older persistent `wrappers_fdw_stats` counters are +flushed only when the iterator reaches EOF; like other Wrappers FDWs, an executor +that stops early may not persist the final delta because SPI is not safe from +`EndForeignScan`. + +The selected value array must have one unique, safe name for every array +dimension, in array order: v2 uses the xarray `_ARRAY_DIMENSIONS` attribute and +v3 uses native `dimension_names`. If a v3 array also carries the legacy xarray +attribute, the two declarations must match exactly. Each name must resolve to a +same-group, same-name coordinate array in the same Zarr format whose shape is +one dimensional and whose length matches the value-array extent. If a +coordinate array declares native or legacy dimension names, it must contain +only its own name. Missing or malformed dimension metadata fails instead of +falling back to an inferred `[y, x]` or `[time, y, x]` layout. + +Dimension names are preserved for PostgreSQL column matching. Coordinate +metadata can classify dimensions as spatial X/Y, latitude/longitude, vertical, +time, band, channel, or unknown; roles do not rename dimensions. Recognized +`standard_name`, `axis`, and unambiguous units take precedence over conservative +name aliases. Incompatible recognized signals fail instead of being guessed. +Names such as `depth`, `height`, `altitude`, `level`, `lev`, and `z` are vertical +aliases, while band and channel remain distinct roles. + +Supported v2 value mappings are ` 'public.spatial_temperature', + point_ewkb => gis.ST_AsEWKB( + gis.ST_SetSRID(gis.ST_MakePoint(110, 20), 3857) + ), + method => 'nearest' +); +``` + +The function signature is: + +```sql +zarr_sample( + foreign_table text, + point_ewkb bytea, + method text default 'nearest' +) +returns table ( + x double precision, + y double precision, + value double precision, + x_index bigint, + y_index bigint, + coordinate_distance double precision, + srid integer +) +``` + +Selector-bearing tables must use the explicit selector-aware overload: + +```sql +zarr_sample( + foreign_table text, + point_ewkb bytea, + method text, + dimension_selectors text +) +``` + +The `dimension_selectors` argument uses the same JSON grammar as the foreign +table option and has no default. Pass `'{}'` when only the table option should +apply. The point lookup owns the horizontal X/Y dimensions, so selectors may +target auxiliary dimensions only. Every auxiliary dimension must resolve to +zero or one exact native index; zero returns no sample row, while more than one +index fails clearly. + +Use a schema-qualified foreign-table name. The caller must have `SELECT` on the +foreign table and `USAGE` on its foreign server. The table must select one Zarr +value array with exactly two discovered, one-dimensional horizontal coordinate +axes and a supported, unambiguous EPSG CRS. The legacy signature requires a +rank-2 array. The selector-aware overload also accepts auxiliary dimensions +when each resolves to zero or one exact native index. EWKB must contain a +nonzero SRID; the point is transformed into the array CRS before coordinate +lookup. + +`nearest` independently selects the closest stored x and y coordinate. A tie is +resolved to the lower logical array index, including on descending axes. +`exact` returns a sample only when the transformed point exactly matches both +stored coordinate values. `coordinate_distance` is zero for an exact match and +otherwise is Euclidean distance in the array CRS units; it is not a geodesic +distance. Neither method interpolates values. Point lookup currently supports +rectilinear cell centers only, not curvilinear coordinates, cell footprints, or +rotated affine grids. + +The scalar result is widened to `double precision`. The function uses the +foreign table's scientific decoding options, so `decode_cf`, fill/missing +values, valid ranges, and scale/offset have the same meaning as an ordinary +scan. A scientifically missing value is returned as SQL `NULL`. + +## Select polygon cells and calculate zonal statistics + +`zarr_cells` returns the cell centers covered by a PostGIS Polygon or +MultiPolygon, while `zarr_zonal_stats` reduces the same selected values inside +the Zarr executor: + +```sql +with region as ( + select gis.ST_AsEWKB( + gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) as ewkb +) +select cells.* +from region +cross join lateral zarr_cells( + foreign_table => 'public.spatial_temperature', + region_ewkb => region.ewkb +) as cells +order by cells.y_index, cells.x_index; + +with region as ( + select gis.ST_AsEWKB( + gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) as ewkb +) +select stats.* +from region +cross join lateral zarr_zonal_stats( + foreign_table => 'public.spatial_temperature', + region_ewkb => region.ewkb +) as stats; +``` + +The function signatures are: + +```sql +zarr_cells(foreign_table text, region_ewkb bytea) +returns table ( + x double precision, + y double precision, + value double precision, + x_index bigint, + y_index bigint, + srid integer +); + +zarr_zonal_stats(foreign_table text, region_ewkb bytea) +returns table ( + count bigint, + valid_count bigint, + min double precision, + max double precision, + sum double precision, + avg double precision, + srid integer +); +``` + +Selector-bearing tables must use the explicit selector-aware reduction +overload: + +```sql +zarr_zonal_stats( + foreign_table text, + region_ewkb bytea, + dimension_selectors text +) +``` + +The polygon owns horizontal X/Y selection. Table and call selectors intersect +for auxiliary dimensions only. Pass `'{}'` to opt into table selectors without +adding call selectors. If an auxiliary dimension resolves to no exact index, +the function returns one empty statistics row; if it resolves to multiple +indexes, the function fails rather than aggregating unlabeled slices. + +Both legacy functions apply the same foreign-table privilege, +rectilinear-grid, CRS, rank-2, and scientific-decoding rules as the legacy +`zarr_sample` signature. The selector-aware `zarr_zonal_stats` overload also +accepts auxiliary dimensions when each resolves to zero or one exact native +index; `zarr_cells` remains rank-2 and has no selector overload. Region EWKB +must contain a valid, non-empty, two-dimensional Polygon or MultiPolygon with a +positive SRID. The geometry is transformed into the array CRS, and its envelope +is used only for conservative coordinate and chunk pruning. + +Exact inclusion uses PostGIS `ST_Covers(region, cell_center)` semantics. A cell +center on the polygon boundary is therefore included. The returned cells are +center samples, not pixel footprints, and no partial-cell area weighting is +performed. Function output order is unspecified; add an `ORDER BY` when stable +ordering is required. + +`count` is the number of covered logical cells, including cells whose decoded +value is SQL `NULL`. `valid_count`, `min`, `max`, `sum`, and `avg` ignore decoded +NULL values. With no valid values, `valid_count` is zero and the numeric +aggregates are NULL. Values and aggregates are widened to `double precision`. + +### Add a time range to polygon queries + +For an array with one discovered time dimension, use +`zarr_cells_by_time` to return covered cells at each stored timestamp, or +`zarr_zonal_stats_by_time` to return one aggregate row per logical time index: + +```sql +create foreign table spatial_temperature_by_time ( + time timestamptz, + y double precision, + x double precision, + temperature real +) +server public_zarr_server +options ( + array_group 'climate/spatial_temperature_by_time', + time_from_attrs 'true' +); + +with region as ( + select gis.ST_AsEWKB( + gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) as ewkb +) +select stats.* +from region +cross join lateral zarr_zonal_stats_by_time( + foreign_table => 'public.spatial_temperature_by_time', + region_ewkb => region.ewkb, + start_time => timestamptz '2025-01-01 00:00:00+00', + end_time => timestamptz '2025-01-02 00:00:00+00' +) as stats +order by stats.time_index; +``` + +The function signatures are: + +```sql +zarr_cells_by_time( + foreign_table text, + region_ewkb bytea, + start_time timestamptz, + end_time timestamptz +) +returns table ( + time timestamptz, + x double precision, + y double precision, + value double precision, + time_index bigint, + x_index bigint, + y_index bigint, + srid integer +); + +zarr_zonal_stats_by_time( + foreign_table text, + region_ewkb bytea, + start_time timestamptz, + end_time timestamptz +) +returns table ( + time timestamptz, + time_index bigint, + count bigint, + valid_count bigint, + min double precision, + max double precision, + sum double precision, + avg double precision, + srid integer +); +``` + +`zarr_zonal_stats_by_time` also has a selector-aware overload: + +```sql +zarr_zonal_stats_by_time( + foreign_table text, + region_ewkb bytea, + start_time timestamptz, + end_time timestamptz, + dimension_selectors text +) +``` + +The polygon owns X/Y and the time range owns the Time dimension; selectors may +target only other auxiliary dimensions. Pass `'{}'` to opt into selector-aware +execution with only table selectors. A non-overlapping time range returns no +rows. A spatial or auxiliary empty selection with matching time indexes returns +one empty statistics row per selected time index. + +Time bounds are required and form a half-open range: `start_time` is included +and `end_time` is excluded. The FDW uses the same manual or attribute-derived +time conversion as an ordinary scan. It discovers Time, X, and Y roles rather +than relying on dimension names or positions, and supports any array-axis order. +Additional dimensions are accepted only when their extent is one; a +non-singleton band, level, channel, or unknown dimension must resolve to zero +or one exact index through the selector-aware overload and is otherwise +rejected. + +Unordered time coordinates are scanned conservatively and checked exactly, so +matching timestamps are not pruned incorrectly. Duplicate stored timestamps +remain distinct rows through `time_index`. Zonal output contains one row for +each selected logical time index; a slice with no covered or valid values has +zero counts and NULL numeric aggregates. If no stored timestamp falls in the +requested range, both functions return no rows. + +The spatial-time candidate window is limited to 10,000,000 logical cells, +`zarr_cells_by_time` returns at most 1,000,000 rows, and at most 1,000,000 time +slices may be selected. These functions preserve the same PostGIS boundary, +CRS, privilege, scientific-decoding, cache, and cancellation rules as the +rank-2 polygon functions. + +## Decode time coordinates from attributes + +By default, raw values from the one coordinate classified as time are +interpreted from the manual table options +`time_unit` and `time_origin`, or as `seconds` since the Unix epoch when those +options are omitted. + +Set `time_from_attrs 'true'` to derive the time conversion from the sibling +coordinate's attributes instead: + +```sql +create foreign table climate_temperature_from_attrs ( + time timestamptz, + y double precision, + x double precision, + temperature real +) +server public_zarr_server +options ( + array_group 'climate/temperature', + time_from_attrs 'true' +); +``` + +This mode is intentionally opt-in and cannot be combined with `time_unit` or +`time_origin`. It works at any supported rank and with any dimension name, but +requires exactly one coordinate classified as time whose attributes contain: + +- `units` as ` since `; +- `calendar` as `proleptic_gregorian`. + +Supported constant-duration units are `seconds`, `milliseconds`, +`microseconds`, `nanoseconds`, `minutes`, `hours`, and `days`. The origin may +be a representable Gregorian date, date-time, or RFC 3339 date-time; an origin +without an explicit timezone is interpreted as UTC. Unsupported calendars, +malformed units, missing metadata, and out-of-range conversions fail clearly. + +The resolved time conversion is used for both emitted `timestamptz` values and +timestamp predicate pruning. For sub-microsecond units, pruning conservatively +covers every raw value that can round to the PostgreSQL timestamp; PostgreSQL +still rechecks the exact predicate. The FDW performs no remote metadata I/O +during planning; metadata is read only when a scan starts. + +## Decode packed scientific values + +Set `decode_cf 'true'` on a foreign table to apply common CF-style missing-data +and packed-value attributes from the selected value array: + +```sql +create foreign table decoded_temperature ( + time timestamptz, + y double precision, + x double precision, + temperature double precision +) +server public_zarr_server +options ( + array_group 'climate/packed_temperature', + time_unit 'seconds', + time_origin 'unix', + decode_cf 'true' +); +``` + +Decoded mode applies this order: + +1. Decode the stored primitive value. +2. Map `_FillValue`, `missing_value`, and values outside `valid_range` or + `valid_min`/`valid_max` to SQL `NULL`. +3. Return `raw * scale_factor + add_offset` as `double precision`. + +Masking and valid-range checks happen before scale/offset, in the packed/raw +domain. A missing Zarr chunk is first materialized with the array's `fill_value`; +it becomes SQL `NULL` only when that raw value also matches the +scientific missing/validity metadata. The option defaults to `false`, which +preserves the raw dtype mappings above. + +For floating-point arrays, the Zarr JSON spellings `"NaN"`, `"Infinity"`, and +`"-Infinity"` are accepted as missing sentinels. Declaring `"NaN"` masks every +NaN payload; an undeclared non-finite value remains a PostgreSQL non-finite +`double precision` value. + +This value-decoding mode is independent from `time_from_attrs`. It does not +convert physical units, transform a CRS, or apply packing attributes to +coordinate arrays. + +## Aggregate pushdown + +The wrapper reduces ungrouped `count`, `sum`, `avg`, `min`, and `max` queries +inside the Zarr chunk scan and returns one result row to PostgreSQL. This avoids +creating one PostgreSQL tuple for every selected array cell: + +```sql +select count(*) as selected_cells, + count(temperature) as valid_cells, + min(temperature), + max(temperature), + sum(temperature), + avg(temperature) +from decoded_temperature +where time >= timestamptz '2025-01-01 00:00:00+00' + and time < timestamptz '2025-02-01 00:00:00+00' + and y between 30 and 31 + and x in (-8.0, -7.5, -7.0); +``` + +Chunk ranges are still selected conservatively, but aggregate mode evaluates +each accepted predicate exactly before updating the reducer. This preserves +strict inequalities, non-contiguous `IN` membership, unordered coordinates, +value-column predicates, missing chunks, edge chunks, and decoded NULL +semantics. `count(*)` includes matching logical cells whose value is NULL; +`count(column)`, `sum`, `avg`, `min`, and `max` ignore NULL. Non-count +aggregates return NULL for an empty or all-NULL selection. + +Pushdown currently applies only to scalar aggregates over plain columns. A +query with `GROUP BY`, `DISTINCT`, an aggregate `FILTER` or `HAVING` clause, an +aggregate expression such as `sum(temperature + 1)`, or a predicate the wrapper +cannot evaluate exactly remains a normal foreign scan with PostgreSQL doing +the aggregation. Use `EXPLAIN` to confirm whether a query is represented by a +single Foreign Scan or retains a local Aggregate node. + +## Current limitations + +- Read-only Zarr v2 and the core Zarr v3 subset described above on + S3-compatible storage, trusted anonymous HTTP(S) object stores, or a secured + local filesystem directory. Authenticated HTTP, redirects, proxies, custom + certificate authorities, mutual TLS, WebDAV, GCS, Azure, and SSH filesystem + URLs are not supported. Plain HTTP requires explicit opt-in and provides no + transport confidentiality or integrity. +- Scans support one value array with rank 1 through 64 and mandatory v2 + `_ARRAY_DIMENSIONS` or v3 `dimension_names`; scalar arrays remain unsupported. +- Ordinary arrays require a same-group, same-name, rank-1 numeric coordinate + array for every dimension. Explicitly selected supported OME-Zarr 0.5 + rank-2 levels instead synthesize `y` and `x` from their scale/translation + metadata. Other synthesized ordinal coordinates, auxiliary or cross-group + coordinates, curvilinear/multidimensional coordinates, and string or + categorical band/channel coordinates are not supported. `dimension_selectors` + accept only numeric coordinate values and zero-based physical indexes; string + and categorical selector values are not supported. +- Coordinate packing, masks, valid ranges, and scale/offset are not decoded. If + a coordinate used by a query declares those attributes, the scan fails rather + than silently ignoring them. +- One temporal dimension is supported. Multiple temporal dimensions and + per-axis calendars are not supported. +- A foreign-table scan still represents one value array and at most one queried + non-dimension value column. Multi-variable scans and functional `bands` + execution are not supported. +- Spatial functions are limited to rank-2 rectilinear arrays with one discovered + horizontal x/longitude axis and one y/latitude axis. Polygon operations use + center coverage only. They do not implement cell-footprint or area-weighted + statistics, geographic distance, interpolation, curvilinear coordinates, + rotated grids, or topology repair. +- Aggregate pushdown is limited to ungrouped `count`, `sum`, `avg`, `min`, and + `max` over plain columns. Grouped, distinct, filtered, ordered, expression, + and user-defined aggregates are computed by PostgreSQL. +- Raw, gzip, zlib, and Blosc/LZ4 chunk compression is supported for v2. The v3 + subset supports the ordered `transpose`, `bytes`, `gzip`, bounded Blosc, or + bounded Zstandard, and `crc32c` pipeline described above. +- Non-empty v2 filters, Fortran order, consolidated metadata, storage + transformers, writes, and OME-Zarr semantics outside the bounded 0.5 + multiscale subset above are not supported. +- `LIMIT` alone does not prevent coordinate metadata loading; use selective + coordinate predicates for large arrays. It does stop later lazy data-chunk + reads after PostgreSQL has accepted enough rows. +- Scan execution limits each loaded coordinate and all loaded coordinates + together to 16,777,216 values. Chunk-index iteration is O(rank); decoded + chunks retain their 256 MiB limit and storage concurrency/cache use the + server byte limits above. These are safety bounds, not Zarr format limits. +- Inspection has hard depth, node, list-page, object-size, and total metadata + limits and fails explicitly rather than returning a truncated hierarchy. + HTTP(S) stores cannot be inspected because this backend deliberately has no + directory-listing protocol. diff --git a/mkdocs.yaml b/mkdocs.yaml index 8b7bdcb55..32c520d0b 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -30,6 +30,7 @@ nav: - Redis: "catalog/redis.md" - S3 (CSV, JSON, Parquet): "catalog/s3.md" - S3 Vectors: "catalog/s3vectors.md" + - Zarr: "catalog/zarr.md" - Stripe: "catalog/stripe.md" - SQL Server: "catalog/mssql.md" - Wasm: diff --git a/supabase-wrappers/src/interface.rs b/supabase-wrappers/src/interface.rs index 10c01ec73..514b161ac 100644 --- a/supabase-wrappers/src/interface.rs +++ b/supabase-wrappers/src/interface.rs @@ -472,6 +472,20 @@ pub enum Value { Array(Vec), } +/// The runtime state of a parameter used by a [`Qual`]. +/// +/// This is separate from [`Value`] so adding SQL NULL support does not add a +/// new variant to that public enum and break existing exhaustive matches. +#[derive(Debug, Clone)] +pub enum ParamValue { + /// The executor has not supplied a value yet. + Unevaluated, + /// The executor supplied SQL NULL. + Null, + /// The executor supplied a non-NULL value. + Value(Value), +} + // Struct for parameter expression value evaluation #[derive(Debug, Clone)] pub(super) struct ExprEval { @@ -497,10 +511,54 @@ pub struct Param { /// parameter value which is evaluated during query execution pub eval_value: Arc>>, + /// Explicit runtime state, including the distinction between unevaluated + /// and SQL NULL. `eval_value` remains available for source compatibility. + pub(super) eval_state: Arc>, + // internal variables for expression evaluation pub(super) expr_eval: ExprEval, } +impl Param { + /// Return the current runtime parameter state. + pub fn evaluated_value(&self) -> ParamValue { + self.eval_state + .lock() + .expect("parameter evaluation state should be locked") + .clone() + } + + pub(super) fn set_evaluated_value(&self, value: ParamValue) { + let legacy_value = match &value { + ParamValue::Value(value) => Some(value.clone()), + ParamValue::Unevaluated | ParamValue::Null => None, + }; + *self + .eval_value + .lock() + .expect("parameter eval value should be locked") = legacy_value; + *self + .eval_state + .lock() + .expect("parameter evaluation state should be locked") = value; + } + + #[cfg(test)] + pub(super) fn clone_for_execution(&self) -> Self { + Self { + kind: self.kind, + id: self.id, + type_oid: self.type_oid, + eval_value: Mutex::new(None).into(), + eval_state: Mutex::new(ParamValue::Unevaluated).into(), + expr_eval: ExprEval { + expr: self.expr_eval.expr, + expr_state: std::ptr::null_mut(), + }, + } + } +} + /// Query restrictions, a.k.a conditions in `WHERE` clause /// /// A Qual defines a simple condition wich can be used by the FDW to restrict the number @@ -554,6 +612,28 @@ pub struct Qual { } impl Qual { + #[cfg(test)] + pub(super) fn clone_for_execution(&self) -> Self { + Self { + field: self.field.clone(), + operator: self.operator.clone(), + value: self.value.clone(), + use_or: self.use_or, + param: self.param.as_ref().map(Param::clone_for_execution), + } + } + + /// Return the effective runtime value of this restriction. + /// + /// Constant restrictions return [`ParamValue::Value`]. Parameterized + /// restrictions expose SQL NULL and the pre-execution state explicitly. + pub fn evaluated_value(&self) -> ParamValue { + self.param + .as_ref() + .map(Param::evaluated_value) + .unwrap_or_else(|| ParamValue::Value(self.value.clone())) + } + pub fn deparse(&self) -> String { let mut formatter = DefaultFormatter::new(); self.deparse_with_fmt(&mut formatter) @@ -851,6 +931,121 @@ impl Aggregate { } } +/// A typed value emitted by an FDW's runtime [`EXPLAIN`](ForeignDataWrapper::explain) hook. +/// +/// Keeping values typed preserves numbers and booleans in structured EXPLAIN +/// formats instead of forcing every FDW-specific property through text. +#[derive(Debug, Clone, PartialEq)] +pub enum ExplainValue { + Text(String), + Integer { + value: i64, + unit: Option, + }, + Unsigned { + value: u64, + unit: Option, + }, + Float { + value: f64, + unit: Option, + digits: i32, + }, + Boolean(bool), +} + +/// One FDW-specific property to append to `EXPLAIN ANALYZE` output. +#[derive(Debug, Clone, PartialEq)] +pub struct ExplainProperty { + pub label: String, + pub value: ExplainValue, +} + +impl ExplainProperty { + pub fn text(label: impl Into, value: impl Into) -> Self { + Self { + label: label.into(), + value: ExplainValue::Text(value.into()), + } + } + + pub fn integer(label: impl Into, value: i64) -> Self { + Self { + label: label.into(), + value: ExplainValue::Integer { value, unit: None }, + } + } + + pub fn integer_with_unit( + label: impl Into, + value: i64, + unit: impl Into, + ) -> Self { + Self { + label: label.into(), + value: ExplainValue::Integer { + value, + unit: Some(unit.into()), + }, + } + } + + pub fn unsigned(label: impl Into, value: u64) -> Self { + Self { + label: label.into(), + value: ExplainValue::Unsigned { value, unit: None }, + } + } + + pub fn unsigned_with_unit( + label: impl Into, + value: u64, + unit: impl Into, + ) -> Self { + Self { + label: label.into(), + value: ExplainValue::Unsigned { + value, + unit: Some(unit.into()), + }, + } + } + + pub fn float(label: impl Into, value: f64, digits: i32) -> Self { + Self { + label: label.into(), + value: ExplainValue::Float { + value, + unit: None, + digits, + }, + } + } + + pub fn float_with_unit( + label: impl Into, + value: f64, + unit: impl Into, + digits: i32, + ) -> Self { + Self { + label: label.into(), + value: ExplainValue::Float { + value, + unit: Some(unit.into()), + digits, + }, + } + } + + pub fn boolean(label: impl Into, value: bool) -> Self { + Self { + label: label.into(), + value: ExplainValue::Boolean(value), + } + } +} + /// The Foreign Data Wrapper trait /// /// This is the main interface for your foreign data wrapper. Required functions @@ -939,6 +1134,15 @@ pub trait ForeignDataWrapper> { /// [See more details](https://www.postgresql.org/docs/current/fdw-callbacks.html#FDW-CALLBACKS-SCAN). fn end_scan(&mut self) -> Result<(), E>; + /// Return FDW-specific properties for a live scan's `EXPLAIN ANALYZE` output. + /// + /// The framework calls this only when execution constructed an FDW + /// instance. Plain `EXPLAIN` remains planning-only and never creates a + /// remote client for this hook. The default keeps existing FDWs unchanged. + fn explain(&self) -> Vec { + Vec::new() + } + /// Called when begin executing a foreign table modification operation. /// /// - `options` - the options defined when `CREATE FOREIGN TABLE` @@ -1054,6 +1258,27 @@ pub trait ForeignDataWrapper> { false } + /// Decide whether a specific aggregate query is safe for this FDW to push down. + /// + /// This query-specific hook runs after the framework has extracted the + /// aggregates and GROUP BY columns, but before it registers a foreign upper + /// path. The default preserves the behavior of existing aggregate-capable + /// FDWs. Implementations that require exact local predicate evaluation or + /// support only selected input/result type pairs can reject unsupported + /// shapes by returning `Ok(false)`. + #[allow(clippy::too_many_arguments)] + fn can_pushdown_aggregate( + &mut self, + _aggregates: &[Aggregate], + _group_by: &[Column], + _quals: &[Qual], + _base_columns: &[Column], + _all_base_quals_extracted: bool, + _options: &HashMap, + ) -> Result { + Ok(true) + } + /// Estimate the size of aggregate query results for query planning. /// /// Called during query planning when aggregate pushdown is being considered. @@ -1171,6 +1396,23 @@ pub trait ForeignDataWrapper> { unreachable!() } + /// Begin aggregate execution with the base columns retained by the planner. + /// + /// The default delegates to [`begin_aggregate_scan`](Self::begin_aggregate_scan) + /// so existing FDWs remain source compatible. Local aggregate engines can + /// override this hook when they need the original input columns after the + /// ForeignScan target list has been replaced by aggregate result columns. + fn begin_aggregate_scan_with_base_columns( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + quals: &[Qual], + _base_columns: &[Column], + options: &HashMap, + ) -> Result<(), E> { + self.begin_aggregate_scan(aggregates, group_by, quals, options) + } + /// Obtain a list of foreign table creation commands /// /// Return a list of string, each of which must contain a CREATE FOREIGN TABLE @@ -1307,6 +1549,129 @@ pub trait ForeignDataWrapper> { mod tests { use super::*; + fn test_param() -> Param { + Param { + kind: pg_sys::ParamKind::PARAM_EXTERN, + id: 1, + type_oid: pg_sys::INT4OID, + eval_value: Mutex::new(None).into(), + eval_state: Mutex::new(ParamValue::Unevaluated).into(), + expr_eval: ExprEval { + expr: std::ptr::null_mut(), + expr_state: std::ptr::null_mut(), + }, + } + } + + #[test] + fn test_parameter_value_distinguishes_unevaluated_null_and_value() { + let param = test_param(); + assert!(matches!(param.evaluated_value(), ParamValue::Unevaluated)); + + param.set_evaluated_value(ParamValue::Null); + assert!(matches!(param.evaluated_value(), ParamValue::Null)); + assert!( + param + .eval_value + .lock() + .expect("legacy eval value should be locked") + .is_none() + ); + + param.set_evaluated_value(ParamValue::Value(Value::Cell(Cell::I32(42)))); + assert!(matches!( + param.evaluated_value(), + ParamValue::Value(Value::Cell(Cell::I32(42))) + )); + assert!(matches!( + &*param + .eval_value + .lock() + .expect("legacy eval value should be locked"), + Some(Value::Cell(Cell::I32(42))) + )); + + param.set_evaluated_value(ParamValue::Null); + assert!(matches!(param.evaluated_value(), ParamValue::Null)); + assert!( + param + .eval_value + .lock() + .expect("legacy eval value should be locked") + .is_none() + ); + } + + #[test] + fn test_qual_evaluated_value_uses_constant_or_parameter_state() { + let constant = Qual { + field: "value".to_string(), + operator: "=".to_string(), + value: Value::Cell(Cell::I32(7)), + use_or: false, + param: None, + }; + assert!(matches!( + constant.evaluated_value(), + ParamValue::Value(Value::Cell(Cell::I32(7))) + )); + + let parameter = test_param(); + parameter.set_evaluated_value(ParamValue::Null); + let parameterized = Qual { + field: "value".to_string(), + operator: "=".to_string(), + value: Value::Cell(Cell::I64(0)), + use_or: false, + param: Some(parameter), + }; + assert!(matches!(parameterized.evaluated_value(), ParamValue::Null)); + } + + #[test] + fn test_qual_execution_clone_has_independent_parameter_state() { + let parameter = test_param(); + parameter.set_evaluated_value(ParamValue::Value(Value::Cell(Cell::I32(11)))); + let planned = Qual { + field: "value".to_string(), + operator: "=".to_string(), + value: Value::Cell(Cell::I64(0)), + use_or: false, + param: Some(parameter), + }; + + let execution = planned.clone_for_execution(); + let planned_param = planned.param.as_ref().expect("planned parameter"); + let execution_param = execution.param.as_ref().expect("execution parameter"); + + assert!(matches!( + planned_param.evaluated_value(), + ParamValue::Value(Value::Cell(Cell::I32(11))) + )); + assert!(matches!( + execution_param.evaluated_value(), + ParamValue::Unevaluated + )); + assert!(!Arc::ptr_eq( + &planned_param.eval_state, + &execution_param.eval_state + )); + assert!(!Arc::ptr_eq( + &planned_param.eval_value, + &execution_param.eval_value + )); + + execution_param.set_evaluated_value(ParamValue::Null); + assert!(matches!( + planned_param.evaluated_value(), + ParamValue::Value(Value::Cell(Cell::I32(11))) + )); + assert!(matches!( + execution_param.evaluated_value(), + ParamValue::Null + )); + } + #[test] fn test_cell_into_datum_type_oid_is_invalid() { assert_eq!(Cell::type_oid(), Oid::INVALID); diff --git a/supabase-wrappers/src/qual.rs b/supabase-wrappers/src/qual.rs index 429240daa..6b2ec113b 100644 --- a/supabase-wrappers/src/qual.rs +++ b/supabase-wrappers/src/qual.rs @@ -14,7 +14,12 @@ use std::os::raw::c_int; use std::ptr; use std::sync::Mutex; -use crate::interface::Param; +use crate::interface::{Param, ParamValue}; + +pub(crate) fn is_builtin_pg_catalog_object(oid: Oid, namespace: Oid) -> bool { + namespace == Oid::from(pg_sys::PG_CATALOG_NAMESPACE) + && oid.to_u32() < pg_sys::FirstNormalObjectId +} /// Parses a Postgres array `Datum` (given its element/array type OID) into a `Vec`. /// @@ -146,19 +151,31 @@ pub unsafe fn form_array_from_datum( } } -pub(crate) unsafe fn get_operator(opno: pg_sys::Oid) -> pg_sys::Form_pg_operator { +struct OperatorInfo { + name: String, + commutator: pg_sys::Oid, +} + +unsafe fn get_builtin_operator(opno: pg_sys::Oid) -> Option { unsafe { let htup = pg_sys::SearchSysCache1( pg_sys::SysCacheIdentifier::OPEROID.try_into().unwrap(), opno.into(), ); if htup.is_null() { - pg_sys::ReleaseSysCache(htup); pgrx::error!("cache lookup operator {:?} failed", opno); } let op = pg_sys::GETSTRUCT(htup) as pg_sys::Form_pg_operator; + let operator = if is_builtin_pg_catalog_object(opno, (*op).oprnamespace) { + Some(OperatorInfo { + name: pgrx::name_data_to_str(&(*op).oprname).to_string(), + commutator: (*op).oprcom, + }) + } else { + None + }; pg_sys::ReleaseSysCache(htup); - op + operator } } @@ -191,11 +208,10 @@ pub(crate) unsafe fn extract_from_op_expr( // get operator let opno = (*expr).opno; - let opr = get_operator(opno); - if opr.is_null() { - report_warning("operator is empty"); + let Some(mut operator) = get_builtin_operator(opno) else { + report_warning("only built-in pg_catalog operators are supported in quals"); return None; - } + }; let mut left = unnest_clause(*args.get(0).unwrap() as _); let mut right = unnest_clause(*args.get(1).unwrap() as _); @@ -203,8 +219,10 @@ pub(crate) unsafe fn extract_from_op_expr( // swap operands if needed if is_a(right, pg_sys::NodeTag::T_Var) && !is_a(left, pg_sys::NodeTag::T_Var) - && (*opr).oprcom != Oid::INVALID + && operator.commutator != Oid::INVALID { + let commutator = get_builtin_operator(operator.commutator)?; + operator = commutator; std::mem::swap(&mut left, &mut right); } @@ -235,6 +253,7 @@ pub(crate) unsafe fn extract_from_op_expr( id: (*right).paramid as _, type_oid: (*right).paramtype, eval_value: Mutex::new(None).into(), + eval_state: Mutex::new(ParamValue::Unevaluated).into(), expr_eval: ExprEval { expr: if (*right).paramkind == pg_sys::ParamKind::PARAM_EXEC { right as _ @@ -252,7 +271,7 @@ pub(crate) unsafe fn extract_from_op_expr( if let Some(value) = value { let qual = Qual { field: CStr::from_ptr(field).to_str().unwrap().to_string(), - operator: pgrx::name_data_to_str(&(*opr).oprname).to_string(), + operator: operator.name, value: Value::Cell(value), use_or: false, param, @@ -318,10 +337,10 @@ pub(crate) unsafe fn extract_from_scalar_array_op_expr( // get operator let opno = (*expr).opno; - let opr = get_operator(opno); - if opr.is_null() { + let Some(operator) = get_builtin_operator(opno) else { + report_warning("only built-in pg_catalog operators are supported in quals"); return None; - } + }; let left = unnest_clause(*args.get(0).unwrap() as _); let right = unnest_clause(*args.get(1).unwrap() as _); @@ -343,7 +362,7 @@ pub(crate) unsafe fn extract_from_scalar_array_op_expr( if let Some(value) = value { let qual = Qual { field: CStr::from_ptr(field).to_str().unwrap().to_string(), - operator: pgrx::name_data_to_str(&(*opr).oprname).to_string(), + operator: operator.name, value: Value::Array(value), use_or: (*expr).useOr, param: None, @@ -462,14 +481,20 @@ pub(crate) unsafe fn extract_from_boolean_test( } } +pub(crate) struct ExtractedQuals { + pub(crate) quals: Vec, + pub(crate) all_extracted: bool, +} + pub(crate) unsafe fn extract_quals( root: *mut pg_sys::PlannerInfo, baserel: *mut pg_sys::RelOptInfo, baserel_id: pg_sys::Oid, -) -> Vec { +) -> ExtractedQuals { unsafe { pgrx::memcx::current_context(|mcx| { let mut quals = Vec::new(); + let mut all_extracted = true; if let Some(conds) = List::<*mut c_void>::downcast_ptr_in_memcx((*baserel).baserestrictinfo, mcx) @@ -502,11 +527,39 @@ pub(crate) unsafe fn extract_quals( if let Some(qual) = extracted { quals.push(qual); + } else { + all_extracted = false; } } + } else if !(*baserel).baserestrictinfo.is_null() { + all_extracted = false; } - quals + ExtractedQuals { + quals, + all_extracted, + } }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_builtin_pg_catalog_object_identity() { + let catalog = Oid::from(pg_sys::PG_CATALOG_NAMESPACE); + let custom_namespace = Oid::from(pg_sys::PG_CATALOG_NAMESPACE + 1); + + assert!(is_builtin_pg_catalog_object(Oid::from(96_u32), catalog)); + assert!(!is_builtin_pg_catalog_object( + Oid::from(96_u32), + custom_namespace + )); + assert!(!is_builtin_pg_catalog_object( + Oid::from(pg_sys::FirstNormalObjectId), + catalog + )); + } +} diff --git a/supabase-wrappers/src/scan.rs b/supabase-wrappers/src/scan.rs index 436c7d7d9..3a2b08efa 100644 --- a/supabase-wrappers/src/scan.rs +++ b/supabase-wrappers/src/scan.rs @@ -1,11 +1,13 @@ use pgrx::FromDatum; use pgrx::{ IntoDatum, PgSqlErrorCode, debug2, + list::List, memcxt::PgMemoryContexts, pg_sys::{Datum, MemoryContext, MemoryContextData, Oid, ParamKind}, prelude::*, }; use std::collections::HashMap; +use std::ffi::c_void; use std::marker::PhantomData; use pgrx::pg_sys::panic::ErrorReport; @@ -13,7 +15,10 @@ use std::os::raw::c_int; use std::ptr; use crate::instance; -use crate::interface::{Aggregate, Cell, Column, Limit, Qual, Row, Sort, Value}; +use crate::interface::{ + Aggregate, AggregateKind, Cell, Column, ExplainProperty, ExplainValue, ExprEval, Limit, Param, + ParamValue, Qual, Row, Sort, Value, +}; use crate::limit::*; use crate::memctx; use crate::options::options_to_hashmap; @@ -21,10 +26,631 @@ use crate::polyfill; use crate::prelude::ForeignDataWrapper; use crate::qual::*; use crate::sort::*; -use crate::utils::{self, ReportableError, SerdeList, report_error}; +use crate::utils::{self, ReportableError, report_error}; + +const FDW_SCAN_PRIVATE_VERSION: &str = "wrappers-scan-v1"; + +unsafe extern "C" { + #[link_name = "datumCopy"] + fn datum_copy(value: Datum, type_by_value: bool, type_len: c_int) -> Datum; +} + +/// CopyObject-safe data stored in `ForeignScan.fdw_private`. +/// +/// PostgreSQL copies cached plans after FDW planning callbacks return, so no +/// Rust pointer may be stored in the plan. This value is serialized entirely +/// as PostgreSQL Lists, Const nodes, and (for PARAM_EXEC) the original Expr +/// node, all of which PostgreSQL's `copyObject` understands. +struct FdwScanPrivate { + foreigntableid: Oid, + quals: Vec, + tgts: Vec, + sorts: Vec, + limit: Option, + opts: HashMap, + aggregates: Vec, + group_by: Vec, + all_base_quals_extracted: bool, + aggregate_base_columns: Vec, +} + +unsafe fn make_node_list(nodes: impl IntoIterator) -> *mut pg_sys::List { + unsafe { + nodes + .into_iter() + .fold(ptr::null_mut(), |list, node| pg_sys::lappend(list, node)) + } +} + +unsafe fn list_nodes(list: *mut pg_sys::List) -> Option> { + if list.is_null() { + return Some(Vec::new()); + } + pgrx::memcx::current_context(|mcx| unsafe { + List::<*mut c_void>::downcast_ptr_in_memcx(list, mcx) + .map(|list| list.iter().copied().collect()) + }) +} + +unsafe fn make_text_const(value: &str) -> *mut pg_sys::Const { + unsafe { + pg_sys::makeConst( + pg_sys::TEXTOID, + -1, + pg_sys::InvalidOid, + -1, + value.to_owned().into_datum().unwrap(), + false, + false, + ) + } +} + +unsafe fn text_from_node(node: *mut c_void) -> Option { + if node.is_null() + || !unsafe { pgrx::is_a(node.cast::(), pg_sys::NodeTag::T_Const) } + { + return None; + } + unsafe { + let constant = &*node.cast::(); + if constant.consttype != pg_sys::TEXTOID { + return None; + } + String::from_datum(constant.constvalue, constant.constisnull) + } +} + +unsafe fn bool_from_node(node: *mut c_void) -> Option { + match unsafe { text_from_node(node) }?.as_str() { + "0" => Some(false), + "1" => Some(true), + _ => None, + } +} + +fn cell_type_oid(cell: &Cell) -> Oid { + match cell { + Cell::Bool(_) => pg_sys::BOOLOID, + Cell::I8(_) => pg_sys::CHAROID, + Cell::I16(_) => pg_sys::INT2OID, + Cell::F32(_) => pg_sys::FLOAT4OID, + Cell::I32(_) => pg_sys::INT4OID, + Cell::F64(_) => pg_sys::FLOAT8OID, + Cell::I64(_) => pg_sys::INT8OID, + Cell::Numeric(_) => pg_sys::NUMERICOID, + Cell::String(_) => pg_sys::TEXTOID, + Cell::Date(_) => pg_sys::DATEOID, + Cell::Time(_) => pg_sys::TIMEOID, + Cell::Timestamp(_) => pg_sys::TIMESTAMPOID, + Cell::Timestamptz(_) => pg_sys::TIMESTAMPTZOID, + Cell::Interval(_) => pg_sys::INTERVALOID, + Cell::Json(_) => pg_sys::JSONBOID, + Cell::Bytea(_) => pg_sys::BYTEAOID, + Cell::Uuid(_) => pg_sys::UUIDOID, + Cell::BoolArray(_) => pg_sys::BOOLARRAYOID, + Cell::I16Array(_) => pg_sys::INT2ARRAYOID, + Cell::I32Array(_) => pg_sys::INT4ARRAYOID, + Cell::I64Array(_) => pg_sys::INT8ARRAYOID, + Cell::F32Array(_) => pg_sys::FLOAT4ARRAYOID, + Cell::F64Array(_) => pg_sys::FLOAT8ARRAYOID, + Cell::StringArray(_) => pg_sys::TEXTARRAYOID, + } +} + +unsafe fn make_cell_const(cell: &Cell) -> *mut pg_sys::Const { + unsafe { + let type_oid = cell_type_oid(cell); + let mut type_len = 0; + let mut type_by_value = false; + pg_sys::get_typlenbyval(type_oid, &mut type_len, &mut type_by_value); + let datum = datum_copy( + cell.clone().into_datum().unwrap(), + type_by_value, + type_len as c_int, + ); + pg_sys::makeConst( + type_oid, + -1, + pg_sys::InvalidOid, + type_len as _, + datum, + false, + type_by_value, + ) + } +} + +unsafe fn cell_from_node(node: *mut c_void) -> Option { + if node.is_null() + || !unsafe { pgrx::is_a(node.cast::(), pg_sys::NodeTag::T_Const) } + { + return None; + } + unsafe { + let constant = &*node.cast::(); + Cell::from_polymorphic_datum( + constant.constvalue, + constant.constisnull, + constant.consttype, + ) + } +} + +unsafe fn serialize_column(column: &Column) -> *mut pg_sys::List { + unsafe { + make_node_list([ + make_text_const(&column.name).cast(), + make_text_const(&column.num.to_string()).cast(), + make_text_const(&column.type_oid.to_u32().to_string()).cast(), + ]) + } +} + +unsafe fn deserialize_column(list: *mut pg_sys::List) -> Option { + unsafe { + let nodes = list_nodes(list)?; + if nodes.len() != 3 { + return None; + } + Some(Column { + name: text_from_node(*nodes.first()?)?, + num: text_from_node(*nodes.get(1)?)?.parse().ok()?, + type_oid: Oid::from(text_from_node(*nodes.get(2)?)?.parse::().ok()?), + }) + } +} + +unsafe fn serialize_qual( + qual: &Qual, + param_exprs: &mut Vec<*mut pg_sys::Expr>, +) -> *mut pg_sys::List { + unsafe { + let (value_is_array, cells) = match &qual.value { + Value::Cell(cell) => (false, vec![cell]), + Value::Array(cells) => (true, cells.iter().collect()), + }; + let value_nodes = + make_node_list(cells.into_iter().map(|cell| make_cell_const(cell).cast())); + + let (has_param, param_kind, param_id, param_type_oid, param_expr_index) = + if let Some(param) = &qual.param { + let kind = match param.kind { + ParamKind::PARAM_EXTERN => "extern", + ParamKind::PARAM_EXEC => "exec", + _ => "unsupported", + }; + let expr_index = + if param.kind == ParamKind::PARAM_EXEC && !param.expr_eval.expr.is_null() { + param_exprs.push(param.expr_eval.expr); + (param_exprs.len() - 1).to_string() + } else { + "-1".to_string() + }; + ( + true, + kind, + param.id.to_string(), + param.type_oid.to_u32().to_string(), + expr_index, + ) + } else { + ( + false, + "none", + "0".to_string(), + "0".to_string(), + "-1".to_string(), + ) + }; + + make_node_list([ + make_text_const(&qual.field).cast(), + make_text_const(&qual.operator).cast(), + make_text_const(if qual.use_or { "1" } else { "0" }).cast(), + make_text_const(if value_is_array { "1" } else { "0" }).cast(), + value_nodes.cast(), + make_text_const(if has_param { "1" } else { "0" }).cast(), + make_text_const(param_kind).cast(), + make_text_const(¶m_id).cast(), + make_text_const(¶m_type_oid).cast(), + make_text_const(¶m_expr_index).cast(), + ]) + } +} + +unsafe fn deserialize_qual( + list: *mut pg_sys::List, + param_exprs: &[*mut pg_sys::Expr], +) -> Option { + unsafe { + let nodes = list_nodes(list)?; + if nodes.len() != 10 { + return None; + } + let value_is_array = bool_from_node(*nodes.get(3)?)?; + let cells = list_nodes(*nodes.get(4)? as *mut pg_sys::List)? + .into_iter() + .map(|node| cell_from_node(node)) + .collect::>>()?; + let value = if value_is_array { + Value::Array(cells) + } else { + if cells.len() != 1 { + return None; + } + Value::Cell(cells.into_iter().next()?) + }; + + let has_param = bool_from_node(*nodes.get(5)?)?; + let param = if has_param { + let kind = match text_from_node(*nodes.get(6)?)?.as_str() { + "extern" => ParamKind::PARAM_EXTERN, + "exec" => ParamKind::PARAM_EXEC, + _ => return None, + }; + let expr = if kind == ParamKind::PARAM_EXEC { + let index = text_from_node(*nodes.get(9)?)?.parse::().ok()?; + let expr = *param_exprs.get(index)?; + if expr.is_null() + || !pgrx::is_a(expr.cast(), pg_sys::NodeTag::T_Param) + || (*expr.cast::()).paramkind != ParamKind::PARAM_EXEC + { + return None; + } + expr + } else { + if text_from_node(*nodes.get(9)?)? != "-1" { + return None; + } + ptr::null_mut() + }; + Some(Param { + kind, + id: text_from_node(*nodes.get(7)?)?.parse().ok()?, + type_oid: Oid::from(text_from_node(*nodes.get(8)?)?.parse::().ok()?), + eval_value: std::sync::Mutex::new(None).into(), + eval_state: std::sync::Mutex::new(ParamValue::Unevaluated).into(), + expr_eval: ExprEval { + expr, + expr_state: ptr::null_mut(), + }, + }) + } else { + if text_from_node(*nodes.get(6)?)? != "none" + || text_from_node(*nodes.get(7)?)? != "0" + || text_from_node(*nodes.get(8)?)? != "0" + || text_from_node(*nodes.get(9)?)? != "-1" + { + return None; + } + None + }; + + Some(Qual { + field: text_from_node(*nodes.first()?)?, + operator: text_from_node(*nodes.get(1)?)?, + use_or: bool_from_node(*nodes.get(2)?)?, + value, + param, + }) + } +} + +unsafe fn serialize_sort(sort: &Sort) -> *mut pg_sys::List { + unsafe { + make_node_list([ + make_text_const(&sort.field).cast(), + make_text_const(&sort.field_no.to_string()).cast(), + make_text_const(if sort.reversed { "1" } else { "0" }).cast(), + make_text_const(if sort.nulls_first { "1" } else { "0" }).cast(), + make_text_const(if sort.collate.is_some() { "1" } else { "0" }).cast(), + make_text_const(sort.collate.as_deref().unwrap_or("")).cast(), + ]) + } +} + +unsafe fn deserialize_sort(list: *mut pg_sys::List) -> Option { + unsafe { + let nodes = list_nodes(list)?; + if nodes.len() != 6 { + return None; + } + let has_collate = bool_from_node(*nodes.get(4)?)?; + let collate = text_from_node(*nodes.get(5)?)?; + if !has_collate && !collate.is_empty() { + return None; + } + Some(Sort { + field: text_from_node(*nodes.first()?)?, + field_no: text_from_node(*nodes.get(1)?)?.parse().ok()?, + reversed: bool_from_node(*nodes.get(2)?)?, + nulls_first: bool_from_node(*nodes.get(3)?)?, + collate: has_collate.then_some(collate), + }) + } +} + +fn aggregate_kind_tag(kind: AggregateKind) -> &'static str { + match kind { + AggregateKind::Count => "count", + AggregateKind::CountColumn => "count-column", + AggregateKind::Sum => "sum", + AggregateKind::Avg => "avg", + AggregateKind::Min => "min", + AggregateKind::Max => "max", + } +} + +fn aggregate_kind_from_tag(tag: &str) -> Option { + match tag { + "count" => Some(AggregateKind::Count), + "count-column" => Some(AggregateKind::CountColumn), + "sum" => Some(AggregateKind::Sum), + "avg" => Some(AggregateKind::Avg), + "min" => Some(AggregateKind::Min), + "max" => Some(AggregateKind::Max), + _ => None, + } +} + +unsafe fn serialize_aggregate(aggregate: &Aggregate) -> *mut pg_sys::List { + unsafe { + let (has_column, name, num, type_oid) = aggregate + .column + .as_ref() + .map(|column| { + ( + true, + column.name.as_str(), + column.num, + column.type_oid.to_u32(), + ) + }) + .unwrap_or((false, "", 0, 0)); + make_node_list([ + make_text_const(aggregate_kind_tag(aggregate.kind)).cast(), + make_text_const(if has_column { "1" } else { "0" }).cast(), + make_text_const(name).cast(), + make_text_const(&num.to_string()).cast(), + make_text_const(&type_oid.to_string()).cast(), + make_text_const(if aggregate.distinct { "1" } else { "0" }).cast(), + make_text_const(&aggregate.alias).cast(), + make_text_const(&aggregate.type_oid.to_u32().to_string()).cast(), + ]) + } +} + +unsafe fn deserialize_aggregate(list: *mut pg_sys::List) -> Option { + unsafe { + let nodes = list_nodes(list)?; + if nodes.len() != 8 { + return None; + } + let column = if bool_from_node(*nodes.get(1)?)? { + Some(Column { + name: text_from_node(*nodes.get(2)?)?, + num: text_from_node(*nodes.get(3)?)?.parse().ok()?, + type_oid: Oid::from(text_from_node(*nodes.get(4)?)?.parse::().ok()?), + }) + } else { + if !text_from_node(*nodes.get(2)?)?.is_empty() + || text_from_node(*nodes.get(3)?)? != "0" + || text_from_node(*nodes.get(4)?)? != "0" + { + return None; + } + None + }; + Some(Aggregate { + kind: aggregate_kind_from_tag(&text_from_node(*nodes.first()?)?)?, + column, + distinct: bool_from_node(*nodes.get(5)?)?, + alias: text_from_node(*nodes.get(6)?)?, + type_oid: Oid::from(text_from_node(*nodes.get(7)?)?.parse::().ok()?), + }) + } +} + +impl FdwScanPrivate { + fn from_state, W: ForeignDataWrapper>(state: &FdwState) -> Self { + Self { + foreigntableid: state.foreigntableid, + quals: state.quals.clone(), + tgts: state.tgts.clone(), + sorts: state.sorts.clone(), + limit: state.limit.clone(), + opts: state.opts.clone(), + aggregates: state.aggregates.clone(), + group_by: state.group_by.clone(), + all_base_quals_extracted: state.all_base_quals_extracted, + aggregate_base_columns: state.aggregate_base_columns.clone(), + } + } + + /// Serialize to nested PostgreSQL Lists containing only copyObject-safe + /// nodes. The positional schema is versioned by element zero. + unsafe fn serialize_to_list(&self) -> (*mut pg_sys::List, *mut pg_sys::List) { + unsafe { + let mut param_exprs = Vec::new(); + let quals = make_node_list( + self.quals + .iter() + .map(|qual| serialize_qual(qual, &mut param_exprs).cast()), + ); + let tgts = make_node_list( + self.tgts + .iter() + .map(|column| serialize_column(column).cast()), + ); + let sorts = make_node_list(self.sorts.iter().map(|sort| serialize_sort(sort).cast())); + let limit = make_node_list([ + make_text_const(if self.limit.is_some() { "1" } else { "0" }).cast(), + make_text_const( + &self + .limit + .as_ref() + .map(|limit| limit.count) + .unwrap_or_default() + .to_string(), + ) + .cast(), + make_text_const( + &self + .limit + .as_ref() + .map(|limit| limit.offset) + .unwrap_or_default() + .to_string(), + ) + .cast(), + ]); + let mut options = self.opts.iter().collect::>(); + options.sort_unstable_by(|left, right| left.0.cmp(right.0)); + let opts = make_node_list(options.into_iter().map(|(key, value)| { + make_node_list([make_text_const(key).cast(), make_text_const(value).cast()]).cast() + })); + let aggregates = make_node_list( + self.aggregates + .iter() + .map(|aggregate| serialize_aggregate(aggregate).cast()), + ); + let group_by = make_node_list( + self.group_by + .iter() + .map(|column| serialize_column(column).cast()), + ); + let aggregate_base_columns = make_node_list( + self.aggregate_base_columns + .iter() + .map(|column| serialize_column(column).cast()), + ); + + let private = make_node_list([ + make_text_const(FDW_SCAN_PRIVATE_VERSION).cast(), + make_text_const(&self.foreigntableid.to_u32().to_string()).cast(), + quals.cast(), + tgts.cast(), + sorts.cast(), + limit.cast(), + opts.cast(), + aggregates.cast(), + group_by.cast(), + make_text_const(if self.all_base_quals_extracted { + "1" + } else { + "0" + }) + .cast(), + aggregate_base_columns.cast(), + ]); + let fdw_exprs = make_node_list(param_exprs.into_iter().map(|expr| expr.cast())); + (private, fdw_exprs) + } + } + + unsafe fn deserialize_from_list( + list: *mut pg_sys::List, + fdw_exprs: *mut pg_sys::List, + ) -> Option { + unsafe { + let nodes = list_nodes(list)?; + if nodes.len() != 11 || text_from_node(*nodes.first()?)? != FDW_SCAN_PRIVATE_VERSION { + return None; + } + let param_exprs = list_nodes(fdw_exprs)? + .into_iter() + .map(|node| { + if node.is_null() || !pgrx::is_a(node.cast(), pg_sys::NodeTag::T_Param) { + None + } else { + Some(node.cast::()) + } + }) + .collect::>>()?; + let quals = list_nodes(*nodes.get(2)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_qual(node.cast(), ¶m_exprs)) + .collect::>>()?; + if quals + .iter() + .filter(|qual| { + qual.param + .as_ref() + .is_some_and(|param| param.kind == ParamKind::PARAM_EXEC) + }) + .count() + != param_exprs.len() + { + return None; + } + let tgts = list_nodes(*nodes.get(3)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_column(node.cast())) + .collect::>>()?; + let sorts = list_nodes(*nodes.get(4)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_sort(node.cast())) + .collect::>>()?; + let limit_nodes = list_nodes(*nodes.get(5)? as *mut pg_sys::List)?; + if limit_nodes.len() != 3 { + return None; + } + let limit = if bool_from_node(*limit_nodes.first()?)? { + Some(Limit { + count: text_from_node(*limit_nodes.get(1)?)?.parse().ok()?, + offset: text_from_node(*limit_nodes.get(2)?)?.parse().ok()?, + }) + } else { + None + }; + let opts = list_nodes(*nodes.get(6)? as *mut pg_sys::List)? + .into_iter() + .map(|node| { + let pair = list_nodes(node.cast())?; + if pair.len() != 2 { + return None; + } + Some(( + text_from_node(*pair.first()?)?, + text_from_node(*pair.get(1)?)?, + )) + }) + .collect::>>()?; + let aggregates = list_nodes(*nodes.get(7)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_aggregate(node.cast())) + .collect::>>()?; + let group_by = list_nodes(*nodes.get(8)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_column(node.cast())) + .collect::>>()?; + let aggregate_base_columns = list_nodes(*nodes.get(10)? as *mut pg_sys::List)? + .into_iter() + .map(|node| deserialize_column(node.cast())) + .collect::>>()?; + + Some(Self { + foreigntableid: Oid::from(text_from_node(*nodes.get(1)?)?.parse::().ok()?), + quals, + tgts, + sorts, + limit, + opts, + aggregates, + group_by, + all_base_quals_extracted: bool_from_node(*nodes.get(9)?)?, + aggregate_base_columns, + }) + } + } +} // Fdw private state for scan pub(crate) struct FdwState, W: ForeignDataWrapper> { + // foreign table used to construct a fresh FDW for each execution + foreigntableid: Oid, + // foreign data wrapper instance pub(crate) instance: Option, @@ -46,6 +672,8 @@ pub(crate) struct FdwState, W: ForeignDataWrapper> { // aggregate pushdown pub(crate) aggregates: Vec, pub(crate) group_by: Vec, + pub(crate) all_base_quals_extracted: bool, + pub(crate) aggregate_base_columns: Vec, // temporary memory context per foreign table, created under Wrappers root // memory context @@ -57,13 +685,23 @@ pub(crate) struct FdwState, W: ForeignDataWrapper> { row: Row, // fingerprint of current parameter values to detect rescan changes param_fingerprint: String, + // whether begin_scan/begin_aggregate_scan ran for this execution + scan_started: bool, _phantom: PhantomData, } impl, W: ForeignDataWrapper> FdwState { unsafe fn new(foreigntableid: Oid, tmp_ctx: MemoryContext) -> Self { + let mut state = Self::new_without_instance(foreigntableid, tmp_ctx); + state.instance = + Some(unsafe { instance::create_fdw_instance_from_table_id(foreigntableid) }); + state + } + + fn new_without_instance(foreigntableid: Oid, tmp_ctx: MemoryContext) -> Self { Self { - instance: Some(unsafe { instance::create_fdw_instance_from_table_id(foreigntableid) }), + foreigntableid, + instance: None, quals: Vec::new(), tgts: Vec::new(), sorts: Vec::new(), @@ -71,11 +709,14 @@ impl, W: ForeignDataWrapper> FdwState { opts: HashMap::new(), aggregates: Vec::new(), group_by: Vec::new(), + all_base_quals_extracted: true, + aggregate_base_columns: Vec::new(), tmp_ctx, values: Vec::new(), nulls: Vec::new(), row: Row::new(), param_fingerprint: String::new(), + scan_started: false, _phantom: PhantomData, } } @@ -100,10 +741,50 @@ impl, W: ForeignDataWrapper> FdwState { !self.aggregates.is_empty() } + #[inline] + pub(crate) fn can_pushdown_aggregate( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + base_columns: &[Column], + ) -> Result { + if let Some(ref mut instance) = self.instance { + instance.can_pushdown_aggregate( + aggregates, + group_by, + &self.quals, + base_columns, + self.all_base_quals_extracted, + &self.opts, + ) + } else { + Ok(false) + } + } + + #[inline] + pub(crate) fn get_aggregate_rel_size( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + ) -> Result<(i64, i32), E> { + if let Some(ref mut instance) = self.instance { + instance.get_aggregate_rel_size(aggregates, group_by, &self.quals, &self.opts) + } else { + Ok((0, 0)) + } + } + #[inline] fn begin_aggregate_scan(&mut self) -> Result<(), E> { if let Some(ref mut instance) = self.instance { - instance.begin_aggregate_scan(&self.aggregates, &self.group_by, &self.quals, &self.opts) + instance.begin_aggregate_scan_with_base_columns( + &self.aggregates, + &self.group_by, + &self.quals, + &self.aggregate_base_columns, + &self.opts, + ) } else { Ok(()) } @@ -152,8 +833,6 @@ impl, W: ForeignDataWrapper> FdwState { } } -impl, W: ForeignDataWrapper> utils::SerdeList for FdwState {} - impl, W: ForeignDataWrapper> Drop for FdwState { fn drop(&mut self) { // drop foreign data wrapper instance @@ -167,14 +846,6 @@ impl, W: ForeignDataWrapper> Drop for FdwState { } } -// drop the scan state, so the inner fdw instance can be dropped too -unsafe fn drop_fdw_state, W: ForeignDataWrapper>( - fdw_state: *mut FdwState, -) { - let boxed_fdw_state = unsafe { Box::from_raw(fdw_state) }; - drop(boxed_fdw_state); -} - #[pg_guard] pub(super) extern "C-unwind" fn get_foreign_rel_size< E: Into, @@ -195,7 +866,9 @@ pub(super) extern "C-unwind" fn get_foreign_rel_size< PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| { // extract qual list - state.quals = extract_quals(root, baserel, foreigntableid); + let extracted_quals = extract_quals(root, baserel, foreigntableid); + state.quals = extracted_quals.quals; + state.all_base_quals_extracted = extracted_quals.all_extracted; // extract target column list from target and restriction expression state.tgts = utils::extract_target_columns(root, baserel); @@ -226,8 +899,11 @@ pub(super) extern "C-unwind" fn get_foreign_rel_size< (*baserel).rows = rows as f64; (*(*baserel).reltarget).width = width; - // save the state for following callbacks - (*baserel).fdw_private = Box::leak(Box::new(state)) as *mut FdwState as _; + // This is planner-only state. Register it on PlannerInfo.planner_cxt + // explicitly so it survives every FDW planning callback but never + // escapes into the copyObject'd CachedPlan. + (*baserel).fdw_private = + PgMemoryContexts::For((*root).planner_cxt).leak_and_drop_on_delete(state) as _; } } @@ -313,6 +989,7 @@ pub(super) extern "C-unwind" fn get_foreign_plan, W: Foreig // begin_scan and the tuple slot is the base-rel's row type. state.aggregates = Vec::new(); state.group_by = Vec::new(); + state.aggregate_base_columns = Vec::new(); } let (final_tlist, agg_fdw_scan_tlist) = if is_agg { @@ -384,16 +1061,14 @@ pub(super) extern "C-unwind" fn get_foreign_plan, W: Foreig (tlist, ptr::null_mut()) }; - // 'serialize' state to list, basically what we're doing here is to store - // the state pointer as an integer constant in the list, so it can be - // `deserialized` when executing the plan later. - // Note that the state itself is not serialized to any memory contexts, - // it just sits in Rust managed Box'ed memory and will be dropped when - // end_foreign_scan() is called. - let fdw_private = - PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| FdwState::serialize_to_list(state)); + // Serialize only PostgreSQL-native nodes into the plan. Cached plans + // are copyObject'd after this callback, so storing the Rust state + // pointer here would leave the copied plan with a dangling pointer. + let private = FdwScanPrivate::from_state(&state); + state.instance.take(); + let (fdw_private, fdw_exprs) = private.serialize_to_list(); - pg_sys::make_foreignscan( + let foreign_scan = pg_sys::make_foreignscan( final_tlist, scan_clauses, (*baserel).relid, @@ -402,7 +1077,11 @@ pub(super) extern "C-unwind" fn get_foreign_plan, W: Foreig agg_fdw_scan_tlist, ptr::null_mut(), outer_plan, - ) + ); + // PARAM_EXEC expressions live in the dedicated ForeignScan expression + // list so core planner fixups and copyObject process them normally. + (*foreign_scan).fdw_exprs = fdw_exprs; + foreign_scan } } @@ -446,6 +1125,53 @@ pub(super) extern "C-unwind" fn explain_foreign_scan< let value = ctx.pstrdup(&format!("group_by = {:?}", state.group_by)); pg_sys::ExplainPropertyText(label, value, es); } + + if let Some(instance) = state.instance.as_ref() { + for property in instance.explain() { + emit_explain_property(&ctx, property, es); + } + } + } +} + +unsafe fn emit_explain_property( + ctx: &PgMemoryContexts, + property: ExplainProperty, + es: *mut pg_sys::ExplainState, +) { + unsafe { + let label = ctx.pstrdup(&property.label); + match property.value { + ExplainValue::Text(value) => { + let value = ctx.pstrdup(&value); + pg_sys::ExplainPropertyText(label, value, es); + } + ExplainValue::Integer { value, unit } => { + let unit = unit + .as_deref() + .map_or(ptr::null_mut(), |unit| ctx.pstrdup(unit)); + pg_sys::ExplainPropertyInteger(label, unit, value, es); + } + ExplainValue::Unsigned { value, unit } => { + let unit = unit + .as_deref() + .map_or(ptr::null_mut(), |unit| ctx.pstrdup(unit)); + pg_sys::ExplainPropertyUInteger(label, unit, value, es); + } + ExplainValue::Float { + value, + unit, + digits, + } => { + let unit = unit + .as_deref() + .map_or(ptr::null_mut(), |unit| ctx.pstrdup(unit)); + pg_sys::ExplainPropertyFloat(label, unit, value, digits, es); + } + ExplainValue::Boolean(value) => { + pg_sys::ExplainPropertyBool(label, value, es); + } + } } } @@ -461,8 +1187,7 @@ unsafe fn assign_parameter_value, W: ForeignDataWrapper> // assign parameter value to qual for qual in &mut state.quals.iter_mut() { if let Some(param) = &mut qual.param { - let mut current_value: Option = None; - match param.kind { + let current_value = match param.kind { ParamKind::PARAM_EXTERN => { // get parameter list in execution state let plist_info = (*estate).es_param_list_info; @@ -471,13 +1196,21 @@ unsafe fn assign_parameter_value, W: ForeignDataWrapper> if param.id > 0 && param.id <= params_cnt { let plist = (*plist_info).params.as_slice(params_cnt); let p: pg_sys::ParamExternData = plist[param.id - 1]; - if let Some(cell) = + if p.isnull { + ParamValue::Null + } else if let Some(cell) = Cell::from_polymorphic_datum(p.value, p.isnull, p.ptype) { qual.value = Value::Cell(cell.clone()); - current_value = Some(Value::Cell(cell)); + ParamValue::Value(Value::Cell(cell)) + } else { + ParamValue::Unevaluated } + } else { + ParamValue::Unevaluated } + } else { + ParamValue::Unevaluated } } ParamKind::PARAM_EXEC => { @@ -487,54 +1220,57 @@ unsafe fn assign_parameter_value, W: ForeignDataWrapper> node as *mut pg_sys::PlanState, ); let mut isnull = false; - if let Some(datum) = polyfill::exec_eval_expr( + match polyfill::exec_eval_expr( param.expr_eval.expr_state, econtext, &mut isnull, - ) && let Some(cell) = - Cell::from_polymorphic_datum(datum, isnull, param.type_oid) - { - qual.value = Value::Cell(cell.clone()); - current_value = Some(Value::Cell(cell)); + ) { + Some(_) if isnull => ParamValue::Null, + Some(datum) => { + if let Some(cell) = + Cell::from_polymorphic_datum(datum, false, param.type_oid) + { + qual.value = Value::Cell(cell.clone()); + ParamValue::Value(Value::Cell(cell)) + } else { + ParamValue::Unevaluated + } + } + None => ParamValue::Unevaluated, } } - _ => {} - } + _ => ParamValue::Unevaluated, + }; - let mut eval_value = param - .eval_value - .lock() - .expect("param.eval_value should be locked"); - *eval_value = current_value; + param.set_evaluated_value(current_value); } } } } +fn parameter_fingerprint(qual: &Qual) -> Option { + qual.param.as_ref().map(|param| { + let eval_value = format!("{:?}", param.evaluated_value()); + format!( + "{}|{}|{}|{}|{}|{}|{}", + qual.field, + qual.operator, + qual.use_or, + param.kind, + param.id, + param.type_oid, + eval_value, + ) + }) +} + fn compute_param_fingerprint, W: ForeignDataWrapper>( state: &FdwState, ) -> String { state .quals .iter() - .filter_map(|qual| { - qual.param.as_ref().map(|param| { - let eval_value = match param.eval_value.lock() { - Ok(value) => format!("{:?}", *value), - Err(_) => "lock_error".to_string(), - }; - format!( - "{}|{}|{}|{}|{}|{}|{}", - qual.field, - qual.operator, - qual.use_or, - param.kind, - param.id, - param.type_oid, - eval_value, - ) - }) - }) + .filter_map(parameter_fingerprint) .collect::>() .join(";") } @@ -551,15 +1287,57 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< unsafe { let scan_state = (*node).ss; let plan = scan_state.ps.plan as *mut pg_sys::ForeignScan; - let mut state = FdwState::::deserialize_from_list((*plan).fdw_private as _); - assert!(!state.is_null()); + let Some(private) = + FdwScanPrivate::deserialize_from_list((*plan).fdw_private, (*plan).fdw_exprs) + else { + report_error( + PgSqlErrorCode::ERRCODE_FDW_ERROR, + "invalid fdw_private data in begin_foreign_scan", + ); + return; + }; - // assign parameter values to qual - assign_parameter_value(node, &mut state); - state.param_fingerprint = compute_param_fingerprint(&state); + // Every execution, including each use of a cached/generic plan, owns a + // fresh Rust state, FDW instance, parameter cells, and temporary + // context. The plan contains no mutable Rust allocation. + let explain_only = eflags & pg_sys::EXEC_FLAG_EXPLAIN_ONLY as c_int > 0; + let ctx_name = format!("Wrappers_scan_exec_{}", private.foreigntableid.to_u32()); + let ctx = memctx::create_wrappers_memctx(&ctx_name); + // Start without a client so the query context can become the owner + // before W::new performs any fallible work. EXPLAIN without ANALYZE + // keeps this metadata-only state for ExplainForeignScan. + let mut state = FdwState::::new_without_instance(private.foreigntableid, ctx); + PgMemoryContexts::For(ctx).switch_to(|_| { + state.quals = private.quals; + state.tgts = private.tgts; + state.sorts = private.sorts; + state.limit = private.limit; + state.opts = private.opts; + state.aggregates = private.aggregates; + state.group_by = private.group_by; + state.all_base_quals_extracted = private.all_base_quals_extracted; + state.aggregate_base_columns = private.aggregate_base_columns; + }); + + // The executor query context is the sole owner. Its reset callback + // drops the state and its Wrappers temp context even when PostgreSQL + // exits through ERROR before EndForeignScan runs. + let estate = scan_state.ps.state; + let state_ptr = + PgMemoryContexts::For((*estate).es_query_cxt).leak_and_drop_on_delete(state); + (*node).fdw_state = state_ptr.cast(); + let mut state = PgBox::>::from_pg(state_ptr); // begin scan if it is not EXPLAIN statement - if eflags & pg_sys::EXEC_FLAG_EXPLAIN_ONLY as c_int <= 0 { + if !explain_only { + state.instance = Some(instance::create_fdw_instance_from_table_id( + state.foreigntableid, + )); + + // assign parameter values to qual + assign_parameter_value(node, &mut state); + state.param_fingerprint = compute_param_fingerprint(&state); + // choose aggregate scan or normal scan based on state let result = if state.is_aggregate_scan() { state.begin_aggregate_scan() @@ -567,10 +1345,10 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< state.begin_scan() }; if result.is_err() { - drop_fdw_state(state.as_ptr()); - (*plan).fdw_private = ptr::null::>() as _; result.report_unwrap(); + return; } + state.scan_started = true; // For aggregate upper-rel scans, scanrelid=0 so ss_currentRelation is // NULL. Use the number of output columns from state.tgts instead. @@ -587,8 +1365,6 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< .extend_from_slice(&vec![0.into_datum().unwrap(); natts]); state.nulls.extend_from_slice(&vec![true; natts]); } - - (*node).fdw_state = state.into_pg() as _; } } @@ -614,10 +1390,6 @@ pub(super) extern "C-unwind" fn iterate_foreign_scan< state.row.clear(); let result = state.iter_scan(); - if result.is_err() { - drop_fdw_state(state.as_ptr()); - (*node).fdw_state = ptr::null::>() as _; - } if result.report_unwrap().is_some() { if state.row.cols.len() != state.tgts.len() { report_error( @@ -688,13 +1460,15 @@ pub(super) extern "C-unwind" fn re_scan_foreign_scan< state.param_fingerprint = next_fingerprint; // end the active scan to release resources before restarting with new params let _ = state.end_scan(); - state.begin_scan() + if state.is_aggregate_scan() { + state.begin_aggregate_scan() + } else { + state.begin_scan() + } } else { state.re_scan() }; if result.is_err() { - drop_fdw_state(state.as_ptr()); - (*node).fdw_state = ptr::null::>() as _; result.report_unwrap(); } } @@ -712,14 +1486,68 @@ pub(super) extern "C-unwind" fn end_foreign_scan, W: Foreig return; } - // the scan state is actually not allocated by PG, but we use 'from_pg()' - // here just to tell PgBox don't free the state, instead we will handle - // drop the state by ourselves let mut state = PgBox::>::from_pg(fdw_state); - let result = state.end_scan(); - drop_fdw_state(state.as_ptr()); + let result = if state.scan_started { + let result = state.end_scan(); + state.scan_started = false; + result + } else { + Ok(()) + }; + // The es_query_cxt callback remains the sole owner and will drop this + // state exactly once after executor teardown. (*node).fdw_state = ptr::null::>() as _; result.report_unwrap(); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::interface::{ExprEval, Param}; + use std::sync::Mutex; + + fn parameterized_qual() -> Qual { + Qual { + field: "value".to_string(), + operator: "=".to_string(), + value: Value::Cell(Cell::I64(0)), + use_or: false, + param: Some(Param { + kind: ParamKind::PARAM_EXTERN, + id: 1, + type_oid: pg_sys::INT8OID, + eval_value: Mutex::new(None).into(), + eval_state: Mutex::new(ParamValue::Unevaluated).into(), + expr_eval: ExprEval { + expr: ptr::null_mut(), + expr_state: ptr::null_mut(), + }, + }), + } + } + + #[test] + fn test_parameter_fingerprint_distinguishes_null_transitions() { + let qual = parameterized_qual(); + let unevaluated = parameter_fingerprint(&qual).expect("parameter fingerprint"); + + qual.param + .as_ref() + .expect("parameter") + .set_evaluated_value(ParamValue::Null); + let null = parameter_fingerprint(&qual).expect("parameter fingerprint"); + + qual.param + .as_ref() + .expect("parameter") + .set_evaluated_value(ParamValue::Value(Value::Cell(Cell::I64(7)))); + let value = parameter_fingerprint(&qual).expect("parameter fingerprint"); + + assert_ne!(unevaluated, null); + assert_ne!(null, value); + assert_ne!(unevaluated, value); + assert!(null.contains("Null")); + } +} diff --git a/supabase-wrappers/src/upper.rs b/supabase-wrappers/src/upper.rs index 35856f7d9..21d059148 100644 --- a/supabase-wrappers/src/upper.rs +++ b/supabase-wrappers/src/upper.rs @@ -9,7 +9,9 @@ use std::ptr; use crate::interface::{Aggregate, AggregateKind, Column}; use crate::prelude::ForeignDataWrapper; +use crate::qual::is_builtin_pg_catalog_object; use crate::scan::FdwState; +use crate::utils::ReportableError; /// Helper to iterate over a pg_sys::List using raw pointer access. /// Returns an iterator over pointers to the list elements. @@ -34,6 +36,10 @@ unsafe fn list_iter(list: *mut pg_sys::List) -> impl Iterator /// by looking up the function name. fn oid_to_aggregate_kind(aggfnoid: pg_sys::Oid) -> Option { unsafe { + if !is_builtin_pg_catalog_object(aggfnoid, pg_sys::get_func_namespace(aggfnoid)) { + return None; + } + let agg_name = pg_sys::get_func_name(aggfnoid); if agg_name.is_null() { return None; @@ -148,6 +154,13 @@ unsafe fn extract_aggregates( return None; } + // Ordered aggregate transition state is not represented by + // Aggregate, so an FDW could not reproduce its semantics. + if !(*aggref).aggorder.is_null() { + debug2!("Aggregate has ORDER BY clause, skipping pushdown"); + return None; + } + // DISTINCT only supported for COUNT(column) if !(*aggref).aggdistinct.is_null() { match kind { @@ -196,6 +209,12 @@ unsafe fn extract_aggregates( alias: format!("agg_{}", resno + 1), type_oid: (*aggref).aggtype, }); + } else if (*expr).type_ != pg_sys::NodeTag::T_Var { + // A Var may be a plain GROUP BY output. Any other expression + // is not represented in the foreign aggregate result contract + // and must remain a PostgreSQL projection/aggregate plan. + debug2!("Non-aggregate upper target expression, skipping pushdown"); + return None; } } @@ -292,6 +311,12 @@ pub(super) extern "C-unwind" fn get_foreign_upper_paths< } unsafe { + let parse = (*root).parse; + if !parse.is_null() && !(*parse).groupingSets.is_null() { + debug2!("Grouping sets are not supported, skipping aggregate pushdown"); + return; + } + // Get the FDW state from the input relation (set during get_foreign_rel_size) let fdw_private = (*input_rel).fdw_private; if fdw_private.is_null() { @@ -349,6 +374,26 @@ pub(super) extern "C-unwind" fn get_foreign_upper_paths< } } + // An aggregate FDW consumes base rows before PostgreSQL can recheck + // restrictions. Never register the upper path if any base restriction + // was omitted from the Qual representation. + if !state.all_base_quals_extracted { + debug2!("Not all base restrictions were extracted, skipping aggregate pushdown"); + return; + } + + // Preserve the original base columns before get_foreign_plan replaces + // state.tgts with the aggregate output aliases. + let base_columns = state.tgts.clone(); + + let can_pushdown = state + .can_pushdown_aggregate(&aggregates, &group_by, &base_columns) + .report_unwrap(); + if !can_pushdown { + debug2!("FDW rejected this aggregate query shape, skipping pushdown"); + return; + } + // Store aggregates and group_by in the FdwState so they survive to // execution. Note: input_rel.fdw_private and output_rel.fdw_private // share the same FdwState pointer, so this mutation is visible through @@ -357,19 +402,23 @@ pub(super) extern "C-unwind" fn get_foreign_upper_paths< // to the base-rel scan with a local Aggregate on top. state.aggregates = aggregates.clone(); state.group_by = group_by.clone(); + state.aggregate_base_columns = base_columns; + + let (rows, width) = state + .get_aggregate_rel_size(&aggregates, &group_by) + .report_unwrap(); + (*output_rel).rows = rows as f64; + (*(*output_rel).reltarget).width = width; - // Cost estimation. We deliberately price the pushdown at ~0 so the - // planner prefers it over the local HashAgg/GroupAgg alternatives that - // also live on grouped_rel — pushdown collapses the row stream at the - // remote side and is essentially always cheaper than fetching rows and - // aggregating locally. - let rows: i64 = 1; + // Prefer an eligible pushdown path over fetching base rows for a local + // HashAgg/GroupAgg. The FDW callback above supplies the result row count + // and width; startup_cost remains the configurable remote fixed cost. let startup_cost = state .opts .get("startup_cost") .and_then(|c| c.parse::().ok()) .unwrap_or(0.0); - let total_cost = startup_cost; + let total_cost = startup_cost + rows.max(0) as f64; debug2!( "Aggregate pushdown cost estimate: rows={rows}, startup={startup_cost}, total={total_cost}" diff --git a/supabase-wrappers/src/utils.rs b/supabase-wrappers/src/utils.rs index ae69cee7c..215f5ed0f 100644 --- a/supabase-wrappers/src/utils.rs +++ b/supabase-wrappers/src/utils.rs @@ -3,7 +3,6 @@ use crate::interface::{Cell, Column, Row}; use pgrx::{ - IntoDatum, list::List, pg_sys::panic::{ErrorReport, ErrorReportable}, spi::Spi, @@ -515,51 +514,6 @@ pub(super) unsafe fn extract_target_columns( } } -// trait for "serialize" and "deserialize" state from specified memory context, -// so that it is safe to be carried between the planning and the execution -pub(super) trait SerdeList { - unsafe fn serialize_to_list(state: PgBox) -> *mut pg_sys::List - where - Self: Sized, - { - unsafe { - memcx::current_context(|mcx| { - let mut ret = List::<*mut c_void>::Nil; - let val = state.into_pg() as i64; - let cst: *mut pg_sys::Const = pg_sys::makeConst( - pg_sys::INT8OID, - -1, - pg_sys::InvalidOid, - 8, - val.into_datum().unwrap(), - false, - true, - ); - ret.unstable_push_in_context(cst as _, mcx); - ret.into_ptr() - }) - } - } - - unsafe fn deserialize_from_list(list: *mut pg_sys::List) -> PgBox - where - Self: Sized, - { - unsafe { - memcx::current_context(|mcx| { - if let Some(list) = List::<*mut c_void>::downcast_ptr_in_memcx(list, mcx) - && let Some(cst) = list.get(0) - { - let cst = *(*cst as *mut pg_sys::Const); - let ptr = i64::from_datum(cst.constvalue, cst.constisnull).unwrap(); - return PgBox::::from_pg(ptr as _); - } - PgBox::::null() - }) - } - } -} - pub(crate) trait ReportableError { type Output; diff --git a/wrappers/.ci/docker-compose-native.yaml b/wrappers/.ci/docker-compose-native.yaml index b9c057d5a..5329bcfd2 100644 --- a/wrappers/.ci/docker-compose-native.yaml +++ b/wrappers/.ci/docker-compose-native.yaml @@ -96,6 +96,25 @@ services: timeout: 5s retries: 3 + zarr-http: + container_name: zarr-http + build: + context: ../dockerfiles/zarr-http + environment: + - ZARR_HTTP_FIXTURE_ROOT=/fixtures + volumes: + - ../dockerfiles/s3/test_data/zarr:/fixtures:ro + ports: + - "127.0.0.1:8787:8787" + read_only: true + security_opt: + - no-new-privileges:true + healthcheck: + test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8787/__health', timeout=2).read()"] + interval: 2s + timeout: 3s + retries: 15 + s3: image: minio/minio container_name: s3 @@ -142,6 +161,8 @@ services: depends_on: iceberg-rest: condition: service_healthy + zarr-http: + condition: service_healthy container_name: s3-init build: context: ../dockerfiles/s3 @@ -154,11 +175,170 @@ services: - ../dockerfiles/s3/iceberg_seed.py:/iceberg_seed.py entrypoint: | /bin/sh -c " + set -eu; until (/mc alias set s3 http://s3:8000 admin password) do echo '...waiting...' && sleep 1; done; - /mc rm -r --force s3/warehouse; - /mc mb s3/warehouse; + /mc rm -r --force s3/warehouse || true; + /mc mb --ignore-existing s3/warehouse; /mc policy set public s3/warehouse; /mc cp --recursive /test_data/* s3/warehouse; + /mc stat s3/warehouse/zarr/e2e.zarr/.zgroup; + /mc stat s3/warehouse/zarr/e2e.zarr/.zattrs; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/raw/.zarray; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/raw/.zattrs; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/raw/0.0.0; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/blosc/0.1.1; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/lazy1m/.zarray; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/spatial2d/.zarray; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/spatial2d/.zattrs; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/spatial_ref/.zattrs; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/x/.zarray; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/x/0; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/y/.zarray; + /mc stat s3/warehouse/zarr/e2e.zarr/nested/y/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/CODEC_FIXTURE_MANIFEST.txt; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/BLOSC_FIXTURE_MANIFEST.txt; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/ZSTD_FIXTURE_MANIFEST.txt; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/SHARDING_FIXTURE_MANIFEST.txt; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_default/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_v2keys/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/pipeline/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/pipeline/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_crc/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_crc/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_gzip/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_gzip/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/oversize/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/oversize/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_v3/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_x/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_x/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_x/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_coord_values/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_blosc/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_pipeline/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_x/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_x/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_x/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_coord_values/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/values/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/2; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_end/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_end/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_start/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_start/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_blosc/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_zstd/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/1/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_sentinel/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_sentinel/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_truncated_index/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_truncated_index/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_oob/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_oob/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_half_sentinel/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_half_sentinel/c/0/0/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/time/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/time/c/0; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/x/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/x/c/1; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/y/zarr.json; + /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/y/c/1; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/zarr.json; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/OME_FIXTURE_MANIFEST.txt; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/zarr.json; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/0/zarr.json; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/0/c/0/0; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/0/c/0/1; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/0/c/1/0; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/0/c/1/1; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/1/zarr.json; + /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/1/c/0/0; + if /mc stat s3/warehouse/zarr/e2e.zarr/nested/spatial2d/0.0; then + echo 'unexpected data chunk 0.0 in fill-only spatial Zarr fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e.zarr/nested/raw/0.1.1; then + echo 'unexpected raw chunk 0.1.1 in sparse Zarr fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_default/c/0/1/1; then + echo 'unexpected default-key chunk c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/raw_v2keys/0.1.1; then + echo 'unexpected v2-key chunk 0.1.1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/pipeline/c/0/1/1; then + echo 'unexpected pipeline chunk c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/1/1; then + echo 'unexpected Blosc chunk c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/1/1; then + echo 'unexpected Zstandard chunk c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_end/c/0/1/1; then + echo 'unexpected end-index shard c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_start/c/0/1/1; then + echo 'unexpected start-index shard c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/1/1; then + echo 'unexpected Blosc shard c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/1/1; then + echo 'unexpected Zstandard shard c/0/1/1 in sparse Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/x/zarr.json; then + echo 'unexpected x coordinate array in OME-Zarr v3 fixture'; + exit 1; + fi; + if /mc stat s3/warehouse/zarr/e2e-ome-v3.zarr/image/y/zarr.json; then + echo 'unexpected y coordinate array in OME-Zarr v3 fixture'; + exit 1; + fi; python /iceberg_seed.py " @@ -272,4 +452,4 @@ services: test: ["CMD", "sh", "-c", "echo 'db.runCommand({ ping: 1 })' | mongosh --quiet"] interval: 10s timeout: 5s - retries: 5 \ No newline at end of file + retries: 5 diff --git a/wrappers/.ci/pgrx-15.Dockerfile b/wrappers/.ci/pgrx-15.Dockerfile new file mode 100644 index 000000000..0b0bb5ac1 --- /dev/null +++ b/wrappers/.ci/pgrx-15.Dockerfile @@ -0,0 +1,79 @@ +# pgrx + PG15 test container for the wrappers crate +# +# Debian bookworm (rust:1.97.1) with PostgreSQL 15, clang, CMake and +# cargo-pgrx 0.19.2, so native FDWs can be compiled and tested against the PG +# version the repo's CI targets (see .github/workflows/test_wrappers.yml). +# +# Image contents: +# - rust:1.97.1 toolchain (matches workspace.package.rust-version) +# - PostgreSQL 15 server + -dev headers, clang (bindgen), sudo +# - cargo-pgrx 0.19.2 (used by the pgrx-tests framework) +# - ~/.pgrx/config.toml mapping pg15 -> /usr/lib/postgresql/15/bin/pg_config +# - non-root `builder` user (uid 1000): PostgreSQL refuses to run as root, so +# pgrx-tests runs its postmaster as `builder` while installing the +# extension as root (CARGO_PGRX_TEST_RUNAS=builder). +# +# Build (from repo root): +# podman build -t wrappers-pgrx-pg15 -f wrappers/.ci/pgrx-15.Dockerfile wrappers +# +# Keep a persistent target dir for incremental repeats. IMPORTANT: put it on +# real disk, NOT under /tmp — /tmp is a tmpfs (typically ~8 GiB) and the full +# pgrx + aws-sdk-s3 + aws-lc-sys/ring build overflows it, which surfaces as a +# spurious pgrx-tests "Could not obtain test mutex" error. Use e.g.: +# mkdir -p /home//pgrx-target/tmp +# and run with: +# -e CARGO_TARGET_DIR=/target -e TMPDIR=/target/tmp \ +# -v /home//pgrx-target:/target +# +# Compile check: +# podman run --rm \ +# -e CARGO_TARGET_DIR=/target -e TMPDIR=/target/tmp \ +# -v /home//pgrx-target:/target \ +# -v "$PWD":/work -w /work wrappers-pgrx-pg15 \ +# cargo check -p wrappers --no-default-features --features zarr_fdw,pg15 +# +# Unit tests (pure #[test] modules): +# podman run --rm \ +# -e CARGO_TARGET_DIR=/tmp/target -v /tmp/pgrx-target:/tmp/target \ +# -v "$PWD":/work -w /work wrappers-pgrx-pg15 \ +# cargo test -p wrappers --no-default-features --features zarr_fdw,pg15 \ +# --lib -- --skip pg_zarr +# +# #[pg_test] DDL tests (throws up a throwaway PG15 serving the extension): +# podman run --rm \ +# -e CARGO_TARGET_DIR=/tmp/target -e CARGO_PGRX_TEST_RUNAS=builder \ +# -v /tmp/pgrx-target:/tmp/target -v "$PWD":/work -w /work wrappers-pgrx-pg15 \ +# cargo test -p wrappers --no-default-features --features zarr_fdw,pg15 \ +# --lib pg_zarr -- --test-threads=1 +# +# Notes: +# - Always launch through `cargo test` (not the test binary directly): the +# pgrx-tests framework walks the process tree to discover the cargo +# `--features` it must hand to `cargo-pgrx install --test`. +# - Do NOT mount the host ~/.cargo over /usr/local/cargo (host rustup keeps +# its bin/ elsewhere, shadowing cargo with an empty directory). +# - sudo + the `builder` user + "chmod o+rx /root" let pgrx install as root +# while the postmaster runs as a regular user. + +FROM rust:1.97.1-bookworm + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + clang \ + cmake \ + postgresql-15 \ + postgresql-15-postgis-3 \ + postgresql-15-postgis-3-scripts \ + postgresql-server-dev-15 \ + sudo \ + && rm -rf /var/lib/apt/lists/* + +RUN chmod o+rx /root \ + && useradd -m -u 1000 builder \ + && mkdir -p /tmp/pgtarget \ + && chown -R builder:builder /tmp/pgtarget + +RUN cargo install --locked cargo-pgrx --version 0.19.2 \ + && cargo pgrx init --pg15 /usr/lib/postgresql/15/bin/pg_config \ + && rustup component add rustfmt clippy \ + && rm -rf /usr/local/cargo/registry/src /usr/local/cargo/registry/cache diff --git a/wrappers/Cargo.toml b/wrappers/Cargo.toml index e8680cdb8..2b9a35c29 100644 --- a/wrappers/Cargo.toml +++ b/wrappers/Cargo.toml @@ -203,6 +203,27 @@ dynamodb_fdw = [ "thiserror", "tokio", ] +zarr_fdw = [ + "aws-config", + "aws-sdk-s3", + "blosc-rs", + "cap-std", + "crc32c", + "chrono", + "futures-util", + "libc", + "lru", + "reqwest", + "tokio", + "tokio-util", + "async-compression", + "zstd", + "serde", + "serde_json", + "thiserror", + "http", + "url", +] # Does not include helloworld_fdw because of its general uselessness native_fdws = [ "airtable_fdw", @@ -222,6 +243,7 @@ native_fdws = [ "iceberg_fdw", "duckdb_fdw", "dynamodb_fdw", + "zarr_fdw", ] all_fdws = ["native_fdws", "wasm_fdw"] [dependencies] @@ -253,7 +275,7 @@ wiremock = { version = "0.5", optional = true } futures = { version = "0.3", optional = true } # for stripe_fdw, firebase_fdw, logflare_fdw, iceberg_fdw and etc. -reqwest = { version = "0.12.12", features = ["json", "gzip"], optional = true } +reqwest = { version = "0.12.12", features = ["json", "gzip", "rustls-tls"], optional = true } reqwest-middleware = { version = "0.4", optional = true } reqwest-retry = { version = "0.8", optional = true } @@ -272,6 +294,14 @@ aws-config = { version = "1.8.6", features = [ # for s3_fdw aws-sdk-s3 = { version = "1.86.0", optional = true } +# for zarr_fdw +blosc-rs = { version = "0.4.0", features = ["lz4"], optional = true } +cap-std = { version = "3.4.5", optional = true } +crc32c = { version = "0.6.8", optional = true } +libc = { version = "0.2", optional = true } +lru = { version = "0.12.5", optional = true } +zstd = { version = "0.13.3", default-features = false, optional = true } + # for s3vectors_fdw aws-sdk-s3vectors = { version = "1.14.0", optional = true } aws-smithy-types = { version = "1.3.2", optional = true } diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/OME_FIXTURE_MANIFEST.txt b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/OME_FIXTURE_MANIFEST.txt new file mode 100644 index 000000000..a3be7d183 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/OME_FIXTURE_MANIFEST.txt @@ -0,0 +1,39 @@ +Deterministic OME-Zarr 0.5 fixture + +All JSON files are compact UTF-8 with one trailing LF. Chunk payloads are raw +little-endian float32 values in C order and have no trailing bytes. + +Path Size SHA-256 +zarr.json 98 eb68e3c30a89eb1883bcdb152468afc6c853d65861f98751fec19847e886465a +image/zarr.json 619 e28c2f7555f7b3da45aa2c6c1e3a490e9dd72288d6ec7bebf7fa2ef82d8d5cc2 +image/0/zarr.json 344 43db3f8fb94961e250b056277dd2fdb1c6e1c6d2fc976339749ae0822c9e24f5 +image/0/c/0/0 36 83cf3898c7093e3e5466e574f8caa41efb7af0a79e96d5b8d23020e38d85d904 +image/0/c/0/1 36 cfee151d5217284f9d1b1d46bb6668b398072fa40b54f76469da02c927f9cc49 +image/0/c/1/0 36 c81547bbf5e968af93c468aa6aff1b1be7973e585f63f2c898806af82f6379b7 +image/0/c/1/1 36 9c8d0085e96b400e2503648181df76a578189a794166b26d6e776fe267f9a6d4 +image/1/zarr.json 344 876066672fde070eabed12258f8d10e1abfcd71cd78b4a22a8dd57766c5e489f +image/1/c/0/0 16 86536b4219044ba9c64ab26ca4e95312a342554c145ed88c96e44f05e1111535 + +Decoded level 0 (shape [4,4], chunks [3,3]; edge padding is -1): + 0 1 2 3 + 4 5 6 7 + 8 9 10 11 + 12 13 14 15 + +Decoded level 1 (shape [2,2], chunks [2,2], 2x2 means): + 2.5 4.5 + 10.5 12.5 + +Dataset transforms are followed by the multiscale transform. Effective +coordinates are: + level 0: scale [4,12], translation [120,260] + y [120,124,128,132], x [260,272,284,296] + level 1: scale [8,24], translation [122,266] + y [122,130], x [266,290] + +Literal payload hex: + image/0/c/0/0 000000000000803f00000040000080400000a0400000c040000000410000104100002041 + image/0/c/0/1 00004040000080bf000080bf0000e040000080bf000080bf00003041000080bf000080bf + image/0/c/1/0 000040410000504100006041000080bf000080bf000080bf000080bf000080bf000080bf + image/0/c/1/1 00007041000080bf000080bf000080bf000080bf000080bf000080bf000080bf000080bf + image/1/c/0/0 00002040000090400000284100004841 diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/0 new file mode 100644 index 000000000..4d2a97a35 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/1 new file mode 100644 index 000000000..8b5ab1dc3 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/0 new file mode 100644 index 000000000..6dfdf6a11 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/1 new file mode 100644 index 000000000..034ef7f16 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/c/1/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/zarr.json new file mode 100644 index 000000000..b791087cb --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/0/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[4,4],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[3,3]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-1.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["y","x"],"attributes":{}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/c/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/c/0/0 new file mode 100644 index 000000000..49eae9b42 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/c/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/zarr.json new file mode 100644 index 000000000..645b84a84 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/1/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,2],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,2]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-1.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["y","x"],"attributes":{}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/zarr.json new file mode 100644 index 000000000..deefb8d88 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/image/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"group","attributes":{"ome":{"version":"0.5","multiscales":[{"name":"mean-pyramid","axes":[{"name":"y","type":"space","unit":"micrometer"},{"name":"x","type":"space","unit":"micrometer"}],"datasets":[{"path":"0","coordinateTransformations":[{"type":"scale","scale":[2.0,4.0]},{"type":"translation","translation":[10.0,20.0]}]},{"path":"1","coordinateTransformations":[{"type":"scale","scale":[4.0,8.0]},{"type":"translation","translation":[11.0,22.0]}]}],"coordinateTransformations":[{"type":"scale","scale":[2.0,3.0]},{"type":"translation","translation":[100.0,200.0]}],"type":"mean"}]}}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/zarr.json new file mode 100644 index 000000000..c45d9ced7 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-ome-v3.zarr/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"group","attributes":{"title":"Deterministic OME-Zarr 0.5 fixture"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/BLOSC_FIXTURE_MANIFEST.txt b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/BLOSC_FIXTURE_MANIFEST.txt new file mode 100644 index 000000000..4c716fd78 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/BLOSC_FIXTURE_MANIFEST.txt @@ -0,0 +1,28 @@ +Generator: zarr-python 3.1.3, numpy 2.5.2, numcodecs 0.16.5 +Container: docker.io/library/python:3.12.5-slim-bookworm@sha256:f362e1c75ff1670f2776d72cff6ad84094029d2fa73b9fcf5bd7b5b07f45271c (arm64) +Determinism: two isolated ephemeral-container generations compared byte-for-byte +Committed metadata: generator JSON plus one terminal newline; hashes below cover committed bytes +Blosc: cname=lz4, clevel=5, shuffle=shuffle, typesize=4 (8 for f64 coordinate), blocksize=0 +Logical cube: value[time,y,x] = 100*time + 10*y + x; shape [2,5,6]; fill -7.5 +Sparse direct/sharded arrays omit c/0/1/1. The coordinate values are 100,110,120,130,140,150. +bad_blosc/c/0/0/0 keeps 15 Blosc header bytes and adds a valid outer CRC32C. +bad_blosc/c/0/0/1 changes the declared uncompressed size from 96 to 100 and adds a valid outer CRC32C. + +SHA256 BYTES PATH +12162ac50f16e41f564401531f2c238b987e86f6916c87628b7452bb729b1af0 19 nested/bad_blosc/c/0/0/0 +fc4ede630266d27dfb47b8f70251470407e6a1a4c88b9912fd78785cd54633da 116 nested/bad_blosc/c/0/0/1 +0e5ea30294c97076f85b1d57f1e3dd08f279d82ec409489238c4a7f27e239e82 661 nested/bad_blosc/zarr.json +4c9c4f354e74153db012329d71c8562ec23e498148174b2c49de58f45d47cdbe 16 nested/blosc_coord_values/c/0 +61e6781b6b1971e032dc9f85ea23b09e68e90cbafb1b2acb635397552273e94c 16 nested/blosc_coord_values/c/1 +d3040712263db0e3d1e183e66d43cb5aa68357d1f2294cabc107371ced9f454f 368 nested/blosc_coord_values/zarr.json +7c2e635c736d4da3c5893ccdf3af228499dfc621f3933e29d055cbfbef1c3f94 112 nested/blosc_v3/c/0/0/0 +0e6729d2277853bdb6f764bef5983a1fad6cb2b381d35ee96e8f00e84d92181f 112 nested/blosc_v3/c/0/0/1 +c2f6f14cb5fb5274b153b63bc3867107a5e6035956bc00963c60ac7490fa27b0 112 nested/blosc_v3/c/0/1/0 +e74afc8b2e017854233e08916d4e5da31b55f4e964b8b26993ee4b78998bab04 643 nested/blosc_v3/zarr.json +c0856a3ee16e598ef6b15caf1f1899fa3b95e04ac5b050cd07f8072b4a3d94f2 48 nested/blosc_x/c/0 +a6f4fa25efd5cc86ff168b1f034b4410f91c5791091dfce3d8bf23a97f3b6459 48 nested/blosc_x/c/1 +0d42a4fb3dd258c0272618b35ff78d4d7f2208c4c72e03c4907de90abd3331d2 538 nested/blosc_x/zarr.json +0b20a9a038c73d1cbdc9e4c96d78c00df53355bbcb9b34a234c40c39697343d0 228 nested/shard_blosc/c/0/0/0 +9390f39372380bd02b2705f7b7379fd1cf52da544b82fccb871df037eab53c70 148 nested/shard_blosc/c/0/0/1 +8213467353f85b55923a4a391f200478c0ac9af3bfcb86e9ea07aba4cdefcdf4 228 nested/shard_blosc/c/0/1/0 +c98676cbae021f0d168089e09d45320c8a0c165351f69096f24ea157cdf7857e 833 nested/shard_blosc/zarr.json diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/CODEC_FIXTURE_MANIFEST.txt b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/CODEC_FIXTURE_MANIFEST.txt new file mode 100644 index 000000000..7ca520c67 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/CODEC_FIXTURE_MANIFEST.txt @@ -0,0 +1,7 @@ +SHA256 BYTES PATH +a9847c35ecb7cf3ec3778eed644039258e5dbdb3a3fd47c8b427015c76486c91 93 nested/pipeline/c/0/0/0 +167b7d675605843f0e35c6cf24ccfa382de12fe3ebfc6cb7e8c4c7f490f11687 67 nested/pipeline/c/0/0/1 +38785f6443a2f1ccd12a701a7c12138c534bb05979dbd7c5754ff2e0c5364de8 78 nested/pipeline/c/0/1/0 +360e8a0ef3b3976547f5e2287fb80966e56b01fbd72472155059e8350b2d26aa 93 nested/bad_crc/c/0/0/0 +76f8fb842435a035a1f12e77147ae147c4bc9accae04264c8c5b69c2e5c64a75 85 nested/bad_gzip/c/0/0/0 +02cde18c039b3402507e1dcd52c98014ed58c03bcf34f32b7311b1db16f78690 94 nested/oversize/c/0/0/0 diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/SHARDING_FIXTURE_MANIFEST.txt b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/SHARDING_FIXTURE_MANIFEST.txt new file mode 100644 index 000000000..545183097 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/SHARDING_FIXTURE_MANIFEST.txt @@ -0,0 +1,23 @@ +Generator: zarr-python 3.1.3, numpy 2.5.2, numcodecs 0.16.5 +Container: docker.io/library/python:3.12.5-slim-bookworm@sha256:f362e1c75ff1670f2776d72cff6ad84094029d2fa73b9fcf5bd7b5b07f45271c (arm64) +Layout: shape [2,5,6], outer shard [2,3,4], inner chunk [1,3,2], 68-byte bytes-little+CRC32C index +First-shard C-order entries, index=end: (0,24), (48,24), (24,24), (72,24) +First-shard C-order entries, index=start: (68,24), (116,24), (92,24), (140,24) +The non-monotonic offsets are the pinned zarr-python Morton payload order. +Mutations from shard_end/c/0/0/0: sentinel entry 0 = (MAX,MAX); bad-index-crc +flips the final checksum byte; truncated-index keeps the first 67 bytes; oob +entry 0 = (165,24); half-sentinel entry 0 = (MAX,24). Mutated indexes except +bad-index-crc and truncated-index have a recomputed valid CRC32C trailer. + +SHA256 BYTES PATH +711994c2fd69ffddd195ff57991d85b98faba365038f6e451a3e0c6f437d5bd4 164 nested/shard_end/c/0/0/0 +b24cb579bafb7e8acc13d6ba78dd224565920c8c4fa720631a7e25dde09717f8 116 nested/shard_end/c/0/0/1 +25db91c2e4248b71c4b9d3ebbbb07f50c3c241e085b7279570cdeab586bbba9e 164 nested/shard_end/c/0/1/0 +9a4e467e39ce4ca9437d8a01f4729933c43b3ff6146a0198a9c9be5b583f12e3 164 nested/shard_start/c/0/0/0 +b97485ec64355d7bcff095a88eb90fcec8935f7a986635b38b6da554ed0bac98 116 nested/shard_start/c/0/0/1 +e963ff51d5c7fa079d0a7970b652e9af871032ce91c22f71799125b1baba7285 164 nested/shard_start/c/0/1/0 +4668127bc1ca0b8b142f8b19543dd8097ede277cc2851f87feecbb5bcccfdf2a 164 nested/shard_sentinel/c/0/0/0 +83a28a36ce98056904786c984412ad5b00c97f689847144f4be4902f8a6a8c31 164 nested/shard_bad_index_crc/c/0/0/0 +0dae47b986c0ad215eee6c7ba0fc19a3cb9715d6b75ecddc58180fbcd274c0f6 67 nested/shard_truncated_index/c/0/0/0 +1f208cae5d6e14b3d6f37a457ddd775abd28d2878243fbd44b64152f18491bfc 164 nested/shard_oob/c/0/0/0 +b0fd6648dae55bac8a30317f089cf40491c7011da3df9fccb0bdd1c9ab24c083 164 nested/shard_half_sentinel/c/0/0/0 diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/ZSTD_FIXTURE_MANIFEST.txt b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/ZSTD_FIXTURE_MANIFEST.txt new file mode 100644 index 000000000..4eae8baa0 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/ZSTD_FIXTURE_MANIFEST.txt @@ -0,0 +1,33 @@ +Generator: zarr-python 3.1.3, numpy 2.5.2, numcodecs 0.16.5 (libzstd 1.5.7) +Container: docker.io/library/python:3.12.5-slim-bookworm@sha256:f362e1c75ff1670f2776d72cff6ad84094029d2fa73b9fcf5bd7b5b07f45271c (arm64) +Determinism: two isolated ephemeral-container generations compared byte-for-byte +Committed metadata: generator JSON plus one terminal newline; hashes below cover committed bytes +Primary pipeline: transpose [2,1,0] -> bytes little -> zstd level=1/checksum=true -> crc32c +Coordinate pipeline: bytes little -> zstd level=1/checksum=false +Sharded pipeline: one [2,3,4] inner chunk per [2,3,4] shard, zstd checksum=true, 20-byte end index +Logical cube: value[time,y,x] = 100*time + 10*y + x; shape [2,5,6]; fill -7.5 +Sparse direct/sharded arrays omit c/0/1/1. The zstd_x coordinate values are 100,110,120,130,140,150. +zstd_bad cases 0..2 are a flipped content checksum, a 16 MiB window, and a nonzero Dictionary_ID_Flag. + +SHA256 BYTES PATH +186f3d6b9c63a431b675a2a53e2e94b4540d9d50e12fd9488a08f519e9777f9a 98 nested/shard_zstd/c/0/0/0 +8631ccb760254df6ca324ec835191bc75f82769a999e7301df1d2ffb228eb18a 98 nested/shard_zstd/c/0/0/1 +5aece0b83a4cc262141d307ca2cd4491afdbc2575474e54afac9191ceae6e7c7 94 nested/shard_zstd/c/0/1/0 +0862bebbfddb4cabb08a9a968e739f3d4f936f63d90432f9dd176f171ac0e584 786 nested/shard_zstd/zarr.json +b0c45303f7f11848cb5e6e5b2af2fb2aecd0b72c28748b88b583ab6bb76df174 24 nested/zstd_bad/failure_case/c/0 +e16c5d4040bf050671561dde17680b678c388fe3bd014aa4c4c7edd30fb9e204 372 nested/zstd_bad/failure_case/zarr.json +577c5a5001f0b43d1f5bdc13c03ee9ea4e831fd478cc8aab8ffdf1c73b66c8a8 17 nested/zstd_bad/values/c/0 +6c2195f2818e3f2ed77c814489dad2557d666c2a37861acc033b2cc3f25d6604 6 nested/zstd_bad/values/c/1 +72de3c89e92c19de606fa699e23c64383048fb462f27f7f0cf2ac6d8fd74caf7 7 nested/zstd_bad/values/c/2 +c523e7a72f6e944961a2230758a3c6a195daa35a2e885269e5ed7a42a1a6abe4 509 nested/zstd_bad/values/zarr.json +299dcc2ba841f9c6ecdea892e044612119e63f9b6cf6de6b6114a5f0dd200e31 54 nested/zstd_bad/zarr.json +4c9c4f354e74153db012329d71c8562ec23e498148174b2c49de58f45d47cdbe 16 nested/zstd_coord_values/c/0 +61e6781b6b1971e032dc9f85ea23b09e68e90cbafb1b2acb635397552273e94c 16 nested/zstd_coord_values/c/1 +d21f17b187aae277e7a812540097c60ce0db9c41964de0712533e0a506cdb032 367 nested/zstd_coord_values/zarr.json +53097d09f2ead2addce462771e3850d07b0d564de7c850b6da1085820da7d89e 82 nested/zstd_pipeline/c/0/0/0 +e6e0b455e6a65c4662446e1ca1f4e9c508d35edc7c5006fab1abed83a7e5b181 77 nested/zstd_pipeline/c/0/0/1 +4065d0b223321e5f833a8764fbe15788d676d1dcde86805872137423c90b5055 77 nested/zstd_pipeline/c/0/1/0 +b47ff6445485d2e7846e4d6ea60e0b702d51b125543962dfb9032ef834884c37 669 nested/zstd_pipeline/zarr.json +73f899386280fcf55ebeebe74ed1ab13eb10261a0777d8b3606b8de78c8b340a 41 nested/zstd_x/c/0 +a0c2a2e4d6aff4b06ecf1028b0c56707308931582056677a92ee57e4a0cc743e 29 nested/zstd_x/c/1 +de0c79097b57f2c0b522a544a5df8f58c36af5e7a1efcf7b33a4f6cc2fdc11ce 491 nested/zstd_x/zarr.json diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/0 new file mode 100644 index 000000000..b67413642 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/1 new file mode 100644 index 000000000..846d57236 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/zarr.json new file mode 100644 index 000000000..b270b1e56 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_blosc/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"blosc","configuration":{"typesize":4,"cname":"lz4","clevel":5,"shuffle":"shuffle","blocksize":0}},{"name":"crc32c"}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/c/0/0/0 new file mode 100644 index 000000000..5c743a729 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/zarr.json new file mode 100644 index 000000000..5689d97a3 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_crc/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default"},"fill_value":-7.5,"codecs":[{"name":"transpose","configuration":{"order":[2,1,0]}},{"name":"bytes","configuration":{"endian":"little"}},{"name":"gzip","configuration":{"level":1}},{"name":"crc32c"}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/c/0/0/0 new file mode 100644 index 000000000..ce14f2e94 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/zarr.json new file mode 100644 index 000000000..5689d97a3 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/bad_gzip/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default"},"fill_value":-7.5,"codecs":[{"name":"transpose","configuration":{"order":[2,1,0]}},{"name":"bytes","configuration":{"endian":"little"}},{"name":"gzip","configuration":{"level":1}},{"name":"crc32c"}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/0 new file mode 100644 index 000000000..ae6bb1d6c Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/1 new file mode 100644 index 000000000..2561fe03f Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/zarr.json new file mode 100644 index 000000000..2ac8741e6 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_coord_values/zarr.json @@ -0,0 +1 @@ +{"shape":[6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"attributes":{},"dimension_names":["blosc_x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/0 new file mode 100644 index 000000000..35e7e6569 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/1 new file mode 100644 index 000000000..3b24878bd Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/1/0 new file mode 100644 index 000000000..79db50ddd Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/zarr.json new file mode 100644 index 000000000..ce6474f2d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_v3/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"blosc","configuration":{"typesize":4,"cname":"lz4","clevel":5,"shuffle":"shuffle","blocksize":0}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/0 new file mode 100644 index 000000000..74f37c0cc Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/1 new file mode 100644 index 000000000..400c3a88c Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/zarr.json new file mode 100644 index 000000000..9d9261b1a --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/blosc_x/zarr.json @@ -0,0 +1 @@ +{"shape":[6],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"blosc","configuration":{"typesize":8,"cname":"lz4","clevel":5,"shuffle":"shuffle","blocksize":0}}],"attributes":{"axis":"X","standard_name":"projection_x_coordinate","units":"m"},"dimension_names":["blosc_x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/c/0/0/0 new file mode 100644 index 000000000..b9d7cbbd3 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/zarr.json new file mode 100644 index 000000000..5689d97a3 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/oversize/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default"},"fill_value":-7.5,"codecs":[{"name":"transpose","configuration":{"order":[2,1,0]}},{"name":"bytes","configuration":{"endian":"little"}},{"name":"gzip","configuration":{"level":1}},{"name":"crc32c"}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/0 new file mode 100644 index 000000000..6107f8cc1 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/1 new file mode 100644 index 000000000..6582777a3 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/1/0 new file mode 100644 index 000000000..6ae6c0615 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/zarr.json new file mode 100644 index 000000000..5689d97a3 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/pipeline/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default"},"fill_value":-7.5,"codecs":[{"name":"transpose","configuration":{"order":[2,1,0]}},{"name":"bytes","configuration":{"endian":"little"}},{"name":"gzip","configuration":{"level":1}},{"name":"crc32c"}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/0 new file mode 100644 index 000000000..650e6cb76 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/1 new file mode 100644 index 000000000..b1dbde051 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/1/0 new file mode 100644 index 000000000..858a513a7 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/zarr.json new file mode 100644 index 000000000..51c5d72d6 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_default/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default"},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.0 new file mode 100644 index 000000000..650e6cb76 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.1 new file mode 100644 index 000000000..b1dbde051 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.0.1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.1.0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.1.0 new file mode 100644 index 000000000..858a513a7 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/0.1.0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/zarr.json new file mode 100644 index 000000000..02a0c42f1 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/raw_v2keys/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"v2","configuration":{"separator":"."}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["time","y","x"],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/c/0/0/0 new file mode 100644 index 000000000..86bb51d47 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_bad_index_crc/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/0 new file mode 100644 index 000000000..6b3a66c37 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/1 new file mode 100644 index 000000000..ed568cedc Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/1/0 new file mode 100644 index 000000000..9e3b4c450 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/zarr.json new file mode 100644 index 000000000..c4608260b --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_blosc/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"blosc","configuration":{"typesize":4,"cname":"lz4","clevel":5,"shuffle":"shuffle","blocksize":0}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/0 new file mode 100644 index 000000000..115ad486c Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/1 new file mode 100644 index 000000000..da34408e4 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/1/0 new file mode 100644 index 000000000..cd4004fde Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_end/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/c/0/0/0 new file mode 100644 index 000000000..a5e9bdedc Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_half_sentinel/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/c/0/0/0 new file mode 100644 index 000000000..c5a827416 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_oob/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/c/0/0/0 new file mode 100644 index 000000000..0b8015f3b Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_sentinel/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/0 new file mode 100644 index 000000000..178b106a2 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/1 new file mode 100644 index 000000000..073441614 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/1/0 new file mode 100644 index 000000000..fe589e819 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/zarr.json new file mode 100644 index 000000000..dd5f10a40 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_start/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"start"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/c/0/0/0 new file mode 100644 index 000000000..f8c4265f1 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/zarr.json new file mode 100644 index 000000000..a9203104d --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_truncated_index/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[1,3,2],"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/0 new file mode 100644 index 000000000..c8fe0e8f6 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/1 new file mode 100644 index 000000000..533d358b3 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/1/0 new file mode 100644 index 000000000..db612812d Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/zarr.json new file mode 100644 index 000000000..752832add --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/shard_zstd/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"sharding_indexed","configuration":{"chunk_shape":[2,3,4],"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"zstd","configuration":{"level":1,"checksum":true}}],"index_codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"crc32c"}],"index_location":"end"}}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/c/0 new file mode 100644 index 000000000..9b7cd39e4 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/zarr.json new file mode 100644 index 000000000..aabebce9b --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/time/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[2],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["time"],"attributes":{"calendar":"proleptic_gregorian","standard_name":"time","units":"milliseconds since 1970-01-01 00:00:00"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/0 new file mode 100644 index 000000000..3816f8d83 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/1 new file mode 100644 index 000000000..3792b21fb Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/zarr.json new file mode 100644 index 000000000..01ef441a2 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/x/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[6],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["x"],"attributes":{"axis":"X","standard_name":"projection_x_coordinate","units":"m"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/0 new file mode 100644 index 000000000..377939808 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/1 new file mode 100644 index 000000000..ddff17016 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/zarr.json new file mode 100644 index 000000000..a27ff518b --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/y/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"array","shape":[5],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[3]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"dimension_names":["y"],"attributes":{"axis":"Y","standard_name":"projection_y_coordinate","units":"m"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zarr.json new file mode 100644 index 000000000..773678b05 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"group","attributes":{"crs":{"properties":{"name":"EPSG:3857"},"type":"name"},"project":"zarr_fdw v3 metadata inspection"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/c/0 new file mode 100644 index 000000000..6f34a8067 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/zarr.json new file mode 100644 index 000000000..2653dec68 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/failure_case/zarr.json @@ -0,0 +1 @@ +{"shape":[3],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[3]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"attributes":{},"dimension_names":["failure_case"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/0 new file mode 100644 index 000000000..2c568df46 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/1 new file mode 100644 index 000000000..e1e3c1700 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/1 @@ -0,0 +1 @@ +(/p \ No newline at end of file diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/2 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/2 new file mode 100644 index 000000000..09eef0cf9 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/c/2 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/zarr.json new file mode 100644 index 000000000..080f90760 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/values/zarr.json @@ -0,0 +1 @@ +{"shape":[3],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[1]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"zstd","configuration":{"level":1,"checksum":true}}],"attributes":{"cases":["checksum corruption","16 MiB decoder window","dictionary ID flag"]},"dimension_names":["failure_case"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/zarr.json new file mode 100644 index 000000000..7cd650921 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_bad/zarr.json @@ -0,0 +1 @@ +{"attributes":{},"zarr_format":3,"node_type":"group"} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/0 new file mode 100644 index 000000000..ae6bb1d6c Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/1 new file mode 100644 index 000000000..2561fe03f Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/zarr.json new file mode 100644 index 000000000..18239295f --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_coord_values/zarr.json @@ -0,0 +1 @@ +{"shape":[6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"bytes","configuration":{"endian":"little"}}],"attributes":{},"dimension_names":["zstd_x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/0 new file mode 100644 index 000000000..c6bef8e54 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/1 new file mode 100644 index 000000000..ec680ad58 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/0/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/1/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/1/0 new file mode 100644 index 000000000..c914aff95 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/c/0/1/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/zarr.json new file mode 100644 index 000000000..908265788 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_pipeline/zarr.json @@ -0,0 +1 @@ +{"shape":[2,5,6],"data_type":"float32","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[2,3,4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":-7.5,"codecs":[{"name":"transpose","configuration":{"order":[2,1,0]}},{"name":"bytes","configuration":{"endian":"little"}},{"name":"zstd","configuration":{"level":1,"checksum":true}},{"name":"crc32c"}],"attributes":{"_FillValue":-7.5,"add_offset":273.15,"long_name":"packed air temperature","missing_value":[42.0],"scale_factor":0.01,"units":"K","valid_range":[0.0,140.0]},"dimension_names":["time","y","x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/0 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/0 new file mode 100644 index 000000000..7b25f8059 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/0 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/1 b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/1 new file mode 100644 index 000000000..7e962c5f0 Binary files /dev/null and b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/c/1 differ diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/zarr.json new file mode 100644 index 000000000..ce376b49f --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/nested/zstd_x/zarr.json @@ -0,0 +1 @@ +{"shape":[6],"data_type":"float64","chunk_grid":{"name":"regular","configuration":{"chunk_shape":[4]}},"chunk_key_encoding":{"name":"default","configuration":{"separator":"/"}},"fill_value":0.0,"codecs":[{"name":"bytes","configuration":{"endian":"little"}},{"name":"zstd","configuration":{"level":1,"checksum":false}}],"attributes":{"axis":"X","standard_name":"projection_x_coordinate","units":"m"},"dimension_names":["zstd_x"],"zarr_format":3,"node_type":"array","storage_transformers":[]} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/zarr.json b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/zarr.json new file mode 100644 index 000000000..1b7c0f393 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e-v3.zarr/zarr.json @@ -0,0 +1 @@ +{"zarr_format":3,"node_type":"group","attributes":{"institution":"Supabase Wrappers test fixture","title":"Deterministic Zarr v3 inspection fixture"}} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zattrs b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zattrs new file mode 100644 index 000000000..0da5ac9ed --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zattrs @@ -0,0 +1 @@ +{"institution":"Supabase Wrappers test fixture","title":"Deterministic Zarr inspection fixture"} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zgroup b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zgroup new file mode 100644 index 000000000..0c3ab9753 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/.zgroup @@ -0,0 +1 @@ +{"zarr_format":2} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zattrs b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zattrs new file mode 100644 index 000000000..9ca7df15a --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zattrs @@ -0,0 +1 @@ +{"crs":{"properties":{"name":"EPSG:3857"},"type":"name"},"project":"zarr_fdw metadata inspection"} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zgroup b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zgroup new file mode 100644 index 000000000..0c3ab9753 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/.zgroup @@ -0,0 +1 @@ +{"zarr_format":2} diff --git a/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/band/.zarray b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/band/.zarray new file mode 100644 index 000000000..6bce508b9 --- /dev/null +++ b/wrappers/dockerfiles/s3/test_data/zarr/e2e.zarr/nested/band/.zarray @@ -0,0 +1 @@ +{"chunks":[4],"compressor":null,"dimension_separator":".","dtype":" None: + self.object_gets = 0 + self.redirect_sink_gets = 0 + self.proxy_sink_gets = 0 + self.forbidden_header_gets = 0 + self.paths: Counter[str] = Counter() + self.ranges: Counter[str] = Counter() + self.if_matches: Counter[str] = Counter() + self.accept_encodings: Counter[str] = Counter() + + def as_dict(self) -> dict[str, object]: + return { + "object_gets": self.object_gets, + "redirect_sink_gets": self.redirect_sink_gets, + "proxy_sink_gets": self.proxy_sink_gets, + "forbidden_header_gets": self.forbidden_header_gets, + "paths": dict(sorted(self.paths.items())), + "ranges": dict(sorted(self.ranges.items())), + "if_matches": dict(sorted(self.if_matches.items())), + "accept_encodings": dict(sorted(self.accept_encodings.items())), + } + + +STATS: dict[str, CaseStats] = {} +STATS_LOCK = threading.Lock() + + +def increment(case: str, field: str) -> None: + with STATS_LOCK: + stats = STATS.setdefault(case, CaseStats()) + setattr(stats, field, getattr(stats, field) + 1) + + +def record_object_request( + case: str, + object_key: str, + range_header: str | None, + if_match: str | None, + accept_encoding: str | None, + forbidden: bool, +) -> None: + with STATS_LOCK: + stats = STATS.setdefault(case, CaseStats()) + stats.object_gets += 1 + stats.paths[object_key] += 1 + if range_header is not None: + stats.ranges[range_header] += 1 + if if_match is not None: + stats.if_matches[if_match] += 1 + stats.accept_encodings[accept_encoding or ""] += 1 + if forbidden: + stats.forbidden_header_gets += 1 + + +def decode_segments(path: str) -> list[str] | None: + try: + decoded = unquote(path, errors="strict") + except UnicodeError: + return None + if "\x00" in decoded or "\\" in decoded: + return None + segments = [segment for segment in decoded.split("/") if segment] + if any(segment in {".", ".."} for segment in segments): + return None + return segments + + +def fixture_object(fixture: str, object_segments: list[str]) -> tuple[Path, str] | None: + if fixture not in FIXTURES or not object_segments: + return None + fixture_root = (FIXTURE_ROOT / fixture).resolve() + candidate = fixture_root.joinpath(*object_segments).resolve(strict=False) + try: + candidate.relative_to(fixture_root) + except ValueError: + return None + return candidate, "/".join(object_segments) + + +def parse_range(value: str, total: int) -> tuple[int, int] | None: + if not value.startswith("bytes=") or "," in value: + return None + spec = value[6:] + if spec.startswith("-"): + length_text = spec[1:] + if not length_text.isdigit(): + return None + length = int(length_text) + if length <= 0 or total <= 0: + return None + length = min(length, total) + return total - length, total - 1 + if "-" not in spec: + return None + start_text, end_text = spec.split("-", 1) + if not start_text.isdigit() or not end_text.isdigit(): + return None + start = int(start_text) + end = int(end_text) + if start > end or start >= total or end >= total: + return None + return start, end + + +class ZarrHttpServer(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, _format: str, *_args: object) -> None: + return + + def send_empty(self, status: int, **headers: str) -> None: + self.send_response(status) + for name, value in headers.items(): + self.send_header(name.replace("_", "-"), value) + self.send_header("Content-Length", "0") + self.end_headers() + + def send_json(self, value: object) -> None: + body = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: + parsed = urlsplit(self.path) + segments = decode_segments(parsed.path) + if segments is None: + self.send_empty(400) + return + if segments == ["__health"] and not parsed.query: + self.send_empty(200) + return + if len(segments) == 2 and segments[0] == "__stats" and not parsed.query: + case = segments[1] + if not CASE_PATTERN.fullmatch(case): + self.send_empty(400) + return + with STATS_LOCK: + value = STATS.setdefault(case, CaseStats()).as_dict() + self.send_json(value) + return + if len(segments) == 2 and segments[0] == "__sink" and not parsed.query: + case = segments[1] + if not CASE_PATTERN.fullmatch(case): + self.send_empty(400) + return + increment(case, "redirect_sink_gets") + self.send_empty(418) + return + if len(segments) == 2 and segments[0] == "__proxy" and not parsed.query: + case = segments[1] + if not CASE_PATTERN.fullmatch(case): + self.send_empty(400) + return + increment(case, "proxy_sink_gets") + self.send_empty(502) + return + if len(segments) < 5 or segments[0] != "stores" or parsed.query: + self.send_empty(404) + return + + case, mode, fixture = segments[1:4] + if not CASE_PATTERN.fullmatch(case) or mode not in MODES: + self.send_empty(404) + return + resolved = fixture_object(fixture, segments[4:]) + if resolved is None: + self.send_empty(400) + return + path, object_key = resolved + range_header = self.headers.get("Range") + if_match = self.headers.get("If-Match") + accept_encoding = self.headers.get("Accept-Encoding") + forbidden = any(self.headers.get(name) is not None for name in FORBIDDEN_HEADERS) + if mode == "anonymous_only" and accept_encoding != "identity": + forbidden = True + record_object_request( + case, + object_key, + range_header, + if_match, + accept_encoding, + forbidden, + ) + + if mode == "deny_all": + self.send_empty(503) + return + if mode == "anonymous_only" and forbidden: + self.send_empty(400) + return + if mode == "redirect_chunk" and object_key == "nested/raw/0.0.0": + self.send_empty( + 302, + Location=f"http://127.0.0.1:{PORT}/__sink/{case}", + ) + return + if mode == "oversize_metadata" and object_key == "nested/raw/.zarray": + self.send_response(200) + self.send_header("Content-Length", str(OVERSIZE_METADATA_BYTES)) + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + return + if not path.is_file(): + self.send_empty(404) + return + + body = path.read_bytes() + etag = f'"{sha256(body).hexdigest()}"' + if mode == "stall_chunk" and object_key == "nested/raw/0.0.0": + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.send_header("ETag", etag) + self.end_headers() + self.wfile.flush() + time.sleep(5) + try: + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + pass + return + + if mode == "mutate_shard" and range_header is not None: + if range_header.startswith("bytes=-"): + etag = '"generation-a"' + elif if_match == '"generation-a"': + self.send_empty(412) + return + + if if_match is not None and if_match != etag: + self.send_empty(412) + return + if range_header is None: + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.send_header("ETag", etag) + self.end_headers() + self.wfile.write(body) + return + if mode == "range_200": + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.send_header("ETag", etag) + self.end_headers() + self.wfile.write(body) + return + + selected = parse_range(range_header, len(body)) + if selected is None: + self.send_empty(416, Content_Range=f"bytes */{len(body)}") + return + start, end = selected + response_body = body[start : end + 1] + content_range = f"bytes {start}-{end}/{len(body)}" + if mode == "bad_content_range": + content_range = f"bytes 0-{len(response_body) - 1}/{len(body)}" + self.send_response(206) + self.send_header("Content-Range", content_range) + self.send_header("Content-Length", str(len(response_body))) + if mode != "no_etag": + self.send_header("ETag", etag) + self.end_headers() + self.wfile.write(response_body) + + +if __name__ == "__main__": + if not FIXTURE_ROOT.is_dir(): + raise SystemExit(f"fixture root is unavailable: {FIXTURE_ROOT}") + server = ZarrHttpServer((HOST, PORT), Handler) + server.serve_forever() diff --git a/wrappers/src/fdw/mod.rs b/wrappers/src/fdw/mod.rs index 3a9991b98..4c1bb7b45 100644 --- a/wrappers/src/fdw/mod.rs +++ b/wrappers/src/fdw/mod.rs @@ -54,3 +54,6 @@ mod duckdb_fdw; #[cfg(feature = "dynamodb_fdw")] mod dynamodb_fdw; + +#[cfg(feature = "zarr_fdw")] +mod zarr_fdw; diff --git a/wrappers/src/fdw/zarr_fdw/aggregate.rs b/wrappers/src/fdw/zarr_fdw/aggregate.rs new file mode 100644 index 000000000..9948eec10 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/aggregate.rs @@ -0,0 +1,982 @@ +use std::cmp::Ordering; + +use pgrx::{AnyNumeric, pg_sys}; +use supabase_wrappers::prelude::{Aggregate, AggregateKind, Cell, ParamValue, Qual, Value}; + +use super::{ZarrFdwError, ZarrFdwResult}; + +const SUPPORTED_SCALAR_OIDS: &[pg_sys::Oid] = &[ + pg_sys::CHAROID, + pg_sys::INT2OID, + pg_sys::INT4OID, + pg_sys::INT8OID, + pg_sys::FLOAT4OID, + pg_sys::FLOAT8OID, + pg_sys::TIMESTAMPTZOID, +]; + +pub(crate) fn aggregate_signature_supported(aggregate: &Aggregate) -> bool { + if aggregate.distinct { + return false; + } + + match aggregate.kind { + AggregateKind::Count => aggregate.column.is_none() && aggregate.type_oid == pg_sys::INT8OID, + AggregateKind::CountColumn => aggregate.column.as_ref().is_some_and(|column| { + SUPPORTED_SCALAR_OIDS.contains(&column.type_oid) + && aggregate.type_oid == pg_sys::INT8OID + }), + AggregateKind::Sum => aggregate.column.as_ref().is_some_and(|column| { + matches!( + (column.type_oid, aggregate.type_oid), + (pg_sys::INT2OID | pg_sys::INT4OID, pg_sys::INT8OID) + | (pg_sys::INT8OID, pg_sys::NUMERICOID) + | (pg_sys::FLOAT4OID, pg_sys::FLOAT4OID) + | (pg_sys::FLOAT8OID, pg_sys::FLOAT8OID) + ) + }), + AggregateKind::Avg => aggregate.column.as_ref().is_some_and(|column| { + matches!( + (column.type_oid, aggregate.type_oid), + ( + pg_sys::INT2OID | pg_sys::INT4OID | pg_sys::INT8OID, + pg_sys::NUMERICOID + ) | (pg_sys::FLOAT4OID | pg_sys::FLOAT8OID, pg_sys::FLOAT8OID) + ) + }), + AggregateKind::Min | AggregateKind::Max => { + aggregate.column.as_ref().is_some_and(|column| { + SUPPORTED_SCALAR_OIDS.contains(&column.type_oid) + && aggregate.type_oid == column.type_oid + }) + } + } +} + +fn qual_cell_supported(cell: &Cell) -> bool { + matches!( + cell, + Cell::I8(_) + | Cell::I16(_) + | Cell::I32(_) + | Cell::I64(_) + | Cell::F32(_) + | Cell::F64(_) + | Cell::Timestamptz(_) + ) +} + +pub(crate) fn qual_shape_supported(qual: &Qual) -> bool { + if qual.use_or { + return qual.operator == "=" + && matches!(&qual.value, Value::Array(values) if values.iter().all(qual_cell_supported)); + } + + match (&*qual.operator, &qual.value) { + ("is" | "is not", Value::Cell(Cell::String(value))) => value == "null", + ("=" | "<>" | "!=" | "<" | "<=" | ">" | ">=", Value::Cell(cell)) => { + qual_cell_supported(cell) + } + _ => false, + } +} + +pub(crate) fn qual_matches(qual: &Qual, source: Option<&Cell>) -> ZarrFdwResult { + if !qual_shape_supported(qual) { + return Err(aggregate_error(format!( + "qualifier on '{}' with operator '{}' is not supported for exact evaluation", + qual.field, qual.operator + ))); + } + + match qual.operator.as_str() { + "is" => return Ok(source.is_none()), + "is not" => return Ok(source.is_some()), + _ => {} + } + + let Some(source) = source else { + // SQL comparisons with NULL evaluate to unknown, which WHERE rejects. + return Ok(false); + }; + + let evaluated_value = qual.param.as_ref().map(|_| qual.evaluated_value()); + let value = match evaluated_value.as_ref() { + None => &qual.value, + Some(ParamValue::Value(value)) => value, + Some(ParamValue::Null) => return Ok(false), + Some(ParamValue::Unevaluated) => { + return Err(aggregate_error(format!( + "parameter for qualifier on '{}' was not evaluated before aggregate execution", + qual.field + ))); + } + }; + + if qual.use_or { + let Value::Array(values) = value else { + unreachable!("qual_shape_supported checked equality IN values") + }; + return values.iter().try_fold(false, |matched, value| { + Ok(matched || compare_cells(source, value)? == Ordering::Equal) + }); + } + + let Value::Cell(value) = value else { + unreachable!("qual_shape_supported checked scalar comparison value") + }; + let ordering = compare_cells(source, value)?; + Ok(match qual.operator.as_str() { + "=" => ordering == Ordering::Equal, + "<>" | "!=" => ordering != Ordering::Equal, + "<" => ordering == Ordering::Less, + "<=" => ordering != Ordering::Greater, + ">" => ordering == Ordering::Greater, + ">=" => ordering != Ordering::Less, + _ => unreachable!("qual_shape_supported checked comparison operator"), + }) +} + +fn compare_f32(left: f32, right: f32) -> Ordering { + match (left.is_nan(), right.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => left + .partial_cmp(&right) + .expect("non-NaN floats always have an ordering"), + } +} + +pub(crate) fn compare_f64(left: f64, right: f64) -> Ordering { + match (left.is_nan(), right.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => left + .partial_cmp(&right) + .expect("non-NaN floats always have an ordering"), + } +} + +fn compare_cells(left: &Cell, right: &Cell) -> ZarrFdwResult { + if let (Some(left), Some(right)) = (integer_cell_as_i128(left), integer_cell_as_i128(right)) { + return Ok(left.cmp(&right)); + } + + let ordering = match (left, right) { + (Cell::I8(left), Cell::I8(right)) => left.cmp(right), + (Cell::F32(left), Cell::F32(right)) => compare_f32(*left, *right), + (Cell::F32(left), Cell::F64(right)) => compare_f64(f64::from(*left), *right), + (Cell::F64(left), Cell::F32(right)) => compare_f64(*left, f64::from(*right)), + (Cell::F64(left), Cell::F64(right)) => compare_f64(*left, *right), + (Cell::Timestamptz(left), Cell::Timestamptz(right)) => { + (*left).into_inner().cmp(&(*right).into_inner()) + } + _ => { + return Err(aggregate_error(format!( + "cannot compare source cell {left:?} with qualifier or aggregate cell {right:?}" + ))); + } + }; + Ok(ordering) +} + +fn integer_cell_as_i128(cell: &Cell) -> Option { + match cell { + Cell::I16(value) => Some(i128::from(*value)), + Cell::I32(value) => Some(i128::from(*value)), + Cell::I64(value) => Some(i128::from(*value)), + _ => None, + } +} + +pub(crate) struct AggregateReducer { + slots: Vec, +} + +struct AggregateSlot { + alias: String, + state: AggregateState, +} + +enum AggregateState { + Count { + count: i64, + nonnull_only: bool, + input_oid: Option, + }, + Min { + input_oid: pg_sys::Oid, + value: Option, + }, + Max { + input_oid: pg_sys::Oid, + value: Option, + }, + SumI64 { + input_oid: pg_sys::Oid, + value: Option, + }, + SumI128(Option), + SumF32(Option), + SumF64(Option), + AvgI128 { + input_oid: pg_sys::Oid, + sum: i128, + count: i64, + }, + AvgF64 { + input_oid: pg_sys::Oid, + sum: Option, + count: i64, + }, +} + +impl AggregateReducer { + pub(crate) fn new(aggregates: &[Aggregate]) -> ZarrFdwResult { + let mut slots = Vec::new(); + slots.try_reserve_exact(aggregates.len()).map_err(|_| { + aggregate_error(format!( + "could not allocate state for {} aggregate expressions", + aggregates.len() + )) + })?; + + for aggregate in aggregates { + if !aggregate_signature_supported(aggregate) { + return Err(aggregate_error(format!( + "unsupported aggregate signature: {} -> PostgreSQL type OID {}", + aggregate.deparse(), + aggregate.type_oid + ))); + } + + let input_oid = aggregate.column.as_ref().map(|column| column.type_oid); + let state = match aggregate.kind { + AggregateKind::Count => AggregateState::Count { + count: 0, + nonnull_only: false, + input_oid: None, + }, + AggregateKind::CountColumn => AggregateState::Count { + count: 0, + nonnull_only: true, + input_oid, + }, + AggregateKind::Min => AggregateState::Min { + input_oid: input_oid.expect("supported MIN has an input column"), + value: None, + }, + AggregateKind::Max => AggregateState::Max { + input_oid: input_oid.expect("supported MAX has an input column"), + value: None, + }, + AggregateKind::Sum => match input_oid { + Some(pg_sys::INT2OID | pg_sys::INT4OID) => AggregateState::SumI64 { + input_oid: input_oid.expect("matched an integer input OID"), + value: None, + }, + Some(pg_sys::INT8OID) => AggregateState::SumI128(None), + Some(pg_sys::FLOAT4OID) => AggregateState::SumF32(None), + Some(pg_sys::FLOAT8OID) => AggregateState::SumF64(None), + _ => unreachable!("aggregate signature was checked above"), + }, + AggregateKind::Avg => match input_oid { + Some(oid @ (pg_sys::INT2OID | pg_sys::INT4OID | pg_sys::INT8OID)) => { + AggregateState::AvgI128 { + input_oid: oid, + sum: 0, + count: 0, + } + } + Some(oid @ (pg_sys::FLOAT4OID | pg_sys::FLOAT8OID)) => AggregateState::AvgF64 { + input_oid: oid, + sum: None, + count: 0, + }, + _ => unreachable!("aggregate signature was checked above"), + }, + }; + slots.push(AggregateSlot { + alias: aggregate.alias.clone(), + state, + }); + } + + Ok(Self { slots }) + } + + pub(crate) fn observe(&mut self, values: &[Option<&Cell>]) -> ZarrFdwResult<()> { + if values.len() != self.slots.len() { + return Err(aggregate_error(format!( + "received {} source values for {} aggregate expressions", + values.len(), + self.slots.len() + ))); + } + + for (slot, value) in self.slots.iter_mut().zip(values) { + slot.state.observe(*value)?; + } + Ok(()) + } + + pub(crate) fn finish(self) -> ZarrFdwResult)>> { + self.slots + .into_iter() + .map(|slot| Ok((slot.alias, slot.state.finish()?))) + .collect() + } +} + +impl AggregateState { + fn observe(&mut self, value: Option<&Cell>) -> ZarrFdwResult<()> { + match self { + Self::Count { + count, + nonnull_only, + input_oid, + } => { + if *nonnull_only && value.is_none() { + return Ok(()); + } + if let (Some(oid), Some(cell)) = (*input_oid, value) { + require_cell_oid(cell, oid)?; + } + *count = count + .checked_add(1) + .ok_or_else(|| aggregate_error("COUNT overflowed bigint"))?; + } + Self::Min { + input_oid, + value: min, + } => { + let Some(cell) = value else { return Ok(()) }; + require_cell_oid(cell, *input_oid)?; + let replace = match min.as_ref() { + Some(current) => compare_cells(cell, current)? != Ordering::Greater, + None => true, + }; + if replace { + *min = Some(cell.clone()); + } + } + Self::Max { + input_oid, + value: max, + } => { + let Some(cell) = value else { return Ok(()) }; + require_cell_oid(cell, *input_oid)?; + let replace = match max.as_ref() { + Some(current) => compare_cells(cell, current)? != Ordering::Less, + None => true, + }; + if replace { + *max = Some(cell.clone()); + } + } + Self::SumI64 { + input_oid, + value: sum, + } => { + let Some(cell) = value else { return Ok(()) }; + let next = integer_cell(cell, *input_oid)?; + *sum = Some(match *sum { + Some(current) => current + .checked_add(next) + .ok_or_else(|| aggregate_error("SUM overflowed bigint"))?, + None => next, + }); + } + Self::SumI128(sum) => { + let Some(cell) = value else { return Ok(()) }; + let Cell::I64(next) = cell else { + return Err(cell_type_error(cell, pg_sys::INT8OID)); + }; + *sum = Some(match *sum { + Some(current) => current + .checked_add(i128::from(*next)) + .ok_or_else(|| aggregate_error("integer SUM accumulator overflowed"))?, + None => i128::from(*next), + }); + } + Self::SumF32(sum) => { + let Some(cell) = value else { return Ok(()) }; + let Cell::F32(next) = cell else { + return Err(cell_type_error(cell, pg_sys::FLOAT4OID)); + }; + *sum = Some(match *sum { + Some(current) => checked_float_add_f32(current, *next)?, + None => *next, + }); + } + Self::SumF64(sum) => { + let Some(cell) = value else { return Ok(()) }; + let Cell::F64(next) = cell else { + return Err(cell_type_error(cell, pg_sys::FLOAT8OID)); + }; + *sum = Some(match *sum { + Some(current) => checked_float_add_f64(current, *next)?, + None => *next, + }); + } + Self::AvgI128 { + input_oid, + sum, + count, + } => { + let Some(cell) = value else { return Ok(()) }; + let next = i128::from(integer_cell(cell, *input_oid)?); + *sum = sum + .checked_add(next) + .ok_or_else(|| aggregate_error("integer AVG accumulator overflowed"))?; + *count = count + .checked_add(1) + .ok_or_else(|| aggregate_error("AVG count overflowed bigint"))?; + } + Self::AvgF64 { + input_oid, + sum, + count, + } => { + let Some(cell) = value else { return Ok(()) }; + let next = match (*input_oid, cell) { + (pg_sys::FLOAT4OID, Cell::F32(value)) => f64::from(*value), + (pg_sys::FLOAT8OID, Cell::F64(value)) => *value, + _ => return Err(cell_type_error(cell, *input_oid)), + }; + *sum = Some(match *sum { + Some(current) => checked_float_add_f64(current, next)?, + None => next, + }); + *count = count + .checked_add(1) + .ok_or_else(|| aggregate_error("AVG count overflowed bigint"))?; + } + } + Ok(()) + } + + fn finish(self) -> ZarrFdwResult> { + match self { + Self::Count { count, .. } => Ok(Some(Cell::I64(count))), + Self::Min { value, .. } | Self::Max { value, .. } => Ok(value), + Self::SumI64 { value, .. } => Ok(value.map(Cell::I64)), + Self::SumI128(value) => value + .map(numeric_from_i128) + .transpose() + .map(|value| value.map(Cell::Numeric)), + Self::SumF32(value) => Ok(value.map(Cell::F32)), + Self::SumF64(value) => Ok(value.map(Cell::F64)), + Self::AvgI128 { sum, count, .. } => { + if count == 0 { + return Ok(None); + } + let numerator = numeric_from_i128(sum)?; + let denominator = AnyNumeric::from(count); + Ok(Some(Cell::Numeric(numerator / denominator))) + } + Self::AvgF64 { sum, count, .. } => Ok(sum.map(|sum| Cell::F64(sum / count as f64))), + } + } +} + +fn require_cell_oid(cell: &Cell, oid: pg_sys::Oid) -> ZarrFdwResult<()> { + if cell_matches_oid(cell, oid) { + Ok(()) + } else { + Err(cell_type_error(cell, oid)) + } +} + +fn cell_matches_oid(cell: &Cell, oid: pg_sys::Oid) -> bool { + matches!( + (oid, cell), + (pg_sys::CHAROID, Cell::I8(_)) + | (pg_sys::INT2OID, Cell::I16(_)) + | (pg_sys::INT4OID, Cell::I32(_)) + | (pg_sys::INT8OID, Cell::I64(_)) + | (pg_sys::FLOAT4OID, Cell::F32(_)) + | (pg_sys::FLOAT8OID, Cell::F64(_)) + | (pg_sys::TIMESTAMPTZOID, Cell::Timestamptz(_)) + ) +} + +fn integer_cell(cell: &Cell, oid: pg_sys::Oid) -> ZarrFdwResult { + match (oid, cell) { + (pg_sys::INT2OID, Cell::I16(value)) => Ok(i64::from(*value)), + (pg_sys::INT4OID, Cell::I32(value)) => Ok(i64::from(*value)), + (pg_sys::INT8OID, Cell::I64(value)) => Ok(*value), + _ => Err(cell_type_error(cell, oid)), + } +} + +fn checked_float_add_f32(left: f32, right: f32) -> ZarrFdwResult { + let sum = left + right; + if left.is_finite() && right.is_finite() && sum.is_infinite() { + Err(aggregate_error("real SUM/AVG accumulator overflowed")) + } else { + Ok(sum) + } +} + +pub(crate) fn checked_float_add_f64(left: f64, right: f64) -> ZarrFdwResult { + let sum = left + right; + if left.is_finite() && right.is_finite() && sum.is_infinite() { + Err(aggregate_error( + "double precision SUM/AVG accumulator overflowed", + )) + } else { + Ok(sum) + } +} + +fn numeric_from_i128(value: i128) -> ZarrFdwResult { + Ok(AnyNumeric::try_from(value.to_string().as_str())?) +} + +fn aggregate_error(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!("aggregate pushdown: {}", message.into())) +} + +fn cell_type_error(cell: &Cell, oid: pg_sys::Oid) -> ZarrFdwError { + aggregate_error(format!( + "source cell {cell:?} does not match PostgreSQL input type OID {oid}" + )) +} + +#[cfg(test)] +mod unit_tests { + use super::*; + use supabase_wrappers::prelude::Column; + + fn aggregate( + kind: AggregateKind, + input_oid: Option, + type_oid: pg_sys::Oid, + alias: &str, + ) -> Aggregate { + Aggregate { + kind, + column: input_oid.map(|type_oid| Column { + name: "value".to_owned(), + num: 1, + type_oid, + }), + distinct: false, + alias: alias.to_owned(), + type_oid, + } + } + + fn qual(operator: &str, value: Value, use_or: bool) -> Qual { + Qual { + field: "value".to_owned(), + operator: operator.to_owned(), + value, + use_or, + param: None, + } + } + + #[test] + fn signature_matrix_matches_postgres_result_types() { + let supported = [ + aggregate(AggregateKind::Count, None, pg_sys::INT8OID, "c"), + aggregate( + AggregateKind::CountColumn, + Some(pg_sys::FLOAT8OID), + pg_sys::INT8OID, + "c", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::INT2OID), + pg_sys::INT8OID, + "s", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::INT4OID), + pg_sys::INT8OID, + "s", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::INT8OID), + pg_sys::NUMERICOID, + "s", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::FLOAT4OID), + pg_sys::FLOAT4OID, + "s", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::FLOAT8OID), + pg_sys::FLOAT8OID, + "s", + ), + aggregate( + AggregateKind::Avg, + Some(pg_sys::INT8OID), + pg_sys::NUMERICOID, + "a", + ), + aggregate( + AggregateKind::Avg, + Some(pg_sys::FLOAT4OID), + pg_sys::FLOAT8OID, + "a", + ), + aggregate( + AggregateKind::Min, + Some(pg_sys::CHAROID), + pg_sys::CHAROID, + "lo", + ), + aggregate( + AggregateKind::Max, + Some(pg_sys::TIMESTAMPTZOID), + pg_sys::TIMESTAMPTZOID, + "hi", + ), + ]; + assert!(supported.iter().all(aggregate_signature_supported)); + + let mut distinct = supported[1].clone(); + distinct.distinct = true; + assert!(!aggregate_signature_supported(&distinct)); + assert!(!aggregate_signature_supported(&aggregate( + AggregateKind::Sum, + Some(pg_sys::INT8OID), + pg_sys::FLOAT8OID, + "bad", + ))); + assert!(!aggregate_signature_supported(&aggregate( + AggregateKind::Avg, + Some(pg_sys::CHAROID), + pg_sys::NUMERICOID, + "bad", + ))); + } + + #[test] + fn qualifier_matching_follows_sql_null_in_and_nan_semantics() { + let is_null = qual("is", Value::Cell(Cell::String("null".to_owned())), false); + let is_not_null = qual( + "is not", + Value::Cell(Cell::String("null".to_owned())), + false, + ); + assert!(qual_matches(&is_null, None).unwrap()); + assert!(!qual_matches(&is_null, Some(&Cell::F64(1.0))).unwrap()); + assert!(!qual_matches(&is_not_null, None).unwrap()); + + let in_values = qual( + "=", + Value::Array(vec![Cell::F64(1.0), Cell::F64(f64::NAN)]), + true, + ); + assert!(qual_matches(&in_values, Some(&Cell::F64(f64::NAN))).unwrap()); + assert!(!qual_matches(&in_values, Some(&Cell::F64(2.0))).unwrap()); + assert!(!qual_matches(&in_values, None).unwrap()); + + let greater = qual(">", Value::Cell(Cell::F64(10.0)), false); + assert!(qual_matches(&greater, Some(&Cell::F64(f64::NAN))).unwrap()); + assert!(qual_matches(&greater, Some(&Cell::F64(f64::INFINITY))).unwrap()); + assert!(!qual_matches(&greater, Some(&Cell::F64(10.0))).unwrap()); + assert!(qual_matches(&greater, Some(&Cell::F64(11.0))).unwrap()); + + let unsupported = qual("~~", Value::Cell(Cell::String("%".to_owned())), false); + assert!(!qual_shape_supported(&unsupported)); + assert!(qual_matches(&unsupported, Some(&Cell::F64(1.0))).is_err()); + + let float_cross_type = qual("=", Value::Cell(Cell::F64(-7.5)), false); + assert!(qual_matches(&float_cross_type, Some(&Cell::F32(-7.5))).unwrap()); + + let integer_cross_type = qual(">", Value::Cell(Cell::I64(32_000)), false); + assert!(qual_matches(&integer_cross_type, Some(&Cell::I16(32_001))).unwrap()); + } + + #[test] + fn reducer_preserves_order_aliases_nulls_and_float_ordering() { + let aggregates = [ + aggregate(AggregateKind::Count, None, pg_sys::INT8OID, "all"), + aggregate( + AggregateKind::CountColumn, + Some(pg_sys::FLOAT8OID), + pg_sys::INT8OID, + "present", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::FLOAT8OID), + pg_sys::FLOAT8OID, + "sum", + ), + aggregate( + AggregateKind::Min, + Some(pg_sys::FLOAT8OID), + pg_sys::FLOAT8OID, + "min", + ), + aggregate( + AggregateKind::Max, + Some(pg_sys::FLOAT8OID), + pg_sys::FLOAT8OID, + "max", + ), + ]; + let mut reducer = AggregateReducer::new(&aggregates).unwrap(); + reducer.observe(&[None, None, None, None, None]).unwrap(); + let one = Cell::F64(-0.0); + reducer + .observe(&[None, Some(&one), Some(&one), Some(&one), Some(&one)]) + .unwrap(); + let nan = Cell::F64(f64::NAN); + reducer + .observe(&[None, Some(&nan), Some(&nan), Some(&nan), Some(&nan)]) + .unwrap(); + + let values = reducer.finish().unwrap(); + assert_eq!( + values + .iter() + .map(|(alias, _)| alias.as_str()) + .collect::>(), + ["all", "present", "sum", "min", "max"] + ); + assert!(matches!(values[0].1, Some(Cell::I64(3)))); + assert!(matches!(values[1].1, Some(Cell::I64(2)))); + assert!(matches!(values[2].1, Some(Cell::F64(value)) if value.is_nan())); + assert!(matches!(values[3].1, Some(Cell::F64(value)) if value == 0.0)); + assert!(matches!(values[4].1, Some(Cell::F64(value)) if value.is_nan())); + } + + #[test] + fn empty_and_all_null_inputs_return_postgres_results() { + let aggregates = [ + aggregate(AggregateKind::Count, None, pg_sys::INT8OID, "all"), + aggregate( + AggregateKind::CountColumn, + Some(pg_sys::INT4OID), + pg_sys::INT8OID, + "present", + ), + aggregate( + AggregateKind::Sum, + Some(pg_sys::INT4OID), + pg_sys::INT8OID, + "sum", + ), + aggregate( + AggregateKind::Avg, + Some(pg_sys::INT4OID), + pg_sys::NUMERICOID, + "avg", + ), + aggregate( + AggregateKind::Min, + Some(pg_sys::INT4OID), + pg_sys::INT4OID, + "min", + ), + ]; + let mut reducer = AggregateReducer::new(&aggregates).unwrap(); + reducer.observe(&[None, None, None, None, None]).unwrap(); + let values = reducer.finish().unwrap(); + assert!(matches!(values[0].1, Some(Cell::I64(1)))); + assert!(matches!(values[1].1, Some(Cell::I64(0)))); + assert!(values[2..].iter().all(|(_, value)| value.is_none())); + + let empty = AggregateReducer::new(&aggregates) + .unwrap() + .finish() + .unwrap(); + assert!(matches!(empty[0].1, Some(Cell::I64(0)))); + assert!(matches!(empty[1].1, Some(Cell::I64(0)))); + assert!(empty[2..].iter().all(|(_, value)| value.is_none())); + } + + #[test] + fn checked_accumulators_report_overflow() { + let mut count = AggregateState::Count { + count: i64::MAX, + nonnull_only: false, + input_oid: None, + }; + assert!(count.observe(None).is_err()); + + let mut sum = AggregateState::SumI64 { + input_oid: pg_sys::INT4OID, + value: Some(i64::MAX), + }; + assert!(sum.observe(Some(&Cell::I32(1))).is_err()); + + let mut float_sum = AggregateState::SumF64(Some(f64::MAX)); + assert!(float_sum.observe(Some(&Cell::F64(f64::MAX))).is_err()); + let mut infinity = AggregateState::SumF64(Some(f64::INFINITY)); + assert!(infinity.observe(Some(&Cell::F64(1.0))).is_ok()); + } +} + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use super::*; + use pgrx::pg_test; + use supabase_wrappers::prelude::Column; + + fn float_min_max(oid: pg_sys::Oid, values: &[Cell]) -> (Cell, Cell) { + let aggregates = [ + Aggregate { + kind: AggregateKind::Min, + column: Some(Column { + name: "value".to_owned(), + num: 1, + type_oid: oid, + }), + distinct: false, + alias: "min".to_owned(), + type_oid: oid, + }, + Aggregate { + kind: AggregateKind::Max, + column: Some(Column { + name: "value".to_owned(), + num: 1, + type_oid: oid, + }), + distinct: false, + alias: "max".to_owned(), + type_oid: oid, + }, + ]; + let mut reducer = AggregateReducer::new(&aggregates).unwrap(); + for value in values { + reducer.observe(&[Some(value), Some(value)]).unwrap(); + } + let values = reducer.finish().unwrap(); + (values[0].1.clone().unwrap(), values[1].1.clone().unwrap()) + } + + #[pg_test] + fn integer_numeric_sum_and_average_are_exact() { + let aggregates = [ + Aggregate { + kind: AggregateKind::Sum, + column: Some(Column { + name: "value".to_owned(), + num: 1, + type_oid: pg_sys::INT8OID, + }), + distinct: false, + alias: "sum".to_owned(), + type_oid: pg_sys::NUMERICOID, + }, + Aggregate { + kind: AggregateKind::Avg, + column: Some(Column { + name: "value".to_owned(), + num: 1, + type_oid: pg_sys::INT8OID, + }), + distinct: false, + alias: "avg".to_owned(), + type_oid: pg_sys::NUMERICOID, + }, + ]; + let mut reducer = AggregateReducer::new(&aggregates).unwrap(); + let first = Cell::I64(9_007_199_254_740_993); + let second = Cell::I64(9_007_199_254_740_994); + reducer.observe(&[Some(&first), Some(&first)]).unwrap(); + reducer.observe(&[Some(&second), Some(&second)]).unwrap(); + let values = reducer.finish().unwrap(); + + let Some(Cell::Numeric(sum)) = &values[0].1 else { + panic!("SUM(bigint) should return numeric") + }; + let Some(Cell::Numeric(avg)) = &values[1].1 else { + panic!("AVG(bigint) should return numeric") + }; + assert_eq!(sum, &AnyNumeric::try_from("18014398509481987").unwrap()); + assert_eq!(avg, &AnyNumeric::try_from("9007199254740993.5").unwrap()); + } + + #[pg_test] + fn float_min_max_signed_zero_matches_postgres() { + for (sql, values) in [ + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['0'::real, '-0'::real]) AS t(v)", + vec![Cell::F32(0.0), Cell::F32(-0.0)], + ), + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['-0'::real, '0'::real]) AS t(v)", + vec![Cell::F32(-0.0), Cell::F32(0.0)], + ), + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['NaN'::real, 'Infinity'::real, '-Infinity'::real]) AS t(v)", + vec![ + Cell::F32(f32::NAN), + Cell::F32(f32::INFINITY), + Cell::F32(f32::NEG_INFINITY), + ], + ), + ] { + let (pg_min, pg_max) = pgrx::Spi::connect(|client| { + client + .select(sql, None, &[]) + .unwrap() + .first() + .get_two::() + .unwrap() + }); + let (Cell::F32(fdw_min), Cell::F32(fdw_max)) = + float_min_max(pg_sys::FLOAT4OID, &values) + else { + panic!("real MIN/MAX should return real cells") + }; + assert_eq!(fdw_min.to_bits(), pg_min.unwrap().to_bits()); + assert_eq!(fdw_max.to_bits(), pg_max.unwrap().to_bits()); + } + + for (sql, values) in [ + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['0'::double precision, '-0'::double precision]) AS t(v)", + vec![Cell::F64(0.0), Cell::F64(-0.0)], + ), + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['-0'::double precision, '0'::double precision]) AS t(v)", + vec![Cell::F64(-0.0), Cell::F64(0.0)], + ), + ( + "SELECT min(v), max(v) FROM unnest(ARRAY['NaN'::double precision, 'Infinity'::double precision, '-Infinity'::double precision]) AS t(v)", + vec![ + Cell::F64(f64::NAN), + Cell::F64(f64::INFINITY), + Cell::F64(f64::NEG_INFINITY), + ], + ), + ] { + let (pg_min, pg_max) = pgrx::Spi::connect(|client| { + client + .select(sql, None, &[]) + .unwrap() + .first() + .get_two::() + .unwrap() + }); + let (Cell::F64(fdw_min), Cell::F64(fdw_max)) = + float_min_max(pg_sys::FLOAT8OID, &values) + else { + panic!("double precision MIN/MAX should return double precision cells") + }; + assert_eq!(fdw_min.to_bits(), pg_min.unwrap().to_bits()); + assert_eq!(fdw_max.to_bits(), pg_max.unwrap().to_bits()); + } + } +} diff --git a/wrappers/src/fdw/zarr_fdw/cache.rs b/wrappers/src/fdw/zarr_fdw/cache.rs new file mode 100644 index 000000000..c9786b98b --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/cache.rs @@ -0,0 +1,231 @@ +//! Query-local cache for complete bounded Zarr storage reads. + +use lru::LruCache; +use std::sync::Arc; + +use super::store::ReadIdentity; + +/// A complete whole-object or exact-range response cached before decoding. +/// +/// Missing objects are cached explicitly because a sparse Zarr chunk uses the +/// array's fill-value semantics. Other storage errors must never enter the +/// cache. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum CachedObject { + Present(Arc<[u8]>), + Missing, +} + +impl CachedObject { + fn resident_bytes(&self) -> usize { + match self { + Self::Present(bytes) => bytes.len(), + Self::Missing => 0, + } + } +} + +/// Byte- and entry-bounded LRU owned by one PostgreSQL scan execution. +/// +/// The cache is deliberately not global: cached bytes must not cross server, +/// role, credential, or query-lifecycle boundaries. A zero byte or entry +/// limit disables it without requiring a separate optional field at call +/// sites. +pub(crate) struct CompressedChunkCache { + entries: LruCache, + resident_bytes: usize, + max_bytes: usize, + max_entries: usize, + evictions: usize, +} + +impl CompressedChunkCache { + pub(crate) fn new(max_bytes: usize, max_entries: usize) -> Self { + Self { + entries: LruCache::unbounded(), + resident_bytes: 0, + max_bytes, + max_entries, + evictions: 0, + } + } + + #[cfg(test)] + pub(crate) fn get(&mut self, key: &str) -> Option { + self.get_identity(&ReadIdentity::whole(key)) + } + + /// Look up bytes by the complete key/range/generation identity. This + /// prevents a whole object, shard index, and inner payload from aliasing. + pub(crate) fn get_identity(&mut self, identity: &ReadIdentity) -> Option { + self.entries.get(identity).cloned() + } + + /// Insert a complete encoded object. Returns `false` when caching is + /// disabled or the object is larger than the entire byte budget. + #[cfg(test)] + pub(crate) fn insert_present(&mut self, key: String, bytes: Arc<[u8]>) -> bool { + self.insert_present_identity(ReadIdentity::whole(key), bytes) + } + + pub(crate) fn insert_present_identity( + &mut self, + identity: ReadIdentity, + bytes: Arc<[u8]>, + ) -> bool { + self.insert(identity, CachedObject::Present(bytes)) + } + + /// Cache an explicit object-not-found response. + #[cfg(test)] + pub(crate) fn insert_missing(&mut self, key: String) -> bool { + self.insert_missing_identity(ReadIdentity::whole(key)) + } + + pub(crate) fn insert_missing_identity(&mut self, identity: ReadIdentity) -> bool { + self.insert(identity, CachedObject::Missing) + } + + #[cfg(test)] + pub(crate) fn clear(&mut self) { + self.entries.clear(); + self.resident_bytes = 0; + } + + pub(crate) fn len(&self) -> usize { + self.entries.len() + } + + #[cfg(test)] + pub(crate) fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + pub(crate) fn resident_bytes(&self) -> usize { + self.resident_bytes + } + + pub(crate) fn evictions(&self) -> usize { + self.evictions + } + + fn insert(&mut self, identity: ReadIdentity, object: CachedObject) -> bool { + if let Some(previous) = self.entries.pop(&identity) { + self.resident_bytes = self + .resident_bytes + .saturating_sub(previous.resident_bytes()); + } + + let object_bytes = object.resident_bytes(); + if self.max_bytes == 0 || self.max_entries == 0 || object_bytes > self.max_bytes { + return false; + } + + while self.entries.len() >= self.max_entries + || self + .resident_bytes + .checked_add(object_bytes) + .is_none_or(|next| next > self.max_bytes) + { + let Some((_key, evicted)) = self.entries.pop_lru() else { + break; + }; + self.resident_bytes = self.resident_bytes.saturating_sub(evicted.resident_bytes()); + self.evictions = self.evictions.saturating_add(1); + } + + self.resident_bytes += object_bytes; + self.entries.put(identity, object); + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn bytes(value: u8, len: usize) -> Arc<[u8]> { + Arc::from(vec![value; len]) + } + + #[test] + fn weighted_lru_evicts_by_bytes_and_promotes_hits() { + let mut cache = CompressedChunkCache::new(6, 3); + assert!(cache.insert_present("a".to_string(), bytes(1, 3))); + assert!(cache.insert_present("b".to_string(), bytes(2, 3))); + + assert_eq!(cache.get("a"), Some(CachedObject::Present(bytes(1, 3)))); + assert!(cache.insert_present("c".to_string(), bytes(3, 3))); + + assert!(cache.get("b").is_none()); + assert!(cache.get("a").is_some()); + assert!(cache.get("c").is_some()); + assert_eq!(cache.resident_bytes(), 6); + } + + #[test] + fn entry_limit_bounds_zero_weight_missing_objects() { + let mut cache = CompressedChunkCache::new(16, 2); + assert!(cache.insert_missing("a".to_string())); + assert!(cache.insert_missing("b".to_string())); + assert!(cache.insert_missing("c".to_string())); + + assert!(cache.get("a").is_none()); + assert_eq!(cache.get("b"), Some(CachedObject::Missing)); + assert_eq!(cache.get("c"), Some(CachedObject::Missing)); + assert_eq!(cache.len(), 2); + assert_eq!(cache.resident_bytes(), 0); + } + + #[test] + fn oversized_replacement_removes_a_stale_value_and_bypasses_cache() { + let mut cache = CompressedChunkCache::new(4, 2); + assert!(cache.insert_present("chunk".to_string(), bytes(1, 4))); + assert!(!cache.insert_present("chunk".to_string(), bytes(2, 5))); + assert!(cache.get("chunk").is_none()); + assert_eq!(cache.resident_bytes(), 0); + } + + #[test] + fn zero_limit_disables_and_clear_resets_accounting() { + let mut disabled = CompressedChunkCache::new(0, 8); + assert!(!disabled.insert_missing("missing".to_string())); + assert_eq!(disabled.len(), 0); + + let mut cache = CompressedChunkCache::new(8, 2); + assert!(cache.insert_present("chunk".to_string(), bytes(1, 8))); + cache.clear(); + assert!(cache.is_empty()); + assert_eq!(cache.resident_bytes(), 0); + } + + #[test] + fn whole_ranges_and_generations_do_not_alias() { + let mut cache = CompressedChunkCache::new(32, 8); + let whole = ReadIdentity::whole("shard"); + let exact = ReadIdentity::exact("shard", 0, 4).unwrap(); + let generation = super::super::store::ObjectGeneration::S3 { + etag: "etag-a".to_string(), + version_id: None, + total_len: 64, + }; + let exact_generation = exact.clone().with_generation(generation); + + assert!(cache.insert_present_identity(whole.clone(), bytes(1, 4))); + assert!(cache.insert_present_identity(exact.clone(), bytes(2, 4))); + assert!(cache.insert_present_identity(exact_generation.clone(), bytes(3, 4))); + + assert_eq!( + cache.get_identity(&whole), + Some(CachedObject::Present(bytes(1, 4))) + ); + assert_eq!( + cache.get_identity(&exact), + Some(CachedObject::Present(bytes(2, 4))) + ); + assert_eq!( + cache.get_identity(&exact_generation), + Some(CachedObject::Present(bytes(3, 4))) + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/chunk.rs b/wrappers/src/fdw/zarr_fdw/chunk.rs new file mode 100644 index 000000000..146f55130 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/chunk.rs @@ -0,0 +1,436 @@ +//! Pure chunk-selection math for querying Zarr cubes. +//! +//! The core idea (borrowed from `duckdb_zarr`): translate a predicate on a +//! dimension (time, x, y) into a *chunk index range* using that dimension's +//! coordinate vector, so that only the chunks that can contain matching cells +//! are fetched from the object store. All functions here are pure and unit +//! testable without any object store. + +use super::meta::{ArrayMeta, ChunkKeyEncoding}; +use super::{ZarrFdwError, ZarrFdwResult}; + +/// Inclusive 1-based? No — plain 0-based inclusive index bounds `(start, end)`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IndexBounds { + pub start: usize, + pub end: usize, // inclusive +} + +impl IndexBounds { + /// Full bounds for an axis of `len` elements. + pub fn full(len: usize) -> Self { + Self { + start: 0, + end: len.saturating_sub(1), + } + } + + pub fn is_empty(&self) -> bool { + self.end < self.start + } +} + +/// Translate a `[lo, hi]` value range into inclusive index bounds over a +/// coordinate vector. The coordinate vector is assumed to be sorted (either +/// ascending or descending); direction is detected from the endpoint values. +/// +/// If `lo`/`hi` are `None`, the bound is unbounded on that side. +/// Returns `None` if the range misses the coordinate vector entirely (no +/// overlap), which callers use to skip the axis/chunk entirely. +pub fn index_bounds_from_value_range( + coords: &[f64], + lo: Option, + hi: Option, +) -> Option { + if coords.is_empty() { + return Some(IndexBounds::full(0)); + } + if let (Some(lo), Some(hi)) = (lo, hi) { + // quick reject: the [lo, hi] window nowhere overlaps the coord span + let (cmin, cmax) = if coords[0] <= coords[coords.len() - 1] { + (coords[0], coords[coords.len() - 1]) + } else { + (coords[coords.len() - 1], coords[0]) + }; + if hi < cmin || lo > cmax { + return None; + } + } + + let ascending = match (coords.first(), coords.last()) { + (Some(&first), Some(&last)) => last >= first, + _ => true, + }; + + let n = coords.len(); + let bounds = if ascending { + // find first index with coord >= lo (or 0), last index with coord <= hi (or n-1) + let start = match lo { + Some(lo) => coords.partition_point(|&c| c < lo), + None => 0, + }; + let end = match hi { + Some(hi) => { + // first index with coord > hi, then step back + let mut p = coords.partition_point(|&c| c <= hi); + p = p.saturating_sub(1); + p + } + None => n - 1, + }; + IndexBounds { start, end } + } else { + // descending: index i holds value descending. value >= lo and <= hi + // correspond to the first index where value drops below hi, up to the + // last index where value is still >= lo. + let start = match hi { + Some(hi) => coords.partition_point(|&c| c > hi), + None => 0, + }; + let end = match lo { + Some(lo) => { + let mut p = coords.partition_point(|&c| c >= lo); + p = p.saturating_sub(1); + p + } + None => n - 1, + }; + IndexBounds { start, end } + }; + + if bounds.is_empty() { + return None; + } + Some(bounds) +} + +impl IndexBounds { + /// Convert index bounds into an inclusive chunk range for an axis: which + /// chunk numbers (0-based) cover `[start, end]`. + pub fn chunk_range(&self, chunk_len: usize) -> ZarrFdwResult<(usize, usize)> { + if chunk_len == 0 { + return Err(ZarrFdwError::InvalidMetadata( + "chunk length must be greater than zero".to_string(), + )); + } + let first = self.start / chunk_len; + let last = self.end / chunk_len; + Ok((first, last)) + } +} + +/// Build the storage key for one logical chunk coordinate. +pub fn chunk_key(encoding: &ChunkKeyEncoding, indices: &[u64]) -> String { + let separator = match encoding { + ChunkKeyEncoding::Default { separator } | ChunkKeyEncoding::V2 { separator } => separator, + }; + let coordinates = indices + .iter() + .map(|i| i.to_string()) + .collect::>() + .join(&separator.to_string()); + match encoding { + ChunkKeyEncoding::Default { .. } if coordinates.is_empty() => "c".to_string(), + ChunkKeyEncoding::Default { .. } => format!("c{separator}{coordinates}"), + ChunkKeyEncoding::V2 { .. } if coordinates.is_empty() => "0".to_string(), + ChunkKeyEncoding::V2 { .. } => coordinates, + } +} + +/// Lazy Cartesian product of per-axis chunk index ranges. +/// +/// The cursor stores only rank-sized state. `next_into` reuses a caller-owned +/// output vector and visits chunks in row-major (C) order, with the last axis +/// varying fastest. +#[derive(Debug, Default)] +pub struct ChunkIndexCursor { + starts: Vec, + ends: Vec, + next: Vec, + has_next: bool, +} + +impl ChunkIndexCursor { + pub fn new(axis_chunk_ranges: &[(usize, usize)]) -> ZarrFdwResult { + if axis_chunk_ranges.is_empty() || axis_chunk_ranges.iter().any(|&(start, end)| start > end) + { + return Ok(Self::default()); + } + + let starts = axis_chunk_ranges + .iter() + .map(|&(start, _)| { + u64::try_from(start).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "chunk index exceeds the Zarr v2 u64 index capacity".to_string(), + ) + }) + }) + .collect::>>()?; + let ends = axis_chunk_ranges + .iter() + .map(|&(_, end)| { + u64::try_from(end).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "chunk index exceeds the Zarr v2 u64 index capacity".to_string(), + ) + }) + }) + .collect::>>()?; + let next = starts.clone(); + + Ok(Self { + starts, + ends, + next, + has_next: true, + }) + } + + /// Copy the next chunk index into `destination`, reusing its allocation. + /// Returns `false` after the cursor is exhausted and clears `destination` + /// so callers cannot accidentally reuse a stale chunk index. + pub fn next_into(&mut self, destination: &mut Vec) -> bool { + if !self.has_next { + destination.clear(); + return false; + } + + destination.clone_from(&self.next); + for axis in (0..self.next.len()).rev() { + if self.next[axis] < self.ends[axis] { + self.next[axis] += 1; + return true; + } + self.next[axis] = self.starts[axis]; + } + + self.has_next = false; + true + } + + pub fn reset(&mut self) { + self.next.clone_from(&self.starts); + self.has_next = !self.starts.is_empty(); + } +} + +/// Compute the per-axis chunk ranges for a cube given per-axis index bounds. +pub fn axis_chunk_ranges( + meta: &ArrayMeta, + bounds: &[Option], +) -> ZarrFdwResult> { + let full = meta.chunks_per_axis(); + bounds + .iter() + .enumerate() + .map(|(axis, b)| match b { + Some(b) => b.chunk_range(meta.chunk_extent(axis)?), + None => { + let last = full[axis].checked_sub(1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("axis {axis} has no addressable chunks")) + })?; + Ok(( + 0, + usize::try_from(last).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk count for axis {axis} exceeds this platform's index capacity" + )) + })?, + )) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::super::meta::ArrayMeta; + use super::*; + + fn meta(shape: Vec, chunks: Vec) -> ArrayMeta { + ArrayMeta { + zarr_format: 2, + shape, + chunks, + dtype: " Vec { + (0..100).map(|i| i as f64 * 10.0).collect() // 0,10,...990 + } + + #[test] + fn test_index_bounds_ascending() { + let c = coords_asc(); + let b = index_bounds_from_value_range(&c, Some(20.0), Some(50.0)).unwrap(); + assert_eq!(b, IndexBounds { start: 2, end: 5 }); + } + + #[test] + fn test_index_bounds_unbounded_hi() { + let c = coords_asc(); + let b = index_bounds_from_value_range(&c, Some(990.0), None).unwrap(); + assert_eq!(b, IndexBounds { start: 99, end: 99 }); + } + + #[test] + fn test_index_bounds_no_overlap() { + let c = coords_asc(); + assert!(index_bounds_from_value_range(&c, Some(5000.0), None).is_none()); + } + + #[test] + fn test_index_bounds_descending() { + let c: Vec = (0..100).rev().map(|i| i as f64 * 10.0).collect(); // 990..0 + let b = index_bounds_from_value_range(&c, Some(20.0), Some(50.0)).unwrap(); + // values 20..50 are at original indices 94..97 + assert_eq!(b, IndexBounds { start: 94, end: 97 }); + } + + #[test] + fn test_chunk_range() { + let b = IndexBounds { start: 2, end: 5 }; + assert_eq!(b.chunk_range(3).unwrap(), (0, 1)); + let b2 = IndexBounds { start: 6, end: 8 }; + assert_eq!(b2.chunk_range(3).unwrap(), (2, 2)); + assert!(b2.chunk_range(0).is_err()); + } + + #[test] + fn test_chunk_key() { + assert_eq!( + chunk_key(&ChunkKeyEncoding::V2 { separator: '.' }, &[3, 14, 22]), + "3.14.22" + ); + assert_eq!( + chunk_key(&ChunkKeyEncoding::V2 { separator: '/' }, &[3, 14, 22]), + "3/14/22" + ); + assert_eq!( + chunk_key(&ChunkKeyEncoding::Default { separator: '/' }, &[3, 14, 22]), + "c/3/14/22" + ); + assert_eq!( + chunk_key(&ChunkKeyEncoding::Default { separator: '.' }, &[3, 14, 22]), + "c.3.14.22" + ); + assert_eq!( + chunk_key(&ChunkKeyEncoding::Default { separator: '/' }, &[]), + "c" + ); + assert_eq!( + chunk_key(&ChunkKeyEncoding::V2 { separator: '.' }, &[]), + "0" + ); + } + + fn collect_chunks(cursor: &mut ChunkIndexCursor) -> Vec> { + let mut chunks = Vec::new(); + let mut current = Vec::new(); + while cursor.next_into(&mut current) { + chunks.push(current.clone()); + } + chunks + } + + #[test] + fn chunk_cursor_is_row_major() { + let ranges = vec![(0, 1), (1, 2), (0, 0)]; + let mut cursor = ChunkIndexCursor::new(&ranges).unwrap(); + let out = collect_chunks(&mut cursor); + let keys: Vec = out + .iter() + .map(|i| chunk_key(&ChunkKeyEncoding::V2 { separator: '.' }, i)) + .collect(); + assert_eq!(keys, vec!["0.1.0", "0.2.0", "1.1.0", "1.2.0"]); + } + + #[test] + fn test_axis_chunk_ranges_from_bounds() { + let m = meta(vec![48, 100, 100], vec![4, 10, 10]); + let bounds = vec![ + Some(IndexBounds { start: 5, end: 11 }), + None, + Some(IndexBounds { start: 0, end: 99 }), + ]; + let ranges = axis_chunk_ranges(&m, &bounds).unwrap(); + assert_eq!(ranges[0], (1, 2)); // indices 5..11 -> chunks 1..2 + assert_eq!(ranges[1], (0, 9)); // full y axis + assert_eq!(ranges[2], (0, 9)); // full x axis + } + + #[test] + fn empty_and_inverted_ranges_are_exhausted() { + for ranges in [vec![], vec![(2, 1)], vec![(0, 1), (3, 2)]] { + let mut cursor = ChunkIndexCursor::new(&ranges).unwrap(); + let mut current = vec![99]; + assert!(!cursor.next_into(&mut current)); + assert!(current.is_empty()); + cursor.reset(); + assert!(!cursor.next_into(&mut current)); + } + } + + #[test] + fn cursor_reset_replays_partial_and_exhausted_ranges() { + let mut cursor = ChunkIndexCursor::new(&[(0, 1), (4, 5)]).unwrap(); + let mut current = Vec::new(); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![0, 4]); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![0, 5]); + + cursor.reset(); + assert_eq!( + collect_chunks(&mut cursor), + vec![vec![0, 4], vec![0, 5], vec![1, 4], vec![1, 5]] + ); + + cursor.reset(); + assert_eq!(collect_chunks(&mut cursor).len(), 4); + } + + #[test] + fn cursor_accepts_more_than_one_million_chunks_without_enumerating_them() { + let mut cursor = ChunkIndexCursor::new(&[(0, 1_000_000)]).unwrap(); + let mut current = Vec::new(); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![0]); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![1]); + } + + #[test] + fn cursor_supports_rank_64_with_rank_sized_state() { + let ranges = vec![(0, 0); 64]; + let mut cursor = ChunkIndexCursor::new(&ranges).unwrap(); + let mut current = Vec::new(); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![0; 64]); + assert!(!cursor.next_into(&mut current)); + assert!(current.is_empty()); + + cursor.reset(); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![0; 64]); + } + + #[test] + fn cursor_exhausts_at_the_platform_maximum_without_overflow() { + let mut cursor = ChunkIndexCursor::new(&[(usize::MAX, usize::MAX)]).unwrap(); + let mut current = Vec::new(); + assert!(cursor.next_into(&mut current)); + assert_eq!(current, vec![u64::try_from(usize::MAX).unwrap()]); + assert!(!cursor.next_into(&mut current)); + assert!(current.is_empty()); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/codec.rs b/wrappers/src/fdw/zarr_fdw/codec.rs new file mode 100644 index 000000000..584a326ea --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/codec.rs @@ -0,0 +1,2320 @@ +//! Format-neutral, ordered Zarr chunk codec pipelines. +//! +//! Zarr v2 compressors and the bounded direct Zarr v3 codec subset are +//! normalized into one execution contract. Pipelines always decode to the +//! executor's logical C-order primitive bytes. + +use std::io::Cursor; + +use serde_json::{Map, Value}; +use tokio::io::{AsyncRead, AsyncReadExt, BufReader}; + +use super::{ZarrFdwError, ZarrFdwResult}; + +const MAX_DECODED_CHUNK_BYTES: usize = 256 * 1024 * 1024; +const COMPRESSED_OVERHEAD_ALLOWANCE: usize = 1024 * 1024; +const BLOSC_HEADER_BYTES: usize = 16; +const CRC32C_BYTES: usize = 4; +const STREAM_POLL_BYTES: usize = 64 * 1024; +const CRC_POLL_BYTES: usize = 1024 * 1024; +const TRANSPOSE_POLL_CELLS: usize = 4096; +const ZSTD_WINDOW_LOG_MAX: u32 = 23; +const ZSTD_WINDOW_BYTES: u64 = 1 << ZSTD_WINDOW_LOG_MAX; +const ZSTD_FRAME_MAGIC: u32 = 0xfd2f_b528; +const ZSTD_SKIPPABLE_MAGIC_START: u32 = 0x184d_2a50; +const ZSTD_SKIPPABLE_MAGIC_MASK: u32 = 0xffff_fff0; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Endian { + Little, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum CodecStage { + Transpose { order: Vec }, + Bytes { endian: Option }, + Gzip, + Crc32c, + // Zlib remains v2-only. A present Blosc configuration identifies the + // strictly validated v3 codec; v2 deliberately retains its permissive + // compressor-metadata behavior through `None`. + Zlib, + Blosc { config: Option }, + Zstd { config: ZstdConfig }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BloscCname { + Blosclz, + Lz4, + Lz4hc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BloscShuffle { + None, + Byte, + Bit, +} + +/// The complete, validated Zarr v3 Blosc encoding configuration. These +/// parameters are retained even though c-blosc's self-describing header owns +/// decompression; retaining them keeps metadata normalization lossless and +/// permits safe header consistency checks where the binary format represents +/// the corresponding setting. +#[allow(dead_code)] // Encoding parameters are retained intentionally; decoding is self-describing. +#[derive(Debug, Clone, PartialEq, Eq)] +struct BloscConfig { + cname: BloscCname, + clevel: u8, + shuffle: BloscShuffle, + typesize: Option, + blocksize: usize, +} + +/// The complete Zarr v3 Zstd encoding configuration. Compression level is +/// retained for lossless metadata normalization even though it does not +/// affect decoding. +#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Eq)] +struct ZstdConfig { + level: i32, + checksum: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ZstdFrameHeader { + checksum: bool, + content_size: Option, +} + +impl CodecStage { + fn label(&self) -> &'static str { + match self { + Self::Transpose { .. } => "transpose", + Self::Bytes { .. } => "bytes", + Self::Gzip => "gzip", + Self::Crc32c => "crc32c", + Self::Zlib => "zlib", + Self::Blosc { .. } => "blosc", + Self::Zstd { .. } => "zstd", + } + } +} + +/// A validated codec sequence in metadata/encoding order. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CodecPipeline { + stages: Vec, +} + +/// Interruptible codec execution result. PostgreSQL's error-raising interrupt +/// handler must only be called by the executor after `Runtime::block_on` has +/// returned and its future has been dropped. +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum CodecDecode { + Decoded(Vec), + Interrupted, +} + +impl CodecPipeline { + /// A raw v2 pipeline, useful for format-neutral metadata construction. + pub(crate) fn raw_v2() -> Self { + Self { stages: Vec::new() } + } + + /// Normalize the existing Zarr v2 compressor representation without + /// tightening its accepted compressor-specific configuration. + pub(crate) fn from_v2(compressor: &Option) -> ZarrFdwResult { + let Some(compressor) = compressor else { + return Ok(Self::raw_v2()); + }; + let id = compressor + .as_object() + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "compressor must be null or a JSON object with a string 'id'".to_string(), + ) + })? + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "compressor object must contain a non-empty string 'id'".to_string(), + ) + })?; + let stage = match id { + "gzip" => CodecStage::Gzip, + "zlib" => CodecStage::Zlib, + "blosc" => CodecStage::Blosc { config: None }, + other => return Err(ZarrFdwError::UnsupportedCompressor(other.to_string())), + }; + Ok(Self { + stages: vec![stage], + }) + } + + /// Parse the council-locked direct Zarr v3 pipeline and return its + /// executor-normalized NumPy dtype. + /// + /// Supported metadata order is exactly: + /// `[transpose]? -> bytes -> [gzip | blosc | zstd]? -> [crc32c]?`. + pub(crate) fn from_v3( + native_dtype: &str, + rank: usize, + codecs: &Value, + ) -> ZarrFdwResult<(Self, String)> { + let (dtype, multi_byte) = normalize_v3_dtype(native_dtype)?; + let codecs = codecs.as_array().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("Zarr v3 codecs must be an array".to_string()) + })?; + if codecs.is_empty() { + return Err(ZarrFdwError::InvalidMetadata( + "Zarr v3 codecs must contain exactly one bytes codec".to_string(), + )); + } + + let mut stages = Vec::with_capacity(codecs.len()); + let mut cursor = 0usize; + if codec_name(codecs, cursor)? == "transpose" { + stages.push(parse_transpose( + codec_object(codecs, cursor)?, + rank, + cursor, + )?); + cursor += 1; + } + + if cursor >= codecs.len() || codec_name(codecs, cursor)? != "bytes" { + return Err(pipeline_metadata_error( + cursor.min(codecs.len().saturating_sub(1)), + "expected exactly one bytes codec after the optional transpose codec", + )); + } + stages.push(parse_bytes( + codec_object(codecs, cursor)?, + multi_byte, + cursor, + )?); + cursor += 1; + + if cursor < codecs.len() { + let stage = match codec_name(codecs, cursor)? { + "gzip" => Some(parse_gzip(codec_object(codecs, cursor)?, cursor)?), + "blosc" => Some(parse_blosc(codec_object(codecs, cursor)?, cursor)?), + "zstd" => Some(parse_zstd(codec_object(codecs, cursor)?, cursor)?), + _ => None, + }; + if let Some(stage) = stage { + stages.push(stage); + cursor += 1; + } + } + if cursor < codecs.len() && codec_name(codecs, cursor)? == "crc32c" { + stages.push(parse_crc32c(codec_object(codecs, cursor)?, cursor)?); + cursor += 1; + } + if cursor != codecs.len() { + let codec = codec_object(codecs, cursor)?; + validate_codec_object(codec, cursor)?; + let name = codec_name(codecs, cursor)?; + match name { + "transpose" | "bytes" | "gzip" | "blosc" | "zstd" | "crc32c" => { + return Err(pipeline_metadata_error( + cursor, + format!("codec '{name}' is duplicated or appears out of supported order"), + )); + } + "sharding_indexed" => { + return Err(unsupported_pipeline_feature( + cursor, + "sharded Zarr v3 chunks are not supported", + )); + } + other => { + return Err(unsupported_pipeline_feature( + cursor, + format!("Zarr v3 codec '{other}' is not supported"), + )); + } + } + } + + Ok((Self { stages }, dtype)) + } + + /// Ordered codec names for `EXPLAIN ANALYZE` and diagnostics. + pub(crate) fn ordered_label(&self) -> String { + if self.stages.is_empty() { + "raw".to_string() + } else { + self.stages + .iter() + .map(CodecStage::label) + .collect::>() + .join(" -> ") + } + } + + /// Maximum complete encoded object accepted for a declared decoded size. + /// This retains the v2 bounds and composes the exact CRC suffix allowance. + pub(crate) fn encoded_read_limit(&self, decoded_bytes: usize) -> ZarrFdwResult { + if decoded_bytes > MAX_DECODED_CHUNK_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "declared chunk decodes to {decoded_bytes} bytes, exceeding the safety limit of {MAX_DECODED_CHUNK_BYTES}" + ))); + } + self.stages.iter().try_fold(decoded_bytes, |limit, stage| { + let allowance = match stage { + CodecStage::Gzip | CodecStage::Zlib | CodecStage::Zstd { .. } => { + COMPRESSED_OVERHEAD_ALLOWANCE + } + CodecStage::Blosc { .. } => BLOSC_HEADER_BYTES, + CodecStage::Crc32c => CRC32C_BYTES, + CodecStage::Transpose { .. } | CodecStage::Bytes { .. } => 0, + }; + limit.checked_add(allowance).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "encoded chunk read limit exceeds this platform's index capacity".to_string(), + ) + }) + }) + } + + /// Decode a complete direct chunk in reverse metadata order. Every stage + /// is bounded by the checked logical layout, and long loops poll the + /// supplied non-raising cancellation callback. + pub(crate) async fn decode_interruptible( + &self, + encoded: Vec, + logical_shape: &[usize], + itemsize: usize, + mut interrupt_pending: F, + ) -> ZarrFdwResult + where + F: FnMut() -> bool, + { + let expected = checked_logical_bytes(logical_shape, itemsize)?; + let encoded_limit = self.encoded_read_limit(expected)?; + if encoded.len() > encoded_limit { + return Err(codec_read_error( + self.stages.len(), + "input", + format!( + "encoded chunk has {} bytes, exceeding its read limit of {encoded_limit}", + encoded.len() + ), + )); + } + if interrupt_pending() { + return Ok(CodecDecode::Interrupted); + } + let mut data = encoded; + for (index, stage) in self.stages.iter().enumerate().rev() { + if interrupt_pending() { + return Ok(CodecDecode::Interrupted); + } + data = match stage { + CodecStage::Crc32c => { + let Some(decoded) = decode_crc32c(data, &mut interrupt_pending, index)? else { + return Ok(CodecDecode::Interrupted); + }; + decoded + } + CodecStage::Gzip => { + let Some(decoded) = decode_stream( + StreamCodec::Gzip, + data, + expected, + &mut interrupt_pending, + index, + ) + .await? + else { + return Ok(CodecDecode::Interrupted); + }; + decoded + } + CodecStage::Zlib => { + let Some(decoded) = decode_stream( + StreamCodec::Zlib, + data, + expected, + &mut interrupt_pending, + index, + ) + .await? + else { + return Ok(CodecDecode::Interrupted); + }; + decoded + } + CodecStage::Blosc { config } => { + let decoded = decode_blosc(data, expected, index, config.as_ref())?; + if interrupt_pending() { + return Ok(CodecDecode::Interrupted); + } + decoded + } + CodecStage::Zstd { config } => { + let Some(decoded) = + decode_zstd(data, expected, config, &mut interrupt_pending, index)? + else { + return Ok(CodecDecode::Interrupted); + }; + decoded + } + CodecStage::Bytes { endian } => { + debug_assert!(endian.is_none() || *endian == Some(Endian::Little)); + require_exact_length(data, expected, index, stage.label())? + } + CodecStage::Transpose { order } => { + let Some(decoded) = inverse_transpose( + data, + logical_shape, + itemsize, + order, + &mut interrupt_pending, + index, + )? + else { + return Ok(CodecDecode::Interrupted); + }; + decoded + } + }; + if interrupt_pending() { + return Ok(CodecDecode::Interrupted); + } + } + require_exact_length(data, expected, self.stages.len(), "pipeline") + .map(CodecDecode::Decoded) + } +} + +fn normalize_v3_dtype(data_type: &str) -> ZarrFdwResult<(String, bool)> { + let normalized = match data_type { + "float32" => (" (" ("|i1", false), + "int16" => (" (" (" return Err(ZarrFdwError::UnsupportedDataType(other.to_string())), + }; + Ok((normalized.0.to_string(), normalized.1)) +} + +fn codec_object(codecs: &[Value], index: usize) -> ZarrFdwResult<&Map> { + codecs + .get(index) + .and_then(Value::as_object) + .ok_or_else(|| pipeline_metadata_error(index, "codec entry must be an object")) +} + +fn codec_name(codecs: &[Value], index: usize) -> ZarrFdwResult<&str> { + codec_object(codecs, index)? + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .ok_or_else(|| pipeline_metadata_error(index, "codec name must be a non-empty string")) +} + +fn validate_codec_object(codec: &Map, index: usize) -> ZarrFdwResult<()> { + validate_fields(codec, &["name", "configuration", "must_understand"], index)?; + if codec + .get("must_understand") + .is_some_and(|value| !value.is_boolean()) + { + return Err(pipeline_metadata_error( + index, + "must_understand must be a boolean", + )); + } + Ok(()) +} + +fn codec_configuration( + codec: &Map, + index: usize, + required: bool, +) -> ZarrFdwResult>> { + let configuration = codec.get("configuration"); + if required && configuration.is_none() { + return Err(pipeline_metadata_error( + index, + "codec configuration is required", + )); + } + configuration + .map(|value| { + value.as_object().ok_or_else(|| { + pipeline_metadata_error(index, "codec configuration must be an object") + }) + }) + .transpose() +} + +fn validate_fields( + object: &Map, + fields: &[&str], + index: usize, +) -> ZarrFdwResult<()> { + if let Some(field) = object + .keys() + .find(|field| !fields.contains(&field.as_str())) + { + return Err(pipeline_metadata_error( + index, + format!("codec configuration contains unsupported field '{field}'"), + )); + } + Ok(()) +} + +fn parse_transpose( + codec: &Map, + rank: usize, + index: usize, +) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + let configuration = codec_configuration(codec, index, true)?.expect("required above"); + validate_fields(configuration, &["order"], index)?; + let order = configuration + .get("order") + .and_then(Value::as_array) + .ok_or_else(|| pipeline_metadata_error(index, "transpose order must be an array"))? + .iter() + .map(|axis| { + axis.as_u64() + .and_then(|axis| usize::try_from(axis).ok()) + .ok_or_else(|| { + pipeline_metadata_error( + index, + "transpose order must contain non-negative platform-sized integers", + ) + }) + }) + .collect::>>()?; + validate_permutation(&order, rank, index)?; + Ok(CodecStage::Transpose { order }) +} + +fn parse_bytes( + codec: &Map, + multi_byte: bool, + index: usize, +) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + let configuration = codec_configuration(codec, index, multi_byte)?; + if let Some(configuration) = configuration { + validate_fields(configuration, &["endian"], index)?; + } + let endian = configuration + .and_then(|configuration| configuration.get("endian")) + .map(|value| { + value + .as_str() + .ok_or_else(|| pipeline_metadata_error(index, "bytes endian must be a string")) + }) + .transpose()?; + if multi_byte { + match endian { + Some("little") => Ok(CodecStage::Bytes { + endian: Some(Endian::Little), + }), + Some("big") => Err(unsupported_pipeline_feature( + index, + "big-endian Zarr v3 bytes are not supported yet", + )), + Some(other) => Err(pipeline_metadata_error( + index, + format!("bytes endian must be 'little', got '{other}'"), + )), + None => Err(pipeline_metadata_error( + index, + "bytes endian is required for multi-byte numeric data", + )), + } + } else { + match endian { + None => Ok(CodecStage::Bytes { endian: None }), + Some("little") => Ok(CodecStage::Bytes { + endian: Some(Endian::Little), + }), + Some("big") => Err(unsupported_pipeline_feature( + index, + "big-endian Zarr v3 bytes are not supported yet", + )), + Some(other) => Err(pipeline_metadata_error( + index, + format!("bytes endian must be 'little', got '{other}'"), + )), + } + } +} + +fn parse_gzip(codec: &Map, index: usize) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + let configuration = codec_configuration(codec, index, true)?.expect("required above"); + validate_fields(configuration, &["level"], index)?; + let level = configuration + .get("level") + .and_then(Value::as_u64) + .ok_or_else(|| { + pipeline_metadata_error(index, "gzip level must be an integer from 0 to 9") + })?; + if level > 9 { + return Err(pipeline_metadata_error( + index, + "gzip level must be an integer from 0 to 9", + )); + } + Ok(CodecStage::Gzip) +} + +fn parse_zstd(codec: &Map, index: usize) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + let configuration = codec_configuration(codec, index, true)?.expect("required above"); + validate_fields(configuration, &["level", "checksum"], index)?; + + let level = configuration + .get("level") + .and_then(Value::as_i64) + .filter(|&level| (-131_072..=22).contains(&level)) + .and_then(|level| i32::try_from(level).ok()) + .ok_or_else(|| { + pipeline_metadata_error(index, "Zstd level must be an integer from -131072 to 22") + })?; + let checksum = configuration + .get("checksum") + .map(|value| { + value + .as_bool() + .ok_or_else(|| pipeline_metadata_error(index, "Zstd checksum must be a boolean")) + }) + .transpose()? + .unwrap_or(false); + + Ok(CodecStage::Zstd { + config: ZstdConfig { level, checksum }, + }) +} + +fn parse_blosc(codec: &Map, index: usize) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + let configuration = codec_configuration(codec, index, true)?.expect("required above"); + validate_fields( + configuration, + &["cname", "clevel", "shuffle", "typesize", "blocksize"], + index, + )?; + + let cname = configuration + .get("cname") + .and_then(Value::as_str) + .ok_or_else(|| pipeline_metadata_error(index, "Blosc cname must be a string"))?; + let (cname, unavailable_cname) = match cname { + "blosclz" => (Some(BloscCname::Blosclz), None), + "lz4" => (Some(BloscCname::Lz4), None), + "lz4hc" => (Some(BloscCname::Lz4hc), None), + "zstd" | "snappy" | "zlib" => (None, Some(cname)), + other => { + return Err(pipeline_metadata_error( + index, + format!("Blosc cname '{other}' is not defined by the Zarr v3 Blosc codec"), + )); + } + }; + + let clevel = configuration + .get("clevel") + .and_then(Value::as_u64) + .filter(|&level| level <= 9) + .ok_or_else(|| { + pipeline_metadata_error(index, "Blosc clevel must be an integer from 0 to 9") + })? as u8; + + let shuffle = configuration + .get("shuffle") + .and_then(Value::as_str) + .ok_or_else(|| { + pipeline_metadata_error( + index, + "Blosc shuffle must be 'noshuffle', 'shuffle', or 'bitshuffle'", + ) + })?; + let shuffle = match shuffle { + "noshuffle" => BloscShuffle::None, + "shuffle" => BloscShuffle::Byte, + "bitshuffle" => BloscShuffle::Bit, + _ => { + return Err(pipeline_metadata_error( + index, + "Blosc shuffle must be 'noshuffle', 'shuffle', or 'bitshuffle'", + )); + } + }; + + let typesize = configuration + .get("typesize") + .map(|value| { + value + .as_u64() + .filter(|&value| value > 0) + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| { + pipeline_metadata_error( + index, + "Blosc typesize must be a positive platform-sized integer", + ) + }) + }) + .transpose()?; + if shuffle != BloscShuffle::None && typesize.is_none() { + return Err(pipeline_metadata_error( + index, + "Blosc typesize is required when shuffle is 'shuffle' or 'bitshuffle'", + )); + } + + let blocksize = configuration + .get("blocksize") + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| { + pipeline_metadata_error( + index, + "Blosc blocksize must be a non-negative platform-sized integer", + ) + })?; + + if let Some(cname) = unavailable_cname { + return Err(unsupported_pipeline_feature( + index, + format!("Blosc cname '{cname}' is not enabled in this build"), + )); + } + let cname = cname.expect("supported Blosc cname resolved above"); + + Ok(CodecStage::Blosc { + config: Some(BloscConfig { + cname, + clevel, + shuffle, + typesize, + blocksize, + }), + }) +} + +fn parse_crc32c(codec: &Map, index: usize) -> ZarrFdwResult { + validate_codec_object(codec, index)?; + if let Some(configuration) = codec_configuration(codec, index, false)? { + validate_fields(configuration, &[], index)?; + } + Ok(CodecStage::Crc32c) +} + +fn validate_permutation(order: &[usize], rank: usize, index: usize) -> ZarrFdwResult<()> { + if order.len() != rank { + return Err(pipeline_metadata_error( + index, + format!( + "transpose order has rank {}, expected array rank {rank}", + order.len() + ), + )); + } + let mut seen = vec![false; rank]; + for &axis in order { + if axis >= rank || seen[axis] { + return Err(pipeline_metadata_error( + index, + format!("transpose order must be a permutation of 0..{rank}"), + )); + } + seen[axis] = true; + } + Ok(()) +} + +fn pipeline_metadata_error(index: usize, reason: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!("Zarr v3 codec index {index}: {}", reason.into())) +} + +fn unsupported_pipeline_feature(index: usize, reason: impl Into) -> ZarrFdwError { + ZarrFdwError::UnsupportedExecutionFeature(format!( + "Zarr v3 codec index {index}: {}", + reason.into() + )) +} + +fn codec_read_error(index: usize, name: &str, reason: impl Into) -> ZarrFdwError { + ZarrFdwError::ReadError(std::io::Error::other(format!( + "codec index {index} ('{name}'): {}", + reason.into() + ))) +} + +fn checked_logical_bytes(shape: &[usize], itemsize: usize) -> ZarrFdwResult { + if shape.is_empty() || shape.contains(&0) || itemsize == 0 { + return Err(ZarrFdwError::InvalidMetadata( + "codec chunk shape and item size must be positive".to_string(), + )); + } + let cells = shape.iter().try_fold(1usize, |cells, &extent| { + cells.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "declared chunk cell count exceeds this platform's index capacity".to_string(), + ) + }) + })?; + let bytes = cells.checked_mul(itemsize).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "declared chunk byte length exceeds this platform's index capacity".to_string(), + ) + })?; + if bytes > MAX_DECODED_CHUNK_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "declared chunk decodes to {bytes} bytes, exceeding the safety limit of {MAX_DECODED_CHUNK_BYTES}" + ))); + } + Ok(bytes) +} + +fn require_exact_length( + data: Vec, + expected: usize, + index: usize, + name: &str, +) -> ZarrFdwResult> { + if data.len() != expected { + return Err(codec_read_error( + index, + name, + format!( + "decoded chunk has {} bytes, expected exactly {expected}", + data.len() + ), + )); + } + Ok(data) +} + +#[derive(Debug, Clone, Copy)] +enum StreamCodec { + Gzip, + Zlib, +} + +impl StreamCodec { + fn label(&self) -> &'static str { + match self { + Self::Gzip => "gzip", + Self::Zlib => "zlib", + } + } +} + +async fn decode_stream( + codec: StreamCodec, + data: Vec, + expected: usize, + interrupt_pending: &mut F, + index: usize, +) -> ZarrFdwResult>> +where + F: FnMut() -> bool, +{ + use async_compression::tokio::bufread::{GzipDecoder, ZlibDecoder}; + + match codec { + StreamCodec::Gzip => { + let mut decoder = BufReader::new(GzipDecoder::new(BufReader::new(Cursor::new(data)))); + read_exact_bounded( + &mut decoder, + expected, + interrupt_pending, + index, + codec.label(), + ) + .await + } + StreamCodec::Zlib => { + let mut decoder = BufReader::new(ZlibDecoder::new(BufReader::new(Cursor::new(data)))); + read_exact_bounded( + &mut decoder, + expected, + interrupt_pending, + index, + codec.label(), + ) + .await + } + } +} + +async fn read_exact_bounded( + reader: &mut R, + expected: usize, + interrupt_pending: &mut F, + index: usize, + label: &str, +) -> ZarrFdwResult>> +where + R: AsyncRead + Unpin, + F: FnMut() -> bool, +{ + let mut decoded = Vec::new(); + decoded.try_reserve_exact(expected).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate a decoded chunk of {expected} bytes" + )) + })?; + let mut buffer = vec![0u8; STREAM_POLL_BYTES.min(expected.max(1))]; + while decoded.len() < expected { + if interrupt_pending() { + return Ok(None); + } + let remaining = expected - decoded.len(); + let read_len = remaining.min(buffer.len()); + let count = reader + .read(&mut buffer[..read_len]) + .await + .map_err(|error| { + codec_read_error(index, label, format!("failed to decode stream: {error}")) + })?; + if count == 0 { + return Err(codec_read_error( + index, + label, + format!( + "decoded chunk has {} bytes, expected exactly {expected}", + decoded.len() + ), + )); + } + decoded.extend_from_slice(&buffer[..count]); + } + if interrupt_pending() { + return Ok(None); + } + let mut extra = [0u8; 1]; + let extra_count = reader.read(&mut extra).await.map_err(|error| { + codec_read_error(index, label, format!("failed to finish stream: {error}")) + })?; + if extra_count != 0 { + return Err(codec_read_error( + index, + label, + format!("decoded chunk has more than {expected} bytes, expected exactly {expected}"), + )); + } + Ok(Some(decoded)) +} + +fn parse_zstd_frame_header(data: &[u8], index: usize) -> ZarrFdwResult { + let magic_bytes: [u8; 4] = data + .get(..4) + .and_then(|bytes| bytes.try_into().ok()) + .ok_or_else(|| zstd_header_error(index, "object is shorter than the 4-byte magic"))?; + let magic = u32::from_le_bytes(magic_bytes); + if magic & ZSTD_SKIPPABLE_MAGIC_MASK == ZSTD_SKIPPABLE_MAGIC_START { + return Err(codec_read_error( + index, + "zstd", + "skippable Zstandard frames are not supported", + )); + } + if magic != ZSTD_FRAME_MAGIC { + return Err(zstd_header_error( + index, + format!("unexpected magic 0x{magic:08x}"), + )); + } + + let descriptor = *data + .get(4) + .ok_or_else(|| zstd_header_error(index, "frame descriptor is missing"))?; + if descriptor & 0b0000_1000 != 0 { + return Err(zstd_header_error(index, "reserved descriptor bit is set")); + } + if descriptor & 0b0000_0011 != 0 { + return Err(codec_read_error( + index, + "zstd", + "Zstandard dictionaries are not supported", + )); + } + + let single_segment = descriptor & 0b0010_0000 != 0; + let checksum = descriptor & 0b0000_0100 != 0; + let mut cursor = 5usize; + let advertised_window = if single_segment { + None + } else { + let window_descriptor = *data + .get(cursor) + .ok_or_else(|| zstd_header_error(index, "window descriptor is missing"))?; + cursor += 1; + let exponent = u32::from(window_descriptor >> 3); + let mantissa = u64::from(window_descriptor & 0b0000_0111); + let window_log = 10u32 + .checked_add(exponent) + .ok_or_else(|| zstd_header_error(index, "window descriptor exponent overflows"))?; + let window_base = 1u64.checked_shl(window_log).ok_or_else(|| { + zstd_header_error( + index, + "window descriptor exceeds the supported integer range", + ) + })?; + let window_add = (window_base / 8) + .checked_mul(mantissa) + .ok_or_else(|| zstd_header_error(index, "window descriptor mantissa overflows"))?; + Some( + window_base + .checked_add(window_add) + .ok_or_else(|| zstd_header_error(index, "advertised window size overflows"))?, + ) + }; + + let content_size_flag = descriptor >> 6; + let content_size_bytes = match content_size_flag { + 0 if single_segment => 1usize, + 0 => 0usize, + 1 => 2usize, + 2 => 4usize, + 3 => 8usize, + _ => unreachable!("two-bit field"), + }; + let content_size = if content_size_bytes == 0 { + None + } else { + let end = cursor + .checked_add(content_size_bytes) + .ok_or_else(|| zstd_header_error(index, "frame content-size field offset overflows"))?; + let bytes = data + .get(cursor..end) + .ok_or_else(|| zstd_header_error(index, "frame content-size field is truncated"))?; + let mut value_bytes = [0u8; 8]; + value_bytes[..content_size_bytes].copy_from_slice(bytes); + let value = u64::from_le_bytes(value_bytes); + Some(if content_size_bytes == 2 { + value + .checked_add(256) + .ok_or_else(|| zstd_header_error(index, "frame content size overflows"))? + } else { + value + }) + }; + + let window_size = if single_segment { + content_size + .ok_or_else(|| zstd_header_error(index, "single-segment frame has no content size"))? + } else { + advertised_window.expect("non-single-segment frame parsed a window descriptor") + }; + if window_size > ZSTD_WINDOW_BYTES { + return Err(codec_read_error( + index, + "zstd", + format!( + "Zstandard frame window {window_size} exceeds the {ZSTD_WINDOW_BYTES}-byte limit" + ), + )); + } + + Ok(ZstdFrameHeader { + checksum, + content_size, + }) +} + +fn zstd_header_error(index: usize, reason: impl Into) -> ZarrFdwError { + codec_read_error( + index, + "zstd", + format!("invalid Zstandard frame header: {}", reason.into()), + ) +} + +fn zstd_native_error(index: usize, operation: &str, code: usize) -> ZarrFdwError { + codec_read_error( + index, + "zstd", + format!("{operation}: {}", zstd::zstd_safe::get_error_name(code)), + ) +} + +fn decode_zstd( + data: Vec, + expected: usize, + config: &ZstdConfig, + interrupt_pending: &mut F, + index: usize, +) -> ZarrFdwResult>> +where + F: FnMut() -> bool, +{ + if interrupt_pending() { + return Ok(None); + } + let header = parse_zstd_frame_header(&data, index)?; + if header.checksum != config.checksum { + return Err(codec_read_error( + index, + "zstd", + format!( + "Zstandard checksum metadata does not match the frame checksum flag (metadata={}, frame={})", + config.checksum, header.checksum + ), + )); + } + if let Some(content_size) = header.content_size { + let expected = u64::try_from(expected).map_err(|_| { + codec_read_error( + index, + "zstd", + "logical decoded size exceeds the supported integer range", + ) + })?; + if content_size != expected { + return Err(codec_read_error( + index, + "zstd", + format!( + "Zstandard frame content size {content_size} does not match expected {expected}" + ), + )); + } + } + + if interrupt_pending() { + return Ok(None); + } + let frame_size = zstd::zstd_safe::find_frame_compressed_size(&data) + .map_err(|code| zstd_native_error(index, "invalid Zstandard frame", code))?; + if frame_size != data.len() { + return Err(codec_read_error( + index, + "zstd", + format!( + "concatenated Zstandard frames and trailing bytes are not supported (first frame has {frame_size} bytes, object has {})", + data.len() + ), + )); + } + + if interrupt_pending() { + return Ok(None); + } + let mut context = zstd::zstd_safe::DCtx::try_create().ok_or_else(|| { + codec_read_error( + index, + "zstd", + "could not allocate a Zstandard decoder context", + ) + })?; + context + .init() + .map_err(|code| zstd_native_error(index, "failed to initialize Zstandard decoder", code))?; + context + .set_parameter(zstd::zstd_safe::DParameter::WindowLogMax( + ZSTD_WINDOW_LOG_MAX, + )) + .map_err(|code| { + zstd_native_error(index, "failed to set the Zstandard window limit", code) + })?; + let mut decoder = + zstd::stream::read::Decoder::with_context(Cursor::new(data), &mut context).single_frame(); + read_exact_bounded_sync(&mut decoder, expected, interrupt_pending, index) +} + +fn read_exact_bounded_sync( + reader: &mut R, + expected: usize, + interrupt_pending: &mut F, + index: usize, +) -> ZarrFdwResult>> +where + R: std::io::Read, + F: FnMut() -> bool, +{ + let mut decoded = Vec::new(); + decoded.try_reserve_exact(expected).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate a decoded chunk of {expected} bytes" + )) + })?; + let mut buffer = vec![0u8; STREAM_POLL_BYTES.min(expected.max(1))]; + while decoded.len() < expected { + if interrupt_pending() { + return Ok(None); + } + let remaining = expected - decoded.len(); + let read_len = remaining.min(buffer.len()); + let count = std::io::Read::read(reader, &mut buffer[..read_len]).map_err(|error| { + codec_read_error( + index, + "zstd", + format!("failed to decode Zstandard frame: {error}"), + ) + })?; + if count == 0 { + return Err(codec_read_error( + index, + "zstd", + format!( + "decoded chunk has {} bytes, expected exactly {expected}", + decoded.len() + ), + )); + } + decoded.extend_from_slice(&buffer[..count]); + } + if interrupt_pending() { + return Ok(None); + } + let mut extra = [0u8; 1]; + let extra_count = std::io::Read::read(reader, &mut extra).map_err(|error| { + codec_read_error( + index, + "zstd", + format!("failed to decode Zstandard frame: {error}"), + ) + })?; + if extra_count != 0 { + return Err(codec_read_error( + index, + "zstd", + format!("decoded chunk has more than {expected} bytes, expected exactly {expected}"), + )); + } + Ok(Some(decoded)) +} + +fn decode_blosc( + data: Vec, + expected: usize, + index: usize, + v3_config: Option<&BloscConfig>, +) -> ZarrFdwResult> { + if v3_config.is_some() { + validate_v3_blosc_header(&data, expected, index)?; + } + let decoder = blosc_rs::Decoder::new(data).map_err(|error| { + codec_read_error(index, "blosc", format!("invalid Blosc chunk: {error}")) + })?; + if decoder.nbytes() != expected { + return Err(codec_read_error( + index, + "blosc", + format!( + "decoded chunk has {} bytes, expected exactly {expected}", + decoder.nbytes() + ), + )); + } + let decoded = decoder.decompress(1).map_err(|error| { + codec_read_error( + index, + "blosc", + format!("failed to decompress Blosc chunk: {error}"), + ) + })?; + require_exact_length(decoded, expected, index, "blosc") +} + +fn validate_v3_blosc_header(data: &[u8], expected: usize, index: usize) -> ZarrFdwResult<()> { + if data.len() < BLOSC_HEADER_BYTES { + return Err(codec_read_error( + index, + "blosc", + "encoded chunk is shorter than the 16-byte Blosc header", + )); + } + let declared_nbytes = u32::from_le_bytes( + data[4..8] + .try_into() + .expect("validated 16-byte Blosc header"), + ) as usize; + if declared_nbytes != expected { + return Err(codec_read_error( + index, + "blosc", + format!( + "Blosc header declares {declared_nbytes} uncompressed bytes, expected exactly {expected}" + ), + )); + } + let declared_cbytes = u32::from_le_bytes( + data[12..16] + .try_into() + .expect("validated 16-byte Blosc header"), + ) as usize; + if declared_cbytes != data.len() { + return Err(codec_read_error( + index, + "blosc", + format!( + "Blosc header declares {declared_cbytes} compressed bytes, actual object has {}", + data.len() + ), + )); + } + Ok(()) +} + +fn decode_crc32c( + mut data: Vec, + interrupt_pending: &mut F, + index: usize, +) -> ZarrFdwResult>> +where + F: FnMut() -> bool, +{ + if data.len() < CRC32C_BYTES { + return Err(codec_read_error( + index, + "crc32c", + "encoded chunk is truncated before the four-byte checksum", + )); + } + let payload_len = data.len() - CRC32C_BYTES; + let expected = u32::from_le_bytes(data[payload_len..].try_into().expect("four bytes")); + let mut actual = 0u32; + for chunk in data[..payload_len].chunks(CRC_POLL_BYTES) { + if interrupt_pending() { + return Ok(None); + } + actual = crc32c::crc32c_append(actual, chunk); + } + if actual != expected { + return Err(codec_read_error( + index, + "crc32c", + format!("checksum mismatch: expected {expected:#010x}, computed {actual:#010x}"), + )); + } + data.truncate(payload_len); + Ok(Some(data)) +} + +fn checked_strides(shape: &[usize]) -> ZarrFdwResult> { + let mut strides = vec![0usize; shape.len()]; + let mut stride = 1usize; + for axis in (0..shape.len()).rev() { + strides[axis] = stride; + stride = stride.checked_mul(shape[axis]).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "transpose stride exceeds this platform's index capacity".to_string(), + ) + })?; + } + Ok(strides) +} + +fn inverse_transpose( + data: Vec, + logical_shape: &[usize], + itemsize: usize, + order: &[usize], + interrupt_pending: &mut F, + index: usize, +) -> ZarrFdwResult>> +where + F: FnMut() -> bool, +{ + validate_permutation(order, logical_shape.len(), index)?; + let expected = checked_logical_bytes(logical_shape, itemsize)?; + let data = require_exact_length(data, expected, index, "transpose")?; + if order.iter().copied().eq(0..order.len()) { + return Ok(Some(data)); + } + + let logical_strides = checked_strides(logical_shape)?; + let encoded_shape = order + .iter() + .map(|&axis| logical_shape[axis]) + .collect::>(); + let encoded_strides = checked_strides(&encoded_shape)?; + let cells = expected / itemsize; + let mut logical = Vec::new(); + logical.try_reserve_exact(expected).map_err(|_| { + codec_read_error( + index, + "transpose", + format!("could not allocate a transposed chunk of {expected} bytes"), + ) + })?; + logical.resize(expected, 0); + for logical_flat in 0..cells { + if logical_flat % TRANSPOSE_POLL_CELLS == 0 && interrupt_pending() { + return Ok(None); + } + let encoded_flat = + order + .iter() + .enumerate() + .try_fold(0usize, |flat, (encoded_axis, &logical_axis)| { + let coordinate = (logical_flat / logical_strides[logical_axis]) + % logical_shape[logical_axis]; + coordinate + .checked_mul(encoded_strides[encoded_axis]) + .and_then(|offset| flat.checked_add(offset)) + .ok_or_else(|| { + codec_read_error( + index, + "transpose", + "transpose element offset exceeds this platform's index capacity", + ) + }) + })?; + let logical_byte = logical_flat + .checked_mul(itemsize) + .ok_or_else(|| codec_read_error(index, "transpose", "logical byte offset overflow"))?; + let encoded_byte = encoded_flat + .checked_mul(itemsize) + .ok_or_else(|| codec_read_error(index, "transpose", "encoded byte offset overflow"))?; + let encoded_end = encoded_byte + .checked_add(itemsize) + .ok_or_else(|| codec_read_error(index, "transpose", "encoded byte range overflow"))?; + let logical_end = logical_byte + .checked_add(itemsize) + .ok_or_else(|| codec_read_error(index, "transpose", "logical byte range overflow"))?; + logical[logical_byte..logical_end].copy_from_slice(&data[encoded_byte..encoded_end]); + } + Ok(Some(logical)) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + fn v3(codecs: Value) -> ZarrFdwResult<(CodecPipeline, String)> { + CodecPipeline::from_v3("float32", 3, &codecs) + } + + fn v3_blosc(cname: &str, shuffle: &str, typesize: Option) -> Value { + let mut configuration = serde_json::json!({ + "cname": cname, + "clevel": 5, + "shuffle": shuffle, + "blocksize": 0 + }); + if let Some(typesize) = typesize { + configuration["typesize"] = typesize; + } + serde_json::json!({"name":"blosc", "configuration":configuration}) + } + + fn v3_zstd(level: Value, checksum: Option) -> Value { + let mut configuration = serde_json::json!({"level":level}); + if let Some(checksum) = checksum { + configuration["checksum"] = checksum; + } + serde_json::json!({"name":"zstd", "configuration":configuration}) + } + + fn zstd_pipeline(checksum: bool) -> CodecPipeline { + CodecPipeline::from_v3( + "int8", + 1, + &serde_json::json!([ + {"name":"bytes","configuration":{}}, + v3_zstd(serde_json::json!(1), Some(serde_json::json!(checksum))) + ]), + ) + .unwrap() + .0 + } + + fn encode_zstd(raw: &[u8], checksum: bool, include_content_size: bool) -> Vec { + let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 1).unwrap(); + encoder.include_checksum(checksum).unwrap(); + encoder.include_contentsize(include_content_size).unwrap(); + let pledged_size = u64::try_from(raw.len()).unwrap(); + encoder + .set_pledged_src_size(include_content_size.then_some(pledged_size)) + .unwrap(); + encoder.write_all(raw).unwrap(); + encoder.finish().unwrap() + } + + #[test] + fn parses_locked_v3_pipeline_and_label() { + let codecs = serde_json::json!([ + {"name":"transpose","configuration":{"order":[2,1,0]}}, + {"name":"bytes","configuration":{"endian":"little"}}, + {"name":"gzip","configuration":{"level":1}}, + {"name":"crc32c"} + ]); + let (pipeline, dtype) = v3(codecs).unwrap(); + assert_eq!(dtype, " bytes -> gzip -> crc32c" + ); + assert_eq!( + pipeline.encoded_read_limit(1024).unwrap(), + 1024 + COMPRESSED_OVERHEAD_ALLOWANCE + CRC32C_BYTES + ); + } + + #[test] + fn normalizes_v2_without_changing_legacy_coverage() { + for (compressor, label, allowance) in [ + (None, "raw", 0), + ( + Some(serde_json::json!({"id":"gzip","level":1})), + "gzip", + COMPRESSED_OVERHEAD_ALLOWANCE, + ), + ( + Some(serde_json::json!({"id":"zlib","level":1})), + "zlib", + COMPRESSED_OVERHEAD_ALLOWANCE, + ), + ( + Some(serde_json::json!({"id":"blosc","cname":"lz4"})), + "blosc", + BLOSC_HEADER_BYTES, + ), + ] { + let pipeline = CodecPipeline::from_v2(&compressor).unwrap(); + assert_eq!(pipeline.ordered_label(), label); + assert_eq!(pipeline.encoded_read_limit(32).unwrap(), 32 + allowance); + } + } + + #[test] + fn parses_supported_v3_zstd_configurations_and_label() { + for level in [-131_072, 0, 22] { + for checksum in [None, Some(false), Some(true)] { + let codecs = serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_zstd( + serde_json::json!(level), + checksum.map(|value| serde_json::json!(value)) + ), + {"name":"crc32c"} + ]); + let (pipeline, dtype) = v3(codecs).unwrap(); + assert_eq!(dtype, " zstd -> crc32c"); + assert_eq!( + pipeline.encoded_read_limit(1024).unwrap(), + 1024 + COMPRESSED_OVERHEAD_ALLOWANCE + CRC32C_BYTES + ); + let CodecStage::Zstd { config } = &pipeline.stages[1] else { + panic!("expected configured v3 Zstd stage"); + }; + assert_eq!(config.level, level); + assert_eq!(config.checksum, checksum.unwrap_or(false)); + } + } + } + + #[test] + fn rejects_malformed_v3_zstd_configurations() { + let invalid = [ + serde_json::json!({"name":"zstd"}), + serde_json::json!({"name":"zstd","configuration":{}}), + serde_json::json!({"name":"zstd","configuration":{"checksum":true}}), + v3_zstd(serde_json::json!(-131_073), None), + v3_zstd(serde_json::json!(23), None), + v3_zstd(serde_json::json!(1.5), None), + v3_zstd(serde_json::json!("1"), None), + v3_zstd(serde_json::json!(1), Some(serde_json::json!(1))), + serde_json::json!({"name":"zstd","configuration":{"level":1,"extra":true}}), + serde_json::json!({"name":"zstd","configuration":{"level":1},"must_understand":"yes"}), + ]; + for codec in invalid { + let codecs = serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + codec + ]); + let error = v3(codecs).unwrap_err(); + assert!(matches!(error, ZarrFdwError::InvalidMetadata(_))); + } + } + + #[test] + fn rejects_duplicate_mixed_and_misordered_v3_zstd() { + let zstd = v3_zstd(serde_json::json!(1), None); + for codecs in [ + serde_json::json!([ + zstd.clone(), + {"name":"bytes","configuration":{"endian":"little"}} + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + zstd.clone(), + zstd.clone() + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + {"name":"gzip","configuration":{"level":1}}, + zstd.clone() + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + zstd.clone(), + v3_blosc("lz4", "noshuffle", None) + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + {"name":"crc32c"}, + zstd + ]), + ] { + assert!(matches!(v3(codecs), Err(ZarrFdwError::InvalidMetadata(_)))); + } + } + + #[test] + fn parses_supported_v3_blosc_configurations_and_label() { + for cname in ["blosclz", "lz4", "lz4hc"] { + for (shuffle, typesize) in [ + ("noshuffle", None), + ("noshuffle", Some(serde_json::json!(4))), + ("shuffle", Some(serde_json::json!(4))), + ("bitshuffle", Some(serde_json::json!(4))), + ] { + let codecs = serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_blosc(cname, shuffle, typesize), + {"name":"crc32c"} + ]); + let (pipeline, dtype) = v3(codecs).unwrap(); + assert_eq!(dtype, " blosc -> crc32c"); + assert_eq!( + pipeline.encoded_read_limit(1024).unwrap(), + 1024 + BLOSC_HEADER_BYTES + CRC32C_BYTES + ); + let CodecStage::Blosc { + config: Some(config), + } = &pipeline.stages[1] + else { + panic!("expected configured v3 Blosc stage"); + }; + assert!((0..=9).contains(&config.clevel)); + } + } + } + + #[test] + fn rejects_malformed_and_unavailable_v3_blosc_configurations() { + let invalid = [ + serde_json::json!({"name":"blosc"}), + serde_json::json!({"name":"blosc","configuration":{}}), + serde_json::json!({"name":"blosc","configuration":{"cname":1,"clevel":5,"shuffle":"noshuffle","blocksize":0}}), + v3_blosc("future", "noshuffle", None), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":-1,"shuffle":"noshuffle","blocksize":0}}), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":10,"shuffle":"noshuffle","blocksize":0}}), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":1.5,"shuffle":"noshuffle","blocksize":0}}), + v3_blosc("lz4", "auto", Some(serde_json::json!(4))), + v3_blosc("lz4", "shuffle", None), + v3_blosc("lz4", "shuffle", Some(serde_json::json!(0))), + v3_blosc("lz4", "shuffle", Some(serde_json::json!(-1))), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":5,"shuffle":"noshuffle"}}), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":5,"shuffle":"noshuffle","blocksize":-1}}), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":5,"shuffle":"noshuffle","blocksize":0,"extra":true}}), + serde_json::json!({"name":"blosc","configuration":{"cname":"lz4","clevel":5,"shuffle":"noshuffle","blocksize":0},"must_understand":"yes"}), + ]; + for codec in invalid { + let codecs = serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + codec + ]); + assert!(matches!(v3(codecs), Err(ZarrFdwError::InvalidMetadata(_)))); + } + + for cname in ["zstd", "snappy", "zlib"] { + let codecs = serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_blosc(cname, "noshuffle", None) + ]); + let error = v3(codecs).unwrap_err(); + assert!(matches!( + &error, + ZarrFdwError::UnsupportedExecutionFeature(_) + )); + assert!(format!("{error}").contains(&format!("Blosc cname '{cname}'"))); + } + } + + #[test] + fn rejects_duplicate_or_misordered_v3_blosc() { + let blosc = v3_blosc("lz4", "shuffle", Some(serde_json::json!(4))); + for codecs in [ + serde_json::json!([ + blosc.clone(), + {"name":"bytes","configuration":{"endian":"little"}} + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + blosc.clone(), + blosc.clone() + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + {"name":"gzip","configuration":{"level":1}}, + blosc.clone() + ]), + serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + blosc, + {"name":"gzip","configuration":{"level":1}} + ]), + ] { + assert!(matches!(v3(codecs), Err(ZarrFdwError::InvalidMetadata(_)))); + } + } + + #[test] + fn rejects_every_unsupported_v3_order_and_codec() { + let cases = [ + serde_json::json!([{"name":"gzip","configuration":{"level":1}}, {"name":"bytes","configuration":{"endian":"little"}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"transpose","configuration":{"order":[0,1,2]}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"crc32c"}, {"name":"gzip","configuration":{"level":1}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"bytes","configuration":{"endian":"little"}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"blosc","configuration":{}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"zstd","configuration":{"level":23}}]), + serde_json::json!([{"name":"sharding_indexed","configuration":{}}]), + ]; + for codecs in cases { + assert!(v3(codecs).is_err()); + } + } + + #[test] + fn validates_codec_configurations_and_endian() { + for codecs in [ + serde_json::json!([{"name":"transpose","configuration":{"order":[0,0,2]}}, {"name":"bytes","configuration":{"endian":"little"}}]), + serde_json::json!([{"name":"transpose","configuration":{"order":[0,1]}}, {"name":"bytes","configuration":{"endian":"little"}}]), + serde_json::json!([{"name":"bytes"}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"big"}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little","extra":1}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"gzip","configuration":{"level":10}}]), + serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"crc32c","configuration":{"seed":0}}]), + ] { + assert!(v3(codecs).is_err()); + } + + let (_, dtype) = CodecPipeline::from_v3( + "int8", + 1, + &serde_json::json!([{"name":"bytes","configuration":{}}]), + ) + .unwrap(); + assert_eq!(dtype, "|i1"); + let (_, dtype) = CodecPipeline::from_v3( + "int8", + 1, + &serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}]), + ) + .unwrap(); + assert_eq!(dtype, "|i1"); + + let (pipeline, _) = CodecPipeline::from_v3( + "float32", + 3, + &serde_json::json!([ + {"name":"transpose","configuration":{"order":[0,1,2]},"must_understand":false}, + {"name":"bytes","configuration":{"endian":"little"},"must_understand":false}, + {"name":"gzip","configuration":{"level":1},"must_understand":false}, + {"name":"crc32c","must_understand":false} + ]), + ) + .unwrap(); + assert_eq!( + pipeline.ordered_label(), + "transpose -> bytes -> gzip -> crc32c" + ); + } + + #[test] + fn checksum_rejects_truncation_and_corruption() { + let payload = b"payload".to_vec(); + let mut encoded = payload.clone(); + encoded.extend_from_slice(&crc32c::crc32c(&payload).to_le_bytes()); + assert_eq!( + decode_crc32c(encoded, &mut || false, 2).unwrap(), + Some(payload) + ); + assert!(decode_crc32c(vec![1, 2, 3], &mut || false, 2).is_err()); + + let mut corrupt = b"payload".to_vec(); + corrupt.extend_from_slice(&0u32.to_le_bytes()); + assert!(decode_crc32c(corrupt, &mut || false, 2).is_err()); + } + + #[test] + fn inverse_transpose_restores_logical_c_order() { + // Logical A shape [2, 2, 3], values 0..12. Encoding with order + // [2, 1, 0] produces B[x, y, z] in C order. + let logical = (0u8..12).collect::>(); + let encoded = vec![0, 6, 3, 9, 1, 7, 4, 10, 2, 8, 5, 11]; + assert_eq!( + inverse_transpose(encoded, &[2, 2, 3], 1, &[2, 1, 0], &mut || false, 0,).unwrap(), + Some(logical) + ); + } + + #[test] + fn pipeline_decodes_crc_then_gzip_then_transpose() { + use async_compression::tokio::write::GzipEncoder; + use tokio::io::AsyncWriteExt; + + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let pipeline = v3(serde_json::json!([ + {"name":"transpose","configuration":{"order":[2,1,0]}}, + {"name":"bytes","configuration":{"endian":"little"}}, + {"name":"gzip","configuration":{"level":1}}, + {"name":"crc32c"} + ])) + .unwrap() + .0; + let encoded_order = vec![0, 6, 3, 9, 1, 7, 4, 10, 2, 8, 5, 11]; + let mut encoder = GzipEncoder::new(Vec::new()); + encoder.write_all(&encoded_order).await.unwrap(); + encoder.shutdown().await.unwrap(); + let mut encoded = encoder.into_inner(); + let checksum = crc32c::crc32c(&encoded); + encoded.extend_from_slice(&checksum.to_le_bytes()); + assert_eq!( + pipeline + .decode_interruptible(encoded, &[2, 2, 3], 1, || false) + .await + .unwrap(), + CodecDecode::Decoded((0u8..12).collect()) + ); + }); + } + + #[test] + fn parses_zstd_frame_header_variants_and_enforces_window_policy() { + let mut unknown_size = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + unknown_size.extend_from_slice(&[0b0001_0000, 13 << 3]); + assert_eq!( + parse_zstd_frame_header(&unknown_size, 1).unwrap(), + ZstdFrameHeader { + checksum: false, + content_size: None, + } + ); + + let mut single_segment = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + single_segment.extend_from_slice(&[0b0010_0100, 4]); + assert_eq!( + parse_zstd_frame_header(&single_segment, 1).unwrap(), + ZstdFrameHeader { + checksum: true, + content_size: Some(4), + } + ); + + for (descriptor, encoded, expected) in [ + (0b0100_0000, 0u64, 256u64), + (0b1000_0000, 65_792u64, 65_792u64), + ( + 0b1100_0000, + u64::from(u32::MAX) + 1, + u64::from(u32::MAX) + 1, + ), + ] { + let width = match descriptor >> 6 { + 1 => 2, + 2 => 4, + 3 => 8, + _ => unreachable!(), + }; + let mut frame = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + frame.extend_from_slice(&[descriptor, 0]); + frame.extend_from_slice(&encoded.to_le_bytes()[..width]); + assert_eq!( + parse_zstd_frame_header(&frame, 1).unwrap().content_size, + Some(expected) + ); + } + + let mut excessive_window = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + excessive_window.extend_from_slice(&[0, (13 << 3) | 1]); + let error = parse_zstd_frame_header(&excessive_window, 1).unwrap_err(); + assert!(format!("{error}").contains("exceeds the 8388608-byte limit")); + + let mut excessive_single_segment = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + excessive_single_segment.push(0b1110_0000); + excessive_single_segment.extend_from_slice(&(ZSTD_WINDOW_BYTES + 1).to_le_bytes()); + let error = parse_zstd_frame_header(&excessive_single_segment, 1).unwrap_err(); + assert!(format!("{error}").contains("exceeds the 8388608-byte limit")); + + let mut reserved = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + reserved.extend_from_slice(&[0b0000_1000, 0]); + assert!( + format!("{}", parse_zstd_frame_header(&reserved, 1).unwrap_err()) + .contains("reserved descriptor bit") + ); + + let mut dictionary = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + dictionary.push(1); + assert!( + format!("{}", parse_zstd_frame_header(&dictionary, 1).unwrap_err()) + .contains("Zstandard dictionaries are not supported") + ); + + let skippable = ZSTD_SKIPPABLE_MAGIC_START.to_le_bytes(); + assert!( + format!("{}", parse_zstd_frame_header(&skippable, 1).unwrap_err()) + .contains("skippable Zstandard frames are not supported") + ); + assert!( + format!("{}", parse_zstd_frame_header(&[1, 2, 3], 1).unwrap_err()) + .contains("invalid Zstandard frame header") + ); + assert!( + format!( + "{}", + parse_zstd_frame_header(&[0, 0, 0, 0, 0, 0], 1).unwrap_err() + ) + .contains("unexpected magic") + ); + assert!( + format!( + "{}", + parse_zstd_frame_header(&ZSTD_FRAME_MAGIC.to_le_bytes(), 1).unwrap_err() + ) + .contains("frame descriptor is missing") + ); + let mut truncated_window = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + truncated_window.push(0); + assert!( + format!( + "{}", + parse_zstd_frame_header(&truncated_window, 1).unwrap_err() + ) + .contains("window descriptor is missing") + ); + let mut truncated_content_size = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + truncated_content_size.push(0b0010_0000); + assert!( + format!( + "{}", + parse_zstd_frame_header(&truncated_content_size, 1).unwrap_err() + ) + .contains("content-size field is truncated") + ); + } + + #[test] + fn zstd_decodes_known_and_unknown_content_sizes_and_polls_cancellation() { + let raw = (0..(STREAM_POLL_BYTES * 2 + 17)) + .map(|index| u8::try_from(index % 251).unwrap()) + .collect::>(); + let known = encode_zstd(&raw, true, true); + let unknown = encode_zstd(&raw, false, false); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + assert_eq!( + zstd_pipeline(true) + .decode_interruptible(known.clone(), &[raw.len()], 1, || false) + .await + .unwrap(), + CodecDecode::Decoded(raw.clone()) + ); + assert_eq!( + zstd_pipeline(false) + .decode_interruptible(unknown.clone(), &[raw.len()], 1, || false) + .await + .unwrap(), + CodecDecode::Decoded(raw.clone()) + ); + }); + + let config = ZstdConfig { + level: 1, + checksum: false, + }; + assert_eq!( + decode_zstd(unknown.clone(), raw.len(), &config, &mut || true, 1).unwrap(), + None + ); + let mut polls = 0usize; + assert_eq!( + decode_zstd( + unknown, + raw.len(), + &config, + &mut || { + polls += 1; + polls >= 5 + }, + 1, + ) + .unwrap(), + None + ); + assert!(polls >= 5); + } + + #[test] + fn zstd_rejects_checksum_mismatch_corruption_and_non_exact_output() { + let raw = (0u8..64).collect::>(); + let checksummed = encode_zstd(&raw, true, true); + let unchecked = encode_zstd(&raw, false, true); + + for (encoded, metadata_checksum) in + [(checksummed.clone(), false), (unchecked.clone(), true)] + { + let config = ZstdConfig { + level: 1, + checksum: metadata_checksum, + }; + let error = decode_zstd(encoded, raw.len(), &config, &mut || false, 1).unwrap_err(); + assert!( + format!("{error}") + .contains("Zstandard checksum metadata does not match the frame checksum flag") + ); + } + + let mut corrupt = checksummed; + *corrupt.last_mut().unwrap() ^= 1; + let error = decode_zstd( + corrupt, + raw.len(), + &ZstdConfig { + level: 1, + checksum: true, + }, + &mut || false, + 1, + ) + .unwrap_err(); + let error = format!("{error}"); + assert!(error.contains("codec index 1 ('zstd')")); + assert!(error.contains("failed to decode Zstandard frame")); + + let unknown = encode_zstd(&raw, false, false); + let config = ZstdConfig { + level: 1, + checksum: false, + }; + let short = + decode_zstd(unknown.clone(), raw.len() + 1, &config, &mut || false, 1).unwrap_err(); + assert!(format!("{short}").contains("expected exactly 65")); + let long = decode_zstd(unknown, raw.len() - 1, &config, &mut || false, 1).unwrap_err(); + assert!(format!("{long}").contains("more than 63 bytes")); + } + + #[test] + fn zstd_rejects_frame_policy_violations_before_decoding() { + let raw = [1u8, 2, 3, 4]; + let frame = encode_zstd(&raw, true, true); + let config = ZstdConfig { + level: 1, + checksum: true, + }; + + let mismatch = decode_zstd(frame.clone(), 5, &config, &mut || false, 1).unwrap_err(); + assert!(format!("{mismatch}").contains("content size 4 does not match expected 5")); + + let mut trailing = frame.clone(); + trailing.push(0); + let error = decode_zstd(trailing, raw.len(), &config, &mut || false, 1).unwrap_err(); + assert!( + format!("{error}") + .contains("concatenated Zstandard frames and trailing bytes are not supported") + ); + + let mut concatenated = frame.clone(); + concatenated.extend_from_slice(&frame); + let error = decode_zstd(concatenated, raw.len(), &config, &mut || false, 1).unwrap_err(); + assert!( + format!("{error}") + .contains("concatenated Zstandard frames and trailing bytes are not supported") + ); + + let mut excessive_window = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + excessive_window.extend_from_slice(&[0b0000_0100, (13 << 3) | 1]); + assert!( + format!( + "{}", + decode_zstd(excessive_window, raw.len(), &config, &mut || false, 1).unwrap_err() + ) + .contains("exceeds the 8388608-byte limit") + ); + + let mut dictionary = ZSTD_FRAME_MAGIC.to_le_bytes().to_vec(); + dictionary.push(0b0000_0101); + assert!( + format!( + "{}", + decode_zstd(dictionary, raw.len(), &config, &mut || false, 1).unwrap_err() + ) + .contains("Zstandard dictionaries are not supported") + ); + } + + #[test] + fn pipeline_decodes_crc_then_zstd_then_transpose() { + let logical = (0u8..12).collect::>(); + let encoded_order = vec![0, 6, 3, 9, 1, 7, 4, 10, 2, 8, 5, 11]; + let compressed = encode_zstd(&encoded_order, true, true); + let mut encoded = compressed.clone(); + encoded.extend_from_slice(&crc32c::crc32c(&compressed).to_le_bytes()); + let pipeline = v3(serde_json::json!([ + {"name":"transpose","configuration":{"order":[2,1,0]}}, + {"name":"bytes","configuration":{"endian":"little"}}, + v3_zstd(serde_json::json!(1), Some(serde_json::json!(true))), + {"name":"crc32c"} + ])) + .unwrap() + .0; + assert_eq!( + pipeline.ordered_label(), + "transpose -> bytes -> zstd -> crc32c" + ); + let rt = tokio::runtime::Runtime::new().unwrap(); + assert_eq!( + rt.block_on(pipeline.decode_interruptible(encoded, &[2, 2, 3], 1, || false)) + .unwrap(), + CodecDecode::Decoded(logical) + ); + } + + #[test] + fn v3_blosc_header_validation_rejects_truncation_mismatch_and_trailing_bytes() { + use blosc_rs::{CompressAlgo, Encoder}; + + let raw = (0..64u32).flat_map(u32::to_le_bytes).collect::>(); + let encoded = Encoder::default() + .compressor(CompressAlgo::Lz4) + .typesize(4.try_into().unwrap()) + .compress(&raw) + .unwrap(); + validate_v3_blosc_header(&encoded, raw.len(), 1).unwrap(); + + let error = validate_v3_blosc_header(&encoded[..15], raw.len(), 1).unwrap_err(); + assert!(format!("{error}").contains("shorter than the 16-byte Blosc header")); + + let mut wrong_nbytes = encoded.clone(); + wrong_nbytes[4..8].copy_from_slice(&1u32.to_le_bytes()); + let error = validate_v3_blosc_header(&wrong_nbytes, raw.len(), 1).unwrap_err(); + assert!(format!("{error}").contains("declares 1 uncompressed bytes")); + + let mut wrong_cbytes = encoded.clone(); + wrong_cbytes[12..16].copy_from_slice(&16u32.to_le_bytes()); + let error = validate_v3_blosc_header(&wrong_cbytes, raw.len(), 1).unwrap_err(); + assert!(format!("{error}").contains("declares 16 compressed bytes")); + + let mut trailing = encoded; + trailing.push(0); + let error = validate_v3_blosc_header(&trailing, raw.len(), 1).unwrap_err(); + assert!(format!("{error}").contains("actual object has")); + } + + #[test] + fn v3_blosc_decodes_every_advertised_backend_and_shuffle() { + use blosc_rs::{CompressAlgo, Encoder, Shuffle}; + + // A repetitive, nontrivial payload is large enough that every enabled + // backend actually compresses it instead of emitting Blosc's memcpy + // fallback. That makes the test prove native unshuffle behavior too. + let raw = (0..4096) + .map(|index| u8::try_from(index % 16).unwrap()) + .collect::>(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + for (cname, algorithm) in [ + ("blosclz", CompressAlgo::Blosclz), + ("lz4", CompressAlgo::Lz4), + ("lz4hc", CompressAlgo::Lz4hc), + ] { + for (shuffle_name, shuffle, expected_flag) in [ + ("noshuffle", Shuffle::None, 0u8), + ("shuffle", Shuffle::Byte, 1u8), + ("bitshuffle", Shuffle::Bit, 4u8), + ] { + let mut encoder = Encoder::default(); + encoder + .compressor(algorithm) + .shuffle(shuffle) + .typesize(4.try_into().unwrap()); + let encoded = encoder.compress(&raw).unwrap(); + assert!(encoded.len() < raw.len() + BLOSC_HEADER_BYTES); + assert_eq!(encoded[2] & 0b111, expected_flag); + let pipeline = v3(serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_blosc(cname, shuffle_name, Some(serde_json::json!(4))) + ])) + .unwrap() + .0; + assert_eq!( + pipeline + .decode_interruptible(encoded, &[1024], 4, || false) + .await + .unwrap(), + CodecDecode::Decoded(raw.clone()), + "backend={cname}, shuffle={shuffle_name}" + ); + } + } + }); + } + + #[test] + fn v3_blosc_decodes_in_reverse_order_and_polls_after_native_call() { + use blosc_rs::{CompressAlgo, Encoder}; + + let raw = (0..64u32).flat_map(u32::to_le_bytes).collect::>(); + let compressed = Encoder::default() + .compressor(CompressAlgo::Lz4) + .typesize(4.try_into().unwrap()) + .compress(&raw) + .unwrap(); + let pipeline = v3(serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_blosc("lz4", "shuffle", Some(serde_json::json!(4))), + {"name":"crc32c"} + ])) + .unwrap() + .0; + let mut encoded = compressed.clone(); + encoded.extend_from_slice(&crc32c::crc32c(&compressed).to_le_bytes()); + + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + assert_eq!( + pipeline + .decode_interruptible(encoded.clone(), &[64], 4, || false) + .await + .unwrap(), + CodecDecode::Decoded(raw.clone()) + ); + + let no_crc = v3(serde_json::json!([ + {"name":"bytes","configuration":{"endian":"little"}}, + v3_blosc("lz4", "shuffle", Some(serde_json::json!(4))) + ])) + .unwrap() + .0; + let mut polls = 0usize; + assert_eq!( + no_crc + .decode_interruptible(compressed.clone(), &[64], 4, || { + polls += 1; + polls >= 3 + }) + .await + .unwrap(), + CodecDecode::Interrupted + ); + + let mut truncated = compressed; + truncated.pop(); + let declared = u32::try_from(truncated.len()).unwrap(); + truncated[12..16].copy_from_slice(&declared.to_le_bytes()); + let error = no_crc + .decode_interruptible(truncated, &[64], 4, || false) + .await + .unwrap_err(); + let error = format!("{error}"); + assert!(error.contains("codec index 1 ('blosc')")); + assert!( + error.contains("invalid Blosc chunk") || error.contains("failed to decompress") + ); + }); + } + + #[test] + fn v2_raw_gzip_zlib_and_blosc_decode_through_one_interface() { + use async_compression::tokio::write::{GzipEncoder, ZlibEncoder}; + use blosc_rs::{CompressAlgo, Encoder}; + use tokio::io::AsyncWriteExt; + + let raw = (0..64u64).flat_map(u64::to_le_bytes).collect::>(); + let blosc = Encoder::default() + .compressor(CompressAlgo::Lz4) + .typesize(8.try_into().unwrap()) + .compress(&raw) + .unwrap(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let mut gzip = GzipEncoder::new(Vec::new()); + gzip.write_all(&raw).await.unwrap(); + gzip.shutdown().await.unwrap(); + let mut zlib = ZlibEncoder::new(Vec::new()); + zlib.write_all(&raw).await.unwrap(); + zlib.shutdown().await.unwrap(); + for (compressor, encoded) in [ + (None, raw.clone()), + ( + Some(serde_json::json!({"id":"gzip","level":1})), + gzip.into_inner(), + ), + ( + Some(serde_json::json!({"id":"zlib","level":1})), + zlib.into_inner(), + ), + (Some(serde_json::json!({"id":"blosc","cname":"lz4"})), blosc), + ] { + let pipeline = CodecPipeline::from_v2(&compressor).unwrap(); + assert_eq!( + pipeline + .decode_interruptible(encoded, &[raw.len()], 1, || false) + .await + .unwrap(), + CodecDecode::Decoded(raw.clone()) + ); + } + }); + } + + #[test] + fn decoding_is_exact_bounded_and_interruptible() { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let raw = CodecPipeline::raw_v2(); + assert!( + raw.decode_interruptible(vec![1, 2, 3], &[2], 1, || false) + .await + .is_err() + ); + assert_eq!( + raw.decode_interruptible(vec![1, 2], &[2], 1, || true) + .await + .unwrap(), + CodecDecode::Interrupted + ); + + let transposed = v3(serde_json::json!([ + {"name":"transpose","configuration":{"order":[2,1,0]}}, + {"name":"bytes","configuration":{"endian":"little"}} + ])) + .unwrap() + .0; + let mut polls = 0; + assert_eq!( + transposed + .decode_interruptible(vec![0; 12], &[2, 2, 3], 1, || { + polls += 1; + polls > 2 + }) + .await + .unwrap(), + CodecDecode::Interrupted + ); + }); + assert!( + CodecPipeline::raw_v2() + .encoded_read_limit(MAX_DECODED_CHUNK_BYTES + 1) + .is_err() + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/dataset/discovery.rs b/wrappers/src/fdw/zarr_fdw/dataset/discovery.rs new file mode 100644 index 000000000..ac98c3d0a --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/dataset/discovery.rs @@ -0,0 +1,838 @@ +use std::collections::HashSet; + +use serde_json::{Map, Value}; + +use super::super::meta::{ArrayNode, ZarrFormat}; +use super::super::ome::ResolvedOmeLevel; +use super::super::{ZarrFdwError, ZarrFdwResult}; +use super::model::{CoordinateRef, CoordinateSource, Dataset, Dimension, DimensionRole}; + +const ARRAY_DIMENSIONS: &str = "_ARRAY_DIMENSIONS"; + +/// Parse xarray dimension names without choosing strict or tolerant behavior. +/// +/// `None` means the attribute is absent. Callers scanning an array turn that +/// into an error, while metadata inspection preserves it as an unknown hint. +pub(crate) fn parse_named_dimensions( + attrs: &Map, + rank: usize, +) -> Result>, String> { + let Some(value) = attrs.get(ARRAY_DIMENSIONS) else { + return Ok(None); + }; + let values = value + .as_array() + .ok_or_else(|| format!("{ARRAY_DIMENSIONS} must be an array of strings"))?; + let dimensions = values + .iter() + .map(Value::as_str) + .collect::>>() + .ok_or_else(|| format!("{ARRAY_DIMENSIONS} must contain only strings"))?; + if dimensions.len() != rank { + return Err(format!( + "{ARRAY_DIMENSIONS} has {} names but the array rank is {rank}", + dimensions.len() + )); + } + for name in &dimensions { + validate_dimension_name(name)?; + } + let unique = dimensions.iter().copied().collect::>(); + if unique.len() != dimensions.len() { + return Err(format!("{ARRAY_DIMENSIONS} names must be unique")); + } + Ok(Some(dimensions.into_iter().map(str::to_string).collect())) +} + +/// Require valid xarray dimension metadata for a scan array. +pub(crate) fn named_dimensions(node: &ArrayNode, array_path: &str) -> ZarrFdwResult> { + let rank = node.meta.shape.len(); + let legacy = parse_named_dimensions(&node.attributes, rank).map_err(|message| { + ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' has invalid {ARRAY_DIMENSIONS}: {message}" + )) + })?; + match node.format { + ZarrFormat::V2 => legacy.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' must define {ARRAY_DIMENSIONS}" + )) + }), + ZarrFormat::V3 => { + let native = strict_native_dimension_names(node, array_path)?; + if legacy.as_ref().is_some_and(|legacy| legacy != &native) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' has conflicting dimension_names and {ARRAY_DIMENSIONS}" + ))); + } + Ok(native) + } + } +} + +/// Build the scan's format-neutral dataset descriptor from named dimensions +/// and the aligned attributes of their same-group coordinate arrays. +pub(crate) fn named_array_dataset( + array_path: &str, + node: &ArrayNode, + names: &[String], + coordinate_nodes: &[ArrayNode], +) -> ZarrFdwResult { + let meta = &node.meta; + if names.len() != meta.shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' has {} discovered dimensions but rank {}", + names.len(), + meta.shape.len() + ))); + } + if coordinate_nodes.len() != names.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' has attributes for {} coordinate arrays but {} dimensions", + coordinate_nodes.len(), + names.len() + ))); + } + + let coordinate_parent = array_parent_path(array_path); + let dimensions = names + .iter() + .zip(meta.shape.iter()) + .zip(coordinate_nodes) + .map(|((name, &length), coordinate)| { + if coordinate.format != node.format { + return Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' and coordinate array '{name}' use different Zarr formats" + ))); + } + validate_coordinate_dimensions(name, coordinate)?; + Ok(Dimension::new( + name.clone(), + length, + CoordinateSource::Stored(CoordinateRef::new( + coordinate_parent.to_string(), + name.clone(), + )), + infer_dimension_role(name, &coordinate.attributes)?, + )) + }) + .collect::>>()?; + + let time_dimensions = dimensions + .iter() + .filter(|dimension| dimension.semantic_role() == DimensionRole::Time) + .count(); + if time_dimensions > 1 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' has multiple dimensions with the Time semantic role" + ))); + } + + Ok(Dataset::new( + dimensions, + array_path.to_string(), + meta.dtype.clone(), + )) +} + +/// Build the initial OME-Zarr 0.5 execution descriptor. +/// +/// This deliberately supports only a two-dimensional `[y, x]` image. The +/// OME axes are authoritative but must agree exactly with the selected v3 +/// array's native `dimension_names`. Coordinate values are synthesized from +/// the already-composed effective transform in `level`. +pub(crate) fn ome_rank2_dataset( + array_path: &str, + node: &ArrayNode, + level: &ResolvedOmeLevel, +) -> ZarrFdwResult { + if node.format != ZarrFormat::V3 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr 0.5 array '{array_path}' must use Zarr v3" + ))); + } + node.meta.validate()?; + if array_path != level.array_path { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr selected array path '{}' does not match loaded array '{array_path}'", + level.array_path + ))); + } + if node.meta.shape.len() != 2 || level.axes.len() != 2 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr rank-2 execution requires axes [y, x], found array rank {} and {} axes", + node.meta.shape.len(), + level.axes.len() + ))); + } + + let expected_names = ["y", "x"]; + for (axis, expected) in level.axes.iter().zip(expected_names) { + if axis.name != expected || axis.kind.as_deref() != Some("space") { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr rank-2 execution requires axes [y, x] with type 'space', found {:?}", + level + .axes + .iter() + .map(|axis| (&axis.name, axis.kind.as_deref())) + .collect::>() + ))); + } + } + + let names = named_dimensions(node, array_path)?; + let ome_names = level + .axes + .iter() + .map(|axis| axis.name.clone()) + .collect::>(); + if names != ome_names { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr axes {ome_names:?} do not match array dimension_names {names:?} for '{array_path}'" + ))); + } + if level.transform.scale.len() != 2 || level.transform.translation.len() != 2 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr rank-2 transform for '{array_path}' must contain two scale and translation values" + ))); + } + + let dimensions = level + .axes + .iter() + .zip(node.meta.shape.iter().copied()) + .zip( + level + .transform + .scale + .iter() + .copied() + .zip(level.transform.translation.iter().copied()), + ) + .enumerate() + .map(|(axis_index, ((axis, length), (scale, translation)))| { + let semantic_role = if axis_index == 0 { + DimensionRole::SpatialY + } else { + DimensionRole::SpatialX + }; + Dimension::new( + axis.name.clone(), + length, + CoordinateSource::Affine { scale, translation }, + semantic_role, + ) + }) + .collect(); + + Ok(Dataset::new( + dimensions, + array_path.to_string(), + node.meta.dtype.clone(), + )) +} + +fn validate_dimension_name(name: &str) -> Result<(), String> { + if name.is_empty() { + return Err(format!("{ARRAY_DIMENSIONS} names must not be empty")); + } + if name.trim() != name || name.chars().any(char::is_whitespace) { + return Err(format!( + "{ARRAY_DIMENSIONS} names must not contain whitespace" + )); + } + if name.chars().any(char::is_control) { + return Err(format!( + "{ARRAY_DIMENSIONS} names must not contain control characters" + )); + } + if name.contains('/') || name.contains('\\') || matches!(name, "." | "..") { + return Err(format!( + "{ARRAY_DIMENSIONS} names must be same-group array names" + )); + } + Ok(()) +} + +fn strict_native_dimension_names(node: &ArrayNode, array_path: &str) -> ZarrFdwResult> { + let native = node.dimension_names.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' must define dimension_names" + )) + })?; + if native.len() != node.meta.shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' has {} dimension_names but rank {}", + native.len(), + node.meta.shape.len() + ))); + } + let names = native + .iter() + .map(|name| { + name.clone().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' has an unnamed dimension" + )) + }) + }) + .collect::>>()?; + for name in &names { + validate_dimension_name(name).map_err(|message| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' has invalid dimension_names: {message}" + )) + })?; + validate_v3_node_name(name).map_err(|message| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' has invalid dimension_names: {message}" + )) + })?; + } + if names.iter().collect::>().len() != names.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' dimension_names must be unique" + ))); + } + Ok(names) +} + +fn validate_v3_node_name(name: &str) -> Result<(), String> { + if name == "zarr.json" + || name.starts_with("__") + || name.chars().all(|character| character == '.') + { + return Err("Zarr v3 dimension names must be valid node names".to_string()); + } + Ok(()) +} + +fn validate_coordinate_dimensions(name: &str, node: &ArrayNode) -> ZarrFdwResult<()> { + if node.format == ZarrFormat::V3 && node.dimension_names.is_some() { + let dimensions = strict_native_dimension_names(node, name)?; + if dimensions.len() != 1 || dimensions[0] != name { + return Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate array '{name}' declares dimension_names {dimensions:?}, expected [\"{name}\"]" + ))); + } + } + match parse_named_dimensions(&node.attributes, 1) { + Ok(None) => Ok(()), + Ok(Some(dimensions)) if dimensions.first().map(String::as_str) == Some(name) => Ok(()), + Ok(Some(dimensions)) => Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate array '{name}' declares {ARRAY_DIMENSIONS} {dimensions:?}, expected [\"{name}\"]" + ))), + Err(message) => Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate array '{name}' has invalid {ARRAY_DIMENSIONS}: {message}" + ))), + } +} + +fn infer_dimension_role(name: &str, attrs: &Map) -> ZarrFdwResult { + let mut resolved: Option<(DimensionRole, &str)> = None; + for (source, role) in [ + ( + "standard_name", + string_attribute_role(attrs, "standard_name", standard_name_role)?, + ), + ("axis", string_attribute_role(attrs, "axis", axis_role)?), + ("units", string_attribute_role(attrs, "units", units_role)?), + ] { + let Some(role) = role else { + continue; + }; + resolved = Some(match resolved { + None => (role, source), + Some((current, current_source)) => { + let merged = merge_compatible_roles(current, role).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate array '{name}' has conflicting semantic-role signals: {current:?} from {current_source} and {role:?} from {source}" + )) + })?; + (merged, current_source) + } + }); + } + Ok(resolved + .map(|(role, _)| role) + .or_else(|| name_role(name)) + .unwrap_or(DimensionRole::Unknown)) +} + +fn string_attribute_role( + attrs: &Map, + attribute: &str, + classify: fn(&str) -> Option, +) -> ZarrFdwResult> { + let Some(value) = attrs.get(attribute) else { + return Ok(None); + }; + let value = value.as_str().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate attribute '{attribute}' must be a string" + )) + })?; + Ok(classify(value)) +} + +fn standard_name_role(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "projection_x_coordinate" => Some(DimensionRole::SpatialX), + "projection_y_coordinate" => Some(DimensionRole::SpatialY), + "latitude" => Some(DimensionRole::Latitude), + "longitude" => Some(DimensionRole::Longitude), + "time" => Some(DimensionRole::Time), + "depth" | "height" | "altitude" | "air_pressure" | "model_level_number" => { + Some(DimensionRole::Vertical) + } + _ => None, + } +} + +fn axis_role(value: &str) -> Option { + match value.trim().to_ascii_uppercase().as_str() { + "X" => Some(DimensionRole::SpatialX), + "Y" => Some(DimensionRole::SpatialY), + "Z" => Some(DimensionRole::Vertical), + "T" => Some(DimensionRole::Time), + _ => None, + } +} + +fn units_role(value: &str) -> Option { + let normalized = value.trim().to_ascii_lowercase(); + match normalized.as_str() { + "degree_east" | "degrees_east" | "degree_e" | "degrees_e" => Some(DimensionRole::Longitude), + "degree_north" | "degrees_north" | "degree_n" | "degrees_n" => { + Some(DimensionRole::Latitude) + } + _ => normalized + .split_once(" since ") + .map(|(unit, _)| unit) + .filter(|unit| { + matches!( + *unit, + "seconds" + | "milliseconds" + | "microseconds" + | "nanoseconds" + | "minutes" + | "hours" + | "days" + ) + }) + .map(|_| DimensionRole::Time), + } +} + +fn name_role(value: &str) -> Option { + match value.to_ascii_lowercase().as_str() { + "x" => Some(DimensionRole::SpatialX), + "y" => Some(DimensionRole::SpatialY), + "lat" | "latitude" => Some(DimensionRole::Latitude), + "lon" | "longitude" => Some(DimensionRole::Longitude), + "time" => Some(DimensionRole::Time), + "depth" | "height" | "altitude" | "level" | "lev" | "z" => Some(DimensionRole::Vertical), + "band" => Some(DimensionRole::Band), + "channel" => Some(DimensionRole::Channel), + _ => None, + } +} + +fn merge_compatible_roles(left: DimensionRole, right: DimensionRole) -> Option { + match (left, right) { + (left, right) if left == right => Some(left), + (DimensionRole::Latitude, DimensionRole::SpatialY) + | (DimensionRole::SpatialY, DimensionRole::Latitude) => Some(DimensionRole::Latitude), + (DimensionRole::Longitude, DimensionRole::SpatialX) + | (DimensionRole::SpatialX, DimensionRole::Longitude) => Some(DimensionRole::Longitude), + _ => None, + } +} + +fn array_parent_path(array_path: &str) -> &str { + array_path + .rsplit_once('/') + .map(|(parent, _)| parent) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::super::super::ome::{AffineTransform, OmeAxis}; + use super::*; + + fn node(shape: Vec, attributes: Map) -> ArrayNode { + ArrayNode { + format: ZarrFormat::V2, + meta: super::super::super::meta::ArrayMeta { + zarr_format: 2, + chunks: vec![1; shape.len()], + shape, + dtype: ">) -> Vec { + values + .into_iter() + .map(|attributes| node(vec![2], attributes)) + .collect() + } + + fn v3_node( + shape: Vec, + dimension_names: Option>>, + attributes: Map, + ) -> ArrayNode { + let mut node = node(shape, attributes); + node.format = ZarrFormat::V3; + node.meta.zarr_format = 3; + node.meta.chunk_key_encoding = + super::super::super::meta::ChunkKeyEncoding::Default { separator: '/' }; + node.native_dtype = "float32".to_string(); + node.native_codecs = json!([{"name":"bytes","configuration":{"endian":"little"}}]); + node.dimension_names = dimension_names.map(|names| { + names + .into_iter() + .map(|name| name.map(str::to_string)) + .collect() + }); + node + } + + fn attrs(value: Value) -> Map { + value.as_object().cloned().unwrap() + } + + #[test] + fn parses_missing_valid_and_invalid_named_dimensions() { + assert_eq!(parse_named_dimensions(&Map::new(), 2), Ok(None)); + assert_eq!( + parse_named_dimensions( + &attrs(json!({"_ARRAY_DIMENSIONS":["__private", "zarr.json", "..."]})), + 3, + ), + Ok(Some(vec![ + "__private".to_string(), + "zarr.json".to_string(), + "...".to_string(), + ])) + ); + assert!(matches!( + named_dimensions(&node(vec![2, 2], Map::new()), "nested/value"), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("array 'nested/value' must define _ARRAY_DIMENSIONS") + )); + assert_eq!( + parse_named_dimensions(&attrs(json!({"_ARRAY_DIMENSIONS": ["row", "column"]})), 2,), + Ok(Some(vec!["row".to_string(), "column".to_string()])) + ); + + for (value, message) in [ + (json!("x"), "must be an array of strings"), + (json!(["x", 1]), "must contain only strings"), + (json!(["x"]), "has 1 names but the array rank is 2"), + (json!(["x", "x"]), "names must be unique"), + (json!(["x", ""]), "names must not be empty"), + (json!(["x", "bad name"]), "must not contain whitespace"), + (json!(["x", "bad\u{7}name"]), "must not contain control"), + (json!(["x", "../y"]), "must be same-group array names"), + ] { + let error = + parse_named_dimensions(&attrs(json!({"_ARRAY_DIMENSIONS": value})), 2).unwrap_err(); + assert!(error.contains(message), "unexpected error: {error}"); + } + } + + #[test] + fn v3_native_dimensions_are_strict_and_must_match_legacy_hints() { + let valid = v3_node( + vec![2, 2], + Some(vec![Some("row"), Some("column")]), + Map::new(), + ); + assert_eq!( + named_dimensions(&valid, "nested/value").unwrap(), + vec!["row", "column"] + ); + + for invalid in [ + v3_node(vec![2], None, Map::new()), + v3_node(vec![2], Some(vec![None]), Map::new()), + v3_node(vec![2, 2], Some(vec![Some("x")]), Map::new()), + v3_node(vec![2, 2], Some(vec![Some("x"), Some("x")]), Map::new()), + v3_node(vec![2], Some(vec![Some("../x")]), Map::new()), + v3_node( + vec![2], + Some(vec![Some("x")]), + attrs(json!({"_ARRAY_DIMENSIONS":["other"]})), + ), + v3_node(vec![2], Some(vec![Some("zarr.json")]), Map::new()), + v3_node(vec![2], Some(vec![Some("__private")]), Map::new()), + v3_node(vec![2], Some(vec![Some("...")]), Map::new()), + ] { + assert!(named_dimensions(&invalid, "nested/value").is_err()); + } + } + + #[test] + fn v3_coordinates_must_match_format_and_name_themselves() { + let names = vec!["x".to_string()]; + let value = v3_node(vec![2], Some(vec![Some("x")]), Map::new()); + let good = vec![v3_node(vec![2], Some(vec![Some("x")]), Map::new())]; + assert!(named_array_dataset("value", &value, &names, &good).is_ok()); + + let unnamed = vec![v3_node(vec![2], None, Map::new())]; + assert!(named_array_dataset("value", &value, &names, &unnamed).is_ok()); + + let wrong_name = vec![v3_node(vec![2], Some(vec![Some("other")]), Map::new())]; + assert!(named_array_dataset("value", &value, &names, &wrong_name).is_err()); + + let mixed_format = vec![node(vec![2], Map::new())]; + assert!(named_array_dataset("value", &value, &names, &mixed_format).is_err()); + } + + fn ome_level(axes: Vec) -> ResolvedOmeLevel { + ResolvedOmeLevel { + group_path: "nested/image".to_string(), + multiscale_index: 0, + multiscale_name: Some("image".to_string()), + level_index: 0, + array_path: "nested/image/0".to_string(), + axes, + transform: AffineTransform { + scale: vec![2.0, 3.0], + translation: vec![10.0, 100.0], + }, + warnings: Vec::new(), + } + } + + fn ome_axes() -> Vec { + ["y", "x"] + .into_iter() + .map(|name| OmeAxis { + name: name.to_string(), + kind: Some("space".to_string()), + unit: Some("micrometer".to_string()), + }) + .collect() + } + + #[test] + fn constructs_rank2_ome_dataset_with_affine_coordinates() { + let value = v3_node(vec![2, 3], Some(vec![Some("y"), Some("x")]), Map::new()); + let dataset = ome_rank2_dataset("nested/image/0", &value, &ome_level(ome_axes())).unwrap(); + + assert_eq!(dataset.axis_names(), vec!["y", "x"]); + assert_eq!( + dataset.dimensions()[0].semantic_role(), + DimensionRole::SpatialY + ); + assert_eq!( + dataset.dimensions()[1].semantic_role(), + DimensionRole::SpatialX + ); + assert_eq!(dataset.dimensions()[0].stored_coordinate(), None); + assert!(matches!( + dataset.dimensions()[0].coordinate_source(), + CoordinateSource::Affine { + scale: 2.0, + translation: 10.0 + } + )); + assert!(matches!( + dataset.dimensions()[1].coordinate_source(), + CoordinateSource::Affine { + scale: 3.0, + translation: 100.0 + } + )); + } + + #[test] + fn ome_dataset_rejects_format_rank_axes_path_and_dimension_mismatches() { + let valid = v3_node(vec![2, 3], Some(vec![Some("y"), Some("x")]), Map::new()); + let mut cases = Vec::new(); + cases.push(( + node(vec![2, 3], Map::new()), + ome_level(ome_axes()), + "must use Zarr v3", + )); + cases.push(( + v3_node( + vec![1, 2, 3], + Some(vec![Some("z"), Some("y"), Some("x")]), + Map::new(), + ), + ome_level(ome_axes()), + "rank-2 execution requires axes [y, x]", + )); + let mut wrong_axes = ome_axes(); + wrong_axes.swap(0, 1); + cases.push((valid.clone(), ome_level(wrong_axes), "requires axes [y, x]")); + cases.push(( + v3_node( + vec![2, 3], + Some(vec![Some("row"), Some("column")]), + Map::new(), + ), + ome_level(ome_axes()), + "do not match array dimension_names", + )); + let mut wrong_path = ome_level(ome_axes()); + wrong_path.array_path = "nested/image/other".to_string(); + cases.push((valid, wrong_path, "does not match loaded array")); + + for (node, level, phrase) in cases { + let error = ome_rank2_dataset("nested/image/0", &node, &level).unwrap_err(); + assert!( + error.to_string().contains(phrase), + "unexpected error: {error}" + ); + } + } + + #[test] + fn constructs_arbitrary_rank_nested_dataset() { + let names = ["forecast_time", "level", "band", "channel"] + .map(str::to_string) + .to_vec(); + let coordinate_attrs = vec![ + attrs(json!({"standard_name": "time", "axis": "T"})), + attrs(json!({"axis": "Z"})), + Map::new(), + Map::new(), + ]; + let value = node(vec![2, 5, 6, 1], attrs(json!({"_ARRAY_DIMENSIONS": names}))); + let coordinate_nodes = coordinate_nodes(coordinate_attrs); + let dataset = + named_array_dataset("nested/generic4d", &value, &names, &coordinate_nodes).unwrap(); + + assert_eq!(dataset.variable().dimensions(), names); + assert_eq!(dataset.dimensions()[0].semantic_role(), DimensionRole::Time); + assert_eq!( + dataset.dimensions()[1].semantic_role(), + DimensionRole::Vertical + ); + assert_eq!(dataset.dimensions()[2].semantic_role(), DimensionRole::Band); + assert_eq!( + dataset.dimensions()[3].semantic_role(), + DimensionRole::Channel + ); + assert_eq!( + dataset.dimensions()[3] + .stored_coordinate() + .unwrap() + .parent(), + "nested" + ); + } + + #[test] + fn infers_specific_and_generic_roles_without_renaming_dimensions() { + let cases = [ + ("x", json!({"axis": "X"}), DimensionRole::SpatialX), + ("y", json!({"axis": "Y"}), DimensionRole::SpatialY), + ( + "lon", + json!({"standard_name": "longitude", "axis": "X"}), + DimensionRole::Longitude, + ), + ( + "lat", + json!({"standard_name": "latitude", "axis": "Y"}), + DimensionRole::Latitude, + ), + ("forecast_time", json!({"axis": "T"}), DimensionRole::Time), + ( + "valid_time", + json!({"units": "hours since 2000-01-01"}), + DimensionRole::Time, + ), + ( + "east_coordinate", + json!({"units": "degrees_east"}), + DimensionRole::Longitude, + ), + ("depth", json!({}), DimensionRole::Vertical), + ("level", json!({}), DimensionRole::Vertical), + ("band", json!({}), DimensionRole::Band), + ("channel", json!({}), DimensionRole::Channel), + ("sample", json!({}), DimensionRole::Unknown), + ("time", json!({"axis": "Z"}), DimensionRole::Vertical), + ( + "band", + json!({"standard_name": "time"}), + DimensionRole::Time, + ), + ]; + for (name, metadata, expected) in cases { + assert_eq!( + infer_dimension_role(name, &attrs(metadata)).unwrap(), + expected + ); + } + } + + #[test] + fn rejects_conflicting_roles_and_multiple_time_dimensions() { + assert!(matches!( + infer_dimension_role( + "latitude", + &attrs(json!({"standard_name": "latitude", "axis": "X"})) + ), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("conflicting semantic-role signals") + )); + assert!(matches!( + infer_dimension_role("east", &attrs(json!({"axis": "X", "units": "degrees_north"}))), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("conflicting semantic-role signals") + )); + + let names = vec!["time".to_string(), "forecast_time".to_string()]; + let value = node(vec![2, 2], attrs(json!({"_ARRAY_DIMENSIONS": names}))); + let coordinate_nodes = coordinate_nodes(vec![Map::new(), attrs(json!({"axis": "T"}))]); + assert!(matches!( + named_array_dataset( + "multiple_times", + &value, + &names, + &coordinate_nodes, + ), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("multiple dimensions with the Time semantic role") + )); + } + + #[test] + fn validates_optional_coordinate_dimension_hint() { + let names = vec!["level".to_string()]; + let value = node(vec![5], attrs(json!({"_ARRAY_DIMENSIONS": names}))); + let invalid = vec![node( + vec![5], + attrs(json!({"_ARRAY_DIMENSIONS": ["other"]})), + )]; + assert!(named_array_dataset("value", &value, &names, &invalid).is_err()); + let missing = vec![node(vec![5], Map::new())]; + assert!(named_array_dataset("value", &value, &names, &missing).is_ok()); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/dataset/mod.rs b/wrappers/src/fdw/zarr_fdw/dataset/mod.rs new file mode 100644 index 000000000..4fec7090a --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/dataset/mod.rs @@ -0,0 +1,12 @@ +//! Generic dataset descriptors and metadata adapters for `zarr_fdw`. +//! +//! Scan execution consumes xarray named dimensions and same-group coordinate +//! metadata through this format-neutral model. + +mod discovery; +pub(crate) mod model; + +pub(crate) use discovery::{ + named_array_dataset, named_dimensions, ome_rank2_dataset, parse_named_dimensions, +}; +pub(crate) use model::{CoordinateSource, Dataset, DimensionRole}; diff --git a/wrappers/src/fdw/zarr_fdw/dataset/model.rs b/wrappers/src/fdw/zarr_fdw/dataset/model.rs new file mode 100644 index 000000000..701977f2b --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/dataset/model.rs @@ -0,0 +1,172 @@ +/// Dataset-level description consumed by the scan executor. +/// +/// The model is deliberately independent of the legacy `x`/`y`/`time` +/// profile. Metadata adapters assign names and semantic roles; the executor +/// works from those descriptors. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Dataset { + dimensions: Vec, + variable: Variable, +} + +impl Dataset { + pub(super) fn new(dimensions: Vec, variable_path: String, dtype: String) -> Self { + let variable_dimensions = dimensions + .iter() + .map(|dimension| dimension.name().to_string()) + .collect(); + Self { + dimensions, + variable: Variable::new(variable_path, dtype, variable_dimensions), + } + } + + pub(crate) fn dimensions(&self) -> &[Dimension] { + &self.dimensions + } + + pub(crate) fn variable(&self) -> &Variable { + &self.variable + } + + pub(crate) fn axis_names(&self) -> Vec { + self.dimensions + .iter() + .map(|dimension| dimension.name().to_string()) + .collect() + } + + pub(crate) fn is_dimension(&self, name: &str) -> bool { + self.dimensions + .iter() + .any(|dimension| dimension.name() == name) + } + + pub(crate) fn dimension(&self, name: &str) -> Option<&Dimension> { + self.dimensions + .iter() + .find(|dimension| dimension.name() == name) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Dimension { + name: String, + length: u64, + coordinate: CoordinateSource, + semantic_role: DimensionRole, +} + +impl Dimension { + pub(super) fn new( + name: String, + length: u64, + coordinate: CoordinateSource, + semantic_role: DimensionRole, + ) -> Self { + Self { + name, + length, + coordinate, + semantic_role, + } + } + + pub(crate) fn name(&self) -> &str { + &self.name + } + + pub(crate) fn length(&self) -> u64 { + self.length + } + + /// Coordinate source selected by the metadata adapter. + pub(crate) fn coordinate_source(&self) -> &CoordinateSource { + &self.coordinate + } + + #[cfg(test)] + pub(crate) fn stored_coordinate(&self) -> Option<&CoordinateRef> { + match &self.coordinate { + CoordinateSource::Stored(coordinate) => Some(coordinate), + CoordinateSource::Affine { .. } => None, + } + } + + pub(crate) fn semantic_role(&self) -> DimensionRole { + self.semantic_role + } +} + +/// Format-neutral source for one dimension's coordinate values. +/// +/// Existing named-array datasets retain their same-group one-dimensional +/// coordinate arrays. OME-Zarr 0.5 levels instead synthesize rectilinear +/// coordinates from their normalized scale and translation. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum CoordinateSource { + Stored(CoordinateRef), + Affine { scale: f64, translation: f64 }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CoordinateRef { + parent: String, + name: String, +} + +impl CoordinateRef { + pub(super) fn new(parent: String, name: String) -> Self { + Self { parent, name } + } + + pub(crate) fn parent(&self) -> &str { + &self.parent + } + + pub(crate) fn name(&self) -> &str { + &self.name + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum DimensionRole { + SpatialX, + SpatialY, + Latitude, + Longitude, + Vertical, + Time, + Band, + Channel, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Variable { + path: String, + dtype: String, + dimensions: Vec, +} + +impl Variable { + fn new(path: String, dtype: String, dimensions: Vec) -> Self { + Self { + path, + dtype, + dimensions, + } + } + + pub(crate) fn path(&self) -> &str { + &self.path + } + + pub(crate) fn dtype(&self) -> &str { + &self.dtype + } + + pub(crate) fn dimensions(&self) -> &[String] { + &self.dimensions + } +} diff --git a/wrappers/src/fdw/zarr_fdw/decode.rs b/wrappers/src/fdw/zarr_fdw/decode.rs new file mode 100644 index 000000000..2310f1f23 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/decode.rs @@ -0,0 +1,584 @@ +//! Primitive dtype, fill, and raw chunk byte decoding. +//! +//! - dtypes: `f4`, `f8`, `i1`, `i2`, `i4`, `i8` (signed little/big endian) +//! - byte order: `<` little-endian, `>` big-endian + +use super::{ZarrFdwError, ZarrFdwResult}; + +/// Parsed numpy-style dtype string, e.g. ` ZarrFdwResult { + // numpy dtype: [byteorder], e.g. "i2" + let (byte_order, ty, size) = numeric_dtype_parts(dtype)?; + // For 1-byte types byte order is irrelevant. + let big_endian = byte_order == '>'; + let dt = match (ty, size) { + ('f', 4) => DType::F32, + ('f', 8) => DType::F64, + ('i', 1) => DType::I8, + ('i', 2) => DType::I16, + ('i', 4) => DType::I32, + ('i', 8) => DType::I64, + _ => return Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())), + }; + if big_endian && dt != DType::I8 { + // reject until we implement endian-swapped reads + return Err(ZarrFdwError::UnsupportedDataType(format!( + "{dtype} (big-endian numeric types not supported yet)" + ))); + } + Ok(dt) + } + + pub fn itemsize(self) -> usize { + match self { + DType::F32 | DType::I32 => 4, + DType::F64 | DType::I64 => 8, + DType::I8 => 1, + DType::I16 => 2, + } + } + + fn name(self) -> &'static str { + match self { + DType::F32 => "f4", + DType::F64 => "f8", + DType::I8 => "i1", + DType::I16 => "i2", + DType::I32 => "i4", + DType::I64 => "i8", + } + } +} + +/// Decode one supported primitive value into the `f64` domain used by CF +/// scale/offset processing. Raw scans keep their exact PostgreSQL primitive +/// type; this conversion is only used when scientific decoding is enabled. +pub fn value_bytes_to_f64(dtype: DType, bytes: &[u8]) -> ZarrFdwResult { + let too_short = |needed: usize| { + ZarrFdwError::ReadError(std::io::Error::other(format!( + "chunk cell data has {} bytes, expected exactly {needed}", + bytes.len() + ))) + }; + Ok(match dtype { + DType::F32 => f32::from_le_bytes(bytes.try_into().map_err(|_| too_short(4))?) as f64, + DType::F64 => f64::from_le_bytes(bytes.try_into().map_err(|_| too_short(8))?), + DType::I8 => bytes.first().copied().ok_or_else(|| too_short(1))? as i8 as f64, + DType::I16 => i16::from_le_bytes(bytes.try_into().map_err(|_| too_short(2))?) as f64, + DType::I32 => i32::from_le_bytes(bytes.try_into().map_err(|_| too_short(4))?) as f64, + DType::I64 => i64::from_le_bytes(bytes.try_into().map_err(|_| too_short(8))?) as f64, + }) +} + +fn invalid_fill(dtype: &str, reason: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!( + "fill_value is invalid for dtype '{dtype}': {}", + reason.into() + )) +} + +fn float_fill(dtype: &str, value: &serde_json::Value) -> ZarrFdwResult { + match value { + serde_json::Value::Number(number) => number + .as_f64() + .ok_or_else(|| invalid_fill(dtype, "expected a numeric value")), + serde_json::Value::String(value) => match value.as_str() { + "NaN" => Ok(f64::NAN), + "Infinity" => Ok(f64::INFINITY), + "-Infinity" => Ok(f64::NEG_INFINITY), + _ => Err(invalid_fill( + dtype, + "expected a number, 'NaN', 'Infinity', or '-Infinity'", + )), + }, + _ => Err(invalid_fill( + dtype, + "expected a number, 'NaN', 'Infinity', or '-Infinity'", + )), + } +} + +fn signed_integer_fill( + dtype: &str, + value: &serde_json::Value, + min: i64, + max: i64, +) -> ZarrFdwResult { + let serde_json::Value::Number(number) = value else { + return Err(invalid_fill(dtype, "expected an integer")); + }; + let parsed = if let Some(value) = number.as_i64() { + value + } else if let Some(value) = number.as_u64() { + i64::try_from(value).map_err(|_| invalid_fill(dtype, "integer is out of range"))? + } else { + let value = number + .as_f64() + .ok_or_else(|| invalid_fill(dtype, "expected an integer"))?; + // The exclusive upper bound avoids accepting 2^63 after an imprecise + // f64 conversion. Plain JSON integers at i64::MAX take the as_i64 path. + const I64_EXCLUSIVE_UPPER: f64 = 9_223_372_036_854_775_808.0; + if !value.is_finite() + || value.fract() != 0.0 + || value < i64::MIN as f64 + || value >= I64_EXCLUSIVE_UPPER + { + return Err(invalid_fill(dtype, "expected an in-range integer")); + } + value as i64 + }; + if parsed < min || parsed > max { + return Err(invalid_fill(dtype, "integer is out of range")); + } + Ok(parsed) +} + +fn unsigned_integer_fill(dtype: &str, value: &serde_json::Value, max: u64) -> ZarrFdwResult { + let serde_json::Value::Number(number) = value else { + return Err(invalid_fill(dtype, "expected a non-negative integer")); + }; + let parsed = if let Some(value) = number.as_u64() { + value + } else if let Some(value) = number.as_i64() { + u64::try_from(value).map_err(|_| invalid_fill(dtype, "expected a non-negative integer"))? + } else { + let value = number + .as_f64() + .ok_or_else(|| invalid_fill(dtype, "expected a non-negative integer"))?; + // As above, plain u64::MAX JSON integers use the exact as_u64 path. + const U64_EXCLUSIVE_UPPER: f64 = 18_446_744_073_709_551_616.0; + if !value.is_finite() + || value.fract() != 0.0 + || !(0.0..U64_EXCLUSIVE_UPPER).contains(&value) + { + return Err(invalid_fill( + dtype, + "expected an in-range non-negative integer", + )); + } + value as u64 + }; + if parsed > max { + return Err(invalid_fill(dtype, "integer is out of range")); + } + Ok(parsed) +} + +/// Parse a cube array's scalar fill into its decoded little-endian bytes. +/// `None` preserves Zarr's explicit-null/undefined missing-chunk semantics. +pub fn fill_value_bytes(dtype: DType, value: &serde_json::Value) -> ZarrFdwResult>> { + if value.is_null() { + return Ok(None); + } + let bytes = match dtype { + DType::F32 => { + let value = float_fill(dtype.name(), value)?; + let narrowed = value as f32; + if value.is_finite() && !narrowed.is_finite() { + return Err(invalid_fill(dtype.name(), "number is out of range")); + } + narrowed.to_le_bytes().to_vec() + } + DType::F64 => float_fill(dtype.name(), value)?.to_le_bytes().to_vec(), + DType::I8 => (signed_integer_fill(dtype.name(), value, i8::MIN as i64, i8::MAX as i64)? + as i8) + .to_le_bytes() + .to_vec(), + DType::I16 => (signed_integer_fill(dtype.name(), value, i16::MIN as i64, i16::MAX as i64)? + as i16) + .to_le_bytes() + .to_vec(), + DType::I32 => (signed_integer_fill(dtype.name(), value, i32::MIN as i64, i32::MAX as i64)? + as i32) + .to_le_bytes() + .to_vec(), + DType::I64 => signed_integer_fill(dtype.name(), value, i64::MIN, i64::MAX)? + .to_le_bytes() + .to_vec(), + }; + Ok(Some(bytes)) +} + +fn numeric_dtype_parts(dtype: &str) -> ZarrFdwResult<(char, char, usize)> { + let mut chars = dtype.chars(); + let byte_order = chars + .next() + .ok_or_else(|| ZarrFdwError::UnsupportedDataType(dtype.to_string()))?; + if !matches!(byte_order, '=' | '|' | '<' | '>') { + return Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())); + } + let kind = chars + .next() + .ok_or_else(|| ZarrFdwError::UnsupportedDataType(dtype.to_string()))?; + let size_text = chars.collect::(); + if size_text.is_empty() || !size_text.chars().all(|c| c.is_ascii_digit()) { + return Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())); + } + let size = size_text + .parse() + .map_err(|_| ZarrFdwError::UnsupportedDataType(dtype.to_string()))?; + Ok((byte_order, kind, size)) +} + +fn exact_i64_coordinate(dtype: &str, value: i64) -> ZarrFdwResult { + let converted = value as f64; + if converted as i128 != value as i128 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate integer {value} for dtype '{dtype}' cannot be represented exactly as double precision" + ))); + } + Ok(converted) +} + +fn exact_u64_coordinate(dtype: &str, value: u64) -> ZarrFdwResult { + let converted = value as f64; + if converted as u128 != value as u128 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate integer {value} for dtype '{dtype}' cannot be represented exactly as double precision" + ))); + } + Ok(converted) +} + +/// Decoded byte width of a supported coordinate scalar. +pub fn coordinate_itemsize(dtype: &str) -> ZarrFdwResult { + let (_, kind, size) = numeric_dtype_parts(dtype)?; + match (kind, size) { + ('f', 4 | 8) | ('i' | 'u', 1 | 2 | 4 | 8) => Ok(size), + _ => Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())), + } +} + +/// Parse a coordinate array's scalar fill into the f64 representation used by +/// the scan. Coordinate dtypes retain the broader signed/unsigned and endian +/// coverage of [`coord_bytes_to_f64`]. +pub fn coord_fill_value_to_f64( + dtype: &str, + value: &serde_json::Value, +) -> ZarrFdwResult> { + if value.is_null() { + return Ok(None); + } + let (_, kind, size) = numeric_dtype_parts(dtype)?; + let parsed = match (kind, size) { + ('f', 4) => { + let value = float_fill(dtype, value)?; + let narrowed = value as f32; + if value.is_finite() && !narrowed.is_finite() { + return Err(invalid_fill(dtype, "number is out of range")); + } + narrowed as f64 + } + ('f', 8) => float_fill(dtype, value)?, + ('i', 1) => exact_i64_coordinate( + dtype, + signed_integer_fill(dtype, value, i8::MIN as i64, i8::MAX as i64)?, + )?, + ('i', 2) => exact_i64_coordinate( + dtype, + signed_integer_fill(dtype, value, i16::MIN as i64, i16::MAX as i64)?, + )?, + ('i', 4) => exact_i64_coordinate( + dtype, + signed_integer_fill(dtype, value, i32::MIN as i64, i32::MAX as i64)?, + )?, + ('i', 8) => exact_i64_coordinate( + dtype, + signed_integer_fill(dtype, value, i64::MIN, i64::MAX)?, + )?, + ('u', 1) => { + exact_u64_coordinate(dtype, unsigned_integer_fill(dtype, value, u8::MAX as u64)?)? + } + ('u', 2) => { + exact_u64_coordinate(dtype, unsigned_integer_fill(dtype, value, u16::MAX as u64)?)? + } + ('u', 4) => { + exact_u64_coordinate(dtype, unsigned_integer_fill(dtype, value, u32::MAX as u64)?)? + } + ('u', 8) => exact_u64_coordinate(dtype, unsigned_integer_fill(dtype, value, u64::MAX)?)?, + _ => return Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())), + }; + Ok(Some(parsed)) +} + +/// Interpret raw `data` bytes as coordinate values (`f64`) using the numpy +/// style `dtype` string of a coordinate array (e.g. `i4`, `|u2`). +/// +/// Coordinates can be stored as floats or (un)signed ints in either byte +/// order. Integer values must be exactly representable as `f64`, because the +/// resulting value is used for both identity and predicate pruning. +pub fn coord_bytes_to_f64(dtype: &str, data: &[u8]) -> ZarrFdwResult> { + if data.is_empty() { + return Ok(Vec::new()); + } + let (byte_order, kind, size) = numeric_dtype_parts(dtype)?; + let item = coordinate_itemsize(dtype)?; + + let read = |b: &[u8]| -> ZarrFdwResult { + match (kind, size) { + ('f', 4) => Ok(f32_bytes(b, byte_order) as f64), + ('f', 8) => Ok(f64_bytes(b, byte_order)), + ('i', 1) => exact_i64_coordinate(dtype, b[0] as i8 as i64), + ('i', 2) => exact_i64_coordinate(dtype, i16_bytes(b, byte_order) as i64), + ('i', 4) => exact_i64_coordinate(dtype, i32_bytes(b, byte_order) as i64), + ('i', 8) => exact_i64_coordinate(dtype, i64_bytes(b, byte_order)), + ('u', 1) => exact_u64_coordinate(dtype, b[0] as u64), + ('u', 2) => exact_u64_coordinate(dtype, u16_bytes(b, byte_order) as u64), + ('u', 4) => exact_u64_coordinate(dtype, u32_bytes(b, byte_order) as u64), + ('u', 8) => exact_u64_coordinate(dtype, u64_bytes(b, byte_order)), + _ => Err(ZarrFdwError::UnsupportedDataType(dtype.to_string())), + } + }; + + if !data.len().is_multiple_of(item) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate data length {} is not a multiple of dtype item size {item}", + data.len() + ))); + } + let mut out = Vec::with_capacity(data.len() / item); + for b in data.chunks(item) { + out.push(read(b)?); + } + Ok(out) +} + +fn f32_bytes(b: &[u8], byte_order: char) -> f32 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + f32::from_be_bytes(a) + } else { + f32::from_le_bytes(a) + } +} +fn f64_bytes(b: &[u8], byte_order: char) -> f64 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + f64::from_be_bytes(a) + } else { + f64::from_le_bytes(a) + } +} +fn i16_bytes(b: &[u8], byte_order: char) -> i16 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + i16::from_be_bytes(a) + } else { + i16::from_le_bytes(a) + } +} +fn i32_bytes(b: &[u8], byte_order: char) -> i32 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + i32::from_be_bytes(a) + } else { + i32::from_le_bytes(a) + } +} +fn i64_bytes(b: &[u8], byte_order: char) -> i64 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + i64::from_be_bytes(a) + } else { + i64::from_le_bytes(a) + } +} +fn u16_bytes(b: &[u8], byte_order: char) -> u16 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + u16::from_be_bytes(a) + } else { + u16::from_le_bytes(a) + } +} +fn u32_bytes(b: &[u8], byte_order: char) -> u32 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + u32::from_be_bytes(a) + } else { + u32::from_le_bytes(a) + } +} +fn u64_bytes(b: &[u8], byte_order: char) -> u64 { + let a = b.try_into().unwrap(); + if byte_order == '>' { + u64::from_be_bytes(a) + } else { + u64::from_le_bytes(a) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_dtypes() { + assert_eq!(DType::parse("i4").is_err()); + assert!(DType::parse("i2", &serde_json::json!(-12)).unwrap(), + Some(-12.0) + ); + assert_eq!( + coord_fill_value_to_f64("i2", &data).unwrap(), vec![255.0]); + } + + #[test] + fn coord_f64_unsigned() { + let data = 300u16.to_le_bytes(); + assert_eq!(coord_bytes_to_f64(", + variable: Option, + zarr_format: Option, + shape: Option, + dimensions: Option>, + dtype: Option, + chunks: Option, + codecs: Option, + units: Option, + fill_value: Option, + scale_factor: Option, + add_offset: Option, + crs: Option, + calendar: Option, + attributes: Value, + warnings: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +struct MultiscaleInspectionRow { + group_path: String, + multiscale_index: i64, + multiscale_name: Option, + level_index: i64, + array_path: String, + axes: Value, + shape: Value, + chunks: Value, + dtype: String, + codecs: Value, + scale: Vec, + translation: Vec, + supported: bool, + warnings: Vec, +} + +#[derive(Debug)] +struct OmeArrayInspection { + shape: Vec, + chunks: Value, + dimensions: Vec, + dtype: String, + codecs: Value, + execution_warning: Option, +} + +#[derive(Debug, Default)] +struct OmeInspectionNodes { + groups: Vec<(String, Map)>, + arrays: HashMap, +} + +impl MultiscaleInspectionRow { + #[allow(clippy::type_complexity)] + fn sql_row( + self, + ) -> ( + String, + i64, + Option, + i64, + String, + JsonB, + JsonB, + JsonB, + String, + JsonB, + Vec, + Vec, + bool, + Vec, + ) { + ( + self.group_path, + self.multiscale_index, + self.multiscale_name, + self.level_index, + self.array_path, + JsonB(self.axes), + JsonB(self.shape), + JsonB(self.chunks), + self.dtype, + JsonB(self.codecs), + self.scale, + self.translation, + self.supported, + self.warnings, + ) + } +} + +impl InspectionRow { + // pgrx represents a named set-returning SQL record as an explicit tuple. + #[allow(clippy::type_complexity)] + fn sql_row( + self, + ) -> ( + String, + String, + Option, + Option, + Option, + Option, + Option>, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + JsonB, + Vec, + ) { + ( + self.path, + self.kind, + self.group_path, + self.variable, + self.zarr_format, + self.shape.map(JsonB), + self.dimensions, + self.dtype, + self.chunks.map(JsonB), + self.codecs.map(JsonB), + self.units, + self.fill_value.map(JsonB), + self.scale_factor, + self.add_offset, + self.crs.map(JsonB), + self.calendar, + JsonB(self.attributes), + self.warnings, + ) + } +} + +// pgrx requires the complete named SQL record in the exported signature. +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace)] +fn zarr_inspect( + server_name: &str, +) -> TableIterator< + 'static, + ( + name!(path, String), + name!(kind, String), + name!(group_path, Option), + name!(variable, Option), + name!(zarr_format, Option), + name!(shape, Option), + name!(dimensions, Option>), + name!(dtype, Option), + name!(chunks, Option), + name!(codecs, Option), + name!(units, Option), + name!(fill_value, Option), + name!(scale_factor, Option), + name!(add_offset, Option), + name!(crs, Option), + name!(calendar, Option), + name!(attributes, JsonB), + name!(warnings, Vec), + ), +> { + let rows = inspect_server(server_name) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(InspectionRow::sql_row)) +} + +/// Discover OME-Zarr 0.5 multiscale levels without reading chunk payloads. +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace)] +fn zarr_multiscales( + server_name: &str, +) -> TableIterator< + 'static, + ( + name!(group_path, String), + name!(multiscale_index, i64), + name!(multiscale_name, Option), + name!(level_index, i64), + name!(array_path, String), + name!(axes, JsonB), + name!(shape, JsonB), + name!(chunks, JsonB), + name!(dtype, String), + name!(codecs, JsonB), + name!(scale, Vec), + name!(translation, Vec), + name!(supported, bool), + name!(warnings, Vec), + ), +> { + let rows = inspect_multiscales(server_name) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(MultiscaleInspectionRow::sql_row)) +} + +fn inspect_multiscales(server_name: &str) -> ZarrFdwResult> { + let server = load_foreign_server(server_name)?; + let store = ZarrStore::new(&server)?; + store.require_listing()?; + inspect_multiscales_store(&store) +} + +fn inspect_multiscales_store(store: &ZarrStore) -> ZarrFdwResult> { + let OmeInspectionNodes { groups, arrays } = inspect_ome_v3_nodes(store)?; + let mut levels = Vec::new(); + let mut derived_bytes = 0usize; + + for (group_path, attributes) in groups { + if !has_ome_multiscales(&attributes) { + continue; + } + let multiscales = parse_ome_05_multiscales(&group_path, &attributes)?; + for (multiscale_index, multiscale) in multiscales.into_iter().enumerate() { + let axes = multiscale.axes; + let axis_names = axes + .iter() + .map(|axis| axis.name.clone()) + .collect::>(); + let axes_json = serde_json::to_value(&axes).map_err(ZarrFdwError::from)?; + let mut previous_shape: Option<&[u64]> = None; + for (level_index, level) in multiscale.levels.into_iter().enumerate() { + let array_path = join_key(&group_path, &level.relative_path); + let row = arrays.get(array_path.as_str()).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset path '{}' does not resolve to a discovered array", + level.relative_path + )) + })?; + if row.dimensions != axis_names { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr axes {axis_names:?} do not match array dimension_names {:?} for '{array_path}'", + row.dimensions + ))); + } + let rank = row.shape.len(); + if rank != axes.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{array_path}' has rank {rank}, but its multiscale declares {} axes", + axes.len() + ))); + } + validate_resolution_order( + &group_path, + multiscale_index, + level_index, + previous_shape, + &row.shape, + )?; + previous_shape = Some(&row.shape); + + let mut warnings = multiscale.warnings.clone(); + if let Some(warning) = &row.execution_warning { + warnings.push(warning.clone()); + } + let supported = ome_level_support( + &axes, + &row.shape, + &level.effective_transform, + row.execution_warning.is_none(), + &mut warnings, + ); + let inspection_row = MultiscaleInspectionRow { + group_path: if group_path.is_empty() { + "/".to_string() + } else { + group_path.clone() + }, + multiscale_index: i64::try_from(multiscale_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "OME-Zarr multiscale index exceeds bigint".to_string(), + ) + })?, + multiscale_name: multiscale.name.clone(), + level_index: i64::try_from(level_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "OME-Zarr level index exceeds bigint".to_string(), + ) + })?, + array_path, + axes: axes_json.clone(), + shape: json!(row.shape), + chunks: row.chunks.clone(), + dtype: row.dtype.clone(), + codecs: row.codecs.clone(), + scale: level.effective_transform.scale, + translation: level.effective_transform.translation, + supported, + warnings, + }; + derived_bytes = checked_multiscale_derived_bytes( + derived_bytes, + &inspection_row, + MAX_MULTISCALE_DERIVED_BYTES, + )?; + levels.push(inspection_row); + } + } + } + levels.sort_by(|left, right| { + (&left.group_path, left.multiscale_index, left.level_index).cmp(&( + &right.group_path, + right.multiscale_index, + right.level_index, + )) + }); + Ok(levels) +} + +fn has_ome_multiscales(attributes: &Map) -> bool { + attributes + .get("ome") + .and_then(Value::as_object) + .is_some_and(|ome| ome.contains_key("multiscales")) +} + +fn inspect_ome_v3_nodes(store: &ZarrStore) -> ZarrFdwResult { + let mut pending = VecDeque::from([(String::new(), 0usize)]); + let mut discovered = HashSet::from([String::new()]); + let mut groups = Vec::new(); + let mut arrays = HashMap::new(); + let mut metadata_bytes = 0usize; + let mut list_pages = 0usize; + + while let Some((path, depth)) = pending.pop_front() { + let key = metadata_key(&path, "zarr.json"); + let Some(bytes) = read_optional_metadata(store, &key, &mut metadata_bytes)? else { + if path.is_empty() { + return Ok(OmeInspectionNodes { groups, arrays }); + } + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 node '{}' must contain explicit zarr.json metadata", + display_path(&path) + ))); + }; + let v2_array = + read_optional_metadata(store, &metadata_key(&path, ".zarray"), &mut metadata_bytes)?; + let v2_group = + read_optional_metadata(store, &metadata_key(&path, ".zgroup"), &mut metadata_bytes)?; + reject_dual_metadata(&path, true, v2_array.is_some(), v2_group.is_some())?; + + match raw_v3_node_type(&bytes, &key)?.as_str() { + "array" => { + let array = parse_ome_array_inspection(&bytes, &path)?; + arrays.insert(path, array); + continue; + } + "group" => { + let group = match parse_v3_node(&bytes)? { + NodeMeta::Group(group) => group, + NodeMeta::Array(_) => unreachable!("raw node_type and strict parser disagree"), + }; + validate_optional_ome_05_attributes(&display_path(&path), Some(&group.attributes))?; + groups.push((path.clone(), group.attributes)); + } + other => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "'{}' node_type must be 'array' or 'group', got '{other}'", + display_path(&key) + ))); + } + } + + let mut child_prefixes = Vec::new(); + let mut continuation_token = None; + loop { + list_pages = list_pages.checked_add(1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "multiscale inspection list-page count overflowed".to_string(), + ) + })?; + if list_pages > MAX_INSPECTION_LIST_PAGES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "multiscale inspection exceeds the limit of {MAX_INSPECTION_LIST_PAGES} storage list pages" + ))); + } + let page = store.list_directory_page_sync(&path, continuation_token)?; + for child in page.child_prefixes { + if discovered.insert(child.clone()) { + if discovered.len() > MAX_INSPECTION_NODES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "multiscale inspection exceeds the limit of {MAX_INSPECTION_NODES} Zarr nodes" + ))); + } + child_prefixes.push(child); + } + } + continuation_token = page.next_continuation_token; + if continuation_token.is_none() { + break; + } + } + if !child_prefixes.is_empty() && depth >= MAX_INSPECTION_DEPTH { + return Err(ZarrFdwError::InvalidMetadata(format!( + "multiscale inspection exceeds the maximum group depth of {MAX_INSPECTION_DEPTH} at '{}'", + display_path(&path) + ))); + } + child_prefixes.sort(); + for child in child_prefixes { + pending.push_back((child, depth + 1)); + } + } + + Ok(OmeInspectionNodes { groups, arrays }) +} + +fn raw_v3_node_type(bytes: &[u8], key: &str) -> ZarrFdwResult { + let value = serde_json::from_slice::(bytes).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!("could not parse '{}': {error}", display_path(key))) + })?; + let object = value.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "'{}' must contain a JSON object", + display_path(key) + )) + })?; + if object.get("zarr_format").and_then(Value::as_u64) != Some(3) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "'{}' must declare zarr_format 3", + display_path(key) + ))); + } + object + .get("node_type") + .and_then(Value::as_str) + .map(str::to_string) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "'{}' must declare node_type as a string", + display_path(key) + )) + }) +} + +fn parse_ome_array_inspection(bytes: &[u8], array_path: &str) -> ZarrFdwResult { + let value = serde_json::from_slice::(bytes)?; + let object = value.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' metadata must be an object", + display_path(array_path) + )) + })?; + if object.get("zarr_format").and_then(Value::as_u64) != Some(3) + || object.get("node_type").and_then(Value::as_str) != Some("array") + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' must be a Zarr v3 array", + display_path(array_path) + ))); + } + + let shape = positive_u64_array(object.get("shape"), "shape", array_path)?; + let dtype = nonempty_string(object.get("data_type"), "data_type", array_path)?.to_string(); + let fill_value = object.get("fill_value").ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' must define fill_value", + display_path(array_path) + )) + })?; + if fill_value.is_null() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' fill_value must not be null", + display_path(array_path) + ))); + } + validate_inspection_fill_value(&dtype, fill_value, array_path)?; + + let chunk_grid = object + .get("chunk_grid") + .and_then(Value::as_object) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' chunk_grid must be an object", + display_path(array_path) + )) + })?; + nonempty_string(chunk_grid.get("name"), "chunk_grid.name", array_path)?; + let chunk_configuration = chunk_grid + .get("configuration") + .and_then(Value::as_object) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' chunk_grid.configuration must be an object", + display_path(array_path) + )) + })?; + let chunks = match chunk_configuration.get("chunk_shape") { + Some(value) => { + let chunks = positive_u64_array(Some(value), "chunk_shape", array_path)?; + if chunks.len() != shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' chunk_shape rank {} does not match shape rank {}", + display_path(array_path), + chunks.len(), + shape.len() + ))); + } + json!(chunks) + } + None => Value::Object(chunk_configuration.clone()), + }; + + let dimensions = object + .get("dimension_names") + .and_then(Value::as_array) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' must define dimension_names", + display_path(array_path) + )) + })? + .iter() + .map(|name| { + name.as_str().filter(|name| !name.is_empty()).map(str::to_string).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' dimension_names entries must be non-empty strings", + display_path(array_path) + )) + }) + }) + .collect::>>()?; + if dimensions.len() != shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' dimension_names rank {} does not match shape rank {}", + display_path(array_path), + dimensions.len(), + shape.len() + ))); + } + let mut unique_dimensions = HashSet::new(); + if dimensions + .iter() + .any(|dimension| !unique_dimensions.insert(dimension.as_str())) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' dimension_names must be unique", + display_path(array_path) + ))); + } + + let codecs = object.get("codecs").cloned().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' must define codecs", + display_path(array_path) + )) + })?; + let codec_entries = codecs + .as_array() + .filter(|entries| !entries.is_empty()) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' codecs must be a non-empty array", + display_path(array_path) + )) + })?; + for (index, codec) in codec_entries.iter().enumerate() { + let codec = codec.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' codec {index} must be an object", + display_path(array_path) + )) + })?; + if let Some(field) = codec + .keys() + .find(|field| !matches!(field.as_str(), "name" | "configuration" | "must_understand")) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' codec {index} contains unsupported field '{field}'", + display_path(array_path) + ))); + } + if codec + .get("must_understand") + .is_some_and(|value| !value.is_boolean()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' codec {index} must_understand must be a boolean", + display_path(array_path) + ))); + } + nonempty_string(codec.get("name"), "codec name", array_path)?; + if codec + .get("configuration") + .is_some_and(|configuration| !configuration.is_object()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' codec {index} configuration must be an object", + display_path(array_path) + ))); + } + } + + let attributes = match object.get("attributes") { + None => None, + Some(value) => Some(value.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' attributes must be an object", + display_path(array_path) + )) + })?), + }; + validate_optional_ome_05_attributes(&display_path(array_path), attributes)?; + + let execution_warning = match parse_v3_node(bytes) { + Ok(NodeMeta::Array(_)) => None, + Ok(NodeMeta::Group(_)) => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset path '{}' resolves to a group", + display_path(array_path) + ))); + } + Err(error) if is_execution_capability_error(&error) => Some(format!( + "level cannot be scanned by the current Zarr executor: {error}" + )), + Err(error) => return Err(error), + }; + + Ok(OmeArrayInspection { + shape, + chunks, + dimensions, + dtype, + codecs, + execution_warning, + }) +} + +fn positive_u64_array( + value: Option<&Value>, + field: &str, + array_path: &str, +) -> ZarrFdwResult> { + let values = value.and_then(Value::as_array).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' {field} must be an array", + display_path(array_path) + )) + })?; + if values.is_empty() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' {field} must not be empty", + display_path(array_path) + ))); + } + values + .iter() + .map(|value| { + value.as_u64().filter(|value| *value > 0).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' {field} entries must be positive integers", + display_path(array_path) + )) + }) + }) + .collect() +} + +fn validate_inspection_fill_value( + dtype: &str, + value: &Value, + array_path: &str, +) -> ZarrFdwResult<()> { + let valid = match dtype { + "float16" | "float32" | "float64" => { + value.is_number() || matches!(value.as_str(), Some("NaN" | "Infinity" | "-Infinity")) + } + "int8" => value + .as_i64() + .is_some_and(|value| i8::try_from(value).is_ok()), + "int16" => value + .as_i64() + .is_some_and(|value| i16::try_from(value).is_ok()), + "int32" => value + .as_i64() + .is_some_and(|value| i32::try_from(value).is_ok()), + "int64" => value.as_i64().is_some(), + "uint8" => value + .as_u64() + .is_some_and(|value| u8::try_from(value).is_ok()), + "uint16" => value + .as_u64() + .is_some_and(|value| u16::try_from(value).is_ok()), + "uint32" => value + .as_u64() + .is_some_and(|value| u32::try_from(value).is_ok()), + "uint64" => value.as_u64().is_some(), + "bool" => value.is_boolean(), + // Extension data types remain discoverable as unsupported. Their fill + // grammar belongs to the extension and cannot be validated here. + _ => true, + }; + if !valid { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' has an invalid fill_value for data_type '{dtype}'", + display_path(array_path) + ))); + } + Ok(()) +} + +fn nonempty_string<'a>( + value: Option<&'a Value>, + field: &str, + array_path: &str, +) -> ZarrFdwResult<&'a str> { + value + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr dataset array '{}' {field} must be a non-empty string", + display_path(array_path) + )) + }) +} + +fn is_execution_capability_error(error: &ZarrFdwError) -> bool { + matches!( + error, + ZarrFdwError::UnsupportedDataType(_) + | ZarrFdwError::UnsupportedCompressor(_) + | ZarrFdwError::UnsupportedExecutionFeature(_) + | ZarrFdwError::UnsupportedRank { .. } + | ZarrFdwError::UnsupportedZarrFormat { .. } + ) +} + +fn validate_resolution_order( + group_path: &str, + multiscale_index: usize, + level_index: usize, + previous: Option<&[u64]>, + current: &[u64], +) -> ZarrFdwResult<()> { + let Some(previous) = previous else { + return Ok(()); + }; + if previous.len() != current.len() + || previous + .iter() + .zip(current) + .any(|(previous, current)| current > previous) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr multiscale {multiscale_index} in group '{}' must order dataset levels from highest/largest resolution to lowest/smallest; level {level_index} shape {current:?} follows {previous:?}", + display_path(group_path) + ))); + } + Ok(()) +} + +fn ome_level_support( + axes: &[OmeAxis], + shape: &[u64], + transform: &AffineTransform, + execution_metadata_supported: bool, + warnings: &mut Vec, +) -> bool { + let mut supported = execution_metadata_supported; + if shape.len() != 2 + || axes.len() != 2 + || axes[0].name != "y" + || axes[1].name != "x" + || axes + .iter() + .any(|axis| axis.kind.as_deref() != Some("space")) + { + warnings.push("rank-2 execution requires axes [y, x] with type 'space'".to_string()); + supported = false; + } + if shape.len() != transform.scale.len() || shape.len() != transform.translation.len() { + warnings.push("affine transform rank does not match the array shape".to_string()); + return false; + } + + let mut total_coordinates = 0usize; + for (index, length) in shape.iter().copied().enumerate() { + let Ok(length) = usize::try_from(length) else { + warnings.push(format!( + "axis {index} length exceeds the executor index range" + )); + supported = false; + continue; + }; + total_coordinates = match total_coordinates.checked_add(length) { + Some(total) => total, + None => { + warnings.push("coordinate count overflowed the executor limit".to_string()); + supported = false; + continue; + } + }; + let endpoint = transform.scale[index].mul_add( + length.saturating_sub(1) as f64, + transform.translation[index], + ); + if !endpoint.is_finite() { + warnings.push(format!("axis {index} affine endpoint is not finite")); + supported = false; + } + } + if total_coordinates > MAX_OME_COORDINATE_VALUES { + warnings.push(format!( + "coordinate values total {total_coordinates} exceeds the executor limit of {MAX_OME_COORDINATE_VALUES}" + )); + supported = false; + } + supported +} + +fn checked_multiscale_derived_bytes( + current: usize, + row: &MultiscaleInspectionRow, + limit: usize, +) -> ZarrFdwResult { + // A conservative multiplier covers serde_json's tree nodes, Vec/String + // capacities, and the materialized PostgreSQL-return row representation. + let serialized = serde_json::to_vec(row).map_err(ZarrFdwError::from)?; + let charge = serialized.len().checked_mul(4).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("multiscale derived-byte charge overflowed".to_string()) + })?; + let next = current.checked_add(charge).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("multiscale derived-byte count overflowed".to_string()) + })?; + if next > limit { + return Err(ZarrFdwError::InvalidMetadata(format!( + "multiscale discovery exceeds the derived output limit of {limit} bytes" + ))); + } + Ok(next) +} + +fn inspect_server(server_name: &str) -> ZarrFdwResult> { + let server = load_foreign_server(server_name)?; + let store = ZarrStore::new(&server)?; + store.require_listing()?; + inspect_store(&store) +} + +fn load_foreign_server(server_name: &str) -> ZarrFdwResult { + let (server_oid, server_type, server_version, options) = Spi::connect(|client| { + let rows = client.select( + "SELECT s.oid::bigint AS server_oid, + s.srvtype::text AS server_type, + s.srvversion::text AS server_version, + option.option_name::text AS option_name, + option.option_value::text AS option_value + FROM pg_catalog.pg_foreign_server AS s + LEFT JOIN LATERAL pg_catalog.pg_options_to_table(s.srvoptions) AS option + ON true + WHERE s.srvname = $1 + AND pg_catalog.has_server_privilege(s.oid, 'USAGE')", + None, + &[server_name.into()], + )?; + + let mut server_oid = None; + let mut server_type = None; + let mut server_version = None; + let mut options = HashMap::new(); + for row in rows { + server_oid = row.get_by_name::("server_oid")?; + server_type = row.get_by_name::("server_type")?; + server_version = row.get_by_name::("server_version")?; + if let (Some(name), Some(value)) = ( + row.get_by_name::("option_name")?, + row.get_by_name::("option_value")?, + ) { + options.insert(name, value); + } + } + + Ok::<_, pgrx::spi::Error>((server_oid, server_type, server_version, options)) + })?; + + let server_oid = server_oid.ok_or_else(|| ZarrFdwError::ServerUnavailable { + server: server_name.to_string(), + })?; + let server_oid = u32::try_from(server_oid).map_err(|_| { + ZarrFdwError::InvalidMetadata("foreign server OID is out of range".to_string()) + })?; + + Ok(ForeignServer { + server_oid: pg_sys::Oid::from_u32(server_oid), + server_name: server_name.to_string(), + server_type, + server_version, + options, + }) +} + +fn inspect_store(store: &ZarrStore) -> ZarrFdwResult> { + let mut pending = VecDeque::from([(String::new(), 0usize, None)]); + let mut discovered = HashSet::from([String::new()]); + let mut rows = Vec::new(); + let mut metadata_bytes = 0usize; + let mut list_pages = 0usize; + + while let Some((path, depth, parent_format)) = pending.pop_front() { + let v3_key = metadata_key(&path, "zarr.json"); + let v2_array_key = metadata_key(&path, ".zarray"); + let v3 = read_optional_metadata(store, &v3_key, &mut metadata_bytes)?; + let v2_array = read_optional_metadata(store, &v2_array_key, &mut metadata_bytes)?; + + let (group_format, group_attributes) = if let Some(bytes) = v3 { + let v2_group_key = metadata_key(&path, ".zgroup"); + let v2_group = read_optional_metadata(store, &v2_group_key, &mut metadata_bytes)?; + reject_dual_metadata(&path, true, v2_array.is_some(), v2_group.is_some())?; + validate_hierarchy_format(&path, parent_format, ZarrFormat::V3)?; + + match parse_v3_node(&bytes).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!( + "could not parse '{}': {error}", + display_path(&v3_key) + )) + })? { + NodeMeta::Array(node) => { + rows.push(array_row(&path, *node)); + continue; + } + NodeMeta::Group(node) => (Some(node.format), node.attributes), + } + } else if let Some(bytes) = v2_array { + validate_hierarchy_format(&path, parent_format, ZarrFormat::V2)?; + let attributes = read_attributes(store, &path, &mut metadata_bytes)?; + let node = parse_v2_array(&bytes, attributes).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!( + "could not parse '{}': {error}", + display_path(&v2_array_key) + )) + })?; + rows.push(array_row(&path, node)); + continue; + } else { + if parent_format == Some(ZarrFormat::V3) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 node '{}' must contain explicit zarr.json metadata", + display_path(&path) + ))); + } + let v2_group_key = metadata_key(&path, ".zgroup"); + let v2_group = read_optional_metadata(store, &v2_group_key, &mut metadata_bytes)?; + let attributes = read_attributes(store, &path, &mut metadata_bytes)?; + let group = v2_group + .map(|bytes| { + parse_v2_group(&bytes, attributes.clone()).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!( + "could not parse '{}': {error}", + display_path(&v2_group_key) + )) + }) + }) + .transpose()?; + if let Some(group) = &group { + validate_hierarchy_format(&path, parent_format, group.format)?; + } + ( + group.as_ref().map(|group| group.format), + group.map(|group| group.attributes).unwrap_or(attributes), + ) + }; + + let mut child_prefixes = Vec::new(); + let mut continuation_token = None; + loop { + list_pages = list_pages.checked_add(1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("inspection list-page count overflowed".to_string()) + })?; + if list_pages > MAX_INSPECTION_LIST_PAGES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inspection exceeds the limit of {MAX_INSPECTION_LIST_PAGES} storage list pages" + ))); + } + let page = store.list_directory_page_sync(&path, continuation_token)?; + for child in page.child_prefixes { + if discovered.insert(child.clone()) { + if discovered.len() > MAX_INSPECTION_NODES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inspection exceeds the limit of {MAX_INSPECTION_NODES} Zarr nodes" + ))); + } + child_prefixes.push(child); + } + } + continuation_token = page.next_continuation_token; + if continuation_token.is_none() { + break; + } + } + + let is_group = group_format.is_some() + || !group_attributes.is_empty() + || !child_prefixes.is_empty() + || path.is_empty(); + if is_group { + let mut warnings = Vec::new(); + if group_format.is_none() { + warnings.push("group has no .zgroup metadata".to_string()); + } + rows.push(group_row( + &path, + group_format.map(zarr_format_number), + Value::Object(group_attributes), + warnings, + )); + } + + if !child_prefixes.is_empty() && depth >= MAX_INSPECTION_DEPTH { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inspection exceeds the maximum group depth of {MAX_INSPECTION_DEPTH} at '{}'", + display_path(&path) + ))); + } + child_prefixes.sort(); + for child in child_prefixes { + pending.push_back((child, depth + 1, group_format.or(parent_format))); + } + } + + rows.sort_by(|left, right| left.path.cmp(&right.path)); + resolve_crs_references(&mut rows); + Ok(rows) +} + +fn reject_dual_metadata( + path: &str, + has_v3: bool, + has_v2_array: bool, + has_v2_group: bool, +) -> ZarrFdwResult<()> { + if has_v3 && (has_v2_array || has_v2_group) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "node '{}' contains both zarr.json and Zarr v2 metadata", + display_path(path) + ))); + } + Ok(()) +} + +fn validate_hierarchy_format( + path: &str, + parent_format: Option, + node_format: ZarrFormat, +) -> ZarrFdwResult<()> { + if node_format == ZarrFormat::V3 && !path.is_empty() && parent_format.is_none() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 node '{}' requires explicit Zarr v3 metadata on every ancestor group", + display_path(path) + ))); + } + if let Some(parent_format) = parent_format + && parent_format != node_format + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "node '{}' uses Zarr v{}, but its parent group uses Zarr v{}", + display_path(path), + zarr_format_number(node_format), + zarr_format_number(parent_format) + ))); + } + Ok(()) +} + +fn zarr_format_number(format: ZarrFormat) -> i64 { + match format { + ZarrFormat::V2 => 2, + ZarrFormat::V3 => 3, + } +} + +fn read_optional_metadata( + store: &ZarrStore, + key: &str, + total_bytes: &mut usize, +) -> ZarrFdwResult>> { + let bytes = store.get_object_optional_sync(key, MAX_METADATA_OBJECT_BYTES)?; + if let Some(bytes) = &bytes { + *total_bytes = total_bytes.checked_add(bytes.len()).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("inspection metadata byte count overflowed".to_string()) + })?; + if *total_bytes > MAX_INSPECTION_METADATA_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inspection exceeds the metadata read limit of {MAX_INSPECTION_METADATA_BYTES} bytes" + ))); + } + } + Ok(bytes) +} + +fn read_attributes( + store: &ZarrStore, + path: &str, + total_bytes: &mut usize, +) -> ZarrFdwResult> { + let key = metadata_key(path, ".zattrs"); + let Some(bytes) = read_optional_metadata(store, &key, total_bytes)? else { + return Ok(Map::new()); + }; + let value = serde_json::from_slice::(&bytes).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!("could not parse '{}': {error}", display_path(&key))) + })?; + value.as_object().cloned().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "'{}' must contain a JSON object", + display_path(&key) + )) + }) +} + +fn array_row(path: &str, node: ArrayNode) -> InspectionRow { + let ArrayNode { + format, + meta, + attributes: attrs, + dimension_names, + native_dtype, + native_codecs, + } = node; + let mut warnings = Vec::new(); + let dimensions = inspected_dimensions( + format, + dimension_names.as_deref(), + &attrs, + meta.shape.len(), + &mut warnings, + ); + let units = string_attribute(&attrs, "units", &mut warnings); + let calendar = string_attribute(&attrs, "calendar", &mut warnings); + let scale_factor = numeric_attribute(&attrs, "scale_factor", &mut warnings); + let add_offset = numeric_attribute(&attrs, "add_offset", &mut warnings); + let crs = crs_attribute(&attrs); + + InspectionRow { + path: display_path(path), + kind: "array".to_string(), + group_path: parent_path(path), + variable: Some(node_name(path)), + zarr_format: Some(zarr_format_number(format)), + shape: Some(json!(meta.shape)), + dimensions, + dtype: Some(native_dtype), + chunks: Some(json!(meta.native_chunk_shape())), + codecs: Some(native_codecs), + units, + fill_value: Some(meta.fill_value), + scale_factor, + add_offset, + crs, + calendar, + attributes: Value::Object(attrs), + warnings, + } +} + +fn group_row( + path: &str, + zarr_format: Option, + attributes: Value, + warnings: Vec, +) -> InspectionRow { + let crs = attributes.as_object().and_then(crs_attribute); + InspectionRow { + path: display_path(path), + kind: "group".to_string(), + group_path: parent_path(path), + variable: None, + zarr_format, + shape: None, + dimensions: None, + dtype: None, + chunks: None, + codecs: None, + units: None, + fill_value: None, + scale_factor: None, + add_offset: None, + crs, + calendar: None, + attributes, + warnings, + } +} + +fn crs_attribute(attrs: &Map) -> Option { + direct_crs_attribute(attrs).or_else(|| attrs.get("grid_mapping").cloned()) +} + +fn direct_crs_attribute(attrs: &Map) -> Option { + ["crs", "spatial_ref", "crs_wkt"] + .iter() + .find_map(|key| attrs.get(*key).cloned()) +} + +fn resolve_crs_references(rows: &mut [InspectionRow]) { + resolve_crs_references_with_limit(rows, MAX_INSPECTION_DERIVED_CRS_BYTES); +} + +fn resolve_crs_references_with_limit(rows: &mut [InspectionRow], max_derived_bytes: usize) { + let mut kinds_by_path = HashMap::new(); + let mut direct_crs_by_array_path = HashMap::new(); + for row in rows.iter() { + kinds_by_path.insert(row.path.clone(), row.kind.clone()); + if row.kind == "array" + && let Some(crs) = row.attributes.as_object().and_then(direct_crs_attribute) + { + direct_crs_by_array_path.insert(row.path.clone(), crs); + } + } + + let mut derived_crs_bytes = 0usize; + for row in rows.iter_mut() { + if row.kind != "array" { + continue; + } + let Some(grid_mapping) = row.attributes.as_object().and_then(|attrs| { + if direct_crs_attribute(attrs).is_some() { + return None; + } + attrs.get("grid_mapping").cloned() + }) else { + continue; + }; + let Some(reference) = grid_mapping_reference(&grid_mapping, &mut row.warnings) else { + continue; + }; + let sibling_path = same_group_sibling_path(row.group_path.as_deref(), reference); + match kinds_by_path.get(&sibling_path).map(String::as_str) { + Some("array") => { + let Some(crs) = direct_crs_by_array_path.get(&sibling_path) else { + row.warnings.push(format!( + "grid_mapping reference '{reference}' resolves to a sibling array without direct CRS metadata" + )); + continue; + }; + let Ok(bytes) = serde_json::to_vec(crs) else { + row.warnings.push(format!( + "grid_mapping reference '{reference}' CRS metadata could not be measured" + )); + continue; + }; + let Some(next_bytes) = derived_crs_bytes.checked_add(bytes.len()) else { + row.warnings.push(format!( + "grid_mapping reference '{reference}' exceeds the derived CRS byte limit of {max_derived_bytes} bytes" + )); + continue; + }; + if next_bytes > max_derived_bytes { + row.warnings.push(format!( + "grid_mapping reference '{reference}' exceeds the derived CRS byte limit of {max_derived_bytes} bytes" + )); + continue; + } + derived_crs_bytes = next_bytes; + row.crs = Some(crs.clone()); + } + _ => row.warnings.push(format!( + "grid_mapping reference '{reference}' does not resolve to a sibling array" + )), + } + } +} + +fn grid_mapping_reference<'a>(value: &'a Value, warnings: &mut Vec) -> Option<&'a str> { + let Some(reference) = value.as_str() else { + warnings.push("grid_mapping reference must be a string".to_string()); + return None; + }; + if reference.trim().is_empty() { + warnings.push("grid_mapping reference must not be empty".to_string()); + return None; + } + if reference.trim() != reference + || reference.chars().any(char::is_whitespace) + || reference.contains('/') + || reference.contains('\\') + || reference == "." + || reference == ".." + { + warnings.push(format!( + "grid_mapping reference '{reference}' must be a same-group array name" + )); + return None; + } + Some(reference) +} + +fn same_group_sibling_path(group_path: Option<&str>, reference: &str) -> String { + match group_path { + Some("/") | None => reference.to_string(), + Some(group_path) => format!("{group_path}/{reference}"), + } +} + +fn named_dimensions( + attrs: &Map, + rank: usize, + warnings: &mut Vec, +) -> Option> { + match parse_named_dimensions(attrs, rank) { + Ok(dimensions) => dimensions, + Err(message) => { + warnings.push(message); + None + } + } +} + +fn inspected_dimensions( + format: ZarrFormat, + native: Option<&[Option]>, + attrs: &Map, + rank: usize, + warnings: &mut Vec, +) -> Option> { + if format == ZarrFormat::V2 { + return named_dimensions(attrs, rank, warnings); + } + + let native = native?; + if native.len() != rank { + warnings.push(format!( + "dimension_names has {} names but the array rank is {rank}", + native.len() + )); + return None; + } + let Some(names) = native.iter().cloned().collect::>>() else { + warnings.push("dimension_names contains an unnamed dimension".to_string()); + return None; + }; + let probe = Map::from_iter([("_ARRAY_DIMENSIONS".to_string(), json!(names))]); + match parse_named_dimensions(&probe, rank) { + Ok(dimensions) => dimensions, + Err(message) => { + warnings.push( + message + .replace("_ARRAY_DIMENSIONS names", "dimension_names") + .replace("_ARRAY_DIMENSIONS", "dimension_names"), + ); + None + } + } +} + +fn string_attribute( + attrs: &Map, + name: &str, + warnings: &mut Vec, +) -> Option { + let value = attrs.get(name)?; + match value.as_str() { + Some(value) => Some(value.to_string()), + None => { + warnings.push(format!("attribute '{name}' must be a string")); + None + } + } +} + +fn numeric_attribute( + attrs: &Map, + name: &str, + warnings: &mut Vec, +) -> Option { + let value = attrs.get(name)?; + match value.as_f64() { + Some(value) if value.is_finite() => Some(value), + _ => { + warnings.push(format!("attribute '{name}' must be a finite number")); + None + } + } +} + +fn metadata_key(path: &str, name: &str) -> String { + join_key(path, name) +} + +fn display_path(path: &str) -> String { + if path.is_empty() { + "/".to_string() + } else { + path.to_string() + } +} + +fn parent_path(path: &str) -> Option { + if path.is_empty() { + return None; + } + Some( + path.rsplit_once('/') + .map(|(parent, _)| display_path(parent)) + .unwrap_or_else(|| "/".to_string()), + ) +} + +fn node_name(path: &str) -> String { + if path.is_empty() { + "/".to_string() + } else { + path.rsplit('/').next().unwrap_or(path).to_string() + } +} + +#[cfg(test)] +mod tests { + use super::super::meta::{ArrayMeta, ChunkKeyEncoding}; + use super::*; + + fn array_node(attributes: Map) -> ArrayNode { + ArrayNode { + format: ZarrFormat::V2, + meta: ArrayMeta { + zarr_format: 2, + shape: vec![2, 5, 6], + chunks: vec![1, 3, 4], + dtype: ") -> ArrayNode { + let mut node = array_node(attributes); + node.format = ZarrFormat::V3; + node.meta.zarr_format = 3; + node.meta.chunk_key_encoding = ChunkKeyEncoding::Default { separator: '/' }; + node.dimension_names = Some(vec![ + Some("time".to_string()), + Some("y".to_string()), + Some("x".to_string()), + ]); + node.native_dtype = "float32".to_string(); + node.native_codecs = json!([{ + "name": "bytes", + "configuration": {"endian": "little"} + }]); + node + } + + #[test] + fn array_row_exposes_scientific_metadata_without_decoding_it() { + let attrs = serde_json::from_value::>(json!({ + "_ARRAY_DIMENSIONS": ["time", "lat", "lon"], + "units": "K", + "calendar": "proleptic_gregorian", + "scale_factor": 0.01, + "add_offset": 273.15, + "grid_mapping": "spatial_ref", + "long_name": "air temperature" + })) + .unwrap(); + + let row = array_row("climate/temperature", array_node(attrs)); + + assert_eq!(row.path, "climate/temperature"); + assert_eq!(row.group_path.as_deref(), Some("climate")); + assert_eq!(row.variable.as_deref(), Some("temperature")); + assert_eq!(row.dimensions.unwrap(), vec!["time", "lat", "lon"]); + assert_eq!(row.units.as_deref(), Some("K")); + assert_eq!(row.calendar.as_deref(), Some("proleptic_gregorian")); + assert_eq!(row.scale_factor, Some(0.01)); + assert_eq!(row.add_offset, Some(273.15)); + assert_eq!(row.crs, Some(json!("spatial_ref"))); + assert_eq!(row.fill_value, Some(json!(-7.5))); + assert_eq!(row.codecs.unwrap()["compressor"]["id"], "blosc"); + assert!(row.warnings.is_empty()); + } + + #[test] + fn malformed_named_dimensions_are_preserved_as_attributes_and_warned() { + let attrs = serde_json::from_value::>(json!({ + "_ARRAY_DIMENSIONS": ["lat", "lat", "lon"] + })) + .unwrap(); + + let row = array_row("temperature", array_node(attrs)); + + assert!(row.dimensions.is_none()); + assert_eq!( + row.attributes["_ARRAY_DIMENSIONS"], + json!(["lat", "lat", "lon"]) + ); + assert_eq!(row.warnings, vec!["_ARRAY_DIMENSIONS names must be unique"]); + } + + #[test] + fn v3_array_row_exposes_native_metadata_and_embedded_attributes() { + let attrs = serde_json::from_value::>(json!({ + "units": "K", + "scale_factor": 0.01, + "add_offset": 273.15 + })) + .unwrap(); + + let row = array_row("nested/raw", v3_array_node(attrs)); + + assert_eq!(row.zarr_format, Some(3)); + assert_eq!(row.dimensions.unwrap(), vec!["time", "y", "x"]); + assert_eq!(row.dtype.as_deref(), Some("float32")); + assert_eq!( + row.codecs.unwrap(), + json!([{ + "name": "bytes", + "configuration": {"endian": "little"} + }]) + ); + assert_eq!(row.attributes["units"], json!("K")); + assert!(row.warnings.is_empty()); + } + + #[test] + fn rejects_dual_or_implicitly_nested_v3_metadata() { + assert!(matches!( + reject_dual_metadata("nested/raw", true, true, false), + Err(ZarrFdwError::InvalidMetadata(message)) + if message == "node 'nested/raw' contains both zarr.json and Zarr v2 metadata" + )); + assert!(matches!( + validate_hierarchy_format("nested/raw", None, ZarrFormat::V3), + Err(ZarrFdwError::InvalidMetadata(message)) + if message == "Zarr v3 node 'nested/raw' requires explicit Zarr v3 metadata on every ancestor group" + )); + validate_hierarchy_format("", None, ZarrFormat::V3).unwrap(); + validate_hierarchy_format("nested/raw", Some(ZarrFormat::V3), ZarrFormat::V3).unwrap(); + } + + #[test] + fn crs_resolution_preserves_direct_precedence() { + let value_attrs = serde_json::from_value::>(json!({ + "crs": "EPSG:4326", + "grid_mapping": "spatial_ref" + })) + .unwrap(); + let ref_attrs = serde_json::from_value::>(json!({ + "spatial_ref": "EPSG:3857" + })) + .unwrap(); + let mut rows = vec![ + array_row("nested/raw", array_node(value_attrs)), + array_row("nested/spatial_ref", array_node(ref_attrs)), + ]; + + resolve_crs_references(&mut rows); + + let raw = rows.iter().find(|row| row.path == "nested/raw").unwrap(); + assert_eq!(raw.crs, Some(json!("EPSG:4326"))); + assert!(raw.warnings.is_empty()); + } + + #[test] + fn crs_resolution_uses_same_group_sibling_array_direct_crs() { + let value_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "spatial_ref" + })) + .unwrap(); + let ref_attrs = serde_json::from_value::>(json!({ + "spatial_ref": {"type": "ProjectedCRS", "name": "EPSG:3857"} + })) + .unwrap(); + let mut rows = vec![ + array_row("nested/raw", array_node(value_attrs)), + array_row("nested/spatial_ref", array_node(ref_attrs)), + ]; + + resolve_crs_references(&mut rows); + + let raw = rows.iter().find(|row| row.path == "nested/raw").unwrap(); + assert_eq!( + raw.crs, + Some(json!({"type": "ProjectedCRS", "name": "EPSG:3857"})) + ); + assert_eq!(raw.attributes["grid_mapping"], json!("spatial_ref")); + assert!(raw.warnings.is_empty()); + } + + #[test] + fn crs_resolution_warns_for_missing_non_array_and_crs_less_references() { + let missing_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "missing" + })) + .unwrap(); + let non_array_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "group_ref" + })) + .unwrap(); + let crs_less_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "no_crs" + })) + .unwrap(); + let group_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "no_crs" + })) + .unwrap(); + let mut rows = vec![ + array_row("nested/missing_raw", array_node(missing_attrs)), + array_row("nested/non_array_raw", array_node(non_array_attrs)), + array_row("nested/crs_less_raw", array_node(crs_less_attrs)), + group_row( + "nested/group_ref", + Some(2), + Value::Object(Map::new()), + vec![], + ), + group_row( + "nested/source_group", + Some(2), + Value::Object(group_attrs), + vec![], + ), + array_row("nested/no_crs", array_node(Map::new())), + ]; + + resolve_crs_references(&mut rows); + + let missing = rows + .iter() + .find(|row| row.path == "nested/missing_raw") + .unwrap(); + assert_eq!(missing.crs, Some(json!("missing"))); + assert_eq!( + missing.warnings, + vec!["grid_mapping reference 'missing' does not resolve to a sibling array"] + ); + let non_array = rows + .iter() + .find(|row| row.path == "nested/non_array_raw") + .unwrap(); + assert_eq!(non_array.crs, Some(json!("group_ref"))); + assert_eq!( + non_array.warnings, + vec!["grid_mapping reference 'group_ref' does not resolve to a sibling array"] + ); + let crs_less = rows + .iter() + .find(|row| row.path == "nested/crs_less_raw") + .unwrap(); + assert_eq!(crs_less.crs, Some(json!("no_crs"))); + assert_eq!( + crs_less.warnings, + vec![ + "grid_mapping reference 'no_crs' resolves to a sibling array without direct CRS metadata" + ] + ); + let source_group = rows + .iter() + .find(|row| row.path == "nested/source_group") + .unwrap(); + assert_eq!(source_group.crs, Some(json!("no_crs"))); + assert!(source_group.warnings.is_empty()); + } + + #[test] + fn crs_resolution_warns_for_invalid_references() { + let non_string_attrs = serde_json::from_value::>(json!({ + "grid_mapping": 7 + })) + .unwrap(); + let empty_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "" + })) + .unwrap(); + let path_like_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "../spatial_ref" + })) + .unwrap(); + let multi_token_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "spatial ref" + })) + .unwrap(); + let mut rows = vec![ + array_row("nested/non_string", array_node(non_string_attrs)), + array_row("nested/empty", array_node(empty_attrs)), + array_row("nested/path_like", array_node(path_like_attrs)), + array_row("nested/multi_token", array_node(multi_token_attrs)), + ]; + + resolve_crs_references(&mut rows); + + let non_string = rows + .iter() + .find(|row| row.path == "nested/non_string") + .unwrap(); + assert_eq!(non_string.crs, Some(json!(7))); + assert_eq!( + non_string.warnings, + vec!["grid_mapping reference must be a string"] + ); + let empty = rows.iter().find(|row| row.path == "nested/empty").unwrap(); + assert_eq!(empty.crs, Some(json!(""))); + assert_eq!( + empty.warnings, + vec!["grid_mapping reference must not be empty"] + ); + let path_like = rows + .iter() + .find(|row| row.path == "nested/path_like") + .unwrap(); + assert_eq!(path_like.crs, Some(json!("../spatial_ref"))); + assert_eq!( + path_like.warnings, + vec!["grid_mapping reference '../spatial_ref' must be a same-group array name"] + ); + let multi_token = rows + .iter() + .find(|row| row.path == "nested/multi_token") + .unwrap(); + assert_eq!(multi_token.crs, Some(json!("spatial ref"))); + assert_eq!( + multi_token.warnings, + vec!["grid_mapping reference 'spatial ref' must be a same-group array name"] + ); + } + + #[test] + fn crs_resolution_enforces_derived_output_byte_cap() { + let value_attrs = serde_json::from_value::>(json!({ + "grid_mapping": "spatial_ref" + })) + .unwrap(); + let ref_attrs = serde_json::from_value::>(json!({ + "spatial_ref": "EPSG:3857" + })) + .unwrap(); + let mut rows = vec![ + array_row("nested/raw", array_node(value_attrs)), + array_row("nested/spatial_ref", array_node(ref_attrs)), + ]; + + resolve_crs_references_with_limit(&mut rows, 4); + + let raw = rows.iter().find(|row| row.path == "nested/raw").unwrap(); + assert_eq!(raw.crs, Some(json!("spatial_ref"))); + assert_eq!( + raw.warnings, + vec![ + "grid_mapping reference 'spatial_ref' exceeds the derived CRS byte limit of 4 bytes" + ] + ); + } + + #[test] + fn valid_unsupported_v3_array_remains_inspectable() { + let metadata_value = json!({ + "zarr_format": 3, + "node_type": "array", + "shape": [2, 2], + "data_type": "uint16", + "chunk_grid": { + "name": "regular", + "configuration": {"chunk_shape": [2, 2]} + }, + "chunk_key_encoding": { + "name": "default", + "configuration": {"separator": "/"} + }, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "dimension_names": ["y", "x"] + }); + let metadata = serde_json::to_vec(&metadata_value).unwrap(); + + let array = parse_ome_array_inspection(&metadata, "image/0").unwrap(); + assert_eq!(array.shape, vec![2, 2]); + assert_eq!(array.dtype, "uint16"); + assert!( + array + .execution_warning + .as_deref() + .is_some_and(|warning| warning.contains("data type 'uint16' is not supported")) + ); + + let mut adversarial = metadata_value.clone(); + adversarial["not supported"] = json!(true); + assert!(matches!( + parse_ome_array_inspection( + &serde_json::to_vec(&adversarial).unwrap(), + "image/adversarial" + ), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("unrecognized field 'not supported'") + )); + + let mut malformed_codec = metadata_value; + malformed_codec["codecs"] = json!([{"name": "bytes", "configuration": {"endian": "little"}}, { + "name": "zstd", + "must_understand": "false" + }]); + assert!(matches!( + parse_ome_array_inspection( + &serde_json::to_vec(&malformed_codec).unwrap(), + "image/malformed-codec" + ), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("must_understand must be a boolean") + )); + } + + #[test] + fn ome_version_only_group_is_not_treated_as_a_multiscale() { + let attributes = serde_json::from_value::>(json!({ + "ome": {"version": "0.5", "series": []} + })) + .unwrap(); + + validate_optional_ome_05_attributes("/", Some(&attributes)).unwrap(); + assert!(!has_ome_multiscales(&attributes)); + } + + #[test] + fn multiscale_resolution_order_never_increases_shape() { + validate_resolution_order("image", 0, 1, Some(&[4, 4]), &[2, 2]).unwrap(); + assert!(matches!( + validate_resolution_order("image", 0, 1, Some(&[2, 2]), &[4, 2]), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("highest/largest resolution to lowest/smallest") + )); + } + + #[test] + fn multiscale_support_checks_coordinate_budget_and_affine_endpoints() { + let axes = vec![ + OmeAxis { + name: "y".to_string(), + kind: Some("space".to_string()), + unit: None, + }, + OmeAxis { + name: "x".to_string(), + kind: Some("space".to_string()), + unit: None, + }, + ]; + let normal = AffineTransform { + scale: vec![1.0, 1.0], + translation: vec![0.0, 0.0], + }; + let mut warnings = Vec::new(); + assert!(ome_level_support( + &axes, + &[2, 2], + &normal, + true, + &mut warnings + )); + assert!(warnings.is_empty()); + + let mut warnings = Vec::new(); + assert!(!ome_level_support( + &axes, + &[MAX_OME_COORDINATE_VALUES as u64, 1], + &normal, + true, + &mut warnings + )); + assert!(warnings.iter().any(|warning| warning.contains("exceeds"))); + + let overflowing = AffineTransform { + scale: vec![f64::MAX, 1.0], + translation: vec![f64::MAX, 0.0], + }; + let mut warnings = Vec::new(); + assert!(!ome_level_support( + &axes, + &[2, 2], + &overflowing, + true, + &mut warnings + )); + assert!( + warnings + .iter() + .any(|warning| warning.contains("affine endpoint is not finite")) + ); + } + + #[test] + fn multiscale_discovery_enforces_derived_output_budget() { + let row = MultiscaleInspectionRow { + group_path: "image".to_string(), + multiscale_index: 0, + multiscale_name: Some("pyramid".to_string()), + level_index: 0, + array_path: "image/0".to_string(), + axes: json!([{"name": "y"}, {"name": "x"}]), + shape: json!([4, 4]), + chunks: json!([2, 2]), + dtype: "float32".to_string(), + codecs: json!([{"name": "bytes"}]), + scale: vec![1.0, 1.0], + translation: vec![0.0, 0.0], + supported: true, + warnings: vec![], + }; + + assert!(matches!( + checked_multiscale_derived_bytes(0, &row, 1), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("derived output limit") + )); + assert!(checked_multiscale_derived_bytes(0, &row, 64 * 1024).is_ok()); + } + + #[test] + fn paths_are_stable_for_root_and_nested_nodes() { + assert_eq!(display_path(""), "/"); + assert_eq!(parent_path(""), None); + assert_eq!(parent_path("temperature"), Some("/".to_string())); + assert_eq!( + parent_path("nested/temperature"), + Some("nested".to_string()) + ); + assert_eq!(node_name(""), "/"); + assert_eq!(node_name("nested/temperature"), "temperature"); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/meta.rs b/wrappers/src/fdw/zarr_fdw/meta.rs new file mode 100644 index 000000000..97e086084 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/meta.rs @@ -0,0 +1,980 @@ +//! Version-neutral Zarr array metadata used by scan execution. +//! +//! Format-specific metadata is normalized here. The executor therefore sees +//! the same shape, chunk, dtype, fill, and chunk-key contract for Zarr v2 and +//! for the bounded direct Zarr v3 subset supported by this module. + +use super::codec::CodecPipeline; +use super::decode::{DType, fill_value_bytes}; +use super::sharding::{ShardingConfig, StorageLayout}; +use super::{ZarrFdwError, ZarrFdwResult}; +use serde::Deserialize; +use serde_json::{Map, Value}; + +const MAX_SCAN_RANK: usize = 64; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ZarrFormat { + V2, + V3, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChunkKeyEncoding { + /// Zarr v3 default encoding: `c/` or `c.`. + Default { separator: char }, + /// Zarr v2 encoding, including when selected by a v3 array. + V2 { separator: char }, +} + +#[derive(Debug, Clone)] +pub struct ArrayMeta { + pub zarr_format: u32, + pub shape: Vec, + pub chunks: Vec, + /// Executor-normalized NumPy dtype, for example `, + /// Validated, format-neutral execution pipeline. + pub codec_pipeline: CodecPipeline, + /// Physical storage mapping. `chunks` remains the executor's logical + /// chunk shape; sharded layouts retain their native outer shape here. + pub storage_layout: StorageLayout, + pub chunk_key_encoding: ChunkKeyEncoding, + pub order: char, + pub filters: Option>, +} + +#[derive(Debug, Clone)] +pub struct ArrayNode { + pub format: ZarrFormat, + pub meta: ArrayMeta, + pub attributes: Map, + /// Native v3 names. `None` distinguishes an absent field from a present + /// list containing null entries. V2 names remain in `_ARRAY_DIMENSIONS`. + pub dimension_names: Option>>, + /// Native spelling retained for truthful inspection. + pub native_dtype: String, + /// Native ordered codec metadata retained for truthful inspection. + pub native_codecs: Value, +} + +#[derive(Debug, Clone)] +pub struct GroupNode { + pub format: ZarrFormat, + pub attributes: Map, +} + +#[derive(Debug, Clone)] +pub enum NodeMeta { + Array(Box), + Group(GroupNode), +} + +#[derive(Debug, Deserialize)] +struct V2ArrayMeta { + zarr_format: u32, + shape: Vec, + chunks: Vec, + dtype: String, + fill_value: Value, + compressor: Option, + #[serde(default = "default_v2_separator")] + dimension_separator: String, + #[serde(default = "default_order")] + order: char, + filters: Option>, +} + +fn default_v2_separator() -> String { + ".".to_string() +} + +fn default_order() -> char { + 'C' +} + +impl ArrayMeta { + pub(crate) fn validate(&self) -> ZarrFdwResult<()> { + if !matches!(self.zarr_format, 2 | 3) { + return Err(ZarrFdwError::UnsupportedZarrFormat { + version: self.zarr_format, + }); + } + if !(1..=MAX_SCAN_RANK).contains(&self.shape.len()) { + return Err(ZarrFdwError::UnsupportedRank { + rank: self.shape.len(), + }); + } + self.validate_common() + } + + fn validate_common(&self) -> ZarrFdwResult<()> { + if self.shape.len() != self.chunks.len() { + return Err(ZarrFdwError::InvalidMetadata( + "shape and chunks lengths differ".to_string(), + )); + } + if let Some(axis) = self.shape.iter().position(|&extent| extent == 0) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "shape dimension {axis} must be greater than zero" + ))); + } + if let Some(axis) = self.chunks.iter().position(|&extent| extent == 0) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "chunk dimension {axis} must be greater than zero" + ))); + } + for axis in 0..self.shape.len() { + self.shape_extent(axis)?; + self.chunk_extent(axis)?; + } + self.chunk_cell_count()?; + if self + .filters + .as_ref() + .is_some_and(|filters| !filters.is_empty()) + { + return Err(ZarrFdwError::InvalidMetadata( + "zarr filters are not supported yet".to_string(), + )); + } + if self.order != 'C' { + return Err(ZarrFdwError::InvalidMetadata(format!( + "only row-major (C order) arrays are supported, got '{}'", + self.order + ))); + } + Ok(()) + } + + pub fn validate_coordinate(&self) -> ZarrFdwResult<()> { + if self.shape.len() != 1 { + return Err(ZarrFdwError::CoordinateReadError { + axis: String::new(), + error: format!("coordinate array must be 1D, got rank {}", self.shape.len()), + }); + } + self.validate_common() + } + + pub fn chunks_per_axis(&self) -> Vec { + self.shape + .iter() + .zip(self.chunks.iter()) + .map(|(shape, chunk)| shape.div_ceil(*chunk)) + .collect() + } + + /// Native regular chunk-grid shape reported by metadata inspection. + /// This is the shard shape for `sharding_indexed`, and the logical chunk + /// shape for direct v2/v3 arrays. + pub(crate) fn native_chunk_shape(&self) -> &[u64] { + match &self.storage_layout { + StorageLayout::Direct => &self.chunks, + StorageLayout::Sharded(config) => &config.shard_shape, + } + } + + pub fn shape_extent(&self, axis: usize) -> ZarrFdwResult { + usize::try_from(self.shape[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "shape dimension {axis} exceeds this platform's index capacity" + )) + }) + } + + pub fn chunk_extent(&self, axis: usize) -> ZarrFdwResult { + usize::try_from(self.chunks[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk dimension {axis} exceeds this platform's index capacity" + )) + }) + } + + pub fn chunk_cell_count(&self) -> ZarrFdwResult { + self.chunks + .iter() + .enumerate() + .try_fold(1usize, |cells, (axis, &extent)| { + let extent = usize::try_from(extent).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk dimension {axis} exceeds this platform's index capacity" + )) + })?; + cells.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "declared chunk cell count exceeds this platform's index capacity" + .to_string(), + ) + }) + }) + } +} + +pub fn parse_v2_array(bytes: &[u8], attributes: Map) -> ZarrFdwResult { + let raw = serde_json::from_slice::(bytes)?; + let native_dtype = raw.dtype.clone(); + let native_codecs = serde_json::json!({ + "filters": raw.filters.clone(), + "compressor": raw.compressor.clone(), + }); + Ok(ArrayNode { + format: ZarrFormat::V2, + meta: v2_meta(raw)?, + attributes, + dimension_names: None, + native_dtype, + native_codecs, + }) +} + +pub fn parse_v2_group(bytes: &[u8], attributes: Map) -> ZarrFdwResult { + let value = serde_json::from_slice::(bytes)?; + let version = required_u64(value.as_object(), "zarr_format")?; + if version != 2 { + return Err(ZarrFdwError::UnsupportedZarrFormat { + version: u32::try_from(version).unwrap_or(u32::MAX), + }); + } + Ok(GroupNode { + format: ZarrFormat::V2, + attributes, + }) +} + +pub fn parse_v3_node(bytes: &[u8]) -> ZarrFdwResult { + let value = serde_json::from_slice::(bytes)?; + let object = value.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("zarr.json must contain a JSON object".to_string()) + })?; + let version = required_u64(Some(object), "zarr_format")?; + if version != 3 { + return Err(ZarrFdwError::UnsupportedZarrFormat { + version: u32::try_from(version).unwrap_or(u32::MAX), + }); + } + let node_type = required_string(object, "node_type")?; + if object.contains_key("consolidated_metadata") { + return Err(ZarrFdwError::InvalidMetadata( + "Zarr v3 consolidated metadata is not supported yet".to_string(), + )); + } + match node_type { + "group" => { + validate_additional_fields( + object, + &["zarr_format", "node_type", "attributes"], + "Zarr v3 group", + )?; + Ok(NodeMeta::Group(GroupNode { + format: ZarrFormat::V3, + attributes: optional_object(object, "attributes")?, + })) + } + "array" => { + validate_additional_fields( + object, + &[ + "zarr_format", + "node_type", + "shape", + "data_type", + "chunk_grid", + "chunk_key_encoding", + "fill_value", + "codecs", + "attributes", + "storage_transformers", + "dimension_names", + ], + "Zarr v3 array", + )?; + parse_v3_array(object, optional_object(object, "attributes")?) + .map(|node| NodeMeta::Array(Box::new(node))) + } + other => Err(ZarrFdwError::InvalidMetadata(format!( + "zarr.json node_type must be 'array' or 'group', got '{other}'" + ))), + } +} + +fn v2_meta(raw: V2ArrayMeta) -> ZarrFdwResult { + if raw.zarr_format != 2 { + return Err(ZarrFdwError::UnsupportedZarrFormat { + version: raw.zarr_format, + }); + } + let separator = parse_separator(&raw.dimension_separator, "dimension_separator")?; + Ok(ArrayMeta { + zarr_format: 2, + shape: raw.shape, + chunks: raw.chunks, + dtype: raw.dtype, + fill_value: raw.fill_value, + compressor: raw.compressor, + // V2 inspection remains permissive: execution resolves and validates + // the legacy compressor through `CodecPipeline::from_v2` at scan + // startup, while v3 stores its already-validated ordered pipeline. + codec_pipeline: CodecPipeline::raw_v2(), + storage_layout: StorageLayout::Direct, + chunk_key_encoding: ChunkKeyEncoding::V2 { separator }, + order: raw.order, + filters: raw.filters, + }) +} + +fn parse_v3_array( + object: &Map, + attributes: Map, +) -> ZarrFdwResult { + let shape = required_u64_array(object, "shape")?; + let native_dtype = required_string(object, "data_type")?.to_string(); + let fill_value = object.get("fill_value").cloned().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("zarr v3 array must define fill_value".to_string()) + })?; + if fill_value.is_null() { + return Err(ZarrFdwError::InvalidMetadata( + "zarr v3 numeric fill_value must not be null".to_string(), + )); + } + + let chunk_grid = required_object(object, "chunk_grid")?; + validate_extension_object(chunk_grid, "chunk_grid")?; + reject_ignorable_extension(chunk_grid, "chunk_grid")?; + if required_string(chunk_grid, "name")? != "regular" { + return Err(ZarrFdwError::UnsupportedExecutionFeature( + "only the Zarr v3 regular chunk grid is supported".to_string(), + )); + } + let chunk_grid_configuration = required_object(chunk_grid, "configuration")?; + validate_exact_fields( + chunk_grid_configuration, + &["chunk_shape"], + "regular chunk grid configuration", + )?; + let chunks = required_u64_array(chunk_grid_configuration, "chunk_shape")?; + + let key = required_object(object, "chunk_key_encoding")?; + validate_extension_object(key, "chunk_key_encoding")?; + reject_ignorable_extension(key, "chunk_key_encoding")?; + let key_name = required_string(key, "name")?; + let key_configuration = optional_object_ref(key, "configuration")?; + if let Some(configuration) = key_configuration { + validate_exact_fields( + configuration, + &["separator"], + "chunk-key encoding configuration", + )?; + } + let default_separator = if key_name == "default" { "/" } else { "." }; + let separator = key_configuration + .and_then(|configuration| configuration.get("separator")) + .map(|value| { + value.as_str().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("chunk-key separator must be a string".to_string()) + }) + }) + .transpose()? + .unwrap_or(default_separator); + let separator = parse_separator(separator, "chunk-key separator")?; + let chunk_key_encoding = match key_name { + "default" => ChunkKeyEncoding::Default { separator }, + "v2" => ChunkKeyEncoding::V2 { separator }, + other => { + return Err(ZarrFdwError::UnsupportedExecutionFeature(format!( + "Zarr v3 chunk-key encoding '{other}' is not supported" + ))); + } + }; + + if let Some(value) = object.get("storage_transformers") { + let transformers = value.as_array().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "Zarr v3 storage_transformers must be an array".to_string(), + ) + })?; + for (index, transformer) in transformers.iter().enumerate() { + let transformer = transformer.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 storage transformer {index} must be an object" + )) + })?; + validate_extension_object(transformer, "storage transformer")?; + required_string(transformer, "name")?; + } + if !transformers.is_empty() { + return Err(ZarrFdwError::UnsupportedExecutionFeature( + "Zarr v3 storage transformers are not supported yet".to_string(), + )); + } + } + + let native_codecs = object.get("codecs").cloned().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("Zarr v3 array must define codecs".to_string()) + })?; + let sharded = native_codecs.as_array().is_some_and(|codecs| { + codecs.len() == 1 + && codecs[0] + .as_object() + .and_then(|codec| codec.get("name")) + .and_then(Value::as_str) + == Some("sharding_indexed") + }); + let (codec_pipeline, storage_layout, dtype, chunks) = if sharded { + let (config, dtype) = ShardingConfig::from_v3(&native_dtype, &chunks, &native_codecs)?; + let codec_pipeline = config.inner_codecs.clone(); + let chunks = config.inner_chunk_shape.clone(); + ( + codec_pipeline, + StorageLayout::Sharded(config), + dtype, + chunks, + ) + } else { + let (codec_pipeline, dtype) = + CodecPipeline::from_v3(&native_dtype, shape.len(), &native_codecs)?; + (codec_pipeline, StorageLayout::Direct, dtype, chunks) + }; + let parsed_dtype = DType::parse(&dtype)?; + fill_value_bytes(parsed_dtype, &fill_value).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!( + "invalid Zarr v3 fill_value for {native_dtype}: {error}" + )) + })?; + + let dimension_names = object + .get("dimension_names") + .map(|value| { + value + .as_array() + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "Zarr v3 dimension_names must be an array".to_string(), + ) + })? + .iter() + .map(|name| { + if name.is_null() { + Ok(None) + } else { + name.as_str() + .map(|name| Some(name.to_string())) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "Zarr v3 dimension_names entries must be strings or null" + .to_string(), + ) + }) + } + }) + .collect::>>() + }) + .transpose()?; + + let meta = ArrayMeta { + zarr_format: 3, + shape, + chunks, + dtype, + fill_value, + compressor: None, + codec_pipeline, + storage_layout, + chunk_key_encoding, + order: 'C', + filters: None, + }; + meta.validate()?; + Ok(ArrayNode { + format: ZarrFormat::V3, + meta, + attributes, + dimension_names, + native_dtype, + native_codecs, + }) +} + +fn parse_separator(value: &str, field: &str) -> ZarrFdwResult { + match value { + "." => Ok('.'), + "/" => Ok('/'), + _ => Err(ZarrFdwError::InvalidMetadata(format!( + "{field} must be '.' or '/', got '{value}'" + ))), + } +} + +fn validate_additional_fields( + object: &Map, + known: &[&str], + context: &str, +) -> ZarrFdwResult<()> { + for (field, value) in object { + if known.contains(&field.as_str()) { + continue; + } + let ignorable = value + .as_object() + .and_then(|extension| extension.get("must_understand")) + == Some(&Value::Bool(false)); + if !ignorable { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} contains unrecognized field '{field}' that is not marked must_understand=false" + ))); + } + } + Ok(()) +} + +fn validate_extension_object(object: &Map, context: &str) -> ZarrFdwResult<()> { + validate_exact_fields( + object, + &["name", "configuration", "must_understand"], + context, + )?; + if object + .get("must_understand") + .is_some_and(|value| !value.is_boolean()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} must_understand must be a boolean" + ))); + } + Ok(()) +} + +fn reject_ignorable_extension(object: &Map, context: &str) -> ZarrFdwResult<()> { + if object.get("must_understand") == Some(&Value::Bool(false)) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} cannot set must_understand=false" + ))); + } + Ok(()) +} + +fn validate_exact_fields( + object: &Map, + allowed: &[&str], + context: &str, +) -> ZarrFdwResult<()> { + if let Some(field) = object + .keys() + .find(|field| !allowed.contains(&field.as_str())) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} contains unsupported field '{field}'" + ))); + } + Ok(()) +} + +fn required_object<'a>( + object: &'a Map, + field: &str, +) -> ZarrFdwResult<&'a Map> { + object.get(field).and_then(Value::as_object).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("Zarr metadata field '{field}' must be an object")) + }) +} + +fn optional_object_ref<'a>( + object: &'a Map, + field: &str, +) -> ZarrFdwResult>> { + object + .get(field) + .map(|value| { + value.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr metadata field '{field}' must be an object" + )) + }) + }) + .transpose() +} + +fn optional_object(object: &Map, field: &str) -> ZarrFdwResult> { + Ok(optional_object_ref(object, field)? + .cloned() + .unwrap_or_default()) +} + +fn required_string<'a>(object: &'a Map, field: &str) -> ZarrFdwResult<&'a str> { + object.get(field).and_then(Value::as_str).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("Zarr metadata field '{field}' must be a string")) + }) +} + +fn required_u64(object: Option<&Map>, field: &str) -> ZarrFdwResult { + object + .and_then(|object| object.get(field)) + .and_then(Value::as_u64) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr metadata field '{field}' must be a non-negative integer" + )) + }) +} + +fn required_u64_array(object: &Map, field: &str) -> ZarrFdwResult> { + object + .get(field) + .and_then(Value::as_array) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("Zarr metadata field '{field}' must be an array")) + })? + .iter() + .map(|value| { + value.as_u64().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr metadata field '{field}' must contain non-negative integers" + )) + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn meta(shape: Vec, chunks: Vec) -> ArrayMeta { + ArrayMeta { + zarr_format: 2, + shape, + chunks, + dtype: " value["chunk_grid"]["name"] = Value::String("rectilinear".into()), + "pipeline" => { + value["codecs"] = serde_json::json!([{"name":"bytes","configuration":{"endian":"little"}}, {"name":"gzip"}]) + } + "transformer" => value["storage_transformers"] = serde_json::json!([{"name":"x"}]), + "endian" => { + value["codecs"][0]["configuration"]["endian"] = Value::String("big".into()) + } + "dtype" => value["data_type"] = Value::String("uint16".into()), + "fill" => value["fill_value"] = Value::Null, + _ => unreachable!(), + } + assert!( + parse_v3_node(serde_json::to_vec(&value).unwrap().as_slice()).is_err(), + "{mutate}" + ); + } + } + + #[test] + fn v3_metadata_is_fail_closed_and_fill_values_are_typed() { + let base = serde_json::json!({ + "zarr_format": 3, "node_type": "array", "shape": [2], + "data_type": "float32", + "chunk_grid": {"name":"regular","configuration":{"chunk_shape":[1]}}, + "chunk_key_encoding": {"name":"default"}, "fill_value": 0, + "codecs": [{"name":"bytes","configuration":{"endian":"little"}}], + "dimension_names": ["x"], "attributes": {} + }); + for (label, mutate) in [ + ("unknown", serde_json::json!({"unexpected": true})), + ( + "grid config", + serde_json::json!({"chunk_grid":{"configuration":{"extra":1}}}), + ), + ( + "key config", + serde_json::json!({"chunk_key_encoding":{"configuration":{"extra":1}}}), + ), + ( + "codec config", + serde_json::json!({"codecs":[{"configuration":{"extra":1}}]}), + ), + ("bad float fill", serde_json::json!({"fill_value":"bogus"})), + ] { + let mut value = base.clone(); + merge_json(&mut value, mutate); + assert!( + parse_v3_node(&serde_json::to_vec(&value).unwrap()).is_err(), + "{label}" + ); + } + + let mut integer = base.clone(); + integer["data_type"] = serde_json::json!("int8"); + integer["codecs"] = serde_json::json!([{"name":"bytes"}]); + integer["fill_value"] = serde_json::json!(128); + assert!(parse_v3_node(&serde_json::to_vec(&integer).unwrap()).is_err()); + + integer["fill_value"] = serde_json::json!(0); + integer["codecs"] = + serde_json::json!([{"name":"bytes","configuration":{"endian":"banana"}}]); + assert!(parse_v3_node(&serde_json::to_vec(&integer).unwrap()).is_err()); + + for field in ["chunk_grid", "chunk_key_encoding"] { + let mut value = base.clone(); + value[field]["must_understand"] = serde_json::json!(false); + assert!(parse_v3_node(&serde_json::to_vec(&value).unwrap()).is_err()); + } + + let mut ignorable = base.clone(); + ignorable["example_extension"] = + serde_json::json!({"name":"example","must_understand":false}); + assert!(parse_v3_node(&serde_json::to_vec(&ignorable).unwrap()).is_ok()); + + let consolidated = serde_json::json!({ + "zarr_format":3, + "node_type":"group", + "attributes":{}, + "consolidated_metadata":{"must_understand":false,"kind":"inline","metadata":{}} + }); + assert!(parse_v3_node(&serde_json::to_vec(&consolidated).unwrap()).is_err()); + + let mut array_consolidated = base.clone(); + array_consolidated["consolidated_metadata"] = consolidated["consolidated_metadata"].clone(); + assert!(parse_v3_node(&serde_json::to_vec(&array_consolidated).unwrap()).is_err()); + } + + fn merge_json(target: &mut Value, patch: Value) { + for (key, value) in patch.as_object().unwrap() { + if let (Some(target_object), Some(patch_object)) = ( + target.get_mut(key).and_then(Value::as_object_mut), + value.as_object(), + ) { + for (nested_key, nested_value) in patch_object { + if let (Some(target_nested), Some(patch_nested)) = ( + target_object + .get_mut(nested_key) + .and_then(Value::as_object_mut), + nested_value.as_object(), + ) { + for (leaf, leaf_value) in patch_nested { + target_nested.insert(leaf.clone(), leaf_value.clone()); + } + } else { + target_object.insert(nested_key.clone(), nested_value.clone()); + } + } + } else { + target[key] = value.clone(); + } + } + } + + #[test] + fn validates_rank_shape_and_chunk_arithmetic() { + for rank in [1, 4, MAX_SCAN_RANK] { + meta(vec![1; rank], vec![1; rank]).validate().unwrap(); + } + for rank in [0, MAX_SCAN_RANK + 1] { + assert!(matches!( + meta(vec![1; rank], vec![1; rank]).validate(), + Err(ZarrFdwError::UnsupportedRank { rank: actual }) if actual == rank + )); + } + assert!(meta(vec![1, 0], vec![1, 1]).validate().is_err()); + assert!(meta(vec![1, 1], vec![1, 0]).validate().is_err()); + assert!(meta(vec![1, 1], vec![1]).validate().is_err()); + assert!( + meta(vec![u64::MAX; 2], vec![u64::MAX; 2]) + .validate() + .is_err() + ); + } + + #[test] + fn accepts_empty_v2_filters_and_rejects_non_empty() { + let mut value = meta(vec![2, 2], vec![1, 1]); + value.filters = Some(vec![]); + value.validate().unwrap(); + value.filters = Some(vec![serde_json::json!({"id":"delta"})]); + assert!(value.validate().is_err()); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/metrics.rs b/wrappers/src/fdw/zarr_fdw/metrics.rs new file mode 100644 index 000000000..efe486fc4 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/metrics.rs @@ -0,0 +1,521 @@ +//! Query-local Zarr execution metrics and their typed EXPLAIN representation. +//! +//! This module contains no PostgreSQL calls. The shared FDW framework owns +//! rendering [`ExplainProperty`] values, which keeps the counters usable in +//! ordinary Rust tests and avoids coupling scan bookkeeping to `pg_sys`. + +use std::time::Duration; + +use supabase_wrappers::prelude::ExplainProperty; + +/// The purpose of an object-store GET made while executing a scan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ReadKind { + Metadata, + Coordinate, + Data, +} + +/// Actual work performed by one query-local Zarr FDW instance. +/// +/// Byte counters use encoded object bytes for remote I/O and decoded bytes for +/// in-memory payloads. All updates saturate: observability must never make an +/// otherwise valid query fail because a counter overflowed. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct ZarrScanMetrics { + pub(crate) metadata_get_calls: u64, + pub(crate) coordinate_get_calls: u64, + pub(crate) data_get_calls: u64, + pub(crate) metadata_encoded_bytes: u64, + pub(crate) coordinate_encoded_bytes: u64, + pub(crate) data_encoded_bytes: u64, + pub(crate) coordinate_decoded_bytes: u64, + pub(crate) data_decoded_bytes: u64, + pub(crate) fill_bytes_synthesized: u64, + pub(crate) chunks_total: u64, + pub(crate) chunks_selected: u64, + pub(crate) chunks_coordinate_pruned: u64, + pub(crate) chunks_requested: u64, + pub(crate) chunks_present: u64, + pub(crate) chunks_missing: u64, + pub(crate) cache_hits: u64, + pub(crate) cache_misses: u64, + pub(crate) cache_evictions: u64, + pub(crate) shard_index_get_calls: u64, + pub(crate) shard_index_encoded_bytes: u64, + pub(crate) shard_payload_get_calls: u64, + pub(crate) shard_payload_encoded_bytes: u64, + pub(crate) shard_index_cache_hits: u64, + pub(crate) shard_index_cache_misses: u64, + pub(crate) shard_index_cache_evictions: u64, + pub(crate) logical_cells_examined: u64, + /// Present only when the FDW evaluates every scan qual exactly. + pub(crate) logical_cells_matched: Option, + pub(crate) tuples_emitted: u64, + pub(crate) rescans: u64, + pub(crate) decompression_micros: u64, + pub(crate) decoding_micros: u64, + pub(crate) aggregate_micros: u64, +} + +/// Runtime metadata and configured resource bounds shown beside actual work. +pub(crate) struct ZarrExplainContext<'a> { + pub(crate) array: &'a str, + pub(crate) dimensions: &'a [String], + pub(crate) shape: &'a [u64], + pub(crate) chunk_shape: &'a [usize], + pub(crate) dtype: &'a str, + pub(crate) codec: &'a str, + pub(crate) storage_backend: &'a str, + pub(crate) storage_layout: &'a str, + pub(crate) shard_shape: Option<&'a [u64]>, + pub(crate) index_location: Option<&'a str>, + pub(crate) aggregate_mode: &'a str, + pub(crate) max_concurrent_reads: usize, + pub(crate) max_inflight_bytes: usize, + pub(crate) compressed_cache_bytes: usize, + pub(crate) cache_entries: usize, + pub(crate) cache_resident_bytes: usize, + pub(crate) shard_index_cache_bytes: usize, + pub(crate) shard_index_cache_entries: usize, + pub(crate) shard_index_cache_resident_bytes: usize, +} + +impl ZarrScanMetrics { + /// Record one actual remote GET. `encoded_bytes` is `None` for a missing + /// object and `Some(0)` for a present empty object. + pub(crate) fn record_remote_get(&mut self, kind: ReadKind, encoded_bytes: Option) { + self.record_remote_request(kind); + if let Some(encoded_bytes) = encoded_bytes { + self.record_remote_response_bytes(kind, encoded_bytes); + } + } + + /// Record a remote request when it starts. Keeping this separate from the + /// response bytes lets cancelled or prefetched-but-unconsumed work remain + /// visible even when no response reaches the scan loop. + pub(crate) fn record_remote_request(&mut self, kind: ReadKind) { + match kind { + ReadKind::Metadata => { + saturating_increment(&mut self.metadata_get_calls); + } + ReadKind::Coordinate => { + saturating_increment(&mut self.coordinate_get_calls); + } + ReadKind::Data => { + saturating_increment(&mut self.data_get_calls); + } + } + } + + /// Record bytes received from a completed remote request without changing + /// its request count. + pub(crate) fn record_remote_response_bytes(&mut self, kind: ReadKind, encoded_bytes: usize) { + let bytes = usize_to_u64(encoded_bytes); + match kind { + ReadKind::Metadata => saturating_add(&mut self.metadata_encoded_bytes, bytes), + ReadKind::Coordinate => saturating_add(&mut self.coordinate_encoded_bytes, bytes), + ReadKind::Data => saturating_add(&mut self.data_encoded_bytes, bytes), + } + } + + /// Set the current candidate-chunk selection. Selection is a property of + /// the current scan bounds rather than cumulative work across rescans. + pub(crate) fn set_chunk_selection(&mut self, total: u64, selected: u64) { + self.chunks_total = total; + self.chunks_selected = selected; + self.chunks_coordinate_pruned = total.saturating_sub(selected); + } + + /// Record one data-chunk request when it enters the ordered window. + pub(crate) fn record_chunk_request(&mut self) { + saturating_increment(&mut self.chunks_requested); + } + + /// Record the outcome of one consumed data-chunk request. + pub(crate) fn record_chunk_result(&mut self, present: bool) { + if present { + saturating_increment(&mut self.chunks_present); + } else { + saturating_increment(&mut self.chunks_missing); + } + } + + pub(crate) fn record_cache_lookup(&mut self, hit: bool) { + if hit { + saturating_increment(&mut self.cache_hits); + } else { + saturating_increment(&mut self.cache_misses); + } + } + + pub(crate) fn record_cache_evictions(&mut self, count: usize) { + saturating_add(&mut self.cache_evictions, usize_to_u64(count)); + } + + pub(crate) fn record_shard_index_get(&mut self, encoded_bytes: Option) { + saturating_increment(&mut self.shard_index_get_calls); + if let Some(encoded_bytes) = encoded_bytes { + saturating_add( + &mut self.shard_index_encoded_bytes, + usize_to_u64(encoded_bytes), + ); + } + } + + pub(crate) fn record_shard_payload_get(&mut self, encoded_bytes: Option) { + saturating_increment(&mut self.shard_payload_get_calls); + if let Some(encoded_bytes) = encoded_bytes { + saturating_add( + &mut self.shard_payload_encoded_bytes, + usize_to_u64(encoded_bytes), + ); + } + } + + pub(crate) fn record_shard_index_cache_lookup(&mut self, hit: bool) { + if hit { + saturating_increment(&mut self.shard_index_cache_hits); + } else { + saturating_increment(&mut self.shard_index_cache_misses); + } + } + + pub(crate) fn record_shard_index_cache_evictions(&mut self, count: usize) { + saturating_add(&mut self.shard_index_cache_evictions, usize_to_u64(count)); + } + + pub(crate) fn record_decoded_bytes( + &mut self, + kind: ReadKind, + decoded_bytes: usize, + synthesized_fill: bool, + ) { + let decoded_bytes = usize_to_u64(decoded_bytes); + if synthesized_fill { + saturating_add(&mut self.fill_bytes_synthesized, decoded_bytes); + } + match kind { + ReadKind::Metadata => {} + ReadKind::Coordinate => { + saturating_add(&mut self.coordinate_decoded_bytes, decoded_bytes); + } + ReadKind::Data => { + saturating_add(&mut self.data_decoded_bytes, decoded_bytes); + } + } + } + + pub(crate) fn record_cells(&mut self, examined: usize, matched: Option) { + saturating_add(&mut self.logical_cells_examined, usize_to_u64(examined)); + if let Some(matched) = matched { + let total = self.logical_cells_matched.get_or_insert(0); + saturating_add(total, usize_to_u64(matched)); + } + } + + pub(crate) fn record_tuple_emitted(&mut self) { + saturating_increment(&mut self.tuples_emitted); + } + + pub(crate) fn record_rescan(&mut self) { + saturating_increment(&mut self.rescans); + } + + pub(crate) fn record_decompression_time(&mut self, elapsed: Duration) { + saturating_add(&mut self.decompression_micros, duration_micros(elapsed)); + } + + pub(crate) fn record_decoding_time(&mut self, elapsed: Duration) { + saturating_add(&mut self.decoding_micros, duration_micros(elapsed)); + } + + pub(crate) fn record_aggregate_time(&mut self, elapsed: Duration) { + saturating_add(&mut self.aggregate_micros, duration_micros(elapsed)); + } + + pub(crate) fn total_get_calls(&self) -> u64 { + self.metadata_get_calls + .saturating_add(self.coordinate_get_calls) + .saturating_add(self.data_get_calls) + } + + pub(crate) fn total_encoded_bytes(&self) -> u64 { + self.metadata_encoded_bytes + .saturating_add(self.coordinate_encoded_bytes) + .saturating_add(self.data_encoded_bytes) + } + + pub(crate) fn total_decoded_bytes(&self) -> u64 { + self.coordinate_decoded_bytes + .saturating_add(self.data_decoded_bytes) + } + + /// Build structured properties for the shared framework EXPLAIN hook. + pub(crate) fn explain_properties( + &self, + context: ZarrExplainContext<'_>, + ) -> Vec { + let mut properties = vec![ + ExplainProperty::text("Zarr Array", context.array), + ExplainProperty::text( + "Zarr Dimensions", + format!("[{}]", context.dimensions.join(", ")), + ), + ExplainProperty::text("Zarr Shape", format!("{:?}", context.shape)), + ExplainProperty::text("Zarr Chunk Shape", format!("{:?}", context.chunk_shape)), + ExplainProperty::text("Zarr Dtype", context.dtype), + ExplainProperty::text("Zarr Codec", context.codec), + ExplainProperty::text("Zarr Storage Backend", context.storage_backend), + ExplainProperty::text("Zarr Storage Layout", context.storage_layout), + ExplainProperty::text("Zarr Aggregate Pushdown", context.aggregate_mode), + ExplainProperty::text("Zarr Chunk-Stat Pruning", "disabled"), + ExplainProperty::unsigned( + "Zarr Max Concurrent Reads", + usize_to_u64(context.max_concurrent_reads), + ), + ExplainProperty::unsigned_with_unit( + "Zarr Max Inflight Bytes", + usize_to_u64(context.max_inflight_bytes), + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Compressed Cache Capacity", + usize_to_u64(context.compressed_cache_bytes), + "bytes", + ), + ExplainProperty::unsigned( + "Zarr Compressed Cache Entries", + usize_to_u64(context.cache_entries), + ), + ExplainProperty::unsigned_with_unit( + "Zarr Compressed Cache Resident", + usize_to_u64(context.cache_resident_bytes), + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Shard Index Cache Capacity", + usize_to_u64(context.shard_index_cache_bytes), + "bytes", + ), + ExplainProperty::unsigned( + "Zarr Shard Index Cache Entries", + usize_to_u64(context.shard_index_cache_entries), + ), + ExplainProperty::unsigned_with_unit( + "Zarr Shard Index Cache Resident", + usize_to_u64(context.shard_index_cache_resident_bytes), + "bytes", + ), + ExplainProperty::unsigned("Zarr Chunks Total", self.chunks_total), + ExplainProperty::unsigned("Zarr Chunks Selected", self.chunks_selected), + ExplainProperty::unsigned( + "Zarr Chunks Coordinate-Pruned", + self.chunks_coordinate_pruned, + ), + ExplainProperty::unsigned("Zarr Chunks Requested", self.chunks_requested), + ExplainProperty::unsigned("Zarr Chunks Present", self.chunks_present), + ExplainProperty::unsigned("Zarr Chunks Missing", self.chunks_missing), + ExplainProperty::unsigned("Zarr Remote GET Calls", self.total_get_calls()), + ExplainProperty::unsigned("Zarr Metadata GET Calls", self.metadata_get_calls), + ExplainProperty::unsigned("Zarr Coordinate GET Calls", self.coordinate_get_calls), + ExplainProperty::unsigned("Zarr Data GET Calls", self.data_get_calls), + ExplainProperty::unsigned("Zarr Cache Hits", self.cache_hits), + ExplainProperty::unsigned("Zarr Cache Misses", self.cache_misses), + ExplainProperty::unsigned("Zarr Cache Evictions", self.cache_evictions), + ExplainProperty::unsigned("Zarr Shard Index GET Calls", self.shard_index_get_calls), + ExplainProperty::unsigned("Zarr Shard Payload GET Calls", self.shard_payload_get_calls), + ExplainProperty::unsigned("Zarr Shard Index Cache Hits", self.shard_index_cache_hits), + ExplainProperty::unsigned( + "Zarr Shard Index Cache Misses", + self.shard_index_cache_misses, + ), + ExplainProperty::unsigned( + "Zarr Shard Index Cache Evictions", + self.shard_index_cache_evictions, + ), + ExplainProperty::unsigned_with_unit( + "Zarr Remote Encoded Bytes", + self.total_encoded_bytes(), + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Metadata Encoded Bytes", + self.metadata_encoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Coordinate Encoded Bytes", + self.coordinate_encoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Data Encoded Bytes", + self.data_encoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Shard Index Encoded Bytes", + self.shard_index_encoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Shard Payload Encoded Bytes", + self.shard_payload_encoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Decoded Bytes", + self.total_decoded_bytes(), + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Coordinate Decoded Bytes", + self.coordinate_decoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Data Decoded Bytes", + self.data_decoded_bytes, + "bytes", + ), + ExplainProperty::unsigned_with_unit( + "Zarr Fill Bytes Synthesized", + self.fill_bytes_synthesized, + "bytes", + ), + ExplainProperty::unsigned("Zarr Logical Cells Examined", self.logical_cells_examined), + ExplainProperty::unsigned("Zarr Tuples Emitted", self.tuples_emitted), + ExplainProperty::unsigned("Zarr Rescans", self.rescans), + ExplainProperty::unsigned_with_unit( + "Zarr Decompression Time", + self.decompression_micros, + "us", + ), + ExplainProperty::unsigned_with_unit("Zarr Decoding Time", self.decoding_micros, "us"), + ExplainProperty::unsigned_with_unit("Zarr Aggregate Time", self.aggregate_micros, "us"), + ]; + + if let Some(shard_shape) = context.shard_shape { + properties.push(ExplainProperty::text( + "Zarr Shard Shape", + format!("{shard_shape:?}"), + )); + } + if let Some(index_location) = context.index_location { + properties.push(ExplainProperty::text( + "Zarr Shard Index Location", + index_location, + )); + } + + if let Some(matched) = self.logical_cells_matched { + properties.push(ExplainProperty::unsigned( + "Zarr Logical Cells Matched", + matched, + )); + } + + properties + } +} + +fn saturating_increment(value: &mut u64) { + *value = value.saturating_add(1); +} + +fn saturating_add(value: &mut u64, increment: u64) { + *value = value.saturating_add(increment); +} + +fn usize_to_u64(value: usize) -> u64 { + u64::try_from(value).unwrap_or(u64::MAX) +} + +fn duration_micros(elapsed: Duration) -> u64 { + u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX) +} + +#[cfg(test)] +mod tests { + use super::*; + use supabase_wrappers::prelude::ExplainValue; + + #[test] + fn records_remote_cache_fill_and_cell_work_without_double_counting_bytes() { + let mut metrics = ZarrScanMetrics::default(); + metrics.set_chunk_selection(4, 4); + metrics.record_remote_get(ReadKind::Metadata, Some(1_127)); + metrics.record_remote_get(ReadKind::Coordinate, Some(64)); + + for _ in 0..3 { + metrics.record_chunk_request(); + metrics.record_cache_lookup(false); + metrics.record_remote_get(ReadKind::Data, Some(96)); + metrics.record_chunk_result(true); + metrics.record_decoded_bytes(ReadKind::Data, 96, false); + } + metrics.record_chunk_request(); + metrics.record_cache_lookup(false); + metrics.record_remote_get(ReadKind::Data, None); + metrics.record_chunk_result(false); + metrics.record_decoded_bytes(ReadKind::Data, 96, true); + metrics.record_decoded_bytes(ReadKind::Coordinate, 64, true); + metrics.record_cells(40, Some(30)); + metrics.record_tuple_emitted(); + + assert_eq!(metrics.total_get_calls(), 6); + assert_eq!(metrics.total_encoded_bytes(), 1_479); + assert_eq!(metrics.total_decoded_bytes(), 448); + assert_eq!(metrics.chunks_requested, 4); + assert_eq!(metrics.chunks_present, 3); + assert_eq!(metrics.chunks_missing, 1); + assert_eq!(metrics.fill_bytes_synthesized, 160); + assert_eq!(metrics.logical_cells_matched, Some(30)); + } + + #[test] + fn explain_properties_preserve_numeric_types_and_unknown_match_count() { + let metrics = ZarrScanMetrics::default(); + let dimensions = vec!["time".to_string(), "y".to_string(), "x".to_string()]; + let shape = [2, 5, 6]; + let chunk_shape = [1, 5, 3]; + let properties = metrics.explain_properties(ZarrExplainContext { + array: "cube/reflectance", + dimensions: &dimensions, + shape: &shape, + chunk_shape: &chunk_shape, + dtype: ">), + + #[error("list request failed: {0}")] + ListRequestError(#[from] Box>), + + #[error("PostgreSQL catalog query failed: {0}")] + SpiError(#[from] pgrx::spi::Error), + + #[error("parse JSON response failed: {0}")] + JsonParseError(#[from] serde_json::Error), + + #[error("read data failed: {0}")] + ReadError(#[from] std::io::Error), + + #[error("coordinate value {0} is out of range for pg timestamptz")] + TimeOutOfRange(f64), + + #[error("invalid CRS metadata for zarr array '{array}': {message}")] + InvalidCrs { array: String, message: String }, + + #[error("PostGIS is unavailable for zarr spatial operations: {0}")] + PostgisUnavailable(String), + + #[error("invalid PostGIS geometry: {0}")] + InvalidGeometry(String), + + #[error("{0}")] + NumericConversionError(#[from] pgrx::numeric::Error), +} + +impl From for ErrorReport { + fn from(value: ZarrFdwError) -> Self { + let code = match &value { + ZarrFdwError::FileStoreDefinitionRequiresSuperuser + | ZarrFdwError::FileStoreOwnerRequiresSuperuser + | ZarrFdwError::HttpStoreDefinitionRequiresSuperuser + | ZarrFdwError::HttpStoreOwnerRequiresSuperuser => { + PgSqlErrorCode::ERRCODE_INSUFFICIENT_PRIVILEGE + } + ZarrFdwError::InvalidCrs { .. } | ZarrFdwError::InvalidGeometry(_) => { + PgSqlErrorCode::ERRCODE_INVALID_PARAMETER_VALUE + } + ZarrFdwError::PostgisUnavailable(_) => PgSqlErrorCode::ERRCODE_FEATURE_NOT_SUPPORTED, + _ => PgSqlErrorCode::ERRCODE_FDW_ERROR, + }; + ErrorReport::new(code, format!("{value}"), "") + } +} + +impl From> for ZarrFdwError { + fn from(value: SdkError) -> Self { + Self::RequestError(value.into()) + } +} + +impl From> for ZarrFdwError { + fn from(value: SdkError) -> Self { + Self::ListRequestError(value.into()) + } +} + +type ZarrFdwResult = Result; diff --git a/wrappers/src/fdw/zarr_fdw/ome.rs b/wrappers/src/fdw/zarr_fdw/ome.rs new file mode 100644 index 000000000..28f447608 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/ome.rs @@ -0,0 +1,991 @@ +//! Strict OME-Zarr 0.5 multiscale metadata adapter. +//! +//! OME-Zarr 0.5 is a Zarr v3 convention. This module parses only the bounded +//! inline affine subset required by the initial rank-2 executor: every level +//! has one scale followed by an optional translation, and an optional group +//! transform with the same grammar is applied afterwards. + +use std::collections::HashSet; + +use serde::Serialize; +use serde_json::{Map, Value}; + +use super::{ZarrFdwError, ZarrFdwResult}; + +const OME_VERSION: &str = "0.5"; +const MAX_MULTISCALES: usize = 1_024; +const MAX_LEVELS: usize = 10_000; + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(crate) struct OmeAxis { + pub(crate) name: String, + #[serde(rename = "type", skip_serializing_if = "Option::is_none")] + pub(crate) kind: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) unit: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct AffineTransform { + pub(crate) scale: Vec, + pub(crate) translation: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct OmeLevel { + pub(crate) relative_path: String, + /// Level transform after composing the optional group transform. + pub(crate) effective_transform: AffineTransform, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct OmeMultiscale { + pub(crate) name: Option, + pub(crate) axes: Vec, + pub(crate) levels: Vec, + /// Non-fatal violations of OME `SHOULD` recommendations. + pub(crate) warnings: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ResolvedOmeLevel { + /// Canonical path relative to the configured store; empty means root. + pub(crate) group_path: String, + pub(crate) multiscale_index: usize, + pub(crate) multiscale_name: Option, + pub(crate) level_index: usize, + pub(crate) array_path: String, + pub(crate) axes: Vec, + pub(crate) transform: AffineTransform, + pub(crate) warnings: Vec, +} + +/// Parse every OME-Zarr 0.5 multiscale declared by one group. +pub(crate) fn parse_ome_05_multiscales( + group_path: &str, + attributes: &Map, +) -> ZarrFdwResult> { + let display_group = display_path(group_path); + validate_optional_ome_05_attributes(group_path, Some(attributes))?; + let ome = attributes + .get("ome") + .and_then(Value::as_object) + .ok_or_else(|| { + invalid(format!( + "OME-Zarr group '{display_group}' must define attributes.ome as an object" + )) + })?; + let raw_multiscales = required_array(ome, "multiscales", "attributes.ome")?; + if raw_multiscales.is_empty() { + return Err(invalid(format!( + "OME-Zarr group '{display_group}' must define a non-empty ome.multiscales array" + ))); + } + if raw_multiscales.len() > MAX_MULTISCALES { + return Err(invalid(format!( + "OME-Zarr group '{display_group}' declares {} multiscales, exceeding the limit of {MAX_MULTISCALES}", + raw_multiscales.len() + ))); + } + + let mut multiscales = Vec::new(); + multiscales + .try_reserve_exact(raw_multiscales.len()) + .map_err(|_| invalid("could not allocate OME-Zarr multiscale metadata"))?; + let mut total_levels = 0usize; + for (multiscale_index, value) in raw_multiscales.iter().enumerate() { + multiscales.push(parse_multiscale( + value, + display_group, + multiscale_index, + &mut total_levels, + )?); + } + Ok(multiscales) +} + +/// Parse and resolve one explicit zero-based multiscale/level selection. +pub(crate) fn resolve_ome_05_level( + group_path: &str, + attributes: &Map, + multiscale_index: usize, + level_index: usize, +) -> ZarrFdwResult { + let canonical_group = canonical_ome_group_path(group_path)?; + let multiscales = parse_ome_05_multiscales(&canonical_group, attributes)?; + let multiscale = multiscales.get(multiscale_index).ok_or_else(|| { + invalid(format!( + "multiscale index {multiscale_index} is outside the {} multiscales declared by OME-Zarr group '{}'", + multiscales.len(), + display_path(&canonical_group) + )) + })?; + let level = multiscale.levels.get(level_index).ok_or_else(|| { + invalid(format!( + "multiscale level {level_index} is outside the {} levels declared by multiscale {multiscale_index} in OME-Zarr group '{}'", + multiscale.levels.len(), + display_path(&canonical_group) + )) + })?; + let array_path = join_relative_path(&canonical_group, &level.relative_path); + + Ok(ResolvedOmeLevel { + group_path: canonical_group, + multiscale_index, + multiscale_name: multiscale.name.clone(), + level_index, + array_path, + axes: multiscale.axes.clone(), + transform: level.effective_transform.clone(), + warnings: multiscale.warnings.clone(), + }) +} + +/// Canonicalize a user-selected OME group path relative to the store root. +/// +/// Empty and `/` both denote the root. Outer slashes are removed for nested +/// groups; every remaining component must be a valid, non-reserved Zarr v3 +/// node component. +pub(crate) fn canonical_ome_group_path(group_path: &str) -> ZarrFdwResult { + if group_path.is_empty() || group_path == "/" { + return Ok(String::new()); + } + let canonical = group_path.trim_matches('/'); + validate_relative_path(canonical).map_err(|message| { + invalid(format!( + "OME-Zarr group path '{group_path}' is invalid: {message}" + )) + })?; + Ok(canonical.to_string()) +} + +/// Validate OME-Zarr version consistency on an optional hierarchy node. +/// +/// Ordinary Zarr v3 ancestor groups do not need OME metadata. When an `ome` +/// block is present, however, it must be an object with version `0.5` even if +/// the group is a plate, well, labels container, or another non-multiscale +/// hierarchy node. +pub(crate) fn validate_optional_ome_05_attributes( + node_path: &str, + attributes: Option<&Map>, +) -> ZarrFdwResult<()> { + let Some(ome_value) = attributes.and_then(|attributes| attributes.get("ome")) else { + return Ok(()); + }; + let ome = ome_value.as_object().ok_or_else(|| { + invalid(format!( + "OME-Zarr metadata at '{}' must define attributes.ome as an object", + display_path(node_path) + )) + })?; + let version = ome.get("version").and_then(Value::as_str).ok_or_else(|| { + invalid(format!( + "OME-Zarr metadata at '{}' must define attributes.ome.version as a string", + display_path(node_path) + )) + })?; + if version != OME_VERSION { + return Err(invalid(format!( + "unsupported OME-Zarr version '{version}' in group '{}'; expected '{OME_VERSION}'", + display_path(node_path) + ))); + } + Ok(()) +} + +fn parse_multiscale( + value: &Value, + display_group: &str, + multiscale_index: usize, + total_levels: &mut usize, +) -> ZarrFdwResult { + let context = format!("OME-Zarr multiscale {multiscale_index} in group '{display_group}'"); + let object = value + .as_object() + .ok_or_else(|| invalid(format!("{context} must be an object")))?; + validate_fields( + object, + &[ + "name", + "axes", + "datasets", + "coordinateTransformations", + "type", + "metadata", + ], + &context, + )?; + + let mut warnings = Vec::new(); + let name = match object.get("name") { + None => { + warnings.push(format!("{context} should define a name")); + None + } + Some(value) => { + let name = value + .as_str() + .ok_or_else(|| invalid(format!("{context} name must be a string")))?; + if name.is_empty() { + return Err(invalid(format!("{context} name must not be empty"))); + } + Some(name.to_string()) + } + }; + if let Some(value) = object.get("type") { + value + .as_str() + .ok_or_else(|| invalid(format!("{context} type must be a string")))?; + } + if let Some(value) = object.get("metadata") { + value + .as_object() + .ok_or_else(|| invalid(format!("{context} metadata must be an object")))?; + } + + let axes = parse_axes( + required_array(object, "axes", &context)?, + &context, + &mut warnings, + )?; + let rank = axes.len(); + let group_transform = object + .get("coordinateTransformations") + .map(|value| parse_transformations(value, rank, "group", &context)) + .transpose()? + .unwrap_or_else(|| AffineTransform::identity(rank)); + + let datasets = required_array(object, "datasets", &context)?; + if datasets.is_empty() { + return Err(invalid(format!("{context} datasets must not be empty"))); + } + *total_levels = total_levels + .checked_add(datasets.len()) + .ok_or_else(|| invalid("OME-Zarr level count overflowed"))?; + if *total_levels > MAX_LEVELS { + return Err(invalid(format!( + "OME-Zarr metadata declares {total_levels} levels, exceeding the limit of {MAX_LEVELS}" + ))); + } + + let mut levels = Vec::new(); + levels + .try_reserve_exact(datasets.len()) + .map_err(|_| invalid("could not allocate OME-Zarr level metadata"))?; + let mut dataset_paths = HashSet::new(); + for (level_index, dataset) in datasets.iter().enumerate() { + let level_context = format!("{context} level {level_index}"); + let dataset = dataset + .as_object() + .ok_or_else(|| invalid(format!("{level_context} must be an object")))?; + validate_fields( + dataset, + &["path", "coordinateTransformations"], + &level_context, + )?; + let relative_path = required_string(dataset, "path", &level_context)?.to_string(); + validate_relative_path(&relative_path).map_err(|message| { + invalid(format!( + "OME-Zarr dataset path '{relative_path}' in {level_context} is not a valid relative Zarr node path: {message}" + )) + })?; + if !dataset_paths.insert(relative_path.clone()) { + return Err(invalid(format!( + "{context} dataset paths must be unique; '{relative_path}' is duplicated" + ))); + } + let level_transform = parse_transformations( + dataset.get("coordinateTransformations").ok_or_else(|| { + invalid(format!( + "{level_context} must define coordinateTransformations" + )) + })?, + rank, + "dataset", + &level_context, + )?; + levels.push(OmeLevel { + relative_path, + effective_transform: compose_affine( + &level_transform, + &group_transform, + &level_context, + )?, + }); + } + + Ok(OmeMultiscale { + name, + axes, + levels, + warnings, + }) +} + +fn parse_axes( + values: &[Value], + context: &str, + warnings: &mut Vec, +) -> ZarrFdwResult> { + if !(2..=5).contains(&values.len()) { + return Err(invalid(format!( + "{context} axes length must be between 2 and 5, got {}", + values.len() + ))); + } + let mut axes = Vec::new(); + axes.try_reserve_exact(values.len()) + .map_err(|_| invalid("could not allocate OME-Zarr axis metadata"))?; + let mut names = HashSet::new(); + for (index, value) in values.iter().enumerate() { + let axis_context = format!("{context} axis {index}"); + let object = value + .as_object() + .ok_or_else(|| invalid(format!("{axis_context} must be an object")))?; + validate_fields(object, &["name", "type", "unit"], &axis_context)?; + let name = required_string(object, "name", &axis_context)?; + validate_axis_name(name) + .map_err(|message| invalid(format!("{axis_context} name is invalid: {message}")))?; + if !names.insert(name) { + return Err(invalid(format!( + "{context} axis names must be unique; '{name}' is duplicated" + ))); + } + let kind = match object.get("type") { + None | Some(Value::Null) => { + warnings.push(format!("{axis_context} should define a type")); + None + } + Some(value) => Some( + value + .as_str() + .ok_or_else(|| { + invalid(format!("{axis_context} type must be a string or null")) + })? + .to_string(), + ), + }; + let unit = match object.get("unit") { + None => { + if matches!(kind.as_deref(), Some("space" | "time")) { + warnings.push(format!("{axis_context} should define a unit")); + } + None + } + Some(value) => { + let unit = value + .as_str() + .ok_or_else(|| invalid(format!("{axis_context} unit must be a string")))?; + if unit.is_empty() { + return Err(invalid(format!("{axis_context} unit must not be empty"))); + } + Some(unit.to_string()) + } + }; + axes.push(OmeAxis { + name: name.to_string(), + kind, + unit, + }); + } + validate_axis_types(&axes, context)?; + Ok(axes) +} + +fn validate_axis_types(axes: &[OmeAxis], context: &str) -> ZarrFdwResult<()> { + let spatial = axes + .iter() + .filter(|axis| axis.kind.as_deref() == Some("space")) + .count(); + if !(2..=3).contains(&spatial) { + return Err(invalid(format!( + "{context} axes must contain 2 or 3 entries with type 'space', found {spatial}" + ))); + } + let time = axes + .iter() + .filter(|axis| axis.kind.as_deref() == Some("time")) + .count(); + if time > 1 { + return Err(invalid(format!( + "{context} axes may contain at most one time axis" + ))); + } + let auxiliary = axes.len() - spatial - time; + if auxiliary > 1 { + return Err(invalid(format!( + "{context} axes may contain at most one channel, custom, or null-type axis" + ))); + } + + let mut previous_order = 0u8; + for axis in axes { + let order = match axis.kind.as_deref() { + Some("time") => 0, + Some("space") => 2, + _ => 1, + }; + if order < previous_order { + return Err(invalid(format!( + "{context} axes must be ordered time, then channel/custom, then space" + ))); + } + previous_order = order; + } + Ok(()) +} + +fn parse_transformations( + value: &Value, + rank: usize, + location: &str, + context: &str, +) -> ZarrFdwResult { + let transforms = value.as_array().ok_or_else(|| { + invalid(format!( + "{context} {location} coordinateTransformations must be an array" + )) + })?; + if !(1..=2).contains(&transforms.len()) { + return Err(invalid(format!( + "{context} {location} coordinateTransformations must contain scale followed by optional translation" + ))); + } + let scale = parse_transform(&transforms[0], "scale", rank, location, context)?; + let translation = if transforms.len() == 2 { + parse_transform(&transforms[1], "translation", rank, location, context)? + } else { + vec![0.0; rank] + }; + Ok(AffineTransform { scale, translation }) +} + +fn parse_transform( + value: &Value, + expected_type: &str, + rank: usize, + location: &str, + context: &str, +) -> ZarrFdwResult> { + let object = value.as_object().ok_or_else(|| { + invalid(format!( + "{context} {location} {expected_type} transform must be an object" + )) + })?; + if object.contains_key("path") { + return Err(invalid(format!( + "path-backed OME-Zarr transforms are not supported in {context}" + ))); + } + let actual_type = required_string(object, "type", context)?; + if actual_type != expected_type { + return Err(invalid(format!( + "{context} {location} coordinateTransformations must contain scale followed by optional translation; expected '{expected_type}', found '{actual_type}'" + ))); + } + validate_fields(object, &["type", expected_type], context)?; + let values = required_array(object, expected_type, context)?; + if values.len() != rank { + return Err(invalid(format!( + "{context} {location} {expected_type} has {} values but the axes rank is {rank}", + values.len() + ))); + } + values + .iter() + .enumerate() + .map(|(axis, value)| { + let value = value.as_f64().filter(|value| value.is_finite()).ok_or_else(|| { + invalid(format!( + "{context} {location} {expected_type} value at axis {axis} must be a finite number" + )) + })?; + if expected_type == "scale" && value <= 0.0 { + return Err(invalid(format!( + "OME-Zarr scale values must be finite and greater than zero; {context} axis {axis} is {value}" + ))); + } + Ok(value) + }) + .collect() +} + +fn compose_affine( + level: &AffineTransform, + group: &AffineTransform, + context: &str, +) -> ZarrFdwResult { + if level.scale.len() != group.scale.len() + || level.translation.len() != group.translation.len() + || level.scale.len() != level.translation.len() + { + return Err(invalid(format!( + "{context} affine transform ranks do not match" + ))); + } + let mut scale = Vec::new(); + let mut translation = Vec::new(); + scale + .try_reserve_exact(level.scale.len()) + .map_err(|_| invalid("could not allocate composed OME-Zarr scale"))?; + translation + .try_reserve_exact(level.scale.len()) + .map_err(|_| invalid("could not allocate composed OME-Zarr translation"))?; + for axis in 0..level.scale.len() { + let effective_scale = group.scale[axis] * level.scale[axis]; + let effective_translation = + group.scale[axis] * level.translation[axis] + group.translation[axis]; + if !effective_scale.is_finite() || effective_scale <= 0.0 { + return Err(invalid(format!( + "{context} composed scale at axis {axis} is not finite and positive" + ))); + } + if !effective_translation.is_finite() { + return Err(invalid(format!( + "{context} composed translation at axis {axis} is not finite" + ))); + } + scale.push(effective_scale); + translation.push(effective_translation); + } + Ok(AffineTransform { scale, translation }) +} + +impl AffineTransform { + fn identity(rank: usize) -> Self { + Self { + scale: vec![1.0; rank], + translation: vec![0.0; rank], + } + } +} + +fn required_string<'a>( + object: &'a Map, + field: &str, + context: &str, +) -> ZarrFdwResult<&'a str> { + object + .get(field) + .and_then(Value::as_str) + .ok_or_else(|| invalid(format!("{context} must define '{field}' as a string"))) +} + +fn required_array<'a>( + object: &'a Map, + field: &str, + context: &str, +) -> ZarrFdwResult<&'a [Value]> { + object + .get(field) + .and_then(Value::as_array) + .map(Vec::as_slice) + .ok_or_else(|| invalid(format!("{context} must define '{field}' as an array"))) +} + +fn validate_fields( + object: &Map, + allowed: &[&str], + context: &str, +) -> ZarrFdwResult<()> { + if let Some(field) = object + .keys() + .find(|field| !allowed.contains(&field.as_str())) + { + return Err(invalid(format!( + "{context} contains unsupported field '{field}'" + ))); + } + Ok(()) +} + +fn validate_axis_name(name: &str) -> Result<(), &'static str> { + if name.is_empty() { + return Err("must not be empty"); + } + if name.trim() != name || name.chars().any(char::is_whitespace) { + return Err("must not contain whitespace"); + } + if name.chars().any(char::is_control) { + return Err("must not contain control characters"); + } + if name.contains('/') || name.contains('\\') || matches!(name, "." | "..") { + return Err("must not contain path components"); + } + if name == "zarr.json" || name.starts_with("__") || name.chars().all(|value| value == '.') { + return Err("must be a valid Zarr v3 dimension name"); + } + Ok(()) +} + +fn validate_relative_path(path: &str) -> Result<(), &'static str> { + if path.is_empty() { + return Err("must not be empty"); + } + if path.starts_with('/') || path.ends_with('/') || path.contains('\\') { + return Err("must be relative and use '/' separators"); + } + for component in path.split('/') { + if component.is_empty() || matches!(component, "." | ".." | "zarr.json") { + return Err("contains an invalid node component"); + } + if component.starts_with("__") || component.chars().all(|value| value == '.') { + return Err("contains a reserved Zarr v3 node component"); + } + if component.chars().any(char::is_control) { + return Err("contains a control character"); + } + } + Ok(()) +} + +fn join_relative_path(group_path: &str, relative_path: &str) -> String { + if group_path.is_empty() { + relative_path.to_string() + } else { + format!("{group_path}/{relative_path}") + } +} + +fn display_path(path: &str) -> &str { + if path.is_empty() { "/" } else { path } +} + +fn invalid(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(message.into()) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn attributes(multiscales: Value) -> Map { + json!({"ome": {"version": "0.5", "multiscales": multiscales}}) + .as_object() + .cloned() + .unwrap() + } + + fn valid_multiscale() -> Value { + json!({ + "name": "image", + "axes": [ + {"name": "y", "type": "space", "unit": "micrometer"}, + {"name": "x", "type": "space", "unit": "micrometer"} + ], + "datasets": [ + {"path": "0", "coordinateTransformations": [ + {"type": "scale", "scale": [2.0, 3.0]}, + {"type": "translation", "translation": [10.0, 100.0]} + ]}, + {"path": "pyramid/1", "coordinateTransformations": [ + {"type": "scale", "scale": [4.0, 6.0]} + ]} + ], + "coordinateTransformations": [ + {"type": "scale", "scale": [0.5, 2.0]}, + {"type": "translation", "translation": [-1.0, 5.0]} + ], + "type": "gaussian", + "metadata": {"method": "fixture"} + }) + } + + #[test] + fn parses_and_composes_level_then_group_transforms() { + let parsed = + parse_ome_05_multiscales("nested/image", &attributes(json!([valid_multiscale()]))) + .unwrap(); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].name.as_deref(), Some("image")); + assert!(parsed[0].warnings.is_empty()); + assert_eq!( + parsed[0].levels[0].effective_transform, + AffineTransform { + scale: vec![1.0, 6.0], + translation: vec![4.0, 205.0], + } + ); + + let resolved = resolve_ome_05_level( + "/nested/image/", + &attributes(json!([valid_multiscale()])), + 0, + 1, + ) + .unwrap(); + assert_eq!(resolved.group_path, "nested/image"); + assert_eq!(resolved.array_path, "nested/image/pyramid/1"); + assert_eq!(resolved.multiscale_name.as_deref(), Some("image")); + assert_eq!(resolved.transform.scale, vec![2.0, 12.0]); + assert_eq!(resolved.transform.translation, vec![-1.0, 5.0]); + } + + #[test] + fn absent_group_transform_is_identity_and_should_gaps_warn() { + let parsed = parse_ome_05_multiscales( + "/", + &attributes(json!([{ + "axes": [ + {"name": "c"}, + {"name": "y", "type": "space"}, + {"name": "x", "type": "space"} + ], + "datasets": [{ + "path": "0", + "coordinateTransformations": [{"type": "scale", "scale": [1, 2, 3]}] + }] + }])), + ) + .unwrap(); + assert_eq!( + parsed[0].levels[0].effective_transform.scale, + vec![1.0, 2.0, 3.0] + ); + assert_eq!( + parsed[0].levels[0].effective_transform.translation, + vec![0.0; 3] + ); + assert_eq!(parsed[0].warnings.len(), 4); + } + + #[test] + fn rejects_missing_or_wrong_ome_envelope() { + for (attributes, phrase) in [ + (Map::new(), "must define attributes.ome"), + ( + json!({"ome": []}).as_object().cloned().unwrap(), + "must define attributes.ome", + ), + ( + json!({"ome": {"version": "0.4", "multiscales": []}}) + .as_object() + .cloned() + .unwrap(), + "unsupported OME-Zarr version '0.4'", + ), + (attributes(json!([])), "non-empty ome.multiscales"), + ] { + let error = parse_ome_05_multiscales("bad", &attributes).unwrap_err(); + assert!( + error.to_string().contains(phrase), + "unexpected error: {error}" + ); + } + } + + #[test] + fn rejects_invalid_axes_and_ordering() { + let invalid_axes = [ + (json!([{"name":"y","type":"space"}]), "between 2 and 5"), + ( + json!([ + {"name":"y","type":"space"}, + {"name":"y","type":"space"} + ]), + "axis names must be unique", + ), + ( + json!([ + {"name":"y","type":"space"}, + {"name":"t","type":"time"}, + {"name":"x","type":"space"} + ]), + "must be ordered", + ), + ( + json!([ + {"name":"t","type":"time"}, + {"name":"c","type":"channel"}, + {"name":"other"}, + {"name":"y","type":"space"}, + {"name":"x","type":"space"} + ]), + "at most one channel, custom, or null-type axis", + ), + ( + json!([ + {"name":"row","type":"custom"}, + {"name":"column","type":"custom"} + ]), + "2 or 3 entries with type 'space'", + ), + ]; + for (axes, phrase) in invalid_axes { + let error = parse_ome_05_multiscales( + "bad", + &attributes(json!([{ + "name": "bad", + "axes": axes, + "datasets": [{ + "path":"0", + "coordinateTransformations":[{"type":"scale","scale":[1,1]}] + }] + }])), + ) + .unwrap_err(); + assert!( + error.to_string().contains(phrase), + "unexpected error: {error}" + ); + } + } + + #[test] + fn rejects_malformed_or_unsupported_transforms() { + let cases = [ + (json!([]), "scale followed by optional translation"), + ( + json!([{"type":"translation","translation":[0,0]}]), + "expected 'scale'", + ), + ( + json!([ + {"type":"scale","scale":[1,1]}, + {"type":"scale","scale":[1,1]} + ]), + "expected 'translation'", + ), + ( + json!([{"type":"scale","scale":[1]}]), + "has 1 values but the axes rank is 2", + ), + (json!([{"type":"scale","scale":[1,0]}]), "greater than zero"), + ( + json!([{"type":"scale","path":"transform/scale"}]), + "path-backed OME-Zarr transforms", + ), + ]; + for (transforms, phrase) in cases { + let error = parse_ome_05_multiscales( + "bad", + &attributes(json!([{ + "name":"bad", + "axes":[ + {"name":"y","type":"space"}, + {"name":"x","type":"space"} + ], + "datasets":[{"path":"0","coordinateTransformations":transforms}] + }])), + ) + .unwrap_err(); + assert!( + error.to_string().contains(phrase), + "unexpected error: {error}" + ); + } + } + + #[test] + fn rejects_unsafe_paths_and_out_of_range_selection() { + for path in ["", "/0", "../0", "a//b", "a\\b", "__private"] { + let mut multiscale = valid_multiscale(); + multiscale["datasets"][0]["path"] = json!(path); + let error = + parse_ome_05_multiscales("bad", &attributes(json!([multiscale]))).unwrap_err(); + assert!( + error + .to_string() + .contains("not a valid relative Zarr node path"), + "unexpected error for {path:?}: {error}" + ); + } + + let attrs = attributes(json!([valid_multiscale()])); + assert!( + resolve_ome_05_level("root", &attrs, 1, 0) + .unwrap_err() + .to_string() + .contains("multiscale index 1 is outside") + ); + assert!( + resolve_ome_05_level("root", &attrs, 0, 2) + .unwrap_err() + .to_string() + .contains("multiscale level 2 is outside") + ); + } + + #[test] + fn rejects_duplicate_dataset_paths() { + let mut multiscale = valid_multiscale(); + multiscale["datasets"][1]["path"] = json!("0"); + let error = parse_ome_05_multiscales("bad", &attributes(json!([multiscale]))).unwrap_err(); + assert!( + error + .to_string() + .contains("dataset paths must be unique; '0' is duplicated"), + "unexpected error: {error}" + ); + } + + #[test] + fn canonicalizes_safe_group_paths_and_rejects_reserved_components() { + assert_eq!(canonical_ome_group_path("").unwrap(), ""); + assert_eq!(canonical_ome_group_path("/").unwrap(), ""); + assert_eq!( + canonical_ome_group_path("/plates/A/1/").unwrap(), + "plates/A/1" + ); + + for path in [ + "//", + "a//b", + "a/./b", + "a/../b", + "a/zarr.json", + "a/__private", + "a/...", + "a\\b", + "a/line\nbreak", + ] { + let error = canonical_ome_group_path(path).unwrap_err(); + assert!( + error.to_string().contains("OME-Zarr group path"), + "unexpected error for {path:?}: {error}" + ); + } + } + + #[test] + fn validates_optional_hierarchy_ome_versions_without_requiring_multiscales() { + assert!(validate_optional_ome_05_attributes("/", None).is_ok()); + assert!(validate_optional_ome_05_attributes("plain", Some(&Map::new())).is_ok()); + + let hierarchy_only = json!({"ome": {"version": "0.5", "plate": {}}}) + .as_object() + .cloned() + .unwrap(); + assert!(validate_optional_ome_05_attributes("plate", Some(&hierarchy_only)).is_ok()); + + for (attributes, phrase) in [ + ( + json!({"ome": []}).as_object().cloned().unwrap(), + "attributes.ome as an object", + ), + ( + json!({"ome": {}}).as_object().cloned().unwrap(), + "attributes.ome.version as a string", + ), + ( + json!({"ome": {"version": 5}}).as_object().cloned().unwrap(), + "attributes.ome.version as a string", + ), + ( + json!({"ome": {"version": "0.4"}}) + .as_object() + .cloned() + .unwrap(), + "unsupported OME-Zarr version '0.4'", + ), + ] { + let error = validate_optional_ome_05_attributes("bad", Some(&attributes)).unwrap_err(); + assert!( + error.to_string().contains(phrase), + "unexpected error: {error}" + ); + } + } +} diff --git a/wrappers/src/fdw/zarr_fdw/prefetch.rs b/wrappers/src/fdw/zarr_fdw/prefetch.rs new file mode 100644 index 000000000..e63253c3c --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/prefetch.rs @@ -0,0 +1,596 @@ +//! Ordered, bounded foreground prefetch for encoded Zarr chunk objects. +//! +//! Futures are polled directly by the PostgreSQL backend's `Runtime::block_on` +//! call and are never spawned. This keeps output deterministic and makes +//! cancellation explicit: every queued future is dropped before +//! [`PrefetchNext::Interrupted`] is returned. The caller may then leave +//! `block_on` and raise PostgreSQL's interrupt on a clean Rust stack. + +use futures_util::FutureExt; +use futures_util::future::LocalBoxFuture; +use futures_util::stream::{FuturesOrdered, StreamExt}; +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; +use thiserror::Error; +use tokio::time::{MissedTickBehavior, interval}; + +use super::cache::{CachedObject, CompressedChunkCache}; +use super::store::ReadIdentity; + +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum PrefetchConfigError { + #[error("max concurrent reads must be greater than zero")] + NoConcurrentReads, + #[error("max inflight bytes must be greater than zero")] + EmptyByteBudget, + #[error("interrupt poll interval must be greater than zero")] + DisabledInterruptPolling, +} + +/// One object fetch plus caller-owned context such as an N-dimensional chunk +/// coordinate. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct PrefetchRequest { + pub context: T, + pub identity: ReadIdentity, + /// Maximum bytes the storage read is allowed to return. The prefetcher + /// reserves this conservative amount before issuing the request. + pub max_bytes: usize, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum ScheduleError { + /// Retry after consuming the next ordered result. + WindowFull(PrefetchRequest), + /// This request can never fit the configured inflight byte budget. + RequestTooLarge { + request: PrefetchRequest, + max_inflight_bytes: usize, + }, + /// A cached object must obey the same bounded-read contract as a remote + /// response, even if another caller previously used a larger limit. + CachedObjectTooLarge { + request: PrefetchRequest, + actual_bytes: usize, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PrefetchSource { + Cache, + Remote, + /// No object body exists: a sparse shard/index entry was resolved to the + /// array fill value before entering the ordered queue. + Synthesized, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct PrefetchedObject { + pub request: PrefetchRequest, + pub object: CachedObject, + pub source: PrefetchSource, + /// Actual encoded bytes fetched from the object store. Cache hits and + /// explicit missing-object responses report zero. + pub remote_bytes: usize, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum PrefetchNext { + Ready(PrefetchedObject), + FetchError { + request: PrefetchRequest, + error: E, + }, + /// All queued futures were dropped before this value was exposed. + Interrupted, + Empty, +} + +enum Completion { + Cached(PrefetchRequest, CachedObject), + Synthesized(PrefetchRequest, CachedObject), + Fetched(PrefetchRequest, Result>, E>), +} + +type Pending = LocalBoxFuture<'static, Completion>; + +/// A deterministic request window with both request-count and encoded-byte +/// backpressure. +pub(crate) struct OrderedPrefetch { + pending: FuturesOrdered>, + max_concurrent_reads: usize, + max_inflight_bytes: usize, + interrupt_poll_interval: Duration, + inflight_reads: usize, + reserved_bytes: usize, +} + +impl OrderedPrefetch { + pub(crate) fn new( + max_concurrent_reads: usize, + max_inflight_bytes: usize, + interrupt_poll_interval: Duration, + ) -> Result { + if max_concurrent_reads == 0 { + return Err(PrefetchConfigError::NoConcurrentReads); + } + if max_inflight_bytes == 0 { + return Err(PrefetchConfigError::EmptyByteBudget); + } + if interrupt_poll_interval.is_zero() { + return Err(PrefetchConfigError::DisabledInterruptPolling); + } + Ok(Self { + pending: FuturesOrdered::new(), + max_concurrent_reads, + max_inflight_bytes, + interrupt_poll_interval, + inflight_reads: 0, + reserved_bytes: 0, + }) + } + + /// Schedule one request without spawning it. + /// + /// `fetch` is not called on a cache hit. It receives an owned key so the + /// returned future can own a cloned storage client and remain independent + /// of the scan struct's borrow. + pub(crate) fn try_schedule( + &mut self, + request: PrefetchRequest, + cache: &mut CompressedChunkCache, + fetch: F, + ) -> Result> + where + F: FnOnce(ReadIdentity, usize) -> Fut, + Fut: Future>, E>> + 'static, + { + if self.pending.len() >= self.max_concurrent_reads { + return Err(ScheduleError::WindowFull(request)); + } + + if let Some(object) = cache.get_identity(&request.identity) { + if let CachedObject::Present(bytes) = &object + && bytes.len() > request.max_bytes + { + return Err(ScheduleError::CachedObjectTooLarge { + request, + actual_bytes: bytes.len(), + }); + } + self.pending + .push_back(async move { Completion::Cached(request, object) }.boxed_local()); + return Ok(PrefetchSource::Cache); + } + + if request.max_bytes > self.max_inflight_bytes { + return Err(ScheduleError::RequestTooLarge { + request, + max_inflight_bytes: self.max_inflight_bytes, + }); + } + if self.inflight_reads >= self.max_concurrent_reads + || self + .reserved_bytes + .checked_add(request.max_bytes) + .is_none_or(|next| next > self.max_inflight_bytes) + { + return Err(ScheduleError::WindowFull(request)); + } + + let identity = request.identity.clone(); + let max_bytes = request.max_bytes; + let future = fetch(identity, max_bytes); + self.pending + .push_back(async move { Completion::Fetched(request, future.await) }.boxed_local()); + self.inflight_reads += 1; + self.reserved_bytes += max_bytes; + Ok(PrefetchSource::Remote) + } + + /// Queue an already-resolved sparse result without performing a cache + /// lookup or creating a storage future. It still participates in ordered + /// delivery and the request-count window. + pub(crate) fn try_schedule_synthesized( + &mut self, + request: PrefetchRequest, + object: CachedObject, + ) -> Result> { + if self.pending.len() >= self.max_concurrent_reads { + return Err(ScheduleError::WindowFull(request)); + } + self.pending + .push_back(async move { Completion::Synthesized(request, object) }.boxed_local()); + Ok(PrefetchSource::Synthesized) + } + + /// Return the next result in scheduling order. + /// + /// `interrupt_requested` must only inspect cancellation state; it must not + /// raise a PostgreSQL error. On `Interrupted`, this method first clears the + /// entire queue and resets accounting. The caller can safely return from + /// `block_on` before invoking PostgreSQL interrupt processing. + pub(crate) async fn next_interruptible( + &mut self, + cache: &mut CompressedChunkCache, + mut interrupt_requested: I, + ) -> PrefetchNext + where + I: FnMut() -> bool, + { + if self.pending.is_empty() { + return PrefetchNext::Empty; + } + if interrupt_requested() { + self.clear(); + return PrefetchNext::Interrupted; + } + + let mut ticker = interval(self.interrupt_poll_interval); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + // Tokio intervals tick immediately once. Consume that tick so the + // queued I/O receives a full polling interval before the next check. + ticker.tick().await; + + loop { + enum Wait { + Completion(Option>), + Tick, + } + + let wait = { + let next = self.pending.next(); + tokio::pin!(next); + tokio::select! { + biased; + _ = ticker.tick() => Wait::Tick, + completion = &mut next => Wait::Completion(completion), + } + }; + + match wait { + Wait::Tick if interrupt_requested() => { + self.clear(); + return PrefetchNext::Interrupted; + } + Wait::Tick => continue, + Wait::Completion(None) => return PrefetchNext::Empty, + Wait::Completion(Some(Completion::Cached(request, object))) => { + return PrefetchNext::Ready(PrefetchedObject { + request, + object, + source: PrefetchSource::Cache, + remote_bytes: 0, + }); + } + Wait::Completion(Some(Completion::Synthesized(request, object))) => { + return PrefetchNext::Ready(PrefetchedObject { + request, + object, + source: PrefetchSource::Synthesized, + remote_bytes: 0, + }); + } + Wait::Completion(Some(Completion::Fetched(request, result))) => { + self.inflight_reads = self.inflight_reads.saturating_sub(1); + self.reserved_bytes = self.reserved_bytes.saturating_sub(request.max_bytes); + return match result { + Ok(Some(bytes)) => { + let remote_bytes = bytes.len(); + let bytes: Arc<[u8]> = Arc::from(bytes); + cache.insert_present_identity( + request.identity.clone(), + Arc::clone(&bytes), + ); + PrefetchNext::Ready(PrefetchedObject { + request, + object: CachedObject::Present(bytes), + source: PrefetchSource::Remote, + remote_bytes, + }) + } + Ok(None) => { + cache.insert_missing_identity(request.identity.clone()); + PrefetchNext::Ready(PrefetchedObject { + request, + object: CachedObject::Missing, + source: PrefetchSource::Remote, + remote_bytes: 0, + }) + } + Err(error) => { + // A storage failure aborts this ordered window. Do + // not leave later requests alive while the caller + // converts the error into a PostgreSQL error. + self.clear(); + PrefetchNext::FetchError { request, error } + } + }; + } + } + } + } + + pub(crate) fn clear(&mut self) { + self.pending = FuturesOrdered::new(); + self.inflight_reads = 0; + self.reserved_bytes = 0; + } + + pub(crate) fn is_empty(&self) -> bool { + self.pending.is_empty() + } + + #[cfg(test)] + pub(crate) fn inflight_reads(&self) -> usize { + self.inflight_reads + } + + #[cfg(test)] + pub(crate) fn reserved_bytes(&self) -> usize { + self.reserved_bytes + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + use std::convert::Infallible; + use std::rc::Rc; + + fn request(id: usize, max_bytes: usize) -> PrefetchRequest { + PrefetchRequest { + context: id, + identity: ReadIdentity::whole(format!("chunk-{id}")), + max_bytes, + } + } + + fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + } + + #[test] + fn invalid_limits_are_rejected() { + assert!(matches!( + OrderedPrefetch::<(), Infallible>::new(0, 1, Duration::from_millis(1)), + Err(PrefetchConfigError::NoConcurrentReads) + )); + assert!(matches!( + OrderedPrefetch::<(), Infallible>::new(1, 0, Duration::from_millis(1)), + Err(PrefetchConfigError::EmptyByteBudget) + )); + assert!(matches!( + OrderedPrefetch::<(), Infallible>::new(1, 1, Duration::ZERO), + Err(PrefetchConfigError::DisabledInterruptPolling) + )); + } + + #[test] + fn results_remain_in_schedule_order_when_completions_do_not() { + let rt = runtime(); + rt.block_on(async { + let mut cache = CompressedChunkCache::new(64, 8); + let mut prefetch = + OrderedPrefetch::::new(3, 30, Duration::from_millis(2)).unwrap(); + for (id, delay_ms) in [(0, 30), (1, 1), (2, 5)] { + prefetch + .try_schedule(request(id, 10), &mut cache, move |_key, _max| async move { + tokio::time::sleep(Duration::from_millis(delay_ms)).await; + Ok(Some(vec![id as u8])) + }) + .unwrap(); + } + + let mut observed = Vec::new(); + while !prefetch.is_empty() { + let PrefetchNext::Ready(value) = + prefetch.next_interruptible(&mut cache, || false).await + else { + panic!("expected a prefetched object"); + }; + observed.push(value.request.context); + } + assert_eq!(observed, vec![0, 1, 2]); + }); + } + + #[test] + fn request_and_byte_limits_apply_before_fetch_creation() { + let mut cache = CompressedChunkCache::new(16, 4); + let mut prefetch = + OrderedPrefetch::::new(2, 10, Duration::from_millis(1)).unwrap(); + prefetch + .try_schedule(request(0, 6), &mut cache, |_key, _max| async { + Ok(Some(vec![0])) + }) + .unwrap(); + assert!(matches!( + prefetch.try_schedule(request(1, 5), &mut cache, |_key, _max| async { + Ok(Some(vec![1])) + }), + Err(ScheduleError::WindowFull(_)) + )); + assert!(matches!( + prefetch.try_schedule(request(2, 11), &mut cache, |_key, _max| async { + Ok(Some(vec![2])) + }), + Err(ScheduleError::RequestTooLarge { .. }) + )); + assert_eq!(prefetch.inflight_reads(), 1); + assert_eq!(prefetch.reserved_bytes(), 6); + } + + #[test] + fn cache_hits_preserve_order_without_invoking_fetch() { + let rt = runtime(); + rt.block_on(async { + let mut cache = CompressedChunkCache::new(16, 4); + let cached: Arc<[u8]> = Arc::from(vec![7_u8]); + cache.insert_present("chunk-1".to_string(), cached); + let fetch_calls = Rc::new(Cell::new(0)); + let mut prefetch = + OrderedPrefetch::::new(1, 8, Duration::from_millis(1)).unwrap(); + let calls = Rc::clone(&fetch_calls); + prefetch + .try_schedule(request(1, 8), &mut cache, move |_key, _max| { + calls.set(calls.get() + 1); + async { Ok(Some(vec![9])) } + }) + .unwrap(); + + let PrefetchNext::Ready(value) = + prefetch.next_interruptible(&mut cache, || false).await + else { + panic!("expected a cache hit"); + }; + assert_eq!(value.source, PrefetchSource::Cache); + assert_eq!(value.remote_bytes, 0); + assert_eq!(fetch_calls.get(), 0); + }); + } + + #[test] + fn cache_hits_still_obey_the_request_read_limit() { + let mut cache = CompressedChunkCache::new(16, 4); + let cached: Arc<[u8]> = Arc::from(vec![1_u8, 2, 3, 4]); + cache.insert_present("chunk-1".to_string(), cached); + let mut prefetch = + OrderedPrefetch::::new(1, 8, Duration::from_millis(1)).unwrap(); + + assert!(matches!( + prefetch.try_schedule(request(1, 3), &mut cache, |_key, _max| async { + Ok(Some(Vec::new())) + }), + Err(ScheduleError::CachedObjectTooLarge { + actual_bytes: 4, + .. + }) + )); + assert!(prefetch.is_empty()); + } + + struct ActiveGuard { + active: Rc>, + dropped: Rc>, + } + + impl ActiveGuard { + fn new(active: Rc>, dropped: Rc>) -> Self { + active.set(active.get() + 1); + Self { active, dropped } + } + } + + impl Drop for ActiveGuard { + fn drop(&mut self) { + self.active.set(self.active.get() - 1); + self.dropped.set(self.dropped.get() + 1); + } + } + + #[test] + fn cancellation_drops_every_future_before_returning_interrupted() { + let rt = runtime(); + rt.block_on(async { + let mut cache = CompressedChunkCache::new(16, 4); + let active = Rc::new(Cell::new(0)); + let dropped = Rc::new(Cell::new(0)); + let mut prefetch = + OrderedPrefetch::::new(2, 16, Duration::from_millis(1)).unwrap(); + for id in 0..2 { + let active = Rc::clone(&active); + let dropped = Rc::clone(&dropped); + prefetch + .try_schedule(request(id, 8), &mut cache, move |_key, _max| async move { + let _guard = ActiveGuard::new(active, dropped); + std::future::pending::>, Infallible>>().await + }) + .unwrap(); + } + + let checks = Cell::new(0); + let result = prefetch + .next_interruptible(&mut cache, || { + checks.set(checks.get() + 1); + checks.get() >= 2 + }) + .await; + + assert_eq!(result, PrefetchNext::Interrupted); + assert_eq!(active.get(), 0); + assert_eq!(dropped.get(), 2); + assert!(prefetch.is_empty()); + assert_eq!(prefetch.inflight_reads(), 0); + assert_eq!(prefetch.reserved_bytes(), 0); + }); + } + + #[test] + fn first_fetch_error_drops_every_later_future_and_resets_accounting() { + let rt = runtime(); + rt.block_on(async { + let mut cache = CompressedChunkCache::new(24, 4); + let active = Rc::new(Cell::new(0)); + let dropped = Rc::new(Cell::new(0)); + let mut prefetch = + OrderedPrefetch::::new(3, 24, Duration::from_millis(1)) + .unwrap(); + prefetch + .try_schedule(request(0, 8), &mut cache, |_key, _max| async { + Err("first fetch failed") + }) + .unwrap(); + for id in 1..=2 { + let guard = ActiveGuard::new(Rc::clone(&active), Rc::clone(&dropped)); + prefetch + .try_schedule(request(id, 8), &mut cache, move |_key, _max| async move { + let _guard = guard; + std::future::pending::>, &'static str>>().await + }) + .unwrap(); + } + + let result = prefetch.next_interruptible(&mut cache, || false).await; + assert!(matches!( + result, + PrefetchNext::FetchError { + error: "first fetch failed", + .. + } + )); + assert_eq!(active.get(), 0); + assert_eq!(dropped.get(), 2); + assert!(prefetch.is_empty()); + assert_eq!(prefetch.inflight_reads(), 0); + assert_eq!(prefetch.reserved_bytes(), 0); + }); + } + + #[test] + fn remote_missing_objects_are_cached_without_bytes() { + let rt = runtime(); + rt.block_on(async { + let mut cache = CompressedChunkCache::new(16, 4); + let mut prefetch = + OrderedPrefetch::::new(1, 8, Duration::from_millis(1)).unwrap(); + prefetch + .try_schedule(request(0, 8), &mut cache, |_key, _max| async { Ok(None) }) + .unwrap(); + let PrefetchNext::Ready(value) = + prefetch.next_interruptible(&mut cache, || false).await + else { + panic!("expected a missing result"); + }; + assert_eq!(value.object, CachedObject::Missing); + assert_eq!(value.source, PrefetchSource::Remote); + assert_eq!(cache.get("chunk-0"), Some(CachedObject::Missing)); + assert_eq!(cache.resident_bytes(), 0); + }); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/scan_plan.rs b/wrappers/src/fdw/zarr_fdw/scan_plan.rs new file mode 100644 index 000000000..037442c4f --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/scan_plan.rs @@ -0,0 +1,335 @@ +//! Pure executor-time planning for conservative coordinate and chunk pruning. +//! +//! Planning happens only after array metadata and required coordinate values +//! are loaded. A plan stores one inclusive chunk range per axis and never +//! materializes the Cartesian set of chunk coordinates. + +use super::chunk::{axis_chunk_ranges, index_bounds_from_value_range}; +use super::meta::ArrayMeta; +use super::selection::Selection; +use super::{ZarrFdwError, ZarrFdwResult}; + +pub(crate) type CoordinateRange = (Option, Option); + +/// Rank-sized, lazily executable scan plan. +/// +/// Exact SQL, temporal, selector, and spatial residual checks remain with +/// their existing execution layers. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ScanPlan { + selection: Selection, + axis_chunk_ranges: Vec<(usize, usize)>, + chunks_total: u64, + chunks_selected: u64, +} + +impl ScanPlan { + #[cfg(test)] + pub(crate) fn selection(&self) -> &Selection { + &self.selection + } + + pub(crate) fn axis_chunk_ranges(&self) -> &[(usize, usize)] { + &self.axis_chunk_ranges + } + + pub(crate) fn chunks_total(&self) -> u64 { + self.chunks_total + } + + pub(crate) fn chunks_selected(&self) -> u64 { + self.chunks_selected + } + + pub(crate) fn into_selection(self) -> Selection { + self.selection + } +} + +pub(crate) struct ScanPlanner<'a> { + meta: &'a ArrayMeta, +} + +impl<'a> ScanPlanner<'a> { + pub(crate) fn new(meta: &'a ArrayMeta) -> Self { + Self { meta } + } + + pub(crate) fn plan(&self, selection: Selection) -> ZarrFdwResult { + selection.validate(self.meta)?; + let axis_chunk_ranges = if selection.is_empty() { + Vec::new() + } else { + axis_chunk_ranges(self.meta, selection.axis_bounds())? + }; + let chunks_total = saturating_product(self.meta.chunks_per_axis()); + let chunks_selected = if selection.is_empty() { + 0 + } else { + saturating_range_product(&axis_chunk_ranges) + }; + + Ok(ScanPlan { + selection, + axis_chunk_ranges, + chunks_total, + chunks_selected, + }) + } + + /// Convert conservative coordinate-space ranges into index bounds, then + /// derive their lazy chunk plan. Unordered coordinates disable pruning for + /// that axis; exact consumers still apply their residual predicates. + #[cfg(test)] + pub(crate) fn plan_coordinate_ranges( + &self, + axis_names: &[String], + coordinate_values: &[Option>], + ranges: &[CoordinateRange], + ) -> ZarrFdwResult { + let selection = + self.selection_from_coordinate_ranges(axis_names, coordinate_values, ranges)?; + self.plan(selection) + } + + pub(crate) fn selection_from_coordinate_ranges( + &self, + axis_names: &[String], + coordinate_values: &[Option>], + ranges: &[CoordinateRange], + ) -> ZarrFdwResult { + let rank = self.meta.shape.len(); + if axis_names.len() != rank || coordinate_values.len() != rank || ranges.len() != rank { + return Err(ZarrFdwError::InvalidMetadata(format!( + "scan planning inputs do not match array rank {rank}" + ))); + } + + let mut bounds = Vec::with_capacity(rank); + let mut empty = false; + for (axis, &(lo, hi)) in ranges.iter().enumerate() { + if lo.is_none() && hi.is_none() { + bounds.push(None); + continue; + } + let coords = coordinate_values[axis].as_deref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate '{}' is required for predicate pruning but was not loaded", + axis_names[axis] + )) + })?; + if !coordinate_values_are_monotonic(coords) { + bounds.push(None); + continue; + } + let axis_bounds = index_bounds_from_value_range(coords, lo, hi); + if axis_bounds.is_none() { + empty = true; + } + bounds.push(axis_bounds); + } + + Ok(if empty { + Selection::empty(rank) + } else { + Selection::from_axis_bounds(bounds) + }) + } +} + +fn coordinate_values_are_monotonic(values: &[f64]) -> bool { + values.windows(2).all(|pair| pair[0] <= pair[1]) + || values.windows(2).all(|pair| pair[0] >= pair[1]) +} + +fn saturating_product(extents: impl IntoIterator) -> u64 { + extents + .into_iter() + .fold(1u64, |total, extent| total.saturating_mul(extent)) +} + +fn saturating_range_product(ranges: &[(usize, usize)]) -> u64 { + ranges.iter().fold(1u64, |total, &(start, end)| { + let extent = end + .checked_sub(start) + .and_then(|span| span.checked_add(1)) + .and_then(|count| u64::try_from(count).ok()) + .unwrap_or(u64::MAX); + total.saturating_mul(extent) + }) +} + +#[cfg(test)] +mod tests { + use serde_json::Value; + + use super::super::chunk::IndexBounds; + use super::super::codec::CodecPipeline; + use super::super::meta::ChunkKeyEncoding; + use super::super::sharding::StorageLayout; + use super::*; + + fn meta(shape: Vec, chunks: Vec) -> ArrayMeta { + ArrayMeta { + zarr_format: 2, + shape, + chunks, + dtype: " ScanPlan { + ScanPlanner::new(meta).plan(selection).unwrap() + } + + #[test] + fn plans_full_and_bounded_selections_without_chunk_enumeration() { + let meta = meta(vec![48, 100, 100], vec![4, 10, 10]); + let full = plan(&meta, Selection::full(3)); + assert_eq!(full.axis_chunk_ranges(), &[(0, 11), (0, 9), (0, 9)]); + assert_eq!(full.chunks_total(), 1_200); + assert_eq!(full.chunks_selected(), 1_200); + + let bounded = plan( + &meta, + Selection::from_axis_bounds(vec![ + Some(IndexBounds { start: 5, end: 11 }), + None, + Some(IndexBounds { start: 20, end: 39 }), + ]), + ); + assert_eq!(bounded.axis_chunk_ranges(), &[(1, 2), (0, 9), (2, 3)]); + assert_eq!(bounded.chunks_total(), 1_200); + assert_eq!(bounded.chunks_selected(), 40); + } + + #[test] + fn explicit_empty_selection_preserves_total_and_selects_zero_chunks() { + let meta = meta(vec![8, 8], vec![4, 4]); + let empty = plan(&meta, Selection::empty(2)); + let full = plan(&meta, Selection::full(2)); + + assert!(empty.selection().is_empty()); + assert!(empty.axis_chunk_ranges().is_empty()); + assert_eq!(empty.chunks_total(), 4); + assert_eq!(empty.chunks_selected(), 0); + assert_eq!(full.axis_chunk_ranges(), &[(0, 1), (0, 1)]); + assert_eq!(full.chunks_selected(), 4); + } + + #[test] + fn coordinate_ranges_plan_ascending_descending_and_no_overlap() { + let meta = meta(vec![5, 5], vec![2, 2]); + let planner = ScanPlanner::new(&meta); + let axes = vec!["ascending".to_string(), "descending".to_string()]; + let coords = vec![ + Some(vec![0.0, 10.0, 20.0, 30.0, 40.0]), + Some(vec![40.0, 30.0, 20.0, 10.0, 0.0]), + ]; + let plan = planner + .plan_coordinate_ranges( + &axes, + &coords, + &[(Some(10.0), Some(30.0)), (Some(10.0), Some(30.0))], + ) + .unwrap(); + + assert_eq!( + plan.selection().axis_bounds(), + &[ + Some(IndexBounds { start: 1, end: 3 }), + Some(IndexBounds { start: 1, end: 3 }) + ] + ); + assert_eq!(plan.axis_chunk_ranges(), &[(0, 1), (0, 1)]); + + let empty = planner + .plan_coordinate_ranges(&axes, &coords, &[(Some(100.0), None), (None, None)]) + .unwrap(); + assert!(empty.selection().is_empty()); + assert_eq!(empty.chunks_selected(), 0); + assert!(empty.axis_chunk_ranges().is_empty()); + } + + #[test] + fn unordered_coordinates_disable_only_that_axis_pruning() { + let meta = meta(vec![4, 4], vec![2, 2]); + let plan = ScanPlanner::new(&meta) + .plan_coordinate_ranges( + &["unordered".to_string(), "ordered".to_string()], + &[ + Some(vec![30.0, 10.0, 20.0, 0.0]), + Some(vec![0.0, 10.0, 20.0, 30.0]), + ], + &[(Some(10.0), Some(20.0)), (Some(10.0), Some(20.0))], + ) + .unwrap(); + + assert_eq!( + plan.selection().axis_bounds(), + &[None, Some(IndexBounds { start: 1, end: 2 })] + ); + assert_eq!(plan.axis_chunk_ranges(), &[(0, 1), (0, 1)]); + } + + #[test] + fn coordinate_range_inputs_must_be_rank_aligned_and_loaded() { + let meta = meta(vec![4], vec![2]); + let planner = ScanPlanner::new(&meta); + + let rank_error = planner.plan_coordinate_ranges(&[], &[], &[]).unwrap_err(); + assert!(rank_error.to_string().contains("do not match array rank 1")); + + let missing_error = planner + .plan_coordinate_ranges(&["x".to_string()], &[None], &[(Some(1.0), Some(2.0))]) + .unwrap_err(); + assert!( + missing_error + .to_string() + .contains("coordinate 'x' is required for predicate pruning but was not loaded") + ); + } + + #[test] + fn rank_64_plan_stays_rank_sized_and_saturates_counts() { + let meta = meta(vec![2; 64], vec![1; 64]); + let plan = plan(&meta, Selection::full(64)); + + assert_eq!(plan.axis_chunk_ranges().len(), 64); + assert!( + plan.axis_chunk_ranges() + .iter() + .all(|range| *range == (0, 1)) + ); + assert_eq!(plan.chunks_total(), u64::MAX); + assert_eq!(plan.chunks_selected(), u64::MAX); + } + + #[test] + fn selection_validation_errors_propagate_before_chunk_math() { + let meta = meta(vec![8, 8], vec![4, 4]); + let rank_error = ScanPlanner::new(&meta) + .plan(Selection::full(3)) + .unwrap_err(); + assert!(rank_error.to_string().contains( + "zarr array metadata missing or invalid: selection rank 3 does not match array rank 2" + )); + + for bounds in [ + IndexBounds { start: 5, end: 4 }, + IndexBounds { start: 0, end: 8 }, + ] { + let error = ScanPlanner::new(&meta) + .plan(Selection::from_axis_bounds(vec![Some(bounds), None])) + .unwrap_err(); + assert!(error.to_string().contains("selection index bounds")); + } + } +} diff --git a/wrappers/src/fdw/zarr_fdw/scientific.rs b/wrappers/src/fdw/zarr_fdw/scientific.rs new file mode 100644 index 000000000..419c05dac --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/scientific.rs @@ -0,0 +1,207 @@ +//! Narrow CF-style scientific value semantics layered above primitive decode. +//! +//! This module deliberately does not know how Zarr metadata or chunks are +//! stored. A future v3 metadata adapter can supply the same attribute map and +//! reuse this decoder unchanged. + +pub(crate) mod time; + +use serde_json::{Map, Value}; + +use super::decode::{DType, fill_value_bytes, value_bytes_to_f64}; +use super::{ZarrFdwError, ZarrFdwResult}; + +const ATTR_FILL_VALUE: &str = "_FillValue"; +const ATTR_MISSING_VALUE: &str = "missing_value"; +const ATTR_VALID_RANGE: &str = "valid_range"; +const ATTR_VALID_MIN: &str = "valid_min"; +const ATTR_VALID_MAX: &str = "valid_max"; +const ATTR_SCALE_FACTOR: &str = "scale_factor"; +const ATTR_ADD_OFFSET: &str = "add_offset"; + +/// Parsed CF-style masking and packing rules for one value variable. +#[derive(Debug, Clone)] +pub(crate) struct ScientificValueDecoder { + dtype: DType, + missing_values: Vec>, + mask_nan: bool, + valid_min: Option, + valid_max: Option, + scale_factor: f64, + add_offset: f64, +} + +impl ScientificValueDecoder { + pub(crate) fn from_attributes( + dtype: DType, + attributes: &Map, + ) -> ZarrFdwResult { + let mut missing_values = Vec::new(); + let mut mask_nan = false; + if let Some(value) = attributes.get(ATTR_FILL_VALUE) { + let (bytes, is_nan) = missing_attribute_bytes(dtype, ATTR_FILL_VALUE, value)?; + missing_values.push(bytes); + mask_nan |= is_nan; + } + if let Some(value) = attributes.get(ATTR_MISSING_VALUE) { + match value { + Value::Array(values) => { + if values.is_empty() { + return Err(attribute_error( + ATTR_MISSING_VALUE, + "must be a numeric scalar or non-empty numeric array", + )); + } + for value in values { + let (bytes, is_nan) = + missing_attribute_bytes(dtype, ATTR_MISSING_VALUE, value)?; + missing_values.push(bytes); + mask_nan |= is_nan; + } + } + _ => { + let (bytes, is_nan) = + missing_attribute_bytes(dtype, ATTR_MISSING_VALUE, value)?; + missing_values.push(bytes); + mask_nan |= is_nan; + } + } + } + missing_values.sort_unstable(); + missing_values.dedup(); + + if attributes.contains_key(ATTR_VALID_RANGE) + && (attributes.contains_key(ATTR_VALID_MIN) || attributes.contains_key(ATTR_VALID_MAX)) + { + return Err(attribute_error( + ATTR_VALID_RANGE, + "cannot be combined with 'valid_min' or 'valid_max'", + )); + } + + let (valid_min, valid_max) = match attributes.get(ATTR_VALID_RANGE) { + Some(Value::Array(values)) if values.len() == 2 => ( + Some(raw_attribute_f64(dtype, ATTR_VALID_RANGE, &values[0])?), + Some(raw_attribute_f64(dtype, ATTR_VALID_RANGE, &values[1])?), + ), + Some(_) => { + return Err(attribute_error( + ATTR_VALID_RANGE, + "must be a two-element numeric array", + )); + } + None => ( + attributes + .get(ATTR_VALID_MIN) + .map(|value| raw_attribute_f64(dtype, ATTR_VALID_MIN, value)) + .transpose()?, + attributes + .get(ATTR_VALID_MAX) + .map(|value| raw_attribute_f64(dtype, ATTR_VALID_MAX, value)) + .transpose()?, + ), + }; + if valid_min + .zip(valid_max) + .is_some_and(|(minimum, maximum)| minimum > maximum) + { + return Err(attribute_error( + ATTR_VALID_RANGE, + "minimum must not exceed maximum", + )); + } + + let scale_factor = finite_attribute(attributes, ATTR_SCALE_FACTOR)?.unwrap_or(1.0); + let add_offset = finite_attribute(attributes, ATTR_ADD_OFFSET)?.unwrap_or(0.0); + + Ok(Self { + dtype, + missing_values, + mask_nan, + valid_min, + valid_max, + scale_factor, + add_offset, + }) + } + + /// Return a decoded physical value, or `None` for semantic missing data. + /// Masking and valid-range checks intentionally happen in the raw packed + /// domain before scale/offset are applied. + pub(crate) fn decode(&self, raw_bytes: &[u8]) -> ZarrFdwResult> { + if self + .missing_values + .iter() + .any(|missing| missing.as_slice() == raw_bytes) + { + return Ok(None); + } + let raw = value_bytes_to_f64(self.dtype, raw_bytes)?; + if self.mask_nan && raw.is_nan() { + return Ok(None); + } + if self.valid_min.is_some_and(|minimum| raw < minimum) + || self.valid_max.is_some_and(|maximum| raw > maximum) + { + return Ok(None); + } + Ok(Some(raw.mul_add(self.scale_factor, self.add_offset))) + } +} + +fn missing_attribute_bytes( + dtype: DType, + name: &str, + value: &Value, +) -> ZarrFdwResult<(Vec, bool)> { + let mask_nan = matches!(value, Value::String(value) if value == "NaN"); + let valid_special_float = matches!( + value, + Value::String(value) if matches!(value.as_str(), "NaN" | "Infinity" | "-Infinity") + ) && matches!(dtype, DType::F32 | DType::F64); + if !value.is_number() && !valid_special_float { + return Err(attribute_error( + name, + "must be numeric, or a supported non-finite string for a floating-point array", + )); + } + let bytes = fill_value_bytes(dtype, value) + .map_err(|error| attribute_error(name, error.to_string()))? + .ok_or_else(|| attribute_error(name, "must not be null"))?; + Ok((bytes, mask_nan)) +} + +fn raw_attribute_bytes(dtype: DType, name: &str, value: &Value) -> ZarrFdwResult> { + if !value.is_number() { + return Err(attribute_error(name, "must be a numeric value")); + } + fill_value_bytes(dtype, value) + .map_err(|error| attribute_error(name, error.to_string()))? + .ok_or_else(|| attribute_error(name, "must not be null")) +} + +fn raw_attribute_f64(dtype: DType, name: &str, value: &Value) -> ZarrFdwResult { + let bytes = raw_attribute_bytes(dtype, name, value)?; + value_bytes_to_f64(dtype, &bytes) +} + +fn finite_attribute(attributes: &Map, name: &str) -> ZarrFdwResult> { + let Some(value) = attributes.get(name) else { + return Ok(None); + }; + let value = value + .as_f64() + .filter(|value| value.is_finite()) + .ok_or_else(|| attribute_error(name, "must be a finite numeric value"))?; + Ok(Some(value)) +} + +fn attribute_error(name: &str, message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!( + "CF attribute '{name}' is invalid: {}", + message.into() + )) +} + +#[cfg(test)] +mod tests; diff --git a/wrappers/src/fdw/zarr_fdw/scientific/tests.rs b/wrappers/src/fdw/zarr_fdw/scientific/tests.rs new file mode 100644 index 000000000..989af2c67 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/scientific/tests.rs @@ -0,0 +1,103 @@ +use super::*; +use serde_json::json; + +fn attributes(value: Value) -> Map { + value.as_object().cloned().unwrap() +} + +#[test] +fn masks_raw_values_before_applying_scale_and_offset() { + let decoder = ScientificValueDecoder::from_attributes( + DType::I16, + &attributes(json!({ + "_FillValue": -32768, + "missing_value": [-9999, -8888], + "valid_range": [0, 500], + "scale_factor": 0.1, + "add_offset": 273.15 + })), + ) + .unwrap(); + + assert_eq!(decoder.decode(&(-32768_i16).to_le_bytes()).unwrap(), None); + assert_eq!(decoder.decode(&(-9999_i16).to_le_bytes()).unwrap(), None); + assert_eq!(decoder.decode(&(501_i16).to_le_bytes()).unwrap(), None); + let decoded = decoder.decode(&(500_i16).to_le_bytes()).unwrap().unwrap(); + assert!((decoded - 323.15).abs() < 1e-10); +} + +#[test] +fn accepts_scalar_missing_and_separate_valid_bounds() { + let decoder = ScientificValueDecoder::from_attributes( + DType::F32, + &attributes(json!({ + "missing_value": -9999.0, + "valid_min": 0.0, + "valid_max": 10.0 + })), + ) + .unwrap(); + + assert_eq!(decoder.decode(&(-9999.0_f32).to_le_bytes()).unwrap(), None); + assert_eq!(decoder.decode(&(-1.0_f32).to_le_bytes()).unwrap(), None); + assert_eq!( + decoder.decode(&(10.0_f32).to_le_bytes()).unwrap(), + Some(10.0) + ); +} + +#[test] +fn rejects_conflicting_or_malformed_attributes() { + for attrs in [ + json!({"valid_range": [0, 1], "valid_min": 0}), + json!({"valid_range": [1, 0]}), + json!({"valid_range": [0]}), + json!({"missing_value": []}), + json!({"scale_factor": "0.1"}), + json!({"_FillValue": null}), + ] { + assert!(ScientificValueDecoder::from_attributes(DType::I16, &attributes(attrs)).is_err()); + } +} + +#[test] +fn identity_semantics_still_promote_to_f64() { + let decoder = ScientificValueDecoder::from_attributes(DType::I32, &Map::new()).unwrap(); + assert_eq!(decoder.decode(&(42_i32).to_le_bytes()).unwrap(), Some(42.0)); +} + +#[test] +fn masks_declared_non_finite_sentinels_without_requiring_one_nan_payload() { + let decoder = ScientificValueDecoder::from_attributes( + DType::F32, + &attributes(json!({ + "_FillValue": "NaN", + "missing_value": ["Infinity", "-Infinity"] + })), + ) + .unwrap(); + + let alternate_nan = f32::from_bits(0x7fc0_0001); + assert_eq!(decoder.decode(&alternate_nan.to_le_bytes()).unwrap(), None); + assert_eq!(decoder.decode(&f32::INFINITY.to_le_bytes()).unwrap(), None); + assert_eq!( + decoder.decode(&f32::NEG_INFINITY.to_le_bytes()).unwrap(), + None + ); +} + +#[test] +fn preserves_undeclared_non_finite_float_values() { + let decoder = ScientificValueDecoder::from_attributes(DType::F64, &Map::new()).unwrap(); + assert!( + decoder + .decode(&f64::NAN.to_le_bytes()) + .unwrap() + .unwrap() + .is_nan() + ); + assert_eq!( + decoder.decode(&f64::INFINITY.to_le_bytes()).unwrap(), + Some(f64::INFINITY) + ); +} diff --git a/wrappers/src/fdw/zarr_fdw/scientific/time.rs b/wrappers/src/fdw/zarr_fdw/scientific/time.rs new file mode 100644 index 000000000..ea1bf2045 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/scientific/time.rs @@ -0,0 +1,216 @@ +use chrono::{DateTime, NaiveDate, NaiveDateTime}; +use serde_json::{Map, Value}; + +use super::super::{ZarrFdwError, ZarrFdwResult}; + +const ATTR_UNITS: &str = "units"; +const ATTR_CALENDAR: &str = "calendar"; +const SUPPORTED_CALENDAR: &str = "proleptic_gregorian"; + +// PostgreSQL epoch (2000-01-01 00:00:00 UTC) in microseconds since Unix epoch. +const PG_EPOCH_MICROS: i64 = 946_684_800_000_000; +const PG_EPOCH_SECONDS: i64 = 946_684_800; +const I64_MIN_AS_F64: f64 = -9_223_372_036_854_775_808.0; +const I64_MAX_EXCLUSIVE_AS_F64: f64 = 9_223_372_036_854_775_808.0; + +/// Unit of raw CF time coordinate values. +#[derive(Debug, Clone, Copy, PartialEq)] +enum TimeUnit { + Seconds, + Milliseconds, + Microseconds, + Nanoseconds, + Minutes, + Hours, + Days, +} + +impl TimeUnit { + fn parse_option(value: &str) -> ZarrFdwResult { + match value { + "seconds" => Ok(Self::Seconds), + "milliseconds" => Ok(Self::Milliseconds), + "microseconds" => Ok(Self::Microseconds), + "nanoseconds" => Ok(Self::Nanoseconds), + "minutes" => Ok(Self::Minutes), + "hours" => Ok(Self::Hours), + "days" => Ok(Self::Days), + _ => Err(ZarrFdwError::InvalidOptionValue { + option: value.to_string(), + message: supported_units_message(), + }), + } + } + + fn parse_cf(value: &str) -> ZarrFdwResult { + Self::parse_option(value).map_err(|_| { + time_metadata_error(format!( + "unit '{value}' is unsupported; {}", + supported_units_message() + )) + }) + } + + fn microseconds_factor(self) -> f64 { + match self { + Self::Seconds => 1_000_000.0, + Self::Milliseconds => 1_000.0, + Self::Microseconds => 1.0, + Self::Nanoseconds => 1e-3, + Self::Minutes => 60_000_000.0, + Self::Hours => 3_600_000_000.0, + Self::Days => 86_400_000_000.0, + } + } +} + +/// Describes how raw `time` coordinate values map to PostgreSQL instants. +#[derive(Debug, Clone, Copy)] +pub(crate) struct TimeSpec { + unit: TimeUnit, + origin_unix_micros: i64, +} + +impl TimeSpec { + pub(crate) fn default() -> Self { + Self { + unit: TimeUnit::Seconds, + origin_unix_micros: 0, + } + } + + pub(crate) fn from_legacy_options( + unit: Option<&str>, + origin: Option<&str>, + ) -> ZarrFdwResult { + let unit = match unit { + Some(value) => TimeUnit::parse_option(value)?, + None => TimeUnit::Seconds, + }; + let origin_unix_micros = match origin { + Some("unix") => 0, + Some("postgres") => PG_EPOCH_SECONDS + .checked_mul(1_000_000) + .ok_or_else(|| time_metadata_error("PostgreSQL epoch overflowed"))?, + Some(other) => { + return Err(ZarrFdwError::InvalidOptionValue { + option: other.to_string(), + message: "must be 'unix' or 'postgres'".to_string(), + }); + } + None => 0, + }; + Ok(Self { + unit, + origin_unix_micros, + }) + } + + pub(crate) fn from_cf_attributes(attributes: &Map) -> ZarrFdwResult { + let units = required_string_attribute(attributes, ATTR_UNITS)?; + let calendar = required_string_attribute(attributes, ATTR_CALENDAR)?; + if calendar != SUPPORTED_CALENDAR { + return Err(time_metadata_error(format!( + "calendar must be '{SUPPORTED_CALENDAR}', got '{calendar}'" + ))); + } + let (unit, origin) = units.split_once(" since ").ok_or_else(|| { + time_metadata_error("units must have the form ' since '") + })?; + if unit.trim() != unit || origin.trim() != origin || unit.is_empty() || origin.is_empty() { + return Err(time_metadata_error( + "units must have the form ' since ' without empty fields", + )); + } + Ok(Self { + unit: TimeUnit::parse_cf(unit)?, + origin_unix_micros: parse_origin_micros(origin)?, + }) + } + + /// Convert a raw coordinate value into PostgreSQL-epoch microseconds. + pub(crate) fn raw_to_pg_micros(&self, raw: f64) -> ZarrFdwResult { + if !raw.is_finite() { + return Err(time_metadata_error(format!( + "raw time coordinate value must be finite, got {raw}" + ))); + } + let unix_micros = raw + .mul_add( + self.unit.microseconds_factor(), + self.origin_unix_micros as f64, + ) + .round(); + let unix_micros = checked_f64_to_i64_micros(unix_micros)?; + unix_micros + .checked_sub(PG_EPOCH_MICROS) + .ok_or_else(|| ZarrFdwError::TimeOutOfRange(raw)) + } + + /// Return the conservative raw-coordinate interval that can round to one + /// PostgreSQL microsecond. Both endpoints are included intentionally: + /// PostgreSQL rechecks the original qual, while inclusive bounds ensure + /// pruning never drops values at `f64::round` tie boundaries. + pub(crate) fn pg_micros_to_raw_bounds(&self, pg_micros: i64) -> Option<(f64, f64)> { + let unix_micros = pg_micros.checked_add(PG_EPOCH_MICROS)?; + let delta = unix_micros.checked_sub(self.origin_unix_micros)? as f64; + let factor = self.unit.microseconds_factor(); + Some(((delta - 0.5) / factor, (delta + 0.5) / factor)) + } +} + +fn required_string_attribute<'a>( + attributes: &'a Map, + name: &str, +) -> ZarrFdwResult<&'a str> { + attributes + .get(name) + .and_then(Value::as_str) + .ok_or_else(|| time_metadata_error(format!("attribute '{name}' must be a string"))) +} + +fn parse_origin_micros(origin: &str) -> ZarrFdwResult { + if let Ok(value) = DateTime::parse_from_rfc3339(origin) { + return Ok(value.timestamp_micros()); + } + for format in ["%Y-%m-%d %H:%M:%S%.f", "%Y-%m-%dT%H:%M:%S%.f"] { + if let Ok(value) = NaiveDateTime::parse_from_str(origin, format) { + return Ok(value.and_utc().timestamp_micros()); + } + } + if let Ok(value) = NaiveDate::parse_from_str(origin, "%Y-%m-%d") { + let Some(value) = value.and_hms_opt(0, 0, 0) else { + return Err(time_metadata_error(format!( + "origin '{origin}' is outside the supported timestamp range" + ))); + }; + return Ok(value.and_utc().timestamp_micros()); + } + Err(time_metadata_error(format!( + "origin '{origin}' must be a Gregorian date, date-time, or RFC 3339 date-time" + ))) +} + +fn checked_f64_to_i64_micros(value: f64) -> ZarrFdwResult { + if !value.is_finite() + || !(I64_MIN_AS_F64..I64_MAX_EXCLUSIVE_AS_F64).contains(&value) + || value.fract() != 0.0 + { + return Err(time_metadata_error( + "time conversion produced a microsecond value outside the supported range", + )); + } + Ok(value as i64) +} + +fn supported_units_message() -> String { + "must be one of: seconds, milliseconds, microseconds, nanoseconds, minutes, hours, days" + .to_string() +} + +fn time_metadata_error(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!("CF time metadata is invalid: {}", message.into())) +} + +#[cfg(test)] +mod tests; diff --git a/wrappers/src/fdw/zarr_fdw/scientific/time/tests.rs b/wrappers/src/fdw/zarr_fdw/scientific/time/tests.rs new file mode 100644 index 000000000..a69309f8d --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/scientific/time/tests.rs @@ -0,0 +1,112 @@ +use super::*; +use serde_json::json; + +fn attrs(value: Value) -> Map { + value.as_object().cloned().unwrap() +} + +#[test] +fn parses_supported_units_and_uses_checked_pg_microseconds() { + for (unit, raw, expected_seconds) in [ + ("seconds", 2.0, 2_i64), + ("milliseconds", 2_000.0, 2), + ("microseconds", 2_000_000.0, 2), + ("nanoseconds", 2_000_000_000.0, 2), + ("minutes", 2.0, 120), + ("hours", 2.0, 7_200), + ("days", 2.0, 172_800), + ] { + let spec = TimeSpec::from_cf_attributes(&attrs(json!({ + "units": format!("{unit} since 1970-01-01"), + "calendar": "proleptic_gregorian" + }))) + .unwrap(); + + assert_eq!( + spec.raw_to_pg_micros(raw).unwrap(), + expected_seconds * 1_000_000 - PG_EPOCH_MICROS + ); + } +} + +#[test] +fn accepts_date_datetime_and_rfc3339_origins_as_utc_instants() { + let origins = [ + "2000-01-01", + "2000-01-01 00:00:00", + "2000-01-01T00:00:00", + "2000-01-01T01:00:00+01:00", + ]; + + for origin in origins { + let spec = TimeSpec::from_cf_attributes(&attrs(json!({ + "units": format!("seconds since {origin}"), + "calendar": "proleptic_gregorian" + }))) + .unwrap(); + assert_eq!(spec.raw_to_pg_micros(0.0).unwrap(), 0); + } +} + +#[test] +fn rejects_missing_malformed_or_unsupported_cf_time_metadata() { + for value in [ + json!({"calendar": "proleptic_gregorian"}), + json!({"units": "seconds since 1970-01-01"}), + json!({ + "units": "months since 1970-01-01", + "calendar": "proleptic_gregorian" + }), + json!({"units": "seconds after 1970-01-01", "calendar": "proleptic_gregorian"}), + json!({"units": "seconds since 1970-01-01", "calendar": "gregorian"}), + json!({"units": 1, "calendar": "proleptic_gregorian"}), + json!({ + "units": "seconds since not-a-date", + "calendar": "proleptic_gregorian" + }), + ] { + assert!(TimeSpec::from_cf_attributes(&attrs(value)).is_err()); + } +} + +#[test] +fn rejects_non_finite_and_overflowing_time_conversions() { + let spec = TimeSpec::from_cf_attributes(&attrs(json!({ + "units": "days since 1970-01-01", + "calendar": "proleptic_gregorian" + }))) + .unwrap(); + + assert!(spec.raw_to_pg_micros(f64::NAN).is_err()); + assert!(spec.raw_to_pg_micros(f64::INFINITY).is_err()); + assert!(spec.raw_to_pg_micros(f64::MAX).is_err()); +} + +#[test] +fn converts_pg_predicate_micros_back_to_raw_with_same_spec() { + let spec = TimeSpec::from_cf_attributes(&attrs(json!({ + "units": "hours since 2000-01-01 00:00:00", + "calendar": "proleptic_gregorian" + }))) + .unwrap(); + + assert_eq!(spec.raw_to_pg_micros(6.0).unwrap(), 21_600_000_000); + let (lo, hi) = spec.pg_micros_to_raw_bounds(21_600_000_000).unwrap(); + assert!(lo < 6.0 && hi > 6.0); +} + +#[test] +fn inverse_bounds_cover_nanoseconds_that_round_to_the_same_pg_microsecond() { + let spec = TimeSpec::from_cf_attributes(&attrs(json!({ + "units": "nanoseconds since 1970-01-01", + "calendar": "proleptic_gregorian" + }))) + .unwrap(); + + let pg_micros = 1 - PG_EPOCH_MICROS; + assert_eq!(spec.raw_to_pg_micros(501.0).unwrap(), pg_micros); + assert_eq!(spec.raw_to_pg_micros(1_499.0).unwrap(), pg_micros); + let (lo, hi) = spec.pg_micros_to_raw_bounds(pg_micros).unwrap(); + assert!(lo <= 501.0); + assert!(hi >= 1_499.0); +} diff --git a/wrappers/src/fdw/zarr_fdw/selection.rs b/wrappers/src/fdw/zarr_fdw/selection.rs new file mode 100644 index 000000000..79f119f45 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/selection.rs @@ -0,0 +1,226 @@ +//! Pure candidate-cell selection for Zarr scans. +//! +//! A selection is a conservative, rank-aligned rectangular window. Exact SQL, +//! temporal, and spatial predicates remain the responsibility of their current +//! execution layers. + +use super::chunk::IndexBounds; +use super::meta::ArrayMeta; +use super::{ZarrFdwError, ZarrFdwResult}; + +/// Conservative array-index bounds selected for one scan. +/// +/// `None` leaves an axis unconstrained. Since that cannot also represent a +/// selection with no cells, emptiness is tracked explicitly for the complete +/// Cartesian product. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) struct Selection { + axis_bounds: Vec>, + empty: bool, +} + +impl Selection { + #[cfg(test)] + pub(crate) fn full(rank: usize) -> Self { + Self { + axis_bounds: vec![None; rank], + empty: false, + } + } + + pub(crate) fn empty(rank: usize) -> Self { + Self { + axis_bounds: vec![None; rank], + empty: true, + } + } + + pub(crate) fn from_axis_bounds(axis_bounds: Vec>) -> Self { + Self { + axis_bounds, + empty: false, + } + } + + pub(crate) fn axis_bounds(&self) -> &[Option] { + &self.axis_bounds + } + + pub(crate) fn rank(&self) -> usize { + self.axis_bounds.len() + } + + pub(crate) fn is_empty(&self) -> bool { + self.empty + } + + pub(crate) fn intersect(self, other: Self) -> Self { + debug_assert_eq!(self.rank(), other.rank()); + let rank = self.rank(); + if self.empty || other.empty || other.rank() != rank { + return Self::empty(rank); + } + + let mut bounds = Vec::with_capacity(rank); + for axis in 0..rank { + let bound = match (self.axis_bounds[axis], other.axis_bounds[axis]) { + (Some(left), Some(right)) => { + let start = left.start.max(right.start); + let end = left.end.min(right.end); + if start > end { + return Self::empty(rank); + } + Some(IndexBounds { start, end }) + } + (Some(bound), None) | (None, Some(bound)) => Some(bound), + (None, None) => None, + }; + bounds.push(bound); + } + Self::from_axis_bounds(bounds) + } + + pub(crate) fn validate(&self, meta: &ArrayMeta) -> ZarrFdwResult<()> { + if self.rank() != meta.shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "selection rank {} does not match array rank {}", + self.rank(), + meta.shape.len() + ))); + } + + for (axis, bounds) in self.axis_bounds.iter().enumerate() { + let Some(bounds) = bounds else { + continue; + }; + let length = meta.shape_extent(axis)?; + if bounds.start > bounds.end || bounds.end >= length { + return Err(ZarrFdwError::InvalidMetadata(format!( + "selection index bounds {}..={} are invalid for dimension {axis} length {length}", + bounds.start, bounds.end + ))); + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use serde_json::Value; + + use super::super::codec::CodecPipeline; + use super::super::meta::ChunkKeyEncoding; + use super::super::sharding::StorageLayout; + use super::*; + + fn meta(shape: Vec, chunks: Vec) -> ArrayMeta { + ArrayMeta { + zarr_format: 2, + shape, + chunks, + dtype: "), + IndexRange { start: usize, stop: usize }, + Value(f64), + Values(Vec), + ValueRange { min: f64, max: f64 }, +} + +impl DimensionSelector { + fn matches_index(&self, index: usize, coordinates: Option<&[f64]>) -> ZarrFdwResult { + match self { + Self::Index(selected) => Ok(index == *selected), + Self::Indices(selected) => Ok(selected.binary_search(&index).is_ok()), + Self::IndexRange { start, stop } => Ok(*start <= index && index < *stop), + Self::Value(selected) => Ok(coordinate_at(coordinates, index)? == *selected), + Self::Values(selected) => { + let coordinate = coordinate_at(coordinates, index)?; + if !coordinate.is_finite() { + return Ok(false); + } + Ok(selected + .binary_search_by(|candidate| { + candidate + .partial_cmp(&coordinate) + .expect("selector and coordinate values are finite") + }) + .is_ok()) + } + Self::ValueRange { min, max } => { + let coordinate = coordinate_at(coordinates, index)?; + Ok(*min <= coordinate && coordinate <= *max) + } + } + } + + fn conservative_bounds( + &self, + length: usize, + coordinates: Option<&[f64]>, + poll_interrupt: &mut impl FnMut() -> ZarrFdwResult<()>, + ) -> ZarrFdwResult> { + match self { + Self::Index(index) => Ok(Some(IndexBounds { + start: *index, + end: *index, + })), + Self::Indices(indices) => Ok(Some(IndexBounds { + start: *indices + .first() + .expect("nonempty list validated while parsing"), + end: *indices + .last() + .expect("nonempty list validated while parsing"), + })), + Self::IndexRange { start, stop } => Ok(Some(IndexBounds { + start: *start, + end: *stop - 1, + })), + Self::Value(_) | Self::Values(_) | Self::ValueRange { .. } => { + let coordinates = coordinates.ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "coordinate required by a value selector was not loaded".to_string(), + ) + })?; + if coordinates.len() != length { + return Err(ZarrFdwError::InvalidMetadata(format!( + "coordinate length {} does not match selected dimension length {length}", + coordinates.len() + ))); + } + let mut first = None; + let mut last = None; + for index in 0..length { + if index % 1_024 == 0 { + poll_interrupt()?; + } + if self.matches_index(index, Some(coordinates))? { + first.get_or_insert(index); + last = Some(index); + } + } + Ok(first.map(|start| IndexBounds { + start, + end: last.expect("a first selector match also sets the last match"), + })) + } + } + } +} + +fn coordinate_at(coordinates: Option<&[f64]>, index: usize) -> ZarrFdwResult { + coordinates + .and_then(|values| values.get(index)) + .copied() + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate required by a value selector has no value at index {index}" + )) + }) +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub(crate) struct DimensionSelectors { + entries: Vec<(String, DimensionSelector)>, +} + +impl DimensionSelectors { + pub(crate) fn parse(raw: Option<&str>) -> ZarrFdwResult { + let Some(raw) = raw else { + return Ok(Self::default()); + }; + if raw.len() > MAX_SELECTOR_DOCUMENT_BYTES { + return Err(invalid_selector_option(format!( + "JSON document has {} bytes, exceeding the {MAX_SELECTOR_DOCUMENT_BYTES}-byte limit", + raw.len() + ))); + } + + let raw_selectors = serde_json::from_str::(raw) + .map_err(|error| invalid_selector_option(error.to_string()))?; + let mut entries = Vec::new(); + entries + .try_reserve_exact(raw_selectors.0.len()) + .map_err(|_| invalid_selector_option("could not allocate dimension selectors"))?; + for (dimension, raw_selector) in raw_selectors.0 { + if raw_selector.0.len() != 1 { + return Err(invalid_selector_option(format!( + "selector for dimension '{dimension}' must contain exactly one supported selector form" + ))); + } + let (kind, value) = raw_selector.0.into_iter().next().expect("length checked"); + let selector = match kind.as_str() { + "index" => DimensionSelector::Index(parse_index(&dimension, value.into_json()?)?), + "value" => DimensionSelector::Value(parse_value(&dimension, value.into_json()?)?), + "indices" => { + DimensionSelector::Indices(parse_indices(&dimension, value.into_list()?)?) + } + "values" => { + DimensionSelector::Values(parse_values(&dimension, value.into_list()?)?) + } + "index_range" => { + let range = value.into_object()?; + let (start, stop) = parse_index_range(&dimension, range)?; + DimensionSelector::IndexRange { start, stop } + } + "value_range" => { + let range = value.into_object()?; + let (min, max) = parse_value_range(&dimension, range)?; + DimensionSelector::ValueRange { min, max } + } + _ => { + return Err(invalid_selector_option(format!( + "selector for dimension '{dimension}' must contain exactly one supported selector form" + ))); + } + }; + entries.push((dimension, selector)); + } + Ok(Self { entries }) + } + + pub(crate) fn bind( + &self, + axis_names: &[String], + shape: &[u64], + ) -> ZarrFdwResult { + if shape.len() != axis_names.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "dimension selector binding rank {} does not match array rank {}", + axis_names.len(), + shape.len() + ))); + } + let mut by_axis = vec![None; axis_names.len()]; + for (dimension, selector) in &self.entries { + let axis = axis_names + .iter() + .position(|axis_name| axis_name == dimension) + .ok_or_else(|| { + invalid_selector_option(format!( + "dimension selector references unknown dimension '{dimension}'" + )) + })?; + let length = usize::try_from(shape[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "dimension '{dimension}' length exceeds this platform's index capacity" + )) + })?; + match selector { + DimensionSelector::Index(index) if *index >= length => { + return Err(invalid_selector_option(format!( + "dimension selector index {index} is outside dimension '{dimension}' length {length}" + ))); + } + DimensionSelector::Indices(indices) => { + if let Some(index) = indices.iter().find(|&&index| index >= length) { + return Err(invalid_selector_option(format!( + "dimension selector index {index} is outside dimension '{dimension}' length {length}" + ))); + } + } + DimensionSelector::IndexRange { stop, .. } if *stop > length => { + return Err(invalid_selector_option(format!( + "dimension selector index range stop {stop} exceeds dimension '{dimension}' length {length}" + ))); + } + _ => {} + } + by_axis[axis] = Some(selector.clone()); + } + Ok(BoundDimensionSelectors { + axis_names: axis_names.to_vec(), + by_axis, + }) + } +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub(crate) struct BoundDimensionSelectors { + axis_names: Vec, + by_axis: Vec>, +} + +impl BoundDimensionSelectors { + pub(crate) fn is_empty(&self) -> bool { + self.by_axis.iter().all(Option::is_none) + } + + pub(crate) fn requires_coordinate(&self, axis: usize) -> bool { + matches!( + self.by_axis.get(axis), + Some(Some( + DimensionSelector::Value(_) + | DimensionSelector::Values(_) + | DimensionSelector::ValueRange { .. } + )) + ) + } + + pub(crate) fn selects_axis(&self, axis: usize) -> bool { + self.by_axis.get(axis).is_some_and(Option::is_some) + } + + pub(crate) fn resolve( + &self, + shape: &[u64], + coordinate_values: &[Option>], + mut poll_interrupt: impl FnMut() -> ZarrFdwResult<()>, + ) -> ZarrFdwResult { + let rank = self.by_axis.len(); + if shape.len() != rank || coordinate_values.len() != rank { + return Err(ZarrFdwError::InvalidMetadata(format!( + "dimension selector inputs do not match array rank {rank}" + ))); + } + + let mut bounds = vec![None; rank]; + for (axis, selector) in self.by_axis.iter().enumerate() { + let Some(selector) = selector else { + continue; + }; + let axis_name = &self.axis_names[axis]; + let length = usize::try_from(shape[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "dimension {axis} length exceeds this platform's index capacity" + )) + })?; + let coordinates = if self.requires_coordinate(axis) { + Some(coordinate_values[axis].as_deref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate for dimension '{axis_name}' is required by a value selector but was not loaded" + )) + })?) + } else { + None + }; + let Some(selector_bounds) = + selector.conservative_bounds(length, coordinates, &mut poll_interrupt)? + else { + return Ok(Selection::empty(rank)); + }; + bounds[axis] = Some(selector_bounds); + } + Ok(Selection::from_axis_bounds(bounds)) + } + + pub(crate) fn matches_axis_index( + &self, + axis: usize, + index: usize, + coordinate_values: &[Option>], + ) -> ZarrFdwResult { + match self.by_axis.get(axis).and_then(Option::as_ref) { + None => Ok(true), + Some(selector) => { + let axis_name = self.axis_names.get(axis).map_or("unknown", String::as_str); + let coordinates = if self.requires_coordinate(axis) { + Some( + coordinate_values + .get(axis) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate for dimension '{axis_name}' is required by a value selector but was not loaded" + )) + })?, + ) + } else { + None + }; + selector.matches_index(index, coordinates) + } + } + } +} + +fn invalid_selector_option(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidOptionValue { + option: OPT_DIMENSION_SELECTORS.to_string(), + message: message.into(), + } +} + +fn parse_index(dimension: &str, value: JsonValue) -> ZarrFdwResult { + let index = value.as_u64().ok_or_else(|| { + invalid_selector_option(format!( + "selector index for dimension '{dimension}' must be a non-negative integer" + )) + })?; + usize::try_from(index).map_err(|_| { + invalid_selector_option(format!( + "selector index for dimension '{dimension}' exceeds this platform's index capacity" + )) + }) +} + +fn parse_value(dimension: &str, value: JsonValue) -> ZarrFdwResult { + let value = value + .as_f64() + .filter(|value| value.is_finite()) + .ok_or_else(|| { + invalid_selector_option(format!( + "selector value for dimension '{dimension}' must be a finite JSON number" + )) + })?; + Ok(normalize_zero(value)) +} + +fn parse_indices(dimension: &str, raw: RawList) -> ZarrFdwResult> { + if raw.0.is_empty() { + return Err(invalid_selector_option(format!( + "selector indices for dimension '{dimension}' must be a nonempty array" + ))); + } + + let mut indices = Vec::new(); + indices + .try_reserve_exact(raw.0.len()) + .map_err(|_| invalid_selector_option("could not allocate selector indices"))?; + for value in raw.0 { + let index = value.as_u64().ok_or_else(|| { + invalid_selector_option(format!( + "each selector index for dimension '{dimension}' must be a non-negative integer" + )) + })?; + indices.push(usize::try_from(index).map_err(|_| { + invalid_selector_option(format!( + "selector index for dimension '{dimension}' exceeds this platform's index capacity" + )) + })?); + } + indices.sort_unstable(); + if let Some(duplicate) = indices.windows(2).find(|pair| pair[0] == pair[1]) { + return Err(invalid_selector_option(format!( + "selector indices for dimension '{dimension}' contain duplicate index {}", + duplicate[0] + ))); + } + Ok(indices) +} + +fn parse_values(dimension: &str, raw: RawList) -> ZarrFdwResult> { + if raw.0.is_empty() { + return Err(invalid_selector_option(format!( + "selector values for dimension '{dimension}' must be a nonempty array" + ))); + } + + let mut values = Vec::new(); + values + .try_reserve_exact(raw.0.len()) + .map_err(|_| invalid_selector_option("could not allocate selector values"))?; + for value in raw.0 { + let value = value + .as_f64() + .filter(|value| value.is_finite()) + .ok_or_else(|| { + invalid_selector_option(format!( + "each selector value for dimension '{dimension}' must be a finite JSON number" + )) + })?; + values.push(normalize_zero(value)); + } + values.sort_by(|left, right| left.partial_cmp(right).expect("selector values are finite")); + if let Some(duplicate) = values.windows(2).find(|pair| pair[0] == pair[1]) { + return Err(invalid_selector_option(format!( + "selector values for dimension '{dimension}' contain duplicate value {}", + duplicate[0] + ))); + } + Ok(values) +} + +fn parse_index_range(dimension: &str, raw: RawRangeObject) -> ZarrFdwResult<(usize, usize)> { + let mut start = None; + let mut stop = None; + for (member, value) in raw.0 { + match member.as_str() { + "start" => start = Some(parse_index_range_member(dimension, "start", value)?), + "stop" => stop = Some(parse_index_range_member(dimension, "stop", value)?), + _ => { + return Err(invalid_selector_option(format!( + "index_range for dimension '{dimension}' must contain exactly 'start' and 'stop'" + ))); + } + } + } + let (Some(start), Some(stop)) = (start, stop) else { + return Err(invalid_selector_option(format!( + "index_range for dimension '{dimension}' must contain exactly 'start' and 'stop'" + ))); + }; + if start >= stop { + return Err(invalid_selector_option(format!( + "index_range for dimension '{dimension}' requires start < stop" + ))); + } + Ok((start, stop)) +} + +fn parse_index_range_member( + dimension: &str, + member: &str, + value: JsonValue, +) -> ZarrFdwResult { + let index = value.as_u64().ok_or_else(|| { + invalid_selector_option(format!( + "index_range member '{member}' for dimension '{dimension}' must be a non-negative integer" + )) + })?; + usize::try_from(index).map_err(|_| { + invalid_selector_option(format!( + "index_range member '{member}' for dimension '{dimension}' exceeds this platform's index capacity" + )) + }) +} + +fn parse_value_range(dimension: &str, raw: RawRangeObject) -> ZarrFdwResult<(f64, f64)> { + let mut min = None; + let mut max = None; + for (member, value) in raw.0 { + match member.as_str() { + "min" => min = Some(parse_value_range_member(dimension, "min", value)?), + "max" => max = Some(parse_value_range_member(dimension, "max", value)?), + _ => { + return Err(invalid_selector_option(format!( + "value_range for dimension '{dimension}' must contain exactly 'min' and 'max'" + ))); + } + } + } + let (Some(min), Some(max)) = (min, max) else { + return Err(invalid_selector_option(format!( + "value_range for dimension '{dimension}' must contain exactly 'min' and 'max'" + ))); + }; + if min > max { + return Err(invalid_selector_option(format!( + "value_range for dimension '{dimension}' requires min <= max" + ))); + } + Ok((min, max)) +} + +fn parse_value_range_member(dimension: &str, member: &str, value: JsonValue) -> ZarrFdwResult { + let value = value.as_f64().filter(|value| value.is_finite()).ok_or_else(|| { + invalid_selector_option(format!( + "value_range member '{member}' for dimension '{dimension}' must be a finite JSON number" + )) + })?; + Ok(normalize_zero(value)) +} + +fn normalize_zero(value: f64) -> f64 { + if value == 0.0 { 0.0 } else { value } +} + +struct RawDimensionSelectors(Vec<(String, RawSelector)>); + +impl<'de> Deserialize<'de> for RawDimensionSelectors { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_map(RawDimensionSelectorsVisitor) + } +} + +struct RawDimensionSelectorsVisitor; + +impl<'de> Visitor<'de> for RawDimensionSelectorsVisitor { + type Value = RawDimensionSelectors; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON object mapping dimension names to selector objects") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut entries = Vec::new(); + entries + .try_reserve_exact(MAX_SELECTOR_DIMENSIONS.min(map.size_hint().unwrap_or(0))) + .map_err(M::Error::custom)?; + let mut names = HashSet::new(); + names + .try_reserve(MAX_SELECTOR_DIMENSIONS.min(map.size_hint().unwrap_or(0))) + .map_err(M::Error::custom)?; + while let Some(dimension) = map.next_key::()? { + if entries.len() == MAX_SELECTOR_DIMENSIONS { + return Err(M::Error::custom(format!( + "dimension selector object exceeds the {MAX_SELECTOR_DIMENSIONS}-dimension limit" + ))); + } + if !names.insert(dimension.clone()) { + return Err(M::Error::custom(format!( + "duplicate dimension selector '{dimension}'" + ))); + } + let selector = map.next_value::()?; + if entries.len() == entries.capacity() { + entries.try_reserve(1).map_err(M::Error::custom)?; + } + entries.push((dimension, selector)); + } + Ok(RawDimensionSelectors(entries)) + } +} + +struct RawSelector(Vec<(String, RawSelectorValue)>); + +enum RawSelectorValue { + Json(JsonValue), + List(RawList), + Object(RawRangeObject), +} + +impl RawSelectorValue { + fn into_json(self) -> ZarrFdwResult { + match self { + Self::Json(value) => Ok(value), + Self::List(_) | Self::Object(_) => Err(invalid_selector_option( + "selector member has an invalid JSON value type", + )), + } + } + + fn into_list(self) -> ZarrFdwResult { + match self { + Self::List(value) => Ok(value), + Self::Json(_) | Self::Object(_) => Err(invalid_selector_option( + "selector list member must be a JSON array", + )), + } + } + + fn into_object(self) -> ZarrFdwResult { + match self { + Self::Object(value) => Ok(value), + Self::Json(_) | Self::List(_) => Err(invalid_selector_option( + "selector range member must be a JSON object", + )), + } + } +} + +impl<'de> Deserialize<'de> for RawSelector { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_map(RawSelectorVisitor) + } +} + +struct RawSelectorVisitor; + +impl<'de> Visitor<'de> for RawSelectorVisitor { + type Value = RawSelector; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a selector object") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut entries = Vec::new(); + entries.try_reserve_exact(1).map_err(M::Error::custom)?; + while let Some(kind) = map.next_key::()? { + if let Some((existing, _)) = entries.first() { + let message = if existing == &kind { + format!("duplicate selector member '{kind}'") + } else { + "selector object must contain exactly one supported selector form".to_string() + }; + return Err(M::Error::custom(message)); + } + let value = match kind.as_str() { + "indices" | "values" => RawSelectorValue::List(map.next_value::()?), + "index_range" | "value_range" => { + RawSelectorValue::Object(map.next_value::()?) + } + _ => RawSelectorValue::Json(map.next_value::()?), + }; + entries.push((kind, value)); + } + Ok(RawSelector(entries)) + } +} + +struct RawList(Vec); + +impl<'de> Deserialize<'de> for RawList { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_seq(RawListVisitor) + } +} + +struct RawListVisitor; + +impl<'de> Visitor<'de> for RawListVisitor { + type Value = RawList; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON array with at most 4096 entries") + } + + fn visit_seq(self, mut sequence: S) -> Result + where + S: SeqAccess<'de>, + { + let initial = sequence.size_hint().unwrap_or(0).min(MAX_SELECTOR_MEMBERS); + let mut values = Vec::new(); + values + .try_reserve_exact(initial) + .map_err(S::Error::custom)?; + while values.len() < MAX_SELECTOR_MEMBERS { + let Some(value) = sequence.next_element::()? else { + return Ok(RawList(values)); + }; + if values.len() == values.capacity() { + values.try_reserve(1).map_err(S::Error::custom)?; + } + values.push(value); + } + if sequence.next_element::()?.is_some() { + return Err(S::Error::custom(format!( + "selector list exceeds the {MAX_SELECTOR_MEMBERS}-entry limit" + ))); + } + Ok(RawList(values)) + } +} + +struct RawRangeObject(Vec<(String, JsonValue)>); + +impl<'de> Deserialize<'de> for RawRangeObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_map(RawRangeObjectVisitor) + } +} + +struct RawRangeObjectVisitor; + +impl<'de> Visitor<'de> for RawRangeObjectVisitor { + type Value = RawRangeObject; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a range selector object") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut entries = Vec::new(); + entries.try_reserve_exact(2).map_err(M::Error::custom)?; + let mut names = HashSet::new(); + names.try_reserve(2).map_err(M::Error::custom)?; + while let Some(name) = map.next_key::()? { + if !names.insert(name.clone()) { + return Err(M::Error::custom(format!( + "duplicate range selector member '{name}'" + ))); + } + if entries.len() == 2 { + return Err(M::Error::custom( + "range selector object contains more than two members", + )); + } + entries.push((name, map.next_value::()?)); + } + Ok(RawRangeObject(entries)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_exact_index_and_value_selectors() { + let selectors = + DimensionSelectors::parse(Some(r#"{"band":{"index":3},"level":{"value":850}}"#)) + .unwrap(); + let bound = selectors + .bind(&["level".to_string(), "band".to_string()], &[3, 4]) + .unwrap(); + + assert!(bound.requires_coordinate(0)); + assert!(!bound.requires_coordinate(1)); + let selection = bound + .resolve( + &[3, 4], + &[Some(vec![1000.0, 850.0, 500.0]), None], + || Ok(()), + ) + .unwrap(); + assert_eq!( + selection.axis_bounds(), + &[ + Some(IndexBounds { start: 1, end: 1 }), + Some(IndexBounds { start: 3, end: 3 }) + ] + ); + assert!( + bound + .matches_axis_index(0, 1, &[Some(vec![1000.0, 850.0, 500.0]), None]) + .unwrap() + ); + assert!( + !bound + .matches_axis_index(0, 0, &[Some(vec![1000.0, 850.0, 500.0]), None]) + .unwrap() + ); + } + + #[test] + fn value_selector_preserves_duplicate_exact_membership() { + let selectors = DimensionSelectors::parse(Some(r#"{"band":{"value":30}}"#)).unwrap(); + let bound = selectors.bind(&["band".to_string()], &[4]).unwrap(); + let selection = bound + .resolve(&[4], &[Some(vec![30.0, 10.0, 30.0, 20.0])], || Ok(())) + .unwrap(); + + assert_eq!( + selection.axis_bounds(), + &[Some(IndexBounds { start: 0, end: 2 })] + ); + let coordinates = [Some(vec![30.0, 10.0, 30.0, 20.0])]; + assert!(bound.matches_axis_index(0, 0, &coordinates).unwrap()); + assert!(!bound.matches_axis_index(0, 1, &coordinates).unwrap()); + assert!(bound.matches_axis_index(0, 2, &coordinates).unwrap()); + } + + #[test] + fn index_lists_and_ranges_use_conservative_bounds_and_exact_membership() { + let selectors = DimensionSelectors::parse(Some( + r#"{"band":{"indices":[4,1]},"level":{"index_range":{"start":2,"stop":5}}}"#, + )) + .unwrap(); + let bound = selectors + .bind(&["band".to_string(), "level".to_string()], &[6, 5]) + .unwrap(); + let selection = bound.resolve(&[6, 5], &[None, None], || Ok(())).unwrap(); + + assert_eq!( + selection.axis_bounds(), + &[ + Some(IndexBounds { start: 1, end: 4 }), + Some(IndexBounds { start: 2, end: 4 }) + ] + ); + assert!(bound.matches_axis_index(0, 1, &[None, None]).unwrap()); + assert!(!bound.matches_axis_index(0, 2, &[None, None]).unwrap()); + assert!(bound.matches_axis_index(0, 4, &[None, None]).unwrap()); + assert!(!bound.matches_axis_index(1, 1, &[None, None]).unwrap()); + assert!(bound.matches_axis_index(1, 4, &[None, None]).unwrap()); + } + + #[test] + fn value_lists_and_ranges_preserve_native_coordinate_membership() { + let selectors = DimensionSelectors::parse(Some( + r#"{"band":{"values":[40,10]},"level":{"value_range":{"min":15,"max":30}}}"#, + )) + .unwrap(); + let bound = selectors + .bind(&["band".to_string(), "level".to_string()], &[5, 4]) + .unwrap(); + let coordinates = [ + Some(vec![30.0, 10.0, 20.0, 10.0, 40.0]), + Some(vec![30.0, 10.0, 20.0, 40.0]), + ]; + let selection = bound.resolve(&[5, 4], &coordinates, || Ok(())).unwrap(); + + assert_eq!( + selection.axis_bounds(), + &[ + Some(IndexBounds { start: 1, end: 4 }), + Some(IndexBounds { start: 0, end: 2 }) + ] + ); + assert!(bound.matches_axis_index(0, 1, &coordinates).unwrap()); + assert!(!bound.matches_axis_index(0, 2, &coordinates).unwrap()); + assert!(bound.matches_axis_index(0, 3, &coordinates).unwrap()); + assert!(bound.matches_axis_index(0, 4, &coordinates).unwrap()); + assert!(bound.matches_axis_index(1, 0, &coordinates).unwrap()); + assert!(!bound.matches_axis_index(1, 1, &coordinates).unwrap()); + assert!(bound.matches_axis_index(1, 2, &coordinates).unwrap()); + } + + #[test] + fn no_matching_value_is_an_empty_selection() { + let selectors = DimensionSelectors::parse(Some(r#"{"band":{"value":999}}"#)).unwrap(); + let bound = selectors.bind(&["band".to_string()], &[2]).unwrap(); + let selection = bound + .resolve(&[2], &[Some(vec![10.0, 20.0])], || Ok(())) + .unwrap(); + + assert!(selection.is_empty()); + } + + #[test] + fn strict_parser_rejects_invalid_and_duplicate_members() { + for raw in [ + "[]", + r#"{"band":null}"#, + r#"{"band":{}}"#, + r#"{"band":{"index":1,"value":2}}"#, + r#"{"band":{"index":1,"index":2}}"#, + r#"{"band":{"unknown":1}}"#, + r#"{"band":{"index":-1}}"#, + r#"{"band":{"index":1.5}}"#, + r#"{"band":{"value":"B04"}}"#, + r#"{"band":{"value":null}}"#, + r#"{"band":{"index":1},"band":{"index":2}}"#, + r#"{"band":{"indices":[]}}"#, + r#"{"band":{"indices":[1,1]}}"#, + r#"{"band":{"values":[]}}"#, + r#"{"band":{"values":[-0.0,0.0]}}"#, + r#"{"band":{"values":[1,1.0]}}"#, + r#"{"band":{"index_range":{"start":2,"stop":2}}}"#, + r#"{"band":{"index_range":{"start":3,"stop":2}}}"#, + r#"{"band":{"index_range":{"start":0}}}"#, + r#"{"band":{"index_range":{"start":0,"start":1,"stop":2}}}"#, + r#"{"band":{"index_range":{"start":0,"stop":2,"step":1}}}"#, + r#"{"band":{"value_range":{"min":2,"max":1}}}"#, + r#"{"band":{"value_range":{"min":1}}}"#, + r#"{"band":{"value_range":{"min":0,"max":1,"step":0.5}}}"#, + ] { + assert!(DimensionSelectors::parse(Some(raw)).is_err(), "{raw}"); + } + } + + #[test] + fn parser_rejects_the_4097th_list_member_during_visitation() { + let members = std::iter::repeat_n("0", MAX_SELECTOR_MEMBERS + 1) + .collect::>() + .join(","); + let raw = format!(r#"{{"band":{{"indices":[{members}]}}}}"#); + + let error = DimensionSelectors::parse(Some(&raw)).unwrap_err(); + assert!(error.to_string().contains("4096-entry limit")); + } + + #[test] + fn parser_bounds_document_and_dimension_counts() { + let oversized = format!(r#"{{"band":{{"index":0}},"pad":"{}"}}"#, "x".repeat(65_536)); + assert!(DimensionSelectors::parse(Some(&oversized)).is_err()); + + let members = (0..65) + .map(|index| format!(r#""d{index}":{{"index":0}}"#)) + .collect::>() + .join(","); + assert!(DimensionSelectors::parse(Some(&format!("{{{members}}}"))).is_err()); + } + + #[test] + fn bind_and_index_resolution_fail_clearly() { + let unknown = DimensionSelectors::parse(Some(r#"{"member":{"index":0}}"#)).unwrap(); + assert!(unknown.bind(&["band".to_string()], &[6]).is_err()); + + let outside = DimensionSelectors::parse(Some(r#"{"band":{"index":6}}"#)).unwrap(); + let error = outside.bind(&["band".to_string()], &[6]).unwrap_err(); + assert!( + error + .to_string() + .contains("dimension selector index 6 is outside dimension 'band' length 6") + ); + + let list_outside = + DimensionSelectors::parse(Some(r#"{"band":{"indices":[0,6]}}"#)).unwrap(); + assert!(list_outside.bind(&["band".to_string()], &[6]).is_err()); + + let range_outside = + DimensionSelectors::parse(Some(r#"{"band":{"index_range":{"start":0,"stop":7}}}"#)) + .unwrap(); + assert!(range_outside.bind(&["band".to_string()], &[6]).is_err()); + + let range_to_extent = + DimensionSelectors::parse(Some(r#"{"band":{"index_range":{"start":0,"stop":6}}}"#)) + .unwrap(); + assert!(range_to_extent.bind(&["band".to_string()], &[6]).is_ok()); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/sharding.rs b/wrappers/src/fdw/zarr_fdw/sharding.rs new file mode 100644 index 000000000..15d4b7ef8 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/sharding.rs @@ -0,0 +1,1210 @@ +//! Bounded Zarr v3 `sharding_indexed` metadata and index execution. +//! +//! Only range-addressable shards are accepted: one top-level sharding codec, +//! the existing direct inner codec pipeline, and a fixed-size little-endian +//! index encoded as `bytes -> [crc32c]?`. Shard objects are never read whole. + +use std::sync::Arc; + +use lru::LruCache; +use serde_json::{Map, Value}; + +use super::codec::CodecPipeline; +use super::store::{RangedObject, ReadIdentity, ReadRange}; +use super::{ZarrFdwError, ZarrFdwResult}; + +const INDEX_ENTRY_BYTES: usize = 16; +const CRC32C_BYTES: usize = 4; +const INDEX_INTERRUPT_POLL_ENTRIES: usize = 4096; +const CRC_INTERRUPT_POLL_BYTES: usize = 1024 * 1024; +pub(crate) const MAX_SHARD_INDEX_BYTES: usize = 64 * 1024 * 1024 + CRC32C_BYTES; +const MISSING_SENTINEL: u64 = u64::MAX; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum IndexLocation { + Start, + End, +} + +impl IndexLocation { + pub(crate) fn label(self) -> &'static str { + match self { + Self::Start => "start", + Self::End => "end", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct ShardIndexCodec { + pub checksum: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ShardingConfig { + /// Native outer regular-grid chunk shape: one storage object per cell. + pub shard_shape: Vec, + /// Executor-visible logical chunk shape within each shard. + pub inner_chunk_shape: Vec, + pub chunks_per_shard: Vec, + pub inner_codecs: CodecPipeline, + pub index_codec: ShardIndexCodec, + pub index_location: IndexLocation, + pub index_entry_count: usize, + pub decoded_index_bytes: usize, + pub encoded_index_bytes: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ShardChunkAddress { + pub shard_indices: Vec, + pub inner_indices: Vec, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) enum StorageLayout { + #[default] + Direct, + Sharded(ShardingConfig), +} + +impl StorageLayout { + pub(crate) fn ordered_label(&self) -> String { + match self { + Self::Direct => "direct".to_string(), + Self::Sharded(config) => format!( + "sharding_indexed (index: {})", + config.index_location.label() + ), + } + } +} + +impl ShardingConfig { + /// Parse the complete top-level v3 codec list for a range-readable shard. + /// The outer list must contain exactly one `sharding_indexed` codec. + /// Returns the config and executor-normalized dtype produced by the inner + /// pipeline. + pub(crate) fn from_v3( + native_dtype: &str, + shard_shape: &[u64], + codecs: &Value, + ) -> ZarrFdwResult<(Self, String)> { + let codecs = codecs.as_array().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("Zarr v3 codecs must be an array".to_string()) + })?; + if codecs.len() != 1 { + return Err(ZarrFdwError::InvalidMetadata( + "range-readable sharded arrays require exactly one top-level sharding_indexed codec" + .to_string(), + )); + } + let codec = codec_object(&codecs[0], "sharding codec")?; + validate_fields( + codec, + &["name", "configuration", "must_understand"], + "sharding codec", + )?; + validate_must_understand(codec, "sharding codec")?; + if required_string(codec, "name", "sharding codec")? != "sharding_indexed" { + return Err(ZarrFdwError::InvalidMetadata( + "range-readable sharded arrays require a top-level sharding_indexed codec" + .to_string(), + )); + } + let configuration = required_object(codec, "configuration", "sharding codec")?; + validate_fields( + configuration, + &["chunk_shape", "codecs", "index_codecs", "index_location"], + "sharding_indexed configuration", + )?; + + if shard_shape.is_empty() { + return Err(ZarrFdwError::InvalidMetadata( + "sharding_indexed requires a non-empty shard shape".to_string(), + )); + } + let inner_chunk_shape = required_u64_array( + configuration, + "chunk_shape", + "sharding_indexed configuration", + )?; + if inner_chunk_shape.len() != shard_shape.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "sharding_indexed inner chunk rank {} does not match shard rank {}", + inner_chunk_shape.len(), + shard_shape.len() + ))); + } + + let mut chunks_per_shard = Vec::with_capacity(shard_shape.len()); + for (axis, (&shard, &inner)) in shard_shape.iter().zip(inner_chunk_shape.iter()).enumerate() + { + if shard == 0 || inner == 0 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "sharding_indexed chunk dimensions must be greater than zero on axis {axis}" + ))); + } + if shard % inner != 0 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "sharding_indexed inner chunk dimension {inner} does not evenly divide shard dimension {shard} on axis {axis}" + ))); + } + chunks_per_shard.push(shard / inner); + } + + let index_entry_count = chunks_per_shard.iter().try_fold(1usize, |count, &extent| { + let extent = usize::try_from(extent).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "shard index shape exceeds this platform's index capacity".to_string(), + ) + })?; + count.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "shard index entry count exceeds this platform's index capacity".to_string(), + ) + }) + })?; + let decoded_index_bytes = index_entry_count + .checked_mul(INDEX_ENTRY_BYTES) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "decoded shard index byte size exceeds this platform's index capacity" + .to_string(), + ) + })?; + + let index_codecs = configuration.get("index_codecs").ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "sharding_indexed configuration is missing 'index_codecs'".to_string(), + ) + })?; + let index_codec = parse_index_codecs(index_codecs)?; + let encoded_index_bytes = decoded_index_bytes + .checked_add(if index_codec.checksum { + CRC32C_BYTES + } else { + 0 + }) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "encoded shard index byte size exceeds this platform's index capacity" + .to_string(), + ) + })?; + if encoded_index_bytes > MAX_SHARD_INDEX_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "encoded shard index requires {encoded_index_bytes} bytes, exceeding the safety limit of {MAX_SHARD_INDEX_BYTES}" + ))); + } + + let inner_codecs = configuration.get("codecs").ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "sharding_indexed configuration is missing 'codecs'".to_string(), + ) + })?; + let (inner_codecs, normalized_dtype) = + CodecPipeline::from_v3(native_dtype, shard_shape.len(), inner_codecs)?; + + let index_location = match configuration.get("index_location") { + None => IndexLocation::End, + Some(Value::String(location)) if location == "start" => IndexLocation::Start, + Some(Value::String(location)) if location == "end" => IndexLocation::End, + Some(Value::String(location)) => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "sharding_indexed index_location must be 'start' or 'end', got '{location}'" + ))); + } + Some(_) => { + return Err(ZarrFdwError::InvalidMetadata( + "sharding_indexed index_location must be a string".to_string(), + )); + } + }; + + Ok(( + Self { + shard_shape: shard_shape.to_vec(), + inner_chunk_shape, + chunks_per_shard, + inner_codecs, + index_codec, + index_location, + index_entry_count, + decoded_index_bytes, + encoded_index_bytes, + }, + normalized_dtype, + )) + } + + /// Exact start range or suffix range used to fetch this shard's index. + pub(crate) fn index_read_identity( + &self, + shard_key: impl Into, + ) -> ZarrFdwResult { + let length = u64::try_from(self.encoded_index_bytes).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "encoded shard index size exceeds the Zarr u64 range capacity".to_string(), + ) + })?; + match self.index_location { + IndexLocation::Start => ReadIdentity::exact(shard_key, 0, length), + IndexLocation::End => ReadIdentity::suffix(shard_key, length), + } + } + + /// Map an executor logical-chunk coordinate to its outer shard key + /// coordinate and C-order inner-index coordinate. + pub(crate) fn chunk_address( + &self, + logical_chunk_indices: &[u64], + ) -> ZarrFdwResult { + let (shard_indices, inner_indices) = self.split_logical_indices(logical_chunk_indices)?; + Ok(ShardChunkAddress { + shard_indices, + inner_indices, + }) + } + + pub(crate) fn split_logical_indices( + &self, + logical_chunk_indices: &[u64], + ) -> ZarrFdwResult<(Vec, Vec)> { + if logical_chunk_indices.len() != self.chunks_per_shard.len() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "logical chunk index rank {} does not match sharded array rank {}", + logical_chunk_indices.len(), + self.chunks_per_shard.len() + ))); + } + let mut shard_indices = Vec::with_capacity(logical_chunk_indices.len()); + let mut inner_indices = Vec::with_capacity(logical_chunk_indices.len()); + for (&logical, &per_shard) in logical_chunk_indices + .iter() + .zip(self.chunks_per_shard.iter()) + { + if per_shard == 0 { + return Err(ZarrFdwError::InvalidMetadata( + "chunks per shard must be greater than zero".to_string(), + )); + } + shard_indices.push(logical / per_shard); + inner_indices.push(logical % per_shard); + } + Ok((shard_indices, inner_indices)) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ShardEntry { + Missing, + Present { offset: u64, nbytes: u64 }, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ShardIndex { + entries: Vec, + chunks_per_shard: Vec, + index_identity: ReadIdentity, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum ShardIndexDecode { + Decoded(ShardIndex), + Interrupted, +} + +impl ShardIndex { + /// Decode and validate one exactly ranged index response. All offsets and + /// lengths are checked against the shard and index regions. Physical + /// payload order and shared/overlapping payload representations remain + /// unconstrained, as required by the sharding specification. + pub(crate) fn decode_interruptible( + config: &ShardingConfig, + response: RangedObject, + mut interrupt_pending: F, + ) -> ZarrFdwResult + where + F: FnMut() -> bool, + { + let shard_key = response.identity.key.clone(); + if response.bytes.len() != config.encoded_index_bytes { + return Err(shard_error( + &shard_key, + format!( + "shard index range returned {} bytes, expected exactly {} bytes", + response.bytes.len(), + config.encoded_index_bytes + ), + )); + } + let generation = response.identity.generation.as_ref().ok_or_else(|| { + shard_error( + &shard_key, + "shard index response is missing its object generation", + ) + })?; + if generation.validator_is_empty() { + return Err(shard_error( + &shard_key, + "shard index response has an empty object-generation validator", + )); + } + if generation.total_len() != response.total_len { + return Err(shard_error( + &shard_key, + format!( + "shard index generation length {} does not match Content-Range total {}", + generation.total_len(), + response.total_len + ), + )); + } + validate_index_range(config, &response)?; + if interrupt_pending() { + return Ok(ShardIndexDecode::Interrupted); + } + + let mut bytes = response.bytes; + if config.index_codec.checksum { + let payload_len = bytes.len().checked_sub(CRC32C_BYTES).ok_or_else(|| { + shard_error( + &shard_key, + "shard index codec index 1 ('crc32c'): index is truncated before the checksum", + ) + })?; + let expected = u32::from_le_bytes( + bytes[payload_len..] + .try_into() + .expect("crc suffix is exactly four bytes"), + ); + let mut actual = 0u32; + for block in bytes[..payload_len].chunks(CRC_INTERRUPT_POLL_BYTES) { + if interrupt_pending() { + return Ok(ShardIndexDecode::Interrupted); + } + actual = crc32c::crc32c_append(actual, block); + } + if actual != expected { + return Err(shard_error( + &shard_key, + format!( + "shard index codec index 1 ('crc32c'): checksum mismatch: expected {expected:#010x}, computed {actual:#010x}" + ), + )); + } + bytes.truncate(payload_len); + } + if bytes.len() != config.decoded_index_bytes { + return Err(shard_error( + &shard_key, + format!( + "decoded shard index has {} bytes, expected exactly {} bytes", + bytes.len(), + config.decoded_index_bytes + ), + )); + } + + let index_region = index_region(config, response.total_len)?; + let mut entries = Vec::new(); + entries + .try_reserve_exact(config.index_entry_count) + .map_err(|_| { + shard_error( + &shard_key, + format!( + "could not allocate {} shard index entries", + config.index_entry_count + ), + ) + })?; + for (index, pair) in bytes.chunks_exact(INDEX_ENTRY_BYTES).enumerate() { + if index % INDEX_INTERRUPT_POLL_ENTRIES == 0 && interrupt_pending() { + return Ok(ShardIndexDecode::Interrupted); + } + let offset = u64::from_le_bytes(pair[..8].try_into().expect("eight offset bytes")); + let nbytes = u64::from_le_bytes(pair[8..].try_into().expect("eight length bytes")); + let entry = match (offset == MISSING_SENTINEL, nbytes == MISSING_SENTINEL) { + (true, true) => ShardEntry::Missing, + (true, false) | (false, true) => { + return Err(shard_error( + &shard_key, + format!( + "shard index entry {index} uses a mixed uint64 missing sentinel; offset and nbytes must both be 2^64 - 1" + ), + )); + } + (false, false) => { + if nbytes == 0 { + return Err(shard_error( + &shard_key, + format!("shard index entry {index} has zero encoded bytes"), + )); + } + let end = offset.checked_add(nbytes).ok_or_else(|| { + shard_error( + &shard_key, + format!("shard index entry {index} byte range overflows u64"), + ) + })?; + if end > response.total_len { + return Err(shard_error( + &shard_key, + format!( + "inner chunk byte range {offset}..{end} from shard index entry {index} exceeds shard object length {}", + response.total_len + ), + )); + } + if offset < index_region.1 && end > index_region.0 { + return Err(shard_error( + &shard_key, + format!( + "inner chunk byte range {offset}..{end} from shard index entry {index} overlaps shard index region {}..{}", + index_region.0, index_region.1 + ), + )); + } + ShardEntry::Present { offset, nbytes } + } + }; + entries.push(entry); + } + if entries.len() != config.index_entry_count { + return Err(shard_error( + &shard_key, + format!( + "decoded shard index has {} entries, expected exactly {}", + entries.len(), + config.index_entry_count + ), + )); + } + if interrupt_pending() { + return Ok(ShardIndexDecode::Interrupted); + } + + Ok(ShardIndexDecode::Decoded(Self { + entries, + chunks_per_shard: config.chunks_per_shard.clone(), + index_identity: response.identity, + })) + } + + pub(crate) fn entry(&self, inner_indices: &[u64]) -> ZarrFdwResult { + if inner_indices.len() != self.chunks_per_shard.len() { + return Err(shard_error( + &self.index_identity.key, + format!( + "inner chunk index rank {} does not match shard rank {}", + inner_indices.len(), + self.chunks_per_shard.len() + ), + )); + } + let flat = inner_indices + .iter() + .zip(self.chunks_per_shard.iter()) + .enumerate() + .try_fold(0usize, |flat, (axis, (&index, &extent))| { + if index >= extent { + return Err(shard_error( + &self.index_identity.key, + format!( + "inner chunk index {index} is outside shard extent {extent} on axis {axis}" + ), + )); + } + let extent = usize::try_from(extent).map_err(|_| { + shard_error( + &self.index_identity.key, + "shard index extent exceeds this platform's index capacity", + ) + })?; + let index = usize::try_from(index).map_err(|_| { + shard_error( + &self.index_identity.key, + "inner chunk index exceeds this platform's index capacity", + ) + })?; + flat.checked_mul(extent) + .and_then(|flat| flat.checked_add(index)) + .ok_or_else(|| { + shard_error( + &self.index_identity.key, + "flat shard index position exceeds this platform's index capacity", + ) + }) + })?; + self.entries.get(flat).copied().ok_or_else(|| { + shard_error( + &self.index_identity.key, + "flat shard index position is outside the decoded index", + ) + }) + } + + /// Build a generation-conditioned range request for one present entry. + pub(crate) fn payload_read_identity( + &self, + entry: ShardEntry, + ) -> ZarrFdwResult> { + let ShardEntry::Present { offset, nbytes } = entry else { + return Ok(None); + }; + let generation = self.index_identity.generation.clone().ok_or_else(|| { + shard_error( + &self.index_identity.key, + "decoded shard index is missing its object generation", + ) + })?; + Ok(Some( + ReadIdentity::exact(self.index_identity.key.clone(), offset, nbytes)? + .with_generation(generation), + )) + } + + pub(crate) fn index_identity(&self) -> &ReadIdentity { + &self.index_identity + } + + fn resident_bytes(&self) -> usize { + self.entries + .capacity() + .saturating_mul(std::mem::size_of::()) + } +} + +/// Query-local byte- and entry-bounded cache of decoded shard indexes. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum CachedShardIndex { + Present(Arc), + Missing, +} + +impl CachedShardIndex { + fn resident_bytes(&self) -> usize { + match self { + Self::Present(index) => index.resident_bytes(), + Self::Missing => 0, + } + } +} + +pub(crate) struct ShardIndexCache { + entries: LruCache, + resident_bytes: usize, + max_bytes: usize, + max_entries: usize, + evictions: usize, +} + +impl ShardIndexCache { + pub(crate) fn new(max_bytes: usize, max_entries: usize) -> Self { + Self { + entries: LruCache::unbounded(), + resident_bytes: 0, + max_bytes, + max_entries, + evictions: 0, + } + } + + pub(crate) fn get(&mut self, request: &ReadIdentity) -> Option { + self.entries.get(request).cloned() + } + + /// Cache an index under the request identity that located it (start exact + /// or end suffix). The cached value retains the resolved exact identity + /// and observed generation needed by payload reads. + pub(crate) fn insert_present(&mut self, request: ReadIdentity, index: Arc) -> bool { + self.insert(request, CachedShardIndex::Present(index)) + } + + /// Cache an explicit absent outer shard. Missing entries consume the + /// entry budget but no byte budget, preventing repeated range GETs for + /// every logical inner chunk in the shard. + pub(crate) fn insert_missing(&mut self, request: ReadIdentity) -> bool { + self.insert(request, CachedShardIndex::Missing) + } + + fn insert(&mut self, request: ReadIdentity, value: CachedShardIndex) -> bool { + if let Some(previous) = self.entries.pop(&request) { + self.resident_bytes = self + .resident_bytes + .saturating_sub(previous.resident_bytes()); + } + let bytes = value.resident_bytes(); + if self.max_bytes == 0 || self.max_entries == 0 || bytes > self.max_bytes { + return false; + } + while self.entries.len() >= self.max_entries + || self + .resident_bytes + .checked_add(bytes) + .is_none_or(|total| total > self.max_bytes) + { + let Some((_key, evicted)) = self.entries.pop_lru() else { + break; + }; + self.resident_bytes = self.resident_bytes.saturating_sub(evicted.resident_bytes()); + self.evictions = self.evictions.saturating_add(1); + } + self.resident_bytes += bytes; + self.entries.put(request, value); + true + } + + pub(crate) fn len(&self) -> usize { + self.entries.len() + } + + pub(crate) fn resident_bytes(&self) -> usize { + self.resident_bytes + } + + pub(crate) fn evictions(&self) -> usize { + self.evictions + } +} + +fn parse_index_codecs(value: &Value) -> ZarrFdwResult { + let codecs = value.as_array().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("sharding_indexed index_codecs must be an array".to_string()) + })?; + if !(1..=2).contains(&codecs.len()) { + return Err(ZarrFdwError::InvalidMetadata( + "supported shard index codecs are exactly bytes -> [crc32c]?".to_string(), + )); + } + let bytes = codec_object(&codecs[0], "shard index codec index 0")?; + validate_fields( + bytes, + &["name", "configuration", "must_understand"], + "shard index codec index 0", + )?; + validate_must_understand(bytes, "shard index codec index 0")?; + if required_string(bytes, "name", "shard index codec index 0")? != "bytes" { + return Err(ZarrFdwError::InvalidMetadata( + "shard index codec index 0 must be 'bytes'".to_string(), + )); + } + let configuration = required_object(bytes, "configuration", "shard index codec index 0")?; + validate_fields( + configuration, + &["endian"], + "shard index bytes configuration", + )?; + if required_string(configuration, "endian", "shard index bytes configuration")? != "little" { + return Err(ZarrFdwError::InvalidMetadata( + "shard index bytes endian must be 'little'".to_string(), + )); + } + + let checksum = codecs.len() == 2; + if checksum { + let crc = codec_object(&codecs[1], "shard index codec index 1")?; + validate_fields( + crc, + &["name", "configuration", "must_understand"], + "shard index codec index 1", + )?; + validate_must_understand(crc, "shard index codec index 1")?; + if required_string(crc, "name", "shard index codec index 1")? != "crc32c" { + return Err(ZarrFdwError::InvalidMetadata( + "shard index codec index 1 must be 'crc32c'".to_string(), + )); + } + if let Some(configuration) = crc.get("configuration") { + let configuration = configuration.as_object().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "shard index crc32c configuration must be an object".to_string(), + ) + })?; + validate_fields(configuration, &[], "shard index crc32c configuration")?; + } + } + Ok(ShardIndexCodec { checksum }) +} + +fn validate_index_range(config: &ShardingConfig, response: &RangedObject) -> ZarrFdwResult<()> { + let length = u64::try_from(config.encoded_index_bytes).map_err(|_| { + shard_error( + &response.identity.key, + "encoded shard index length exceeds u64", + ) + })?; + let ReadRange::Exact { + start, + length: actual_length, + } = &response.identity.range + else { + return Err(shard_error( + &response.identity.key, + "shard index response was not normalized to an exact byte range", + )); + }; + if *actual_length != length { + return Err(shard_error( + &response.identity.key, + format!("shard index range length is {actual_length}, expected exactly {length} bytes"), + )); + } + let expected_start = match config.index_location { + IndexLocation::Start => 0, + IndexLocation::End => response.total_len.checked_sub(length).ok_or_else(|| { + shard_error( + &response.identity.key, + format!( + "shard object length {} is smaller than its {length}-byte index", + response.total_len + ), + ) + })?, + }; + if *start != expected_start { + return Err(shard_error( + &response.identity.key, + format!( + "shard index starts at byte {start}, expected byte {expected_start} for index_location '{}'", + config.index_location.label() + ), + )); + } + Ok(()) +} + +fn index_region(config: &ShardingConfig, total_len: u64) -> ZarrFdwResult<(u64, u64)> { + let length = u64::try_from(config.encoded_index_bytes).map_err(|_| { + ZarrFdwError::InvalidMetadata("encoded shard index length exceeds u64".to_string()) + })?; + match config.index_location { + IndexLocation::Start => Ok((0, length)), + IndexLocation::End => total_len + .checked_sub(length) + .map(|start| (start, total_len)) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "shard object length {total_len} is smaller than its {length}-byte index" + )) + }), + } +} + +fn codec_object<'a>(value: &'a Value, context: &str) -> ZarrFdwResult<&'a Map> { + value + .as_object() + .ok_or_else(|| ZarrFdwError::InvalidMetadata(format!("{context} must be an object"))) +} + +fn required_object<'a>( + object: &'a Map, + field: &str, + context: &str, +) -> ZarrFdwResult<&'a Map> { + object.get(field).and_then(Value::as_object).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("{context} field '{field}' must be an object")) + }) +} + +fn required_string<'a>( + object: &'a Map, + field: &str, + context: &str, +) -> ZarrFdwResult<&'a str> { + object.get(field).and_then(Value::as_str).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("{context} field '{field}' must be a string")) + }) +} + +fn required_u64_array( + object: &Map, + field: &str, + context: &str, +) -> ZarrFdwResult> { + object + .get(field) + .and_then(Value::as_array) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("{context} field '{field}' must be an array")) + })? + .iter() + .map(|value| { + value.as_u64().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "{context} field '{field}' must contain non-negative integers" + )) + }) + }) + .collect() +} + +fn validate_fields( + object: &Map, + allowed: &[&str], + context: &str, +) -> ZarrFdwResult<()> { + if let Some(field) = object + .keys() + .find(|field| !allowed.contains(&field.as_str())) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} contains unsupported field '{field}'" + ))); + } + Ok(()) +} + +fn validate_must_understand(object: &Map, context: &str) -> ZarrFdwResult<()> { + if object + .get("must_understand") + .is_some_and(|value| !value.is_boolean()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "{context} must_understand must be a boolean" + ))); + } + Ok(()) +} + +fn shard_error(key: &str, message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!("shard '{key}': {}", message.into())) +} + +#[cfg(test)] +mod tests { + use super::super::store::ObjectGeneration; + use super::*; + + fn codecs(index_location: &str, checksum: bool) -> Value { + let mut index_codecs = vec![serde_json::json!({ + "name":"bytes", + "configuration":{"endian":"little"} + })]; + if checksum { + index_codecs.push(serde_json::json!({"name":"crc32c"})); + } + serde_json::json!([{ + "name":"sharding_indexed", + "configuration":{ + "chunk_shape":[1,3,2], + "codecs":[{"name":"bytes","configuration":{"endian":"little"}}], + "index_codecs":index_codecs, + "index_location":index_location + } + }]) + } + + fn config(location: &str, checksum: bool) -> ShardingConfig { + ShardingConfig::from_v3("float32", &[2, 3, 4], &codecs(location, checksum)) + .unwrap() + .0 + } + + fn encoded_index(entries: &[(u64, u64)], checksum: bool) -> Vec { + let mut bytes = Vec::new(); + for &(offset, nbytes) in entries { + bytes.extend_from_slice(&offset.to_le_bytes()); + bytes.extend_from_slice(&nbytes.to_le_bytes()); + } + if checksum { + bytes.extend_from_slice(&crc32c::crc32c(&bytes).to_le_bytes()); + } + bytes + } + + fn response(config: &ShardingConfig, total_len: u64, bytes: Vec) -> RangedObject { + let index_len = config.encoded_index_bytes as u64; + let start = match config.index_location { + IndexLocation::Start => 0, + IndexLocation::End => total_len - index_len, + }; + RangedObject { + identity: ReadIdentity::exact("array/c/0/0/0", start, index_len) + .unwrap() + .with_generation(ObjectGeneration::S3 { + etag: "\"etag-1\"".to_string(), + version_id: Some("version-observed-only".to_string()), + total_len, + }), + total_len, + bytes, + } + } + + #[test] + fn parses_start_end_and_checked_index_size() { + let end = config("end", true); + assert_eq!(end.chunks_per_shard, vec![2, 1, 2]); + assert_eq!(end.index_entry_count, 4); + assert_eq!(end.decoded_index_bytes, 64); + assert_eq!(end.encoded_index_bytes, 68); + assert_eq!(end.index_location, IndexLocation::End); + assert!(matches!( + end.index_read_identity("shard").unwrap().range, + ReadRange::Suffix { length: 68 } + )); + assert_eq!( + end.chunk_address(&[3, 0, 5]).unwrap(), + ShardChunkAddress { + shard_indices: vec![1, 0, 2], + inner_indices: vec![1, 0, 1] + } + ); + + let start = config("start", false); + assert_eq!(start.encoded_index_bytes, 64); + assert!(matches!( + start.index_read_identity("shard").unwrap().range, + ReadRange::Exact { + start: 0, + length: 64 + } + )); + + let mut default_end = codecs("end", true); + default_end[0]["configuration"] + .as_object_mut() + .unwrap() + .remove("index_location"); + assert_eq!( + ShardingConfig::from_v3("float32", &[2, 3, 4], &default_end) + .unwrap() + .0 + .index_location, + IndexLocation::End + ); + } + + #[test] + fn rejects_invalid_sharding_metadata() { + let mut invalid = codecs("end", true); + invalid[0]["configuration"]["chunk_shape"] = serde_json::json!([1, 2, 2]); + assert!(ShardingConfig::from_v3("float32", &[2, 3, 4], &invalid).is_err()); + + let mut invalid = codecs("middle", true); + assert!(ShardingConfig::from_v3("float32", &[2, 3, 4], &invalid).is_err()); + invalid[0]["configuration"]["index_location"] = serde_json::json!("end"); + invalid[0]["configuration"]["index_codecs"] = + serde_json::json!([{"name":"gzip","configuration":{"level":1}}]); + assert!(ShardingConfig::from_v3("float32", &[2, 3, 4], &invalid).is_err()); + + let extra_outer = serde_json::json!([ + codecs("end", true)[0].clone(), + {"name":"crc32c"} + ]); + assert!(ShardingConfig::from_v3("float32", &[2, 3, 4], &extra_outer).is_err()); + } + + #[test] + fn decodes_c_order_entries_and_generation_conditioned_payload_ranges() { + let config = config("end", true); + let entries = [(0, 24), (24, 24), (48, 24), (72, 24)]; + let bytes = encoded_index(&entries, true); + let total_len = 96 + config.encoded_index_bytes as u64; + let decoded = match ShardIndex::decode_interruptible( + &config, + response(&config, total_len, bytes), + || false, + ) + .unwrap() + { + ShardIndexDecode::Decoded(index) => index, + ShardIndexDecode::Interrupted => panic!("unexpected interrupt"), + }; + assert_eq!( + decoded.entry(&[1, 0, 0]).unwrap(), + ShardEntry::Present { + offset: 48, + nbytes: 24 + } + ); + let payload = decoded + .payload_read_identity(decoded.entry(&[0, 0, 1]).unwrap()) + .unwrap() + .unwrap(); + assert_eq!( + payload.range, + ReadRange::Exact { + start: 24, + length: 24 + } + ); + assert_eq!(payload.generation.unwrap().s3_etag(), Some("\"etag-1\"")); + } + + #[test] + fn both_uint64_max_values_are_the_only_missing_sentinel() { + let config = config("end", false); + let valid = [(MISSING_SENTINEL, MISSING_SENTINEL), (0, 1), (1, 1), (2, 1)]; + let total_len = 3 + config.encoded_index_bytes as u64; + let decoded = ShardIndex::decode_interruptible( + &config, + response(&config, total_len, encoded_index(&valid, false)), + || false, + ) + .unwrap(); + let ShardIndexDecode::Decoded(decoded) = decoded else { + panic!("unexpected interrupt") + }; + assert_eq!(decoded.entry(&[0, 0, 0]).unwrap(), ShardEntry::Missing); + assert!( + decoded + .payload_read_identity(ShardEntry::Missing) + .unwrap() + .is_none() + ); + + for mixed in [(MISSING_SENTINEL, 1), (1, MISSING_SENTINEL)] { + let entries = [mixed, (0, 1), (1, 1), (2, 1)]; + assert!( + ShardIndex::decode_interruptible( + &config, + response(&config, total_len, encoded_index(&entries, false)), + || false, + ) + .is_err() + ); + } + } + + #[test] + fn rejects_checksum_oob_index_overlap_and_overflow() { + let config = config("end", true); + let total_len = 100 + config.encoded_index_bytes as u64; + let valid = [(0, 10), (10, 10), (20, 10), (30, 10)]; + + let mut bad_crc = encoded_index(&valid, true); + let last = bad_crc.len() - 1; + bad_crc[last] ^= 0xff; + assert!( + ShardIndex::decode_interruptible( + &config, + response(&config, total_len, bad_crc), + || false, + ) + .is_err() + ); + + for entries in [ + [(0, 10), (10, 10), (20, 10), (99, 2)], + [(0, 10), (10, 10), (20, 10), (100, 1)], + [(0, 10), (10, 10), (20, 10), (u64::MAX - 1, 4)], + ] { + assert!( + ShardIndex::decode_interruptible( + &config, + response(&config, total_len, encoded_index(&entries, true)), + || false, + ) + .is_err() + ); + } + } + + #[test] + fn start_index_forbids_payload_overlap_and_decode_is_interruptible() { + let config = config("start", false); + let index_len = config.encoded_index_bytes as u64; + let overlap = [ + (0, 1), + (index_len, 1), + (index_len + 1, 1), + (index_len + 2, 1), + ]; + assert!( + ShardIndex::decode_interruptible( + &config, + response(&config, index_len + 3, encoded_index(&overlap, false)), + || false, + ) + .is_err() + ); + + let valid = [ + (index_len, 1), + (index_len + 1, 1), + (index_len + 2, 1), + (index_len + 3, 1), + ]; + assert_eq!( + ShardIndex::decode_interruptible( + &config, + response(&config, index_len + 4, encoded_index(&valid, false)), + || true, + ) + .unwrap(), + ShardIndexDecode::Interrupted + ); + } + + #[test] + fn index_cache_is_bounded_and_keys_start_and_end_requests_distinctly() { + let config = config("end", false); + let entries = [(0, 1), (1, 1), (2, 1), (3, 1)]; + let total_len = 4 + config.encoded_index_bytes as u64; + let ShardIndexDecode::Decoded(index) = ShardIndex::decode_interruptible( + &config, + response(&config, total_len, encoded_index(&entries, false)), + || false, + ) + .unwrap() else { + panic!("unexpected interrupt") + }; + let index = Arc::new(index); + let request = config.index_read_identity("array/c/0/0/0").unwrap(); + let mut cache = ShardIndexCache::new(index.resident_bytes() * 2, 1); + assert!(cache.insert_present(request.clone(), Arc::clone(&index))); + assert!(matches!( + cache.get(&request), + Some(CachedShardIndex::Present(_)) + )); + assert_eq!(cache.len(), 1); + assert!(cache.resident_bytes() > 0); + assert_eq!(cache.evictions(), 0); + + let other = config.index_read_identity("array/c/0/0/1").unwrap(); + assert!(cache.insert_present(other.clone(), index)); + assert!(cache.get(&request).is_none()); + assert!(cache.get(&other).is_some()); + assert_eq!(cache.evictions(), 1); + + let missing = config.index_read_identity("array/c/0/0/2").unwrap(); + assert!(cache.insert_missing(missing.clone())); + assert_eq!(cache.get(&missing), Some(CachedShardIndex::Missing)); + assert_eq!(cache.resident_bytes(), 0); + } + + #[test] + fn index_cache_charges_the_actual_entry_allocation() { + let config = config("end", true); + let entries = [(0, 24), (24, 24), (48, 24), (72, 24)]; + let total_len = 96 + config.encoded_index_bytes as u64; + let index = match ShardIndex::decode_interruptible( + &config, + response(&config, total_len, encoded_index(&entries, true)), + || false, + ) + .unwrap() + { + ShardIndexDecode::Decoded(index) => index, + ShardIndexDecode::Interrupted => panic!("decode was not interrupted"), + }; + let actual_allocation = index + .entries + .capacity() + .checked_mul(std::mem::size_of::()) + .unwrap(); + assert_eq!(index.resident_bytes(), actual_allocation); + assert!(actual_allocation > config.decoded_index_bytes); + + let mut underfunded = ShardIndexCache::new(config.decoded_index_bytes, 1); + assert!(!underfunded.insert_present( + config.index_read_identity("array/c/0/0/0").unwrap(), + Arc::new(index), + )); + assert_eq!(underfunded.resident_bytes(), 0); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/catalog.rs b/wrappers/src/fdw/zarr_fdw/spatial/catalog.rs new file mode 100644 index 000000000..39ab502b7 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/catalog.rs @@ -0,0 +1,189 @@ +//! Security-invoker catalog loading for spatial operations over a Zarr table. + +use std::collections::HashMap; + +use pgrx::{pg_sys, prelude::*}; +use supabase_wrappers::prelude::{Column, ForeignServer}; + +use super::super::selectors::OPT_DIMENSION_SELECTORS; +use super::super::{ZarrFdwError, ZarrFdwResult}; + +#[derive(Debug)] +pub(crate) struct SpatialForeignTable { + pub(crate) server: ForeignServer, + pub(crate) options: HashMap, + pub(crate) columns: Vec, +} + +/// Resolve a caller-visible relation and load only the catalog state needed to +/// construct a fresh Zarr executor. The privilege predicates deliberately make +/// missing and inaccessible relations indistinguishable. +pub(crate) fn load_zarr_foreign_table(relation_name: &str) -> ZarrFdwResult { + load_zarr_foreign_table_impl(relation_name, false) +} + +/// Resolve a caller-visible relation for a selector-aware spatial overload. +/// +/// The same catalog identity and privilege checks run before the caller parses +/// selector input or constructs a storage backend. Only the explicit overloads +/// opt out of the legacy fail-closed table-option guard. +pub(crate) fn load_zarr_foreign_table_with_selectors( + relation_name: &str, +) -> ZarrFdwResult { + load_zarr_foreign_table_impl(relation_name, true) +} + +fn load_zarr_foreign_table_impl( + relation_name: &str, + allow_dimension_selectors: bool, +) -> ZarrFdwResult { + let (table_oid, server_oid, server_name, server_type, server_version) = Spi::connect( + |client| { + let mut rows = client.select( + "SELECT c.oid::bigint AS table_oid, + s.oid::bigint AS server_oid, + s.srvname::text AS server_name, + s.srvtype::text AS server_type, + s.srvversion::text AS server_version + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_foreign_table AS ft ON ft.ftrelid = c.oid + JOIN pg_catalog.pg_foreign_server AS s ON s.oid = ft.ftserver + JOIN pg_catalog.pg_foreign_data_wrapper AS w ON w.oid = s.srvfdw + JOIN pg_catalog.pg_proc AS handler ON handler.oid = w.fdwhandler + WHERE c.oid = pg_catalog.to_regclass($1) + AND pg_catalog.has_table_privilege(c.oid, 'SELECT') + AND pg_catalog.has_server_privilege(s.oid, 'USAGE') + AND handler.proname = 'zarr_fdw_handler' + AND EXISTS ( + SELECT 1 + FROM pg_catalog.pg_depend AS dependency + JOIN pg_catalog.pg_extension AS extension + ON extension.oid = dependency.refobjid + WHERE dependency.classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND dependency.objid = handler.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.deptype = 'e' + AND extension.extname = 'wrappers' + )", + Some(1), + &[relation_name.into()], + )?; + let Some(row) = rows.next() else { + return Ok::<_, pgrx::spi::Error>((None, None, None, None, None)); + }; + Ok(( + row.get_by_name::("table_oid")?, + row.get_by_name::("server_oid")?, + row.get_by_name::("server_name")?, + row.get_by_name::("server_type")?, + row.get_by_name::("server_version")?, + )) + }, + )?; + + let (table_oid, server_oid, server_name) = match (table_oid, server_oid, server_name) { + (Some(table_oid), Some(server_oid), Some(server_name)) => { + (table_oid, server_oid, server_name) + } + _ => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "foreign table '{relation_name}' does not exist or is not accessible" + ))); + } + }; + + let table_options = load_options( + "SELECT option.option_name::text, option.option_value::text + FROM pg_catalog.pg_foreign_table AS table_catalog + CROSS JOIN LATERAL pg_catalog.pg_options_to_table(table_catalog.ftoptions) AS option + WHERE table_catalog.ftrelid = $1::bigint::pg_catalog.oid", + table_oid, + )?; + if !allow_dimension_selectors && table_options.contains_key(OPT_DIMENSION_SELECTORS) { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_DIMENSION_SELECTORS.to_string(), + message: "spatial operations require a selector-aware function overload".to_string(), + }); + } + let server_options = load_options( + "SELECT option.option_name::text, option.option_value::text + FROM pg_catalog.pg_foreign_server AS server_catalog + CROSS JOIN LATERAL pg_catalog.pg_options_to_table(server_catalog.srvoptions) AS option + WHERE server_catalog.oid = $1::bigint::pg_catalog.oid", + server_oid, + )?; + let columns = load_columns(table_oid)?; + if columns.is_empty() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "foreign table '{relation_name}' has no readable columns" + ))); + } + + let server_oid = u32::try_from(server_oid).map_err(|_| { + ZarrFdwError::InvalidMetadata("foreign server OID is out of range".to_string()) + })?; + Ok(SpatialForeignTable { + server: ForeignServer { + server_oid: pg_sys::Oid::from_u32(server_oid), + server_name, + server_type, + server_version, + options: server_options, + }, + options: table_options, + columns, + }) +} + +fn load_options(sql: &str, owner_oid: i64) -> ZarrFdwResult> { + Spi::connect(|client| { + let rows = client.select(sql, None, &[owner_oid.into()])?; + let mut options = HashMap::new(); + for row in rows { + if let (Some(name), Some(value)) = (row.get::(1)?, row.get::(2)?) { + options.insert(name, value); + } + } + Ok::<_, pgrx::spi::Error>(options) + }) + .map_err(Into::into) +} + +fn load_columns(table_oid: i64) -> ZarrFdwResult> { + Spi::connect(|client| { + let rows = client.select( + "SELECT attribute.attname::text AS name, + attribute.attnum::integer AS number, + attribute.atttypid::bigint AS type_oid + FROM pg_catalog.pg_attribute AS attribute + WHERE attribute.attrelid = $1::bigint::pg_catalog.oid + AND attribute.attnum > 0 + AND NOT attribute.attisdropped + ORDER BY attribute.attnum", + None, + &[table_oid.into()], + )?; + let mut columns = Vec::new(); + for row in rows { + let Some(name) = row.get_by_name::("name")? else { + continue; + }; + let number = row + .get_by_name::("number")? + .ok_or(pgrx::spi::Error::InvalidPosition)?; + let type_oid = row + .get_by_name::("type_oid")? + .ok_or(pgrx::spi::Error::InvalidPosition)?; + let number = usize::try_from(number).map_err(|_| pgrx::spi::Error::InvalidPosition)?; + let type_oid = + u32::try_from(type_oid).map_err(|_| pgrx::spi::Error::InvalidPosition)?; + columns.push(Column { + name, + num: number, + type_oid: pg_sys::Oid::from_u32(type_oid), + }); + } + Ok::<_, pgrx::spi::Error>(columns) + }) + .map_err(Into::into) +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/crs.rs b/wrappers/src/fdw/zarr_fdw/spatial/crs.rs new file mode 100644 index 000000000..fa2be2803 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/crs.rs @@ -0,0 +1,513 @@ +//! Strict CRS metadata resolution for spatial execution. +//! +//! `zarr_inspect` intentionally exposes best-effort raw CRS metadata. Spatial +//! execution instead requires one unambiguous EPSG identifier and a valid +//! same-group `grid_mapping` reference when one is declared. + +use serde_json::{Map, Value}; + +use super::super::{ZarrFdwError, ZarrFdwResult}; + +const GRID_MAPPING: &str = "grid_mapping"; +const GEO_TRANSFORM: &str = "GeoTransform"; + +/// Metadata loaded from the sibling array named by `grid_mapping`. +#[derive(Debug, Clone, Copy)] +pub(crate) struct GridMappingMetadata<'a> { + pub(crate) path: &'a str, + pub(crate) attributes: &'a Map, +} + +/// CRS information accepted by the initial rectilinear spatial engine. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ResolvedCrs { + pub(crate) epsg: i32, + pub(crate) wkt: Option, + pub(crate) geotransform: Option<[f64; 6]>, +} + +#[derive(Debug)] +struct EpsgCandidate { + code: i32, + source: String, +} + +#[derive(Default)] +struct CollectedCrs { + epsg: Vec, + wkt: Vec, + geotransforms: Vec<([f64; 6], String)>, +} + +/// Return the same-group sibling path named by an array's `grid_mapping`. +/// +/// This separate step lets an executor validate the reference before reading +/// the one sibling metadata object needed to resolve it. +pub(crate) fn grid_mapping_sibling_path( + array_path: &str, + array_attributes: &Map, +) -> ZarrFdwResult> { + let Some(value) = array_attributes.get(GRID_MAPPING) else { + return Ok(None); + }; + let reference = value.as_str().ok_or_else(|| { + invalid_crs( + array_path, + "grid_mapping reference must be a string naming a same-group array", + ) + })?; + validate_same_group_name(array_path, reference)?; + Ok(Some(same_group_sibling_path(array_path, reference))) +} + +/// Resolve CRS data from CF grid-mapping, direct array, and group metadata. +/// +/// A declared grid mapping must be supplied and is the authoritative source. +/// Direct array and group EPSG candidates are still checked for conflicts. +pub(crate) fn resolve_crs( + array_path: &str, + array_attributes: &Map, + group_attributes: Option<&Map>, + grid_mapping: Option>, +) -> ZarrFdwResult { + let expected_mapping_path = grid_mapping_sibling_path(array_path, array_attributes)?; + match (expected_mapping_path.as_deref(), grid_mapping.as_ref()) { + (Some(expected), Some(actual)) if actual.path != expected => { + return Err(invalid_crs( + array_path, + format!( + "grid_mapping resolves to sibling array '{expected}', but metadata for '{}' was supplied", + actual.path + ), + )); + } + (Some(expected), None) => { + return Err(invalid_crs( + array_path, + format!("grid_mapping sibling array '{expected}' was not found"), + )); + } + (None, Some(actual)) => { + return Err(invalid_crs( + array_path, + format!( + "CRS metadata for sibling array '{}' was supplied without a grid_mapping reference", + actual.path + ), + )); + } + _ => {} + } + + let mut collected = CollectedCrs::default(); + if let Some(mapping) = grid_mapping { + collect_attributes( + array_path, + &format!("grid_mapping array '{}'", mapping.path), + mapping.attributes, + &mut collected, + )?; + } + collect_attributes( + array_path, + "selected array", + array_attributes, + &mut collected, + )?; + if let Some(attributes) = group_attributes { + collect_attributes( + array_path, + "same-group metadata", + attributes, + &mut collected, + )?; + } + + let Some(first) = collected.epsg.first() else { + let message = if collected.wkt.is_empty() { + "no supported EPSG identifier was found in grid-mapping, array, or group metadata" + } else { + "CRS metadata contains WKT but no supported EPSG identifier; WKT-only CRS resolution is not supported yet" + }; + return Err(invalid_crs(array_path, message)); + }; + if let Some(conflict) = collected + .epsg + .iter() + .skip(1) + .find(|candidate| candidate.code != first.code) + { + return Err(invalid_crs( + array_path, + format!( + "conflicting EPSG identifiers: EPSG:{} from {} and EPSG:{} from {}", + first.code, first.source, conflict.code, conflict.source + ), + )); + } + + let geotransform = collected + .geotransforms + .first() + .map(|(transform, _)| *transform); + if let Some((first_transform, first_source)) = collected.geotransforms.first() + && let Some((_, conflict_source)) = collected + .geotransforms + .iter() + .skip(1) + .find(|(transform, _)| transform != first_transform) + { + return Err(invalid_crs( + array_path, + format!("conflicting GeoTransform values in {first_source} and {conflict_source}"), + )); + } + + Ok(ResolvedCrs { + epsg: first.code, + wkt: collected.wkt.into_iter().next(), + geotransform, + }) +} + +fn collect_attributes( + array_path: &str, + source: &str, + attributes: &Map, + collected: &mut CollectedCrs, +) -> ZarrFdwResult<()> { + if let Some(value) = attributes.get("epsg_code") { + let code = parse_epsg_value(value).ok_or_else(|| { + invalid_crs( + array_path, + format!("epsg_code in {source} must be a positive integer or 'EPSG:'"), + ) + })?; + collected.epsg.push(EpsgCandidate { + code, + source: format!("{source}.epsg_code"), + }); + } + + if let Some(value) = attributes.get("crs") { + collect_crs_value(array_path, source, value, collected)?; + } + for attribute in ["spatial_ref", "crs_wkt"] { + let Some(value) = attributes.get(attribute) else { + continue; + }; + let text = nonempty_string(array_path, source, attribute, value)?; + if let Some(code) = parse_epsg_label(text) { + collected.epsg.push(EpsgCandidate { + code, + source: format!("{source}.{attribute}"), + }); + } else { + collected.wkt.push(text.to_string()); + } + } + + if let Some(value) = attributes.get(GEO_TRANSFORM) { + let transform = parse_geotransform(array_path, source, value)?; + if transform[2] != 0.0 || transform[4] != 0.0 { + return Err(invalid_crs( + array_path, + format!( + "{GEO_TRANSFORM} in {source} describes a rotated grid; only zero-rotation rectilinear grids are supported" + ), + )); + } + collected + .geotransforms + .push((transform, source.to_string())); + } + Ok(()) +} + +fn collect_crs_value( + array_path: &str, + source: &str, + value: &Value, + collected: &mut CollectedCrs, +) -> ZarrFdwResult<()> { + let text = match value { + Value::String(value) if !value.trim().is_empty() => value.as_str(), + Value::Object(object) if object.get("type").and_then(Value::as_str) == Some("name") => { + object + .get("properties") + .and_then(Value::as_object) + .and_then(|properties| properties.get("name")) + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + invalid_crs( + array_path, + format!( + "crs in {source} with type 'name' must contain a non-empty properties.name string" + ), + ) + })? + } + _ => { + return Err(invalid_crs( + array_path, + format!("crs in {source} must be a non-empty string or a named CRS object"), + )); + } + }; + + if let Some(code) = parse_epsg_label(text) { + collected.epsg.push(EpsgCandidate { + code, + source: format!("{source}.crs"), + }); + } else { + collected.wkt.push(text.to_string()); + } + Ok(()) +} + +fn parse_epsg_value(value: &Value) -> Option { + match value { + Value::Number(number) => number + .as_i64() + .and_then(|value| i32::try_from(value).ok()) + .filter(|value| *value > 0), + Value::String(value) => parse_epsg_label(value), + _ => None, + } +} + +fn parse_epsg_label(value: &str) -> Option { + let (authority, code) = value.split_once(':')?; + if !authority.eq_ignore_ascii_case("EPSG") + || code.is_empty() + || !code.bytes().all(|byte| byte.is_ascii_digit()) + { + return None; + } + code.parse::().ok().filter(|code| *code > 0) +} + +fn nonempty_string<'a>( + array_path: &str, + source: &str, + attribute: &str, + value: &'a Value, +) -> ZarrFdwResult<&'a str> { + value + .as_str() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + invalid_crs( + array_path, + format!("{attribute} in {source} must be a non-empty string"), + ) + }) +} + +fn parse_geotransform(array_path: &str, source: &str, value: &Value) -> ZarrFdwResult<[f64; 6]> { + let values = match value { + Value::String(value) => value + .split_whitespace() + .map(str::parse::) + .collect::, _>>() + .map_err(|_| { + invalid_crs( + array_path, + format!("{GEO_TRANSFORM} in {source} must contain six numbers"), + ) + })?, + Value::Array(values) => values + .iter() + .map(Value::as_f64) + .collect::>>() + .ok_or_else(|| { + invalid_crs( + array_path, + format!("{GEO_TRANSFORM} in {source} must contain six numbers"), + ) + })?, + _ => { + return Err(invalid_crs( + array_path, + format!( + "{GEO_TRANSFORM} in {source} must be a whitespace-separated string or numeric array" + ), + )); + } + }; + let transform: [f64; 6] = values.try_into().map_err(|values: Vec| { + invalid_crs( + array_path, + format!( + "{GEO_TRANSFORM} in {source} has {} values; expected exactly six", + values.len() + ), + ) + })?; + if transform.iter().any(|value| !value.is_finite()) { + return Err(invalid_crs( + array_path, + format!("{GEO_TRANSFORM} in {source} must contain only finite numbers"), + )); + } + Ok(transform) +} + +fn validate_same_group_name(array_path: &str, reference: &str) -> ZarrFdwResult<()> { + if reference.is_empty() + || reference.trim() != reference + || reference.chars().any(char::is_whitespace) + || reference.chars().any(char::is_control) + || reference.contains('/') + || reference.contains('\\') + || matches!(reference, "." | "..") + { + return Err(invalid_crs( + array_path, + format!( + "grid_mapping reference '{reference}' must be a non-empty same-group array name" + ), + )); + } + Ok(()) +} + +fn same_group_sibling_path(array_path: &str, reference: &str) -> String { + array_path + .rsplit_once('/') + .map(|(parent, _)| format!("{parent}/{reference}")) + .unwrap_or_else(|| reference.to_string()) +} + +fn invalid_crs(array: &str, message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidCrs { + array: array.to_string(), + message: message.into(), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn attributes(value: Value) -> Map { + value.as_object().unwrap().clone() + } + + #[test] + fn resolves_consistent_grid_mapping_array_and_group_candidates() { + let array = attributes(json!({ + "grid_mapping": "spatial_ref", + "crs": {"type": "name", "properties": {"name": "epsg:3857"}} + })); + let mapping = attributes(json!({ + "epsg_code": "EPSG:3857", + "spatial_ref": "PROJCRS[\"WGS 84 / Pseudo-Mercator\"]", + "GeoTransform": "100 10 0 50 0 -10" + })); + let group = attributes(json!({"crs": "EPSG:3857"})); + + let resolved = resolve_crs( + "nested/raw", + &array, + Some(&group), + Some(GridMappingMetadata { + path: "nested/spatial_ref", + attributes: &mapping, + }), + ) + .unwrap(); + + assert_eq!(resolved.epsg, 3857); + assert_eq!( + resolved.wkt.as_deref(), + Some("PROJCRS[\"WGS 84 / Pseudo-Mercator\"]") + ); + assert_eq!( + resolved.geotransform, + Some([100.0, 10.0, 0.0, 50.0, 0.0, -10.0]) + ); + } + + #[test] + fn root_array_grid_mapping_resolves_to_root_sibling() { + let array = attributes(json!({"grid_mapping": "crs"})); + assert_eq!( + grid_mapping_sibling_path("temperature", &array).unwrap(), + Some("crs".to_string()) + ); + } + + #[test] + fn rejects_invalid_or_mismatched_grid_mapping_references() { + for reference in [json!(7), json!(""), json!("../crs"), json!("spatial ref")] { + let array = attributes(json!({"grid_mapping": reference})); + assert!(grid_mapping_sibling_path("nested/raw", &array).is_err()); + } + + let array = attributes(json!({"grid_mapping": "spatial_ref"})); + let mapping = attributes(json!({"epsg_code": 3857})); + assert!( + resolve_crs( + "nested/raw", + &array, + None, + Some(GridMappingMetadata { + path: "other/spatial_ref", + attributes: &mapping, + }), + ) + .is_err() + ); + assert!(resolve_crs("nested/raw", &array, None, None).is_err()); + } + + #[test] + fn rejects_conflicting_epsg_candidates() { + let array = attributes(json!({"crs": "EPSG:4326"})); + let group = attributes(json!({"epsg_code": 3857})); + let error = resolve_crs("nested/raw", &array, Some(&group), None).unwrap_err(); + assert!(error.to_string().contains("conflicting EPSG identifiers")); + } + + #[test] + fn rejects_wkt_only_or_invalid_epsg_metadata() { + let wkt = attributes(json!({"crs_wkt": "GEOGCRS[\"WGS 84\"]"})); + assert!(resolve_crs("nested/raw", &wkt, None, None).is_err()); + + for value in [json!(0), json!(-1), json!("EPSG:0"), json!("EPSG:abc")] { + let invalid = attributes(json!({"epsg_code": value})); + assert!(resolve_crs("nested/raw", &invalid, None, None).is_err()); + } + } + + #[test] + fn validates_geotransform_shape_finiteness_rotation_and_conflicts() { + for value in [ + json!("0 1 0 2 0"), + json!("0 1 NaN 2 0 -1"), + json!([0, 1, 0.5, 2, 0, -1]), + ] { + let invalid = attributes(json!({ + "epsg_code": 3857, + "GeoTransform": value + })); + assert!(resolve_crs("nested/raw", &invalid, None, None).is_err()); + } + + let array = attributes(json!({ + "epsg_code": 3857, + "GeoTransform": [0, 1, 0, 2, 0, -1] + })); + let group = attributes(json!({ + "epsg_code": 3857, + "GeoTransform": [10, 1, 0, 2, 0, -1] + })); + let error = resolve_crs("nested/raw", &array, Some(&group), None).unwrap_err(); + assert!(error.to_string().contains("conflicting GeoTransform")); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/grid.rs b/wrappers/src/fdw/zarr_fdw/spatial/grid.rs new file mode 100644 index 000000000..4e70c928c --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/grid.rs @@ -0,0 +1,667 @@ +use super::super::chunk::IndexBounds; +use super::super::dataset::DimensionRole; +use super::super::{ZarrFdwError, ZarrFdwResult}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum HorizontalGridKind { + Projected, + Geographic, +} + +/// Array-axis positions for one unambiguous horizontal coordinate pair. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct HorizontalAxes { + pub(crate) x: usize, + pub(crate) y: usize, + pub(crate) kind: HorizontalGridKind, +} + +pub(crate) fn discover_horizontal_axes_from_roles( + roles: impl IntoIterator, +) -> ZarrFdwResult { + let roles = roles.into_iter().collect::>(); + let positions = |wanted| { + roles + .iter() + .enumerate() + .filter_map(|(axis, role)| (*role == wanted).then_some(axis)) + .collect::>() + }; + let projected_x = positions(DimensionRole::SpatialX); + let projected_y = positions(DimensionRole::SpatialY); + let longitude = positions(DimensionRole::Longitude); + let latitude = positions(DimensionRole::Latitude); + let horizontal_count = projected_x + .len() + .saturating_add(projected_y.len()) + .saturating_add(longitude.len()) + .saturating_add(latitude.len()); + + let projected = projected_x.len() == 1 + && projected_y.len() == 1 + && longitude.is_empty() + && latitude.is_empty(); + if projected { + return Ok(HorizontalAxes { + x: projected_x[0], + y: projected_y[0], + kind: HorizontalGridKind::Projected, + }); + } + + let geographic = longitude.len() == 1 + && latitude.len() == 1 + && projected_x.is_empty() + && projected_y.is_empty(); + if geographic { + return Ok(HorizontalAxes { + x: longitude[0], + y: latitude[0], + kind: HorizontalGridKind::Geographic, + }); + } + + Err(ZarrFdwError::InvalidMetadata(format!( + "spatial execution requires exactly one compatible horizontal pair (SpatialX/SpatialY or Longitude/Latitude), found {horizontal_count} horizontal role assignments" + ))) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AxisOrder { + Ascending, + Descending, + Unordered, +} + +/// Classify strict axis direction. Empty and singleton axes are ascending; +/// repeated values are unordered because they do not have a strict direction. +pub(crate) fn axis_order(values: &[f64]) -> AxisOrder { + if values.windows(2).all(|pair| pair[0] < pair[1]) { + AxisOrder::Ascending + } else if values.windows(2).all(|pair| pair[0] > pair[1]) { + AxisOrder::Descending + } else { + AxisOrder::Unordered + } +} + +fn validate_finite_values(name: &str, values: &[f64]) -> ZarrFdwResult<()> { + if let Some((index, value)) = values + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{name}' contains non-finite value {value} at index {index}" + ))); + } + Ok(()) +} + +fn validate_finite_target(name: &str, value: f64) -> ZarrFdwResult<()> { + if !value.is_finite() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial {name} must be finite, got {value}" + ))); + } + Ok(()) +} + +/// Return an inclusive, conservative array-index interval containing every +/// center in `[lo, hi]`. Unordered axes may over-select between the first and +/// last matching global indexes, but never omit a matching center. +pub(crate) fn inclusive_center_bounds( + values: &[f64], + lo: f64, + hi: f64, +) -> ZarrFdwResult> { + validate_finite_values("axis", values)?; + validate_finite_target("lower bound", lo)?; + validate_finite_target("upper bound", hi)?; + if values.is_empty() || lo > hi { + return Ok(None); + } + + let bounds = match axis_order(values) { + AxisOrder::Ascending => { + let start = values.partition_point(|value| *value < lo); + let end_exclusive = values.partition_point(|value| *value <= hi); + (start < end_exclusive).then(|| IndexBounds { + start, + end: end_exclusive - 1, + }) + } + AxisOrder::Descending => { + let start = values.partition_point(|value| *value > hi); + let end_exclusive = values.partition_point(|value| *value >= lo); + (start < end_exclusive).then(|| IndexBounds { + start, + end: end_exclusive - 1, + }) + } + AxisOrder::Unordered => { + let mut matches = values + .iter() + .enumerate() + .filter_map(|(index, value)| (lo <= *value && *value <= hi).then_some(index)); + matches.next().map(|start| IndexBounds { + start, + end: matches.next_back().unwrap_or(start), + }) + } + }; + Ok(bounds) +} + +/// Locate an exactly equal center, preferring the lowest global array index. +pub(crate) fn exact_center_index(values: &[f64], target: f64) -> ZarrFdwResult> { + validate_finite_values("axis", values)?; + validate_finite_target("coordinate", target)?; + Ok(values.iter().position(|value| *value == target)) +} + +/// Locate the nearest center, preferring the lowest global array index on a +/// distance tie. The linear pass is also correct for unordered coordinates. +pub(crate) fn nearest_center_index(values: &[f64], target: f64) -> ZarrFdwResult> { + validate_finite_values("axis", values)?; + validate_finite_target("coordinate", target)?; + let mut best: Option<(usize, f64)> = None; + for (index, value) in values.iter().enumerate() { + let distance = (*value - target).abs(); + if best.is_none_or(|(_, best_distance)| distance < best_distance) { + best = Some((index, distance)); + } + } + Ok(best.map(|(index, _)| index)) +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct GridCell { + /// Indexes in the original two-dimensional array axis order. + pub(crate) array_indices: [usize; 2], + pub(crate) x_index: usize, + pub(crate) y_index: usize, + pub(crate) x: f64, + pub(crate) y: f64, + pub(crate) distance: f64, +} + +/// One resolved horizontal cell independent of the selected array's rank. +/// +/// Native array-axis placement remains with `HorizontalAxes`; this value keeps +/// only the semantic X/Y indexes and coordinates needed by operation results +/// and exact PostGIS masking. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct HorizontalCell { + pub(crate) x_index: usize, + pub(crate) y_index: usize, + pub(crate) x: f64, + pub(crate) y: f64, + pub(crate) distance: f64, +} + +impl From for HorizontalCell { + fn from(cell: GridCell) -> Self { + Self { + x_index: cell.x_index, + y_index: cell.y_index, + x: cell.x, + y: cell.y, + distance: cell.distance, + } + } +} + +/// Finite transformed geometry bounds in the grid's coordinate reference +/// system. Bounds are inclusive because exact polygon masking happens after +/// this conservative center-window selection. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct CoordinateEnvelope { + x_min: f64, + y_min: f64, + x_max: f64, + y_max: f64, +} + +impl CoordinateEnvelope { + pub(crate) fn new(x_min: f64, y_min: f64, x_max: f64, y_max: f64) -> ZarrFdwResult { + for (name, value) in [ + ("x minimum", x_min), + ("y minimum", y_min), + ("x maximum", x_max), + ("y maximum", y_max), + ] { + if !value.is_finite() { + return Err(ZarrFdwError::InvalidGeometry(format!( + "transformed envelope {name} must be finite, got {value}" + ))); + } + } + if x_min > x_max || y_min > y_max { + return Err(ZarrFdwError::InvalidGeometry(format!( + "spatial envelope minimums must not exceed maximums, got ({x_min}, {y_min})..({x_max}, {y_max})" + ))); + } + Ok(Self { + x_min, + y_min, + x_max, + y_max, + }) + } +} + +/// Checked index bounds for one conservative rank-2 center window. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct GridWindowPlan { + bounds: [IndexBounds; 2], + total: usize, +} + +impl GridWindowPlan { + fn new(bounds: [IndexBounds; 2]) -> ZarrFdwResult { + let total = checked_window_cell_count(&bounds)?; + Ok(Self { bounds, total }) + } + + pub(crate) fn bounds(&self) -> [IndexBounds; 2] { + self.bounds + } + + pub(crate) fn total_cells(&self) -> usize { + self.total + } +} + +fn checked_window_cell_count(bounds: &[IndexBounds; 2]) -> ZarrFdwResult { + bounds.iter().try_fold(1usize, |total, bounds| { + let extent = bounds + .end + .checked_sub(bounds.start) + .and_then(|span| span.checked_add(1)) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial window extent exceeds this platform's index capacity".to_string(), + ) + })?; + total.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial window cell count exceeds this platform's index capacity".to_string(), + ) + }) + }) +} + +/// A two-dimensional rectilinear grid whose semantic X/Y axes may occur in +/// either array order. +#[derive(Debug)] +pub(crate) struct RectilinearGrid<'a> { + axes: HorizontalAxes, + x: &'a [f64], + y: &'a [f64], +} + +impl<'a> RectilinearGrid<'a> { + pub(crate) fn new(axes: HorizontalAxes, x: &'a [f64], y: &'a [f64]) -> ZarrFdwResult { + if axes.x > 1 || axes.y > 1 || axes.x == axes.y { + return Err(ZarrFdwError::InvalidMetadata(format!( + "rectilinear grid requires two distinct rank-2 array axes, got x axis {} and y axis {}", + axes.x, axes.y + ))); + } + if x.is_empty() || y.is_empty() { + return Err(ZarrFdwError::InvalidMetadata( + "rectilinear grid coordinate axes must not be empty".to_string(), + )); + } + validate_finite_values("x", x)?; + validate_finite_values("y", y)?; + Ok(Self { axes, x, y }) + } + + pub(crate) fn axes(&self) -> HorizontalAxes { + self.axes + } + + pub(crate) fn exact(&self, x: f64, y: f64) -> ZarrFdwResult> { + let Some(x_index) = exact_center_index(self.x, x)? else { + return Ok(None); + }; + let Some(y_index) = exact_center_index(self.y, y)? else { + return Ok(None); + }; + self.cell(x_index, y_index, x, y).map(Some) + } + + pub(crate) fn nearest(&self, x: f64, y: f64) -> ZarrFdwResult { + let x_index = nearest_center_index(self.x, x)?.expect("grid x axis is non-empty"); + let y_index = nearest_center_index(self.y, y)?.expect("grid y axis is non-empty"); + self.cell(x_index, y_index, x, y) + } + + pub(crate) fn cell( + &self, + x_index: usize, + y_index: usize, + target_x: f64, + target_y: f64, + ) -> ZarrFdwResult { + validate_finite_target("x coordinate", target_x)?; + validate_finite_target("y coordinate", target_y)?; + let x = *self.x.get(x_index).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "x index {x_index} is outside coordinate length {}", + self.x.len() + )) + })?; + let y = *self.y.get(y_index).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "y index {y_index} is outside coordinate length {}", + self.y.len() + )) + })?; + let mut array_indices = [0; 2]; + array_indices[self.axes.x] = x_index; + array_indices[self.axes.y] = y_index; + Ok(GridCell { + array_indices, + x_index, + y_index, + x, + y, + distance: (x - target_x).hypot(y - target_y), + }) + } + + /// Return X/Y center bounds placed in original array axis order. + pub(crate) fn inclusive_bounds( + &self, + x_lo: f64, + y_lo: f64, + x_hi: f64, + y_hi: f64, + ) -> ZarrFdwResult> { + let Some(x_bounds) = inclusive_center_bounds(self.x, x_lo, x_hi)? else { + return Ok(None); + }; + let Some(y_bounds) = inclusive_center_bounds(self.y, y_lo, y_hi)? else { + return Ok(None); + }; + let mut array_bounds = [x_bounds; 2]; + array_bounds[self.axes.x] = x_bounds; + array_bounds[self.axes.y] = y_bounds; + Ok(Some(array_bounds)) + } + + /// Plan a bounded, lazy C-order stream of candidate centers for one + /// transformed geometry envelope. `None` means no grid center can match. + pub(crate) fn window_plan( + &self, + envelope: CoordinateEnvelope, + ) -> ZarrFdwResult> { + let Some(bounds) = self.inclusive_bounds( + envelope.x_min, + envelope.y_min, + envelope.x_max, + envelope.y_max, + )? + else { + return Ok(None); + }; + GridWindowPlan::new(bounds).map(Some) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn projected(x: usize, y: usize) -> HorizontalAxes { + HorizontalAxes { + x, + y, + kind: HorizontalGridKind::Projected, + } + } + + #[test] + fn discovers_only_complete_compatible_horizontal_pairs() { + assert_eq!( + discover_horizontal_axes_from_roles([ + DimensionRole::SpatialY, + DimensionRole::Time, + DimensionRole::SpatialX, + ]) + .unwrap(), + HorizontalAxes { + x: 2, + y: 0, + kind: HorizontalGridKind::Projected, + } + ); + assert_eq!( + discover_horizontal_axes_from_roles([ + DimensionRole::Latitude, + DimensionRole::Longitude, + ]) + .unwrap(), + HorizontalAxes { + x: 1, + y: 0, + kind: HorizontalGridKind::Geographic, + } + ); + + for roles in [ + vec![], + vec![DimensionRole::SpatialX], + vec![DimensionRole::SpatialX, DimensionRole::Latitude], + vec![ + DimensionRole::SpatialX, + DimensionRole::SpatialX, + DimensionRole::SpatialY, + ], + vec![ + DimensionRole::SpatialX, + DimensionRole::SpatialY, + DimensionRole::Longitude, + DimensionRole::Latitude, + ], + ] { + assert!(discover_horizontal_axes_from_roles(roles).is_err()); + } + } + + #[test] + fn classifies_strict_axis_order() { + assert_eq!(axis_order(&[]), AxisOrder::Ascending); + assert_eq!(axis_order(&[7.0]), AxisOrder::Ascending); + assert_eq!(axis_order(&[1.0, 2.0, 3.0]), AxisOrder::Ascending); + assert_eq!(axis_order(&[3.0, 2.0, 1.0]), AxisOrder::Descending); + assert_eq!(axis_order(&[1.0, 1.0, 2.0]), AxisOrder::Unordered); + assert_eq!(axis_order(&[1.0, 3.0, 2.0]), AxisOrder::Unordered); + } + + #[test] + fn computes_inclusive_bounds_for_every_axis_order() { + assert_eq!( + inclusive_center_bounds(&[10.0, 20.0, 30.0, 40.0], 20.0, 30.0).unwrap(), + Some(IndexBounds { start: 1, end: 2 }) + ); + assert_eq!( + inclusive_center_bounds(&[40.0, 30.0, 20.0, 10.0], 20.0, 30.0).unwrap(), + Some(IndexBounds { start: 1, end: 2 }) + ); + assert_eq!( + inclusive_center_bounds(&[50.0, 10.0, 40.0, 20.0, 30.0], 20.0, 40.0).unwrap(), + Some(IndexBounds { start: 2, end: 4 }) + ); + assert_eq!( + inclusive_center_bounds(&[50.0, 10.0, 40.0, 20.0, 30.0], 21.0, 29.0).unwrap(), + None + ); + assert_eq!( + inclusive_center_bounds(&[10.0, 20.0], 21.0, 20.0).unwrap(), + None + ); + } + + #[test] + fn exact_and_nearest_choose_lowest_global_index() { + assert_eq!( + exact_center_index(&[30.0, 20.0, 20.0, 10.0], 20.0).unwrap(), + Some(1) + ); + assert_eq!( + nearest_center_index(&[50.0, 40.0, 30.0, 20.0], 35.0).unwrap(), + Some(1) + ); + assert_eq!( + nearest_center_index(&[0.0, 100.0, 20.0, 40.0], 30.0).unwrap(), + Some(2) + ); + assert_eq!(nearest_center_index(&[], 1.0).unwrap(), None); + } + + #[test] + fn grid_preserves_array_axis_order_and_reports_distance() { + let y = [50.0, 40.0, 30.0, 20.0, 10.0]; + let x = [100.0, 110.0, 120.0, 130.0, 140.0, 150.0]; + let grid = RectilinearGrid::new(projected(1, 0), &x, &y).unwrap(); + + assert_eq!(axis_order(&x), AxisOrder::Ascending); + assert_eq!(axis_order(&y), AxisOrder::Descending); + let nearest = grid.nearest(121.0, 39.0).unwrap(); + assert_eq!(nearest.array_indices, [1, 2]); + assert_eq!((nearest.x_index, nearest.y_index), (2, 1)); + assert_eq!((nearest.x, nearest.y), (120.0, 40.0)); + assert_eq!(nearest.distance, 2.0_f64.sqrt()); + + let tie = grid.nearest(125.0, 35.0).unwrap(); + assert_eq!(tie.array_indices, [1, 2]); + assert_eq!((tie.x, tie.y), (120.0, 40.0)); + assert_eq!( + grid.exact(120.0, 40.0).unwrap(), + Some(grid.cell(2, 1, 120.0, 40.0).unwrap()) + ); + assert_eq!(grid.exact(121.0, 40.0).unwrap(), None); + } + + #[test] + fn grid_bounds_follow_array_axis_order() { + let grid = RectilinearGrid::new( + projected(0, 1), + &[100.0, 110.0, 120.0, 130.0], + &[40.0, 30.0, 20.0, 10.0], + ) + .unwrap(); + assert_eq!( + grid.inclusive_bounds(110.0, 20.0, 120.0, 30.0).unwrap(), + Some([ + IndexBounds { start: 1, end: 2 }, + IndexBounds { start: 1, end: 2 }, + ]) + ); + } + + #[test] + fn grid_rejects_invalid_rank_axes_values_and_indexes() { + assert!(RectilinearGrid::new(projected(0, 0), &[1.0], &[2.0]).is_err()); + assert!(RectilinearGrid::new(projected(0, 2), &[1.0], &[2.0]).is_err()); + assert!(RectilinearGrid::new(projected(0, 1), &[], &[2.0]).is_err()); + assert!(RectilinearGrid::new(projected(0, 1), &[f64::NAN], &[2.0]).is_err()); + + let grid = RectilinearGrid::new(projected(0, 1), &[1.0], &[2.0]).unwrap(); + assert!(grid.cell(1, 0, 1.0, 2.0).is_err()); + assert!(grid.nearest(f64::INFINITY, 2.0).is_err()); + } + + #[test] + fn envelope_requires_finite_ordered_bounds() { + assert_eq!( + CoordinateEnvelope::new(1.0, 2.0, 3.0, 4.0).unwrap(), + CoordinateEnvelope { + x_min: 1.0, + y_min: 2.0, + x_max: 3.0, + y_max: 4.0, + } + ); + assert!(CoordinateEnvelope::new(3.0, 2.0, 1.0, 4.0).is_err()); + assert!(CoordinateEnvelope::new(1.0, 4.0, 3.0, 2.0).is_err()); + assert!(CoordinateEnvelope::new(f64::NAN, 2.0, 3.0, 4.0).is_err()); + assert!(CoordinateEnvelope::new(1.0, 2.0, f64::INFINITY, 4.0).is_err()); + } + + #[test] + fn window_plan_preserves_array_order_and_checked_count() { + let grid = RectilinearGrid::new( + projected(1, 0), + &[100.0, 110.0, 120.0, 130.0], + &[40.0, 30.0, 20.0, 10.0], + ) + .unwrap(); + let envelope = CoordinateEnvelope::new(110.0, 20.0, 120.0, 30.0).unwrap(); + let plan = grid.window_plan(envelope).unwrap().unwrap(); + assert_eq!( + plan.bounds(), + [ + IndexBounds { start: 1, end: 2 }, + IndexBounds { start: 1, end: 2 }, + ] + ); + assert_eq!(plan.total_cells(), 4); + } + + #[test] + fn unordered_windows_overfetch_conservatively() { + let grid = RectilinearGrid::new( + projected(0, 1), + &[300.0, 999.0, 100.0, 200.0, 400.0], + &[10.0], + ) + .unwrap(); + let envelope = CoordinateEnvelope::new(200.0, 10.0, 400.0, 10.0).unwrap(); + let plan = grid.window_plan(envelope).unwrap().unwrap(); + assert_eq!( + plan.bounds(), + [ + IndexBounds { start: 0, end: 4 }, + IndexBounds { start: 0, end: 0 }, + ] + ); + assert_eq!(plan.total_cells(), 5); + } + + #[test] + fn window_plan_handles_empty_singleton_and_checked_limits() { + let grid = RectilinearGrid::new(projected(0, 1), &[1.0], &[2.0]).unwrap(); + let empty_envelope = CoordinateEnvelope::new(3.0, 2.0, 4.0, 2.0).unwrap(); + assert!(grid.window_plan(empty_envelope).unwrap().is_none()); + + let envelope = CoordinateEnvelope::new(1.0, 2.0, 1.0, 2.0).unwrap(); + let plan = grid.window_plan(envelope).unwrap().unwrap(); + assert_eq!(plan.total_cells(), 1); + + assert!( + checked_window_cell_count(&[ + IndexBounds { + start: 0, + end: usize::MAX, + }, + IndexBounds { start: 0, end: 0 }, + ]) + .is_err() + ); + assert!( + checked_window_cell_count(&[ + IndexBounds { + start: 1, + end: usize::MAX, + }, + IndexBounds { start: 0, end: 1 }, + ]) + .is_err() + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/mod.rs b/wrappers/src/fdw/zarr_fdw/spatial/mod.rs new file mode 100644 index 000000000..10ffaacb3 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/mod.rs @@ -0,0 +1,12 @@ +//! Spatial metadata, optional PostGIS adapters, and rectilinear-grid math. +//! +//! PostGIS adapters and spatial execution build on these helpers, but this +//! module does not depend on geometry types, SPI, storage, or scan state. + +mod catalog; +pub(crate) mod crs; +pub(crate) mod grid; +mod point; +pub(crate) mod postgis; +mod temporal; +mod zonal; diff --git a/wrappers/src/fdw/zarr_fdw/spatial/point.rs b/wrappers/src/fdw/zarr_fdw/spatial/point.rs new file mode 100644 index 000000000..63fb60c01 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/point.rs @@ -0,0 +1,302 @@ +//! Read-only point sampling over one rank-2 rectilinear Zarr foreign table. + +use pgrx::pg_sys::panic::{ErrorReport, ErrorReportable}; +use pgrx::prelude::*; +use supabase_wrappers::prelude::{Cell, ForeignDataWrapper, Row}; + +use super::super::selectors::{DimensionSelectors, OPT_DIMENSION_SELECTORS}; +use super::super::zarr_fdw::ZarrFdw; +use super::super::{ZarrFdwError, ZarrFdwResult}; +use super::catalog::{load_zarr_foreign_table, load_zarr_foreign_table_with_selectors}; +use super::grid::HorizontalCell; +use super::postgis::PostgisCatalog; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SampleMethod { + Exact, + Nearest, +} + +impl SampleMethod { + fn parse(value: &str) -> ZarrFdwResult { + match value { + "exact" => Ok(Self::Exact), + "nearest" => Ok(Self::Nearest), + _ => Err(ZarrFdwError::InvalidGeometry(format!( + "point sampling method must be 'exact' or 'nearest', got '{value}'" + ))), + } + } +} + +#[derive(Debug, PartialEq)] +struct SampleRow { + x: f64, + y: f64, + value: Option, + x_index: i64, + y_index: i64, + coordinate_distance: f64, + srid: i32, +} + +impl SampleRow { + fn sql_row(self) -> (f64, f64, Option, i64, i64, f64, i32) { + ( + self.x, + self.y, + self.value, + self.x_index, + self.y_index, + self.coordinate_distance, + self.srid, + ) + } +} + +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace, volatile, parallel_unsafe)] +fn zarr_sample( + foreign_table: &str, + point_ewkb: &[u8], + method: default!(&str, "'nearest'"), +) -> TableIterator< + 'static, + ( + name!(x, f64), + name!(y, f64), + name!(value, Option), + name!(x_index, i64), + name!(y_index, i64), + name!(coordinate_distance, f64), + name!(srid, i32), + ), +> { + let rows = sample_foreign_table(foreign_table, point_ewkb, method, None) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(SampleRow::sql_row)) +} + +#[allow(clippy::type_complexity)] +#[pg_extern(name = "zarr_sample", create_or_replace, volatile, parallel_unsafe)] +fn zarr_sample_with_selectors( + foreign_table: &str, + point_ewkb: &[u8], + method: &str, + dimension_selectors: &str, +) -> TableIterator< + 'static, + ( + name!(x, f64), + name!(y, f64), + name!(value, Option), + name!(x_index, i64), + name!(y_index, i64), + name!(coordinate_distance, f64), + name!(srid, i32), + ), +> { + let rows = sample_foreign_table(foreign_table, point_ewkb, method, Some(dimension_selectors)) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(SampleRow::sql_row)) +} + +fn sample_foreign_table( + foreign_table: &str, + point_ewkb: &[u8], + method: &str, + call_selectors: Option<&str>, +) -> ZarrFdwResult> { + let method = SampleMethod::parse(method)?; + let selector_aware = call_selectors.is_some(); + // Fail before remote metadata reads when the optional spatial dependency is + // absent or incomplete. + let postgis = PostgisCatalog::require()?; + let table = if selector_aware { + load_zarr_foreign_table_with_selectors(foreign_table)? + } else { + load_zarr_foreign_table(foreign_table)? + }; + let call_selectors = match call_selectors { + Some(raw) => { + DimensionSelectors::parse( + table + .options + .get(OPT_DIMENSION_SELECTORS) + .map(String::as_str), + )?; + DimensionSelectors::parse(Some(raw))? + } + None => DimensionSelectors::default(), + }; + let mut fdw = ZarrFdw::new(table.server)?; + fdw.set_call_dimension_selectors(call_selectors)?; + let begin = >::begin_scan( + &mut fdw, + &[], + &table.columns, + &[], + &None, + &table.options, + ); + if let Err(error) = begin { + let _ = >::end_scan(&mut fdw); + return Err(error); + } + + let result = sample_prepared_scan(&mut fdw, &postgis, point_ewkb, method, selector_aware); + let cleanup = >::end_scan(&mut fdw); + match (result, cleanup) { + (Ok(value), Ok(())) => Ok(value), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +fn sample_prepared_scan( + fdw: &mut ZarrFdw, + postgis: &PostgisCatalog, + point_ewkb: &[u8], + method: SampleMethod, + selector_aware: bool, +) -> ZarrFdwResult> { + // Validate the table contract before exact lookup can return no rows. + let value_column = fdw.spatial_value_column()?.to_string(); + let crs = fdw.resolved_spatial_crs()?; + let point = postgis.transform_ewkb_point(point_ewkb, fdw.spatial_array_path(), crs.epsg)?; + let (axes, cell) = if selector_aware { + let axes = fdw.spatial_horizontal_axes()?; + if !fdw.apply_spatial_dimension_selectors("zarr_sample", &[axes.x, axes.y])? { + return Ok(Vec::new()); + } + match method { + SampleMethod::Exact => { + let Some((axes, cell)) = fdw.spatial_exact_horizontal_cell(point.x, point.y)? + else { + return Ok(Vec::new()); + }; + (axes, cell) + } + SampleMethod::Nearest => fdw.spatial_nearest_horizontal_cell(point.x, point.y)?, + } + } else { + let grid = fdw.rectilinear_grid()?; + let axes = grid.axes(); + let cell = match method { + SampleMethod::Exact => grid.exact(point.x, point.y)?, + SampleMethod::Nearest => Some(grid.nearest(point.x, point.y)?), + }; + let Some(cell) = cell else { + return Ok(Vec::new()); + }; + fdw.restrict_to_spatial_cell(cell.array_indices)?; + (axes, HorizontalCell::from(cell)) + }; + if selector_aware { + fdw.restrict_to_horizontal_cell(axes, cell.x_index, cell.y_index)?; + } + + let mut row = Row::new(); + if >::iter_scan(fdw, &mut row)?.is_none() { + return Err(ZarrFdwError::InvalidMetadata( + "point sampling selected a logical cell but the Zarr executor returned no row" + .to_string(), + )); + } + let value = row + .iter() + .find(|(name, _)| name.as_str() == value_column) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "point sampling result did not contain value column '{value_column}'" + )) + })? + .1 + .as_ref() + .map(numeric_cell_to_f64) + .transpose()?; + + row.clear(); + if >::iter_scan(fdw, &mut row)?.is_some() { + return Err(ZarrFdwError::InvalidMetadata( + "point sampling selected more than one logical cell".to_string(), + )); + } + Ok(vec![sample_row(cell, value, crs.epsg)?]) +} + +fn sample_row(cell: HorizontalCell, value: Option, srid: i32) -> ZarrFdwResult { + Ok(SampleRow { + x: cell.x, + y: cell.y, + value, + x_index: i64::try_from(cell.x_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "spatial x index exceeds PostgreSQL bigint range".to_string(), + ) + })?, + y_index: i64::try_from(cell.y_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "spatial y index exceeds PostgreSQL bigint range".to_string(), + ) + })?, + coordinate_distance: cell.distance, + srid, + }) +} + +pub(super) fn numeric_cell_to_f64(cell: &Cell) -> ZarrFdwResult { + let value = match cell { + Cell::I8(value) => f64::from(*value), + Cell::I16(value) => f64::from(*value), + Cell::I32(value) => f64::from(*value), + Cell::I64(value) => { + let converted = *value as f64; + if converted as i128 != i128::from(*value) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "point sample bigint value {value} cannot be represented exactly as double precision" + ))); + } + converted + } + Cell::F32(value) => f64::from(*value), + Cell::F64(value) => *value, + Cell::Numeric(value) => value.to_string().parse::().map_err(|_| { + ZarrFdwError::InvalidMetadata( + "point sample numeric value cannot be represented as double precision".to_string(), + ) + })?, + _ => { + return Err(ZarrFdwError::InvalidMetadata( + "point sampling supports only numeric Zarr value columns".to_string(), + )); + } + }; + Ok(value) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_only_documented_sampling_methods() { + assert_eq!(SampleMethod::parse("exact").unwrap(), SampleMethod::Exact); + assert_eq!( + SampleMethod::parse("nearest").unwrap(), + SampleMethod::Nearest + ); + assert!(SampleMethod::parse("Nearest").is_err()); + assert!(SampleMethod::parse("").is_err()); + } + + #[test] + fn widens_supported_numeric_cells_without_silent_bigint_rounding() { + assert_eq!(numeric_cell_to_f64(&Cell::F32(1.5)).unwrap(), 1.5); + assert_eq!(numeric_cell_to_f64(&Cell::I32(42)).unwrap(), 42.0); + assert!(numeric_cell_to_f64(&Cell::I64(9_007_199_254_740_993)).is_err()); + assert!(numeric_cell_to_f64(&Cell::String("42".to_string())).is_err()); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/postgis.rs b/wrappers/src/fdw/zarr_fdw/spatial/postgis.rs new file mode 100644 index 000000000..6dcbcc42e --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/postgis.rs @@ -0,0 +1,1101 @@ +//! Optional, query-local PostGIS catalog and geometry transformation adapter. +//! +//! The core extension has no PostGIS link-time dependency. This module first +//! proves that the geometry type and every called function are owned by the +//! installed `postgis` extension, then invokes only schema-qualified functions +//! with parameterized values. + +use pgrx::{ + Spi, + pg_sys::{PgTryBuilder, errcodes::PgSqlErrorCode}, + spi::quote_identifier, +}; + +use super::super::{ZarrFdwError, ZarrFdwResult}; + +const MAX_EWKB_BYTES: usize = 8 * 1024 * 1024; +pub(crate) const MAX_COVERAGE_CANDIDATES: usize = 65_536; + +const DISCOVER_POSTGIS_SQL: &str = r#" + SELECT extension.oid::bigint AS extension_oid, + namespace.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS extension + JOIN pg_catalog.pg_namespace AS namespace + ON namespace.oid = extension.extnamespace + WHERE extension.extname = 'postgis' +"#; + +// Verify extension ownership as well as names/signatures. A same-name object in +// the PostGIS schema is not accepted unless pg_depend records it as an extension +// member. Built-in argument/result OIDs are stable PostgreSQL catalog OIDs. +const VERIFY_POSTGIS_MEMBERS_SQL: &str = r#" + WITH owned_geometry AS ( + SELECT type.oid AS geometry_oid + FROM pg_catalog.pg_type AS type + JOIN pg_catalog.pg_extension AS extension + ON extension.oid = $1::bigint::oid + AND type.typnamespace = extension.extnamespace + JOIN pg_catalog.pg_depend AS dependency + ON dependency.classid = 'pg_catalog.pg_type'::pg_catalog.regclass + AND dependency.objid = type.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.refobjid = extension.oid + AND dependency.deptype = 'e' + WHERE type.typname = 'geometry' + ), + owned_box3d AS ( + SELECT type.oid AS box3d_oid + FROM pg_catalog.pg_type AS type + JOIN pg_catalog.pg_extension AS extension + ON extension.oid = $1::bigint::oid + AND type.typnamespace = extension.extnamespace + JOIN pg_catalog.pg_depend AS dependency + ON dependency.classid = 'pg_catalog.pg_type'::pg_catalog.regclass + AND dependency.objid = type.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.refobjid = extension.oid + AND dependency.deptype = 'e' + WHERE type.typname = 'box3d' + ), + owned_functions AS ( + SELECT procedure.proname, + procedure.pronargs, + procedure.proargtypes, + procedure.prorettype, + geometry.geometry_oid, + box3d.box3d_oid + FROM pg_catalog.pg_proc AS procedure + JOIN pg_catalog.pg_extension AS extension + ON extension.oid = $1::bigint::oid + AND procedure.pronamespace = extension.extnamespace + CROSS JOIN owned_geometry AS geometry + CROSS JOIN owned_box3d AS box3d + JOIN pg_catalog.pg_depend AS dependency + ON dependency.classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND dependency.objid = procedure.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.refobjid = extension.oid + AND dependency.deptype = 'e' + ), + owned_spatial_ref_sys AS ( + SELECT relation.oid + FROM pg_catalog.pg_class AS relation + JOIN pg_catalog.pg_extension AS extension + ON extension.oid = $1::bigint::oid + AND relation.relnamespace = extension.extnamespace + JOIN pg_catalog.pg_depend AS dependency + ON dependency.classid = 'pg_catalog.pg_class'::pg_catalog.regclass + AND dependency.objid = relation.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.refobjid = extension.oid + AND dependency.deptype = 'e' + WHERE relation.relname = 'spatial_ref_sys' + ) + SELECT EXISTS (SELECT 1 FROM owned_geometry) + AND EXISTS (SELECT 1 FROM owned_box3d) + AND EXISTS (SELECT 1 FROM owned_spatial_ref_sys) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_geomfromewkb' AND pronargs = 1 + AND proargtypes[0] = 17 AND prorettype = geometry_oid + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_srid' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 23 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_geometrytype' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 25 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_isempty' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 16 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_isvalid' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 16 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_ndims' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 21 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_transform' AND pronargs = 2 + AND proargtypes[0] = geometry_oid AND proargtypes[1] = 23 + AND prorettype = geometry_oid + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_x' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_y' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'box3d' AND pronargs = 1 + AND proargtypes[0] = geometry_oid AND prorettype = box3d_oid + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_xmin' AND pronargs = 1 + AND proargtypes[0] = box3d_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_ymin' AND pronargs = 1 + AND proargtypes[0] = box3d_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_xmax' AND pronargs = 1 + AND proargtypes[0] = box3d_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_ymax' AND pronargs = 1 + AND proargtypes[0] = box3d_oid AND prorettype = 701 + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_makepoint' AND pronargs = 2 + AND proargtypes[0] = 701 AND proargtypes[1] = 701 + AND prorettype = geometry_oid + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_setsrid' AND pronargs = 2 + AND proargtypes[0] = geometry_oid AND proargtypes[1] = 23 + AND prorettype = geometry_oid + ) + AND EXISTS ( + SELECT 1 FROM owned_functions + WHERE proname = 'st_covers' AND pronargs = 2 + AND proargtypes[0] = geometry_oid AND proargtypes[1] = geometry_oid + AND prorettype = 16 + ) AS valid +"#; + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct TransformedPoint { + pub(crate) x: f64, + pub(crate) y: f64, + pub(crate) source_srid: i32, + pub(crate) target_srid: i32, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +#[allow(dead_code)] +pub(crate) struct GeometryEnvelope { + pub(crate) xmin: f64, + pub(crate) ymin: f64, + pub(crate) xmax: f64, + pub(crate) ymax: f64, + pub(crate) source_srid: i32, + pub(crate) target_srid: i32, +} + +/// A validated, query-local reference to the installed PostGIS extension. +#[derive(Debug, Clone)] +pub(crate) struct PostgisCatalog { + schema: String, +} + +#[derive(Debug)] +struct PointDescription { + source_srid: i32, + geometry_type: String, + dimensions: i32, + empty: bool, + x: Option, + y: Option, + transformed_srid: Option, +} + +#[derive(Debug)] +struct PolygonDescription { + source_srid: i32, + source_srid_known: bool, + geometry_type: String, + dimensions: i32, + empty: bool, + valid: bool, + transformed_srid: Option, +} + +#[derive(Debug)] +struct EnvelopeDescription { + polygon: PolygonDescription, + xmin: Option, + ymin: Option, + xmax: Option, + ymax: Option, +} + +impl PostgisCatalog { + /// Discover and authenticate the installed PostGIS catalog. `None` means + /// the extension is not installed; a present but incomplete or shadowed + /// installation is rejected. + pub(crate) fn discover() -> ZarrFdwResult> { + let discovered = Spi::connect(|client| { + let mut rows = client.select(DISCOVER_POSTGIS_SQL, None, &[])?; + let Some(row) = rows.next() else { + return Ok::<_, pgrx::spi::Error>(None); + }; + Ok::<_, pgrx::spi::Error>(Some(( + row.get_by_name::("extension_oid")? + .expect("pg_extension.oid is not null"), + row.get_by_name::("schema_name")? + .expect("pg_namespace.nspname is not null"), + ))) + })?; + let Some((extension_oid, schema)) = discovered else { + return Ok(None); + }; + + let members_valid = + Spi::get_one_with_args::(VERIFY_POSTGIS_MEMBERS_SQL, &[extension_oid.into()])? + .unwrap_or(false); + if !members_valid { + return Err(ZarrFdwError::PostgisUnavailable( + "the installed extension does not own the required spatial types, spatial_ref_sys catalog, and geometry functions" + .to_string(), + )); + } + Ok(Some(Self { schema })) + } + + pub(crate) fn require() -> ZarrFdwResult { + Self::discover()?.ok_or_else(|| { + ZarrFdwError::PostgisUnavailable("the postgis extension is not installed".to_string()) + }) + } + + /// Parse one EWKB point, transform it to `target_epsg`, and return only + /// built-in numeric values. Geometry datums never cross the Rust boundary. + pub(crate) fn transform_ewkb_point( + &self, + ewkb: &[u8], + array_path: &str, + target_epsg: i32, + ) -> ZarrFdwResult { + validate_ewkb_size(ewkb)?; + if target_epsg <= 0 { + return Err(invalid_crs( + array_path, + format!("target EPSG code must be positive, got {target_epsg}"), + )); + } + self.require_known_srid(array_path, target_epsg)?; + + let sql = point_transform_sql(&self.schema); + // PostGIS 3.x is not relocatable and some transformation internals + // resolve its extension-owned spatial_ref_sys through search_path. Set + // the catalog-discovered schema only for this transaction-local call, + // then restore the caller's setting on every ordinary Result path. + let prior_search_path = + Spi::get_one::("SELECT pg_catalog.current_setting('search_path')")? + .unwrap_or_default(); + let operation_search_path = format!("{}, pg_catalog", quote_identifier(&self.schema)); + Spi::get_one_with_args::( + "SELECT pg_catalog.set_config('search_path', $1, true)", + &[operation_search_path.into()], + )?; + let description = Spi::connect(|client| { + let mut rows = client + .select(&sql, Some(1), &[ewkb.to_vec().into(), target_epsg.into()]) + .map_err(|_| { + ZarrFdwError::InvalidGeometry( + "PostGIS could not parse or transform the supplied EWKB".to_string(), + ) + })?; + let row = rows.next().ok_or_else(|| { + ZarrFdwError::InvalidGeometry( + "PostGIS returned no description for the supplied EWKB".to_string(), + ) + })?; + Ok::<_, ZarrFdwError>(PointDescription { + source_srid: row + .get_by_name::("source_srid")? + .expect("ST_SRID result is not null"), + geometry_type: row + .get_by_name::("geometry_type")? + .expect("ST_GeometryType result is not null"), + dimensions: row + .get_by_name::("dimensions")? + .expect("ST_NDims result is not null"), + empty: row + .get_by_name::("is_empty")? + .expect("ST_IsEmpty result is not null"), + x: row.get_by_name::("x")?, + y: row.get_by_name::("y")?, + transformed_srid: row.get_by_name::("transformed_srid")?, + }) + }); + let restore = Spi::get_one_with_args::( + "SELECT pg_catalog.set_config('search_path', $1, true)", + &[prior_search_path.into()], + ); + match (description, restore) { + (Ok(description), Ok(_)) => validate_point_description(description, target_epsg), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error.into()), + } + } + + /// Validate and transform a polygonal EWKB geometry, then return its + /// finite target-CRS envelope using only built-in Rust scalar types. + #[allow(dead_code)] + pub(crate) fn transform_ewkb_geometry_envelope( + &self, + ewkb: &[u8], + array_path: &str, + target_epsg: i32, + ) -> ZarrFdwResult { + validate_ewkb_size(ewkb)?; + if target_epsg <= 0 { + return Err(invalid_crs( + array_path, + format!("target EPSG code must be positive, got {target_epsg}"), + )); + } + self.require_known_srid(array_path, target_epsg)?; + + let sql = polygon_envelope_sql(&self.schema); + let description = self.with_extension_search_path(|| { + PgTryBuilder::new(|| { + Spi::connect(|client| { + let mut rows = client.select( + &sql, + Some(1), + &[ewkb.to_vec().into(), target_epsg.into()], + )?; + let row = rows.next().ok_or_else(|| { + ZarrFdwError::InvalidGeometry( + "PostGIS returned no description for the supplied polygon".to_string(), + ) + })?; + Ok::<_, ZarrFdwError>(EnvelopeDescription { + polygon: PolygonDescription { + source_srid: row + .get_by_name::("source_srid")? + .expect("ST_SRID result is not null"), + source_srid_known: row + .get_by_name::("source_srid_known")? + .expect("source SRID catalog check is not null"), + geometry_type: row + .get_by_name::("geometry_type")? + .expect("ST_GeometryType result is not null"), + dimensions: row + .get_by_name::("dimensions")? + .expect("ST_NDims result is not null"), + empty: row + .get_by_name::("is_empty")? + .expect("ST_IsEmpty result is not null"), + valid: row + .get_by_name::("is_valid")? + .expect("ST_IsValid result is not null"), + transformed_srid: row.get_by_name::("transformed_srid")?, + }, + xmin: row.get_by_name::("xmin")?, + ymin: row.get_by_name::("ymin")?, + xmax: row.get_by_name::("xmax")?, + ymax: row.get_by_name::("ymax")?, + }) + }) + }) + // PostGIS reports malformed EWKB and invalid transforms as XX000. + // Catch only that code so cancellation, permissions, and all other + // PostgreSQL errors retain their native behavior. + .catch_when(PgSqlErrorCode::ERRCODE_INTERNAL_ERROR, |_| { + Err(ZarrFdwError::InvalidGeometry( + "PostGIS could not parse or transform the supplied polygon".to_string(), + )) + }) + .execute() + })?; + validate_envelope_description(description, array_path, target_epsg) + } + + /// Test a bounded batch of target-CRS cell centers against one polygonal + /// EWKB geometry. One SPI query evaluates the whole batch in input order. + #[allow(dead_code)] + pub(crate) fn covers_ewkb_geometry_points( + &self, + ewkb: &[u8], + array_path: &str, + target_epsg: i32, + candidates: &[(f64, f64)], + ) -> ZarrFdwResult> { + validate_ewkb_size(ewkb)?; + validate_coverage_candidates(candidates)?; + if target_epsg <= 0 { + return Err(invalid_crs( + array_path, + format!("target EPSG code must be positive, got {target_epsg}"), + )); + } + self.require_known_srid(array_path, target_epsg)?; + let candidate_json = serde_json::to_string(candidates).map_err(|_| { + ZarrFdwError::InvalidGeometry( + "could not encode polygon coverage candidates".to_string(), + ) + })?; + + let sql = polygon_coverage_sql(&self.schema); + let (description, coverage_json) = self.with_extension_search_path(|| { + PgTryBuilder::new(|| { + Spi::connect(|client| { + let mut rows = client.select( + &sql, + Some(1), + &[ + ewkb.to_vec().into(), + target_epsg.into(), + candidate_json.into(), + ], + )?; + let row = rows.next().ok_or_else(|| { + ZarrFdwError::InvalidGeometry( + "PostGIS returned no coverage result for the supplied polygon" + .to_string(), + ) + })?; + let description = PolygonDescription { + source_srid: row + .get_by_name::("source_srid")? + .expect("ST_SRID result is not null"), + source_srid_known: row + .get_by_name::("source_srid_known")? + .expect("source SRID catalog check is not null"), + geometry_type: row + .get_by_name::("geometry_type")? + .expect("ST_GeometryType result is not null"), + dimensions: row + .get_by_name::("dimensions")? + .expect("ST_NDims result is not null"), + empty: row + .get_by_name::("is_empty")? + .expect("ST_IsEmpty result is not null"), + valid: row + .get_by_name::("is_valid")? + .expect("ST_IsValid result is not null"), + transformed_srid: row.get_by_name::("transformed_srid")?, + }; + let coverage_json = row + .get_by_name::("coverage_json")? + .expect("coverage JSON result is not null"); + Ok::<_, ZarrFdwError>((description, coverage_json)) + }) + }) + .catch_when(PgSqlErrorCode::ERRCODE_INTERNAL_ERROR, |_| { + Err(ZarrFdwError::InvalidGeometry( + "PostGIS could not parse, transform, or mask the supplied polygon".to_string(), + )) + }) + .execute() + })?; + validate_polygon_description(&description, array_path, target_epsg)?; + parse_coverage_json(&coverage_json, candidates.len()) + } + + fn with_extension_search_path( + &self, + operation: impl FnOnce() -> ZarrFdwResult, + ) -> ZarrFdwResult { + let prior_search_path = + Spi::get_one::("SELECT pg_catalog.current_setting('search_path')")? + .unwrap_or_default(); + let operation_search_path = format!("{}, pg_catalog", quote_identifier(&self.schema)); + Spi::get_one_with_args::( + "SELECT pg_catalog.set_config('search_path', $1, true)", + &[operation_search_path.into()], + )?; + let result = operation(); + let restore = Spi::get_one_with_args::( + "SELECT pg_catalog.set_config('search_path', $1, true)", + &[prior_search_path.into()], + ); + match (result, restore) { + (Ok(value), Ok(_)) => Ok(value), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error.into()), + } + } + + fn require_known_srid(&self, array_path: &str, epsg: i32) -> ZarrFdwResult<()> { + let schema = quote_identifier(&self.schema); + let sql = format!("SELECT EXISTS (SELECT 1 FROM {schema}.spatial_ref_sys WHERE srid = $1)"); + let exists = Spi::get_one_with_args::(&sql, &[epsg.into()])?.unwrap_or(false); + if !exists { + return Err(invalid_crs( + array_path, + format!("EPSG:{epsg} is not present in the installed PostGIS spatial_ref_sys"), + )); + } + Ok(()) + } +} + +fn polygon_transform_ctes(schema: &str) -> String { + let schema = quote_identifier(schema); + format!( + r#" + WITH parsed AS MATERIALIZED ( + SELECT {schema}.ST_GeomFromEWKB($1) AS geometry + ), + described AS MATERIALIZED ( + SELECT geometry, + {schema}.ST_SRID(geometry) AS source_srid, + EXISTS ( + SELECT 1 + FROM {schema}.spatial_ref_sys AS source_crs + WHERE source_crs.srid = {schema}.ST_SRID(geometry) + ) AS source_srid_known, + {schema}.ST_GeometryType(geometry) AS geometry_type, + {schema}.ST_NDims(geometry)::integer AS dimensions, + {schema}.ST_IsEmpty(geometry) AS is_empty, + {schema}.ST_IsValid(geometry) AS is_valid + FROM parsed + ), + transformed AS MATERIALIZED ( + SELECT source_srid, source_srid_known, geometry_type, dimensions, is_empty, is_valid, + CASE + WHEN source_srid > 0 + AND source_srid_known + AND geometry_type IN ('ST_Polygon', 'ST_MultiPolygon') + AND dimensions = 2 + AND NOT is_empty + AND is_valid + THEN {schema}.ST_Transform(geometry, $2) + END AS geometry + FROM described + ) + "# + ) +} + +fn polygon_envelope_sql(schema: &str) -> String { + let schema_identifier = quote_identifier(schema); + let transform = polygon_transform_ctes(schema); + format!( + r#" + {transform}, + bounded AS ( + SELECT source_srid, source_srid_known, geometry_type, dimensions, is_empty, is_valid, geometry, + CASE WHEN geometry IS NOT NULL + THEN {schema_identifier}.Box3D(geometry) + END AS bounds + FROM transformed + ) + SELECT source_srid, source_srid_known, geometry_type, dimensions, is_empty, is_valid, + CASE WHEN geometry IS NOT NULL + THEN {schema_identifier}.ST_SRID(geometry) + END AS transformed_srid, + CASE WHEN bounds IS NOT NULL + THEN {schema_identifier}.ST_XMin(bounds) + END AS xmin, + CASE WHEN bounds IS NOT NULL + THEN {schema_identifier}.ST_YMin(bounds) + END AS ymin, + CASE WHEN bounds IS NOT NULL + THEN {schema_identifier}.ST_XMax(bounds) + END AS xmax, + CASE WHEN bounds IS NOT NULL + THEN {schema_identifier}.ST_YMax(bounds) + END AS ymax + FROM bounded + "# + ) +} + +fn polygon_coverage_sql(schema: &str) -> String { + let schema_identifier = quote_identifier(schema); + let transform = polygon_transform_ctes(schema); + format!( + r#" + {transform}, + candidate_points AS ( + SELECT candidate.ordinality, + pg_catalog.jsonb_extract_path_text(candidate.value, '0')::pg_catalog.float8 AS x, + pg_catalog.jsonb_extract_path_text(candidate.value, '1')::pg_catalog.float8 AS y + FROM pg_catalog.jsonb_array_elements($3::pg_catalog.jsonb) + WITH ORDINALITY AS candidate(value, ordinality) + ), + coverage AS ( + SELECT candidate.ordinality, + CASE WHEN transformed.geometry IS NOT NULL THEN + {schema_identifier}.ST_Covers( + transformed.geometry, + {schema_identifier}.ST_SetSRID( + {schema_identifier}.ST_MakePoint(candidate.x, candidate.y), + $2 + ) + ) + END AS is_covered + FROM transformed + CROSS JOIN candidate_points AS candidate + ) + SELECT transformed.source_srid, + transformed.source_srid_known, + transformed.geometry_type, + transformed.dimensions, + transformed.is_empty, + transformed.is_valid, + CASE WHEN transformed.geometry IS NOT NULL + THEN {schema_identifier}.ST_SRID(transformed.geometry) + END AS transformed_srid, + COALESCE( + ( + SELECT pg_catalog.jsonb_agg( + coverage.is_covered ORDER BY coverage.ordinality + ) + FROM coverage + ), + '[]'::pg_catalog.jsonb + )::pg_catalog.text AS coverage_json + FROM transformed + "# + ) +} + +fn point_transform_sql(schema: &str) -> String { + let schema = quote_identifier(schema); + format!( + r#" + WITH parsed AS ( + SELECT {schema}.ST_GeomFromEWKB($1) AS geometry + ), + described AS ( + SELECT geometry, + {schema}.ST_SRID(geometry) AS source_srid, + {schema}.ST_GeometryType(geometry) AS geometry_type, + {schema}.ST_NDims(geometry)::integer AS dimensions, + {schema}.ST_IsEmpty(geometry) AS is_empty + FROM parsed + ), + transformed AS ( + SELECT source_srid, geometry_type, dimensions, is_empty, + CASE + WHEN source_srid > 0 + AND geometry_type = 'ST_Point' + AND dimensions = 2 + AND NOT is_empty + THEN {schema}.ST_Transform(geometry, $2) + END AS geometry + FROM described + ) + SELECT source_srid, geometry_type, dimensions, is_empty, + CASE WHEN geometry IS NOT NULL THEN {schema}.ST_X(geometry) END AS x, + CASE WHEN geometry IS NOT NULL THEN {schema}.ST_Y(geometry) END AS y, + CASE WHEN geometry IS NOT NULL THEN {schema}.ST_SRID(geometry) END AS transformed_srid + FROM transformed + "# + ) +} + +fn validate_point_description( + point: PointDescription, + target_srid: i32, +) -> ZarrFdwResult { + if point.source_srid <= 0 { + return Err(ZarrFdwError::InvalidGeometry( + "EWKB point must declare a positive SRID".to_string(), + )); + } + if point.geometry_type != "ST_Point" { + return Err(ZarrFdwError::InvalidGeometry(format!( + "expected a Point, got {}", + point.geometry_type + ))); + } + if point.dimensions != 2 { + return Err(ZarrFdwError::InvalidGeometry(format!( + "expected a two-dimensional Point, got {} dimensions", + point.dimensions + ))); + } + if point.empty { + return Err(ZarrFdwError::InvalidGeometry( + "point must not be empty".to_string(), + )); + } + let (Some(x), Some(y), Some(transformed_srid)) = (point.x, point.y, point.transformed_srid) + else { + return Err(ZarrFdwError::InvalidGeometry( + "PostGIS did not return transformed point coordinates".to_string(), + )); + }; + if !x.is_finite() || !y.is_finite() { + return Err(ZarrFdwError::InvalidGeometry( + "transformed point coordinates must be finite".to_string(), + )); + } + if transformed_srid != target_srid { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned SRID {transformed_srid}, expected {target_srid}" + ))); + } + Ok(TransformedPoint { + x, + y, + source_srid: point.source_srid, + target_srid, + }) +} + +fn validate_ewkb_size(ewkb: &[u8]) -> ZarrFdwResult<()> { + if ewkb.len() > MAX_EWKB_BYTES { + return Err(ZarrFdwError::InvalidGeometry(format!( + "EWKB input is {} bytes, exceeding the {MAX_EWKB_BYTES}-byte limit", + ewkb.len() + ))); + } + Ok(()) +} + +fn validate_coverage_candidates(candidates: &[(f64, f64)]) -> ZarrFdwResult<()> { + if candidates.len() > MAX_COVERAGE_CANDIDATES { + return Err(ZarrFdwError::InvalidGeometry(format!( + "polygon coverage batch contains {} candidates, exceeding the {MAX_COVERAGE_CANDIDATES}-candidate limit", + candidates.len() + ))); + } + if let Some((index, (x, y))) = candidates + .iter() + .enumerate() + .find(|(_, (x, y))| !x.is_finite() || !y.is_finite()) + { + return Err(ZarrFdwError::InvalidGeometry(format!( + "polygon coverage candidate {index} must contain finite coordinates, got ({x}, {y})" + ))); + } + Ok(()) +} + +fn validate_polygon_description( + polygon: &PolygonDescription, + array_path: &str, + target_srid: i32, +) -> ZarrFdwResult<()> { + if polygon.source_srid <= 0 { + return Err(ZarrFdwError::InvalidGeometry( + "EWKB polygon must declare a positive SRID".to_string(), + )); + } + if !polygon.source_srid_known { + return Err(invalid_crs( + array_path, + format!( + "EPSG:{} is not present in the installed PostGIS spatial_ref_sys", + polygon.source_srid + ), + )); + } + if !matches!( + polygon.geometry_type.as_str(), + "ST_Polygon" | "ST_MultiPolygon" + ) { + return Err(ZarrFdwError::InvalidGeometry(format!( + "expected a Polygon or MultiPolygon, got {}", + polygon.geometry_type + ))); + } + if polygon.dimensions != 2 { + return Err(ZarrFdwError::InvalidGeometry(format!( + "expected a two-dimensional polygon, got {} dimensions", + polygon.dimensions + ))); + } + if polygon.empty { + return Err(ZarrFdwError::InvalidGeometry( + "polygon must not be empty".to_string(), + )); + } + if !polygon.valid { + return Err(ZarrFdwError::InvalidGeometry( + "polygon must be valid according to PostGIS ST_IsValid".to_string(), + )); + } + let Some(transformed_srid) = polygon.transformed_srid else { + return Err(ZarrFdwError::InvalidGeometry( + "PostGIS did not return a transformed polygon".to_string(), + )); + }; + if transformed_srid != target_srid { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned SRID {transformed_srid}, expected {target_srid}" + ))); + } + Ok(()) +} + +fn validate_envelope_description( + description: EnvelopeDescription, + array_path: &str, + target_srid: i32, +) -> ZarrFdwResult { + validate_polygon_description(&description.polygon, array_path, target_srid)?; + let (Some(xmin), Some(ymin), Some(xmax), Some(ymax)) = ( + description.xmin, + description.ymin, + description.xmax, + description.ymax, + ) else { + return Err(ZarrFdwError::InvalidGeometry( + "PostGIS did not return a polygon envelope".to_string(), + )); + }; + if [xmin, ymin, xmax, ymax] + .iter() + .any(|value| !value.is_finite()) + { + return Err(ZarrFdwError::InvalidGeometry( + "transformed polygon envelope must contain only finite coordinates".to_string(), + )); + } + if xmin > xmax || ymin > ymax { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned an invalid polygon envelope ({xmin}, {ymin}, {xmax}, {ymax})" + ))); + } + Ok(GeometryEnvelope { + xmin, + ymin, + xmax, + ymax, + source_srid: description.polygon.source_srid, + target_srid, + }) +} + +fn parse_coverage_json(value: &str, expected_len: usize) -> ZarrFdwResult> { + let coverage = serde_json::from_str::>(value).map_err(|_| { + ZarrFdwError::InvalidGeometry( + "PostGIS returned malformed polygon coverage results".to_string(), + ) + })?; + if coverage.len() != expected_len { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned {} polygon coverage results for {expected_len} candidates", + coverage.len() + ))); + } + Ok(coverage) +} + +fn invalid_crs(array: &str, message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidCrs { + array: array.to_string(), + message: message.into(), + } +} + +#[cfg(test)] +mod unit_tests { + use super::*; + + fn point() -> PointDescription { + PointDescription { + source_srid: 4326, + geometry_type: "ST_Point".to_string(), + dimensions: 2, + empty: false, + x: Some(100.0), + y: Some(20.0), + transformed_srid: Some(3857), + } + } + + fn polygon() -> PolygonDescription { + PolygonDescription { + source_srid: 4326, + source_srid_known: true, + geometry_type: "ST_Polygon".to_string(), + dimensions: 2, + empty: false, + valid: true, + transformed_srid: Some(3857), + } + } + + fn envelope() -> EnvelopeDescription { + EnvelopeDescription { + polygon: polygon(), + xmin: Some(10.0), + ymin: Some(20.0), + xmax: Some(30.0), + ymax: Some(40.0), + } + } + + #[test] + fn validates_transformed_point_description() { + assert_eq!( + validate_point_description(point(), 3857).unwrap(), + TransformedPoint { + x: 100.0, + y: 20.0, + source_srid: 4326, + target_srid: 3857, + } + ); + + let mut invalid = point(); + invalid.source_srid = 0; + assert!(validate_point_description(invalid, 3857).is_err()); + + let mut invalid = point(); + invalid.geometry_type = "ST_LineString".to_string(); + assert!(validate_point_description(invalid, 3857).is_err()); + + let mut invalid = point(); + invalid.dimensions = 3; + assert!(validate_point_description(invalid, 3857).is_err()); + + let mut invalid = point(); + invalid.empty = true; + assert!(validate_point_description(invalid, 3857).is_err()); + + let mut invalid = point(); + invalid.x = Some(f64::NAN); + assert!(validate_point_description(invalid, 3857).is_err()); + + let mut invalid = point(); + invalid.transformed_srid = Some(4326); + assert!(validate_point_description(invalid, 3857).is_err()); + } + + #[test] + fn validates_polygon_contract_and_finite_envelope() { + assert_eq!( + validate_envelope_description(envelope(), "nested/spatial2d", 3857).unwrap(), + GeometryEnvelope { + xmin: 10.0, + ymin: 20.0, + xmax: 30.0, + ymax: 40.0, + source_srid: 4326, + target_srid: 3857, + } + ); + + let mut invalid = polygon(); + invalid.source_srid = 0; + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut invalid = polygon(); + invalid.source_srid_known = false; + assert!(matches!( + validate_polygon_description(&invalid, "array", 3857), + Err(ZarrFdwError::InvalidCrs { .. }) + )); + + let mut invalid = polygon(); + invalid.geometry_type = "ST_LineString".to_string(); + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut multipolygon = polygon(); + multipolygon.geometry_type = "ST_MultiPolygon".to_string(); + assert!(validate_polygon_description(&multipolygon, "array", 3857).is_ok()); + + let mut invalid = polygon(); + invalid.dimensions = 3; + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut invalid = polygon(); + invalid.empty = true; + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut invalid = polygon(); + invalid.valid = false; + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut invalid = polygon(); + invalid.transformed_srid = Some(4326); + assert!(validate_polygon_description(&invalid, "array", 3857).is_err()); + + let mut invalid = envelope(); + invalid.xmin = Some(f64::NAN); + assert!(validate_envelope_description(invalid, "array", 3857).is_err()); + + let mut invalid = envelope(); + invalid.xmin = Some(31.0); + assert!(validate_envelope_description(invalid, "array", 3857).is_err()); + } + + #[test] + fn bounds_and_validates_coverage_candidates_and_results() { + assert!(validate_coverage_candidates(&[]).is_ok()); + assert!(validate_coverage_candidates(&[(1.0, 2.0)]).is_ok()); + assert!(validate_coverage_candidates(&[(f64::INFINITY, 2.0)]).is_err()); + assert!( + validate_coverage_candidates(&vec![(0.0, 0.0); MAX_COVERAGE_CANDIDATES + 1]).is_err() + ); + + assert_eq!( + parse_coverage_json("[true,false,true]", 3).unwrap(), + vec![true, false, true] + ); + assert!(parse_coverage_json("[true]", 2).is_err()); + assert!(parse_coverage_json("[true,null]", 2).is_err()); + } +} + +// SQL construction quotes identifiers with PostgreSQL's own routine, so these +// checks run in a backend rather than initializing pgrx from a Rust test thread. +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use super::*; + use pgrx::pg_test; + + #[pg_test] + fn point_sql_is_qualified_and_parameterized() { + let sql = point_transform_sql("post\"gis"); + assert!(sql.contains("\"post\"\"gis\".ST_GeomFromEWKB($1)")); + assert!(sql.contains("\"post\"\"gis\".ST_Transform(geometry, $2)")); + assert!(!sql.contains("4326")); + assert!(!sql.contains("3857")); + } + + #[pg_test] + fn polygon_sql_is_qualified_parameterized_and_batched() { + let envelope_sql = polygon_envelope_sql("post\"gis"); + assert!(envelope_sql.contains("\"post\"\"gis\".ST_IsValid(geometry)")); + assert!(envelope_sql.contains("\"post\"\"gis\".Box3D(geometry)")); + assert!(envelope_sql.contains("\"post\"\"gis\".ST_XMin(bounds)")); + assert!(envelope_sql.contains("ST_Transform(geometry, $2)")); + + let coverage_sql = polygon_coverage_sql("post\"gis"); + assert!(coverage_sql.contains("jsonb_array_elements($3::pg_catalog.jsonb)")); + assert!(coverage_sql.contains("\"post\"\"gis\".ST_Covers(")); + assert!(coverage_sql.contains("\"post\"\"gis\".ST_MakePoint(")); + assert!(coverage_sql.contains("ORDER BY coverage.ordinality")); + assert!(!coverage_sql.contains("4326")); + assert!(!coverage_sql.contains("3857")); + + for function in [ + "st_isvalid", + "box3d", + "st_xmin", + "st_ymin", + "st_xmax", + "st_ymax", + "st_makepoint", + "st_setsrid", + "st_covers", + ] { + assert!(VERIFY_POSTGIS_MEMBERS_SQL.contains(&format!("proname = '{function}'"))); + } + } +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/temporal.rs b/wrappers/src/fdw/zarr_fdw/spatial/temporal.rs new file mode 100644 index 000000000..4204d07f6 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/temporal.rs @@ -0,0 +1,501 @@ +//! Time-aware polygon cell extraction and zonal statistics for rank-3+ grids. + +use std::collections::BTreeMap; + +use pgrx::datum::TimestampWithTimeZone; +use pgrx::pg_sys::panic::{ErrorReport, ErrorReportable}; +use pgrx::prelude::*; +use supabase_wrappers::prelude::{ForeignDataWrapper, Row}; + +use super::super::selectors::{DimensionSelectors, OPT_DIMENSION_SELECTORS}; +use super::super::zarr_fdw::ZarrFdw; +use super::super::{ZarrFdwError, ZarrFdwResult}; +use super::catalog::{load_zarr_foreign_table, load_zarr_foreign_table_with_selectors}; +use super::point::numeric_cell_to_f64; +use super::postgis::{MAX_COVERAGE_CANDIDATES, PostgisCatalog}; +use super::zonal::{ + MAX_SPATIAL_CANDIDATE_CELLS, MAX_SPATIAL_OUTPUT_CELLS, SPATIAL_INTERRUPT_POLL_CELLS, + ZonalAccumulator, +}; + +#[derive(Debug)] +struct TemporalCellRow { + time: TimestampWithTimeZone, + x: f64, + y: f64, + value: Option, + time_index: i64, + x_index: i64, + y_index: i64, + srid: i32, +} + +impl TemporalCellRow { + #[allow(clippy::type_complexity)] + fn sql_row( + self, + ) -> ( + TimestampWithTimeZone, + f64, + f64, + Option, + i64, + i64, + i64, + i32, + ) { + ( + self.time, + self.x, + self.y, + self.value, + self.time_index, + self.x_index, + self.y_index, + self.srid, + ) + } +} + +#[derive(Debug)] +struct TemporalZonalStatsRow { + time: TimestampWithTimeZone, + time_index: i64, + count: i64, + valid_count: i64, + min: Option, + max: Option, + sum: Option, + avg: Option, + srid: i32, +} + +impl TemporalZonalStatsRow { + #[allow(clippy::type_complexity)] + fn sql_row( + self, + ) -> ( + TimestampWithTimeZone, + i64, + i64, + i64, + Option, + Option, + Option, + Option, + i32, + ) { + ( + self.time, + self.time_index, + self.count, + self.valid_count, + self.min, + self.max, + self.sum, + self.avg, + self.srid, + ) + } +} + +#[derive(Debug)] +struct TemporalVisitSummary { + times: Vec<(TimestampWithTimeZone, i64)>, + srid: i32, +} + +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace, volatile, parallel_unsafe)] +fn zarr_cells_by_time( + foreign_table: &str, + region_ewkb: &[u8], + start_time: TimestampWithTimeZone, + end_time: TimestampWithTimeZone, +) -> TableIterator< + 'static, + ( + name!(time, TimestampWithTimeZone), + name!(x, f64), + name!(y, f64), + name!(value, Option), + name!(time_index, i64), + name!(x_index, i64), + name!(y_index, i64), + name!(srid, i32), + ), +> { + let rows = cells_foreign_table(foreign_table, region_ewkb, start_time, end_time) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(TemporalCellRow::sql_row)) +} + +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace, volatile, parallel_unsafe)] +fn zarr_zonal_stats_by_time( + foreign_table: &str, + region_ewkb: &[u8], + start_time: TimestampWithTimeZone, + end_time: TimestampWithTimeZone, +) -> TableIterator< + 'static, + ( + name!(time, TimestampWithTimeZone), + name!(time_index, i64), + name!(count, i64), + name!(valid_count, i64), + name!(min, Option), + name!(max, Option), + name!(sum, Option), + name!(avg, Option), + name!(srid, i32), + ), +> { + let rows = zonal_foreign_table(foreign_table, region_ewkb, start_time, end_time, None) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(TemporalZonalStatsRow::sql_row)) +} + +#[allow(clippy::type_complexity)] +#[pg_extern( + name = "zarr_zonal_stats_by_time", + create_or_replace, + volatile, + parallel_unsafe +)] +fn zarr_zonal_stats_by_time_with_selectors( + foreign_table: &str, + region_ewkb: &[u8], + start_time: TimestampWithTimeZone, + end_time: TimestampWithTimeZone, + dimension_selectors: &str, +) -> TableIterator< + 'static, + ( + name!(time, TimestampWithTimeZone), + name!(time_index, i64), + name!(count, i64), + name!(valid_count, i64), + name!(min, Option), + name!(max, Option), + name!(sum, Option), + name!(avg, Option), + name!(srid, i32), + ), +> { + let rows = zonal_foreign_table( + foreign_table, + region_ewkb, + start_time, + end_time, + Some(dimension_selectors), + ) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(rows.into_iter().map(TemporalZonalStatsRow::sql_row)) +} + +fn cells_foreign_table( + foreign_table: &str, + region_ewkb: &[u8], + start: TimestampWithTimeZone, + end: TimestampWithTimeZone, +) -> ZarrFdwResult> { + let mut cells = Vec::new(); + visit_foreign_table_cells(foreign_table, region_ewkb, start, end, None, |cell| { + if cells.len() >= MAX_SPATIAL_OUTPUT_CELLS { + return Err(ZarrFdwError::InvalidGeometry(format!( + "spatiotemporal polygon selects more than the {MAX_SPATIAL_OUTPUT_CELLS}-cell output limit" + ))); + } + cells.push(cell); + Ok(()) + })?; + Ok(cells) +} + +fn zonal_foreign_table( + foreign_table: &str, + region_ewkb: &[u8], + start: TimestampWithTimeZone, + end: TimestampWithTimeZone, + call_selectors: Option<&str>, +) -> ZarrFdwResult> { + let mut accumulators = BTreeMap::::new(); + let summary = visit_foreign_table_cells( + foreign_table, + region_ewkb, + start, + end, + call_selectors, + |cell| { + accumulators + .entry(cell.time_index) + .or_default() + .observe(cell.value) + }, + )?; + + summary + .times + .into_iter() + .map(|(time, time_index)| { + let stats = accumulators + .remove(&time_index) + .unwrap_or_default() + .finish(summary.srid)?; + Ok(TemporalZonalStatsRow { + time, + time_index, + count: stats.count, + valid_count: stats.valid_count, + min: stats.min, + max: stats.max, + sum: stats.sum, + avg: stats.avg, + srid: stats.srid, + }) + }) + .collect() +} + +fn visit_foreign_table_cells( + foreign_table: &str, + region_ewkb: &[u8], + start: TimestampWithTimeZone, + end: TimestampWithTimeZone, + call_selectors: Option<&str>, + mut visitor: impl FnMut(TemporalCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult { + // Fail before any remote metadata request when the optional spatial + // dependency is absent or incomplete. + let postgis = PostgisCatalog::require()?; + let selector_aware = call_selectors.is_some(); + let table = if selector_aware { + load_zarr_foreign_table_with_selectors(foreign_table)? + } else { + load_zarr_foreign_table(foreign_table)? + }; + let call_selectors = match call_selectors { + Some(raw) => { + DimensionSelectors::parse( + table + .options + .get(OPT_DIMENSION_SELECTORS) + .map(String::as_str), + )?; + DimensionSelectors::parse(Some(raw))? + } + None => DimensionSelectors::default(), + }; + let mut fdw = ZarrFdw::new(table.server)?; + fdw.set_call_dimension_selectors(call_selectors)?; + let begin = >::begin_scan( + &mut fdw, + &[], + &table.columns, + &[], + &None, + &table.options, + ); + if let Err(error) = begin { + let _ = >::end_scan(&mut fdw); + return Err(error); + } + + let result = visit_prepared_scan(&mut fdw, &postgis, region_ewkb, start, end, &mut visitor); + let cleanup = >::end_scan(&mut fdw); + match (result, cleanup) { + (Ok(value), Ok(())) => Ok(value), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +fn visit_prepared_scan( + fdw: &mut ZarrFdw, + postgis: &PostgisCatalog, + region_ewkb: &[u8], + start: TimestampWithTimeZone, + end: TimestampWithTimeZone, + visitor: &mut impl FnMut(TemporalCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult { + // Validate the table contract before a non-overlapping envelope can + // return an otherwise-successful empty result. + let value_column = fdw.spatial_value_column()?.to_string(); + let operation_layout = fdw.spatial_time_layout()?; + let time_axis = operation_layout.time; + let time_indices = fdw.spatial_time_indices(start, end, MAX_SPATIAL_OUTPUT_CELLS)?; + let mut times = Vec::with_capacity(time_indices.len()); + for (selected_position, &time_index) in time_indices.iter().enumerate() { + if selected_position.is_multiple_of(SPATIAL_INTERRUPT_POLL_CELLS) { + fdw.spatial_check_for_interrupt()?; + } + times.push(( + fdw.spatial_time_at_index(time_axis, time_index)?, + pg_index("time", time_index)?, + )); + } + let crs = fdw.resolved_spatial_crs()?; + let envelope = postgis.transform_ewkb_geometry_envelope( + region_ewkb, + fdw.spatial_array_path(), + crs.epsg, + )?; + if !fdw.apply_spatial_dimension_selectors( + "zarr_zonal_stats_by_time", + &[ + operation_layout.time, + operation_layout.horizontal.x, + operation_layout.horizontal.y, + ], + )? { + return Ok(TemporalVisitSummary { + times, + srid: crs.epsg, + }); + } + let Some((window_axes, horizontal_bounds, _spatial_candidates)) = fdw + .spatial_time_horizontal_window( + envelope.xmin, + envelope.ymin, + envelope.xmax, + envelope.ymax, + )? + else { + return Ok(TemporalVisitSummary { + times, + srid: crs.epsg, + }); + }; + if time_indices.is_empty() { + return Ok(TemporalVisitSummary { + times, + srid: crs.epsg, + }); + } + let selection = + fdw.spatial_time_selection(time_indices, horizontal_bounds, MAX_SPATIAL_CANDIDATE_CELLS)?; + debug_assert_eq!(window_axes, selection.layout.horizontal); + debug_assert!(selection.candidate_cells <= MAX_SPATIAL_CANDIDATE_CELLS); + let layout = selection.layout; + let exact_time_indices = selection.time_indices; + fdw.restrict_to_axis_bounds(selection.bounds)?; + + let mut candidates = Vec::with_capacity(MAX_COVERAGE_CANDIDATES); + let mut row = Row::new(); + let mut visited = 0usize; + while >::iter_scan(fdw, &mut row)?.is_some() { + if visited.is_multiple_of(SPATIAL_INTERRUPT_POLL_CELLS) { + fdw.spatial_check_for_interrupt()?; + } + visited = visited.saturating_add(1); + + let indices = fdw.spatial_last_emitted_global_indices()?; + let time_index = axis_index(indices, layout.time, "time")?; + let Ok(time_position) = exact_time_indices.binary_search(&time_index) else { + continue; + }; + let x_index = axis_index(indices, layout.horizontal.x, "x")?; + let y_index = axis_index(indices, layout.horizontal.y, "y")?; + let value = row + .iter() + .find(|(name, _)| name.as_str() == value_column) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatiotemporal result did not contain value column '{value_column}'" + )) + })? + .1 + .as_ref() + .map(numeric_cell_to_f64) + .transpose()?; + candidates.push(TemporalCellRow { + time: times[time_position].0, + x: fdw.spatial_coordinate_at_index(layout.horizontal.x, x_index)?, + y: fdw.spatial_coordinate_at_index(layout.horizontal.y, y_index)?, + value, + time_index: pg_index("time", time_index)?, + x_index: pg_index("x", x_index)?, + y_index: pg_index("y", y_index)?, + srid: crs.epsg, + }); + if candidates.len() == MAX_COVERAGE_CANDIDATES { + visit_covered_batch( + postgis, + region_ewkb, + fdw, + crs.epsg, + &mut candidates, + visitor, + )?; + } + } + if !candidates.is_empty() { + visit_covered_batch( + postgis, + region_ewkb, + fdw, + crs.epsg, + &mut candidates, + visitor, + )?; + } + + Ok(TemporalVisitSummary { + times, + srid: crs.epsg, + }) +} + +fn visit_covered_batch( + postgis: &PostgisCatalog, + region_ewkb: &[u8], + fdw: &mut ZarrFdw, + srid: i32, + candidates: &mut Vec, + visitor: &mut impl FnMut(TemporalCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult<()> { + fdw.spatial_check_for_interrupt()?; + let centers = candidates + .iter() + .map(|cell| (cell.x, cell.y)) + .collect::>(); + let covered = postgis.covers_ewkb_geometry_points( + region_ewkb, + fdw.spatial_array_path(), + srid, + ¢ers, + )?; + if covered.len() != candidates.len() { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned {} polygon coverage results for {} spatiotemporal candidate cells", + covered.len(), + candidates.len() + ))); + } + for (cell, covered) in candidates.drain(..).zip(covered) { + if covered { + visitor(cell)?; + } + } + Ok(()) +} + +fn axis_index(indices: &[usize], axis: usize, name: &str) -> ZarrFdwResult { + indices.get(axis).copied().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatiotemporal {name} axis {axis} is outside emitted array indexes" + )) + }) +} + +fn pg_index(name: &str, index: usize) -> ZarrFdwResult { + i64::try_from(index).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "spatiotemporal {name} index exceeds PostgreSQL bigint range" + )) + }) +} diff --git a/wrappers/src/fdw/zarr_fdw/spatial/zonal.rs b/wrappers/src/fdw/zarr_fdw/spatial/zonal.rs new file mode 100644 index 000000000..6195d64b1 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/spatial/zonal.rs @@ -0,0 +1,489 @@ +//! Polygon-masked cell extraction and zonal statistics for rank-2 grids. + +use std::cmp::Ordering; + +use pgrx::pg_sys::panic::{ErrorReport, ErrorReportable}; +use pgrx::prelude::*; +use supabase_wrappers::prelude::{ForeignDataWrapper, Row}; + +use super::super::aggregate::{checked_float_add_f64, compare_f64}; +use super::super::selectors::{DimensionSelectors, OPT_DIMENSION_SELECTORS}; +use super::super::zarr_fdw::ZarrFdw; +use super::super::{ZarrFdwError, ZarrFdwResult}; +use super::catalog::{load_zarr_foreign_table, load_zarr_foreign_table_with_selectors}; +use super::grid::{CoordinateEnvelope, HorizontalCell}; +use super::point::numeric_cell_to_f64; +use super::postgis::{MAX_COVERAGE_CANDIDATES, PostgisCatalog}; + +pub(super) const MAX_SPATIAL_CANDIDATE_CELLS: usize = 10_000_000; +pub(super) const MAX_SPATIAL_OUTPUT_CELLS: usize = 1_000_000; +pub(super) const SPATIAL_INTERRUPT_POLL_CELLS: usize = 1_024; + +#[derive(Debug, Clone, PartialEq)] +struct SpatialCellRow { + x: f64, + y: f64, + value: Option, + x_index: i64, + y_index: i64, + srid: i32, +} + +impl SpatialCellRow { + fn sql_row(self) -> (f64, f64, Option, i64, i64, i32) { + ( + self.x, + self.y, + self.value, + self.x_index, + self.y_index, + self.srid, + ) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct ZonalStatsRow { + pub(super) count: i64, + pub(super) valid_count: i64, + pub(super) min: Option, + pub(super) max: Option, + pub(super) sum: Option, + pub(super) avg: Option, + pub(super) srid: i32, +} + +impl ZonalStatsRow { + #[allow(clippy::type_complexity)] + fn sql_row( + self, + ) -> ( + i64, + i64, + Option, + Option, + Option, + Option, + i32, + ) { + ( + self.count, + self.valid_count, + self.min, + self.max, + self.sum, + self.avg, + self.srid, + ) + } +} + +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace, volatile, parallel_unsafe)] +fn zarr_cells( + foreign_table: &str, + region_ewkb: &[u8], +) -> TableIterator< + 'static, + ( + name!(x, f64), + name!(y, f64), + name!(value, Option), + name!(x_index, i64), + name!(y_index, i64), + name!(srid, i32), + ), +> { + let mut cells = Vec::new(); + visit_foreign_table_cells(foreign_table, region_ewkb, None, |cell| { + if cells.len() >= MAX_SPATIAL_OUTPUT_CELLS { + return Err(ZarrFdwError::InvalidGeometry(format!( + "polygon selects more than the {MAX_SPATIAL_OUTPUT_CELLS}-cell output limit" + ))); + } + cells.push(cell); + Ok(()) + }) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::new(cells.into_iter().map(SpatialCellRow::sql_row)) +} + +#[allow(clippy::type_complexity)] +#[pg_extern(create_or_replace, volatile, parallel_unsafe)] +fn zarr_zonal_stats( + foreign_table: &str, + region_ewkb: &[u8], +) -> TableIterator< + 'static, + ( + name!(count, i64), + name!(valid_count, i64), + name!(min, Option), + name!(max, Option), + name!(sum, Option), + name!(avg, Option), + name!(srid, i32), + ), +> { + let mut accumulator = ZonalAccumulator::default(); + let srid = visit_foreign_table_cells(foreign_table, region_ewkb, None, |cell| { + accumulator.observe(cell.value) + }) + .map_err(ErrorReport::from) + .unwrap_or_report(); + let row = accumulator + .finish(srid) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::once(row.sql_row()) +} + +#[allow(clippy::type_complexity)] +#[pg_extern( + name = "zarr_zonal_stats", + create_or_replace, + volatile, + parallel_unsafe +)] +fn zarr_zonal_stats_with_selectors( + foreign_table: &str, + region_ewkb: &[u8], + dimension_selectors: &str, +) -> TableIterator< + 'static, + ( + name!(count, i64), + name!(valid_count, i64), + name!(min, Option), + name!(max, Option), + name!(sum, Option), + name!(avg, Option), + name!(srid, i32), + ), +> { + let mut accumulator = ZonalAccumulator::default(); + let srid = visit_foreign_table_cells( + foreign_table, + region_ewkb, + Some(dimension_selectors), + |cell| accumulator.observe(cell.value), + ) + .map_err(ErrorReport::from) + .unwrap_or_report(); + let row = accumulator + .finish(srid) + .map_err(ErrorReport::from) + .unwrap_or_report(); + TableIterator::once(row.sql_row()) +} + +fn visit_foreign_table_cells( + foreign_table: &str, + region_ewkb: &[u8], + call_selectors: Option<&str>, + mut visitor: impl FnMut(SpatialCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult { + // Fail before any remote metadata request when the optional spatial + // dependency is absent or incomplete. + let postgis = PostgisCatalog::require()?; + let selector_aware = call_selectors.is_some(); + let table = if selector_aware { + load_zarr_foreign_table_with_selectors(foreign_table)? + } else { + load_zarr_foreign_table(foreign_table)? + }; + let call_selectors = match call_selectors { + Some(raw) => { + DimensionSelectors::parse( + table + .options + .get(OPT_DIMENSION_SELECTORS) + .map(String::as_str), + )?; + DimensionSelectors::parse(Some(raw))? + } + None => DimensionSelectors::default(), + }; + let mut fdw = ZarrFdw::new(table.server)?; + fdw.set_call_dimension_selectors(call_selectors)?; + let begin = >::begin_scan( + &mut fdw, + &[], + &table.columns, + &[], + &None, + &table.options, + ); + if let Err(error) = begin { + let _ = >::end_scan(&mut fdw); + return Err(error); + } + + let result = visit_prepared_scan( + &mut fdw, + &postgis, + region_ewkb, + selector_aware, + &mut visitor, + ); + let cleanup = >::end_scan(&mut fdw); + match (result, cleanup) { + (Ok(value), Ok(())) => Ok(value), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +fn visit_prepared_scan( + fdw: &mut ZarrFdw, + postgis: &PostgisCatalog, + region_ewkb: &[u8], + selector_aware: bool, + visitor: &mut impl FnMut(SpatialCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult { + // Validate the table contract before a non-overlapping envelope can + // return an otherwise-successful empty result. + let value_column = fdw.spatial_value_column()?.to_string(); + let crs = fdw.resolved_spatial_crs()?; + let envelope = postgis.transform_ewkb_geometry_envelope( + region_ewkb, + fdw.spatial_array_path(), + crs.epsg, + )?; + let validated_envelope = + CoordinateEnvelope::new(envelope.xmin, envelope.ymin, envelope.xmax, envelope.ymax)?; + let (axes, bounds, candidate_count) = if selector_aware { + let axes = fdw.spatial_horizontal_axes()?; + if !fdw.apply_spatial_dimension_selectors("zarr_zonal_stats", &[axes.x, axes.y])? { + return Ok(crs.epsg); + } + let plan = fdw.spatial_horizontal_window( + envelope.xmin, + envelope.ymin, + envelope.xmax, + envelope.ymax, + )?; + ( + axes, + plan.as_ref().map(|(_, bounds, _)| *bounds), + plan.as_ref().map_or(0, |(_, _, total)| *total), + ) + } else { + let grid = fdw.rectilinear_grid()?; + let plan = grid.window_plan(validated_envelope)?; + let bounds = plan.as_ref().map(|plan| plan.bounds()); + let candidate_count = plan.as_ref().map_or(0, |plan| plan.total_cells()); + (grid.axes(), bounds, candidate_count) + }; + if candidate_count > MAX_SPATIAL_CANDIDATE_CELLS { + return Err(ZarrFdwError::InvalidGeometry(format!( + "polygon envelope selects {candidate_count} candidate cells, exceeding the {MAX_SPATIAL_CANDIDATE_CELLS}-cell safety limit" + ))); + } + let Some(bounds) = bounds else { + return Ok(crs.epsg); + }; + if selector_aware { + fdw.restrict_to_horizontal_bounds(axes, bounds)?; + } else { + fdw.restrict_to_spatial_bounds(bounds)?; + } + + let mut candidates = Vec::with_capacity(MAX_COVERAGE_CANDIDATES); + let mut row = Row::new(); + let mut visited = 0usize; + while >::iter_scan(fdw, &mut row)?.is_some() { + if visited.is_multiple_of(SPATIAL_INTERRUPT_POLL_CELLS) { + fdw.spatial_check_for_interrupt()?; + } + visited = visited.saturating_add(1); + let cell = fdw.spatial_last_emitted_cell(axes)?; + let value = row + .iter() + .find(|(name, _)| name.as_str() == value_column) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial result did not contain value column '{value_column}'" + )) + })? + .1 + .as_ref() + .map(numeric_cell_to_f64) + .transpose()?; + candidates.push((cell, value)); + if candidates.len() == MAX_COVERAGE_CANDIDATES { + visit_covered_batch( + postgis, + region_ewkb, + fdw, + crs.epsg, + &mut candidates, + visitor, + )?; + } + } + if !candidates.is_empty() { + visit_covered_batch( + postgis, + region_ewkb, + fdw, + crs.epsg, + &mut candidates, + visitor, + )?; + } + Ok(crs.epsg) +} + +fn visit_covered_batch( + postgis: &PostgisCatalog, + region_ewkb: &[u8], + fdw: &mut ZarrFdw, + srid: i32, + candidates: &mut Vec<(HorizontalCell, Option)>, + visitor: &mut impl FnMut(SpatialCellRow) -> ZarrFdwResult<()>, +) -> ZarrFdwResult<()> { + fdw.spatial_check_for_interrupt()?; + let centers = candidates + .iter() + .map(|(cell, _)| (cell.x, cell.y)) + .collect::>(); + let covered = postgis.covers_ewkb_geometry_points( + region_ewkb, + fdw.spatial_array_path(), + srid, + ¢ers, + )?; + if covered.len() != candidates.len() { + return Err(ZarrFdwError::InvalidGeometry(format!( + "PostGIS returned {} mask results for {} candidate cells", + covered.len(), + candidates.len() + ))); + } + for ((cell, value), covered) in candidates.drain(..).zip(covered) { + if covered { + visitor(spatial_cell_row(cell, value, srid)?)?; + } + } + Ok(()) +} + +fn spatial_cell_row( + cell: HorizontalCell, + value: Option, + srid: i32, +) -> ZarrFdwResult { + Ok(SpatialCellRow { + x: cell.x, + y: cell.y, + value, + x_index: i64::try_from(cell.x_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "spatial x index exceeds PostgreSQL bigint range".to_string(), + ) + })?, + y_index: i64::try_from(cell.y_index).map_err(|_| { + ZarrFdwError::InvalidMetadata( + "spatial y index exceeds PostgreSQL bigint range".to_string(), + ) + })?, + srid, + }) +} + +#[derive(Debug, Default)] +pub(super) struct ZonalAccumulator { + count: i64, + valid_count: i64, + min: Option, + max: Option, + sum: Option, +} + +impl ZonalAccumulator { + pub(super) fn observe(&mut self, value: Option) -> ZarrFdwResult<()> { + self.count = self.count.checked_add(1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("zonal COUNT overflowed bigint".to_string()) + })?; + let Some(value) = value else { return Ok(()) }; + self.valid_count = self.valid_count.checked_add(1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("zonal valid COUNT overflowed bigint".to_string()) + })?; + if self + .min + .is_none_or(|current| compare_f64(value, current) != Ordering::Greater) + { + self.min = Some(value); + } + if self + .max + .is_none_or(|current| compare_f64(value, current) != Ordering::Less) + { + self.max = Some(value); + } + self.sum = Some(match self.sum { + Some(current) => checked_float_add_f64(current, value)?, + None => value, + }); + Ok(()) + } + + pub(super) fn finish(self, srid: i32) -> ZarrFdwResult { + let avg = match (self.sum, self.valid_count) { + (Some(sum), count) if count > 0 => Some(sum / count as f64), + _ => None, + }; + Ok(ZonalStatsRow { + count: self.count, + valid_count: self.valid_count, + min: self.min, + max: self.max, + sum: self.sum, + avg, + srid, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn zonal_accumulator_preserves_null_and_float_semantics() { + let mut accumulator = ZonalAccumulator::default(); + for value in [Some(2.0), None, Some(-1.0), Some(3.0)] { + accumulator.observe(value).unwrap(); + } + assert_eq!( + accumulator.finish(3857).unwrap(), + ZonalStatsRow { + count: 4, + valid_count: 3, + min: Some(-1.0), + max: Some(3.0), + sum: Some(4.0), + avg: Some(4.0 / 3.0), + srid: 3857, + } + ); + } + + #[test] + fn empty_zonal_accumulator_returns_sql_null_statistics() { + assert_eq!( + ZonalAccumulator::default().finish(4326).unwrap(), + ZonalStatsRow { + count: 0, + valid_count: 0, + min: None, + max: None, + sum: None, + avg: None, + srid: 4326, + } + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/store.rs b/wrappers/src/fdw/zarr_fdw/store.rs new file mode 100644 index 000000000..c8bd6019e --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/store.rs @@ -0,0 +1,1644 @@ +//! Bounded storage access for Zarr reading. +//! +//! The scan-facing [`ZarrStore`] owns PostgreSQL interruption handling while +//! backend implementations own format-specific object and directory access. + +mod http; +mod local; + +use ::http::Uri; +use aws_config::BehaviorVersion; +use aws_sdk_s3 as s3; +use futures_util::FutureExt; +use futures_util::future::LocalBoxFuture; +use pgrx::pg_sys; +use std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::time::{MissedTickBehavior, interval}; + +use supabase_wrappers::prelude::*; + +use self::http::HttpBackend; +use self::local::LocalBackend; +use super::{ZarrFdwError, ZarrFdwResult}; + +/// Metadata objects are tiny JSON documents. Chunk callers provide a tighter +/// limit derived from their declared decoded layout. +pub(crate) const MAX_METADATA_OBJECT_BYTES: usize = 1024 * 1024; +const INTERRUPT_POLL_INTERVAL: Duration = Duration::from_millis(25); + +pub(crate) type StoreFuture = LocalBoxFuture<'static, ZarrFdwResult>; + +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +pub(crate) enum StorageBackendKind { + S3, + Local, + Http, +} + +impl StorageBackendKind { + pub(crate) fn label(self) -> &'static str { + match self { + Self::S3 => "s3", + Self::Local => "local", + Self::Http => "http", + } + } + + pub(crate) fn effective_max_concurrent_reads(self, configured: usize) -> usize { + match self { + Self::S3 | Self::Http => configured, + // Local reads are poll-driven foreground file I/O. Scheduling + // more than one cannot create useful kernel concurrency. + Self::Local => 1, + } + } +} + +pub(crate) trait StorageBackend: Send + Sync { + fn kind(&self) -> StorageBackendKind; + + fn get_object_owned(&self, key: String, max_bytes: usize) -> StoreFuture>>; + + fn get_range_owned(&self, identity: ReadIdentity) -> StoreFuture>; + + fn list_directory_page_owned( + &self, + path: String, + continuation_token: Option, + ) -> StoreFuture; +} + +/// The exact object bytes represented by one storage read/cache entry. +/// +/// Suffix requests are useful for end-located shard indexes. Successful +/// suffix reads are normalized to [`ReadRange::Exact`] in [`RangedObject`]. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub(crate) enum ReadRange { + Whole, + Exact { start: u64, length: u64 }, + Suffix { length: u64 }, +} + +/// Observed identity of one backend object generation. +/// +/// S3's `version_id` is deliberately observational. Follow-up S3 and HTTP +/// reads use `If-Match`; local reads compare a capability-relative file +/// fingerprint. Variants make cross-backend cache/generation reuse impossible. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub(crate) enum ObjectGeneration { + S3 { + etag: String, + version_id: Option, + total_len: u64, + }, + Http { + etag: String, + total_len: u64, + }, + Local { + fingerprint: String, + total_len: u64, + }, +} + +impl ObjectGeneration { + pub(crate) fn backend_kind(&self) -> StorageBackendKind { + match self { + Self::S3 { .. } => StorageBackendKind::S3, + Self::Local { .. } => StorageBackendKind::Local, + Self::Http { .. } => StorageBackendKind::Http, + } + } + + pub(crate) fn total_len(&self) -> u64 { + match self { + Self::S3 { total_len, .. } + | Self::Local { total_len, .. } + | Self::Http { total_len, .. } => *total_len, + } + } + + pub(crate) fn validator_is_empty(&self) -> bool { + match self { + Self::S3 { etag, .. } | Self::Http { etag, .. } => etag.is_empty(), + Self::Local { fingerprint, .. } => fingerprint.is_empty(), + } + } + + pub(crate) fn s3_etag(&self) -> Option<&str> { + match self { + Self::S3 { etag, .. } => Some(etag), + _ => None, + } + } + + pub(crate) fn http_etag(&self) -> Option<&str> { + match self { + Self::Http { etag, .. } => Some(etag), + _ => None, + } + } +} + +/// Complete identity for a query-local object read/cache entry. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub(crate) struct ReadIdentity { + pub key: String, + pub range: ReadRange, + pub generation: Option, +} + +impl ReadIdentity { + pub(crate) fn whole(key: impl Into) -> Self { + Self { + key: key.into(), + range: ReadRange::Whole, + generation: None, + } + } + + pub(crate) fn exact(key: impl Into, start: u64, length: u64) -> ZarrFdwResult { + validate_nonempty_range(length)?; + start.checked_add(length - 1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("storage byte range end overflows u64".to_string()) + })?; + Ok(Self { + key: key.into(), + range: ReadRange::Exact { start, length }, + generation: None, + }) + } + + pub(crate) fn suffix(key: impl Into, length: u64) -> ZarrFdwResult { + validate_nonempty_range(length)?; + Ok(Self { + key: key.into(), + range: ReadRange::Suffix { length }, + generation: None, + }) + } + + /// Apply an observed generation to a follow-up exact read. The resulting + /// request is sent with `If-Match` so an index and payload can never come + /// from different shard generations. + pub(crate) fn with_generation(mut self, generation: ObjectGeneration) -> Self { + self.generation = Some(generation); + self + } +} + +/// One exactly validated backend range response. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct RangedObject { + /// Resolved exact range plus the generation observed in the response. + pub identity: ReadIdentity, + pub total_len: u64, + pub bytes: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct ContentRange { + start: u64, + end: u64, + total: u64, +} + +enum Interruptible { + Ready(T), + Interrupted, +} + +fn postgres_interrupt_pending() -> bool { + // PostgreSQL declares this as volatile sig_atomic_t because signal + // handlers update it asynchronously. + unsafe { std::ptr::read_volatile(&raw const pg_sys::InterruptPending) != 0 } +} + +fn process_postgres_interrupts() { + unsafe { + if postgres_interrupt_pending() { + pg_sys::ProcessInterrupts(); + } + } +} + +async fn await_interruptibly(future: F) -> Interruptible +where + F: Future, +{ + tokio::pin!(future); + if postgres_interrupt_pending() { + return Interruptible::Interrupted; + } + + let mut ticker = interval(INTERRUPT_POLL_INTERVAL); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + ticker.tick().await; + loop { + tokio::select! { + biased; + _ = ticker.tick() => { + if postgres_interrupt_pending() { + return Interruptible::Interrupted; + } + } + result = &mut future => return Interruptible::Ready(result), + } + } +} + +/// One bounded page of immediate child prefixes below a Zarr group. +/// +/// S3's delimiter keeps discovery at the metadata hierarchy level. Array +/// nodes are identified with an exact `.zarray` GET before listing, so slash- +/// separated chunk keys are never traversed as groups. +pub(crate) struct DirectoryPage { + pub child_prefixes: Vec, + pub next_continuation_token: Option, +} + +/// A parsed S3 `s3://bucket/prefix` location. +#[derive(Debug, Clone)] +pub(crate) struct StoreUrl { + pub bucket: String, + /// Object-key prefix (no leading `/`), e.g. `sentinel2/2025.zarr`. + pub prefix: String, +} + +impl StoreUrl { + pub fn parse(s: &str) -> ZarrFdwResult { + let uri = s.parse::().map_err(|_| { + ZarrFdwError::InvalidStoreUrl("expected s3://bucket/prefix".to_string()) + })?; + if uri.scheme_str() != Some("s3") || uri.host().is_none() { + return Err(ZarrFdwError::InvalidStoreUrl( + "expected s3://bucket/prefix".to_string(), + )); + } + let bucket = uri.host().expect("host checked above").to_owned(); + let prefix = uri.path().trim_matches('/').to_string(); + Ok(Self { bucket, prefix }) + } + + fn object_key(&self, key: &str) -> String { + join_key(&self.prefix, key) + } + + fn relative_key(&self, key: &str) -> Option { + let key = key.trim_matches('/'); + if self.prefix.is_empty() { + return Some(key.to_string()); + } + if key == self.prefix { + return Some(String::new()); + } + key.strip_prefix(&format!("{}/", self.prefix)) + .map(str::to_string) + } +} + +/// Join two S3 object-key fragments without introducing or duplicating `/`. +pub(crate) fn join_key(prefix: &str, key: &str) -> String { + let prefix = prefix.trim_matches('/'); + let key = key.trim_matches('/'); + match (prefix.is_empty(), key.is_empty()) { + (true, _) => key.to_string(), + (_, true) => prefix.to_string(), + (false, false) => format!("{prefix}/{key}"), + } +} + +enum ClientAuth { + Anonymous, + Static { + access_key: String, + secret_key: String, + }, + ProviderChain, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AuthMode { + Anonymous, + Direct, + Vault, + ProviderChain, +} + +fn boolean_option(options: &HashMap, name: &str) -> ZarrFdwResult { + match options.get(name).map(String::as_str) { + None | Some("false") => Ok(false), + Some("true") => Ok(true), + Some(_) => Err(ZarrFdwError::InvalidOptionValue { + option: name.to_string(), + message: "must be 'true' or 'false'".to_string(), + }), + } +} + +/// Validate the mutually exclusive authentication modes and return the one to +/// use. `anonymous=false` is neutral, so an option-free configuration keeps +/// the AWS SDK provider-chain behavior. +pub(crate) fn validate_auth_options(options: &HashMap) -> ZarrFdwResult { + let anonymous = boolean_option(options, "anonymous")?; + let _path_style_url = boolean_option(options, "path_style_url")?; + + let direct_id = options.contains_key("aws_access_key_id"); + let direct_secret = options.contains_key("aws_secret_access_key"); + let vault_id = options.contains_key("vault_access_key_id"); + let vault_secret = options.contains_key("vault_secret_access_key"); + + match (direct_id, direct_secret) { + (true, false) => { + require_option("aws_secret_access_key", options)?; + } + (false, true) => { + require_option("aws_access_key_id", options)?; + } + _ => {} + } + match (vault_id, vault_secret) { + (true, false) => { + require_option("vault_secret_access_key", options)?; + } + (false, true) => { + require_option("vault_access_key_id", options)?; + } + _ => {} + } + + for name in [ + "aws_access_key_id", + "aws_secret_access_key", + "vault_access_key_id", + "vault_secret_access_key", + ] { + if options + .get(name) + .is_some_and(|value| value.trim().is_empty()) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: name.to_string(), + message: "must not be empty".to_string(), + }); + } + } + + let has_direct = direct_id && direct_secret; + let has_vault = vault_id && vault_secret; + if anonymous && (has_direct || has_vault) { + return Err(ZarrFdwError::InvalidAuthenticationOptions( + "anonymous authentication cannot be combined with explicit credentials".to_string(), + )); + } + if has_direct && has_vault { + return Err(ZarrFdwError::InvalidAuthenticationOptions( + "direct and Vault credentials cannot both be configured".to_string(), + )); + } + + Ok(if anonymous { + AuthMode::Anonymous + } else if has_direct { + AuthMode::Direct + } else if has_vault { + AuthMode::Vault + } else { + AuthMode::ProviderChain + }) +} + +const S3_ONLY_OPTIONS: &[&str] = &[ + "anonymous", + "aws_access_key_id", + "aws_secret_access_key", + "vault_access_key_id", + "vault_secret_access_key", + "aws_region", + "endpoint_url", + "path_style_url", +]; + +const HTTP_INSECURE_OPTION: &str = "allow_insecure_http"; + +/// Validate the configured storage scheme and its backend-specific options. +/// +/// PostgreSQL privilege checks are deliberately separate: the DDL validator +/// checks the current user, while runtime construction checks the cataloged +/// foreign-server owner so delegated `USAGE` continues to work. +pub(crate) fn validate_store_options( + options: &HashMap, +) -> ZarrFdwResult { + let store_url = require_option("store_url", options)?; + if store_url.starts_with("s3://") { + if options.contains_key(HTTP_INSECURE_OPTION) { + return Err(ZarrFdwError::InvalidOptionValue { + option: HTTP_INSECURE_OPTION.to_string(), + message: "is only valid for http:// stores".to_string(), + }); + } + StoreUrl::parse(store_url)?; + validate_auth_options(options)?; + return Ok(StorageBackendKind::S3); + } + if store_url.starts_with("file:") { + LocalBackend::validate_url(store_url)?; + if let Some(option) = S3_ONLY_OPTIONS + .iter() + .find(|option| options.contains_key(**option)) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: (*option).to_string(), + message: "is only valid for s3:// stores".to_string(), + }); + } + if options.contains_key(HTTP_INSECURE_OPTION) { + return Err(ZarrFdwError::InvalidOptionValue { + option: HTTP_INSECURE_OPTION.to_string(), + message: "is only valid for http:// stores".to_string(), + }); + } + return Ok(StorageBackendKind::Local); + } + let scheme = store_url + .split_once(':') + .map(|(scheme, _)| scheme) + .unwrap_or_default(); + if scheme.eq_ignore_ascii_case("https") || scheme.eq_ignore_ascii_case("http") { + if let Some(option) = S3_ONLY_OPTIONS + .iter() + .find(|option| options.contains_key(**option)) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: (*option).to_string(), + message: "is only valid for s3:// stores".to_string(), + }); + } + let allow_insecure_http = boolean_option(options, HTTP_INSECURE_OPTION)?; + HttpBackend::validate_url(store_url, allow_insecure_http)?; + return Ok(StorageBackendKind::Http); + } + Err(ZarrFdwError::InvalidStoreUrl( + "storage URL scheme is unsupported".to_string(), + )) +} + +/// Enforce the `CREATE/ALTER SERVER` privilege boundary for trusted stores. +pub(crate) fn validate_store_definition_privilege(kind: StorageBackendKind) -> ZarrFdwResult<()> { + if !unsafe { pg_sys::superuser() } { + return match kind { + StorageBackendKind::Local => Err(ZarrFdwError::FileStoreDefinitionRequiresSuperuser), + StorageBackendKind::Http => Err(ZarrFdwError::HttpStoreDefinitionRequiresSuperuser), + StorageBackendKind::S3 => Ok(()), + }; + } + Ok(()) +} + +fn validate_file_server_owner(server_oid: pg_sys::Oid) -> ZarrFdwResult<()> { + if server_oid == pg_sys::Oid::INVALID { + return Err(ZarrFdwError::FileStoreOwnerRequiresSuperuser); + } + let server = unsafe { pg_sys::GetForeignServer(server_oid) }; + if server.is_null() || !unsafe { pg_sys::superuser_arg((*server).owner) } { + return Err(ZarrFdwError::FileStoreOwnerRequiresSuperuser); + } + Ok(()) +} + +fn validate_http_server_owner(server_oid: pg_sys::Oid) -> ZarrFdwResult<()> { + if server_oid == pg_sys::Oid::INVALID { + return Err(ZarrFdwError::HttpStoreOwnerRequiresSuperuser); + } + let server = unsafe { pg_sys::GetForeignServer(server_oid) }; + if server.is_null() || !unsafe { pg_sys::superuser_arg((*server).owner) } { + return Err(ZarrFdwError::HttpStoreOwnerRequiresSuperuser); + } + Ok(()) +} + +struct S3Backend { + client: s3::Client, + url: StoreUrl, +} + +impl StorageBackend for S3Backend { + fn kind(&self) -> StorageBackendKind { + StorageBackendKind::S3 + } + + fn get_object_owned(&self, key: String, max_bytes: usize) -> StoreFuture>> { + get_object_optional_owned(self.client.clone(), self.url.clone(), key, max_bytes) + .boxed_local() + } + + fn get_range_owned(&self, identity: ReadIdentity) -> StoreFuture> { + get_object_range_owned(self.client.clone(), self.url.clone(), identity).boxed_local() + } + + fn list_directory_page_owned( + &self, + path: String, + continuation_token: Option, + ) -> StoreFuture { + list_directory_page_owned( + self.client.clone(), + self.url.clone(), + path, + continuation_token, + ) + .boxed_local() + } +} + +/// Query-local storage coordinator. +pub(crate) struct ZarrStore { + pub rt: Runtime, + backend: Arc, +} + +impl ZarrStore { + fn block_on_interruptibly(&self, future: F) -> ZarrFdwResult + where + F: Future>, + { + match self.rt.block_on(await_interruptibly(future)) { + Interruptible::Ready(result) => result, + Interruptible::Interrupted => { + // `await_interruptibly` and its owned SDK future have been + // dropped before PostgreSQL is allowed to raise ERROR. + process_postgres_interrupts(); + Err(ZarrFdwError::InvalidMetadata( + "query interruption was requested".to_string(), + )) + } + } + } + + /// Build the configured store from `CREATE SERVER` options. + pub fn new(server: &ForeignServer) -> ZarrFdwResult { + // Cannot use create_async_runtime() as the runtime needs multiple threads + let rt = tokio::runtime::Runtime::new() + .map_err(CreateRuntimeError::FailedToCreateAsyncRuntime)?; + + let kind = validate_store_options(&server.options)?; + let store_url = require_option("store_url", &server.options)?; + let backend: Arc = match kind { + StorageBackendKind::S3 => { + let url = StoreUrl::parse(store_url)?; + let auth_mode = validate_auth_options(&server.options)?; + let client = match auth_mode { + AuthMode::Anonymous => { + Self::build_client(&rt, &server.options, ClientAuth::Anonymous) + } + AuthMode::Direct => { + let access_key = + require_option("aws_access_key_id", &server.options)?.to_string(); + let secret_key = + require_option("aws_secret_access_key", &server.options)?.to_string(); + Self::build_client( + &rt, + &server.options, + ClientAuth::Static { + access_key, + secret_key, + }, + ) + } + AuthMode::Vault => { + let vault_access_key_id = + require_option("vault_access_key_id", &server.options)?; + let vault_secret_access_key = + require_option("vault_secret_access_key", &server.options)?; + let access_key = + get_vault_secret(vault_access_key_id).ok_or_else(|| { + ZarrFdwError::VaultSecretNotFound { + option: "vault_access_key_id".to_string(), + } + })?; + let secret_key = + get_vault_secret(vault_secret_access_key).ok_or_else(|| { + ZarrFdwError::VaultSecretNotFound { + option: "vault_secret_access_key".to_string(), + } + })?; + Self::build_client( + &rt, + &server.options, + ClientAuth::Static { + access_key, + secret_key, + }, + ) + } + AuthMode::ProviderChain => { + Self::build_client(&rt, &server.options, ClientAuth::ProviderChain) + } + }; + Arc::new(S3Backend { client, url }) + } + StorageBackendKind::Local => { + validate_file_server_owner(server.server_oid)?; + Arc::new(LocalBackend::new(store_url)?) + } + StorageBackendKind::Http => { + validate_http_server_owner(server.server_oid)?; + let allow_insecure_http = boolean_option(&server.options, HTTP_INSECURE_OPTION)?; + Arc::new(HttpBackend::new(store_url, allow_insecure_http)?) + } + }; + + Ok(Self { rt, backend }) + } + + pub(crate) fn backend_kind(&self) -> StorageBackendKind { + self.backend.kind() + } + + pub(crate) fn backend_label(&self) -> &'static str { + self.backend_kind().label() + } + + pub(crate) fn effective_max_concurrent_reads(&self, configured: usize) -> usize { + self.backend_kind() + .effective_max_concurrent_reads(configured) + } + + /// Fail before any backend request when hierarchy discovery is requested + /// from a readable-but-non-listable HTTP store. + pub(crate) fn require_listing(&self) -> ZarrFdwResult<()> { + if self.backend_kind() == StorageBackendKind::Http { + return Err(ZarrFdwError::UnsupportedExecutionFeature( + http::HTTP_LISTING_UNSUPPORTED.to_string(), + )); + } + Ok(()) + } + + fn build_client(rt: &Runtime, opts: &HashMap, auth: ClientAuth) -> s3::Client { + let region = require_option_or("aws_region", opts, "us-east-1"); + let mut config_loader = aws_config::defaults(BehaviorVersion::latest()) + .region(s3::config::Region::new(region.to_string())); + config_loader = match auth { + ClientAuth::Anonymous => config_loader.no_credentials(), + ClientAuth::Static { + access_key, + secret_key, + } => config_loader.credentials_provider(s3::config::Credentials::new( + access_key, secret_key, None, None, "zarr_fdw", + )), + ClientAuth::ProviderChain => config_loader, + }; + // endpoint_url not supported as env var in rust https://github.com/awslabs/aws-sdk-rust/issues/932 + if let Some(endpoint_url) = opts.get("endpoint_url") { + if endpoint_url.ends_with('/') { + config_loader = config_loader.endpoint_url(endpoint_url); + } else { + config_loader = config_loader.endpoint_url(format!("{endpoint_url}/")); + }; + } + + let config = rt.block_on(config_loader.load()); + let path_style_url = opts.get("path_style_url").map(|s| s.as_str()) == Some("true"); + let mut s3_config_builder = s3::config::Builder::from(&config); + s3_config_builder = s3_config_builder.force_path_style(path_style_url); + s3::Client::from_conf(s3_config_builder.build()) + } + + /// Fetch a full object by key (relative to the store root), returning + /// `None` only when the backend explicitly reports that it is absent. + pub async fn get_object_optional( + &self, + key: &str, + max_bytes: usize, + ) -> ZarrFdwResult>> { + self.get_object_optional_owned(key.to_string(), max_bytes) + .await + } + + /// Create an owned object fetch future suitable for the foreground + /// prefetch window. The future owns cloned backend state and therefore + /// never borrows the PostgreSQL scan object while it is queued. + pub fn get_object_optional_owned( + &self, + key: String, + max_bytes: usize, + ) -> StoreFuture>> { + self.backend.get_object_owned(key, max_bytes) + } + + /// Synchronous optional fetch used for sparse Zarr chunks. + pub fn get_object_optional_sync( + &self, + key: &str, + max_bytes: usize, + ) -> ZarrFdwResult>> { + self.block_on_interruptibly(self.get_object_optional(key, max_bytes)) + } + + /// Create an owned, exactly bounded storage range request. This method + /// never falls back to reading the complete object. + pub(crate) fn get_object_range_owned( + &self, + identity: ReadIdentity, + ) -> StoreFuture> { + self.backend.get_range_owned(identity) + } + + /// Synchronous range fetch for eager coordinate and shard-index reads. + pub(crate) fn get_object_range_sync( + &self, + identity: ReadIdentity, + ) -> ZarrFdwResult> { + self.block_on_interruptibly(self.get_object_range_owned(identity)) + } + + /// List one bounded page of immediate child prefixes below `path`. + /// + /// The caller owns pagination and global discovery limits. Only common + /// prefixes are returned; ordinary objects (including dot-separated chunk + /// keys) are deliberately ignored. + pub async fn list_directory_page( + &self, + path: &str, + continuation_token: Option, + ) -> ZarrFdwResult { + self.backend + .list_directory_page_owned(path.to_string(), continuation_token) + .await + } + + pub fn list_directory_page_sync( + &self, + path: &str, + continuation_token: Option, + ) -> ZarrFdwResult { + self.block_on_interruptibly(self.list_directory_page(path, continuation_token)) + } +} + +async fn get_object_optional_owned( + client: s3::Client, + url: StoreUrl, + key: String, + max_bytes: usize, +) -> ZarrFdwResult>> { + let full_key = url.object_key(&key); + + let resp = match client + .get_object() + .bucket(&url.bucket) + .key(&full_key) + .send() + .await + { + Ok(response) => response, + Err(error) => { + let modeled_no_such_key = error + .as_service_error() + .is_some_and(|error| error.is_no_such_key()); + let status = error + .raw_response() + .map(|response| response.status().as_u16()); + if is_missing_object_response(modeled_no_such_key, status) { + return Ok(None); + } + return Err(error.into()); + } + }; + if let Some(content_length) = resp.content_length { + let content_length = usize::try_from(content_length).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' length exceeds this platform's index capacity" + )) + })?; + if content_length > max_bytes { + return Err(ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' is {content_length} bytes, exceeding the read limit of {max_bytes}" + ))); + } + } + Ok(Some( + read_bounded_object(resp.body.into_async_read(), max_bytes, &full_key).await?, + )) +} + +async fn list_directory_page_owned( + client: s3::Client, + url: StoreUrl, + path: String, + continuation_token: Option, +) -> ZarrFdwResult { + let key = url.object_key(&path); + let prefix = if key.is_empty() { + String::new() + } else { + format!("{}/", key.trim_end_matches('/')) + }; + let response = client + .list_objects_v2() + .bucket(&url.bucket) + .prefix(prefix) + .delimiter("/") + .max_keys(1000) + .set_continuation_token(continuation_token) + .send() + .await?; + + let mut child_prefixes = response + .common_prefixes + .unwrap_or_default() + .into_iter() + .filter_map(|entry| entry.prefix) + .filter_map(|entry| url.relative_key(&entry)) + .map(|entry| entry.trim_matches('/').to_string()) + .collect::>(); + child_prefixes.sort(); + child_prefixes.dedup(); + + let next_continuation_token = if response.is_truncated.unwrap_or(false) { + response.next_continuation_token + } else { + None + }; + if response.is_truncated.unwrap_or(false) && next_continuation_token.is_none() { + return Err(ZarrFdwError::InvalidMetadata( + "S3 returned a truncated listing without a continuation token".to_string(), + )); + } + + Ok(DirectoryPage { + child_prefixes, + next_continuation_token, + }) +} + +async fn get_object_range_owned( + client: s3::Client, + url: StoreUrl, + identity: ReadIdentity, +) -> ZarrFdwResult> { + if identity + .generation + .as_ref() + .is_some_and(|generation| generation.backend_kind() != StorageBackendKind::S3) + { + return Err(ZarrFdwError::InvalidMetadata( + "storage object generation belongs to a different backend".to_string(), + )); + } + let full_key = url.object_key(&identity.key); + let range_header = range_header(&identity.range)?; + let expected_length = range_length(&identity.range)?; + let expected_length_usize = usize::try_from(expected_length).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "range read for object '{full_key}' exceeds this platform's index capacity" + )) + })?; + + let mut request = client + .get_object() + .bucket(&url.bucket) + .key(&full_key) + .range(range_header); + if let Some(generation) = &identity.generation { + let etag = generation.s3_etag().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "storage object generation belongs to a different backend".to_string(), + ) + })?; + request = request.if_match(etag.to_string()); + } + let resp = match request.send().await { + Ok(response) => response, + Err(error) => { + let modeled_no_such_key = error + .as_service_error() + .is_some_and(|error| error.is_no_such_key()); + let status = error + .raw_response() + .map(|response| response.status().as_u16()); + if is_missing_object_response(modeled_no_such_key, status) { + return if identity.generation.is_some() { + Err(ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' changed while reading a shard (generation-conditioned S3 range is now missing)" + ))) + } else { + Ok(None) + }; + } + if status == Some(412) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' changed while reading a shard (S3 If-Match precondition failed)" + ))); + } + return Err(error.into()); + } + }; + + let content_range = + required_content_range(resp.content_range.as_deref()).map_err(|message| { + ZarrFdwError::InvalidMetadata(format!( + "invalid Content-Range for object '{full_key}': {message}" + )) + })?; + validate_content_range(&identity.range, content_range).map_err(|message| { + ZarrFdwError::InvalidMetadata(format!( + "invalid Content-Range for object '{full_key}': {message}" + )) + })?; + + let content_length = resp.content_length.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "range response for object '{full_key}' omitted Content-Length" + )) + })?; + let content_length = u64::try_from(content_length).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "range response for object '{full_key}' has a negative Content-Length" + )) + })?; + if content_length != expected_length { + return Err(ZarrFdwError::InvalidMetadata(format!( + "range response for object '{full_key}' has Content-Length {content_length}, expected exactly {expected_length} bytes" + ))); + } + + let etag = resp + .e_tag + .clone() + .filter(|etag| !etag.is_empty()) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "range response for object '{full_key}' omitted ETag required for shard consistency" + )) + })?; + if let Some(expected) = &identity.generation + && expected.s3_etag() != Some(etag.as_str()) + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' changed while reading a shard (response ETag did not match If-Match)" + ))); + } + if let Some(expected) = &identity.generation + && expected.total_len() != content_range.total + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "object '{full_key}' changed while reading a shard (Content-Range total {} did not match indexed object length {})", + content_range.total, + expected.total_len() + ))); + } + let generation = ObjectGeneration::S3 { + etag, + version_id: resp.version_id.clone(), + total_len: content_range.total, + }; + let bytes = read_bounded_object( + resp.body.into_async_read(), + expected_length_usize, + &full_key, + ) + .await?; + if bytes.len() != expected_length_usize { + return Err(ZarrFdwError::InvalidMetadata(format!( + "range response for object '{full_key}' returned {} body bytes, expected exactly {expected_length} bytes", + bytes.len() + ))); + } + + Ok(Some(RangedObject { + identity: ReadIdentity { + key: identity.key, + range: ReadRange::Exact { + start: content_range.start, + length: expected_length, + }, + generation: Some(generation), + }, + total_len: content_range.total, + bytes, + })) +} + +fn validate_nonempty_range(length: u64) -> ZarrFdwResult<()> { + if length == 0 { + return Err(ZarrFdwError::InvalidMetadata( + "storage byte range length must be greater than zero".to_string(), + )); + } + Ok(()) +} + +fn range_length(range: &ReadRange) -> ZarrFdwResult { + match range { + ReadRange::Whole => Err(ZarrFdwError::InvalidMetadata( + "whole-object identity cannot be used for a storage range read".to_string(), + )), + ReadRange::Exact { length, .. } | ReadRange::Suffix { length } => { + validate_nonempty_range(*length)?; + Ok(*length) + } + } +} + +fn range_header(range: &ReadRange) -> ZarrFdwResult { + match range { + ReadRange::Whole => Err(ZarrFdwError::InvalidMetadata( + "whole-object identity cannot be used for a storage range read".to_string(), + )), + ReadRange::Exact { start, length } => { + validate_nonempty_range(*length)?; + let end = start.checked_add(length - 1).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("storage byte range end overflows u64".to_string()) + })?; + Ok(format!("bytes={start}-{end}")) + } + ReadRange::Suffix { length } => { + validate_nonempty_range(*length)?; + Ok(format!("bytes=-{length}")) + } + } +} + +fn parse_content_range(value: &str) -> Result { + let bytes = value + .strip_prefix("bytes ") + .ok_or_else(|| "expected canonical 'bytes START-END/TOTAL'".to_string())?; + let (range, total) = bytes + .split_once('/') + .ok_or_else(|| "expected canonical 'bytes START-END/TOTAL'".to_string())?; + if total == "*" { + return Err("wildcard total length is not accepted".to_string()); + } + let (start, end) = range + .split_once('-') + .ok_or_else(|| "expected canonical 'bytes START-END/TOTAL'".to_string())?; + let start = parse_canonical_u64(start, "start")?; + let end = parse_canonical_u64(end, "end")?; + let total = parse_canonical_u64(total, "total")?; + if start > end { + return Err("range start exceeds range end".to_string()); + } + if end >= total { + return Err("range end must be smaller than total object length".to_string()); + } + Ok(ContentRange { start, end, total }) +} + +fn required_content_range(value: Option<&str>) -> Result { + let value = value.ok_or_else(|| { + "header is absent; refusing a full-object fallback for a range request".to_string() + })?; + parse_content_range(value) +} + +fn parse_canonical_u64(value: &str, label: &str) -> Result { + if value.is_empty() + || (value.len() > 1 && value.starts_with('0')) + || !value.bytes().all(|byte| byte.is_ascii_digit()) + { + return Err(format!("{label} is not a canonical unsigned integer")); + } + value + .parse::() + .map_err(|_| format!("{label} exceeds u64")) +} + +fn validate_content_range(range: &ReadRange, actual: ContentRange) -> Result<(), String> { + let actual_length = actual + .end + .checked_sub(actual.start) + .and_then(|length| length.checked_add(1)) + .ok_or_else(|| "returned byte length overflows u64".to_string())?; + match range { + ReadRange::Whole => Err("whole-object request cannot have a range response".to_string()), + ReadRange::Exact { start, length } => { + let expected_end = start + .checked_add(length.checked_sub(1).ok_or_else(|| { + "requested byte range length must be greater than zero".to_string() + })?) + .ok_or_else(|| "requested byte range end overflows u64".to_string())?; + if actual.start != *start || actual.end != expected_end || actual_length != *length { + return Err(format!( + "returned bytes {}-{}/{}, expected exactly {start}-{expected_end}", + actual.start, actual.end, actual.total + )); + } + Ok(()) + } + ReadRange::Suffix { length } => { + if actual_length != *length || actual.end.checked_add(1) != Some(actual.total) { + return Err(format!( + "returned suffix bytes {}-{}/{}, expected exactly the final {length} bytes", + actual.start, actual.end, actual.total + )); + } + Ok(()) + } + } +} + +async fn read_bounded_object( + mut reader: R, + max_bytes: usize, + key: &str, +) -> ZarrFdwResult> +where + R: tokio::io::AsyncRead + Unpin, +{ + let mut buf = Vec::new(); + let mut block = [0u8; 64 * 1024]; + while buf.len() < max_bytes { + let remaining = max_bytes - buf.len(); + let read_len = remaining.min(block.len()); + let count = reader.read(&mut block[..read_len]).await?; + if count == 0 { + return Ok(buf); + } + buf.try_reserve_exact(count).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not grow the object read buffer for '{key}'" + )) + })?; + buf.extend_from_slice(&block[..count]); + } + + // Probe one byte beyond the cap without growing the result vector. + let mut extra = [0u8; 1]; + if reader.read(&mut extra).await? != 0 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "object '{key}' exceeds the read limit of {max_bytes} bytes" + ))); + } + Ok(buf) +} + +fn is_missing_object_response(modeled_no_such_key: bool, status: Option) -> bool { + modeled_no_such_key || status == Some(404) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn options(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect() + } + + #[test] + fn bounded_object_reader_accepts_limit_and_rejects_limit_plus_one() { + let rt = tokio::runtime::Runtime::new().unwrap(); + assert_eq!( + rt.block_on(read_bounded_object( + std::io::Cursor::new(b"1234"), + 4, + "exact" + )) + .unwrap(), + b"1234" + ); + assert!( + rt.block_on(read_bounded_object( + std::io::Cursor::new(b"12345"), + 4, + "too-large" + )) + .is_err() + ); + } + + #[test] + fn range_headers_use_checked_inclusive_http_bounds() { + assert_eq!(ReadIdentity::whole("/array/c/0/").key, "/array/c/0/"); + assert_eq!( + range_header(&ReadRange::Exact { + start: 10, + length: 4 + }) + .unwrap(), + "bytes=10-13" + ); + assert_eq!( + range_header(&ReadRange::Suffix { length: 68 }).unwrap(), + "bytes=-68" + ); + assert!( + range_header(&ReadRange::Exact { + start: 0, + length: 0 + }) + .is_err() + ); + assert!( + range_header(&ReadRange::Exact { + start: u64::MAX, + length: 2 + }) + .is_err() + ); + } + + #[test] + fn content_range_parser_is_canonical_and_checked() { + assert_eq!( + parse_content_range("bytes 10-13/100").unwrap(), + ContentRange { + start: 10, + end: 13, + total: 100 + } + ); + for invalid in [ + "bytes 10-13/*", + "bytes 010-13/100", + "Bytes 10-13/100", + "bytes 13-10/100", + "bytes 10-100/100", + "bytes 10-13/", + ] { + assert!(parse_content_range(invalid).is_err(), "{invalid}"); + } + assert!(required_content_range(None).is_err()); + assert!(required_content_range(Some("not-a-content-range")).is_err()); + } + + #[test] + fn content_range_must_match_exact_or_full_suffix_request() { + let exact = ReadRange::Exact { + start: 10, + length: 4, + }; + assert!( + validate_content_range( + &exact, + ContentRange { + start: 10, + end: 13, + total: 100 + } + ) + .is_ok() + ); + assert!( + validate_content_range( + &exact, + ContentRange { + start: 10, + end: 12, + total: 100 + } + ) + .is_err() + ); + + let suffix = ReadRange::Suffix { length: 4 }; + assert!( + validate_content_range( + &suffix, + ContentRange { + start: 96, + end: 99, + total: 100 + } + ) + .is_ok() + ); + assert!( + validate_content_range( + &suffix, + ContentRange { + start: 0, + end: 2, + total: 3 + } + ) + .is_err() + ); + } + + #[test] + fn parse_store_url() { + let u = StoreUrl::parse("s3://cinecube/sentinel2/2025.zarr").unwrap(); + assert_eq!(u.bucket, "cinecube"); + assert_eq!(u.prefix, "sentinel2/2025.zarr"); + } + + #[test] + fn validates_backend_specific_store_options() { + assert_eq!( + validate_store_options(&options(&[("store_url", "s3://bucket/root")])).unwrap(), + StorageBackendKind::S3 + ); + assert_eq!( + validate_store_options(&options(&[("store_url", "file:///tmp/zarr")])).unwrap(), + StorageBackendKind::Local + ); + assert_eq!( + validate_store_options(&options(&[( + "store_url", + "https://objects.example.test/root.zarr" + )])) + .unwrap(), + StorageBackendKind::Http + ); + assert_eq!( + validate_store_options(&options(&[ + ("store_url", "http://objects.example.test/root.zarr"), + ("allow_insecure_http", "true"), + ])) + .unwrap(), + StorageBackendKind::Http + ); + let error = validate_store_options(&options(&[ + ("store_url", "file:///tmp/zarr"), + ("anonymous", "true"), + ])) + .unwrap_err(); + assert_eq!( + error.to_string(), + "invalid value for option 'anonymous': is only valid for s3:// stores" + ); + assert!( + validate_store_options(&options(&[( + "store_url", + "http://objects.example.test/root.zarr" + )])) + .is_err() + ); + assert!( + validate_store_options(&options(&[ + ("store_url", "https://objects.example.test/root.zarr"), + ("anonymous", "true"), + ])) + .is_err() + ); + let error = validate_store_options(&options(&[( + "store_url", + "HTTPS://user:secret@objects.example.test/root.zarr", + )])) + .unwrap_err() + .to_string(); + assert!(!error.contains("user")); + assert!(!error.contains("secret")); + } + + #[test] + fn backend_read_concurrency_is_truthful() { + assert_eq!(StorageBackendKind::S3.effective_max_concurrent_reads(8), 8); + assert_eq!( + StorageBackendKind::Http.effective_max_concurrent_reads(8), + 8 + ); + assert_eq!( + StorageBackendKind::Local.effective_max_concurrent_reads(8), + 1 + ); + } + + #[test] + fn generation_backend_tags_are_disjoint() { + let s3 = ObjectGeneration::S3 { + etag: "\"etag\"".to_string(), + version_id: None, + total_len: 1, + }; + let local = ObjectGeneration::Local { + fingerprint: "1:2:1:3:4:5:6".to_string(), + total_len: 1, + }; + let http = ObjectGeneration::Http { + etag: "\"etag\"".to_string(), + total_len: 1, + }; + assert_eq!(s3.backend_kind(), StorageBackendKind::S3); + assert_eq!(local.backend_kind(), StorageBackendKind::Local); + assert_eq!(http.backend_kind(), StorageBackendKind::Http); + assert_ne!(s3, local); + assert_ne!(s3, http); + } + + #[test] + fn parse_store_url_root() { + let u = StoreUrl::parse("s3://cinecube/").unwrap(); + assert_eq!(u.bucket, "cinecube"); + assert_eq!(u.prefix, ""); + } + + #[test] + fn parse_store_url_without_slash() { + let u = StoreUrl::parse("s3://cinecube").unwrap(); + assert_eq!(u.bucket, "cinecube"); + assert_eq!(u.prefix, ""); + } + + #[test] + fn reject_non_s3() { + assert!(StoreUrl::parse("https://example.com/x").is_err()); + } + + #[test] + fn parse_store_url_trailing_slash() { + let u = StoreUrl::parse("s3://bucket/k.alt/").unwrap(); + assert_eq!(u.prefix, "k.alt"); + } + + #[test] + fn joins_object_key_prefix_once() { + let u = StoreUrl::parse("s3://bucket/grid/data.zarr").unwrap(); + assert_eq!( + u.object_key("longitude/.zarray"), + "grid/data.zarr/longitude/.zarray" + ); + assert_eq!(u.object_key("/x/0/"), "grid/data.zarr/x/0"); + } + + #[test] + fn joins_root_object_key_without_leading_slash() { + let u = StoreUrl::parse("s3://bucket/").unwrap(); + assert_eq!(u.object_key("/.zarray"), ".zarray"); + } + + #[test] + fn converts_list_prefixes_back_to_store_relative_paths() { + let url = StoreUrl::parse("s3://warehouse/zarr/e2e.zarr").unwrap(); + + assert_eq!( + url.relative_key("zarr/e2e.zarr/nested/raw/"), + Some("nested/raw".to_string()) + ); + assert_eq!(url.relative_key("zarr/e2e.zarr"), Some(String::new())); + assert_eq!(url.relative_key("other/prefix"), None); + } + + #[test] + fn only_no_such_key_or_http_404_is_optional_absence() { + assert!(is_missing_object_response(true, None)); + assert!(is_missing_object_response(false, Some(404))); + assert!(!is_missing_object_response(false, None)); + assert!(!is_missing_object_response(false, Some(403))); + assert!(!is_missing_object_response(false, Some(500))); + } + + #[test] + fn selects_each_supported_auth_mode() { + assert_eq!( + validate_auth_options(&options(&[])).unwrap(), + AuthMode::ProviderChain + ); + assert_eq!( + validate_auth_options(&options(&[("anonymous", "true")])).unwrap(), + AuthMode::Anonymous + ); + assert_eq!( + validate_auth_options(&options(&[("anonymous", "false")])).unwrap(), + AuthMode::ProviderChain + ); + assert_eq!( + validate_auth_options(&options(&[ + ("aws_access_key_id", "key"), + ("aws_secret_access_key", "secret"), + ])) + .unwrap(), + AuthMode::Direct + ); + assert_eq!( + validate_auth_options(&options(&[ + ("vault_access_key_id", "key-id"), + ("vault_secret_access_key", "secret-id"), + ])) + .unwrap(), + AuthMode::Vault + ); + assert_eq!( + validate_auth_options(&options(&[ + ("anonymous", "false"), + ("aws_access_key_id", "key"), + ("aws_secret_access_key", "secret"), + ])) + .unwrap(), + AuthMode::Direct + ); + assert_eq!( + validate_auth_options(&options(&[ + ("anonymous", "false"), + ("vault_access_key_id", "key-id"), + ("vault_secret_access_key", "secret-id"), + ])) + .unwrap(), + AuthMode::Vault + ); + } + + #[test] + fn rejects_invalid_boolean_options() { + for (name, value) in [("anonymous", "TRUE"), ("path_style_url", "yes")] { + let err = validate_auth_options(&options(&[(name, value)])).unwrap_err(); + assert_eq!( + err.to_string(), + format!("invalid value for option '{name}': must be 'true' or 'false'") + ); + } + } + + #[test] + fn rejects_partial_credential_pairs() { + for (present, missing) in [ + ("aws_access_key_id", "aws_secret_access_key"), + ("aws_secret_access_key", "aws_access_key_id"), + ("vault_access_key_id", "vault_secret_access_key"), + ("vault_secret_access_key", "vault_access_key_id"), + ] { + let err = validate_auth_options(&options(&[(present, "value")])).unwrap_err(); + assert_eq!( + err.to_string(), + format!("required option `{missing}` is not specified") + ); + } + } + + #[test] + fn rejects_empty_credentials_without_echoing_values() { + for (pairs, empty_option) in [ + ( + vec![ + ("aws_access_key_id", ""), + ("aws_secret_access_key", "secret"), + ], + "aws_access_key_id", + ), + ( + vec![ + ("vault_access_key_id", "key-id"), + ("vault_secret_access_key", " "), + ], + "vault_secret_access_key", + ), + ] { + let err = validate_auth_options(&options(&pairs)).unwrap_err(); + assert_eq!( + err.to_string(), + format!("invalid value for option '{empty_option}': must not be empty") + ); + } + } + + #[test] + fn rejects_conflicting_authentication_modes() { + let anonymous_and_direct = options(&[ + ("anonymous", "true"), + ("aws_access_key_id", "key"), + ("aws_secret_access_key", "secret"), + ]); + assert_eq!( + validate_auth_options(&anonymous_and_direct) + .unwrap_err() + .to_string(), + "invalid authentication options: anonymous authentication cannot be combined with explicit credentials" + ); + + let anonymous_and_vault = options(&[ + ("anonymous", "true"), + ("vault_access_key_id", "key-id"), + ("vault_secret_access_key", "secret-id"), + ]); + assert_eq!( + validate_auth_options(&anonymous_and_vault) + .unwrap_err() + .to_string(), + "invalid authentication options: anonymous authentication cannot be combined with explicit credentials" + ); + + let direct_and_vault = options(&[ + ("aws_access_key_id", "key"), + ("aws_secret_access_key", "secret"), + ("vault_access_key_id", "key-id"), + ("vault_secret_access_key", "secret-id"), + ]); + assert_eq!( + validate_auth_options(&direct_and_vault) + .unwrap_err() + .to_string(), + "invalid authentication options: direct and Vault credentials cannot both be configured" + ); + } + + #[test] + fn routing_options_do_not_change_authentication_mode() { + let opts = options(&[ + ("endpoint_url", "http://localhost:9000"), + ("path_style_url", "true"), + ]); + assert_eq!( + validate_auth_options(&opts).unwrap(), + AuthMode::ProviderChain + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/store/http.rs b/wrappers/src/fdw/zarr_fdw/store/http.rs new file mode 100644 index 000000000..4c02546f0 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/store/http.rs @@ -0,0 +1,625 @@ +//! Trusted, anonymous, read-only HTTP(S) object storage. + +use futures_util::FutureExt; +use reqwest::header::{ + ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_RANGE, ETAG, IF_MATCH, RANGE, +}; +use reqwest::{Client, Response, StatusCode, Url, redirect}; +use std::time::Duration; + +use super::{ + ContentRange, DirectoryPage, ObjectGeneration, RangedObject, ReadIdentity, ReadRange, + StorageBackend, StorageBackendKind, StoreFuture, parse_canonical_u64, range_header, + range_length, required_content_range, validate_content_range, +}; +use crate::fdw::zarr_fdw::{ZarrFdwError, ZarrFdwResult}; + +const MAX_CONFIGURED_URL_BYTES: usize = 8 * 1024; +const MAX_STORAGE_KEY_BYTES: usize = 8 * 1024; +const CONNECT_TIMEOUT: Duration = Duration::from_secs(30); + +pub(crate) const HTTP_LISTING_UNSUPPORTED: &str = "HTTP(S) Zarr stores do not support hierarchy listing; configure an explicit array path or OME multiscale selection"; + +#[derive(Clone)] +pub(super) struct HttpBackend { + client: Client, + root: HttpStoreUrl, +} + +#[derive(Clone, Debug)] +struct HttpStoreUrl { + root: Url, +} + +impl HttpBackend { + pub(super) fn validate_url(raw: &str, allow_insecure_http: bool) -> ZarrFdwResult<()> { + HttpStoreUrl::parse(raw, allow_insecure_http).map(|_| ()) + } + + pub(super) fn new(raw: &str, allow_insecure_http: bool) -> ZarrFdwResult { + let root = HttpStoreUrl::parse(raw, allow_insecure_http)?; + let client = Client::builder() + .use_rustls_tls() + .min_tls_version(reqwest::tls::Version::TLS_1_2) + .https_only(!allow_insecure_http) + .redirect(redirect::Policy::none()) + .referer(false) + .no_proxy() + .no_gzip() + .no_brotli() + .no_deflate() + .no_zstd() + .retry(reqwest::retry::never()) + .connect_timeout(CONNECT_TIMEOUT) + .user_agent(concat!("wrappers-zarr/", env!("CARGO_PKG_VERSION"))) + .build() + .map_err(|_| http_access_error("", "client initialization failed"))?; + Ok(Self { client, root }) + } + + async fn read_whole(self, key: String, max_bytes: usize) -> ZarrFdwResult>> { + let object_url = self.root.object_url(&key)?; + let response = self + .client + .get(object_url) + .header(ACCEPT_ENCODING, "identity") + .send() + .await + .map_err(|_| http_access_error(&key, "transport error"))?; + + match response.status() { + StatusCode::NOT_FOUND => return Ok(None), + StatusCode::OK => {} + status if status.is_redirection() => return Err(redirect_error(&key)), + status => { + return Err(protocol_error( + &key, + format!( + "whole-object response returned status {}, expected 200 or 404", + status.as_u16() + ), + )); + } + } + + validate_identity_encoding(&response, &key)?; + let declared_length = optional_content_length(&response, &key, "whole-object")?; + if declared_length.is_some_and(|length| length > max_bytes as u64) { + return Err(protocol_error( + &key, + format!( + "object length {} exceeds the read limit of {max_bytes} bytes", + declared_length.expect("checked Some above") + ), + )); + } + let bytes = read_bounded_response(response, max_bytes, &key).await?; + if let Some(declared_length) = declared_length + && u64::try_from(bytes.len()).ok() != Some(declared_length) + { + return Err(protocol_error( + &key, + format!( + "whole-object body returned {} bytes, Content-Length declared {declared_length}", + bytes.len() + ), + )); + } + Ok(Some(bytes)) + } + + async fn read_range(self, identity: ReadIdentity) -> ZarrFdwResult> { + if identity + .generation + .as_ref() + .is_some_and(|generation| generation.backend_kind() != StorageBackendKind::Http) + { + return Err(ZarrFdwError::InvalidMetadata( + "storage object generation belongs to a different backend".to_string(), + )); + } + + let object_url = self.root.object_url(&identity.key)?; + let requested_range = range_header(&identity.range)?; + let expected_length = range_length(&identity.range)?; + let expected_length_usize = usize::try_from(expected_length).map_err(|_| { + protocol_error( + &identity.key, + "range length exceeds this platform's index capacity", + ) + })?; + let mut request = self + .client + .get(object_url) + .header(ACCEPT_ENCODING, "identity") + .header(RANGE, requested_range); + if let Some(generation) = &identity.generation { + let etag = generation.http_etag().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "storage object generation belongs to a different backend".to_string(), + ) + })?; + request = request.header(IF_MATCH, etag); + } + let response = request + .send() + .await + .map_err(|_| http_access_error(&identity.key, "transport error"))?; + + match response.status() { + StatusCode::PARTIAL_CONTENT => {} + StatusCode::NOT_FOUND if identity.generation.is_none() => return Ok(None), + StatusCode::NOT_FOUND | StatusCode::PRECONDITION_FAILED + if identity.generation.is_some() => + { + return Err(changed_while_reading( + &identity.key, + "generation-conditioned object is missing or If-Match failed", + )); + } + StatusCode::RANGE_NOT_SATISFIABLE if identity.generation.is_some() => { + return Err(changed_while_reading( + &identity.key, + "generation-conditioned range is no longer satisfiable", + )); + } + StatusCode::OK => { + return Err(protocol_error( + &identity.key, + "range response returned status 200; server ignored Range", + )); + } + status if status.is_redirection() => return Err(redirect_error(&identity.key)), + status => { + return Err(protocol_error( + &identity.key, + format!( + "range response returned status {}, expected 206 or 404", + status.as_u16() + ), + )); + } + } + + validate_identity_encoding(&response, &identity.key)?; + let content_range = response_content_range(&response, &identity.key)?; + validate_content_range(&identity.range, content_range).map_err(|message| { + protocol_error(&identity.key, format!("invalid Content-Range: {message}")) + })?; + if let Some(content_length) = optional_content_length(&response, &identity.key, "range")? + && content_length != expected_length + { + return Err(protocol_error( + &identity.key, + format!( + "range Content-Length {content_length} does not equal the requested {expected_length} bytes" + ), + )); + } + let etag = required_strong_etag(&response, &identity.key)?; + if let Some(expected) = &identity.generation { + if expected.http_etag() != Some(etag.as_str()) { + return Err(changed_while_reading( + &identity.key, + "response ETag did not match If-Match", + )); + } + if expected.total_len() != content_range.total { + return Err(changed_while_reading( + &identity.key, + "Content-Range total changed after the shard index read", + )); + } + } + + let bytes = read_bounded_response(response, expected_length_usize, &identity.key).await?; + if bytes.len() != expected_length_usize { + return Err(protocol_error( + &identity.key, + format!( + "range body returned {} bytes, expected exactly {expected_length} bytes", + bytes.len() + ), + )); + } + let generation = ObjectGeneration::Http { + etag, + total_len: content_range.total, + }; + Ok(Some(RangedObject { + identity: ReadIdentity { + key: identity.key, + range: ReadRange::Exact { + start: content_range.start, + length: expected_length, + }, + generation: Some(generation), + }, + total_len: content_range.total, + bytes, + })) + } +} + +impl StorageBackend for HttpBackend { + fn kind(&self) -> StorageBackendKind { + StorageBackendKind::Http + } + + fn get_object_owned(&self, key: String, max_bytes: usize) -> StoreFuture>> { + let backend = self.clone(); + async move { backend.read_whole(key, max_bytes).await }.boxed_local() + } + + fn get_range_owned(&self, identity: ReadIdentity) -> StoreFuture> { + let backend = self.clone(); + async move { backend.read_range(identity).await }.boxed_local() + } + + fn list_directory_page_owned( + &self, + _path: String, + _continuation_token: Option, + ) -> StoreFuture { + async move { + Err(ZarrFdwError::UnsupportedExecutionFeature( + HTTP_LISTING_UNSUPPORTED.to_string(), + )) + } + .boxed_local() + } +} + +impl HttpStoreUrl { + fn parse(raw: &str, allow_insecure_http: bool) -> ZarrFdwResult { + if raw.len() > MAX_CONFIGURED_URL_BYTES + || raw.chars().any(char::is_control) + || raw.contains('\\') + { + return Err(invalid_url( + "URL is too long or contains backslash or control characters", + )); + } + reject_ambiguous_percent_encoding(raw) + .map_err(|message| invalid_url(format!("configured path {message}")))?; + let mut root = Url::parse(raw).map_err(|_| invalid_url("URL is not valid"))?; + match root.scheme() { + "https" if allow_insecure_http => { + return Err(ZarrFdwError::InvalidOptionValue { + option: "allow_insecure_http".to_string(), + message: "may be 'true' only for http:// stores".to_string(), + }); + } + "https" => {} + "http" if allow_insecure_http => {} + "http" => { + return Err(ZarrFdwError::InvalidOptionValue { + option: "allow_insecure_http".to_string(), + message: "must be 'true' for http:// stores".to_string(), + }); + } + _ => { + return Err(invalid_url( + "scheme must be https:// or explicitly enabled http://", + )); + } + } + if root.cannot_be_a_base() + || root.host_str().is_none() + || !root.username().is_empty() + || root.password().is_some() + || root.query().is_some() + || root.fragment().is_some() + { + return Err(invalid_url( + "host is required; credentials, query, and fragment are not allowed", + )); + } + let raw_path = configured_raw_path(raw); + if root.port().is_none() && root.port_or_known_default().is_none() + || raw_path + .split('/') + .any(|component| matches!(component, "." | "..")) + || root.path()[1..].contains("//") + { + return Err(invalid_url("port or path is invalid")); + } + root.path_segments_mut() + .map_err(|_| invalid_url("URL cannot be used as an object root"))? + .pop_if_empty() + .push(""); + Ok(Self { root }) + } + + fn object_url(&self, key: &str) -> ZarrFdwResult { + validate_http_key(key)?; + let mut url = self.root.clone(); + { + let mut segments = url + .path_segments_mut() + .map_err(|_| invalid_url("URL cannot be used as an object root"))?; + segments.pop_if_empty(); + for component in key.split('/') { + segments.push(component); + } + } + Ok(url) + } +} + +fn configured_raw_path(raw: &str) -> &str { + let Some((_, after_scheme)) = raw.split_once("://") else { + return ""; + }; + let Some(path_start) = after_scheme.find('/') else { + return ""; + }; + let path = &after_scheme[path_start..]; + let end = path.find(['?', '#']).unwrap_or(path.len()); + &path[..end] +} + +fn validate_http_key(key: &str) -> ZarrFdwResult<()> { + if key.is_empty() || key.len() > MAX_STORAGE_KEY_BYTES { + return Err(invalid_key("must be nonempty and at most 8192 bytes")); + } + if key.starts_with('/') || key.ends_with('/') { + return Err(invalid_key("must be a relative object path")); + } + if key + .chars() + .any(|character| matches!(character, '\\' | '?' | '#') || character.is_control()) + { + return Err(invalid_key( + "backslash, query, fragment, and control characters are not allowed", + )); + } + if key + .split('/') + .any(|component| component.is_empty() || matches!(component, "." | "..")) + { + return Err(invalid_key( + "empty, '.' and '..' components are not allowed", + )); + } + reject_ambiguous_percent_encoding(key).map_err(invalid_key) +} + +fn reject_ambiguous_percent_encoding(value: &str) -> Result<(), &'static str> { + let lower = value.to_ascii_lowercase(); + if lower.contains("%2f") || lower.contains("%5c") || lower.contains("%2e") { + return Err("must not contain percent-encoded slash, backslash, or dot"); + } + Ok(()) +} + +fn response_content_range(response: &Response, key: &str) -> ZarrFdwResult { + let value = single_header(response, CONTENT_RANGE.as_str(), key, "Content-Range")?; + required_content_range(value.as_deref()) + .map_err(|message| protocol_error(key, format!("invalid Content-Range: {message}"))) +} + +fn optional_content_length( + response: &Response, + key: &str, + phase: &str, +) -> ZarrFdwResult> { + let Some(value) = single_header(response, CONTENT_LENGTH.as_str(), key, "Content-Length")? + else { + return Ok(None); + }; + parse_canonical_u64(&value, "Content-Length") + .map(Some) + .map_err(|message| { + protocol_error(key, format!("invalid {phase} Content-Length: {message}")) + }) +} + +fn validate_identity_encoding(response: &Response, key: &str) -> ZarrFdwResult<()> { + let Some(value) = single_header(response, CONTENT_ENCODING.as_str(), key, "Content-Encoding")? + else { + return Ok(()); + }; + if !value.eq_ignore_ascii_case("identity") { + return Err(protocol_error( + key, + "response Content-Encoding is not identity", + )); + } + Ok(()) +} + +fn required_strong_etag(response: &Response, key: &str) -> ZarrFdwResult { + let value = single_header(response, ETAG.as_str(), key, "ETag")?.ok_or_else(|| { + protocol_error( + key, + "range response omitted strong ETag required for shard consistency", + ) + })?; + validate_strong_etag(&value) + .map_err(|message| protocol_error(key, format!("invalid range ETag: {message}")))?; + Ok(value) +} + +fn validate_strong_etag(value: &str) -> Result<(), &'static str> { + if value.starts_with("W/") { + return Err("weak ETag is not accepted"); + } + let bytes = value.as_bytes(); + if bytes.len() < 2 || bytes.first() != Some(&b'"') || bytes.last() != Some(&b'"') { + return Err("expected a quoted strong ETag"); + } + if bytes[1..bytes.len() - 1] + .iter() + .any(|byte| *byte == b'"' || *byte < 0x21 || *byte == 0x7f) + { + return Err("strong ETag contains invalid opaque-tag bytes"); + } + Ok(()) +} + +fn single_header( + response: &Response, + name: &str, + key: &str, + display_name: &str, +) -> ZarrFdwResult> { + let mut values = response.headers().get_all(name).iter(); + let Some(value) = values.next() else { + return Ok(None); + }; + if values.next().is_some() { + return Err(protocol_error( + key, + format!("response has multiple {display_name} headers"), + )); + } + let value = value + .to_str() + .map_err(|_| protocol_error(key, format!("response has invalid {display_name} header")))?; + Ok(Some(value.to_string())) +} + +async fn read_bounded_response( + mut response: Response, + max_bytes: usize, + key: &str, +) -> ZarrFdwResult> { + let mut bytes = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .map_err(|_| http_access_error(key, "response body read failed"))? + { + let next = bytes + .len() + .checked_add(chunk.len()) + .ok_or_else(|| protocol_error(key, "response body length exceeds platform capacity"))?; + if next > max_bytes { + return Err(protocol_error( + key, + format!("response body exceeds the read limit of {max_bytes} bytes"), + )); + } + bytes + .try_reserve_exact(chunk.len()) + .map_err(|_| protocol_error(key, "could not grow the bounded HTTP response buffer"))?; + bytes.extend_from_slice(&chunk); + tokio::task::yield_now().await; + } + Ok(bytes) +} + +fn invalid_url(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidHttpStoreUrl(message.into()) +} + +fn invalid_key(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidHttpStorageKey(message.into()) +} + +fn protocol_error(key: &str, message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!( + "HTTP storage protocol error for object '{key}': {}", + message.into() + )) +} + +fn redirect_error(key: &str) -> ZarrFdwError { + protocol_error(key, "redirect response was rejected") +} + +fn changed_while_reading(key: &str, category: &'static str) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!( + "HTTP storage object '{key}' changed while reading a shard: {category}" + )) +} + +fn http_access_error(key: impl Into, category: &'static str) -> ZarrFdwError { + ZarrFdwError::HttpStorageAccess { + key: key.into(), + category, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn validates_https_and_explicit_insecure_http() { + assert!(HttpStoreUrl::parse("https://example.test/root.zarr", false).is_ok()); + assert!(HttpStoreUrl::parse("http://example.test/root.zarr", false).is_err()); + assert!(HttpStoreUrl::parse("http://example.test/root.zarr", true).is_ok()); + assert!(HttpStoreUrl::parse("https://example.test/root.zarr", true).is_err()); + } + + #[test] + fn rejects_credentials_query_fragment_and_ambiguous_paths() { + for invalid in [ + "https://user@example.test/root.zarr", + "https://example.test/root.zarr?token=secret", + "https://example.test/root.zarr#fragment", + "https://example.test/root%2fzarr", + "https://example.test/root//zarr", + ] { + assert!(HttpStoreUrl::parse(invalid, false).is_err(), "{invalid}"); + } + } + + #[test] + fn object_keys_append_components_without_replacing_origin() { + let root = HttpStoreUrl::parse("https://example.test/base/root.zarr", false).unwrap(); + let url = root.object_url("nested values/c/0").unwrap(); + assert_eq!( + url.as_str(), + "https://example.test/base/root.zarr/nested%20values/c/0" + ); + assert_eq!(url.host_str(), Some("example.test")); + for invalid in [ + "", + "/absolute", + "trailing/", + "../escape", + "a//b", + "a\\b", + "a?query", + "%2e%2e/escape", + ] { + assert!(root.object_url(invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn strong_etags_are_quoted_and_not_weak() { + assert!(validate_strong_etag("\"generation-1\"").is_ok()); + for invalid in [ + "generation-1", + "W/\"generation-1\"", + "\"bad value\"", + "\"a\"b\"", + ] { + assert!(validate_strong_etag(invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn backend_generation_is_typed() { + let generation = ObjectGeneration::Http { + etag: "\"generation-1\"".to_string(), + total_len: 10, + }; + assert_eq!(generation.backend_kind(), StorageBackendKind::Http); + assert_eq!(generation.http_etag(), Some("\"generation-1\"")); + assert_eq!(generation.s3_etag(), None); + assert_eq!(generation.total_len(), 10); + } + + #[test] + fn listing_error_is_stable() { + assert_eq!( + HTTP_LISTING_UNSUPPORTED, + "HTTP(S) Zarr stores do not support hierarchy listing; configure an explicit array path or OME multiscale selection" + ); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/store/local.rs b/wrappers/src/fdw/zarr_fdw/store/local.rs new file mode 100644 index 000000000..22dd3c8f9 --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/store/local.rs @@ -0,0 +1,795 @@ +//! Capability-confined read-only local filesystem storage. + +use cap_std::ambient_authority; +use cap_std::fs::{Dir, Metadata, MetadataExt, OpenOptions, OpenOptionsExt}; +use futures_util::FutureExt; +use std::collections::BinaryHeap; +use std::io::{ErrorKind, Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use url::Url; + +use super::{ + DirectoryPage, ObjectGeneration, RangedObject, ReadIdentity, ReadRange, StorageBackend, + StorageBackendKind, StoreFuture, join_key, +}; +use crate::fdw::zarr_fdw::{ZarrFdwError, ZarrFdwResult}; + +const READ_BLOCK_BYTES: usize = 64 * 1024; +const LIST_PAGE_ENTRIES: usize = 1_000; +const LIST_YIELD_ENTRIES: usize = 64; + +#[derive(Clone)] +pub(super) struct LocalBackend { + root_path: PathBuf, + root: Arc>>>, +} + +impl LocalBackend { + pub(super) fn validate_url(raw: &str) -> ZarrFdwResult { + parse_file_url(raw) + } + + pub(super) fn new(raw: &str) -> ZarrFdwResult { + Ok(Self { + root_path: parse_file_url(raw)?, + root: Arc::new(Mutex::new(None)), + }) + } + + fn root(&self) -> ZarrFdwResult> { + let mut root = self + .root + .lock() + .map_err(|_| local_access_error("", "configured root lock was poisoned"))?; + if let Some(root) = root.as_ref() { + return Ok(Arc::clone(root)); + } + + let opened = + Dir::open_ambient_dir(&self.root_path, ambient_authority()).map_err(|error| { + local_access_error("", format!("could not open configured root: {error}")) + })?; + let metadata = opened.dir_metadata().map_err(|error| { + local_access_error( + "", + format!("could not inspect configured root: {error}"), + ) + })?; + if !metadata.is_dir() { + return Err(local_access_error( + "", + "configured root is not a directory", + )); + } + let opened = Arc::new(opened); + *root = Some(Arc::clone(&opened)); + Ok(opened) + } + + async fn read_whole(self, key: String, max_bytes: usize) -> ZarrFdwResult>> { + validate_storage_key(&key, false)?; + let root = self.root()?; + let Some((mut file, before)) = open_regular_file(&root, &key)? else { + return Ok(None); + }; + let generation = local_generation(&before); + let length = usize::try_from(generation.total_len()).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "local storage object '{key}' length exceeds this platform's index capacity" + )) + })?; + if length > max_bytes { + return Err(ZarrFdwError::InvalidMetadata(format!( + "local storage object '{key}' is {length} bytes, exceeding the read limit of {max_bytes}" + ))); + } + + let bytes = read_exact_blocks(&mut file, length, &key).await?; + let mut extra = [0_u8; 1]; + if file + .read(&mut extra) + .map_err(|error| local_access_error(&key, error.to_string()))? + != 0 + { + return Err(changed_while_reading(&key)); + } + let after = file.metadata().map_err(|error| { + local_access_error(&key, format!("could not inspect open file: {error}")) + })?; + ensure_generation(&key, &generation, &after)?; + Ok(Some(bytes)) + } + + async fn read_range(self, identity: ReadIdentity) -> ZarrFdwResult> { + validate_storage_key(&identity.key, false)?; + let root = self.root()?; + let Some((mut file, before)) = open_regular_file(&root, &identity.key)? else { + return if identity.generation.is_some() { + Err(changed_while_reading(&identity.key)) + } else { + Ok(None) + }; + }; + let generation = local_generation(&before); + if let Some(expected) = &identity.generation { + ensure_same_generation(&identity.key, expected, &generation)?; + } + + let (start, length) = + resolved_range(&identity.range, generation.total_len(), &identity.key)?; + let length_usize = usize::try_from(length).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "local range read for object '{}' exceeds this platform's index capacity", + identity.key + )) + })?; + file.seek(SeekFrom::Start(start)) + .map_err(|error| local_access_error(&identity.key, error.to_string()))?; + let bytes = read_exact_blocks(&mut file, length_usize, &identity.key).await?; + let after = file.metadata().map_err(|error| { + local_access_error( + &identity.key, + format!("could not inspect open file: {error}"), + ) + })?; + ensure_generation(&identity.key, &generation, &after)?; + + Ok(Some(RangedObject { + identity: ReadIdentity { + key: identity.key, + range: ReadRange::Exact { start, length }, + generation: Some(generation.clone()), + }, + total_len: generation.total_len(), + bytes, + })) + } + + async fn list_page( + self, + path: String, + continuation_token: Option, + ) -> ZarrFdwResult { + validate_storage_key(&path, true)?; + let after = continuation_child(&path, continuation_token.as_deref())?; + let root = self.root()?; + let entries = if path.is_empty() { + root.entries() + } else { + root.read_dir(&path) + }; + let entries = match entries { + Ok(entries) => entries, + Err(error) if error.kind() == ErrorKind::NotFound => { + return Ok(DirectoryPage { + child_prefixes: Vec::new(), + next_continuation_token: None, + }); + } + Err(error) => { + return Err(local_access_error( + display_path(&path), + format!("could not list directory: {error}"), + )); + } + }; + + let mut names = BinaryHeap::with_capacity(LIST_PAGE_ENTRIES + 1); + let mut scanned = 0_usize; + for entry in entries { + let entry = entry.map_err(|error| { + local_access_error( + display_path(&path), + format!("could not read directory entry: {error}"), + ) + })?; + scanned = scanned.saturating_add(1); + if scanned.is_multiple_of(LIST_YIELD_ENTRIES) { + tokio::task::yield_now().await; + } + let file_type = entry.file_type().map_err(|error| { + local_access_error( + display_path(&path), + format!("could not inspect directory entry: {error}"), + ) + })?; + if !file_type.is_dir() { + continue; + } + let Ok(name) = entry.file_name().into_string() else { + continue; + }; + if validate_child_name(&name).is_err() + || after + .as_ref() + .is_some_and(|after| name.as_str() <= after.as_str()) + { + continue; + } + if names.len() < LIST_PAGE_ENTRIES + 1 { + names.push(name); + } else if names + .peek() + .is_some_and(|largest| name.as_str() < largest.as_str()) + { + names.pop(); + names.push(name); + } + } + + let mut names = names.into_sorted_vec(); + let has_more = names.len() > LIST_PAGE_ENTRIES; + if has_more { + names.truncate(LIST_PAGE_ENTRIES); + } + let child_prefixes = names + .iter() + .map(|name| join_key(&path, name)) + .collect::>(); + let next_continuation_token = if has_more { + child_prefixes.last().cloned() + } else { + None + }; + Ok(DirectoryPage { + child_prefixes, + next_continuation_token, + }) + } +} + +impl StorageBackend for LocalBackend { + fn kind(&self) -> StorageBackendKind { + StorageBackendKind::Local + } + + fn get_object_owned(&self, key: String, max_bytes: usize) -> StoreFuture>> { + let backend = self.clone(); + async move { backend.read_whole(key, max_bytes).await }.boxed_local() + } + + fn get_range_owned(&self, identity: ReadIdentity) -> StoreFuture> { + let backend = self.clone(); + async move { backend.read_range(identity).await }.boxed_local() + } + + fn list_directory_page_owned( + &self, + path: String, + continuation_token: Option, + ) -> StoreFuture { + let backend = self.clone(); + async move { backend.list_page(path, continuation_token).await }.boxed_local() + } +} + +fn parse_file_url(raw: &str) -> ZarrFdwResult { + if !raw.starts_with("file:///") { + return Err(ZarrFdwError::InvalidFileStoreUrl( + "expected file:///absolute/path".to_string(), + )); + } + let url = Url::parse(raw).map_err(|_| { + ZarrFdwError::InvalidFileStoreUrl("expected file:///absolute/path".to_string()) + })?; + if url.scheme() != "file" + || url.cannot_be_a_base() + || url.host_str().is_some() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err(ZarrFdwError::InvalidFileStoreUrl( + "host, credentials, query, and fragment are not allowed".to_string(), + )); + } + let path = url.to_file_path().map_err(|_| { + ZarrFdwError::InvalidFileStoreUrl("path must be absolute UTF-8".to_string()) + })?; + let path_text = path.to_str().ok_or_else(|| { + ZarrFdwError::InvalidFileStoreUrl("path must be absolute UTF-8".to_string()) + })?; + if !path.is_absolute() || path_text.contains('\0') { + return Err(ZarrFdwError::InvalidFileStoreUrl( + "path must be absolute UTF-8".to_string(), + )); + } + Ok(path) +} + +fn validate_storage_key(key: &str, allow_empty: bool) -> ZarrFdwResult<()> { + if key.is_empty() { + return if allow_empty { + Ok(()) + } else { + Err(invalid_key("must not be empty")) + }; + } + if key.starts_with('/') || key.ends_with('/') || Path::new(key).is_absolute() { + return Err(invalid_key("must be a relative object path")); + } + if key.contains('\\') || key.chars().any(|character| character.is_ascii_control()) { + return Err(invalid_key( + "backslash and ASCII control characters are not allowed", + )); + } + if key + .split('/') + .any(|component| component.is_empty() || matches!(component, "." | "..")) + { + return Err(invalid_key( + "empty, '.' and '..' components are not allowed", + )); + } + Ok(()) +} + +fn validate_child_name(name: &str) -> ZarrFdwResult<()> { + validate_storage_key(name, false)?; + if name.contains('/') { + return Err(invalid_key( + "directory child and continuation token must contain one path component", + )); + } + Ok(()) +} + +fn continuation_child(path: &str, token: Option<&str>) -> ZarrFdwResult> { + let Some(token) = token else { + return Ok(None); + }; + validate_storage_key(token, false)?; + let child = if path.is_empty() { + token + } else { + token + .strip_prefix(&format!("{path}/")) + .ok_or_else(|| invalid_key("continuation token does not belong to the directory"))? + }; + validate_child_name(child)?; + Ok(Some(child.to_string())) +} + +fn open_regular_file( + root: &Dir, + key: &str, +) -> ZarrFdwResult> { + let mut options = OpenOptions::new(); + options.read(true).custom_flags(libc::O_NONBLOCK); + let file = match root.open_with(key, &options) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(local_access_error(key, error.to_string())), + }; + let metadata = file.metadata().map_err(|error| { + local_access_error(key, format!("could not inspect open file: {error}")) + })?; + if !metadata.is_file() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "local storage object '{key}' is not a regular file" + ))); + } + Ok(Some((file, metadata))) +} + +async fn read_exact_blocks( + file: &mut cap_std::fs::File, + expected: usize, + key: &str, +) -> ZarrFdwResult> { + let mut bytes = Vec::new(); + bytes.try_reserve_exact(expected).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate {expected} bytes for local storage object '{key}'" + )) + })?; + let mut block = [0_u8; READ_BLOCK_BYTES]; + while bytes.len() < expected { + let remaining = expected - bytes.len(); + let read_len = remaining.min(block.len()); + let count = file + .read(&mut block[..read_len]) + .map_err(|error| local_access_error(key, error.to_string()))?; + if count == 0 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "local storage object '{key}' returned {} bytes, expected exactly {expected} bytes", + bytes.len() + ))); + } + bytes.extend_from_slice(&block[..count]); + tokio::task::yield_now().await; + } + Ok(bytes) +} + +fn resolved_range(range: &ReadRange, total: u64, key: &str) -> ZarrFdwResult<(u64, u64)> { + let (start, length) = match range { + ReadRange::Whole => { + return Err(ZarrFdwError::InvalidMetadata( + "whole-object identity cannot be used for a storage range read".to_string(), + )); + } + ReadRange::Exact { start, length } => (*start, *length), + ReadRange::Suffix { length } => { + let start = total.checked_sub(*length).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "local suffix range for object '{key}' exceeds object length {total}" + )) + })?; + (start, *length) + } + }; + if length == 0 { + return Err(ZarrFdwError::InvalidMetadata( + "storage byte range length must be greater than zero".to_string(), + )); + } + let end = start.checked_add(length).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("storage byte range end overflows u64".to_string()) + })?; + if end > total { + return Err(ZarrFdwError::InvalidMetadata(format!( + "local byte range {start}..{end} for object '{key}' exceeds object length {total}" + ))); + } + Ok((start, length)) +} + +fn local_generation(metadata: &Metadata) -> ObjectGeneration { + let total_len = metadata.size(); + ObjectGeneration::Local { + fingerprint: format!( + "{}:{}:{total_len}:{}:{}:{}:{}", + metadata.dev(), + metadata.ino(), + metadata.mtime(), + metadata.mtime_nsec(), + metadata.ctime(), + metadata.ctime_nsec() + ), + total_len, + } +} + +fn ensure_generation( + key: &str, + expected: &ObjectGeneration, + actual: &Metadata, +) -> ZarrFdwResult<()> { + ensure_same_generation(key, expected, &local_generation(actual)) +} + +fn ensure_same_generation( + key: &str, + expected: &ObjectGeneration, + actual: &ObjectGeneration, +) -> ZarrFdwResult<()> { + if expected.backend_kind() != StorageBackendKind::Local + || actual.backend_kind() != StorageBackendKind::Local + || expected != actual + { + return Err(changed_while_reading(key)); + } + Ok(()) +} + +fn changed_while_reading(key: &str) -> ZarrFdwError { + ZarrFdwError::InvalidMetadata(format!( + "local storage object '{key}' changed while reading" + )) +} + +fn invalid_key(message: impl Into) -> ZarrFdwError { + ZarrFdwError::InvalidLocalStorageKey(message.into()) +} + +fn local_access_error(key: impl Into, message: impl Into) -> ZarrFdwError { + ZarrFdwError::LocalStorageAccess { + key: key.into(), + message: message.into(), + } +} + +fn display_path(path: &str) -> &str { + if path.is_empty() { "/" } else { path } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::sync::atomic::{AtomicU64, Ordering}; + + static NEXT_ROOT: AtomicU64 = AtomicU64::new(0); + + struct TestRoot { + path: PathBuf, + } + + impl TestRoot { + fn new() -> Self { + let id = NEXT_ROOT.fetch_add(1, Ordering::Relaxed); + let path = + std::env::temp_dir().join(format!("zarr-fdw-local-{}-{id}", std::process::id())); + fs::create_dir(&path).unwrap(); + Self { path } + } + + fn backend(&self) -> LocalBackend { + let url = Url::from_directory_path(&self.path).unwrap(); + LocalBackend::new(url.as_str()).unwrap() + } + } + + impl Drop for TestRoot { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + } + } + + #[test] + fn file_urls_are_absolute_and_authority_free() { + assert!(parse_file_url("file:///tmp/example.zarr").is_ok()); + for invalid in [ + "file:relative", + "file://host/tmp/zarr", + "file:///tmp/zarr?query=1", + "file:///tmp/zarr#fragment", + "https://host/tmp/zarr", + ] { + assert!(parse_file_url(invalid).is_err(), "{invalid}"); + } + } + + #[test] + fn keys_reject_escaping_and_ambiguous_components() { + assert!(validate_storage_key("nested/array/zarr.json", false).is_ok()); + assert!(validate_storage_key("", true).is_ok()); + for invalid in [ + "", + "/absolute", + "trailing/", + "double//separator", + "./child", + "parent/../child", + "back\\slash", + "control\ncharacter", + "nul\0character", + ] { + assert!(validate_storage_key(invalid, false).is_err(), "{invalid:?}"); + } + assert_eq!( + continuation_child("nested", Some("nested/child")).unwrap(), + Some("child".to_string()) + ); + assert!(continuation_child("nested", Some("other/child")).is_err()); + assert!(continuation_child("", Some("nested/child")).is_err()); + } + + #[test] + fn whole_reads_are_bounded_and_missing_is_optional() { + let root = TestRoot::new(); + fs::write(root.path.join("object"), b"abcdef").unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + + assert_eq!( + runtime + .block_on(backend.clone().read_whole("object".to_string(), 6)) + .unwrap(), + Some(b"abcdef".to_vec()) + ); + assert!( + runtime + .block_on(backend.clone().read_whole("object".to_string(), 5)) + .is_err() + ); + assert_eq!( + runtime + .block_on(backend.read_whole("missing".to_string(), 6)) + .unwrap(), + None + ); + } + + #[test] + fn exact_and_suffix_reads_normalize_to_exact_ranges() { + let root = TestRoot::new(); + fs::write(root.path.join("object"), b"0123456789").unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + + let exact = runtime + .block_on( + backend + .clone() + .read_range(ReadIdentity::exact("object", 2, 4).unwrap()), + ) + .unwrap() + .unwrap(); + assert_eq!(exact.bytes.as_slice(), b"2345"); + assert_eq!( + exact.identity.range, + ReadRange::Exact { + start: 2, + length: 4 + } + ); + assert_eq!(exact.total_len, 10); + + let suffix = runtime + .block_on(backend.read_range(ReadIdentity::suffix("object", 3).unwrap())) + .unwrap() + .unwrap(); + assert_eq!(suffix.bytes.as_slice(), b"789"); + assert_eq!( + suffix.identity.range, + ReadRange::Exact { + start: 7, + length: 3 + } + ); + } + + #[test] + fn ranges_reject_bounds_nonregular_files_and_absolute_keys() { + let root = TestRoot::new(); + fs::write(root.path.join("object"), b"0123").unwrap(); + fs::create_dir(root.path.join("directory")).unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + + assert!( + runtime + .block_on( + backend + .clone() + .read_range(ReadIdentity::exact("object", 3, 2).unwrap()) + ) + .is_err() + ); + assert!( + runtime + .block_on( + backend + .clone() + .read_range(ReadIdentity::suffix("object", 5).unwrap()) + ) + .is_err() + ); + let error = runtime + .block_on( + backend + .clone() + .read_range(ReadIdentity::exact("directory", 0, 1).unwrap()), + ) + .unwrap_err(); + assert!(error.to_string().contains("is not a regular file")); + assert!( + runtime + .block_on(backend.read_range(ReadIdentity::exact("/object", 0, 1).unwrap())) + .is_err() + ); + } + + #[test] + fn generation_condition_detects_replacement_and_disappearance() { + let root = TestRoot::new(); + fs::write(root.path.join("shard"), b"abcdefgh").unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let indexed = runtime + .block_on( + backend + .clone() + .read_range(ReadIdentity::suffix("shard", 4).unwrap()), + ) + .unwrap() + .unwrap(); + let generation = indexed.identity.generation.unwrap(); + + fs::write(root.path.join("replacement"), b"ABCDEFGH").unwrap(); + fs::rename(root.path.join("replacement"), root.path.join("shard")).unwrap(); + let replaced = ReadIdentity::exact("shard", 0, 4) + .unwrap() + .with_generation(generation.clone()); + assert!( + runtime + .block_on(backend.clone().read_range(replaced)) + .unwrap_err() + .to_string() + .contains("changed while reading") + ); + + fs::remove_file(root.path.join("shard")).unwrap(); + let missing = ReadIdentity::exact("shard", 0, 4) + .unwrap() + .with_generation(generation); + assert!( + runtime + .block_on(backend.read_range(missing)) + .unwrap_err() + .to_string() + .contains("changed while reading") + ); + } + + #[test] + fn local_ranges_reject_s3_generations() { + let root = TestRoot::new(); + fs::write(root.path.join("object"), b"0123").unwrap(); + let backend = root.backend(); + let identity = ReadIdentity::exact("object", 0, 1) + .unwrap() + .with_generation(ObjectGeneration::S3 { + etag: "\"s3-etag\"".to_string(), + version_id: None, + total_len: 4, + }); + let runtime = tokio::runtime::Runtime::new().unwrap(); + assert!(runtime.block_on(backend.read_range(identity)).is_err()); + } + + #[test] + fn directory_listing_is_sorted_bounded_and_paginated() { + let root = TestRoot::new(); + for index in (0..=LIST_PAGE_ENTRIES).rev() { + fs::create_dir(root.path.join(format!("child-{index:04}"))).unwrap(); + } + fs::write(root.path.join("ordinary-object"), b"ignored").unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let first = runtime + .block_on(backend.clone().list_page(String::new(), None)) + .unwrap(); + assert_eq!(first.child_prefixes.len(), LIST_PAGE_ENTRIES); + assert_eq!(first.child_prefixes.first().unwrap(), "child-0000"); + assert_eq!(first.child_prefixes.last().unwrap(), "child-0999"); + assert_eq!(first.next_continuation_token.as_deref(), Some("child-0999")); + + let second = runtime + .block_on(backend.list_page(String::new(), first.next_continuation_token)) + .unwrap(); + assert_eq!(second.child_prefixes, vec!["child-1000".to_string()]); + assert!(second.next_continuation_token.is_none()); + } + + #[cfg(unix)] + #[test] + fn capability_root_rejects_symlink_escape() { + use std::os::unix::fs::symlink; + + let root = TestRoot::new(); + let outside = TestRoot::new(); + fs::write(outside.path.join("secret"), b"outside").unwrap(); + symlink(&outside.path, root.path.join("escape")).unwrap(); + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + assert!( + runtime + .block_on(backend.read_whole("escape/secret".to_string(), 64)) + .is_err() + ); + } + + #[cfg(unix)] + #[test] + fn nonblocking_open_rejects_fifo_without_reading_it() { + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt; + + let root = TestRoot::new(); + let fifo = root.path.join("fifo"); + let fifo = CString::new(fifo.as_os_str().as_bytes()).unwrap(); + assert_eq!(unsafe { libc::mkfifo(fifo.as_ptr(), 0o600) }, 0); + + let backend = root.backend(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let error = runtime + .block_on(backend.read_whole("fifo".to_string(), 64)) + .unwrap_err(); + assert!(error.to_string().contains("is not a regular file")); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/tests.rs b/wrappers/src/fdw/zarr_fdw/tests.rs new file mode 100644 index 000000000..e650f160b --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/tests.rs @@ -0,0 +1,7746 @@ +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::JsonB; + use pgrx::pg_test; + use pgrx::prelude::*; + + fn create_minio_e2e_server() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_e2e_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_e2e_server + FOREIGN DATA WRAPPER zarr_e2e_wrapper + OPTIONS ( + store_url 's3://warehouse/zarr/e2e.zarr', + aws_access_key_id 'admin', + aws_secret_access_key 'password', + aws_region 'us-east-1', + endpoint_url 'http://localhost:8000', + path_style_url 'true' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_v3_e2e_server() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_v3_e2e_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_v3_e2e_server + FOREIGN DATA WRAPPER zarr_v3_e2e_wrapper + OPTIONS ( + store_url 's3://warehouse/zarr/e2e-v3.zarr', + aws_access_key_id 'admin', + aws_secret_access_key 'password', + aws_region 'us-east-1', + endpoint_url 'http://localhost:8000', + path_style_url 'true' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_v3_e2e_table(table: &str, array_group: &str, decode_cf: bool) { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server(table, array_group, decode_cf); + } + + fn create_minio_v3_e2e_table_on_server(table: &str, array_group: &str, decode_cf: bool) { + let decode_option = if decode_cf { + ",\n decode_cf 'true'" + } else { + "" + }; + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + time timestamp with time zone, + y double precision, + x double precision, + value {value_type} + ) + SERVER zarr_v3_e2e_server + OPTIONS ( + array_group '{array_group}', + time_from_attrs 'true'{decode_option} + )"#, + value_type = if decode_cf { + "double precision" + } else { + "real" + }, + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_ome_v3_e2e_server() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_ome_v3_e2e_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_ome_v3_e2e_server + FOREIGN DATA WRAPPER zarr_ome_v3_e2e_wrapper + OPTIONS ( + store_url 's3://warehouse/zarr/e2e-ome-v3.zarr', + aws_access_key_id 'admin', + aws_secret_access_key 'password', + aws_region 'us-east-1', + endpoint_url 'http://localhost:8000', + path_style_url 'true' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_ome_v3_e2e_table(table: &str, level: usize) { + create_minio_ome_v3_e2e_server(); + create_minio_ome_v3_e2e_table_on_server(table, level); + } + + fn create_minio_ome_v3_e2e_table_on_server(table: &str, level: usize) { + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + y double precision, + x double precision, + value real + ) + SERVER zarr_ome_v3_e2e_server + OPTIONS ( + multiscale_group 'image', + multiscale_index '0', + multiscale_level '{level}' + )"# + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn local_fixture_url(fixture: &str) -> String { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("dockerfiles/s3/test_data/zarr") + .join(fixture) + .canonicalize() + .unwrap_or_else(|error| { + panic!("local Zarr fixture '{fixture}' is unavailable: {error}") + }); + let path = path + .to_str() + .expect("local Zarr fixture path must be valid UTF-8"); + assert!(path.starts_with('/'), "local Zarr fixture must be absolute"); + assert!( + !path.contains('\''), + "local Zarr fixture path cannot contain a SQL quote" + ); + format!("file://{path}") + } + + fn create_local_e2e_wrapper() { + Spi::run( + r#"CREATE FOREIGN DATA WRAPPER zarr_local_e2e_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + ) + .unwrap(); + } + + fn create_local_e2e_server(server: &str, fixture: &str) { + let store_url = local_fixture_url(fixture); + Spi::run(&format!( + r#"CREATE SERVER {server} + FOREIGN DATA WRAPPER zarr_local_e2e_wrapper + OPTIONS (store_url '{store_url}')"# + )) + .unwrap(); + } + + fn create_local_time_y_x_table(table: &str, server: &str, array_group: &str, decode_cf: bool) { + let (value_type, decode_option) = if decode_cf { + ( + "double precision", + ",\n decode_cf 'true'", + ) + } else { + ("real", "") + }; + Spi::run(&format!( + r#"CREATE FOREIGN TABLE {table} ( + time timestamp with time zone, + y double precision, + x double precision, + value {value_type} + ) + SERVER {server} + OPTIONS ( + array_group '{array_group}', + time_from_attrs 'true'{decode_option} + )"# + )) + .unwrap(); + } + + fn create_http_e2e_wrapper() { + Spi::run( + r#"CREATE FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + ) + .unwrap(); + } + + fn create_http_e2e_server(server: &str, case: &str, mode: &str, fixture: &str) { + let store_url = format!("http://127.0.0.1:8787/stores/{case}/{mode}/{fixture}"); + Spi::run(&format!( + r#"CREATE SERVER {server} + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS ( + store_url '{store_url}', + allow_insecure_http 'true' + )"# + )) + .unwrap(); + } + + fn create_http_time_y_x_table(table: &str, server: &str, array_group: &str, decode_cf: bool) { + let (value_type, decode_option) = if decode_cf { + ( + "double precision", + ",\n decode_cf 'true'", + ) + } else { + ("real", "") + }; + Spi::run(&format!( + r#"CREATE FOREIGN TABLE {table} ( + time timestamp with time zone, + y double precision, + x double precision, + value {value_type} + ) + SERVER {server} + OPTIONS ( + array_group '{array_group}', + time_from_attrs 'true'{decode_option} + )"# + )) + .unwrap(); + } + + fn create_http_ome_table(table: &str, server: &str, level: usize) { + Spi::run(&format!( + r#"CREATE FOREIGN TABLE {table} ( + y double precision, + x double precision, + value real + ) + SERVER {server} + OPTIONS ( + multiscale_group 'image', + multiscale_index '0', + multiscale_level '{level}' + )"# + )) + .unwrap(); + } + + fn http_case_stats(case: &str) -> serde_json::Value { + use std::io::{Read, Write}; + + let mut stream = std::net::TcpStream::connect(("127.0.0.1", 8787)).unwrap(); + stream + .set_read_timeout(Some(std::time::Duration::from_secs(2))) + .unwrap(); + write!( + stream, + "GET /__stats/{case} HTTP/1.1\r\nHost: 127.0.0.1:8787\r\nConnection: close\r\n\r\n" + ) + .unwrap(); + let mut response = Vec::new(); + stream.read_to_end(&mut response).unwrap(); + let separator = response + .windows(4) + .position(|window| window == b"\r\n\r\n") + .expect("HTTP stats response must contain a header terminator"); + let headers = std::str::from_utf8(&response[..separator]).unwrap(); + assert!(headers.starts_with("HTTP/1.1 200 "), "response: {headers}"); + serde_json::from_slice(&response[separator + 4..]).unwrap() + } + + fn capture_query_error(statement: &str) -> String { + Spi::connect_mut(|c| { + c.update( + r#"CREATE OR REPLACE FUNCTION pg_temp.capture_zarr_error(statement text) + RETURNS text + LANGUAGE plpgsql + AS $function$ + BEGIN + EXECUTE statement; + RETURN NULL; + EXCEPTION WHEN query_canceled THEN + RETURN SQLERRM; + WHEN OTHERS THEN + RETURN SQLERRM; + END + $function$"#, + None, + &[], + ) + .unwrap(); + c.select( + "SELECT pg_temp.capture_zarr_error($1) AS message", + None, + &[statement.into()], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("message") + .unwrap() + .expect("query must fail") + }) + } + + fn capture_query_error_as_role(role: &str, statement: &str) -> String { + Spi::connect_mut(|c| { + c.update( + r#"CREATE OR REPLACE FUNCTION pg_temp.capture_zarr_error(statement text) + RETURNS text + LANGUAGE plpgsql + AS $function$ + BEGIN + EXECUTE statement; + RETURN NULL; + EXCEPTION WHEN query_canceled THEN + RETURN SQLERRM; + WHEN OTHERS THEN + RETURN SQLERRM; + END + $function$"#, + None, + &[], + ) + .unwrap(); + c.update(&format!("SET ROLE {role}"), None, &[]).unwrap(); + let message = c + .select( + "SELECT pg_temp.capture_zarr_error($1) AS message", + None, + &[statement.into()], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("message") + .unwrap() + .expect("query must fail"); + c.update("RESET ROLE", None, &[]).unwrap(); + message + }) + } + + fn create_minio_e2e_table(table: &str, array_group: &str, value_type: &str) { + create_minio_e2e_table_with_cf(table, array_group, value_type, false); + } + + fn create_minio_e2e_table_with_cf( + table: &str, + array_group: &str, + value_type: &str, + decode_cf: bool, + ) { + let decode_option = if decode_cf { + ",\n decode_cf 'true'" + } else { + "" + }; + create_minio_e2e_table_with_options( + table, + array_group, + value_type, + &format!( + ",\n time_unit 'seconds',\n time_origin 'unix'{decode_option}" + ), + ); + } + + fn create_minio_e2e_table_with_options( + table: &str, + array_group: &str, + value_type: &str, + options: &str, + ) { + create_minio_e2e_server(); + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + x double precision, + y double precision, + time timestamp with time zone, + value {value_type} + ) + SERVER zarr_e2e_server + OPTIONS ( + array_group '{array_group}'{options} + )"# + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_generic4d_table(table: &str, band_type: &str) { + create_minio_generic4d_table_with_options(table, band_type, ""); + } + + fn create_minio_generic4d_table_with_options(table: &str, band_type: &str, options: &str) { + create_minio_e2e_server(); + create_minio_generic4d_table_on_server(table, band_type, options); + } + + fn create_minio_generic4d_table_on_server(table: &str, band_type: &str, options: &str) { + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + forecast_time timestamp with time zone, + level double precision, + band {band_type}, + channel double precision, + measurement real + ) + SERVER zarr_e2e_server + OPTIONS ( + array_group 'nested/generic4d', + time_from_attrs 'true'{options} + )"# + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_spatial2d_table(table: &str) { + create_minio_e2e_server(); + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + y double precision, + x double precision, + value real + ) + SERVER zarr_e2e_server + OPTIONS (array_group 'nested/spatial2d')"# + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_spatial2d_coordinate_only_table(table: &str) { + create_minio_e2e_server(); + Spi::connect_mut(|c| { + c.update( + &format!( + r#"CREATE FOREIGN TABLE {table} ( + y double precision, + x double precision + ) + SERVER zarr_e2e_server + OPTIONS (array_group 'nested/spatial2d')"# + ), + None, + &[], + ) + .unwrap(); + }); + } + + fn create_minio_spatial_time_table(table: &str, decode_cf: bool) { + let (value_type, decode_option) = if decode_cf { + ( + "double precision", + ",\n decode_cf 'true'", + ) + } else { + ("real", "") + }; + create_minio_e2e_table_with_options( + table, + "nested/raw", + value_type, + &format!(",\n time_from_attrs 'true'{decode_option}"), + ); + } + + fn create_minio_spatial_time_table_with_options(table: &str, options: &str) { + create_minio_e2e_table_with_options(table, "nested/raw", "real", options); + } + + fn install_postgis_in_test_schema() { + Spi::run("CREATE SCHEMA zarr_gis; CREATE EXTENSION postgis WITH SCHEMA zarr_gis").unwrap(); + } + + fn explain_lines(sql: &str) -> Vec { + Spi::connect(|c| { + c.select(&format!("EXPLAIN {sql}"), None, &[]) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect() + }) + } + + fn assert_aggregate_pushed_down(sql: &str) { + let plan = explain_lines(sql); + assert!( + plan.iter().any(|line| line.contains("Foreign Scan")), + "expected a Foreign Scan in plan: {plan:?}" + ); + assert!( + !plan + .iter() + .any(|line| line.contains("Aggregate") && line.contains("(cost=")), + "expected no local Aggregate plan node: {plan:?}" + ); + assert!( + plan.iter() + .any(|line| line.contains("Wrappers") && line.contains("aggregates =")), + "expected aggregate details on the Foreign Scan: {plan:?}" + ); + } + + fn assert_sparse_cube_cf_aggregate(table: &str) { + let sql = format!( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM {table}"# + ); + assert_aggregate_pushed_down(&sql); + + Spi::connect(|c| { + let row = c.select(&sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + row.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + row.get_by_name::("value_count").unwrap().unwrap(), + 48 + ); + assert!( + (row.get_by_name::("value_sum").unwrap().unwrap() - 13_142.86).abs() < 1e-8 + ); + assert!( + (row.get_by_name::("value_avg").unwrap().unwrap() - 273.809_583_333_333_36) + .abs() + < 1e-10 + ); + assert!( + (row.get_by_name::("value_min").unwrap().unwrap() - 273.15).abs() < 1e-10 + ); + assert!( + (row.get_by_name::("value_max").unwrap().unwrap() - 274.55).abs() < 1e-10 + ); + }); + } + + fn assert_aggregate_falls_back(sql: &str) { + let plan = explain_lines(sql); + assert!( + plan.iter() + .any(|line| line.contains("Aggregate") && line.contains("(cost=")), + "expected a local Aggregate plan node: {plan:?}" + ); + } + + // DDL-only smoke test. The MinIO-backed cases below cover actual scans + // against the fixture seeded by .ci/docker-compose-native.yaml. + #[pg_test] + fn zarr_ddl_smoketest() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_test_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS ( + store_url 's3://zarr-test/sentinel2/2025.zarr', + aws_access_key_id 'test-key', + aws_secret_access_key 'test-secret', + aws_region 'us-east-1' + )"#, + None, + &[], + ) + .unwrap(); + c.update( + r#" + CREATE FOREIGN TABLE zarr_test_cells ( + x double precision, + y double precision, + time timestamptz, + b04 real + ) + SERVER zarr_test_server + OPTIONS ( + array_group 'reflectance', + time_unit 'seconds', + time_origin 'unix' + ) + "#, + None, + &[], + ) + .unwrap(); + }); + } + + // Building a server without a store_url must be rejected by the validator. + #[pg_test(error = "required option `store_url` is not specified")] + fn zarr_validator_requires_store_url() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_bad_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS ( + endpoint_url 'http://localhost:9000' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + // Invalid time_unit values must be rejected at CREATE FOREIGN TABLE time. + #[pg_test( + error = "invalid value for option 'fortnights': must be one of: seconds, milliseconds, microseconds, nanoseconds, minutes, hours, days" + )] + fn zarr_validator_rejects_bad_time_unit() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_test_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS (store_url 's3://zarr-test/x.zarr')"#, + None, + &[], + ) + .unwrap(); + c.update( + r#" + CREATE FOREIGN TABLE zarr_bad_time ( + x double precision, + y double precision, + b04 real + ) + SERVER zarr_test_server + OPTIONS (time_unit 'fortnights') + "#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test(error = "invalid value for option 'decode_cf': must be 'true' or 'false'")] + fn zarr_validator_rejects_bad_cf_decode_boolean() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_cf_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS (store_url 's3://zarr-test/x.zarr')"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE FOREIGN TABLE zarr_bad_cf_option ( + x double precision, + y double precision, + value double precision + ) + SERVER zarr_cf_server + OPTIONS (array_group 'value', decode_cf 'yes')"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test(error = "invalid value for option 'time_from_attrs': must be 'true' or 'false'")] + fn zarr_validator_rejects_bad_time_from_attrs_boolean() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_time_attrs_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS (store_url 's3://zarr-test/x.zarr')"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE FOREIGN TABLE zarr_bad_time_attrs_option ( + x double precision, + y double precision, + time timestamptz, + value real + ) + SERVER zarr_time_attrs_server + OPTIONS (array_group 'value', time_from_attrs 'yes')"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test( + error = "invalid value for option 'time_from_attrs': cannot be combined with 'time_unit' or 'time_origin'" + )] + fn zarr_validator_rejects_time_from_attrs_with_manual_time_options() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_time_attrs_conflict_server + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS (store_url 's3://zarr-test/x.zarr')"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE FOREIGN TABLE zarr_bad_time_attrs_conflict ( + x double precision, + y double precision, + time timestamptz, + value real + ) + SERVER zarr_time_attrs_conflict_server + OPTIONS ( + array_group 'value', + time_from_attrs 'true', + time_unit 'seconds' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test(error = "invalid value for option 'path_style_url': must be 'true' or 'false'")] + fn zarr_validator_rejects_bad_path_style_boolean() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_bad_path_style + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS ( + store_url 's3://zarr-test/x.zarr', + path_style_url 'yes' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test(error = "required option `aws_secret_access_key` is not specified")] + fn zarr_validator_rejects_partial_direct_credentials() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_partial_credentials + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS ( + store_url 's3://zarr-test/x.zarr', + aws_access_key_id 'test-key' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test( + error = "invalid authentication options: anonymous authentication cannot be combined with explicit credentials" + )] + fn zarr_validator_rejects_conflicting_authentication() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_conflicting_auth + FOREIGN DATA WRAPPER zarr_wrapper + OPTIONS ( + store_url 's3://zarr-test/x.zarr', + anonymous 'true', + aws_access_key_id 'test-key', + aws_secret_access_key 'test-secret' + )"#, + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test] + fn zarr_local_v2_scan_sparse_fill_cf_aggregate_and_explain_e2e() { + create_local_e2e_wrapper(); + create_local_e2e_server("zarr_local_v2_server", "e2e.zarr"); + create_local_time_y_x_table( + "zarr_local_v2_raw", + "zarr_local_v2_server", + "nested/raw", + false, + ); + create_local_time_y_x_table( + "zarr_local_v2_cf", + "zarr_local_v2_server", + "nested/raw", + true, + ); + assert_sparse_cube_cf_aggregate("zarr_local_v2_cf"); + + Spi::connect(|c| { + let summary = c + .select( + r#"SELECT count(*) AS total_count, + sum(value)::double precision AS value_sum, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max, + count(*) FILTER (WHERE value = -7.5) AS fill_count + FROM zarr_local_v2_raw"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3_574.0 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + assert_eq!( + summary + .get_by_name::("fill_count") + .unwrap() + .unwrap(), + 8 + ); + + let probes = c + .select( + r#"SELECT value + FROM zarr_local_v2_raw + WHERE (time, y, x) IN ( + ('1970-01-01 00:00:00+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:00+00'::timestamptz, 50, 130), + ('1970-01-01 00:00:03.6+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:03.6+00'::timestamptz, 50, 150) + ) + ORDER BY time, y, x"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("value").unwrap().unwrap()) + .collect::>(); + assert_eq!(probes, vec![11.0, 43.0, 111.0, -7.5]); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT count(*), sum(value) + FROM zarr_local_v2_raw"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Storage Backend: local"), "plan: {plan:?}"); + assert!(has("Zarr Max Concurrent Reads: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 3"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Missing: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 288 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_local_v3_sharding_range_and_rescan_cache_e2e() { + create_local_e2e_wrapper(); + create_local_e2e_server("zarr_local_v3_server", "e2e-v3.zarr"); + create_local_time_y_x_table( + "zarr_local_v3_shard", + "zarr_local_v3_server", + "nested/shard_end", + false, + ); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT x, value + FROM zarr_local_v3_shard + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130 + ORDER BY x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!(values, vec![(110.0, 11.0), (120.0, 12.0), (130.0, 13.0)]); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT ordinal, + (SELECT count(*) + FROM zarr_local_v3_shard + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND upper_x) AS selected + FROM (VALUES (1, 130.0::double precision), + (2, 130.0::double precision)) AS limits(ordinal, upper_x) + ORDER BY ordinal"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Storage Backend: local"), "plan: {plan:?}"); + assert!(has("Zarr Max Concurrent Reads: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Storage Layout: sharding_indexed (index: end)"), + "plan: {plan:?}" + ); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 3"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 116 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Cache Hits: 2"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 3"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Shard Index Encoded Bytes: 68 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Shard Payload Encoded Bytes: 48 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Rescans: 1"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_local_inspect_and_multiscales_e2e() { + create_local_e2e_wrapper(); + create_local_e2e_server("zarr_local_inspect_server", "e2e.zarr"); + create_local_e2e_server("zarr_local_ome_server", "e2e-ome-v3.zarr"); + + Spi::connect(|c| { + let paths = c + .select( + "SELECT path FROM zarr_inspect('zarr_local_inspect_server') ORDER BY path", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("path").unwrap()) + .collect::>(); + assert_eq!( + paths, + vec![ + "/", + "nested", + "nested/band", + "nested/blosc", + "nested/channel", + "nested/forecast_time", + "nested/generic4d", + "nested/lazy1m", + "nested/level", + "nested/raw", + "nested/sample", + "nested/spatial2d", + "nested/spatial_ref", + "nested/time", + "nested/x", + "nested/y", + ] + ); + + let raw = c + .select( + r#"SELECT zarr_format, shape, dimensions, dtype, chunks + FROM zarr_inspect('zarr_local_inspect_server') + WHERE path = 'nested/raw'"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(raw.get_by_name::("zarr_format").unwrap(), Some(2)); + assert_eq!( + raw.get_by_name::("shape").unwrap().unwrap().0, + serde_json::json!([2, 5, 6]) + ); + assert_eq!( + raw.get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time", "y", "x"] + ); + assert_eq!( + raw.get_by_name::("dtype").unwrap(), + Some("("chunks").unwrap().unwrap().0, + serde_json::json!([2, 3, 4]) + ); + + let multiscales = c + .select( + r#"SELECT level_index, array_path, shape, chunks, scale, + translation, supported, warnings + FROM zarr_multiscales('zarr_local_ome_server') + ORDER BY level_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("level_index").unwrap().unwrap(), + row.get_by_name::("array_path").unwrap().unwrap(), + row.get_by_name::("shape").unwrap().unwrap().0, + row.get_by_name::("chunks").unwrap().unwrap().0, + row.get_by_name::, _>("scale").unwrap().unwrap(), + row.get_by_name::, _>("translation") + .unwrap() + .unwrap(), + row.get_by_name::("supported").unwrap().unwrap(), + row.get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + ) + }) + .collect::>(); + assert_eq!( + multiscales, + vec![ + ( + 0, + "image/0".to_string(), + serde_json::json!([4, 4]), + serde_json::json!([3, 3]), + vec![4.0, 12.0], + vec![120.0, 260.0], + true, + Vec::new(), + ), + ( + 1, + "image/1".to_string(), + serde_json::json!([2, 2]), + serde_json::json!([2, 2]), + vec![8.0, 24.0], + vec![122.0, 266.0], + true, + Vec::new(), + ), + ] + ); + + let ome_paths = c + .select( + "SELECT path FROM zarr_inspect('zarr_local_ome_server') ORDER BY path", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("path").unwrap()) + .collect::>(); + assert_eq!(ome_paths, vec!["/", "image", "image/0", "image/1"]); + }); + } + + #[pg_test] + fn zarr_local_privileges_owner_transfer_and_runtime_guard_e2e() { + create_local_e2e_wrapper(); + let store_url = local_fixture_url("e2e.zarr"); + Spi::run( + r#"CREATE ROLE zarr_local_nonsupervisor NOSUPERUSER; + GRANT USAGE ON FOREIGN DATA WRAPPER zarr_local_e2e_wrapper + TO zarr_local_nonsupervisor"#, + ) + .unwrap(); + + Spi::run("SET ROLE zarr_local_nonsupervisor").unwrap(); + let create_error = capture_query_error(&format!( + r#"CREATE SERVER zarr_local_forbidden + FOREIGN DATA WRAPPER zarr_local_e2e_wrapper + OPTIONS (store_url '{store_url}')"# + )); + assert!( + create_error.contains( + "file:// Zarr stores may only be created or altered by a PostgreSQL superuser" + ), + "message: {create_error}" + ); + Spi::run("RESET ROLE").unwrap(); + + create_local_e2e_server("zarr_local_delegated_server", "e2e.zarr"); + create_local_time_y_x_table( + "zarr_local_delegated", + "zarr_local_delegated_server", + "nested/raw", + false, + ); + Spi::run( + r#"GRANT USAGE ON FOREIGN SERVER zarr_local_delegated_server + TO zarr_local_nonsupervisor; + GRANT SELECT ON zarr_local_delegated + TO zarr_local_nonsupervisor; + GRANT SELECT, INSERT, UPDATE ON public.wrappers_fdw_stats + TO zarr_local_nonsupervisor; + SET ROLE zarr_local_nonsupervisor"#, + ) + .unwrap(); + let delegated = Spi::get_one::( + r#"SELECT value + FROM zarr_local_delegated + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ) + .unwrap(); + assert_eq!(delegated, Some(11.0)); + Spi::run("RESET ROLE").unwrap(); + + Spi::run("ALTER SERVER zarr_local_delegated_server OWNER TO zarr_local_nonsupervisor") + .unwrap(); + Spi::run("SET ROLE zarr_local_nonsupervisor").unwrap(); + let alter_error = capture_query_error( + "ALTER SERVER zarr_local_delegated_server OPTIONS (ADD max_concurrent_reads '2')", + ); + assert!( + alter_error.contains( + "file:// Zarr stores may only be created or altered by a PostgreSQL superuser" + ), + "message: {alter_error}" + ); + let runtime_error = capture_query_error( + r#"SELECT value + FROM zarr_local_delegated + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + runtime_error.contains( + "file:// Zarr store foreign server must be owned by a PostgreSQL superuser" + ), + "message: {runtime_error}" + ); + Spi::run("RESET ROLE").unwrap(); + } + + #[pg_test] + fn zarr_local_table_path_traversal_is_rejected_before_io() { + create_local_e2e_wrapper(); + create_local_e2e_server("zarr_local_traversal_server", "e2e.zarr"); + let message = capture_query_error( + r#"CREATE FOREIGN TABLE zarr_local_traversal ( + time timestamp with time zone, + y double precision, + x double precision, + value real + ) + SERVER zarr_local_traversal_server + OPTIONS (array_group '../nested/raw')"#, + ); + assert!( + message.contains("must be a non-empty array path inside the store"), + "message: {message}" + ); + } + + #[pg_test] + fn zarr_local_plain_explain_does_not_open_missing_root_e2e() { + create_local_e2e_wrapper(); + let store_url = format!("{}/missing-root", local_fixture_url("e2e.zarr")); + Spi::run(&format!( + r#"CREATE SERVER zarr_local_missing_root_server + FOREIGN DATA WRAPPER zarr_local_e2e_wrapper + OPTIONS (store_url '{store_url}')"# + )) + .unwrap(); + create_local_time_y_x_table( + "zarr_local_missing_root", + "zarr_local_missing_root_server", + "nested/raw", + false, + ); + + let plan = Spi::connect(|c| { + c.select( + "EXPLAIN (COSTS OFF) SELECT * FROM zarr_local_missing_root", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>() + }); + assert!( + plan.iter().any(|line| line.contains("Foreign Scan")), + "plan: {plan:?}" + ); + } + + #[pg_test] + fn zarr_http_validator_requires_insecure_opt_in() { + create_http_e2e_wrapper(); + for (statement, expected) in [ + ( + r#"CREATE SERVER zarr_http_missing_opt_in + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS (store_url 'http://127.0.0.1:8787/root.zarr')"#, + "invalid value for option 'allow_insecure_http': must be 'true' for http:// stores", + ), + ( + r#"CREATE SERVER zarr_https_bad_opt_in + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS ( + store_url 'https://datasets.example.test/root.zarr', + allow_insecure_http 'true' + )"#, + "invalid value for option 'allow_insecure_http': may be 'true' only for http:// stores", + ), + ( + r#"CREATE SERVER zarr_http_credentials + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS ( + store_url 'http://user:secret@127.0.0.1:8787/root.zarr', + allow_insecure_http 'true' + )"#, + "credentials, query, and fragment are not allowed", + ), + ( + r#"CREATE SERVER zarr_http_s3_auth + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS ( + store_url 'http://127.0.0.1:8787/root.zarr', + allow_insecure_http 'true', + anonymous 'true' + )"#, + "invalid value for option 'anonymous': is only valid for s3:// stores", + ), + ] { + let message = capture_query_error(statement); + assert!(message.contains(expected), "message: {message}"); + } + Spi::run( + r#"CREATE SERVER zarr_https_default + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS (store_url 'https://datasets.example.test/root.zarr')"#, + ) + .unwrap(); + } + + #[pg_test] + fn zarr_http_v2_sparse_cf_aggregate_and_anonymous_transport_e2e() { + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_v2_server", + "v2_transport", + "anonymous_only", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_v2_raw", + "zarr_http_v2_server", + "nested/raw", + false, + ); + create_http_time_y_x_table("zarr_http_v2_cf", "zarr_http_v2_server", "nested/raw", true); + assert_sparse_cube_cf_aggregate("zarr_http_v2_cf"); + + Spi::connect(|c| { + let summary = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value)::double precision AS value_sum, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max, + count(*) FILTER (WHERE value = -7.5) AS fill_count + FROM zarr_http_v2_raw"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary.get_by_name::("total_count").unwrap(), + Some(60) + ); + assert_eq!( + summary.get_by_name::("value_count").unwrap(), + Some(60) + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap(), + Some(3_574.0) + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap(), + Some(-7.5) + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap(), + Some(143.0) + ); + assert_eq!( + summary.get_by_name::("fill_count").unwrap(), + Some(8) + ); + + let raw = c + .select( + r#"SELECT value + FROM zarr_http_v2_raw + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(raw, Some(11.0)); + + let decoded = c + .select( + r#"SELECT value + FROM zarr_http_v2_cf + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap() + .unwrap(); + assert!((decoded - 273.26).abs() < 1e-10); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT count(*), sum(value) + FROM zarr_http_v2_raw"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Storage Backend: http"), "plan: {plan:?}"); + assert!(has("Zarr Max Concurrent Reads: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 3"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Missing: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 288 bytes"), "plan: {plan:?}"); + assert!( + has("Zarr Fill Bytes Synthesized: 96 bytes"), + "plan: {plan:?}" + ); + }); + + let stats = http_case_stats("v2_transport"); + assert_eq!(stats["forbidden_header_gets"], serde_json::json!(0)); + let encodings = stats["accept_encodings"].as_object().unwrap(); + assert_eq!(encodings.len(), 1, "stats: {stats}"); + assert!( + encodings + .get("identity") + .and_then(serde_json::Value::as_u64) + .is_some_and(|count| count > 0), + "stats: {stats}" + ); + } + + #[pg_test] + fn zarr_http_v3_direct_coordinate_aggregate_and_ome_e2e() { + create_http_e2e_wrapper(); + create_http_e2e_server("zarr_http_v3_server", "v3_direct", "plain", "e2e-v3.zarr"); + create_http_time_y_x_table( + "zarr_http_v3_direct", + "zarr_http_v3_server", + "nested/raw_default", + false, + ); + create_http_e2e_server( + "zarr_http_ome_server", + "ome_explicit", + "plain", + "e2e-ome-v3.zarr", + ); + create_http_ome_table("zarr_http_ome_level1", "zarr_http_ome_server", 1); + + Spi::connect(|c| { + let summary = c + .select( + r#"SELECT count(*) AS cells, + sum(value)::double precision AS value_sum, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max, + count(*) FILTER (WHERE value = -7.5) AS fill_count + FROM zarr_http_v3_direct"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(summary.get_by_name::("cells").unwrap(), Some(60)); + assert_eq!( + summary.get_by_name::("value_sum").unwrap(), + Some(3_574.0) + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap(), + Some(-7.5) + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap(), + Some(143.0) + ); + assert_eq!( + summary.get_by_name::("fill_count").unwrap(), + Some(8) + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_http_v3_direct + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Storage Backend: http"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Coordinate-Pruned: 3"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 5"), "plan: {plan:?}"); + assert!( + has("Zarr Coordinate Encoded Bytes: 128 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 96 bytes"), "plan: {plan:?}"); + + let level1 = c + .select( + "SELECT y, x, value FROM zarr_http_ome_level1 ORDER BY y, x", + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("y").unwrap().unwrap(), + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + level1, + vec![ + (122.0, 266.0, 2.5), + (122.0, 290.0, 4.5), + (130.0, 266.0, 10.5), + (130.0, 290.0, 12.5), + ] + ); + }); + } + + #[pg_test] + fn zarr_http_v3_sharding_range_etag_rescan_cache_e2e() { + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_shard_values_server", + "shard_values", + "plain", + "e2e-v3.zarr", + ); + create_http_time_y_x_table( + "zarr_http_shard_values", + "zarr_http_shard_values_server", + "nested/shard_end", + false, + ); + create_http_e2e_server( + "zarr_http_shard_rescan_server", + "shard_rescan", + "plain", + "e2e-v3.zarr", + ); + create_http_time_y_x_table( + "zarr_http_shard_rescan", + "zarr_http_shard_rescan_server", + "nested/shard_end", + false, + ); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT x, value + FROM zarr_http_shard_values + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130 + ORDER BY x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!(values, vec![(110.0, 11.0), (120.0, 12.0), (130.0, 13.0)]); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT ordinal, + (SELECT count(*) + FROM zarr_http_shard_rescan + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND upper_x) AS selected + FROM (VALUES (1, 130.0::double precision), + (2, 130.0::double precision)) AS limits(ordinal, upper_x) + ORDER BY ordinal"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Storage Backend: http"), "plan: {plan:?}"); + assert!( + has("Zarr Storage Layout: sharding_indexed (index: end)"), + "plan: {plan:?}" + ); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 3"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 116 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Cache Hits: 2"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 3"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!(has("Zarr Rescans: 1"), "plan: {plan:?}"); + }); + + let stats = http_case_stats("shard_rescan"); + assert_eq!( + stats["ranges"], + serde_json::json!({ + "bytes=-68": 1, + "bytes=0-23": 1, + "bytes=48-71": 1 + }), + "stats: {stats}" + ); + assert_eq!( + stats["if_matches"], + serde_json::json!({ + "\"711994c2fd69ffddd195ff57991d85b98faba365038f6e451a3e0c6f437d5bd4\"": 2 + }), + "stats: {stats}" + ); + } + + #[pg_test] + fn zarr_http_privileges_delegated_usage_and_owner_guard_e2e() { + create_http_e2e_wrapper(); + Spi::run( + r#"CREATE ROLE zarr_http_reader NOSUPERUSER; + GRANT USAGE ON FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + TO zarr_http_reader; + SET ROLE zarr_http_reader"#, + ) + .unwrap(); + let create_error = capture_query_error( + r#"CREATE SERVER zarr_http_forbidden + FOREIGN DATA WRAPPER zarr_http_e2e_wrapper + OPTIONS ( + store_url 'http://127.0.0.1:8787/stores/forbidden/plain/e2e.zarr', + allow_insecure_http 'true' + )"#, + ); + assert!( + create_error.contains( + "HTTP(S) Zarr stores may only be created or altered by a PostgreSQL superuser" + ), + "message: {create_error}" + ); + Spi::run("RESET ROLE").unwrap(); + + create_http_e2e_server( + "zarr_http_delegated_server", + "delegated", + "anonymous_only", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_delegated", + "zarr_http_delegated_server", + "nested/raw", + false, + ); + Spi::run( + r#"GRANT USAGE ON FOREIGN SERVER zarr_http_delegated_server + TO zarr_http_reader; + GRANT SELECT ON zarr_http_delegated TO zarr_http_reader; + GRANT SELECT, INSERT, UPDATE ON public.wrappers_fdw_stats + TO zarr_http_reader; + SET ROLE zarr_http_reader"#, + ) + .unwrap(); + let value = Spi::get_one::( + r#"SELECT value + FROM zarr_http_delegated + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ) + .unwrap(); + assert_eq!(value, Some(11.0)); + Spi::run("RESET ROLE").unwrap(); + + let stats = http_case_stats("delegated"); + assert_eq!(stats["forbidden_header_gets"], serde_json::json!(0)); + + Spi::run("ALTER SERVER zarr_http_delegated_server OWNER TO zarr_http_reader").unwrap(); + Spi::run("SET ROLE zarr_http_reader").unwrap(); + let alter_error = capture_query_error( + "ALTER SERVER zarr_http_delegated_server OPTIONS (ADD max_concurrent_reads '2')", + ); + assert!( + alter_error.contains( + "HTTP(S) Zarr stores may only be created or altered by a PostgreSQL superuser" + ), + "message: {alter_error}" + ); + let runtime_error = capture_query_error( + r#"SELECT value + FROM zarr_http_delegated + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + runtime_error.contains( + "HTTP(S) Zarr store foreign server must be owned by a PostgreSQL superuser" + ), + "message: {runtime_error}" + ); + Spi::run("RESET ROLE").unwrap(); + } + + #[pg_test] + fn zarr_http_plain_explain_and_listing_fail_without_object_io_e2e() { + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_explain_server", + "explain_zero", + "deny_all", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_explain", + "zarr_http_explain_server", + "nested/raw", + false, + ); + create_http_e2e_server( + "zarr_http_inspect_server", + "inspect_zero", + "deny_all", + "e2e.zarr", + ); + create_http_e2e_server( + "zarr_http_multiscales_server", + "multiscales_zero", + "deny_all", + "e2e-ome-v3.zarr", + ); + + let plan = explain_lines("SELECT * FROM zarr_http_explain"); + assert!( + plan.iter().any(|line| line.contains("Foreign Scan")), + "plan: {plan:?}" + ); + assert_eq!( + http_case_stats("explain_zero")["object_gets"], + serde_json::json!(0) + ); + + let inspect_error = + capture_query_error("SELECT * FROM zarr_inspect('zarr_http_inspect_server')"); + assert!( + inspect_error.contains( + "HTTP(S) Zarr stores do not support hierarchy listing; configure an explicit array path or OME multiscale selection" + ), + "message: {inspect_error}" + ); + assert_eq!( + http_case_stats("inspect_zero")["object_gets"], + serde_json::json!(0) + ); + + let multiscales_error = + capture_query_error("SELECT * FROM zarr_multiscales('zarr_http_multiscales_server')"); + assert!( + multiscales_error.contains( + "HTTP(S) Zarr stores do not support hierarchy listing; configure an explicit array path or OME multiscale selection" + ), + "message: {multiscales_error}" + ); + assert_eq!( + http_case_stats("multiscales_zero")["object_gets"], + serde_json::json!(0) + ); + } + + #[pg_test] + fn zarr_http_redirect_and_oversize_responses_fail_closed_e2e() { + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_redirect_server", + "redirect_rejected", + "redirect_chunk", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_redirect", + "zarr_http_redirect_server", + "nested/raw", + false, + ); + create_http_e2e_server( + "zarr_http_oversize_server", + "oversize_rejected", + "oversize_metadata", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_oversize", + "zarr_http_oversize_server", + "nested/raw", + false, + ); + + let redirect_error = capture_query_error( + r#"SELECT value + FROM zarr_http_redirect + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + redirect_error.contains("nested/raw/0.0.0") + && redirect_error.contains("redirect response was rejected"), + "message: {redirect_error}" + ); + assert_eq!( + http_case_stats("redirect_rejected")["redirect_sink_gets"], + serde_json::json!(0) + ); + + let oversize_error = capture_query_error("SELECT value FROM zarr_http_oversize LIMIT 1"); + assert!( + oversize_error.contains("nested/raw/.zarray") + && oversize_error + .contains("object length 1048577 exceeds the read limit of 1048576 bytes"), + "message: {oversize_error}" + ); + } + + #[pg_test] + fn zarr_http_range_and_generation_failures_are_explicit_e2e() { + create_http_e2e_wrapper(); + for (server, table, case, mode, expected) in [ + ( + "zarr_http_range_200_server", + "zarr_http_range_200", + "range_200_rejected", + "range_200", + "range response returned status 200; server ignored Range", + ), + ( + "zarr_http_bad_range_server", + "zarr_http_bad_range", + "bad_content_range_rejected", + "bad_content_range", + "invalid Content-Range", + ), + ( + "zarr_http_no_etag_server", + "zarr_http_no_etag", + "missing_etag_rejected", + "no_etag", + "range response omitted strong ETag required for shard consistency", + ), + ( + "zarr_http_mutated_server", + "zarr_http_mutated", + "mutation_rejected", + "mutate_shard", + "changed while reading a shard: generation-conditioned object is missing or If-Match failed", + ), + ] { + create_http_e2e_server(server, case, mode, "e2e-v3.zarr"); + create_http_time_y_x_table(table, server, "nested/shard_end", false); + let message = capture_query_error(&format!( + r#"SELECT value + FROM {table} + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"# + )); + assert!( + message.contains("nested/shard_end/c/0/0/0"), + "message: {message}" + ); + assert!(message.contains(expected), "message: {message}"); + } + + let mutation = http_case_stats("mutation_rejected"); + assert_eq!( + mutation["ranges"], + serde_json::json!({"bytes=-68": 1, "bytes=0-23": 1}), + "stats: {mutation}" + ); + assert_eq!( + mutation["if_matches"], + serde_json::json!({"\"generation-a\"": 1}), + "stats: {mutation}" + ); + } + + #[pg_test] + fn zarr_http_stalled_body_honors_statement_timeout_e2e() { + // `SET statement_timeout` inside this pg_test would be too late: the + // outer test statement has already started. Arm PostgreSQL's existing + // statement-timeout handler directly so it fires while reqwest is + // awaiting the stalled body. + unsafe extern "C" { + fn enable_timeout_after(id: std::ffi::c_int, delay_ms: std::ffi::c_int); + fn disable_timeout(id: std::ffi::c_int, keep_indicator: bool); + } + const STATEMENT_TIMEOUT_ID: std::ffi::c_int = 3; + + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_stall_server", + "stalled_body", + "stall_chunk", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_stall", + "zarr_http_stall_server", + "nested/raw", + false, + ); + unsafe { enable_timeout_after(STATEMENT_TIMEOUT_ID, 1_500) }; + let message = capture_query_error( + r#"SELECT value + FROM zarr_http_stall + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + unsafe { disable_timeout(STATEMENT_TIMEOUT_ID, false) }; + assert!( + message.contains("canceling statement due to statement timeout"), + "message: {message}" + ); + assert_eq!(Spi::get_one::("SELECT 1").unwrap(), Some(1)); + } + + #[pg_test] + fn zarr_minio_raw_scan_e2e() { + create_minio_e2e_table("zarr_e2e_raw", "nested/raw", "real"); + + Spi::connect_mut(|c| { + let summary = c + .select( + r#"SELECT count(value) AS row_count, + sum(value)::double precision AS value_sum, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max + FROM zarr_e2e_raw"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary.get_by_name::("row_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3574.0 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + + let fill_count = c + .select( + "SELECT count(*) AS fill_count FROM zarr_e2e_raw WHERE value = -7.5", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("fill_count") + .unwrap() + .unwrap(); + assert_eq!(fill_count, 8); + + // Only `value` is projected and only `x` is restricted: `y` must + // remain internal scan state rather than a required SQL column. + let x_only = c + .select( + r#"SELECT count(value) AS row_count, + sum(value)::double precision AS value_sum + FROM zarr_e2e_raw + WHERE x = 120"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + x_only.get_by_name::("row_count").unwrap().unwrap(), + 10 + ); + assert_eq!( + x_only.get_by_name::("value_sum").unwrap().unwrap(), + 720.0 + ); + + let boundary = c + .select( + r#"SELECT value + FROM zarr_e2e_raw + WHERE time = '1970-01-01 01:00:00+00'::timestamptz + AND y = 50 + AND x = 150"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(boundary, vec![-7.5]); + }); + } + + #[pg_test] + fn zarr_minio_lazy_chunk_cursor_starts_large_selection() { + create_minio_e2e_server(); + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN TABLE zarr_e2e_lazy1m (value real) + SERVER zarr_e2e_server + OPTIONS (array_group 'nested/lazy1m')"#, + None, + &[], + ) + .unwrap(); + + let values = c + .select("SELECT value FROM zarr_e2e_lazy1m LIMIT 1", None, &[]) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(values, vec![42.0]); + }); + } + + #[pg_test] + fn zarr_minio_execution_metrics_are_explained() { + create_minio_e2e_table("zarr_e2e_explain_runtime", "nested/raw", "real"); + Spi::connect(|c| { + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT count(*) AS cells, sum(value) AS total + FROM zarr_e2e_explain_runtime"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Total: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 3"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Missing: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 4"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 288 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Data Decoded Bytes: 384 bytes"), "plan: {plan:?}"); + assert!( + has("Zarr Fill Bytes Synthesized: 96 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Logical Cells Examined: 60"), "plan: {plan:?}"); + assert!(has("Zarr Logical Cells Matched: 60"), "plan: {plan:?}"); + assert!(has("Zarr Tuples Emitted: 1"), "plan: {plan:?}"); + assert!(has("Zarr Max Concurrent Reads: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunk-Stat Pruning: disabled"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_rescan_reuses_compressed_chunk_cache() { + create_minio_e2e_table("zarr_e2e_cache_rescan", "nested/raw", "real"); + Spi::connect(|c| { + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT ordinal, + (SELECT count(*) + FROM zarr_e2e_cache_rescan + WHERE x > threshold) AS selected + FROM (VALUES (1, 120.0::double precision), + (2, NULL::double precision), + (3, 140.0::double precision)) AS limits(ordinal, threshold) + ORDER BY ordinal"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Data GET Calls: 4"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 4"), "plan: {plan:?}"); + assert!(has("Zarr Cache Hits: 6"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_scalar_aggregate_pushdown_e2e() { + create_minio_e2e_table("zarr_e2e_aggregate", "nested/raw", "real"); + + let whole_array_sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_aggregate"#; + assert_aggregate_pushed_down(whole_array_sql); + assert_aggregate_pushed_down("SELECT count(*) FROM zarr_e2e_aggregate WHERE 120 < x"); + for operator in ["<", "<=", "=", ">", ">="] { + assert_aggregate_pushed_down(&format!( + "SELECT count(*) FROM zarr_e2e_aggregate WHERE x {operator} 'NaN'::double precision" + )); + } + assert_aggregate_pushed_down( + "SELECT count(*) FROM zarr_e2e_aggregate WHERE x IN ('NaN'::double precision, 110)", + ); + + Spi::connect_mut(|c| { + let whole = c + .select(whole_array_sql, None, &[]) + .unwrap() + .next() + .unwrap(); + assert_eq!( + whole.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + whole.get_by_name::("value_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + whole.get_by_name::("value_sum").unwrap().unwrap(), + 3574.0 + ); + assert!( + (whole.get_by_name::("value_avg").unwrap().unwrap() + - 59.566_666_666_666_67) + .abs() + < 1e-12 + ); + assert_eq!( + whole.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + whole.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + + let strict = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_aggregate + WHERE x > 120"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + strict + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 30 + ); + assert_eq!( + strict + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 30 + ); + assert_eq!( + strict.get_by_name::("value_sum").unwrap().unwrap(), + 1444.0 + ); + assert!( + (strict.get_by_name::("value_avg").unwrap().unwrap() + - 48.133_333_333_333_33) + .abs() + < 1e-12 + ); + assert_eq!( + strict.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + strict.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + + let reversed_operand = c + .select( + "SELECT count(*) FROM zarr_e2e_aggregate WHERE 120 < x", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap() + .unwrap(); + assert_eq!(reversed_operand, 30); + + let coordinate_aggregates = c + .select( + r#"SELECT count(x) AS x_count, + sum(x) AS x_sum, + avg(x) AS x_avg, + min(x) AS x_min, + max(x) AS x_max + FROM zarr_e2e_aggregate"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + coordinate_aggregates + .get_by_name::("x_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + coordinate_aggregates + .get_by_name::("x_sum") + .unwrap() + .unwrap(), + 7500.0 + ); + assert_eq!( + coordinate_aggregates + .get_by_name::("x_avg") + .unwrap() + .unwrap(), + 125.0 + ); + assert_eq!( + coordinate_aggregates + .get_by_name::("x_min") + .unwrap() + .unwrap(), + 100.0 + ); + assert_eq!( + coordinate_aggregates + .get_by_name::("x_max") + .unwrap() + .unwrap(), + 150.0 + ); + + let sparse = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_aggregate + WHERE y >= 40 AND x >= 140"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + sparse + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 8 + ); + assert_eq!( + sparse + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 8 + ); + assert_eq!( + sparse.get_by_name::("value_sum").unwrap().unwrap(), + -60.0 + ); + assert_eq!( + sparse.get_by_name::("value_avg").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + sparse.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + sparse.get_by_name::("value_max").unwrap().unwrap(), + -7.5 + ); + + let membership = c + .select( + r#"SELECT count(*) AS total_count, + sum(value) AS value_sum, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_aggregate + WHERE x IN (110, 150)"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + membership + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 20 + ); + assert_eq!( + membership + .get_by_name::("value_sum") + .unwrap() + .unwrap(), + 1070.0 + ); + assert_eq!( + membership + .get_by_name::("value_min") + .unwrap() + .unwrap(), + -7.5 + ); + assert_eq!( + membership + .get_by_name::("value_max") + .unwrap() + .unwrap(), + 141.0 + ); + + let value_qual = c + .select( + r#"SELECT count(*) AS total_count, sum(value) AS value_sum + FROM zarr_e2e_aggregate + WHERE value = -7.5"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + value_qual + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 8 + ); + assert_eq!( + value_qual + .get_by_name::("value_sum") + .unwrap() + .unwrap(), + -60.0 + ); + + let empty = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_aggregate + WHERE x > 999"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + empty.get_by_name::("total_count").unwrap().unwrap(), + 0 + ); + assert_eq!( + empty.get_by_name::("value_count").unwrap().unwrap(), + 0 + ); + assert_eq!(empty.get_by_name::("value_sum").unwrap(), None); + assert_eq!(empty.get_by_name::("value_avg").unwrap(), None); + assert_eq!(empty.get_by_name::("value_min").unwrap(), None); + assert_eq!(empty.get_by_name::("value_max").unwrap(), None); + + let empty_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT count(*), sum(value) + FROM zarr_e2e_aggregate + WHERE x > 999"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let empty_has = |text: &str| empty_plan.iter().any(|line| line.contains(text)); + assert!(empty_has("Zarr Chunks Selected: 0"), "plan: {empty_plan:?}"); + assert!( + empty_has("Zarr Chunks Requested: 0"), + "plan: {empty_plan:?}" + ); + assert!(empty_has("Zarr Data GET Calls: 0"), "plan: {empty_plan:?}"); + assert!( + empty_has("Zarr Logical Cells Examined: 0"), + "plan: {empty_plan:?}" + ); + assert!(empty_has("Zarr Tuples Emitted: 1"), "plan: {empty_plan:?}"); + + for (operator, expected) in [("<", 60), ("<=", 60), ("=", 0), (">", 0), (">=", 0)] { + let count = c + .select( + &format!( + "SELECT count(*) FROM zarr_e2e_aggregate WHERE x {operator} 'NaN'::double precision" + ), + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap() + .unwrap(); + assert_eq!(count, expected, "unexpected result for x {operator} NaN"); + } + assert_eq!( + c.select( + "SELECT count(*) FROM zarr_e2e_aggregate WHERE x IN ('NaN'::double precision, 110)", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(10) + ); + + let rescans = c + .select( + r#"SELECT ordinal, threshold, + (SELECT count(*) + FROM zarr_e2e_aggregate + WHERE x > threshold) AS selected + FROM (VALUES (1, 120.0::double precision), + (2, NULL::double precision), + (3, 140.0::double precision)) AS limits(ordinal, threshold) + ORDER BY ordinal"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("threshold").unwrap(), + row.get_by_name::("selected").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + rescans, + vec![(Some(120.0), 30), (None, 0), (Some(140.0), 10)] + ); + + c.update("SET LOCAL plan_cache_mode = force_generic_plan", None, &[]) + .unwrap(); + c.update( + "PREPARE zarr_aggregate_threshold(double precision) AS SELECT count(*) FROM zarr_e2e_aggregate WHERE x > $1", + None, + &[], + ) + .unwrap(); + for (argument, expected) in [("120", 30), ("NULL", 0), ("140", 10)] { + let count = c + .select( + &format!("EXECUTE zarr_aggregate_threshold({argument})"), + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap() + .unwrap(); + assert_eq!(count, expected, "unexpected prepared result for {argument}"); + } + c.update("DEALLOCATE zarr_aggregate_threshold", None, &[]) + .unwrap(); + }); + } + + #[pg_test] + fn zarr_scalar_aggregate_unsupported_shapes_fall_back() { + create_minio_e2e_table("zarr_e2e_aggregate_fallback", "nested/raw", "real"); + Spi::connect_mut(|c| { + c.update("CREATE SCHEMA zarr_aggregate_custom", None, &[]) + .unwrap(); + c.update( + r#"CREATE FUNCTION zarr_aggregate_custom.add_hundred(real, real) + RETURNS real + LANGUAGE sql IMMUTABLE STRICT + AS 'SELECT $1 + 100::real'"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE AGGREGATE zarr_aggregate_custom.sum(real) ( + SFUNC = zarr_aggregate_custom.add_hundred, + STYPE = real, + INITCOND = '0' + )"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE FUNCTION zarr_aggregate_custom.always_true(double precision, double precision) + RETURNS boolean + LANGUAGE plpgsql IMMUTABLE + AS $$ BEGIN RETURN true; END $$"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE OPERATOR zarr_aggregate_custom.= ( + LEFTARG = double precision, + RIGHTARG = double precision, + FUNCTION = zarr_aggregate_custom.always_true + )"#, + None, + &[], + ) + .unwrap(); + }); + + let cases = [ + "SELECT count(DISTINCT value) FROM zarr_e2e_aggregate_fallback", + "SELECT sum(value + 1) FROM zarr_e2e_aggregate_fallback", + "SELECT count(*) FILTER (WHERE value = -7.5) FROM zarr_e2e_aggregate_fallback", + "SELECT sum(value ORDER BY x) FROM zarr_e2e_aggregate_fallback", + "SELECT count(*) FROM zarr_e2e_aggregate_fallback WHERE value + 1 > 0", + "SELECT time, count(*) FROM zarr_e2e_aggregate_fallback GROUP BY time", + "SELECT count(*), 1 FROM zarr_e2e_aggregate_fallback", + "SELECT count(*) FROM zarr_e2e_aggregate_fallback HAVING count(*) > 0", + "SELECT zarr_aggregate_custom.sum(value) FROM zarr_e2e_aggregate_fallback", + "SELECT count(*) FROM zarr_e2e_aggregate_fallback WHERE x OPERATOR(zarr_aggregate_custom.=) 999::double precision", + "SELECT count(*) FROM zarr_e2e_aggregate_fallback GROUP BY GROUPING SETS ((), ())", + ]; + for sql in cases { + assert_aggregate_falls_back(sql); + } + + Spi::connect(|c| { + assert_eq!( + c.select( + "SELECT count(DISTINCT value) FROM zarr_e2e_aggregate_fallback", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(53) + ); + assert_eq!( + c.select( + "SELECT sum(value + 1) FROM zarr_e2e_aggregate_fallback", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(3634.0) + ); + assert_eq!( + c.select( + "SELECT count(*) FILTER (WHERE value = -7.5) FROM zarr_e2e_aggregate_fallback", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(8) + ); + assert_eq!( + c.select( + "SELECT count(*) FROM zarr_e2e_aggregate_fallback WHERE value + 1 > 0", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(52) + ); + assert_eq!( + c.select( + "SELECT zarr_aggregate_custom.sum(value) FROM zarr_e2e_aggregate_fallback", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(6000.0) + ); + assert_eq!( + c.select( + "SELECT count(*) FROM zarr_e2e_aggregate_fallback WHERE x OPERATOR(zarr_aggregate_custom.=) 999::double precision", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(), + Some(60) + ); + let grouping_sets = c + .select( + "SELECT count(*) FROM zarr_e2e_aggregate_fallback GROUP BY GROUPING SETS ((), ())", + None, + &[], + ) + .unwrap() + .map(|row| row.get::(1).unwrap().unwrap()) + .collect::>(); + assert_eq!(grouping_sets, vec![60, 60]); + }); + } + + #[pg_test] + fn zarr_minio_blosc_scan_e2e() { + create_minio_e2e_table("zarr_e2e_blosc", "nested/blosc", "real"); + + assert_aggregate_pushed_down( + r#"SELECT count(*) AS total_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_blosc + WHERE time = '1970-01-01 01:00:00+00'::timestamptz + AND y >= 40 + AND x BETWEEN 130 AND 150"#, + ); + + Spi::connect_mut(|c| { + let summary = c + .select( + r#"SELECT count(value) AS row_count, + sum(value)::double precision AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_blosc + WHERE time = '1970-01-01 01:00:00+00'::timestamptz + AND y >= 40 + AND x BETWEEN 130 AND 150"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary.get_by_name::("row_count").unwrap().unwrap(), + 6 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 834.0 + ); + assert_eq!( + summary.get_by_name::("value_avg").unwrap().unwrap(), + 139.0 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + 133.0 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 145.0 + ); + + let boundary = c + .select( + r#"SELECT value + FROM zarr_e2e_blosc + WHERE time = '1970-01-01 01:00:00+00'::timestamptz + AND y = 50 + AND x = 150"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(boundary, vec![145.0]); + }); + } + + #[pg_test( + error = "column 'value' has incompatible PostgreSQL type OID 701; expected real (OID 700)" + )] + fn zarr_minio_rejects_wrong_value_type() { + create_minio_e2e_table("zarr_e2e_bad_value_type", "nested/raw", "double precision"); + + Spi::connect_mut(|c| { + c.select("SELECT value FROM zarr_e2e_bad_value_type", None, &[]) + .unwrap(); + }); + } + + #[pg_test( + error = "column 'value' has incompatible PostgreSQL type OID 700; expected double precision (OID 701)" + )] + fn zarr_minio_cf_decode_rejects_non_float8_value_type() { + create_minio_e2e_table_with_cf("zarr_e2e_bad_cf_type", "nested/raw", "real", true); + + Spi::connect_mut(|c| { + c.select("SELECT value FROM zarr_e2e_bad_cf_type", None, &[]) + .unwrap(); + }); + } + + #[pg_test] + fn zarr_minio_cf_value_decode_e2e() { + create_minio_e2e_table_with_cf( + "zarr_e2e_cf_decoded", + "nested/raw", + "double precision", + true, + ); + + assert_aggregate_pushed_down( + r#"SELECT count(*) AS total_count, + count(value) AS valid_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_cf_decoded"#, + ); + + Spi::connect(|c| { + let summary = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS valid_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_cf_decoded"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary + .get_by_name::("valid_count") + .unwrap() + .unwrap(), + 48 + ); + let value_sum = summary.get_by_name::("value_sum").unwrap().unwrap(); + let value_avg = summary.get_by_name::("value_avg").unwrap().unwrap(); + let value_min = summary.get_by_name::("value_min").unwrap().unwrap(); + let value_max = summary.get_by_name::("value_max").unwrap().unwrap(); + assert!((value_sum - 13_142.86).abs() < 1e-8); + assert!((value_avg - 273.809_583_333_333_36).abs() < 1e-10); + assert!((value_min - 273.15).abs() < 1e-10); + assert!((value_max - 274.55).abs() < 1e-10); + + let sparse = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS valid_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_e2e_cf_decoded + WHERE y >= 40 AND x >= 140"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + sparse + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 8 + ); + assert_eq!( + sparse + .get_by_name::("valid_count") + .unwrap() + .unwrap(), + 0 + ); + assert_eq!(sparse.get_by_name::("value_sum").unwrap(), None); + assert_eq!(sparse.get_by_name::("value_avg").unwrap(), None); + assert_eq!(sparse.get_by_name::("value_min").unwrap(), None); + assert_eq!(sparse.get_by_name::("value_max").unwrap(), None); + + let decoded_nulls = c + .select( + "SELECT count(*) FROM zarr_e2e_cf_decoded WHERE value IS NULL", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap() + .unwrap(); + assert_eq!(decoded_nulls, 12); + + for predicate in [ + "time = '1970-01-01 00:00:00+00'::timestamptz AND y = 50 AND x = 120", + "time = '1970-01-01 01:00:00+00'::timestamptz AND y = 50 AND x = 110", + "time = '1970-01-01 01:00:00+00'::timestamptz AND y = 50 AND x = 150", + ] { + let sql = format!("SELECT value FROM zarr_e2e_cf_decoded WHERE {predicate}"); + let value = c + .select(&sql, None, &[]) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(value, None); + } + + let valid_boundary = c + .select( + r#"SELECT value + FROM zarr_e2e_cf_decoded + WHERE time = '1970-01-01 01:00:00+00'::timestamptz + AND y = 50 + AND x = 100"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap() + .unwrap(); + assert!((valid_boundary - 274.55).abs() < 1e-10); + }); + } + + #[pg_test] + fn zarr_minio_time_from_attrs_e2e() { + create_minio_e2e_table_with_options( + "zarr_e2e_time_from_attrs", + "nested/raw", + "real", + ",\n time_from_attrs 'true'", + ); + + Spi::connect(|c| { + let times = c + .select( + r#"SELECT string_agg(ts, ',' ORDER BY ts) AS times + FROM ( + SELECT DISTINCT to_char( + time AT TIME ZONE 'UTC', + 'YYYY-MM-DD HH24:MI:SS.MS' + ) AS ts + FROM zarr_e2e_time_from_attrs + ) t"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("times") + .unwrap() + .unwrap(); + assert_eq!(times, "1970-01-01 00:00:00.000,1970-01-01 00:00:03.600"); + + let selected = c + .select( + r#"SELECT value + FROM zarr_e2e_time_from_attrs + WHERE time = '1970-01-01 00:00:03.6+00'::timestamptz + AND y = 50 + AND x = 100"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(selected, vec![140.0]); + }); + } + + #[pg_test] + fn zarr_minio_generic_dimensions_scan_e2e() { + create_minio_generic4d_table("zarr_e2e_generic4d", "double precision"); + + let generic_aggregate_sql = r#"SELECT count(*) AS total_count, + sum(measurement) AS value_sum, + avg(measurement) AS value_avg, + min(measurement) AS value_min, + max(measurement) AS value_max + FROM zarr_e2e_generic4d + WHERE forecast_time = '1970-01-01 00:00:03.6+00'::timestamptz + AND level >= 40 + AND band > 120 + AND channel = 7"#; + assert_aggregate_pushed_down(generic_aggregate_sql); + + Spi::connect(|c| { + // No dimension is projected or restricted. The executor must still + // return the complete logical value array without requiring any + // coordinate chunk values. + let summary = c + .select( + r#"SELECT count(measurement) AS row_count, + sum(measurement)::double precision AS value_sum, + min(measurement)::double precision AS value_min, + max(measurement)::double precision AS value_max + FROM zarr_e2e_generic4d"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary.get_by_name::("row_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3574.0 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + + let fill_count = c + .select( + "SELECT count(*) AS fill_count FROM zarr_e2e_generic4d WHERE measurement = -7.5", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("fill_count") + .unwrap() + .unwrap(); + assert_eq!(fill_count, 8); + + let boundary = c + .select( + r#"SELECT to_char( + forecast_time AT TIME ZONE 'UTC', + 'YYYY-MM-DD HH24:MI:SS.MS' + ) AS forecast_time, + level, + band, + channel, + measurement + FROM zarr_e2e_generic4d + WHERE forecast_time = '1970-01-01 00:00:03.6+00'::timestamptz + AND level = 50 + AND band = 130 + AND channel = 7"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + boundary + .get_by_name::("forecast_time") + .unwrap() + .unwrap(), + "1970-01-01 00:00:03.600" + ); + assert_eq!( + boundary.get_by_name::("level").unwrap().unwrap(), + 50.0 + ); + assert_eq!( + boundary.get_by_name::("band").unwrap().unwrap(), + 130.0 + ); + assert_eq!( + boundary.get_by_name::("channel").unwrap().unwrap(), + 7.0 + ); + assert_eq!( + boundary + .get_by_name::("measurement") + .unwrap() + .unwrap(), + 143.0 + ); + + let missing_boundary = c + .select( + r#"SELECT measurement + FROM zarr_e2e_generic4d + WHERE forecast_time = '1970-01-01 00:00:03.6+00'::timestamptz + AND level = 50 + AND band = 150 + AND channel = 7"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("measurement") + .unwrap() + .unwrap(); + assert_eq!(missing_boundary, -7.5); + + let aggregate = c + .select(generic_aggregate_sql, None, &[]) + .unwrap() + .next() + .unwrap(); + assert_eq!( + aggregate + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 6 + ); + assert_eq!( + aggregate + .get_by_name::("value_sum") + .unwrap() + .unwrap(), + 246.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_avg") + .unwrap() + .unwrap(), + 41.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_min") + .unwrap() + .unwrap(), + -7.5 + ); + assert_eq!( + aggregate + .get_by_name::("value_max") + .unwrap() + .unwrap(), + 143.0 + ); + }); + } + + #[pg_test] + fn zarr_minio_index_selectors_no_coordinate_reads() { + create_minio_generic4d_table_with_options( + "zarr_e2e_index_selectors", + "double precision", + r#", + dimension_selectors '{"forecast_time":{"index":1},"level":{"index":4},"band":{"index":3},"channel":{"index":0}}'"#, + ); + + let aggregate_sql = r#"SELECT count(*) AS total_count, + sum(measurement) AS value_sum, + avg(measurement) AS value_avg, + min(measurement) AS value_min, + max(measurement) AS value_max + FROM zarr_e2e_index_selectors"#; + assert_aggregate_pushed_down(aggregate_sql); + + Spi::connect(|c| { + let values = c + .select( + "SELECT measurement FROM zarr_e2e_index_selectors", + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("measurement").unwrap().unwrap()) + .collect::>(); + assert_eq!(values, vec![143.0]); + + let aggregate = c.select(aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + aggregate.get_by_name::("total_count").unwrap(), + Some(1) + ); + assert_eq!( + aggregate.get_by_name::("value_sum").unwrap(), + Some(143.0) + ); + assert_eq!( + aggregate.get_by_name::("value_avg").unwrap(), + Some(143.0) + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT measurement FROM zarr_e2e_index_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Total: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 0"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 96 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Logical Cells Examined: 1"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_exact_value_dimension_selectors_and_with_sql_quals() { + create_minio_e2e_server(); + create_minio_generic4d_table_on_server( + "zarr_e2e_value_selectors", + "double precision", + r#", + dimension_selectors '{"forecast_time":{"value":3600},"level":{"value":50},"band":{"value":130},"channel":{"value":7}}'"#, + ); + create_minio_generic4d_table_on_server( + "zarr_e2e_missing_selector", + "double precision", + r#", + dimension_selectors '{"band":{"value":999}}'"#, + ); + create_minio_generic4d_table_on_server( + "zarr_e2e_conflicting_selector", + "double precision", + r#", + dimension_selectors '{"band":{"index":3}}'"#, + ); + + Spi::connect(|c| { + let value = c + .select( + "SELECT measurement FROM zarr_e2e_value_selectors", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("measurement") + .unwrap(); + assert_eq!(value, Some(143.0)); + + let value_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT measurement FROM zarr_e2e_value_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let value_has = |text: &str| value_plan.iter().any(|line| line.contains(text)); + assert!(value_has("Zarr Chunks Selected: 1"), "plan: {value_plan:?}"); + assert!( + value_has("Zarr Coordinate GET Calls: 6"), + "plan: {value_plan:?}" + ); + assert!( + value_has("Zarr Coordinate Encoded Bytes: 128 bytes"), + "plan: {value_plan:?}" + ); + assert!(value_has("Zarr Data GET Calls: 1"), "plan: {value_plan:?}"); + + let missing = c + .select( + "SELECT count(*) AS count FROM zarr_e2e_missing_selector", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("count") + .unwrap(); + assert_eq!(missing, Some(0)); + let missing_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT count(*) FROM zarr_e2e_missing_selector"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let missing_has = |text: &str| missing_plan.iter().any(|line| line.contains(text)); + assert!( + missing_has("Zarr Chunks Selected: 0"), + "plan: {missing_plan:?}" + ); + assert!( + missing_has("Zarr Chunks Requested: 0"), + "plan: {missing_plan:?}" + ); + assert!( + missing_has("Zarr Coordinate GET Calls: 2"), + "plan: {missing_plan:?}" + ); + assert!( + missing_has("Zarr Coordinate Encoded Bytes: 64 bytes"), + "plan: {missing_plan:?}" + ); + assert!( + missing_has("Zarr Data GET Calls: 0"), + "plan: {missing_plan:?}" + ); + + let conflict = c + .select( + "SELECT count(*) AS count FROM zarr_e2e_conflicting_selector WHERE band = 120", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("count") + .unwrap(); + assert_eq!(conflict, Some(0)); + }); + } + + #[pg_test] + fn zarr_minio_value_list_range_selectors_e2e() { + create_minio_e2e_server(); + create_minio_generic4d_table_on_server( + "zarr_e2e_value_list_range_selectors", + "double precision", + r#", + dimension_selectors '{"forecast_time":{"index":1},"level":{"value_range":{"min":20,"max":40}},"band":{"values":[130,110]},"channel":{"index":0}}'"#, + ); + + let aggregate_sql = r#"SELECT count(*) AS total_count, + sum(measurement) AS value_sum, + avg(measurement) AS value_avg, + min(measurement) AS value_min, + max(measurement) AS value_max + FROM zarr_e2e_value_list_range_selectors"#; + assert_aggregate_pushed_down(aggregate_sql); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT measurement + FROM zarr_e2e_value_list_range_selectors"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("measurement").unwrap().unwrap()) + .collect::>(); + assert_eq!(values, vec![111.0, 113.0, 121.0, 123.0, 131.0, 133.0]); + + let sql_intersection = c + .select( + r#"SELECT measurement + FROM zarr_e2e_value_list_range_selectors + WHERE band = 130"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("measurement").unwrap().unwrap()) + .collect::>(); + assert_eq!(sql_intersection, vec![113.0, 123.0, 133.0]); + + let aggregate = c.select(aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + aggregate.get_by_name::("total_count").unwrap(), + Some(6) + ); + assert_eq!( + aggregate.get_by_name::("value_sum").unwrap(), + Some(732.0) + ); + assert_eq!( + aggregate.get_by_name::("value_avg").unwrap(), + Some(122.0) + ); + assert_eq!( + aggregate.get_by_name::("value_min").unwrap(), + Some(111.0) + ); + assert_eq!( + aggregate.get_by_name::("value_max").unwrap(), + Some(133.0) + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT measurement FROM zarr_e2e_value_list_range_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Selected: 2"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 4"), "plan: {plan:?}"); + assert!( + has("Zarr Coordinate Encoded Bytes: 112 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Data GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 192 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Logical Cells Examined: 9"), "plan: {plan:?}"); + assert!(has("Zarr Tuples Emitted: 6"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_index_list_range_sparse_fill_e2e() { + create_minio_e2e_server(); + create_minio_generic4d_table_on_server( + "zarr_e2e_index_list_range_selectors", + "double precision", + r#", + dimension_selectors '{"forecast_time":{"index":1},"level":{"index_range":{"start":1,"stop":4}},"band":{"indices":[3,1]},"channel":{"index":0}}'"#, + ); + create_minio_generic4d_table_on_server( + "zarr_e2e_sparse_index_selectors", + "double precision", + r#", + dimension_selectors '{"forecast_time":{"index":1},"level":{"index":4},"band":{"indices":[1,3,5]},"channel":{"index":0}}'"#, + ); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT measurement + FROM zarr_e2e_index_list_range_selectors"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("measurement").unwrap().unwrap()) + .collect::>(); + assert_eq!(values, vec![111.0, 113.0, 121.0, 123.0, 131.0, 133.0]); + + let index_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT measurement FROM zarr_e2e_index_list_range_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let index_has = |text: &str| index_plan.iter().any(|line| line.contains(text)); + assert!( + index_has("Zarr Coordinate GET Calls: 0"), + "plan: {index_plan:?}" + ); + + let sparse_values = c + .select( + r#"SELECT measurement + FROM zarr_e2e_sparse_index_selectors"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("measurement").unwrap().unwrap()) + .collect::>(); + assert_eq!(sparse_values, vec![141.0, 143.0, -7.5]); + + let sparse_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT measurement FROM zarr_e2e_sparse_index_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let sparse_has = |text: &str| sparse_plan.iter().any(|line| line.contains(text)); + assert!( + sparse_has("Zarr Chunks Selected: 2"), + "plan: {sparse_plan:?}" + ); + assert!( + sparse_has("Zarr Chunks Requested: 2"), + "plan: {sparse_plan:?}" + ); + assert!( + sparse_has("Zarr Chunks Present: 1"), + "plan: {sparse_plan:?}" + ); + assert!( + sparse_has("Zarr Chunks Missing: 1"), + "plan: {sparse_plan:?}" + ); + assert!( + sparse_has("Zarr Coordinate GET Calls: 0"), + "plan: {sparse_plan:?}" + ); + assert!( + sparse_has("Zarr Data Encoded Bytes: 96 bytes"), + "plan: {sparse_plan:?}" + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_sharded_list_range_selectors_e2e() { + create_minio_v3_e2e_server(); + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN TABLE zarr_v3_shard_selectors (value real) + SERVER zarr_v3_e2e_server + OPTIONS ( + array_group 'nested/shard_end', + dimension_selectors '{"time":{"indices":[0]},"y":{"index_range":{"start":1,"stop":2}},"x":{"indices":[1]}}' + )"#, + None, + &[], + ) + .unwrap(); + }); + + Spi::connect(|c| { + let value = c + .select("SELECT value FROM zarr_v3_shard_selectors", None, &[]) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(value, Some(11.0)); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value FROM zarr_v3_shard_selectors"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Total: 12"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 0"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 92 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 1"), "plan: {plan:?}"); + }); + } + + #[pg_test( + error = "invalid value for option 'dimension_selectors': spatial operations require a selector-aware function overload" + )] + fn zarr_spatial_table_selectors_fail_closed() { + install_postgis_in_test_schema(); + create_minio_e2e_server(); + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN TABLE zarr_e2e_spatial_selector ( + y double precision, + x double precision, + value real + ) + SERVER zarr_e2e_server + OPTIONS ( + array_group 'nested/spatial2d', + dimension_selectors '{"x":{"index":1}}' + )"#, + None, + &[], + ) + .unwrap(); + }); + Spi::run( + r#"SELECT * + FROM zarr_sample( + 'zarr_e2e_spatial_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(999, 999), 3857) + ), + 'exact' + )"#, + ) + .unwrap(); + } + + #[pg_test] + fn zarr_spatial_selector_overload_catalog_contract() { + let overloads_owned_by_extension = Spi::get_one::( + r#"WITH targets(signature, expected_defaults) AS ( + VALUES + ('zarr_sample(text,bytea,text)'::pg_catalog.regprocedure, 1), + ('zarr_sample(text,bytea,text,text)'::pg_catalog.regprocedure, 0), + ('zarr_zonal_stats(text,bytea)'::pg_catalog.regprocedure, 0), + ('zarr_zonal_stats(text,bytea,text)'::pg_catalog.regprocedure, 0), + ('zarr_zonal_stats_by_time(text,bytea,timestamp with time zone,timestamp with time zone)'::pg_catalog.regprocedure, 0), + ('zarr_zonal_stats_by_time(text,bytea,timestamp with time zone,timestamp with time zone,text)'::pg_catalog.regprocedure, 0) + ) + SELECT pg_catalog.count(*) = 6 + AND pg_catalog.bool_and(proc.pronargdefaults::integer = targets.expected_defaults) + AND pg_catalog.bool_and( + extension.extname IS NOT DISTINCT FROM 'wrappers' + ) + FROM targets + JOIN pg_catalog.pg_proc AS proc ON proc.oid = targets.signature + LEFT JOIN pg_catalog.pg_depend AS dependency + ON dependency.classid = 'pg_catalog.pg_proc'::pg_catalog.regclass + AND dependency.objid = proc.oid + AND dependency.refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass + AND dependency.deptype = 'e' + LEFT JOIN pg_catalog.pg_extension AS extension + ON extension.oid = dependency.refobjid"#, + ) + .unwrap(); + assert_eq!(overloads_owned_by_extension, Some(true)); + + let no_helper_names = Spi::get_one::( + r#"SELECT NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_proc + WHERE proname IN ( + 'zarr_sample_with_selectors', + 'zarr_zonal_stats_with_selectors', + 'zarr_zonal_stats_by_time_with_selectors' + ) + ) + AND pg_catalog.to_regprocedure('zarr_cells(text,bytea,text)') IS NULL + AND pg_catalog.to_regprocedure( + 'zarr_cells_by_time(text,bytea,timestamp with time zone,timestamp with time zone,text)' + ) IS NULL"#, + ) + .unwrap(); + assert_eq!(no_helper_names, Some(true)); + } + + #[pg_test] + fn zarr_postgis_selector_overloads_on_rank3_time_auxiliary() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_selector_ops", false); + + Spi::connect(|c| { + let sample = c + .select( + r#"SELECT sample.* + FROM zarr_sample( + 'zarr_e2e_spatial_time_selector_ops', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"time":{"index":1}}' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(sample.get_by_name::("x").unwrap(), Some(110.0)); + assert_eq!(sample.get_by_name::("y").unwrap(), Some(20.0)); + assert_eq!(sample.get_by_name::("value").unwrap(), Some(111.0)); + assert_eq!(sample.get_by_name::("x_index").unwrap(), Some(1)); + assert_eq!(sample.get_by_name::("y_index").unwrap(), Some(1)); + assert_eq!( + sample.get_by_name::("coordinate_distance").unwrap(), + Some(0.0) + ); + + let stats = c + .select( + r#"SELECT stats.* + FROM zarr_zonal_stats( + 'zarr_e2e_spatial_time_selector_ops', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + '{"time":{"index":1}}' + ) AS stats"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(stats.get_by_name::("count").unwrap(), Some(9)); + assert_eq!(stats.get_by_name::("valid_count").unwrap(), Some(9)); + assert_eq!(stats.get_by_name::("min").unwrap(), Some(111.0)); + assert_eq!(stats.get_by_name::("max").unwrap(), Some(133.0)); + assert_eq!(stats.get_by_name::("sum").unwrap(), Some(1098.0)); + assert_eq!(stats.get_by_name::("avg").unwrap(), Some(122.0)); + + let by_time = c + .select( + r#"SELECT extract(epoch FROM stats.time)::double precision AS epoch, + stats.time_index, + stats.count, + stats.valid_count, + stats.min, + stats.max, + stats.sum, + stats.avg + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_selector_ops', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00', + '{}' + ) AS stats + ORDER BY stats.time_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("epoch").unwrap().unwrap(), + row.get_by_name::("time_index").unwrap().unwrap(), + row.get_by_name::("count").unwrap().unwrap(), + row.get_by_name::("valid_count").unwrap().unwrap(), + row.get_by_name::("min").unwrap(), + row.get_by_name::("max").unwrap(), + row.get_by_name::("sum").unwrap(), + row.get_by_name::("avg").unwrap(), + ) + }) + .collect::>(); + assert_eq!( + by_time, + vec![ + ( + 0.0, + 0, + 9, + 9, + Some(11.0), + Some(33.0), + Some(198.0), + Some(22.0) + ), + ( + 3.6, + 1, + 9, + 9, + Some(111.0), + Some(133.0), + Some(1098.0), + Some(122.0) + ) + ] + ); + }); + } + + #[pg_test] + fn zarr_postgis_list_range_selector_cardinality() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table_with_options( + "zarr_e2e_spatial_time_table_list_selector", + r#", + time_from_attrs 'true', + dimension_selectors '{"time":{"indices":[0,1]}}'"#, + ); + Spi::run( + r#"CREATE FOREIGN TABLE zarr_e2e_spatial_time_call_list_selector ( + time timestamp with time zone, + y double precision, + x double precision, + value real + ) + SERVER zarr_e2e_server + OPTIONS ( + array_group 'nested/raw', + time_from_attrs 'true' + )"#, + ) + .unwrap(); + + Spi::connect(|c| { + let intersecting = c + .select( + r#"SELECT sample.value + FROM zarr_sample( + 'zarr_e2e_spatial_time_table_list_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"time":{"index":1}}' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(intersecting, Some(111.0)); + + let empty = c + .select( + r#"SELECT pg_catalog.count(*)::bigint AS count + FROM zarr_sample( + 'zarr_e2e_spatial_time_call_list_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"time":{"values":[999]}}' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("count") + .unwrap(); + assert_eq!(empty, Some(0)); + }); + + for (sql, expected) in [ + ( + r#"SELECT * + FROM zarr_zonal_stats( + 'zarr_e2e_spatial_time_call_list_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + '{"time":{"indices":[0,1]}}' + )"#, + "auxiliary dimension 'time' resolves to more than one index; spatial operations require zero or one", + ), + ( + r#"SELECT * + FROM zarr_sample( + 'zarr_e2e_spatial_time_call_list_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"x":{"indices":[1]}}' + )"#, + "zarr_sample owns dimension 'x'; selectors may target auxiliary dimensions only", + ), + ] { + let message = capture_query_error(sql); + assert!(message.contains(expected), "message: {message}"); + } + } + + #[pg_test] + fn zarr_postgis_selector_overloads_intersections_and_rejections() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table_with_options( + "zarr_e2e_spatial_time_table_selector", + r#", + time_from_attrs 'true', + dimension_selectors '{"time":{"index":1}}'"#, + ); + Spi::run( + r#"CREATE FOREIGN TABLE zarr_e2e_spatial_time_unselected ( + time timestamp with time zone, + y double precision, + x double precision, + value real + ) + SERVER zarr_e2e_server + OPTIONS ( + array_group 'nested/raw', + time_from_attrs 'true' + )"#, + ) + .unwrap(); + + Spi::connect(|c| { + let intersecting = c + .select( + r#"SELECT sample.value + FROM zarr_sample( + 'zarr_e2e_spatial_time_table_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"time":{"value":3600}}' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(intersecting, Some(111.0)); + + let conflicting = c + .select( + r#"SELECT pg_catalog.count(*)::bigint AS count + FROM zarr_sample( + 'zarr_e2e_spatial_time_table_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"time":{"index":0}}' + )"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("count") + .unwrap(); + assert_eq!(conflicting, Some(0)); + + let empty = c + .select( + r#"SELECT stats.count, + stats.valid_count, + stats.min, + stats.max, + stats.sum, + stats.avg + FROM zarr_zonal_stats( + 'zarr_e2e_spatial_time_table_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(1000, 1000, 1010, 1010, 3857) + ), + '{"time":{"value":3600}}' + ) AS stats"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(empty.get_by_name::("count").unwrap(), Some(0)); + assert_eq!(empty.get_by_name::("valid_count").unwrap(), Some(0)); + assert_eq!(empty.get_by_name::("min").unwrap(), None); + assert_eq!(empty.get_by_name::("max").unwrap(), None); + assert_eq!(empty.get_by_name::("sum").unwrap(), None); + assert_eq!(empty.get_by_name::("avg").unwrap(), None); + + let by_time_empty = c + .select( + r#"SELECT pg_catalog.count(*)::bigint AS rows, + pg_catalog.count(*) FILTER ( + WHERE stats.count = 0 + AND stats.valid_count = 0 + AND stats.min IS NULL + AND stats.max IS NULL + AND stats.sum IS NULL + AND stats.avg IS NULL + )::bigint AS empty_rows + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_unselected', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(1000, 1000, 1010, 1010, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00', + '{}' + ) AS stats"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + by_time_empty.get_by_name::("rows").unwrap(), + Some(2) + ); + assert_eq!( + by_time_empty.get_by_name::("empty_rows").unwrap(), + Some(2) + ); + }); + + for (sql, expected) in [ + ( + r#"SELECT * + FROM zarr_sample( + 'zarr_e2e_spatial_time_unselected', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact', + '{"x":{"index":1}}' + )"#, + "zarr_sample owns dimension 'x'; selectors may target auxiliary dimensions only", + ), + ( + r#"SELECT * + FROM zarr_zonal_stats( + 'zarr_e2e_spatial_time_unselected', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + '{"y":{"index":1}}' + )"#, + "zarr_zonal_stats owns dimension 'y'; selectors may target auxiliary dimensions only", + ), + ( + r#"SELECT * + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_table_selector', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00', + '{"time":{"index":1}}' + )"#, + "zarr_zonal_stats_by_time owns dimension 'time'; selectors may target auxiliary dimensions only", + ), + ( + r#"SELECT * + FROM zarr_zonal_stats( + 'zarr_e2e_spatial_time_unselected', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + '{}' + )"#, + "auxiliary dimension 'time' resolves to more than one index; spatial operations require zero or one", + ), + ] { + let message = capture_query_error(sql); + assert!(message.contains(expected), "message: {message}"); + } + } + + #[pg_test] + fn zarr_spatial_selector_overloads_preserve_no_remote_io_ordering() { + create_http_e2e_wrapper(); + create_http_e2e_server( + "zarr_http_selector_order_server", + "selector_order", + "normal", + "e2e.zarr", + ); + create_http_time_y_x_table( + "zarr_http_selector_order", + "zarr_http_selector_order_server", + "nested/raw", + false, + ); + + let plan = explain_lines( + r#"SELECT * + FROM zarr_sample( + 'zarr_http_selector_order', + '\x'::bytea, + 'exact', + '{"time":{"index":0}}' + )"#, + ); + assert!( + plan.iter().any(|line| line.contains("Function Scan")), + "plan: {plan:?}" + ); + assert_eq!( + http_case_stats("selector_order")["object_gets"], + serde_json::json!(0) + ); + + let missing_postgis = capture_query_error( + r#"SELECT * + FROM zarr_zonal_stats( + 'zarr_http_selector_order', + '\x'::bytea, + '{"time":{"index":0}}' + )"#, + ); + assert!( + missing_postgis.contains( + "PostGIS is unavailable for zarr spatial operations: the postgis extension is not installed" + ), + "message: {missing_postgis}" + ); + assert_eq!( + http_case_stats("selector_order")["object_gets"], + serde_json::json!(0) + ); + + install_postgis_in_test_schema(); + Spi::run( + r#"CREATE ROLE zarr_selector_no_access; + GRANT USAGE ON SCHEMA zarr_gis TO zarr_selector_no_access"#, + ) + .unwrap(); + let inaccessible = capture_query_error_as_role( + "zarr_selector_no_access", + r#"SELECT * + FROM zarr_zonal_stats( + 'zarr_http_selector_order', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + '{"time":{"index":0}}' + )"#, + ); + assert!( + inaccessible.contains( + "foreign table 'zarr_http_selector_order' does not exist or is not accessible" + ), + "message: {inaccessible}" + ); + assert_eq!( + http_case_stats("selector_order")["object_gets"], + serde_json::json!(0) + ); + } + + #[pg_test( + error = "column 'band' has incompatible PostgreSQL type OID 23; expected double precision (OID 701)" + )] + fn zarr_minio_generic_dimension_rejects_wrong_type() { + create_minio_generic4d_table("zarr_e2e_bad_generic_dimension", "integer"); + + Spi::connect(|c| { + c.select( + "SELECT measurement FROM zarr_e2e_bad_generic_dimension WHERE band = 130", + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test] + fn zarr_postgis_point_sample_exact_and_transformed_nearest() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d"); + + Spi::connect(|c| { + c.select( + "SELECT pg_catalog.set_config('search_path', 'zarr_gis, public, pg_catalog', true)", + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + let exact = c + .select( + r#"SELECT sample.* + FROM zarr_sample( + 'zarr_e2e_spatial2d', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'exact' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(exact.get_by_name::("x").unwrap(), Some(110.0)); + assert_eq!(exact.get_by_name::("y").unwrap(), Some(20.0)); + assert_eq!(exact.get_by_name::("value").unwrap(), Some(42.0)); + assert_eq!(exact.get_by_name::("x_index").unwrap(), Some(1)); + assert_eq!(exact.get_by_name::("y_index").unwrap(), Some(1)); + assert_eq!( + exact.get_by_name::("coordinate_distance").unwrap(), + Some(0.0) + ); + assert_eq!(exact.get_by_name::("srid").unwrap(), Some(3857)); + + let transformed = c + .select( + r#"SELECT sample.* + FROM zarr_sample( + 'zarr_e2e_spatial2d', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID( + zarr_gis.ST_Point( + 0.000988146812531, + 0.000179663056824 + ), + 4326 + ) + ), + 'nearest' + ) AS sample"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(transformed.get_by_name::("x").unwrap(), Some(110.0)); + assert_eq!(transformed.get_by_name::("y").unwrap(), Some(20.0)); + assert_eq!( + transformed.get_by_name::("value").unwrap(), + Some(42.0) + ); + + let misses = c + .select( + r#"SELECT count(*)::bigint AS count + FROM zarr_sample( + 'zarr_e2e_spatial2d', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(111, 20), 3857) + ), + 'exact' + )"#, + Some(1), + &[], + ) + .unwrap() + .first() + .get_by_name::("count") + .unwrap(); + assert_eq!(misses, Some(0)); + }); + } + + #[pg_test( + error = "invalid PostGIS geometry: point sampling method must be 'exact' or 'nearest', got 'bilinear'" + )] + fn zarr_postgis_point_sample_rejects_unknown_method() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_bad_method"); + Spi::run( + r#"SELECT * + FROM zarr_sample( + 'zarr_e2e_spatial2d_bad_method', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ), + 'bilinear' + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "zarr array metadata missing or invalid: spatial operations require exactly one value column" + )] + fn zarr_postgis_point_exact_miss_still_validates_value_column() { + install_postgis_in_test_schema(); + create_minio_spatial2d_coordinate_only_table("zarr_e2e_spatial2d_point_no_value"); + Spi::run( + r#"SELECT * + FROM zarr_sample( + 'zarr_e2e_spatial2d_point_no_value', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(111, 20), 3857) + ), + 'exact' + )"#, + ) + .unwrap(); + } + + #[pg_test] + fn zarr_postgis_cells_include_boundary_and_transform_region_crs() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_cells"); + + let rows = Spi::connect(|c| { + c.select( + "SELECT pg_catalog.set_config('search_path', 'zarr_gis, public, pg_catalog', true)", + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + c.select( + r#"WITH regions(label, region_ewkb) AS ( + VALUES + ( + 'boundary', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) + ), + ( + 'transformed', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope( + 0.000943231048325, + 0.000134747292628, + 0.001212725633561, + 0.000404241877857, + 4326 + ) + ) + ) + ) + SELECT regions.label, cells.* + FROM regions + CROSS JOIN LATERAL zarr_cells( + 'zarr_e2e_spatial2d_cells', + regions.region_ewkb + ) AS cells + ORDER BY regions.label, cells.y_index, cells.x_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("label").unwrap().unwrap(), + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("y").unwrap().unwrap(), + row.get_by_name::("value").unwrap(), + row.get_by_name::("x_index").unwrap().unwrap(), + row.get_by_name::("y_index").unwrap().unwrap(), + row.get_by_name::("srid").unwrap().unwrap(), + ) + }) + .collect::>() + }); + + let mut expected = Vec::new(); + for label in ["boundary", "transformed"] { + for y_index in 1_i64..=3 { + for x_index in 1_i64..=3 { + expected.push(( + label.to_string(), + 100.0 + 10.0 * x_index as f64, + 10.0 + 10.0 * y_index as f64, + Some(42.0), + x_index, + y_index, + 3857, + )); + } + } + } + assert_eq!(rows, expected); + } + + #[pg_test] + fn zarr_postgis_zonal_stats_preserve_fill_value_semantics() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_zonal"); + + Spi::connect(|c| { + c.select( + "SELECT pg_catalog.set_config('search_path', 'zarr_gis, public, pg_catalog', true)", + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + let stats = c + .select( + r#"SELECT stats.* + FROM zarr_zonal_stats( + 'zarr_e2e_spatial2d_zonal', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) + ) AS stats"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(stats.get_by_name::("count").unwrap(), Some(9)); + assert_eq!(stats.get_by_name::("valid_count").unwrap(), Some(9)); + assert_eq!(stats.get_by_name::("min").unwrap(), Some(42.0)); + assert_eq!(stats.get_by_name::("max").unwrap(), Some(42.0)); + assert_eq!(stats.get_by_name::("sum").unwrap(), Some(378.0)); + assert_eq!(stats.get_by_name::("avg").unwrap(), Some(42.0)); + assert_eq!(stats.get_by_name::("srid").unwrap(), Some(3857)); + }); + } + + #[pg_test] + fn zarr_postgis_cells_reject_non_polygon_geometry() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_bad_region"); + Spi::run( + "SELECT pg_catalog.set_config('search_path', 'zarr_gis, public, pg_catalog', true)", + ) + .unwrap(); + + let result = std::panic::catch_unwind(|| { + Spi::run( + r#"SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_bad_region', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID(zarr_gis.ST_Point(110, 20), 3857) + ) + )"#, + ) + }); + let failed = match result { + Ok(result) => result.is_err(), + Err(_) => true, + }; + assert!(failed, "expected zarr_cells to reject a Point region"); + } + + #[pg_test( + error = "zarr array metadata missing or invalid: spatial operations require exactly one value column" + )] + fn zarr_postgis_nonoverlap_still_validates_value_column() { + install_postgis_in_test_schema(); + create_minio_spatial2d_coordinate_only_table("zarr_e2e_spatial2d_region_no_value"); + Spi::run( + r#"SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_region_no_value', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(1000, 1000, 1010, 1010, 3857) + ) + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "invalid CRS metadata for zarr array 'nested/spatial2d': EPSG:999999 is not present in the installed PostGIS spatial_ref_sys" + )] + fn zarr_postgis_cells_reject_unknown_source_srid() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_unknown_srid"); + Spi::run( + r#"SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_unknown_srid', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_SetSRID( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857), + 999999 + ) + ) + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "invalid PostGIS geometry: PostGIS could not parse or transform the supplied polygon" + )] + fn zarr_postgis_cells_normalize_malformed_ewkb_error() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_malformed_ewkb"); + Spi::run( + r#"SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_malformed_ewkb', + '\x0102'::bytea + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "PostGIS is unavailable for zarr spatial operations: the postgis extension is not installed" + )] + fn zarr_postgis_cells_fail_cleanly_without_postgis() { + create_minio_spatial2d_table("zarr_e2e_spatial2d_no_postgis"); + Spi::run( + r#"SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_no_postgis', + '\x'::bytea + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "zarr array metadata missing or invalid: foreign table 'zarr_e2e_spatial2d_private' does not exist or is not accessible" + )] + fn zarr_postgis_cells_enforce_foreign_table_privileges() { + install_postgis_in_test_schema(); + create_minio_spatial2d_table("zarr_e2e_spatial2d_private"); + Spi::run( + r#"CREATE ROLE zarr_spatial_no_access; + GRANT USAGE ON SCHEMA zarr_gis TO zarr_spatial_no_access; + SET ROLE zarr_spatial_no_access; + SELECT * + FROM zarr_cells( + 'zarr_e2e_spatial2d_private', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ) + )"#, + ) + .unwrap(); + } + + #[pg_test] + fn zarr_postgis_cells_by_time_honor_half_open_bounds() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_cells", false); + + let rows = Spi::connect(|c| { + c.select( + r#"SELECT extract(epoch FROM cells.time)::double precision AS epoch, + cells.time_index, + pg_catalog.count(*)::bigint AS count, + pg_catalog.min(cells.value) AS min, + pg_catalog.max(cells.value) AS max, + pg_catalog.sum(cells.value)::double precision AS sum + FROM zarr_cells_by_time( + 'zarr_e2e_spatial_time_cells', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00' + ) AS cells + GROUP BY cells.time, cells.time_index + ORDER BY cells.time_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("epoch").unwrap().unwrap(), + row.get_by_name::("time_index").unwrap().unwrap(), + row.get_by_name::("count").unwrap().unwrap(), + row.get_by_name::("min").unwrap().unwrap(), + row.get_by_name::("max").unwrap().unwrap(), + row.get_by_name::("sum").unwrap().unwrap(), + ) + }) + .collect::>() + }); + + assert_eq!( + rows, + vec![ + (0.0, 0, 9, 11.0, 33.0, 198.0), + (3.6, 1, 9, 111.0, 133.0, 1098.0) + ] + ); + + let first_slice_count = Spi::get_one::( + r#"SELECT pg_catalog.count(*)::bigint + FROM zarr_cells_by_time( + 'zarr_e2e_spatial_time_cells', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.6+00' + )"#, + ) + .unwrap(); + assert_eq!(first_slice_count, Some(9)); + } + + #[pg_test] + fn zarr_postgis_zonal_stats_by_time_preserve_scientific_semantics() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_zonal", false); + + let rows = Spi::connect(|c| { + c.select( + r#"SELECT extract(epoch FROM stats.time)::double precision AS epoch, + stats.time_index, + stats.count, + stats.valid_count, + stats.min, + stats.max, + stats.sum, + stats.avg, + stats.srid + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_zonal', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00' + ) AS stats + ORDER BY stats.time_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("epoch").unwrap().unwrap(), + row.get_by_name::("time_index").unwrap().unwrap(), + row.get_by_name::("count").unwrap().unwrap(), + row.get_by_name::("valid_count").unwrap().unwrap(), + row.get_by_name::("min").unwrap(), + row.get_by_name::("max").unwrap(), + row.get_by_name::("sum").unwrap(), + row.get_by_name::("avg").unwrap(), + row.get_by_name::("srid").unwrap().unwrap(), + ) + }) + .collect::>() + }); + + assert_eq!( + rows, + vec![ + ( + 0.0, + 0, + 9, + 9, + Some(11.0), + Some(33.0), + Some(198.0), + Some(22.0), + 3857, + ), + ( + 3.6, + 1, + 9, + 9, + Some(111.0), + Some(133.0), + Some(1098.0), + Some(122.0), + 3857, + ), + ] + ); + } + + #[pg_test] + fn zarr_postgis_zonal_stats_by_time_keep_decoded_null_slices() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_decoded", true); + + let rows = Spi::connect(|c| { + c.select( + r#"SELECT stats.time_index, + stats.count, + stats.valid_count, + stats.min, + stats.max, + stats.sum, + stats.avg + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_decoded', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(140, 40, 150, 50, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00' + ) AS stats + ORDER BY stats.time_index"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("time_index").unwrap().unwrap(), + row.get_by_name::("count").unwrap().unwrap(), + row.get_by_name::("valid_count").unwrap().unwrap(), + row.get_by_name::("min").unwrap(), + row.get_by_name::("max").unwrap(), + row.get_by_name::("sum").unwrap(), + row.get_by_name::("avg").unwrap(), + ) + }) + .collect::>() + }); + + assert_eq!( + rows, + vec![ + (0, 4, 0, None, None, None, None), + (1, 4, 0, None, None, None, None), + ] + ); + } + + #[pg_test] + fn zarr_postgis_time_range_validation_and_empty_selection() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_range", false); + + let no_rows = Spi::get_one::( + r#"SELECT pg_catalog.count(*)::bigint + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_range', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:10+00', + TIMESTAMPTZ '1970-01-01 00:00:20+00' + )"#, + ) + .unwrap(); + assert_eq!(no_rows, Some(0)); + + Spi::connect(|c| { + let empty = c + .select( + r#"SELECT pg_catalog.count(*)::bigint AS rows, + pg_catalog.count(*) FILTER ( + WHERE stats.count = 0 + AND stats.valid_count = 0 + AND stats.min IS NULL + AND stats.max IS NULL + AND stats.sum IS NULL + AND stats.avg IS NULL + )::bigint AS empty_rows + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_range', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(1000, 1000, 1010, 1010, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.600001+00' + ) AS stats"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(empty.get_by_name::("rows").unwrap(), Some(2)); + assert_eq!(empty.get_by_name::("empty_rows").unwrap(), Some(2)); + }); + + let invalid = std::panic::catch_unwind(|| { + Spi::run( + r#"SELECT * + FROM zarr_zonal_stats_by_time( + 'zarr_e2e_spatial_time_range', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:03.6+00', + TIMESTAMPTZ '1970-01-01 00:00:03.6+00' + )"#, + ) + }); + let failed = match invalid { + Ok(result) => result.is_err(), + Err(_) => true, + }; + assert!(failed, "expected an empty time range to be rejected"); + } + + #[pg_test( + error = "PostGIS is unavailable for zarr spatial operations: the postgis extension is not installed" + )] + fn zarr_postgis_cells_by_time_fail_cleanly_without_postgis() { + create_minio_spatial_time_table("zarr_e2e_spatial_time_no_postgis", false); + Spi::run( + r#"SELECT * + FROM zarr_cells_by_time( + 'zarr_e2e_spatial_time_no_postgis', + '\x'::bytea, + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.6+00' + )"#, + ) + .unwrap(); + } + + #[pg_test( + error = "zarr array metadata missing or invalid: foreign table 'zarr_e2e_spatial_time_private' does not exist or is not accessible" + )] + fn zarr_postgis_cells_by_time_enforce_foreign_table_privileges() { + install_postgis_in_test_schema(); + create_minio_spatial_time_table("zarr_e2e_spatial_time_private", false); + Spi::run( + r#"CREATE ROLE zarr_spatial_time_no_access; + GRANT USAGE ON SCHEMA zarr_gis TO zarr_spatial_time_no_access; + SET ROLE zarr_spatial_time_no_access; + SELECT * + FROM zarr_cells_by_time( + 'zarr_e2e_spatial_time_private', + zarr_gis.ST_AsEWKB( + zarr_gis.ST_MakeEnvelope(110, 20, 130, 40, 3857) + ), + TIMESTAMPTZ '1970-01-01 00:00:00+00', + TIMESTAMPTZ '1970-01-01 00:00:03.6+00' + )"#, + ) + .unwrap(); + } + + #[pg_test] + fn zarr_multiscales_minio_ome_v05_discovery_e2e() { + create_minio_ome_v3_e2e_server(); + + Spi::connect(|c| { + let rows = c + .select( + r#"SELECT group_path, multiscale_index, multiscale_name, + level_index, array_path, axes, shape, chunks, + dtype, codecs, scale, translation, supported, + warnings + FROM zarr_multiscales('zarr_ome_v3_e2e_server') + ORDER BY group_path, multiscale_index, level_index"#, + None, + &[], + ) + .unwrap() + .collect::>(); + assert_eq!(rows.len(), 2); + + let expected = [ + ( + 0_i64, + "image/0", + serde_json::json!([4, 4]), + serde_json::json!([3, 3]), + vec![4.0, 12.0], + vec![120.0, 260.0], + ), + ( + 1_i64, + "image/1", + serde_json::json!([2, 2]), + serde_json::json!([2, 2]), + vec![8.0, 24.0], + vec![122.0, 266.0], + ), + ]; + for (row, (level, array_path, shape, chunks, scale, translation)) in + rows.iter().zip(expected) + { + assert_eq!( + row.get_by_name::("group_path").unwrap(), + Some("image".to_string()) + ); + assert_eq!( + row.get_by_name::("multiscale_index").unwrap(), + Some(0) + ); + assert_eq!( + row.get_by_name::("multiscale_name").unwrap(), + Some("mean-pyramid".to_string()) + ); + assert_eq!( + row.get_by_name::("level_index").unwrap(), + Some(level) + ); + assert_eq!( + row.get_by_name::("array_path").unwrap(), + Some(array_path.to_string()) + ); + assert_eq!( + row.get_by_name::("axes").unwrap().unwrap().0, + serde_json::json!([ + {"name": "y", "type": "space", "unit": "micrometer"}, + {"name": "x", "type": "space", "unit": "micrometer"} + ]) + ); + assert_eq!( + row.get_by_name::("shape").unwrap().unwrap().0, + shape + ); + assert_eq!( + row.get_by_name::("chunks").unwrap().unwrap().0, + chunks + ); + assert_eq!( + row.get_by_name::("dtype").unwrap(), + Some("float32".to_string()) + ); + assert_eq!( + row.get_by_name::("codecs").unwrap().unwrap().0, + serde_json::json!([{ + "name": "bytes", + "configuration": {"endian": "little"} + }]) + ); + assert_eq!( + row.get_by_name::, _>("scale").unwrap(), + Some(scale) + ); + assert_eq!( + row.get_by_name::, _>("translation").unwrap(), + Some(translation) + ); + assert_eq!(row.get_by_name::("supported").unwrap(), Some(true)); + assert_eq!( + row.get_by_name::, _>("warnings").unwrap(), + Some(Vec::new()) + ); + } + + let paths = c + .select( + "SELECT path FROM zarr_inspect('zarr_ome_v3_e2e_server') ORDER BY path", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("path").unwrap()) + .collect::>(); + assert_eq!(paths, vec!["/", "image", "image/0", "image/1"]); + }); + } + + #[pg_test] + fn zarr_minio_ome_v05_explicit_level_scans_e2e() { + create_minio_ome_v3_e2e_server(); + create_minio_ome_v3_e2e_table_on_server("zarr_ome_v05_level0", 0); + create_minio_ome_v3_e2e_table_on_server("zarr_ome_v05_level1", 1); + + Spi::connect(|c| { + let level0 = c + .select( + "SELECT y, x, value FROM zarr_ome_v05_level0 ORDER BY y, x", + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("y").unwrap().unwrap(), + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + level0, + vec![ + (120.0, 260.0, 0.0), + (120.0, 272.0, 1.0), + (120.0, 284.0, 2.0), + (120.0, 296.0, 3.0), + (124.0, 260.0, 4.0), + (124.0, 272.0, 5.0), + (124.0, 284.0, 6.0), + (124.0, 296.0, 7.0), + (128.0, 260.0, 8.0), + (128.0, 272.0, 9.0), + (128.0, 284.0, 10.0), + (128.0, 296.0, 11.0), + (132.0, 260.0, 12.0), + (132.0, 272.0, 13.0), + (132.0, 284.0, 14.0), + (132.0, 296.0, 15.0), + ] + ); + + let level1 = c + .select( + "SELECT y, x, value FROM zarr_ome_v05_level1 ORDER BY y, x", + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("y").unwrap().unwrap(), + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + level1, + vec![ + (122.0, 266.0, 2.5), + (122.0, 290.0, 4.5), + (130.0, 266.0, 10.5), + (130.0, 290.0, 12.5), + ] + ); + }); + } + + #[pg_test] + fn zarr_minio_ome_v05_affine_pruning_metrics_e2e() { + create_minio_ome_v3_e2e_table("zarr_ome_v05_pruning", 0); + + Spi::connect(|c| { + let value = c + .select( + "SELECT value FROM zarr_ome_v05_pruning WHERE y = 132 AND x = 296", + None, + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(value, Some(15.0)); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_ome_v05_pruning + WHERE y = 132 AND x = 296"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Array: image/0"), "plan: {plan:?}"); + assert!(has("Zarr Dimensions: [y, x]"), "plan: {plan:?}"); + assert!(has("Zarr Shape: [4, 4]"), "plan: {plan:?}"); + assert!(has("Zarr Chunk Shape: [3, 3]"), "plan: {plan:?}"); + assert!(has("Zarr OME Group: image"), "plan: {plan:?}"); + assert!(has("Zarr OME Multiscale Index: 0"), "plan: {plan:?}"); + assert!(has("Zarr OME Level Index: 0"), "plan: {plan:?}"); + assert!( + has("Zarr OME Effective Scale: [4.0, 12.0]"), + "plan: {plan:?}" + ); + assert!( + has("Zarr OME Effective Translation: [120.0, 260.0]"), + "plan: {plan:?}" + ); + assert!(has("Zarr Chunks Total: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Coordinate-Pruned: 3"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 1"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 0"), "plan: {plan:?}"); + assert!( + has("Zarr Coordinate Encoded Bytes: 0 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 36 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Data Decoded Bytes: 36 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_ome_v05_aggregate_pushdown_e2e() { + create_minio_ome_v3_e2e_server(); + create_minio_ome_v3_e2e_table_on_server("zarr_ome_v05_aggregate0", 0); + create_minio_ome_v3_e2e_table_on_server("zarr_ome_v05_aggregate1", 1); + + for (table, count, sum, minimum, maximum) in [ + ( + "zarr_ome_v05_aggregate0", + 16_i64, + 120.0_f32, + 0.0_f32, + 15.0_f32, + ), + ( + "zarr_ome_v05_aggregate1", + 4_i64, + 30.0_f32, + 2.5_f32, + 12.5_f32, + ), + ] { + let sql = format!( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM {table}"# + ); + assert_aggregate_pushed_down(&sql); + Spi::connect(|c| { + let row = c.select(&sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + row.get_by_name::("total_count").unwrap(), + Some(count) + ); + assert_eq!( + row.get_by_name::("value_count").unwrap(), + Some(count) + ); + assert_eq!(row.get_by_name::("value_sum").unwrap(), Some(sum)); + assert_eq!(row.get_by_name::("value_avg").unwrap(), Some(7.5)); + assert_eq!( + row.get_by_name::("value_min").unwrap(), + Some(minimum) + ); + assert_eq!( + row.get_by_name::("value_max").unwrap(), + Some(maximum) + ); + }); + } + + for table in ["zarr_ome_v05_aggregate0", "zarr_ome_v05_aggregate1"] { + Spi::connect(|c| { + let row = c + .select( + &format!( + r#"SELECT count(*) AS cells, + sum(value) AS value_sum + FROM {table} + WHERE y BETWEEN 122 AND 130 + AND x BETWEEN 266 AND 290"# + ), + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!(row.get_by_name::("cells").unwrap(), Some(4)); + assert_eq!(row.get_by_name::("value_sum").unwrap(), Some(30.0)); + }); + } + } + + #[pg_test] + fn zarr_minio_ome_v05_selector_rejections_e2e() { + create_minio_ome_v3_e2e_server(); + + let partial = capture_query_error( + r#"CREATE FOREIGN TABLE zarr_ome_v05_partial ( + y double precision, x double precision, value real + ) SERVER zarr_ome_v3_e2e_server + OPTIONS (multiscale_group 'image')"#, + ); + assert!( + partial.contains( + "multiscale_group, multiscale_index, and multiscale_level must be provided together" + ), + "message: {partial}" + ); + + let conflicting = capture_query_error( + r#"CREATE FOREIGN TABLE zarr_ome_v05_conflicting ( + y double precision, x double precision, value real + ) SERVER zarr_ome_v3_e2e_server + OPTIONS ( + array_group 'image/0', + multiscale_group 'image', + multiscale_index '0', + multiscale_level '0' + )"#, + ); + assert!( + conflicting + .contains("array_group cannot be combined with multiscale selection options"), + "message: {conflicting}" + ); + + create_minio_ome_v3_e2e_table_on_server("zarr_ome_v05_bad_level", 2); + let bad_level = capture_query_error("SELECT * FROM zarr_ome_v05_bad_level"); + assert!( + bad_level.contains("multiscale level 2 is outside"), + "message: {bad_level}" + ); + + Spi::run( + r#"CREATE FOREIGN TABLE zarr_ome_v05_bad_index ( + y double precision, x double precision, value real + ) SERVER zarr_ome_v3_e2e_server + OPTIONS ( + multiscale_group 'image', + multiscale_index '1', + multiscale_level '0' + )"#, + ) + .unwrap(); + let bad_index = capture_query_error("SELECT * FROM zarr_ome_v05_bad_index"); + assert!( + bad_index.contains("multiscale index 1 is outside"), + "message: {bad_index}" + ); + } + + #[pg_test] + fn zarr_minio_v3_default_and_v2_chunk_keys_scan_e2e() { + create_minio_v3_e2e_server(); + for (table, array_group) in [ + ("zarr_v3_default_keys", "nested/raw_default"), + ("zarr_v3_v2_keys", "nested/raw_v2keys"), + ] { + create_minio_v3_e2e_table_on_server(table, array_group, false); + Spi::connect(|c| { + let summary = c + .select( + &format!( + r#"SELECT count(*) AS row_count, + count(value) AS value_count, + sum(value)::double precision AS value_sum, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max + FROM {table}"# + ), + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary.get_by_name::("row_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + summary + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3574.0 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + + let boundary = c + .select( + &format!( + r#"SELECT value + FROM {table} + WHERE time = '1970-01-01 00:00:03.6+00'::timestamptz + AND y = 50 + AND x = 150"# + ), + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(boundary, vec![-7.5]); + }); + } + } + + #[pg_test] + fn zarr_minio_v3_sharding_start_end_scan_and_sparse_fill_e2e() { + create_minio_v3_e2e_server(); + for (table, array_group) in [ + ("zarr_v3_shard_end", "nested/shard_end"), + ("zarr_v3_shard_start", "nested/shard_start"), + ] { + create_minio_v3_e2e_table_on_server(table, array_group, false); + let aggregate_sql = format!( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM {table}"# + ); + assert_aggregate_pushed_down(&aggregate_sql); + + Spi::connect(|c| { + let summary = c.select(&aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + summary + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3_574.0 + ); + assert!( + (summary.get_by_name::("value_avg").unwrap().unwrap() + - 59.566_666_666_666_67) + .abs() + < 1e-12 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + let fill_count = c + .select( + &format!("SELECT count(*) FROM {table} WHERE value = -7.5"), + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(); + assert_eq!(fill_count, Some(8)); + + // The pinned start-index shard stores these two logical inner + // chunks out of C-order physically. Correct values prove that + // the index offsets, rather than payload order, drive reads. + let morton_order_probe = c + .select( + &format!( + r#"SELECT x, value + FROM {table} + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130 + ORDER BY x"# + ), + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + morton_order_probe, + vec![(110.0, 11.0), (120.0, 12.0), (130.0, 13.0)] + ); + + let absent_shard_fill = c + .select( + &format!( + r#"SELECT value + FROM {table} + WHERE time = '1970-01-01 00:00:03.6+00'::timestamptz + AND y = 50 + AND x = 150"# + ), + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(absent_shard_fill, Some(-7.5)); + }); + } + } + + #[pg_test] + fn zarr_minio_v3_sharding_scientific_aggregate_pushdown_e2e() { + create_minio_v3_e2e_table("zarr_v3_shard_cf", "nested/shard_end", true); + let sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_shard_cf"#; + assert_aggregate_pushed_down(sql); + + Spi::connect(|c| { + let row = c.select(sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + row.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + row.get_by_name::("value_count").unwrap().unwrap(), + 48 + ); + assert!( + (row.get_by_name::("value_sum").unwrap().unwrap() - 13_142.86).abs() < 1e-8 + ); + assert!( + (row.get_by_name::("value_avg").unwrap().unwrap() - 273.809_583_333_333_36) + .abs() + < 1e-10 + ); + assert!( + (row.get_by_name::("value_min").unwrap().unwrap() - 273.15).abs() < 1e-10 + ); + assert!( + (row.get_by_name::("value_max").unwrap().unwrap() - 274.55).abs() < 1e-10 + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_sharding_missing_inner_sentinel_uses_fill_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server( + "zarr_v3_shard_sentinel_raw", + "nested/shard_sentinel", + false, + ); + create_minio_v3_e2e_table_on_server( + "zarr_v3_shard_sentinel_cf", + "nested/shard_sentinel", + true, + ); + + Spi::connect(|c| { + let raw = c + .select( + r#"SELECT value + FROM zarr_v3_shard_sentinel_raw + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(raw, Some(-7.5)); + + let decoded = c + .select( + r#"SELECT value + FROM zarr_v3_shard_sentinel_cf + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(decoded, None); + }); + } + + #[pg_test] + fn zarr_minio_v3_sharding_index_corruption_fails_closed_e2e() { + create_minio_v3_e2e_server(); + let cases: [(&str, &str, &[&str]); 4] = [ + ( + "zarr_v3_shard_bad_index_crc", + "nested/shard_bad_index_crc", + &["shard index codec index 1 ('crc32c')", "checksum mismatch"], + ), + ( + "zarr_v3_shard_truncated_index", + "nested/shard_truncated_index", + &["expected exactly the final 68 bytes"], + ), + ( + "zarr_v3_shard_oob", + "nested/shard_oob", + &["inner chunk byte range", "exceeds shard object length"], + ), + ( + "zarr_v3_shard_half_sentinel", + "nested/shard_half_sentinel", + &[ + "uses a mixed uint64 missing sentinel", + "offset and nbytes must both be 2^64 - 1", + ], + ), + ]; + + for (table, array_group, expected_phrases) in cases { + create_minio_v3_e2e_table_on_server(table, array_group, false); + let message = capture_query_error(&format!( + r#"SELECT value + FROM {table} + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"# + )); + assert!( + message.contains(&format!("{array_group}/c/0/0/0")), + "message: {message}" + ); + for phrase in expected_phrases { + assert!(message.contains(*phrase), "message: {message}"); + } + } + } + + #[pg_test] + fn zarr_minio_v3_sharding_bounded_range_metrics_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_shard_ranges", "nested/shard_end", false); + create_minio_v3_e2e_table_on_server( + "zarr_v3_shard_start_range", + "nested/shard_start", + false, + ); + + Spi::connect(|c| { + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_shard_ranges + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!( + has("Zarr Storage Layout: sharding_indexed (index: end)"), + "plan: {plan:?}" + ); + assert!(has("Zarr Shard Shape: [2, 3, 4]"), "plan: {plan:?}"); + assert!(has("Zarr Chunk Shape: [1, 3, 2]"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Location: end"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 2"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 2"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 2"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 3"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 116 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Cache Hits: 0"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Shard Index Encoded Bytes: 68 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Shard Payload Encoded Bytes: 48 bytes"), + "plan: {plan:?}" + ); + + let start_plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_shard_start_range + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 130"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let start_has = |text: &str| start_plan.iter().any(|line| line.contains(text)); + assert!( + start_has("Zarr Storage Layout: sharding_indexed (index: start)"), + "plan: {start_plan:?}" + ); + assert!( + start_has("Zarr Shard Index Location: start"), + "plan: {start_plan:?}" + ); + assert!(start_has("Zarr Data GET Calls: 2"), "plan: {start_plan:?}"); + assert!( + start_has("Zarr Data Encoded Bytes: 92 bytes"), + "plan: {start_plan:?}" + ); + assert!( + start_has("Zarr Shard Index GET Calls: 1"), + "plan: {start_plan:?}" + ); + assert!( + start_has("Zarr Shard Payload GET Calls: 1"), + "plan: {start_plan:?}" + ); + assert!( + start_has("Zarr Shard Index Encoded Bytes: 68 bytes"), + "plan: {start_plan:?}" + ); + assert!( + start_has("Zarr Shard Payload Encoded Bytes: 24 bytes"), + "plan: {start_plan:?}" + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_sharding_rescan_cache_e2e() { + create_minio_v3_e2e_table("zarr_v3_shard_rescan", "nested/shard_end", false); + + Spi::connect(|c| { + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT ordinal, + (SELECT count(*) + FROM zarr_v3_shard_rescan + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND upper_x) AS selected + FROM (VALUES (1, 130.0::double precision), + (2, 130.0::double precision)) AS limits(ordinal, upper_x) + ORDER BY ordinal"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Requested: 4"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 4"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 3"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 116 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Cache Hits: 2"), "plan: {plan:?}"); + assert!(has("Zarr Cache Misses: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 3"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!(has("Zarr Rescans: 1"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_v3_ordered_codec_pipeline_scan_e2e() { + create_minio_v3_e2e_table("zarr_v3_pipeline", "nested/pipeline", false); + + Spi::connect(|c| { + let summary = c + .select( + r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value)::double precision AS value_sum, + avg(value) AS value_avg, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max, + count(*) FILTER (WHERE value = -7.5) AS fill_count + FROM zarr_v3_pipeline"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + summary + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + summary.get_by_name::("value_sum").unwrap().unwrap(), + 3_574.0 + ); + assert!( + (summary.get_by_name::("value_avg").unwrap().unwrap() + - 59.566_666_666_666_67) + .abs() + < 1e-12 + ); + assert_eq!( + summary.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + summary.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + assert_eq!( + summary + .get_by_name::("fill_count") + .unwrap() + .unwrap(), + 8 + ); + + for (time, y, x, expected) in [ + ("1970-01-01 00:00:00+00", 20, 110, 11.0_f32), + ("1970-01-01 00:00:00+00", 50, 130, 43.0_f32), + ("1970-01-01 00:00:03.6+00", 20, 110, 111.0_f32), + ("1970-01-01 00:00:03.6+00", 30, 150, 125.0_f32), + ("1970-01-01 00:00:03.6+00", 50, 150, -7.5_f32), + ] { + let values = c + .select( + &format!( + "SELECT value FROM zarr_v3_pipeline \ + WHERE time = '{time}'::timestamptz AND y = {y} AND x = {x}" + ), + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("value").unwrap()) + .collect::>(); + assert_eq!(values, vec![expected]); + } + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_pipeline + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Codec: transpose -> bytes -> gzip -> crc32c"), + "plan: {plan:?}" + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_codec_pipeline_aggregate_pushdown_e2e() { + create_minio_v3_e2e_table("zarr_v3_pipeline_cf", "nested/pipeline", true); + let sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_pipeline_cf"#; + assert_aggregate_pushed_down(sql); + + Spi::connect(|c| { + let row = c.select(sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + row.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + row.get_by_name::("value_count").unwrap().unwrap(), + 48 + ); + let value_sum = row.get_by_name::("value_sum").unwrap().unwrap(); + let value_avg = row.get_by_name::("value_avg").unwrap().unwrap(); + let value_min = row.get_by_name::("value_min").unwrap().unwrap(); + let value_max = row.get_by_name::("value_max").unwrap().unwrap(); + assert!((value_sum - 13_142.86).abs() < 1e-8); + assert!((value_avg - 273.809_583_333_333_36).abs() < 1e-10); + assert!((value_min - 273.15).abs() < 1e-10); + assert!((value_max - 274.55).abs() < 1e-10); + }); + } + + #[pg_test] + fn zarr_minio_v3_zstd_pipeline_scan_and_aggregate_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_zstd_pipeline", "nested/zstd_pipeline", false); + create_minio_v3_e2e_table_on_server( + "zarr_v3_zstd_pipeline_cf", + "nested/zstd_pipeline", + true, + ); + assert_sparse_cube_cf_aggregate("zarr_v3_zstd_pipeline_cf"); + + Spi::connect(|c| { + let probes = c + .select( + r#"SELECT time, y, x, value + FROM zarr_v3_zstd_pipeline + WHERE (time, y, x) IN ( + ('1970-01-01 00:00:00+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:00+00'::timestamptz, 50, 130), + ('1970-01-01 00:00:03.6+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:03.6+00'::timestamptz, 50, 150) + ) + ORDER BY time, y, x"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("value").unwrap().unwrap()) + .collect::>(); + assert_eq!(probes, vec![11.0, 43.0, 111.0, -7.5]); + + let fill_count = c + .select( + "SELECT count(*) FROM zarr_v3_zstd_pipeline WHERE value = -7.5", + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(); + assert_eq!(fill_count, Some(8)); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_zstd_pipeline + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!( + has("Zarr Codec: transpose -> bytes -> zstd -> crc32c"), + "plan: {plan:?}" + ); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 82 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Data Decoded Bytes: 96 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_v3_zstd_coordinate_decode_and_pruning_e2e() { + create_minio_v3_e2e_server(); + Spi::run( + r#"CREATE FOREIGN TABLE zarr_v3_zstd_coordinate ( + zstd_x double precision, + value real + ) + SERVER zarr_v3_e2e_server + OPTIONS (array_group 'nested/zstd_coord_values')"#, + ) + .unwrap(); + + let aggregate_sql = r#"SELECT count(*) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_zstd_coordinate + WHERE zstd_x BETWEEN 110 AND 140"#; + assert_aggregate_pushed_down(aggregate_sql); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT zstd_x, value + FROM zarr_v3_zstd_coordinate + WHERE zstd_x BETWEEN 110 AND 140 + ORDER BY zstd_x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("zstd_x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + values, + vec![(110.0, 1.0), (120.0, 2.0), (130.0, 3.0), (140.0, 4.0)] + ); + + let aggregate = c.select(aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + aggregate + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 4 + ); + assert_eq!( + aggregate + .get_by_name::("value_sum") + .unwrap() + .unwrap(), + 10.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_avg") + .unwrap() + .unwrap(), + 2.5 + ); + assert_eq!( + aggregate + .get_by_name::("value_min") + .unwrap() + .unwrap(), + 1.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_max") + .unwrap() + .unwrap(), + 4.0 + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_zstd_coordinate + WHERE zstd_x = 150"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Coordinate-Pruned: 1"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 2"), "plan: {plan:?}"); + assert!( + has("Zarr Coordinate Encoded Bytes: 70 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Coordinate Decoded Bytes: 64 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 16 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_v3_zstd_sharded_inner_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_shard_zstd", "nested/shard_zstd", false); + create_minio_v3_e2e_table_on_server("zarr_v3_shard_zstd_cf", "nested/shard_zstd", true); + assert_sparse_cube_cf_aggregate("zarr_v3_shard_zstd_cf"); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT x, value + FROM zarr_v3_shard_zstd + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130 + ORDER BY x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!(values, vec![(110.0, 11.0), (120.0, 12.0), (130.0, 13.0)]); + + let sparse_fill = c + .select( + r#"SELECT value + FROM zarr_v3_shard_zstd + WHERE time = '1970-01-01 00:00:03.6+00'::timestamptz + AND y = 50 + AND x = 150"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(sparse_fill, Some(-7.5)); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_shard_zstd + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Codec: bytes -> zstd"), "plan: {plan:?}"); + assert!( + has("Zarr Storage Layout: sharding_indexed (index: end)"), + "plan: {plan:?}" + ); + assert!(has("Zarr Shard Shape: [2, 3, 4]"), "plan: {plan:?}"); + assert!(has("Zarr Chunk Shape: [2, 3, 4]"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 98 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 0"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Shard Index Encoded Bytes: 20 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Shard Payload Encoded Bytes: 78 bytes"), + "plan: {plan:?}" + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_zstd_frame_policy_fails_closed_e2e() { + create_minio_v3_e2e_server(); + Spi::run( + r#"CREATE FOREIGN TABLE zarr_v3_bad_zstd ( + failure_case double precision, + value real + ) + SERVER zarr_v3_e2e_server + OPTIONS (array_group 'nested/zstd_bad/values')"#, + ) + .unwrap(); + + for (failure_case, key, reason) in [ + ( + 0, + "nested/zstd_bad/values/c/0", + "failed to decode Zstandard frame", + ), + ( + 1, + "nested/zstd_bad/values/c/1", + "Zstandard frame window 16777216 exceeds the 8388608-byte limit", + ), + ( + 2, + "nested/zstd_bad/values/c/2", + "Zstandard dictionaries are not supported", + ), + ] { + let message = capture_query_error(&format!( + "SELECT value FROM zarr_v3_bad_zstd WHERE failure_case = {failure_case}" + )); + assert!(message.contains(key), "message: {message}"); + assert!( + message.contains("codec index 1 ('zstd')"), + "message: {message}" + ); + assert!(message.contains(reason), "message: {message}"); + } + } + + #[pg_test] + fn zarr_minio_v3_blosc_direct_scan_and_aggregate_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_blosc_direct", "nested/blosc_v3", false); + create_minio_v3_e2e_table_on_server("zarr_v3_blosc_direct_cf", "nested/blosc_v3", true); + + let raw_sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value)::double precision AS value_sum, + avg(value) AS value_avg, + min(value)::double precision AS value_min, + max(value)::double precision AS value_max + FROM zarr_v3_blosc_direct"#; + let decoded_sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_blosc_direct_cf"#; + assert_aggregate_pushed_down(decoded_sql); + + Spi::connect(|c| { + let raw = c.select(raw_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + raw.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + raw.get_by_name::("value_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + raw.get_by_name::("value_sum").unwrap().unwrap(), + 3_574.0 + ); + assert!( + (raw.get_by_name::("value_avg").unwrap().unwrap() - 59.566_666_666_666_67) + .abs() + < 1e-12 + ); + assert_eq!( + raw.get_by_name::("value_min").unwrap().unwrap(), + -7.5 + ); + assert_eq!( + raw.get_by_name::("value_max").unwrap().unwrap(), + 143.0 + ); + let fill_count = c + .select( + "SELECT count(*) FROM zarr_v3_blosc_direct WHERE value = -7.5", + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get::(1) + .unwrap(); + assert_eq!(fill_count, Some(8)); + + let decoded = c.select(decoded_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + decoded + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + decoded + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 48 + ); + assert!( + (decoded.get_by_name::("value_sum").unwrap().unwrap() - 13_142.86).abs() + < 1e-8 + ); + assert!( + (decoded.get_by_name::("value_avg").unwrap().unwrap() + - 273.809_583_333_333_36) + .abs() + < 1e-10 + ); + assert!( + (decoded.get_by_name::("value_min").unwrap().unwrap() - 273.15).abs() + < 1e-10 + ); + assert!( + (decoded.get_by_name::("value_max").unwrap().unwrap() - 274.55).abs() + < 1e-10 + ); + + let probes = c + .select( + r#"SELECT time, y, x, value + FROM zarr_v3_blosc_direct + WHERE (time, y, x) IN ( + ('1970-01-01 00:00:00+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:00+00'::timestamptz, 50, 130), + ('1970-01-01 00:00:03.6+00'::timestamptz, 20, 110), + ('1970-01-01 00:00:03.6+00'::timestamptz, 50, 150) + ) + ORDER BY time, y, x"#, + None, + &[], + ) + .unwrap() + .map(|row| row.get_by_name::("value").unwrap().unwrap()) + .collect::>(); + assert_eq!(probes, vec![11.0, 43.0, 111.0, -7.5]); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_blosc_direct + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Codec: bytes -> blosc"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 112 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Data Decoded Bytes: 96 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_v3_blosc_coordinate_decode_and_pruning_e2e() { + create_minio_v3_e2e_server(); + Spi::run( + r#"CREATE FOREIGN TABLE zarr_v3_blosc_coordinate ( + blosc_x double precision, + value real + ) + SERVER zarr_v3_e2e_server + OPTIONS (array_group 'nested/blosc_coord_values')"#, + ) + .unwrap(); + + let aggregate_sql = r#"SELECT count(*) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_blosc_coordinate + WHERE blosc_x BETWEEN 110 AND 140"#; + assert_aggregate_pushed_down(aggregate_sql); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT blosc_x, value + FROM zarr_v3_blosc_coordinate + WHERE blosc_x BETWEEN 110 AND 140 + ORDER BY blosc_x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("blosc_x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!( + values, + vec![(110.0, 1.0), (120.0, 2.0), (130.0, 3.0), (140.0, 4.0)] + ); + + let aggregate = c.select(aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + aggregate + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 4 + ); + assert_eq!( + aggregate + .get_by_name::("value_sum") + .unwrap() + .unwrap(), + 10.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_avg") + .unwrap() + .unwrap(), + 2.5 + ); + assert_eq!( + aggregate + .get_by_name::("value_min") + .unwrap() + .unwrap(), + 1.0 + ); + assert_eq!( + aggregate + .get_by_name::("value_max") + .unwrap() + .unwrap(), + 4.0 + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_blosc_coordinate + WHERE blosc_x = 150"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Chunks Selected: 1"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Coordinate-Pruned: 1"), "plan: {plan:?}"); + assert!(has("Zarr Coordinate GET Calls: 2"), "plan: {plan:?}"); + assert!( + has("Zarr Coordinate Encoded Bytes: 96 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Coordinate Decoded Bytes: 64 bytes"), + "plan: {plan:?}" + ); + assert!(has("Zarr Data GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 16 bytes"), "plan: {plan:?}"); + }); + } + + #[pg_test] + fn zarr_minio_v3_blosc_sharded_inner_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_shard_blosc", "nested/shard_blosc", false); + create_minio_v3_e2e_table_on_server("zarr_v3_shard_blosc_cf", "nested/shard_blosc", true); + + let aggregate_sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_shard_blosc_cf"#; + assert_aggregate_pushed_down(aggregate_sql); + + Spi::connect(|c| { + let values = c + .select( + r#"SELECT x, value + FROM zarr_v3_shard_blosc + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130 + ORDER BY x"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("x").unwrap().unwrap(), + row.get_by_name::("value").unwrap().unwrap(), + ) + }) + .collect::>(); + assert_eq!(values, vec![(110.0, 11.0), (120.0, 12.0), (130.0, 13.0)]); + + let sparse_fill = c + .select( + r#"SELECT value + FROM zarr_v3_shard_blosc + WHERE time = '1970-01-01 00:00:03.6+00'::timestamptz + AND y = 50 + AND x = 150"#, + Some(1), + &[], + ) + .unwrap() + .next() + .unwrap() + .get_by_name::("value") + .unwrap(); + assert_eq!(sparse_fill, Some(-7.5)); + + let aggregate = c.select(aggregate_sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + aggregate + .get_by_name::("total_count") + .unwrap() + .unwrap(), + 60 + ); + assert_eq!( + aggregate + .get_by_name::("value_count") + .unwrap() + .unwrap(), + 48 + ); + assert!( + (aggregate + .get_by_name::("value_sum") + .unwrap() + .unwrap() + - 13_142.86) + .abs() + < 1e-8 + ); + assert!( + (aggregate + .get_by_name::("value_avg") + .unwrap() + .unwrap() + - 273.809_583_333_333_36) + .abs() + < 1e-10 + ); + assert!( + (aggregate + .get_by_name::("value_min") + .unwrap() + .unwrap() + - 273.15) + .abs() + < 1e-10 + ); + assert!( + (aggregate + .get_by_name::("value_max") + .unwrap() + .unwrap() + - 274.55) + .abs() + < 1e-10 + ); + + let plan = c + .select( + r#"EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) + SELECT value + FROM zarr_v3_shard_blosc + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x BETWEEN 110 AND 130"#, + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + let has = |text: &str| plan.iter().any(|line| line.contains(text)); + assert!(has("Zarr Codec: bytes -> blosc"), "plan: {plan:?}"); + assert!( + has("Zarr Storage Layout: sharding_indexed (index: end)"), + "plan: {plan:?}" + ); + assert!(has("Zarr Shard Shape: [2, 3, 4]"), "plan: {plan:?}"); + assert!(has("Zarr Chunk Shape: [1, 3, 2]"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Selected: 2"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Requested: 2"), "plan: {plan:?}"); + assert!(has("Zarr Chunks Present: 2"), "plan: {plan:?}"); + assert!(has("Zarr Data GET Calls: 3"), "plan: {plan:?}"); + assert!(has("Zarr Data Encoded Bytes: 148 bytes"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index GET Calls: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Payload GET Calls: 2"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Hits: 1"), "plan: {plan:?}"); + assert!(has("Zarr Shard Index Cache Misses: 1"), "plan: {plan:?}"); + assert!( + has("Zarr Shard Index Encoded Bytes: 68 bytes"), + "plan: {plan:?}" + ); + assert!( + has("Zarr Shard Payload Encoded Bytes: 80 bytes"), + "plan: {plan:?}" + ); + }); + } + + #[pg_test] + fn zarr_minio_v3_blosc_corrupt_chunks_fail_closed_e2e() { + create_minio_v3_e2e_server(); + create_minio_v3_e2e_table_on_server("zarr_v3_bad_blosc", "nested/bad_blosc", false); + + for (x, key, reason) in [ + ( + 110, + "nested/bad_blosc/c/0/0/0", + "encoded chunk is shorter than the 16-byte Blosc header", + ), + ( + 150, + "nested/bad_blosc/c/0/0/1", + "Blosc header declares 100 uncompressed bytes, expected exactly 96", + ), + ] { + let message = capture_query_error(&format!( + r#"SELECT value + FROM zarr_v3_bad_blosc + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = {x}"# + )); + assert!(message.contains(key), "message: {message}"); + assert!( + message.contains("codec index 1 ('blosc')"), + "message: {message}" + ); + assert!(message.contains(reason), "message: {message}"); + } + } + + #[pg_test] + fn zarr_minio_v3_crc32c_corruption_fails_closed_e2e() { + create_minio_v3_e2e_table("zarr_v3_bad_crc", "nested/bad_crc", false); + let message = capture_query_error( + r#"SELECT value + FROM zarr_v3_bad_crc + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + message.contains("nested/bad_crc/c/0/0/0"), + "message: {message}" + ); + assert!( + message.contains("codec index 3 ('crc32c')"), + "message: {message}" + ); + assert!( + message.to_ascii_lowercase().contains("checksum mismatch"), + "message: {message}" + ); + } + + #[pg_test] + fn zarr_minio_v3_truncated_gzip_fails_closed_e2e() { + create_minio_v3_e2e_table("zarr_v3_bad_gzip", "nested/bad_gzip", false); + let message = capture_query_error( + r#"SELECT value + FROM zarr_v3_bad_gzip + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + message.contains("nested/bad_gzip/c/0/0/0"), + "message: {message}" + ); + assert!( + message.contains("codec index 2 ('gzip')"), + "message: {message}" + ); + } + + #[pg_test] + fn zarr_minio_v3_overexpanding_gzip_fails_closed_e2e() { + create_minio_v3_e2e_table("zarr_v3_oversize", "nested/oversize", false); + let message = capture_query_error( + r#"SELECT value + FROM zarr_v3_oversize + WHERE time = '1970-01-01 00:00:00+00'::timestamptz + AND y = 20 + AND x = 110"#, + ); + assert!( + message.contains("nested/oversize/c/0/0/0"), + "message: {message}" + ); + assert!( + message.contains("codec index 2 ('gzip')") + && message.contains("decoded chunk has more than 96 bytes, expected exactly 96"), + "message: {message}" + ); + } + + #[pg_test] + fn zarr_minio_v3_scalar_aggregate_pushdown_e2e() { + create_minio_v3_e2e_table("zarr_v3_aggregate", "nested/raw_default", true); + let sql = r#"SELECT count(*) AS total_count, + count(value) AS value_count, + sum(value) AS value_sum, + avg(value) AS value_avg, + min(value) AS value_min, + max(value) AS value_max + FROM zarr_v3_aggregate"#; + assert_aggregate_pushed_down(sql); + + Spi::connect(|c| { + let row = c.select(sql, None, &[]).unwrap().next().unwrap(); + assert_eq!( + row.get_by_name::("total_count").unwrap().unwrap(), + 60 + ); + assert_eq!( + row.get_by_name::("value_count").unwrap().unwrap(), + 48 + ); + let value_sum = row.get_by_name::("value_sum").unwrap().unwrap(); + let value_avg = row.get_by_name::("value_avg").unwrap().unwrap(); + let value_min = row.get_by_name::("value_min").unwrap().unwrap(); + let value_max = row.get_by_name::("value_max").unwrap().unwrap(); + assert!((value_sum - 13_142.86).abs() < 1e-8); + assert!((value_avg - 273.809_583_333_333_36).abs() < 1e-10); + assert!((value_min - 273.15).abs() < 1e-10); + assert!((value_max - 274.55).abs() < 1e-10); + }); + } + + #[pg_test] + fn zarr_inspect_minio_v3_metadata_e2e() { + create_minio_v3_e2e_server(); + + Spi::connect(|c| { + let paths = c + .select( + "SELECT path FROM zarr_inspect('zarr_v3_e2e_server') ORDER BY path", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("path").unwrap()) + .collect::>(); + assert_eq!( + paths, + vec![ + "/", + "nested", + "nested/bad_blosc", + "nested/bad_crc", + "nested/bad_gzip", + "nested/blosc_coord_values", + "nested/blosc_v3", + "nested/blosc_x", + "nested/oversize", + "nested/pipeline", + "nested/raw_default", + "nested/raw_v2keys", + "nested/shard_bad_index_crc", + "nested/shard_blosc", + "nested/shard_end", + "nested/shard_half_sentinel", + "nested/shard_oob", + "nested/shard_sentinel", + "nested/shard_start", + "nested/shard_truncated_index", + "nested/shard_zstd", + "nested/time", + "nested/x", + "nested/y", + "nested/zstd_bad", + "nested/zstd_bad/failure_case", + "nested/zstd_bad/values", + "nested/zstd_coord_values", + "nested/zstd_pipeline", + "nested/zstd_x", + ] + ); + + let root = c + .select( + "SELECT kind, zarr_format, attributes, warnings FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = '/'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + root.get_by_name::("kind").unwrap().unwrap(), + "group" + ); + assert_eq!( + root.get_by_name::("zarr_format").unwrap().unwrap(), + 3 + ); + assert_eq!( + root.get_by_name::("attributes") + .unwrap() + .unwrap() + .0["title"], + serde_json::json!("Deterministic Zarr v3 inspection fixture") + ); + assert_eq!( + root.get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + Vec::::new() + ); + + let nested = c + .select( + "SELECT kind, zarr_format, crs, warnings FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + nested.get_by_name::("kind").unwrap().unwrap(), + "group" + ); + assert_eq!( + nested + .get_by_name::("zarr_format") + .unwrap() + .unwrap(), + 3 + ); + let nested_crs = nested.get_by_name::("crs").unwrap().unwrap().0; + assert_eq!( + nested_crs["properties"]["name"], + serde_json::json!("EPSG:3857") + ); + assert_eq!( + nested + .get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + Vec::::new() + ); + + let raw = c + .select( + r#"SELECT kind, group_path, variable, zarr_format, shape, + dimensions, dtype, chunks, codecs, units, + fill_value, scale_factor, add_offset, attributes, + warnings + FROM zarr_inspect('zarr_v3_e2e_server') + WHERE path = 'nested/raw_default'"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + raw.get_by_name::("kind").unwrap().unwrap(), + "array" + ); + assert_eq!( + raw.get_by_name::("group_path").unwrap().unwrap(), + "nested" + ); + assert_eq!( + raw.get_by_name::("variable").unwrap().unwrap(), + "raw_default" + ); + assert_eq!( + raw.get_by_name::("zarr_format").unwrap().unwrap(), + 3 + ); + assert_eq!( + raw.get_by_name::("shape").unwrap().unwrap().0, + serde_json::json!([2, 5, 6]) + ); + assert_eq!( + raw.get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time", "y", "x"] + ); + assert_eq!( + raw.get_by_name::("dtype").unwrap().unwrap(), + "float32" + ); + assert_eq!( + raw.get_by_name::("chunks").unwrap().unwrap().0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + raw.get_by_name::("codecs").unwrap().unwrap().0, + serde_json::json!([{ + "name": "bytes", + "configuration": {"endian": "little"} + }]) + ); + assert_eq!( + raw.get_by_name::("fill_value") + .unwrap() + .unwrap() + .0, + serde_json::json!(-7.5) + ); + assert_eq!(raw.get_by_name::("units").unwrap().unwrap(), "K"); + assert_eq!( + raw.get_by_name::("scale_factor").unwrap().unwrap(), + 0.01 + ); + assert_eq!( + raw.get_by_name::("add_offset").unwrap().unwrap(), + 273.15 + ); + assert_eq!( + raw.get_by_name::("attributes") + .unwrap() + .unwrap() + .0["missing_value"], + serde_json::json!([42.0]) + ); + assert_eq!( + raw.get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + Vec::::new() + ); + + let alternate = c + .select( + "SELECT zarr_format, dimensions, dtype, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/raw_v2keys'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + alternate + .get_by_name::("zarr_format") + .unwrap() + .unwrap(), + 3 + ); + assert_eq!( + alternate + .get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time", "y", "x"] + ); + assert_eq!( + alternate + .get_by_name::("dtype") + .unwrap() + .unwrap(), + "float32" + ); + assert_eq!( + alternate + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([{ + "name": "bytes", + "configuration": {"endian": "little"} + }]) + ); + + let pipeline = c + .select( + "SELECT codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/pipeline'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + pipeline + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([ + {"name": "transpose", "configuration": {"order": [2, 1, 0]}}, + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "gzip", "configuration": {"level": 1}}, + {"name": "crc32c"} + ]) + ); + + let zstd_pipeline = c + .select( + "SELECT chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/zstd_pipeline'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + zstd_pipeline + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + zstd_pipeline + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([ + {"name": "transpose", "configuration": {"order": [2, 1, 0]}}, + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 1, "checksum": true}}, + {"name": "crc32c"} + ]) + ); + + let zstd_coordinate = c + .select( + "SELECT dimensions, dtype, chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/zstd_x'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + zstd_coordinate + .get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["zstd_x"] + ); + assert_eq!( + zstd_coordinate + .get_by_name::("dtype") + .unwrap() + .unwrap(), + "float64" + ); + assert_eq!( + zstd_coordinate + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([4]) + ); + assert_eq!( + zstd_coordinate + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": 1, "checksum": false}} + ]) + ); + + let blosc = c + .select( + "SELECT chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/blosc_v3'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + blosc.get_by_name::("chunks").unwrap().unwrap().0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + blosc.get_by_name::("codecs").unwrap().unwrap().0, + serde_json::json!([ + {"name": "bytes", "configuration": {"endian": "little"}}, + { + "name": "blosc", + "configuration": { + "typesize": 4, + "cname": "lz4", + "clevel": 5, + "shuffle": "shuffle", + "blocksize": 0 + } + } + ]) + ); + + let blosc_coordinate = c + .select( + "SELECT dimensions, dtype, chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/blosc_x'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + blosc_coordinate + .get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["blosc_x"] + ); + assert_eq!( + blosc_coordinate + .get_by_name::("dtype") + .unwrap() + .unwrap(), + "float64" + ); + assert_eq!( + blosc_coordinate + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([4]) + ); + assert_eq!( + blosc_coordinate + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([ + {"name": "bytes", "configuration": {"endian": "little"}}, + { + "name": "blosc", + "configuration": { + "typesize": 8, + "cname": "lz4", + "clevel": 5, + "shuffle": "shuffle", + "blocksize": 0 + } + } + ]) + ); + + let sharded = c + .select( + "SELECT chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/shard_end'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + sharded + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + sharded + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([{ + "name": "sharding_indexed", + "configuration": { + "chunk_shape": [1, 3, 2], + "codecs": [{ + "name": "bytes", + "configuration": {"endian": "little"} + }], + "index_codecs": [ + { + "name": "bytes", + "configuration": {"endian": "little"} + }, + {"name": "crc32c"} + ], + "index_location": "end" + } + }]) + ); + + let sharded_blosc = c + .select( + "SELECT chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/shard_blosc'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + sharded_blosc + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + sharded_blosc + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([{ + "name": "sharding_indexed", + "configuration": { + "chunk_shape": [1, 3, 2], + "codecs": [ + { + "name": "bytes", + "configuration": {"endian": "little"} + }, + { + "name": "blosc", + "configuration": { + "typesize": 4, + "cname": "lz4", + "clevel": 5, + "shuffle": "shuffle", + "blocksize": 0 + } + } + ], + "index_codecs": [ + { + "name": "bytes", + "configuration": {"endian": "little"} + }, + {"name": "crc32c"} + ], + "index_location": "end" + } + }]) + ); + + let sharded_zstd = c + .select( + "SELECT chunks, codecs FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/shard_zstd'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + sharded_zstd + .get_by_name::("chunks") + .unwrap() + .unwrap() + .0, + serde_json::json!([2, 3, 4]) + ); + assert_eq!( + sharded_zstd + .get_by_name::("codecs") + .unwrap() + .unwrap() + .0, + serde_json::json!([{ + "name": "sharding_indexed", + "configuration": { + "chunk_shape": [2, 3, 4], + "codecs": [ + { + "name": "bytes", + "configuration": {"endian": "little"} + }, + { + "name": "zstd", + "configuration": {"level": 1, "checksum": true} + } + ], + "index_codecs": [ + { + "name": "bytes", + "configuration": {"endian": "little"} + }, + {"name": "crc32c"} + ], + "index_location": "end" + } + }]) + ); + + let time = c + .select( + "SELECT dimensions, units, calendar FROM zarr_inspect('zarr_v3_e2e_server') WHERE path = 'nested/time'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + time.get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time"] + ); + assert_eq!( + time.get_by_name::("units").unwrap().unwrap(), + "milliseconds since 1970-01-01 00:00:00" + ); + assert_eq!( + time.get_by_name::("calendar").unwrap().unwrap(), + "proleptic_gregorian" + ); + }); + } + + #[pg_test] + fn zarr_inspect_minio_metadata_e2e() { + create_minio_e2e_server(); + + Spi::connect(|c| { + let paths = c + .select( + "SELECT path FROM zarr_inspect('zarr_e2e_server') ORDER BY path", + None, + &[], + ) + .unwrap() + .filter_map(|row| row.get_by_name::("path").unwrap()) + .collect::>(); + assert_eq!( + paths, + vec![ + "/", + "nested", + "nested/band", + "nested/blosc", + "nested/channel", + "nested/forecast_time", + "nested/generic4d", + "nested/lazy1m", + "nested/level", + "nested/raw", + "nested/sample", + "nested/spatial2d", + "nested/spatial_ref", + "nested/time", + "nested/x", + "nested/y", + ] + ); + + let root = c + .select( + "SELECT kind, attributes FROM zarr_inspect('zarr_e2e_server') WHERE path = '/'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + root.get_by_name::("kind").unwrap().unwrap(), + "group" + ); + assert_eq!( + root.get_by_name::("attributes") + .unwrap() + .unwrap() + .0["title"], + "Deterministic Zarr inspection fixture" + ); + + let group = c + .select( + "SELECT crs FROM zarr_inspect('zarr_e2e_server') WHERE path = 'nested'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + group.get_by_name::("crs").unwrap().unwrap().0["properties"]["name"], + "EPSG:3857" + ); + + let raw = c + .select( + r#"SELECT kind, group_path, variable, shape, dimensions, dtype, + chunks, codecs, units, fill_value, scale_factor, + add_offset, crs, attributes, warnings + FROM zarr_inspect('zarr_e2e_server') + WHERE path = 'nested/raw'"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + raw.get_by_name::("kind").unwrap().unwrap(), + "array" + ); + assert_eq!( + raw.get_by_name::("group_path").unwrap().unwrap(), + "nested" + ); + assert_eq!( + raw.get_by_name::("variable").unwrap().unwrap(), + "raw" + ); + assert_eq!( + raw.get_by_name::("shape").unwrap().unwrap().0, + serde_json::json!([2, 5, 6]) + ); + assert_eq!( + raw.get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time", "y", "x"] + ); + assert_eq!( + raw.get_by_name::("dtype").unwrap().unwrap(), + "("chunks").unwrap().unwrap().0, + serde_json::json!([2, 3, 4]) + ); + assert!( + raw.get_by_name::("codecs").unwrap().unwrap().0["compressor"].is_null() + ); + assert_eq!(raw.get_by_name::("units").unwrap().unwrap(), "K"); + assert_eq!( + raw.get_by_name::("fill_value") + .unwrap() + .unwrap() + .0, + serde_json::json!(-7.5) + ); + assert_eq!( + raw.get_by_name::("scale_factor").unwrap().unwrap(), + 0.01 + ); + assert_eq!( + raw.get_by_name::("add_offset").unwrap().unwrap(), + 273.15 + ); + assert_eq!( + raw.get_by_name::("crs").unwrap().unwrap().0, + serde_json::json!("PROJCRS[\"WGS 84 / Pseudo-Mercator\"]") + ); + assert_eq!( + raw.get_by_name::("attributes") + .unwrap() + .unwrap() + .0["grid_mapping"], + serde_json::json!("spatial_ref") + ); + assert_eq!( + raw.get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + Vec::::new() + ); + + let blosc = c + .select( + "SELECT codecs FROM zarr_inspect('zarr_e2e_server') WHERE path = 'nested/blosc'", + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + blosc.get_by_name::("codecs").unwrap().unwrap().0["compressor"]["id"], + "blosc" + ); + + let spatial_ref = c + .select( + r#"SELECT kind, group_path, variable, dimensions, dtype, + crs, attributes, warnings + FROM zarr_inspect('zarr_e2e_server') + WHERE path = 'nested/spatial_ref'"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + spatial_ref + .get_by_name::("kind") + .unwrap() + .unwrap(), + "array" + ); + assert_eq!( + spatial_ref + .get_by_name::("group_path") + .unwrap() + .unwrap(), + "nested" + ); + assert_eq!( + spatial_ref + .get_by_name::("variable") + .unwrap() + .unwrap(), + "spatial_ref" + ); + assert_eq!( + spatial_ref + .get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + Vec::::new() + ); + assert_eq!( + spatial_ref + .get_by_name::("dtype") + .unwrap() + .unwrap(), + "|i1" + ); + let spatial_ref_crs = spatial_ref + .get_by_name::("crs") + .unwrap() + .unwrap() + .0; + assert_eq!( + spatial_ref_crs, + serde_json::json!("PROJCRS[\"WGS 84 / Pseudo-Mercator\"]") + ); + let spatial_ref_attrs = spatial_ref + .get_by_name::("attributes") + .unwrap() + .unwrap() + .0; + assert_eq!( + spatial_ref_attrs["grid_mapping_name"], + serde_json::json!("mercator") + ); + assert_eq!( + spatial_ref_attrs["epsg_code"], + serde_json::json!("EPSG:3857") + ); + assert_eq!( + spatial_ref_attrs["crs_wkt"], + serde_json::json!("PROJCRS[\"WGS 84 / Pseudo-Mercator\"]") + ); + assert_eq!( + spatial_ref_attrs["GeoTransform"], + serde_json::json!("100 10 0 50 0 -10") + ); + assert_eq!( + spatial_ref + .get_by_name::, _>("warnings") + .unwrap() + .unwrap(), + Vec::::new() + ); + + let time = c + .select( + r#"SELECT dimensions, units, calendar + FROM zarr_inspect('zarr_e2e_server') + WHERE path = 'nested/time'"#, + None, + &[], + ) + .unwrap() + .next() + .unwrap(); + assert_eq!( + time.get_by_name::, _>("dimensions") + .unwrap() + .unwrap(), + vec!["time"] + ); + assert_eq!( + time.get_by_name::("units").unwrap().unwrap(), + "milliseconds since 1970-01-01 00:00:00" + ); + assert_eq!( + time.get_by_name::("calendar").unwrap().unwrap(), + "proleptic_gregorian" + ); + + let projected_axes = c + .select( + r#"SELECT path, attributes + FROM zarr_inspect('zarr_e2e_server') + WHERE path IN ('nested/x', 'nested/y') + ORDER BY path"#, + None, + &[], + ) + .unwrap() + .map(|row| { + ( + row.get_by_name::("path").unwrap().unwrap(), + row.get_by_name::("attributes") + .unwrap() + .unwrap() + .0, + ) + }) + .collect::>(); + assert_eq!(projected_axes.len(), 2); + assert_eq!(projected_axes[0].0, "nested/x"); + assert_eq!(projected_axes[0].1["axis"], serde_json::json!("X")); + assert_eq!( + projected_axes[0].1["standard_name"], + serde_json::json!("projection_x_coordinate") + ); + assert_eq!(projected_axes[0].1["units"], serde_json::json!("m")); + assert_eq!(projected_axes[1].0, "nested/y"); + assert_eq!(projected_axes[1].1["axis"], serde_json::json!("Y")); + assert_eq!( + projected_axes[1].1["standard_name"], + serde_json::json!("projection_y_coordinate") + ); + assert_eq!(projected_axes[1].1["units"], serde_json::json!("m")); + }); + } + + #[pg_test(error = "foreign server 'missing_zarr_server' does not exist or is not accessible")] + fn zarr_inspect_rejects_missing_server() { + Spi::run("SELECT * FROM zarr_inspect('missing_zarr_server')").unwrap(); + } + + #[pg_test(error = "foreign server 'zarr_private_server' does not exist or is not accessible")] + fn zarr_inspect_requires_server_usage() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_private_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_private_server + FOREIGN DATA WRAPPER zarr_private_wrapper + OPTIONS ( + store_url 's3://warehouse/zarr/e2e.zarr', + anonymous 'true' + )"#, + None, + &[], + ) + .unwrap(); + c.update("CREATE ROLE zarr_inspect_no_usage", None, &[]) + .unwrap(); + c.update("SET ROLE zarr_inspect_no_usage", None, &[]) + .unwrap(); + c.select( + "SELECT * FROM zarr_inspect('zarr_private_server')", + None, + &[], + ) + .unwrap(); + }); + } + + #[pg_test] + fn zarr_explain_uses_network_free_positive_estimate() { + Spi::connect_mut(|c| { + c.update( + r#"CREATE FOREIGN DATA WRAPPER zarr_plan_wrapper + HANDLER zarr_fdw_handler VALIDATOR zarr_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE SERVER zarr_plan_server + FOREIGN DATA WRAPPER zarr_plan_wrapper + OPTIONS ( + store_url 's3://zarr-test/does-not-exist.zarr', + anonymous 'true' + )"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"CREATE FOREIGN TABLE zarr_plan_table ( + x double precision, + y double precision, + value real + ) + SERVER zarr_plan_server + OPTIONS (array_group 'value')"#, + None, + &[], + ) + .unwrap(); + + let plan = c + .select("EXPLAIN SELECT value FROM zarr_plan_table", None, &[]) + .unwrap() + .filter_map(|row| row.get::<&str>(1).unwrap().map(str::to_string)) + .collect::>(); + assert!( + plan.iter() + .any(|line| line.contains("rows=1000000 width=4")), + "expected a positive network-free Zarr estimate, got {plan:?}" + ); + }); + } +} diff --git a/wrappers/src/fdw/zarr_fdw/zarr_fdw.rs b/wrappers/src/fdw/zarr_fdw/zarr_fdw.rs new file mode 100644 index 000000000..a82e4c98e --- /dev/null +++ b/wrappers/src/fdw/zarr_fdw/zarr_fdw.rs @@ -0,0 +1,4395 @@ +//! Main `zarr_fdw` implementation. +//! +//! Given a query plan's pushed-down quals (WHERE) and target columns, this FDW +//! translates them into a lazy *chunk fetch stream* against S3, decompresses the +//! chunks and streams flat rows back to Postgres. Data model (MVP): +//! +//! - a single normalized named-dimension Zarr v2 or direct v3 array in C order, +//! - one same-group, same-name 1D numeric coordinate array per dimension, +//! - flat row output where every non-dimension target column receives the +//! selected array's scalar value. +//! +//! Pushdown: a qual on any finite monotonic coordinate is converted into an +//! index range over that dimension's coordinate vector, which prunes the chunk +//! list before any data chunk is fetched. A time qual is interpreted via either the +//! `time_unit`/`time_origin` table options or, when `time_from_attrs` is true, +//! the discovered time coordinate's CF `units`/`calendar` attributes. +//! +//! Spatial PostGIS predicates (`ST_Intersects`, `geom && box`) do *not* reach +//! this code as `Qual`s — the framework only extracts simple Var-op-Const +//! expressions — so strict geometry pushdown is deferred to v1 (chunk-extent +//! catalog table); the MVP prunes on the `x`/`y`/`time` columns directly. + +use crate::stats; +use futures_util::FutureExt; +use pgrx::datum::TimestampWithTimeZone; +use pgrx::pg_sys; +use serde_json::{Map, Value as JsonValue}; +use std::collections::HashMap; +use std::future::Future; +use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, +}; +use std::time::{Duration, Instant}; + +use supabase_wrappers::prelude::*; + +use super::aggregate::{ + AggregateReducer, aggregate_signature_supported, qual_matches, qual_shape_supported, +}; +use super::cache::{CachedObject, CompressedChunkCache}; +use super::chunk::{ChunkIndexCursor, IndexBounds, chunk_key}; +use super::codec::{CodecDecode, CodecPipeline}; +use super::dataset::{ + CoordinateSource, Dataset, DimensionRole, named_array_dataset, named_dimensions, + ome_rank2_dataset, +}; +use super::decode::{ + DType, coord_bytes_to_f64, coord_fill_value_to_f64, coordinate_itemsize, fill_value_bytes, +}; +use super::meta::{ + ArrayMeta, ArrayNode, NodeMeta, ZarrFormat, parse_v2_array, parse_v2_group, parse_v3_node, +}; +use super::metrics::{ReadKind, ZarrExplainContext, ZarrScanMetrics}; +use super::ome::{ + ResolvedOmeLevel, canonical_ome_group_path, resolve_ome_05_level, + validate_optional_ome_05_attributes, +}; +use super::prefetch::{ + OrderedPrefetch, PrefetchNext, PrefetchRequest, PrefetchSource, ScheduleError, +}; +use super::scan_plan::{CoordinateRange, ScanPlan, ScanPlanner}; +use super::scientific::{ScientificValueDecoder, time::TimeSpec}; +use super::selection::Selection; +use super::selectors::{BoundDimensionSelectors, DimensionSelectors, OPT_DIMENSION_SELECTORS}; +use super::sharding::{ + CachedShardIndex, MAX_SHARD_INDEX_BYTES, ShardIndex, ShardIndexCache, ShardIndexDecode, + ShardingConfig, StorageLayout, +}; +use super::spatial::crs::{ + GridMappingMetadata, ResolvedCrs, grid_mapping_sibling_path, resolve_crs, +}; +use super::spatial::grid::{ + HorizontalAxes, HorizontalCell, RectilinearGrid, discover_horizontal_axes_from_roles, + exact_center_index, inclusive_center_bounds, nearest_center_index, +}; +use super::store::{ + MAX_METADATA_OBJECT_BYTES, RangedObject, ReadIdentity, ReadRange, ZarrStore, join_key, + validate_store_definition_privilege, validate_store_options, +}; +use super::{ZarrFdwError, ZarrFdwResult}; + +const FDW_NAME: &str = "ZarrFdw"; + +// Table option names. +const OPT_ARRAY_GROUP: &str = "array_group"; +const OPT_MULTISCALE_GROUP: &str = "multiscale_group"; +const OPT_MULTISCALE_INDEX: &str = "multiscale_index"; +const OPT_MULTISCALE_LEVEL: &str = "multiscale_level"; +const OPT_BANDS: &str = "bands"; +const OPT_TIME_UNIT: &str = "time_unit"; +const OPT_TIME_ORIGIN: &str = "time_origin"; +const OPT_TIME_FROM_ATTRS: &str = "time_from_attrs"; +const OPT_DECODE_CF: &str = "decode_cf"; +const OPT_MAX_CONCURRENT_READS: &str = "max_concurrent_reads"; +const OPT_MAX_INFLIGHT_BYTES: &str = "max_inflight_bytes"; +const OPT_COMPRESSED_CACHE_BYTES: &str = "compressed_cache_bytes"; + +const DEFAULT_MAX_CONCURRENT_READS: usize = 4; +const MAX_CONCURRENT_READS: usize = 32; +// One maximum-size decoded chunk, gzip/zlib's bounded framing allowance, and +// the optional v3 CRC32C trailer must all fit under the default request budget. +const DEFAULT_MAX_INFLIGHT_BYTES: usize = 257 * 1024 * 1024 + 4; +const MIN_MAX_INFLIGHT_BYTES: usize = 1024 * 1024; +const MAX_MAX_INFLIGHT_BYTES: usize = 1024 * 1024 * 1024; +const DEFAULT_COMPRESSED_CACHE_BYTES: usize = 64 * 1024 * 1024; +const MAX_COMPRESSED_CACHE_BYTES: usize = 1024 * 1024 * 1024; +const MAX_COMPRESSED_CACHE_ENTRIES: usize = 4096; +const SHARD_INDEX_CACHE_FRACTION: usize = 4; +const INTERRUPT_POLL_INTERVAL: Duration = Duration::from_millis(25); + +// Planning must stay deterministic and network-free. Until metadata-backed or +// configured estimates are available, use a deliberately non-zero cardinality +// for remote arrays so PostgreSQL does not price every scan at startup cost. +const DEFAULT_PLANNER_ROWS: i64 = 1_000_000; +const DEFAULT_EMPTY_PROJECTION_WIDTH: i32 = 8; +const DEFAULT_UNKNOWN_TYPE_WIDTH: i32 = 32; + +// The executor decodes one data chunk and every required coordinate vector in +// a PostgreSQL backend. Chunk coordinates themselves are streamed lazily. +// Keep the remaining remote-metadata-driven allocations bounded. +const MAX_DECODED_CHUNK_BYTES: usize = 256 * 1024 * 1024; +const MAX_COORDINATE_VALUES: usize = 16 * 1024 * 1024; +const MAX_TOTAL_COORDINATE_VALUES: usize = MAX_COORDINATE_VALUES; +const SPATIAL_TIME_INTERRUPT_POLL_VALUES: usize = 1_024; +const SELECTOR_INTERRUPT_POLL_CELLS: usize = 1_024; +const UNSUPPORTED_COORDINATE_DECODING_ATTRIBUTES: [&str; 7] = [ + "_FillValue", + "missing_value", + "valid_range", + "valid_min", + "valid_max", + "scale_factor", + "add_offset", +]; + +#[derive(Clone, Debug, PartialEq, Eq)] +struct MultiscaleSelectionOptions { + group: String, + index: usize, + level: usize, +} + +fn multiscale_selection_options( + options: &HashMap, +) -> ZarrFdwResult> { + let group = options.get(OPT_MULTISCALE_GROUP); + let index = options.get(OPT_MULTISCALE_INDEX); + let level = options.get(OPT_MULTISCALE_LEVEL); + let present = [group.is_some(), index.is_some(), level.is_some()]; + if !present.iter().any(|present| *present) { + return Ok(None); + } + if !present.iter().all(|present| *present) { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_MULTISCALE_GROUP.to_string(), + message: + "multiscale_group, multiscale_index, and multiscale_level must be provided together" + .to_string(), + }); + } + if options.contains_key(OPT_ARRAY_GROUP) { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_MULTISCALE_GROUP.to_string(), + message: "array_group cannot be combined with multiscale selection options".to_string(), + }); + } + + let raw_group = group.expect("all multiscale options were checked"); + let group = + canonical_ome_group_path(raw_group).map_err(|_| ZarrFdwError::InvalidOptionValue { + option: OPT_MULTISCALE_GROUP.to_string(), + message: "must be '/' or a safe relative Zarr group path".to_string(), + })?; + + let parse_index = |option: &'static str, value: &str| { + value + .parse::() + .map_err(|_| ZarrFdwError::InvalidOptionValue { + option: option.to_string(), + message: "must be a zero-based non-negative integer".to_string(), + }) + }; + Ok(Some(MultiscaleSelectionOptions { + group, + index: parse_index( + OPT_MULTISCALE_INDEX, + index.expect("all multiscale options were checked"), + )?, + level: parse_index( + OPT_MULTISCALE_LEVEL, + level.expect("all multiscale options were checked"), + )?, + })) +} + +enum ArrayMetadataDocument { + V2(Vec), + V3(Vec), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ChunkFetchContext { + logical_indices: Vec, + object_key: String, +} + +enum ResolvedChunkRequest { + Fetch(PrefetchRequest), + Synthesized(PrefetchRequest), +} + +enum DeferredChunkRequest { + Logical(Vec), + Resolved(ResolvedChunkRequest), +} + +enum ShardIndexResolution { + Ready(Option>), + WouldBlock, +} + +/// Array-axis positions required by a spatial-time operation. +/// +/// Horizontal bounds exposed by the operation layer remain in semantic +/// `[x, y]` order. `horizontal` maps them back to the selected array's actual +/// dimension order. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct SpatialTimeLayout { + pub(super) time: usize, + pub(super) horizontal: HorizontalAxes, +} + +/// Exact selected time indexes plus the conservative full-rank scan window. +/// +/// `time_indices` contains only coordinates inside the requested half-open +/// interval. `bounds` may span rejected indexes on an unordered time axis, so +/// callers must use `time_indices` for exact row acceptance. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct SpatialTimeSelection { + pub(super) layout: SpatialTimeLayout, + pub(super) time_indices: Vec, + pub(super) bounds: Vec>, + pub(super) candidate_cells: usize, +} + +#[cfg(test)] +fn discover_spatial_time_layout( + rank: usize, + axis_roles: &[DimensionRole], + shape: &[u64], +) -> ZarrFdwResult { + let layout = discover_spatial_time_layout_with_auxiliary_dimensions(rank, axis_roles, shape)?; + for (axis, &extent) in shape.iter().enumerate() { + if axis != layout.time + && axis != layout.horizontal.x + && axis != layout.horizontal.y + && extent != 1 + { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time execution requires auxiliary dimension {axis} to have extent 1, found {extent}" + ))); + } + } + Ok(layout) +} + +fn discover_spatial_time_layout_with_auxiliary_dimensions( + rank: usize, + axis_roles: &[DimensionRole], + shape: &[u64], +) -> ZarrFdwResult { + if rank < 3 || shape.len() != rank || axis_roles.len() != rank { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time execution requires an array of rank 3 or greater, found rank {rank}" + ))); + } + + let horizontal = discover_horizontal_axes_from_roles(axis_roles.iter().copied())?; + let time_axes = axis_roles + .iter() + .enumerate() + .filter_map(|(axis, role)| (*role == DimensionRole::Time).then_some(axis)) + .collect::>(); + if time_axes.len() != 1 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time execution requires exactly one time axis, found {}", + time_axes.len() + ))); + } + let time = time_axes[0]; + if time == horizontal.x || time == horizontal.y { + return Err(ZarrFdwError::InvalidMetadata( + "spatial-time axes must be distinct".to_string(), + )); + } + + Ok(SpatialTimeLayout { time, horizontal }) +} + +fn spatial_time_value_in_range( + time_spec: TimeSpec, + raw: f64, + start_micros: i64, + end_micros: i64, +) -> ZarrFdwResult { + let micros = time_spec.raw_to_pg_micros(raw)?; + Ok(start_micros <= micros && micros < end_micros) +} + +#[wrappers_fdw( + version = "0.0.1", + author = "MPSY", + website = "https://github.com/supabase/wrappers/tree/main/wrappers/src/fdw/zarr_fdw", + error_type = "ZarrFdwError" +)] +pub(crate) struct ZarrFdw { + store: ZarrStore, + + // --- scan state, (re)built in begin_scan ------------------------------ + tgt_cols: Vec, + // object-key path of the cube array, relative to the store prefix + array_dir: String, + // discovered dimension names in array order + axes: Vec, + // selected value array attributes retained for operation-layer metadata + // such as strict CRS resolution + array_attributes: Map, + // explicit OME-Zarr 0.5 selection, retained for EXPLAIN ANALYZE only + selected_ome_level: Option, + // scientific meaning assigned by the metadata adapter, in array order + axis_roles: Vec, + rank: usize, + axis_meta: Option, + dtype: Option, + codec: Option, + scientific_decoder: Option, + // one decoded scalar, repeated when a data chunk is absent + fill_bytes: Option>, + // coordinate values per axis when required by projection or restrictions + coords: Vec>>, + // conservative physical cell window shared by ordinary and spatial scans + selection: Selection, + // Persistent foreign-table selectors and one optional selector-aware + // spatial call are kept as separate AND sources so same-axis constraints + // retain exact residual semantics. + dimension_selectors: BoundDimensionSelectors, + call_dimension_selectors: DimensionSelectors, + bound_call_dimension_selectors: BoundDimensionSelectors, + // lazy chunk indexes to read, in row-major order + chunk_cursor: ChunkIndexCursor, + current_chunk: Vec, + current_object_key: String, + deferred_prefetch: Option, + prefetch: OrderedPrefetch, + compressed_cache: CompressedChunkCache, + shard_index_cache: ShardIndexCache, + payload_cache_bytes: usize, + shard_index_cache_bytes: usize, + cache_layout_sharded: bool, + max_concurrent_reads: usize, + max_inflight_bytes: usize, + compressed_cache_bytes: usize, + metrics: ZarrScanMetrics, + remote_data_get_calls: Arc, + remote_data_encoded_bytes: Arc, + remote_shard_payload_get_calls: Arc, + remote_shard_payload_encoded_bytes: Arc, + flushed_encoded_bytes: u64, + flushed_cells: u64, + flushed_tuples: u64, + + // --- per-chunk iteration state --------------------------------------- + chunk_bytes: Vec, + chunk_shape: Vec, + sub_lo: Vec, + sub_hi: Vec, + sub_idx: Vec, + capture_spatial_indices: bool, + last_emitted_indices: Option>, + pending: bool, + + // --- scalar aggregate execution state ------------------------------- + aggregate_defs: Vec, + aggregate_quals: Vec, + aggregate_reducer: Option, + aggregate_emitted: bool, + + time_spec: TimeSpec, + rows_out: i64, +} + +fn zeroed_scan_cursors(rank: usize) -> [Vec; 3] { + std::array::from_fn(|_| vec![0; rank]) +} + +fn postgres_interrupt_pending() -> bool { + // PostgreSQL's signal handlers update this `volatile sig_atomic_t` + // asynchronously. Preserve the C macro's volatile-read semantics here. + unsafe { std::ptr::read_volatile(&raw const pg_sys::InterruptPending) != 0 } +} + +fn process_postgres_interrupts() { + unsafe { + if std::ptr::read_volatile(&raw const pg_sys::InterruptPending) != 0 { + pg_sys::ProcessInterrupts(); + } + } +} + +fn atomic_saturating_add(counter: &AtomicU64, value: u64) { + let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_add(value)) + }); +} + +async fn observe_data_fetch( + future: F, + remote_get_calls: Arc, + remote_encoded_bytes: Arc, +) -> ZarrFdwResult>> +where + F: Future>>>, +{ + // The async body is lazy: this is an initiated GET, not merely a future + // placed behind an earlier request in FuturesOrdered. + atomic_saturating_add(&remote_get_calls, 1); + let result = future.await; + if let Ok(Some(bytes)) = &result { + atomic_saturating_add( + &remote_encoded_bytes, + u64::try_from(bytes.len()).unwrap_or(u64::MAX), + ); + } + result +} + +async fn observe_shard_payload_fetch( + future: F, + shard_key: String, + remote_get_calls: Arc, + remote_encoded_bytes: Arc, + shard_payload_get_calls: Arc, + shard_payload_encoded_bytes: Arc, +) -> ZarrFdwResult>> +where + F: Future>>, +{ + atomic_saturating_add(&remote_get_calls, 1); + atomic_saturating_add(&shard_payload_get_calls, 1); + let response = future.await?.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "indexed shard object '{shard_key}' disappeared before its payload range was read" + )) + })?; + let bytes = response.bytes; + let len = u64::try_from(bytes.len()).unwrap_or(u64::MAX); + atomic_saturating_add(&remote_encoded_bytes, len); + atomic_saturating_add(&shard_payload_encoded_bytes, len); + Ok(Some(bytes)) +} + +fn checked_chunk_layout( + meta: &ArrayMeta, + itemsize: usize, +) -> ZarrFdwResult<(Vec, usize, usize)> { + let storage_shape = (0..meta.chunks.len()) + .map(|axis| meta.chunk_extent(axis)) + .collect::>>()?; + let storage_cells = meta.chunk_cell_count()?; + let decoded_bytes = storage_cells.checked_mul(itemsize).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "declared chunk byte length exceeds this platform's index capacity".to_string(), + ) + })?; + if decoded_bytes > MAX_DECODED_CHUNK_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "declared chunk decodes to {decoded_bytes} bytes, exceeding the safety limit of {MAX_DECODED_CHUNK_BYTES}" + ))); + } + Ok((storage_shape, storage_cells, decoded_bytes)) +} + +fn checked_flat_offset(indices: &[usize], shape: &[usize]) -> ZarrFdwResult { + if indices.len() != shape.len() { + return Err(ZarrFdwError::InvalidMetadata( + "chunk index rank does not match the declared chunk shape".to_string(), + )); + } + let mut offset = 0usize; + let mut stride = 1usize; + for axis in (0..shape.len()).rev() { + if indices[axis] >= shape[axis] { + return Err(ZarrFdwError::InvalidMetadata(format!( + "within-chunk index {} is outside dimension {axis} with extent {}", + indices[axis], shape[axis] + ))); + } + let contribution = indices[axis].checked_mul(stride).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("chunk cell offset overflow".to_string()) + })?; + offset = offset.checked_add(contribution).ok_or_else(|| { + ZarrFdwError::InvalidMetadata("chunk cell offset overflow".to_string()) + })?; + stride = stride + .checked_mul(shape[axis]) + .ok_or_else(|| ZarrFdwError::InvalidMetadata("chunk stride overflow".to_string()))?; + } + Ok(offset) +} + +fn checked_chunk_byte_range( + cell_offset: usize, + itemsize: usize, + available: usize, +) -> ZarrFdwResult> { + let start = cell_offset + .checked_mul(itemsize) + .ok_or_else(|| ZarrFdwError::InvalidMetadata("chunk byte offset overflow".to_string()))?; + let end = start + .checked_add(itemsize) + .ok_or_else(|| ZarrFdwError::InvalidMetadata("chunk byte range overflow".to_string()))?; + if end > available { + return Err(ZarrFdwError::ReadError(std::io::Error::other(format!( + "chunk cell byte range {start}..{end} exceeds decoded length {available}" + )))); + } + Ok(start..end) +} + +fn require_exact_decoded_len(key: &str, actual: usize, expected: usize) -> ZarrFdwResult<()> { + if actual != expected { + return Err(ZarrFdwError::ReadError(std::io::Error::other(format!( + "chunk '{key}' decoded to {actual} bytes, expected exactly {expected}" + )))); + } + Ok(()) +} + +fn filled_chunk_bytes( + fill_bytes: Option<&[u8]>, + cell_count: usize, + key: &str, +) -> ZarrFdwResult> { + let fill_bytes = fill_bytes.ok_or_else(|| ZarrFdwError::MissingChunkWithoutFillValue { + key: key.to_string(), + })?; + if fill_bytes.is_empty() { + return Err(ZarrFdwError::InvalidMetadata( + "decoded fill value must contain at least one byte".to_string(), + )); + } + let byte_count = fill_bytes.len().checked_mul(cell_count).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "filled chunk byte length exceeds this platform's index capacity".to_string(), + ) + })?; + if byte_count > MAX_DECODED_CHUNK_BYTES { + return Err(ZarrFdwError::InvalidMetadata(format!( + "filled chunk requires {byte_count} bytes, exceeding the safety limit of {MAX_DECODED_CHUNK_BYTES}" + ))); + } + let mut bytes = Vec::new(); + bytes.try_reserve_exact(byte_count).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate a filled chunk of {byte_count} bytes" + )) + })?; + for _ in 0..cell_count { + bytes.extend_from_slice(fill_bytes); + } + Ok(bytes) +} + +fn filled_coordinate_values( + fill_value: Option, + cell_count: usize, + key: &str, + axis: &str, +) -> ZarrFdwResult> { + let fill_value = fill_value.ok_or_else(|| ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: ZarrFdwError::MissingChunkWithoutFillValue { + key: key.to_string(), + } + .to_string(), + })?; + if cell_count > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!( + "coordinate chunk has {cell_count} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}" + ), + }); + } + let mut values = Vec::new(); + values + .try_reserve_exact(cell_count) + .map_err(|_| ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!("could not allocate a coordinate chunk of {cell_count} values"), + })?; + values.resize(cell_count, fill_value); + Ok(values) +} + +fn affine_coordinate_values( + axis: &str, + length: u64, + scale: f64, + translation: f64, +) -> ZarrFdwResult> { + if !scale.is_finite() || scale <= 0.0 || !translation.is_finite() { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: + "OME-Zarr affine coordinates require a finite positive scale and finite translation" + .to_string(), + }); + } + let length = usize::try_from(length).map_err(|_| ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: "coordinate length exceeds this platform's index capacity".to_string(), + })?; + if length > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!( + "coordinate has {length} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}" + ), + }); + } + let mut values = Vec::new(); + values + .try_reserve_exact(length) + .map_err(|_| ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!("could not allocate {length} synthesized coordinate values"), + })?; + for index in 0..length { + if index % SPATIAL_TIME_INTERRUPT_POLL_VALUES == 0 { + process_postgres_interrupts(); + } + let value = scale.mul_add(index as f64, translation); + if !value.is_finite() { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!("synthesized coordinate at index {index} is not finite"), + }); + } + values.push(value); + } + Ok(values) +} + +fn expected_value_pg_type(dtype: DType, decode_cf: bool) -> (pg_sys::Oid, &'static str) { + if decode_cf { + return (pg_sys::FLOAT8OID, "double precision"); + } + match dtype { + DType::F32 => (pg_sys::FLOAT4OID, "real"), + DType::F64 => (pg_sys::FLOAT8OID, "double precision"), + DType::I8 => (pg_sys::CHAROID, r#""char""#), + DType::I16 => (pg_sys::INT2OID, "smallint"), + DType::I32 => (pg_sys::INT4OID, "integer"), + DType::I64 => (pg_sys::INT8OID, "bigint"), + } +} + +fn require_column_type( + column: &Column, + expected_oid: pg_sys::Oid, + expected_name: &'static str, +) -> ZarrFdwResult<()> { + if column.type_oid != expected_oid { + return Err(ZarrFdwError::ColumnTypeMismatch { + column: column.name.clone(), + actual: column.type_oid.to_u32(), + expected: expected_name, + expected_oid: expected_oid.to_u32(), + }); + } + Ok(()) +} + +fn validate_column_types( + columns: &[Column], + dataset: &Dataset, + dtype: DType, + decode_cf: bool, +) -> ZarrFdwResult<()> { + let (value_oid, value_name) = expected_value_pg_type(dtype, decode_cf); + for column in columns { + match dataset.dimension(&column.name) { + Some(dimension) if dimension.semantic_role() == DimensionRole::Time => { + require_column_type(column, pg_sys::TIMESTAMPTZOID, "timestamp with time zone")?; + } + Some(_) => { + require_column_type(column, pg_sys::FLOAT8OID, "double precision")?; + } + None => require_column_type(column, value_oid, value_name)?, + } + } + Ok(()) +} + +fn validate_coordinate_values(axis: &str, values: &[f64]) -> ZarrFdwResult<()> { + if let Some((index, value)) = values + .iter() + .enumerate() + .find(|(_, value)| !value.is_finite()) + { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!("coordinate value at index {index} is not finite ({value})"), + }); + } + + Ok(()) +} + +fn validate_coordinate_decoding_attributes( + axis: &str, + attributes: &Map, +) -> ZarrFdwResult<()> { + if let Some(attribute) = UNSUPPORTED_COORDINATE_DECODING_ATTRIBUTES + .iter() + .copied() + .find(|attribute| attributes.contains_key(*attribute)) + { + return Err(ZarrFdwError::CoordinateReadError { + axis: axis.to_string(), + error: format!( + "attribute '{attribute}' requires coordinate decoding, which is not supported yet" + ), + }); + } + Ok(()) +} + +fn checked_total_coordinate_values<'a>( + dimensions: impl IntoIterator, + limit: usize, +) -> ZarrFdwResult { + let mut total = 0usize; + for (name, length, required) in dimensions { + if !required { + continue; + } + let length = usize::try_from(length).map_err(|_| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "coordinate length exceeds this platform's index capacity".to_string(), + })?; + total = total.checked_add(length).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "total required coordinate value count overflowed".to_string(), + ) + })?; + if total > limit { + return Err(ZarrFdwError::InvalidMetadata(format!( + "required coordinate arrays contain {total} values, exceeding the cumulative safety limit of {limit}" + ))); + } + } + Ok(total) +} + +fn estimated_pg_type_width(type_oid: pg_sys::Oid) -> i32 { + match type_oid { + pg_sys::CHAROID => 1, + pg_sys::INT2OID => 2, + pg_sys::FLOAT4OID | pg_sys::INT4OID => 4, + pg_sys::FLOAT8OID | pg_sys::INT8OID | pg_sys::TIMESTAMPTZOID => 8, + _ => DEFAULT_UNKNOWN_TYPE_WIDTH, + } +} + +fn conservative_rel_size(columns: &[Column]) -> (i64, i32) { + let width = if columns.is_empty() { + DEFAULT_EMPTY_PROJECTION_WIDTH + } else { + columns.iter().fold(0_i32, |sum, column| { + sum.saturating_add(estimated_pg_type_width(column.type_oid)) + }) + }; + (DEFAULT_PLANNER_ROWS, width) +} + +impl ZarrFdw { + fn value_cell(dt: DType, b: &[u8]) -> ZarrFdwResult { + let ok = |n: usize| { + ZarrFdwError::ReadError(std::io::Error::other(format!( + "chunk cell data too short: need {n} bytes, got {}", + b.len() + ))) + }; + Ok(match dt { + DType::F32 => Cell::F32(f32::from_le_bytes(b.try_into().map_err(|_| ok(4))?)), + DType::F64 => Cell::F64(f64::from_le_bytes(b.try_into().map_err(|_| ok(8))?)), + DType::I8 => Cell::I8(b.first().copied().ok_or_else(|| ok(1))? as i8), + DType::I16 => Cell::I16(i16::from_le_bytes(b.try_into().map_err(|_| ok(2))?)), + DType::I32 => Cell::I32(i32::from_le_bytes(b.try_into().map_err(|_| ok(4))?)), + DType::I64 => Cell::I64(i64::from_le_bytes(b.try_into().map_err(|_| ok(8))?)), + }) + } + + /// Install one already-parsed selector document supplied by an explicit + /// spatial overload. It remains a separate AND source from the foreign + /// table option throughout binding, pruning, and exact residual checks. + pub(super) fn set_call_dimension_selectors( + &mut self, + selectors: DimensionSelectors, + ) -> ZarrFdwResult<()> { + if self.axis_meta.is_some() { + return Err(ZarrFdwError::InvalidMetadata( + "spatial call selectors must be installed before scan initialization".to_string(), + )); + } + self.call_dimension_selectors = selectors; + Ok(()) + } + + fn spatial_horizontal_coordinates(&self) -> ZarrFdwResult<(HorizontalAxes, &[f64], &[f64])> { + if self.rank < 2 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial execution requires an array of rank 2 or greater, found rank {}", + self.rank + ))); + } + let axes = discover_horizontal_axes_from_roles(self.axis_roles.iter().copied())?; + let x = self + .coords + .get(axes.x) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{}' was not loaded", + self.axes.get(axes.x).map_or("x", String::as_str) + )) + })?; + let y = self + .coords + .get(axes.y) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{}' was not loaded", + self.axes.get(axes.y).map_or("y", String::as_str) + )) + })?; + Ok((axes, x, y)) + } + + pub(super) fn spatial_horizontal_axes(&self) -> ZarrFdwResult { + self.spatial_horizontal_coordinates() + .map(|(axes, _, _)| axes) + } + + /// Resolve one exact horizontal center on a rank-2-or-greater array. + pub(super) fn spatial_exact_horizontal_cell( + &self, + target_x: f64, + target_y: f64, + ) -> ZarrFdwResult> { + let (axes, x, y) = self.spatial_horizontal_coordinates()?; + let Some(x_index) = exact_center_index(x, target_x)? else { + return Ok(None); + }; + let Some(y_index) = exact_center_index(y, target_y)? else { + return Ok(None); + }; + Ok(Some(( + axes, + HorizontalCell { + x_index, + y_index, + x: x[x_index], + y: y[y_index], + distance: 0.0, + }, + ))) + } + + /// Resolve one nearest horizontal center on a rank-2-or-greater array. + pub(super) fn spatial_nearest_horizontal_cell( + &self, + target_x: f64, + target_y: f64, + ) -> ZarrFdwResult<(HorizontalAxes, HorizontalCell)> { + let (axes, x, y) = self.spatial_horizontal_coordinates()?; + let x_index = nearest_center_index(x, target_x)?.ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial x coordinate is empty".to_string()) + })?; + let y_index = nearest_center_index(y, target_y)?.ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial y coordinate is empty".to_string()) + })?; + let selected_x = x[x_index]; + let selected_y = y[y_index]; + Ok(( + axes, + HorizontalCell { + x_index, + y_index, + x: selected_x, + y: selected_y, + distance: (selected_x - target_x).hypot(selected_y - target_y), + }, + )) + } + + /// Convert a transformed geometry envelope to inclusive semantic `[x, y]` + /// bounds for a rank-2-or-greater array. + pub(super) fn spatial_horizontal_window( + &self, + xmin: f64, + ymin: f64, + xmax: f64, + ymax: f64, + ) -> ZarrFdwResult> { + let (axes, x, y) = self.spatial_horizontal_coordinates()?; + let Some(x_bounds) = inclusive_center_bounds(x, xmin, xmax)? else { + return Ok(None); + }; + let Some(y_bounds) = inclusive_center_bounds(y, ymin, ymax)? else { + return Ok(None); + }; + let candidate_cells = [x_bounds, y_bounds] + .iter() + .try_fold(1usize, |total, bounds| { + let extent = bounds + .end + .checked_sub(bounds.start) + .and_then(|extent| extent.checked_add(1)) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial horizontal candidate count overflowed".to_string(), + ) + })?; + total.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial horizontal candidate count overflowed".to_string(), + ) + }) + })?; + Ok(Some((axes, [x_bounds, y_bounds], candidate_cells))) + } + + /// Borrow the fully decoded horizontal coordinate grid prepared by + /// `begin_scan`. Spatial operations use this view to choose cells without + /// creating a second storage or metadata path. + pub(super) fn rectilinear_grid(&self) -> ZarrFdwResult> { + if self.rank != 2 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "point sampling requires a rank-2 array, found rank {}", + self.rank + ))); + } + let axes = discover_horizontal_axes_from_roles(self.axis_roles.iter().copied())?; + let x = self + .coords + .get(axes.x) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{}' was not loaded", + self.axes.get(axes.x).map_or("x", String::as_str) + )) + })?; + let y = self + .coords + .get(axes.y) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{}' was not loaded", + self.axes.get(axes.y).map_or("y", String::as_str) + )) + })?; + RectilinearGrid::new(axes, x, y) + } + + /// Validate the dimension contract shared by spatial-time operations. + /// + /// Arrays may contain singleton auxiliary dimensions, but execution needs + /// exactly one time axis and one unambiguous horizontal pair. + pub(super) fn spatial_time_layout(&self) -> ZarrFdwResult { + let meta = self.axis_meta.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial-time execution requires an initialized array scan".to_string(), + ) + })?; + let layout = discover_spatial_time_layout_with_auxiliary_dimensions( + self.rank, + &self.axis_roles, + &meta.shape, + )?; + + for axis in [layout.time, layout.horizontal.x, layout.horizontal.y] { + if self.coords.get(axis).and_then(Option::as_deref).is_none() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time coordinate '{}' was not loaded", + self.axes.get(axis).map_or("unknown", String::as_str) + ))); + } + } + + Ok(layout) + } + + /// Convert a transformed geometry envelope to inclusive horizontal index + /// bounds in semantic `[x, y]` order. The complete spatial-time layout is + /// validated before an empty window is returned. + pub(super) fn spatial_time_horizontal_window( + &self, + xmin: f64, + ymin: f64, + xmax: f64, + ymax: f64, + ) -> ZarrFdwResult> { + let layout = self.spatial_time_layout()?; + let x = self.coords[layout.horizontal.x] + .as_deref() + .expect("spatial_time_layout validated the x coordinate"); + let y = self.coords[layout.horizontal.y] + .as_deref() + .expect("spatial_time_layout validated the y coordinate"); + let Some(x_bounds) = inclusive_center_bounds(x, xmin, xmax)? else { + return Ok(None); + }; + let Some(y_bounds) = inclusive_center_bounds(y, ymin, ymax)? else { + return Ok(None); + }; + let candidate_cells = x_bounds + .end + .checked_sub(x_bounds.start) + .and_then(|extent| extent.checked_add(1)) + .and_then(|x_extent| { + y_bounds + .end + .checked_sub(y_bounds.start) + .and_then(|extent| extent.checked_add(1)) + .and_then(|y_extent| x_extent.checked_mul(y_extent)) + }) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial-time horizontal candidate count overflowed".to_string(), + ) + })?; + Ok(Some(( + layout.horizontal, + [x_bounds, y_bounds], + candidate_cells, + ))) + } + + /// Build a conservative full-rank scan window from exact time indexes + /// previously returned by `spatial_time_indices`. Unordered time + /// coordinates are accepted; exact row filtering must still use + /// `time_indices` after the conservative scan. + pub(super) fn spatial_time_selection( + &self, + time_indices: Vec, + horizontal_bounds: [IndexBounds; 2], + max_candidates: usize, + ) -> ZarrFdwResult { + let layout = self.spatial_time_layout()?; + + if time_indices.is_empty() { + return Ok(SpatialTimeSelection { + layout, + time_indices, + bounds: vec![None; self.rank], + candidate_cells: 0, + }); + } + + let meta = self + .axis_meta + .as_ref() + .expect("spatial_time_layout validated metadata"); + let time_bounds = IndexBounds { + start: *time_indices + .first() + .expect("non-empty time selection has a first index"), + end: *time_indices + .last() + .expect("non-empty time selection has a last index"), + }; + let mut bounds = self.selection.axis_bounds().to_vec(); + bounds[layout.time] = Some(time_bounds); + bounds[layout.horizontal.x] = Some(horizontal_bounds[0]); + bounds[layout.horizontal.y] = Some(horizontal_bounds[1]); + + let mut candidate_cells = 1usize; + for (axis, axis_bounds) in bounds.iter().enumerate() { + let length = meta.shape_extent(axis)?; + let extent = match axis_bounds { + Some(axis_bounds) + if axis_bounds.start <= axis_bounds.end && axis_bounds.end < length => + { + axis_bounds + .end + .checked_sub(axis_bounds.start) + .and_then(|extent| extent.checked_add(1)) + .expect("validated inclusive bounds have a positive extent") + } + Some(axis_bounds) => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time index bounds {}..={} are invalid for dimension {axis} length {length}", + axis_bounds.start, axis_bounds.end + ))); + } + None => length, + }; + candidate_cells = candidate_cells.checked_mul(extent).ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial-time candidate cell count overflowed".to_string(), + ) + })?; + if candidate_cells > max_candidates { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time request has {candidate_cells} candidate cells, exceeding the limit of {max_candidates}" + ))); + } + } + + Ok(SpatialTimeSelection { + layout, + time_indices, + bounds, + candidate_cells, + }) + } + + /// Resolve exact native time-axis indexes inside the requested half-open + /// interval independently of any horizontal overlap. + pub(super) fn spatial_time_indices( + &mut self, + start: TimestampWithTimeZone, + end: TimestampWithTimeZone, + max_time_slices: usize, + ) -> ZarrFdwResult> { + let layout = self.spatial_time_layout()?; + let start_micros = start.into_inner(); + let end_micros = end.into_inner(); + if start_micros >= end_micros { + return Err(ZarrFdwError::InvalidMetadata( + "spatial-time start must be earlier than end".to_string(), + )); + } + + let time_value_count = self.coords[layout.time] + .as_deref() + .expect("spatial_time_layout validated the time coordinate") + .len(); + let mut time_indices = Vec::new(); + for index in 0..time_value_count { + if index.is_multiple_of(SPATIAL_TIME_INTERRUPT_POLL_VALUES) { + self.spatial_check_for_interrupt()?; + } + let raw = self.coords[layout.time] + .as_deref() + .expect("spatial_time_layout validated the time coordinate")[index]; + if spatial_time_value_in_range(self.time_spec, raw, start_micros, end_micros)? { + if time_indices.len() >= max_time_slices { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time request exceeds the limit of {max_time_slices} time slices" + ))); + } + time_indices.push(index); + } + } + Ok(time_indices) + } + + /// Strictly resolve the selected array's operation-time CRS. The + /// inspection surface remains permissive; spatial execution requires an + /// explicit, conflict-free EPSG identifier. + pub(super) fn resolved_spatial_crs(&mut self) -> ZarrFdwResult { + if self.axis_meta.is_none() { + return Err(ZarrFdwError::InvalidMetadata( + "spatial CRS resolution requires an initialized array scan".to_string(), + )); + } + let array_path = self.array_dir.clone(); + let array_attributes = self.array_attributes.clone(); + let group_path = array_parent_path(&array_path).to_string(); + let group_attributes = + read_array_attributes_optional(&self.store, &mut self.metrics, &group_path)?; + let mapping_path = grid_mapping_sibling_path(&array_path, &array_attributes)?; + let mapping_attributes = match mapping_path.as_deref() { + Some(path) => read_array_attributes_optional(&self.store, &mut self.metrics, path)?, + None => None, + }; + resolve_crs( + &array_path, + &array_attributes, + group_attributes.as_ref(), + mapping_path + .as_deref() + .zip(mapping_attributes.as_ref()) + .map(|(path, attributes)| GridMappingMetadata { path, attributes }), + ) + } + + /// Return the one non-dimension column selected by the foreign table. + pub(super) fn spatial_value_column(&self) -> ZarrFdwResult<&str> { + let mut values = self + .tgt_cols + .iter() + .filter(|column| !self.axes.iter().any(|axis| axis == &column.name)); + let value = values.next().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial operations require exactly one value column".to_string(), + ) + })?; + if values.next().is_some() { + return Err(ZarrFdwError::InvalidMetadata( + "spatial operations require exactly one value column".to_string(), + )); + } + Ok(&value.name) + } + + pub(super) fn spatial_array_path(&self) -> &str { + &self.array_dir + } + + /// Global array indexes for the row most recently returned by + /// `iter_scan`, in the selected array's native dimension order. + pub(super) fn spatial_last_emitted_global_indices(&self) -> ZarrFdwResult<&[usize]> { + self.last_emitted_indices.as_deref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "spatial row indexes are unavailable before a row is emitted".to_string(), + ) + }) + } + + /// Resolve one loaded coordinate by native array-axis and global index. + pub(super) fn spatial_coordinate_at_index( + &self, + axis: usize, + index: usize, + ) -> ZarrFdwResult { + let axis_name = self.axes.get(axis).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate axis {axis} is outside array rank {}", + self.rank + )) + })?; + let values = self + .coords + .get(axis) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate '{axis_name}' was not loaded" + )) + })?; + values.get(index).copied().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial coordinate index {index} is outside axis {axis} length {}", + values.len() + )) + }) + } + + /// Convert one coordinate on an already-resolved time axis to PostgreSQL's + /// timestamptz representation without rediscovering the complete layout + /// for every emitted cell. + pub(super) fn spatial_time_at_index( + &self, + time_axis: usize, + index: usize, + ) -> ZarrFdwResult { + if self.axis_roles.get(time_axis) != Some(&DimensionRole::Time) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial-time axis {time_axis} is not the discovered time axis" + ))); + } + let raw = self.spatial_coordinate_at_index(time_axis, index)?; + let micros = self.time_spec.raw_to_pg_micros(raw)?; + TimestampWithTimeZone::try_from(micros).map_err(|_| ZarrFdwError::TimeOutOfRange(raw)) + } + + /// Resolve the most recently emitted row to one horizontal cell in constant + /// time. The coordinate vectors and horizontal axes were already fully + /// validated when the spatial window was prepared, so polygon execution + /// must not rebuild and revalidate the complete grid for every row. + pub(super) fn spatial_last_emitted_cell( + &self, + axes: HorizontalAxes, + ) -> ZarrFdwResult { + let array_indices = self.spatial_last_emitted_global_indices()?; + let x_index = *array_indices.get(axes.x).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial x axis {} is outside rank-{} array indexes", + axes.x, self.rank + )) + })?; + let y_index = *array_indices.get(axes.y).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial y axis {} is outside rank-{} array indexes", + axes.y, self.rank + )) + })?; + let x_values = self + .coords + .get(axes.x) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial x coordinate was not loaded".to_string()) + })?; + let y_values = self + .coords + .get(axes.y) + .and_then(Option::as_deref) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial y coordinate was not loaded".to_string()) + })?; + let x = *x_values.get(x_index).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial x index {x_index} is outside coordinate length {}", + x_values.len() + )) + })?; + let y = *y_values.get(y_index).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "spatial y index {y_index} is outside coordinate length {}", + y_values.len() + )) + })?; + Ok(HorizontalCell { + x_index, + y_index, + x, + y, + distance: 0.0, + }) + } + + /// Poll PostgreSQL cancellation while a spatial SRF is consuming many + /// rows from one decoded chunk. This preserves the prefetch cleanup + /// invariant enforced by the ordinary scan path. + pub(super) fn spatial_check_for_interrupt(&mut self) -> ZarrFdwResult<()> { + self.process_pending_interrupt() + } + + /// Resolve every operation-auxiliary dimension to zero or one exact native + /// index under the intersection of table and call selectors. + /// + /// Operation-owned axes cannot be named by either selector source. All + /// auxiliary axes are checked even after one resolves empty so an ambiguous + /// axis cannot be hidden by an unrelated no-match selector. + pub(super) fn apply_spatial_dimension_selectors( + &mut self, + operation: &str, + operation_axes: &[usize], + ) -> ZarrFdwResult { + let meta = self.axis_meta.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial scan was not initialized".to_string()) + })?; + if self.axes.len() != self.rank || meta.shape.len() != self.rank { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial dimension metadata does not match array rank {}", + self.rank + ))); + } + let axis_lengths = (0..self.rank) + .map(|axis| meta.shape_extent(axis)) + .collect::>>()?; + + let mut owned = vec![false; self.rank]; + for &axis in operation_axes { + let Some(slot) = owned.get_mut(axis) else { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial operation axis {axis} is outside array rank {}", + self.rank + ))); + }; + *slot = true; + if self.dimension_selectors.selects_axis(axis) + || self.bound_call_dimension_selectors.selects_axis(axis) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_DIMENSION_SELECTORS.to_string(), + message: format!( + "{operation} owns dimension '{}'; selectors may target auxiliary dimensions only", + self.axes[axis] + ), + }); + } + } + + let mut auxiliary_bounds = vec![None; self.rank]; + let mut any_empty = false; + for axis in 0..self.rank { + if owned[axis] { + continue; + } + let length = axis_lengths[axis]; + let mut selected = None; + for index in 0..length { + if index.is_multiple_of(SELECTOR_INTERRUPT_POLL_CELLS) { + self.process_pending_interrupt()?; + } + if !self + .dimension_selectors + .matches_axis_index(axis, index, &self.coords)? + || !self.bound_call_dimension_selectors.matches_axis_index( + axis, + index, + &self.coords, + )? + { + continue; + } + if selected.is_some() { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_DIMENSION_SELECTORS.to_string(), + message: format!( + "auxiliary dimension '{}' resolves to more than one index; spatial operations require zero or one", + self.axes[axis] + ), + }); + } + selected = Some(index); + } + match selected { + Some(index) => { + auxiliary_bounds[axis] = Some(IndexBounds { + start: index, + end: index, + }); + } + None => any_empty = true, + } + } + + let selection = if any_empty { + Selection::empty(self.rank) + } else { + self.selection + .clone() + .intersect(Selection::from_axis_bounds(auxiliary_bounds)) + }; + let nonempty = !selection.is_empty(); + self.apply_selection(selection, true)?; + Ok(nonempty) + } + + /// Narrow an already-prepared rank-2 scan to one global array cell. This + /// preserves the existing chunk loader, missing-chunk semantics, cache, + /// cancellation, metrics, and scientific decoder. + pub(super) fn restrict_to_spatial_cell( + &mut self, + array_indices: [usize; 2], + ) -> ZarrFdwResult<()> { + let array_bounds = array_indices.map(|index| IndexBounds { + start: index, + end: index, + }); + self.restrict_to_spatial_bounds(array_bounds) + } + + /// Narrow an already-prepared rank-2 scan to inclusive global array-index + /// bounds. Spatial polygon operations derive these bounds from the + /// transformed geometry envelope, then apply an exact PostGIS mask to the + /// candidate cell centers. + pub(super) fn restrict_to_spatial_bounds( + &mut self, + array_bounds: [IndexBounds; 2], + ) -> ZarrFdwResult<()> { + let meta = self.axis_meta.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial scan was not initialized".to_string()) + })?; + if self.rank != 2 || meta.shape.len() != 2 { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial execution requires a rank-2 array, found rank {}", + self.rank + ))); + } + self.restrict_to_axis_bounds(array_bounds.into_iter().map(Some).collect()) + } + + /// Narrow a rank-2-or-greater scan to one semantic horizontal cell while + /// preserving the exact singleton bounds already chosen for auxiliaries. + pub(super) fn restrict_to_horizontal_cell( + &mut self, + axes: HorizontalAxes, + x_index: usize, + y_index: usize, + ) -> ZarrFdwResult<()> { + self.restrict_to_horizontal_bounds( + axes, + [ + IndexBounds { + start: x_index, + end: x_index, + }, + IndexBounds { + start: y_index, + end: y_index, + }, + ], + ) + } + + /// Narrow a rank-2-or-greater scan with semantic `[x, y]` bounds. + pub(super) fn restrict_to_horizontal_bounds( + &mut self, + axes: HorizontalAxes, + horizontal_bounds: [IndexBounds; 2], + ) -> ZarrFdwResult<()> { + if axes.x >= self.rank || axes.y >= self.rank || axes.x == axes.y { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial horizontal axes ({}, {}) are invalid for array rank {}", + axes.x, axes.y, self.rank + ))); + } + let mut axis_bounds = vec![None; self.rank]; + axis_bounds[axes.x] = Some(horizontal_bounds[0]); + axis_bounds[axes.y] = Some(horizontal_bounds[1]); + self.restrict_to_axis_bounds(axis_bounds) + } + + /// Narrow an initialized scan to optional inclusive bounds for every + /// native array axis. `None` preserves the complete extent of that axis. + pub(super) fn restrict_to_axis_bounds( + &mut self, + axis_bounds: Vec>, + ) -> ZarrFdwResult<()> { + let meta = self.axis_meta.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata("spatial scan was not initialized".to_string()) + })?; + if axis_bounds.len() != self.rank || meta.shape.len() != self.rank { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial bounds rank {} does not match array rank {}", + axis_bounds.len(), + self.rank + ))); + } + for (axis, bounds) in axis_bounds.iter().enumerate() { + let Some(bounds) = bounds else { + continue; + }; + let length = meta.shape_extent(axis)?; + if bounds.start > bounds.end || bounds.end >= length { + return Err(ZarrFdwError::InvalidMetadata(format!( + "spatial index bounds {}..={} are invalid for dimension {axis} length {length}", + bounds.start, bounds.end + ))); + } + } + + let selection = self + .selection + .clone() + .intersect(Selection::from_axis_bounds(axis_bounds)); + self.apply_selection(selection, true) + } + + /// Install a conservative candidate window and reset every cursor/buffer + /// derived from the previous window. Exact SQL, temporal, and PostGIS + /// residual checks remain with their existing owners. + fn apply_selection( + &mut self, + selection: Selection, + capture_emitted_indices: bool, + ) -> ZarrFdwResult<()> { + let meta = self.axis_meta.as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "scan selection requires initialized metadata".to_string(), + ) + })?; + let plan = ScanPlanner::new(meta).plan(selection)?; + self.apply_scan_plan(plan, capture_emitted_indices) + } + + /// Install an already validated executor-time plan. The plan contains only + /// rank-sized axis ranges; `ChunkIndexCursor` remains the sole lazy chunk + /// enumerator. + fn apply_scan_plan( + &mut self, + plan: ScanPlan, + capture_emitted_indices: bool, + ) -> ZarrFdwResult<()> { + let chunk_cursor = ChunkIndexCursor::new(plan.axis_chunk_ranges())?; + self.metrics + .set_chunk_selection(plan.chunks_total(), plan.chunks_selected()); + self.selection = plan.into_selection(); + self.chunk_cursor = chunk_cursor; + self.current_chunk.clear(); + self.prefetch.clear(); + self.deferred_prefetch = None; + self.chunk_bytes.clear(); + self.chunk_shape.clear(); + [self.sub_lo, self.sub_hi, self.sub_idx] = zeroed_scan_cursors(self.rank); + self.capture_spatial_indices = capture_emitted_indices; + self.last_emitted_indices = None; + self.pending = false; + self.rows_out = 0; + Ok(()) + } + + fn value_cell_at_cursor(&self) -> ZarrFdwResult> { + let dt = self.dtype.expect("dtype set in begin_scan"); + let offset = checked_flat_offset(&self.sub_idx, &self.chunk_shape)?; + let byte_range = checked_chunk_byte_range(offset, dt.itemsize(), self.chunk_bytes.len())?; + let raw_value = &self.chunk_bytes[byte_range]; + match &self.scientific_decoder { + Some(decoder) => Ok(decoder.decode(raw_value)?.map(Cell::F64)), + None => Ok(Some(Self::value_cell(dt, raw_value)?)), + } + } + + fn coordinate_cell_at_cursor(&self, axis: usize) -> ZarrFdwResult { + let meta = self + .axis_meta + .as_ref() + .expect("begin_scan must be called before iter_scan"); + let chunk_indices = &self.current_chunk; + let chunk_index = usize::try_from(chunk_indices[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk index for axis {axis} exceeds this platform's index capacity" + )) + })?; + let chunk_len = meta.chunk_extent(axis)?; + let global = chunk_index + .checked_mul(chunk_len) + .and_then(|base| base.checked_add(self.sub_idx[axis])) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("coordinate index overflow on axis {axis}")) + })?; + let coords = self.coords[axis].as_deref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate '{}' is required for row output or predicate evaluation but was not loaded", + self.axes[axis] + )) + })?; + let coord = *coords.get(global).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "coordinate index {global} is outside axis {axis} length {}", + coords.len() + )) + })?; + if self.axis_roles[axis] == DimensionRole::Time { + let micros = self.time_spec.raw_to_pg_micros(coord)?; + let timestamp = TimestampWithTimeZone::try_from(micros) + .map_err(|_| ZarrFdwError::TimeOutOfRange(coord))?; + Ok(Cell::Timestamptz(timestamp)) + } else { + Ok(Cell::F64(coord)) + } + } + + fn global_index_at_cursor(&self, axis: usize) -> ZarrFdwResult { + let meta = self + .axis_meta + .as_ref() + .expect("begin_scan must be called before iter_scan"); + if self.current_chunk.len() != self.rank || self.sub_idx.len() != self.rank { + return Err(ZarrFdwError::InvalidMetadata( + "scan cursor rank does not match array rank".to_string(), + )); + } + let chunk_index = usize::try_from(self.current_chunk[axis]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk index for axis {axis} exceeds this platform's index capacity" + )) + })?; + let chunk_len = meta.chunk_extent(axis)?; + chunk_index + .checked_mul(chunk_len) + .and_then(|base| base.checked_add(self.sub_idx[axis])) + .ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "global coordinate index overflow on axis {axis}" + )) + }) + } + + fn column_cell_at_cursor( + &self, + column_name: &str, + value_cell: Option<&Cell>, + ) -> ZarrFdwResult> { + match self.axes.iter().position(|axis| axis == column_name) { + Some(axis) => self.coordinate_cell_at_cursor(axis).map(Some), + None => Ok(value_cell.cloned()), + } + } + + fn advance_cursor(&mut self) { + // Advance in C order (last axis varies fastest). + for axis in (0..self.rank).rev() { + if self.sub_idx[axis] < self.sub_hi[axis] { + self.sub_idx[axis] += 1; + return; + } + self.sub_idx[axis] = self.sub_lo[axis]; + } + self.pending = false; + } + + fn configure_sharded_cache_budget(&mut self) { + if self.cache_layout_sharded { + return; + } + let index_bytes = + (self.compressed_cache_bytes / SHARD_INDEX_CACHE_FRACTION).min(MAX_SHARD_INDEX_BYTES); + let payload_bytes = self.compressed_cache_bytes.saturating_sub(index_bytes); + let index_entries = if index_bytes == 0 { + 0 + } else { + MAX_COMPRESSED_CACHE_ENTRIES / SHARD_INDEX_CACHE_FRACTION + }; + let payload_entries = MAX_COMPRESSED_CACHE_ENTRIES.saturating_sub(index_entries); + self.compressed_cache = CompressedChunkCache::new(payload_bytes, payload_entries); + self.shard_index_cache = ShardIndexCache::new(index_bytes, index_entries); + self.payload_cache_bytes = payload_bytes; + self.shard_index_cache_bytes = index_bytes; + self.cache_layout_sharded = true; + } + + fn resolve_shard_index( + &mut self, + config: &ShardingConfig, + shard_key: &str, + read_kind: ReadKind, + allow_remote: bool, + ) -> ZarrFdwResult { + let request = config.index_read_identity(shard_key.to_string())?; + if let Some(cached) = self.shard_index_cache.get(&request) { + self.metrics.record_shard_index_cache_lookup(true); + return Ok(ShardIndexResolution::Ready(match cached { + CachedShardIndex::Present(index) => Some(index), + CachedShardIndex::Missing => None, + })); + } + if !allow_remote { + return Ok(ShardIndexResolution::WouldBlock); + } + if config.encoded_index_bytes > self.max_inflight_bytes { + return Err(ZarrFdwError::InvalidMetadata(format!( + "shard index read limit {} exceeds max_inflight_bytes {}", + config.encoded_index_bytes, self.max_inflight_bytes + ))); + } + + self.metrics.record_shard_index_cache_lookup(false); + self.metrics.record_remote_request(read_kind); + let response = self.store.get_object_range_sync(request.clone())?; + let response_bytes = response.as_ref().map(|response| response.bytes.len()); + if let Some(bytes) = response_bytes { + self.metrics.record_remote_response_bytes(read_kind, bytes); + } + self.metrics.record_shard_index_get(response_bytes); + + let evictions_before = self.shard_index_cache.evictions(); + let Some(response) = response else { + self.shard_index_cache.insert_missing(request); + self.metrics.record_shard_index_cache_evictions( + self.shard_index_cache + .evictions() + .saturating_sub(evictions_before), + ); + return Ok(ShardIndexResolution::Ready(None)); + }; + let index = + match ShardIndex::decode_interruptible(config, response, postgres_interrupt_pending)? { + ShardIndexDecode::Decoded(index) => Arc::new(index), + ShardIndexDecode::Interrupted => { + self.process_pending_interrupt()?; + return Err(ZarrFdwError::InvalidMetadata( + "query interruption was requested".to_string(), + )); + } + }; + self.shard_index_cache + .insert_present(request, Arc::clone(&index)); + self.metrics.record_shard_index_cache_evictions( + self.shard_index_cache + .evictions() + .saturating_sub(evictions_before), + ); + Ok(ShardIndexResolution::Ready(Some(index))) + } + + fn chunk_request( + &mut self, + indices: Vec, + allow_remote_index: bool, + ) -> ZarrFdwResult> { + let meta = self + .axis_meta + .as_ref() + .expect("begin_scan must be called before iter_scan") + .clone(); + let dtype = self.dtype.expect("dtype set in begin_scan"); + let codec = self.codec.as_ref().expect("codec set in begin_scan"); + let (_, _, expected) = checked_chunk_layout(&meta, dtype.itemsize())?; + let encoded_limit = codec.encoded_read_limit(expected)?; + let context = |object_key: String| ChunkFetchContext { + logical_indices: indices.clone(), + object_key, + }; + + match &meta.storage_layout { + StorageLayout::Direct => { + if encoded_limit > self.max_inflight_bytes { + return Err(ZarrFdwError::InvalidMetadata(format!( + "encoded chunk read limit {encoded_limit} exceeds max_inflight_bytes {}", + self.max_inflight_bytes + ))); + } + let key = join_key( + &self.array_dir, + &chunk_key(&meta.chunk_key_encoding, &indices), + ); + Ok(Some(ResolvedChunkRequest::Fetch(PrefetchRequest { + context: context(key.clone()), + identity: ReadIdentity::whole(key), + max_bytes: encoded_limit, + }))) + } + StorageLayout::Sharded(config) => { + let address = config.chunk_address(&indices)?; + let shard_key = join_key( + &self.array_dir, + &chunk_key(&meta.chunk_key_encoding, &address.shard_indices), + ); + let index = match self.resolve_shard_index( + config, + &shard_key, + ReadKind::Data, + allow_remote_index, + )? { + ShardIndexResolution::Ready(index) => index, + ShardIndexResolution::WouldBlock => return Ok(None), + }; + let request_context = context(shard_key.clone()); + let Some(index) = index else { + return Ok(Some(ResolvedChunkRequest::Synthesized(PrefetchRequest { + context: request_context, + identity: config.index_read_identity(shard_key)?, + max_bytes: 0, + }))); + }; + let entry = index.entry(&address.inner_indices)?; + let Some(identity) = index.payload_read_identity(entry)? else { + return Ok(Some(ResolvedChunkRequest::Synthesized(PrefetchRequest { + context: request_context, + identity: index.index_identity().clone(), + max_bytes: 0, + }))); + }; + let max_bytes = match &identity.range { + ReadRange::Exact { length, .. } => usize::try_from(*length).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "inner chunk range in shard '{shard_key}' exceeds this platform's index capacity" + )) + })?, + ReadRange::Whole | ReadRange::Suffix { .. } => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inner chunk in shard '{shard_key}' did not resolve to an exact byte range" + ))); + } + }; + if max_bytes > encoded_limit { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inner chunk range in shard '{shard_key}' is {max_bytes} bytes, exceeding its encoded read limit of {encoded_limit}" + ))); + } + if max_bytes > self.max_inflight_bytes { + return Err(ZarrFdwError::InvalidMetadata(format!( + "inner chunk range in shard '{shard_key}' is {max_bytes} bytes, exceeding max_inflight_bytes {}", + self.max_inflight_bytes + ))); + } + Ok(Some(ResolvedChunkRequest::Fetch(PrefetchRequest { + context: request_context, + identity, + max_bytes, + }))) + } + } + } + + fn schedule_chunk_request(&mut self, resolved: ResolvedChunkRequest) -> ZarrFdwResult { + let result = match resolved { + ResolvedChunkRequest::Synthesized(request) => self + .prefetch + .try_schedule_synthesized(request, CachedObject::Missing) + .map(|source| (source, false)) + .map_err(|error| (error, true)), + ResolvedChunkRequest::Fetch(request) => { + let store = &self.store; + let remote_get_calls = Arc::clone(&self.remote_data_get_calls); + let remote_encoded_bytes = Arc::clone(&self.remote_data_encoded_bytes); + let shard_payload_get_calls = Arc::clone(&self.remote_shard_payload_get_calls); + let shard_payload_encoded_bytes = + Arc::clone(&self.remote_shard_payload_encoded_bytes); + self.prefetch + .try_schedule( + request, + &mut self.compressed_cache, + move |identity, max_bytes| match identity.range.clone() { + ReadRange::Whole => { + let fetch = + store.get_object_optional_owned(identity.key, max_bytes); + observe_data_fetch(fetch, remote_get_calls, remote_encoded_bytes) + .boxed_local() + } + ReadRange::Exact { .. } | ReadRange::Suffix { .. } => { + let shard_key = identity.key.clone(); + let fetch = store.get_object_range_owned(identity); + observe_shard_payload_fetch( + fetch, + shard_key, + remote_get_calls, + remote_encoded_bytes, + shard_payload_get_calls, + shard_payload_encoded_bytes, + ) + .boxed_local() + } + }, + ) + .map(|source| (source, true)) + .map_err(|error| (error, false)) + } + }; + + match result { + Ok((source, has_payload_cache_lookup)) => { + self.metrics.record_chunk_request(); + if has_payload_cache_lookup { + self.metrics + .record_cache_lookup(source == PrefetchSource::Cache); + } + Ok(true) + } + Err((ScheduleError::WindowFull(request), synthesized)) => { + let request = if synthesized { + ResolvedChunkRequest::Synthesized(request) + } else { + ResolvedChunkRequest::Fetch(request) + }; + self.deferred_prefetch = Some(DeferredChunkRequest::Resolved(request)); + Ok(false) + } + Err(( + ScheduleError::RequestTooLarge { + request, + max_inflight_bytes, + }, + _, + )) => Err(ZarrFdwError::InvalidMetadata(format!( + "object '{}' read limit {} exceeds max_inflight_bytes {max_inflight_bytes}", + request.identity.key, request.max_bytes + ))), + Err(( + ScheduleError::CachedObjectTooLarge { + request, + actual_bytes, + }, + _, + )) => Err(ZarrFdwError::InvalidMetadata(format!( + "cached object '{}' is {actual_bytes} bytes, exceeding its read limit of {}", + request.identity.key, request.max_bytes + ))), + } + } + + fn fill_prefetch_window(&mut self) -> ZarrFdwResult<()> { + loop { + let logical = if let Some(deferred) = self.deferred_prefetch.take() { + match deferred { + DeferredChunkRequest::Logical(indices) => indices, + DeferredChunkRequest::Resolved(request) => { + if !self.schedule_chunk_request(request)? { + break; + } + continue; + } + } + } else { + let mut indices = Vec::new(); + if !self.chunk_cursor.next_into(&mut indices) { + break; + } + indices + }; + let Some(request) = self.chunk_request(logical.clone(), self.prefetch.is_empty())? + else { + self.deferred_prefetch = Some(DeferredChunkRequest::Logical(logical)); + break; + }; + if !self.schedule_chunk_request(request)? { + break; + } + } + Ok(()) + } + + fn process_pending_interrupt(&mut self) -> ZarrFdwResult<()> { + if postgres_interrupt_pending() { + // No Rust future may remain owned by scan state when PostgreSQL's + // cancellation path raises ERROR through the backend stack. + self.prefetch.clear(); + self.deferred_prefetch = None; + process_postgres_interrupts(); + // PostgreSQL can defer interrupts while they are held. Never + // continue after dropping an already-advanced prefetch window. + return Err(ZarrFdwError::InvalidMetadata( + "query interruption was requested".to_string(), + )); + } + Ok(()) + } + + fn next_prefetched_chunk(&mut self) -> ZarrFdwResult> { + self.fill_prefetch_window()?; + let evictions_before = self.compressed_cache.evictions(); + let outcome = { + let runtime = &self.store.rt; + let prefetch = &mut self.prefetch; + let cache = &mut self.compressed_cache; + runtime.block_on(prefetch.next_interruptible(cache, postgres_interrupt_pending)) + }; + self.metrics.record_cache_evictions( + self.compressed_cache + .evictions() + .saturating_sub(evictions_before), + ); + + match outcome { + PrefetchNext::Ready(result) => { + self.current_chunk = result.request.context.logical_indices; + self.current_object_key = result.request.context.object_key; + self.metrics + .record_chunk_result(matches!(&result.object, CachedObject::Present(_))); + Ok(Some(result.object)) + } + PrefetchNext::FetchError { error, .. } => Err(error), + PrefetchNext::Interrupted => { + // All queued Rust futures have been dropped. Raise PostgreSQL's + // canonical cancellation only after leaving Runtime::block_on. + self.deferred_prefetch = None; + self.process_pending_interrupt()?; + Err(ZarrFdwError::InvalidMetadata( + "query interruption was requested".to_string(), + )) + } + PrefetchNext::Empty => Ok(None), + } + } + + /// Decode `self.current_chunk`, priming the within-chunk + /// index window from the active selection. + fn load_chunk(&mut self, encoded: CachedObject) -> ZarrFdwResult<()> { + self.process_pending_interrupt()?; + let meta = self + .axis_meta + .as_ref() + .expect("begin_scan must be called before iter_scan"); + let dt = self.dtype.expect("dtype set in begin_scan"); + let codec = self.codec.as_ref().expect("codec set in begin_scan"); + let ci = self.current_chunk.clone(); + debug_assert_eq!(self.sub_lo.len(), self.rank); + debug_assert_eq!(self.sub_hi.len(), self.rank); + debug_assert_eq!(self.sub_idx.len(), self.rank); + + // Effective (edge) chunk shape. Regular Zarr chunks retain the full + // declared shape; `eff` only controls which logical cells are emitted. + if ci.len() != self.rank { + return Err(ZarrFdwError::InvalidMetadata( + "chunk index rank does not match the array rank".to_string(), + )); + } + let (storage_shape, storage_cells, expected) = checked_chunk_layout(meta, dt.itemsize())?; + let mut eff = Vec::with_capacity(self.rank); + for (axis, &ci_d) in ci.iter().enumerate() { + let dim = meta.shape_extent(axis)?; + let chunk_len = storage_shape[axis]; + let chunk_index = usize::try_from(ci_d).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk index for axis {axis} exceeds this platform's index capacity" + )) + })?; + let start = chunk_index.checked_mul(chunk_len).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("chunk start offset overflow on axis {axis}")) + })?; + if start >= dim { + return Err(ZarrFdwError::InvalidMetadata(format!( + "chunk index {ci_d} starts outside array dimension {axis} with extent {dim}" + ))); + } + let remaining = dim.checked_sub(start).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "chunk start offset exceeds array dimension {axis}" + )) + })?; + eff.push(remaining.min(chunk_len)); + } + + if eff.contains(&0) { + // zero-length dimension: yield nothing from this chunk + self.sub_lo = (0..self.rank).map(|_| 1).collect(); + self.sub_hi = vec![0; self.rank]; + self.chunk_bytes.clear(); + return Ok(()); + } + + // Regular edge chunks retain the declared chunk shape. Use that full + // shape for byte validation and C-order strides; `eff` ignores the + // out-of-array region when deciding which cells to emit. + let object_key = self.current_object_key.clone(); + let (decoded, synthesized_fill) = match encoded { + CachedObject::Present(raw) => { + let started = Instant::now(); + let decoded = self + .store + .rt + .block_on(codec.decode_interruptible( + raw.as_ref().to_vec(), + &storage_shape, + dt.itemsize(), + postgres_interrupt_pending, + )) + .map_err(|error| { + ZarrFdwError::ReadError(std::io::Error::other(format!( + "chunk '{object_key}': {error}" + ))) + })?; + self.metrics.record_decompression_time(started.elapsed()); + let decoded = match decoded { + CodecDecode::Decoded(decoded) => decoded, + CodecDecode::Interrupted => { + self.process_pending_interrupt()?; + return Err(ZarrFdwError::InvalidMetadata( + "query interruption was requested".to_string(), + )); + } + }; + (decoded, false) + } + CachedObject::Missing => ( + filled_chunk_bytes(self.fill_bytes.as_deref(), storage_cells, &object_key)?, + true, + ), + }; + require_exact_decoded_len(&object_key, decoded.len(), expected)?; + self.metrics + .record_decoded_bytes(ReadKind::Data, decoded.len(), synthesized_fill); + self.chunk_bytes.clear(); + self.chunk_bytes.try_reserve_exact(expected).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate a decoded chunk of {expected} bytes" + )) + })?; + self.chunk_bytes.extend_from_slice(&decoded[..expected]); + self.chunk_shape = storage_shape; + + // within-chunk index window for this chunk + for d in 0..self.rank { + let chunk_len = self.chunk_shape[d]; + let chunk_index = usize::try_from(ci[d]).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "chunk index for axis {d} exceeds this platform's index capacity" + )) + })?; + let base = chunk_index.checked_mul(chunk_len).ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("chunk start offset overflow on axis {d}")) + })?; + let hi_def = eff[d].saturating_sub(1); + match &self.selection.axis_bounds()[d] { + Some(b) => { + self.sub_lo[d] = b.start.saturating_sub(base).min(hi_def); + self.sub_hi[d] = b.end.saturating_sub(base).min(hi_def); + } + None => { + self.sub_lo[d] = 0; + self.sub_hi[d] = hi_def; + } + } + } + self.sub_idx.clone_from(&self.sub_lo); + Ok(()) + } + + /// Emit the row at the current `sub_idx`, then advance to the next cell. + /// Returns `Ok(())`; the row should be consumed regardless. + fn emit_and_advance(&mut self, row: &mut Row) -> ZarrFdwResult<()> { + let decode_started = Instant::now(); + let value_cell = self.value_cell_at_cursor()?; + self.metrics.record_decoding_time(decode_started.elapsed()); + if self.capture_spatial_indices { + let mut global_indices = self.last_emitted_indices.take().unwrap_or_default(); + global_indices.clear(); + global_indices.try_reserve(self.rank).map_err(|_| { + ZarrFdwError::InvalidMetadata(format!( + "could not allocate {0} emitted array indexes", + self.rank + )) + })?; + for axis in 0..self.rank { + global_indices.push(self.global_index_at_cursor(axis)?); + } + self.last_emitted_indices = Some(global_indices); + } + + for col in &self.tgt_cols { + let cell = self.column_cell_at_cursor(&col.name, value_cell.as_ref())?; + row.push(col.name.as_str(), cell); + } + + self.metrics.record_cells(1, None); + self.metrics.record_tuple_emitted(); + self.advance_cursor(); + Ok(()) + } + + fn reduce_current_cell(&mut self) -> ZarrFdwResult<()> { + let aggregate_started = Instant::now(); + let decode_started = Instant::now(); + let value_cell = self.value_cell_at_cursor()?; + self.metrics.record_decoding_time(decode_started.elapsed()); + let mut matches = true; + for qual in &self.aggregate_quals { + let cell = self.column_cell_at_cursor(&qual.field, value_cell.as_ref())?; + if !qual_matches(qual, cell.as_ref())? { + matches = false; + break; + } + } + if !matches { + self.metrics.record_cells(1, Some(0)); + self.metrics + .record_aggregate_time(aggregate_started.elapsed()); + self.advance_cursor(); + return Ok(()); + } + + let values = self + .aggregate_defs + .iter() + .map(|aggregate| { + aggregate + .column + .as_ref() + .map(|column| self.column_cell_at_cursor(&column.name, value_cell.as_ref())) + .transpose() + .map(Option::flatten) + }) + .collect::>>()?; + let value_refs = values.iter().map(Option::as_ref).collect::>(); + self.aggregate_reducer + .as_mut() + .expect("aggregate reducer set in begin_aggregate_scan") + .observe(&value_refs)?; + self.metrics.record_cells(1, Some(1)); + self.metrics + .record_aggregate_time(aggregate_started.elapsed()); + self.advance_cursor(); + Ok(()) + } + + fn flush_persistent_stats_at_eof(&mut self) { + let metrics = self.metrics_snapshot(); + let encoded = metrics + .total_encoded_bytes() + .saturating_sub(self.flushed_encoded_bytes); + let cells = metrics + .logical_cells_examined + .saturating_sub(self.flushed_cells); + let tuples = metrics.tuples_emitted.saturating_sub(self.flushed_tuples); + let as_i64 = |value: u64| i64::try_from(value).unwrap_or(i64::MAX); + if encoded > 0 { + stats::inc_stats(FDW_NAME, stats::Metric::BytesIn, as_i64(encoded)); + } + if cells > 0 { + stats::inc_stats(FDW_NAME, stats::Metric::RowsIn, as_i64(cells)); + } + if tuples > 0 { + stats::inc_stats(FDW_NAME, stats::Metric::RowsOut, as_i64(tuples)); + } + self.flushed_encoded_bytes = metrics.total_encoded_bytes(); + self.flushed_cells = metrics.logical_cells_examined; + self.flushed_tuples = metrics.tuples_emitted; + self.rows_out = 0; + } + + fn metrics_snapshot(&self) -> ZarrScanMetrics { + let mut metrics = self.metrics.clone(); + metrics.data_get_calls = metrics + .data_get_calls + .saturating_add(self.remote_data_get_calls.load(Ordering::Relaxed)); + metrics.data_encoded_bytes = metrics + .data_encoded_bytes + .saturating_add(self.remote_data_encoded_bytes.load(Ordering::Relaxed)); + metrics.shard_payload_get_calls = metrics + .shard_payload_get_calls + .saturating_add(self.remote_shard_payload_get_calls.load(Ordering::Relaxed)); + metrics.shard_payload_encoded_bytes = metrics.shard_payload_encoded_bytes.saturating_add( + self.remote_shard_payload_encoded_bytes + .load(Ordering::Relaxed), + ); + metrics + } + + /// Make the next selected cell ready for a consumer. + /// + /// This is the shared chunk-execution state machine for tuple and aggregate + /// scans: preserve an already pending cell, otherwise fetch/cache/decode + /// chunks lazily until one has a non-empty selected window. Consumers own + /// cell decoding, residual filtering, cursor advancement, and EOF output. + fn ensure_cell_ready(&mut self) -> ZarrFdwResult { + loop { + if self.pending { + return Ok(true); + } + let Some(encoded) = self.next_prefetched_chunk()? else { + return Ok(false); + }; + self.load_chunk(encoded)?; + let empty_window = (0..self.rank).any(|axis| self.sub_lo[axis] > self.sub_hi[axis]); + if empty_window { + continue; + } + self.pending = true; + return Ok(true); + } + } + + fn current_cell_matches_selection(&self) -> ZarrFdwResult { + if self.dimension_selectors.is_empty() && self.bound_call_dimension_selectors.is_empty() { + return Ok(true); + } + for axis in 0..self.rank { + let index = self.global_index_at_cursor(axis)?; + if !self + .dimension_selectors + .matches_axis_index(axis, index, &self.coords)? + || !self.bound_call_dimension_selectors.matches_axis_index( + axis, + index, + &self.coords, + )? + { + return Ok(false); + } + } + Ok(true) + } + + fn ensure_selected_cell_ready(&mut self) -> ZarrFdwResult { + let mut skipped = 0usize; + while self.ensure_cell_ready()? { + if self.current_cell_matches_selection()? { + return Ok(true); + } + let matched = (!self.aggregate_defs.is_empty()).then_some(0); + self.metrics.record_cells(1, matched); + self.advance_cursor(); + skipped = skipped.saturating_add(1); + if skipped.is_multiple_of(SELECTOR_INTERRUPT_POLL_CELLS) { + self.process_pending_interrupt()?; + } + } + Ok(false) + } + + fn iter_aggregate_scan(&mut self, row: &mut Row) -> ZarrFdwResult> { + if self.aggregate_emitted { + self.flush_persistent_stats_at_eof(); + return Ok(None); + } + + while self.ensure_selected_cell_ready()? { + self.reduce_current_cell()?; + } + + let results = self + .aggregate_reducer + .take() + .expect("aggregate reducer set in begin_aggregate_scan") + .finish()?; + row.clear(); + for (alias, cell) in results { + row.push(&alias, cell); + } + self.aggregate_emitted = true; + self.rows_out = 1; + self.metrics.record_tuple_emitted(); + Ok(Some(())) + } + + fn iter_scalar_scan(&mut self, row: &mut Row) -> ZarrFdwResult> { + if !self.ensure_selected_cell_ready()? { + self.flush_persistent_stats_at_eof(); + return Ok(None); + } + row.clear(); + self.emit_and_advance(row)?; + self.rows_out += 1; + Ok(Some(())) + } +} + +fn cell_to_f64_bounds(cell: &Cell, is_time: bool, spec: TimeSpec) -> Option<(f64, f64)> { + if is_time { + match cell { + Cell::Timestamptz(v) => spec.pg_micros_to_raw_bounds((*v).into_inner()), + Cell::Timestamp(v) => spec.pg_micros_to_raw_bounds((*v).into_inner()), + _ => None, + } + } else { + let value = match cell { + Cell::F64(v) => Some(*v), + Cell::F32(v) => Some(*v as f64), + Cell::I64(v) => Some(*v as f64), + Cell::I32(v) => Some(*v as f64), + Cell::I16(v) => Some(*v as f64), + Cell::I8(v) => Some(*v as f64), + _ => None, + }?; + // Coordinate vectors are finite, while PostgreSQL gives NaN a total + // ordering above every non-NaN float. Binary-search range math uses + // ordinary IEEE comparisons, so a NaN bound could narrow the scan + // incorrectly (for example, every finite x satisfies x < NaN in + // PostgreSQL). Disable pruning and let the exact/local qual decide. + if value.is_nan() { + return None; + } + Some((value, value)) + } +} + +/// Translate a single qual into an optional `(lo, hi)` value range over an +/// axis's coordinate space. Returns `None` when the qual cannot be (or should +/// not be) used for pruning. +fn qual_to_range( + q: &Qual, + is_time: bool, + spec: TimeSpec, +) -> ZarrFdwResult, Option)>> { + let evaluated_value = q.param.as_ref().map(|_| q.evaluated_value()); + let value = match evaluated_value.as_ref() { + None => &q.value, + Some(ParamValue::Value(value)) => value, + // NULL comparisons cannot select a row, but a full scan is the safe + // pruning choice for normal scans whose clauses PostgreSQL rechecks. + Some(ParamValue::Null | ParamValue::Unevaluated) => return Ok(None), + }; + + if q.use_or { + // `IN (...)` -> bounding box over the values (over-approximated) + let Value::Array(cells) = value else { + return Ok(None); + }; + let mut lo = f64::INFINITY; + let mut hi = f64::NEG_INFINITY; + let mut found = false; + for c in cells { + if let Some((value_lo, value_hi)) = cell_to_f64_bounds(c, is_time, spec) { + lo = lo.min(value_lo); + hi = hi.max(value_hi); + found = true; + } + } + return if found { + Ok(Some((Some(lo), Some(hi)))) + } else { + Ok(None) + }; + } + + let Value::Cell(cell) = value else { + return Ok(None); + }; + let Some((value_lo, value_hi)) = cell_to_f64_bounds(cell, is_time, spec) else { + return Ok(None); + }; + let r = match q.operator.as_str() { + "=" => (Some(value_lo), Some(value_hi)), + ">" => (Some(value_hi), None), + ">=" => (Some(value_lo), None), + "<" => (None, Some(value_lo)), + "<=" => (None, Some(value_hi)), + // `<>`/LIKE/etc. cannot prune this axis + _ => return Ok(None), + }; + Ok(Some(r)) +} + +fn array_parent_path(array_path: &str) -> &str { + array_path + .rsplit_once('/') + .map(|(parent, _)| parent) + .unwrap_or_default() +} + +fn select_array_metadata_document( + array_path: &str, + v3: Option>, + v2: Option>, +) -> ZarrFdwResult { + match (v3, v2) { + (Some(_), Some(_)) => Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' contains both zarr.json and .zarray metadata" + ))), + (Some(bytes), None) => Ok(ArrayMetadataDocument::V3(bytes)), + (None, Some(bytes)) => Ok(ArrayMetadataDocument::V2(bytes)), + (None, None) => Err(ZarrFdwError::InvalidMetadata(format!( + "array '{array_path}' contains neither zarr.json nor .zarray metadata" + ))), + } +} + +fn read_optional_metadata_object( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + key: &str, +) -> ZarrFdwResult>> { + let bytes = store.get_object_optional_sync(key, MAX_METADATA_OBJECT_BYTES)?; + metrics.record_remote_get(ReadKind::Metadata, bytes.as_ref().map(Vec::len)); + Ok(bytes) +} + +/// Read and normalize exactly one array node without exposing format-specific +/// metadata to the scan executor. +fn read_array_node( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + array_path: &str, +) -> ZarrFdwResult { + let v3_key = join_key(array_path, "zarr.json"); + let v2_key = join_key(array_path, ".zarray"); + let v3 = read_optional_metadata_object(store, metrics, &v3_key)?; + let v2 = read_optional_metadata_object(store, metrics, &v2_key)?; + + let node = match select_array_metadata_document(array_path, v3, v2)? { + ArrayMetadataDocument::V3(bytes) => match parse_v3_node(&bytes)? { + NodeMeta::Array(node) => Ok(*node), + NodeMeta::Group(_) => Err(ZarrFdwError::InvalidMetadata(format!( + "node '{array_path}' is a Zarr v3 group, expected an array" + ))), + }, + ArrayMetadataDocument::V2(bytes) => { + let attributes = + read_array_attributes_optional(store, metrics, array_path)?.unwrap_or_default(); + parse_v2_array(&bytes, attributes) + } + }?; + validate_array_ancestors(store, metrics, array_path, node.format)?; + Ok(node) +} + +fn read_ome_group_attributes( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + group_path: &str, +) -> ZarrFdwResult> { + let v3_key = join_key(group_path, "zarr.json"); + let v2_group_key = join_key(group_path, ".zgroup"); + let v2_array_key = join_key(group_path, ".zarray"); + let v3 = read_optional_metadata_object(store, metrics, &v3_key)?; + let v2_group = read_optional_metadata_object(store, metrics, &v2_group_key)?; + let v2_array = read_optional_metadata_object(store, metrics, &v2_array_key)?; + if v2_group.is_some() || v2_array.is_some() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr group '{}' must be a Zarr v3 group", + if group_path.is_empty() { + "/" + } else { + group_path + } + ))); + } + let bytes = v3.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr group '{}' must be a Zarr v3 group", + if group_path.is_empty() { + "/" + } else { + group_path + } + )) + })?; + let group = match parse_v3_node(&bytes)? { + NodeMeta::Group(group) => group, + NodeMeta::Array(_) => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr group '{}' must be a Zarr v3 group", + if group_path.is_empty() { + "/" + } else { + group_path + } + ))); + } + }; + validate_array_ancestors(store, metrics, group_path, ZarrFormat::V3)?; + Ok(group.attributes) +} + +fn array_ancestor_paths(array_path: &str) -> Vec { + let components = array_path + .split('/') + .filter(|component| !component.is_empty()) + .collect::>(); + if components.is_empty() { + return Vec::new(); + } + let mut paths = vec![String::new()]; + let mut current = String::new(); + for component in components.iter().take(components.len() - 1) { + current = join_key(¤t, component); + paths.push(current.clone()); + } + paths +} + +fn validate_array_ancestors( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + array_path: &str, + format: ZarrFormat, +) -> ZarrFdwResult<()> { + for ancestor in array_ancestor_paths(array_path) { + let v3_key = join_key(&ancestor, "zarr.json"); + let v2_group_key = join_key(&ancestor, ".zgroup"); + let v2_array_key = join_key(&ancestor, ".zarray"); + let v3 = read_optional_metadata_object(store, metrics, &v3_key)?; + let v2_group = read_optional_metadata_object(store, metrics, &v2_group_key)?; + let v2_array = read_optional_metadata_object(store, metrics, &v2_array_key)?; + if v2_array.is_some() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "ancestor '{}' is an array, expected a group", + if ancestor.is_empty() { "/" } else { &ancestor } + ))); + } + match format { + ZarrFormat::V3 => { + if v2_group.is_some() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' has a Zarr v2 ancestor group '{}'", + if ancestor.is_empty() { "/" } else { &ancestor } + ))); + } + let bytes = v3.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "Zarr v3 array '{array_path}' requires explicit zarr.json metadata on ancestor group '{}'", + if ancestor.is_empty() { "/" } else { &ancestor } + )) + })?; + if !matches!(parse_v3_node(&bytes)?, NodeMeta::Group(_)) { + return Err(ZarrFdwError::InvalidMetadata(format!( + "ancestor '{}' is not a Zarr v3 group", + if ancestor.is_empty() { "/" } else { &ancestor } + ))); + } + } + ZarrFormat::V2 => { + if v3.is_some() { + return Err(ZarrFdwError::InvalidMetadata(format!( + "Zarr v2 array '{array_path}' has a Zarr v3 ancestor group '{}'", + if ancestor.is_empty() { "/" } else { &ancestor } + ))); + } + if let Some(bytes) = v2_group { + parse_v2_group(&bytes, Map::new())?; + } + } + } + } + Ok(()) +} + +fn validate_ome_hierarchy_versions( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + array_path: &str, +) -> ZarrFdwResult<()> { + for ancestor in array_ancestor_paths(array_path) { + let key = join_key(&ancestor, "zarr.json"); + let bytes = read_optional_metadata_object(store, metrics, &key)?.ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr array '{array_path}' requires explicit zarr.json metadata on ancestor group '{}'", + if ancestor.is_empty() { "/" } else { &ancestor } + )) + })?; + let attributes = match parse_v3_node(&bytes)? { + NodeMeta::Group(group) => group.attributes, + NodeMeta::Array(_) => { + return Err(ZarrFdwError::InvalidMetadata(format!( + "OME-Zarr ancestor '{}' is an array, expected a group", + if ancestor.is_empty() { "/" } else { &ancestor } + ))); + } + }; + validate_optional_ome_05_attributes( + if ancestor.is_empty() { "/" } else { &ancestor }, + Some(&attributes), + )?; + } + Ok(()) +} + +fn codec_pipeline_for_execution(meta: &ArrayMeta) -> ZarrFdwResult { + if meta.zarr_format == 2 { + CodecPipeline::from_v2(&meta.compressor) + } else { + Ok(meta.codec_pipeline.clone()) + } +} + +fn read_coordinate_metadata( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + prefix: &str, + name: &str, + expected_length: u64, +) -> ZarrFdwResult<(ArrayNode, Option)> { + let dir = join_key(prefix, name); + let node = + read_array_node(store, metrics, &dir).map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("coordinate array metadata: {e}"), + })?; + let meta = &node.meta; + meta.validate_coordinate() + .map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: e.to_string(), + })?; + + let itemsize = + coordinate_itemsize(&meta.dtype).map_err(|error| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("coordinate array dtype: {error}"), + })?; + let fill_value = coord_fill_value_to_f64(&meta.dtype, &meta.fill_value).map_err(|error| { + ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("coordinate array fill value: {error}"), + } + })?; + let coordinate_len = + meta.shape_extent(0) + .map_err(|error| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: error.to_string(), + })?; + if coordinate_len > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate array has {coordinate_len} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}" + ), + }); + } + let (storage_shape, _, _) = checked_chunk_layout(meta, itemsize).map_err(|error| { + ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: error.to_string(), + } + })?; + if storage_shape[0] > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate chunk has {} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}", + storage_shape[0] + ), + }); + } + codec_pipeline_for_execution(meta).map_err(|error| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("coordinate array codec pipeline: {error}"), + })?; + if meta.shape[0] != expected_length { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate array has shape {} but the dimension has shape {expected_length}", + meta.shape[0] + ), + }); + } + Ok((node, fill_value)) +} + +/// Read a validated 1D numeric coordinate array and return its values as `f64`. +fn read_coordinate_values( + fdw: &mut ZarrFdw, + prefix: &str, + name: &str, + coordinate: &ArrayNode, + fill_value: Option, +) -> ZarrFdwResult> { + let dir = join_key(prefix, name); + let meta = &coordinate.meta; + + let coordinate_len = meta + .shape_extent(0) + .map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: e.to_string(), + })?; + if coordinate_len > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate array has {coordinate_len} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}" + ), + }); + } + let itemsize = + coordinate_itemsize(&meta.dtype).map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: e.to_string(), + })?; + let (storage_shape, _, expected_bytes) = + checked_chunk_layout(meta, itemsize).map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: e.to_string(), + })?; + let chunk_len = storage_shape[0]; + if chunk_len > MAX_COORDINATE_VALUES { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate chunk has {chunk_len} values, exceeding the safety limit of {MAX_COORDINATE_VALUES}" + ), + }); + } + + let codec = codec_pipeline_for_execution(meta)?; + let per_axis = meta.chunks_per_axis(); + let chunk_count = + usize::try_from(per_axis[0]).map_err(|_| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "coordinate chunk count exceeds this platform's index capacity".to_string(), + })?; + let mut values = Vec::new(); + values + .try_reserve_exact(coordinate_len) + .map_err(|_| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("could not allocate {coordinate_len} coordinate values"), + })?; + for chunk_index in 0..chunk_count { + let ci = u64::try_from(chunk_index).map_err(|_| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "coordinate chunk index exceeds the Zarr u64 index capacity".to_string(), + })?; + let encoded_limit = codec.encoded_read_limit(expected_bytes)?; + let (object_key, encoded) = match &meta.storage_layout { + StorageLayout::Direct => { + let chunk = chunk_key(&meta.chunk_key_encoding, &[ci]); + let object_key = join_key(&dir, &chunk); + let encoded = fdw + .store + .get_object_optional_sync(&object_key, encoded_limit)?; + fdw.metrics + .record_remote_get(ReadKind::Coordinate, encoded.as_ref().map(Vec::len)); + (object_key, encoded) + } + StorageLayout::Sharded(config) => { + let address = config.chunk_address(&[ci])?; + let shard_key = join_key( + &dir, + &chunk_key(&meta.chunk_key_encoding, &address.shard_indices), + ); + let index = match fdw.resolve_shard_index( + config, + &shard_key, + ReadKind::Coordinate, + true, + )? { + ShardIndexResolution::Ready(index) => index, + ShardIndexResolution::WouldBlock => { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "coordinate shard index resolution unexpectedly blocked" + .to_string(), + }); + } + }; + match index { + None => (shard_key, None), + Some(index) => { + let entry = index.entry(&address.inner_indices)?; + match index.payload_read_identity(entry)? { + None => (shard_key, None), + Some(identity) => { + let payload_bytes = match &identity.range { + ReadRange::Exact { length, .. } => usize::try_from(*length) + .map_err(|_| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "inner chunk range in shard '{shard_key}' exceeds this platform's index capacity" + ), + })?, + ReadRange::Whole | ReadRange::Suffix { .. } => { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "inner chunk in shard '{shard_key}' did not resolve to an exact byte range" + ), + }); + } + }; + if payload_bytes > encoded_limit { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "inner chunk range in shard '{shard_key}' is {payload_bytes} bytes, exceeding its encoded read limit of {encoded_limit}" + ), + }); + } + if payload_bytes > fdw.max_inflight_bytes { + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "inner chunk range in shard '{shard_key}' is {payload_bytes} bytes, exceeding max_inflight_bytes {}", + fdw.max_inflight_bytes + ), + }); + } + fdw.metrics.record_remote_request(ReadKind::Coordinate); + let response = fdw.store.get_object_range_sync(identity)?; + let response_bytes = + response.as_ref().map(|response| response.bytes.len()); + if let Some(bytes) = response_bytes { + fdw.metrics + .record_remote_response_bytes(ReadKind::Coordinate, bytes); + } + fdw.metrics.record_shard_payload_get(response_bytes); + let response = response.ok_or_else(|| { + ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "indexed shard object '{shard_key}' disappeared before its payload range was read" + ), + } + })?; + (shard_key, Some(response.bytes)) + } + } + } + } + } + }; + let decoded_values = match encoded { + Some(raw) => { + let started = Instant::now(); + let decoded = fdw + .store + .rt + .block_on(codec.decode_interruptible( + raw, + &storage_shape, + itemsize, + postgres_interrupt_pending, + )) + .map_err(|e| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!("chunk '{object_key}': {e}"), + })?; + let decoded = match decoded { + CodecDecode::Decoded(decoded) => decoded, + CodecDecode::Interrupted => { + process_postgres_interrupts(); + return Err(ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "query interruption was requested".to_string(), + }); + } + }; + fdw.metrics.record_decompression_time(started.elapsed()); + fdw.metrics + .record_decoded_bytes(ReadKind::Coordinate, decoded.len(), false); + coord_bytes_to_f64(&meta.dtype, &decoded[..expected_bytes])? + } + None => { + fdw.metrics + .record_decoded_bytes(ReadKind::Coordinate, expected_bytes, true); + filled_coordinate_values(fill_value, chunk_len, &object_key, name)? + } + }; + let start = chunk_index.checked_mul(chunk_len).ok_or_else(|| { + ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: "coordinate chunk start offset overflow".to_string(), + } + })?; + let remaining = + coordinate_len + .checked_sub(start) + .ok_or_else(|| ZarrFdwError::CoordinateReadError { + axis: name.to_string(), + error: format!( + "coordinate chunk '{object_key}' starts beyond the declared array length" + ), + })?; + let effective_len = remaining.min(chunk_len); + values.extend_from_slice(&decoded_values[..effective_len]); + } + Ok(values) +} + +fn read_array_attributes_optional( + store: &ZarrStore, + metrics: &mut ZarrScanMetrics, + array_dir: &str, +) -> ZarrFdwResult>> { + let key = join_key(array_dir, ".zattrs"); + let Some(bytes) = store.get_object_optional_sync(&key, MAX_METADATA_OBJECT_BYTES)? else { + metrics.record_remote_get(ReadKind::Metadata, None); + return Ok(None); + }; + metrics.record_remote_get(ReadKind::Metadata, Some(bytes.len())); + let value = serde_json::from_slice::(&bytes).map_err(|error| { + ZarrFdwError::InvalidMetadata(format!("could not parse '{key}': {error}")) + })?; + let attributes = value.as_object().cloned().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!("'{key}' must contain a JSON object")) + })?; + Ok(Some(attributes)) +} + +fn parse_boolean_option(name: &str, value: &str) -> ZarrFdwResult { + match value { + "true" => Ok(true), + "false" => Ok(false), + _ => Err(ZarrFdwError::InvalidOptionValue { + option: name.to_string(), + message: "must be 'true' or 'false'".to_string(), + }), + } +} + +fn boolean_table_option(options: &HashMap, name: &str) -> ZarrFdwResult { + options + .get(name) + .map(|value| parse_boolean_option(name, value)) + .transpose() + .map(|value| value.unwrap_or(false)) +} + +fn bounded_server_usize_option( + options: &HashMap, + name: &str, + default: usize, + min: usize, + max: usize, +) -> ZarrFdwResult { + let Some(raw) = options.get(name) else { + return Ok(default); + }; + let value = raw + .parse::() + .map_err(|_| ZarrFdwError::InvalidOptionValue { + option: name.to_string(), + message: format!("must be an integer between {min} and {max}"), + })?; + if !(min..=max).contains(&value) { + return Err(ZarrFdwError::InvalidOptionValue { + option: name.to_string(), + message: format!("must be between {min} and {max}"), + }); + } + Ok(value) +} + +fn compressed_cache_bytes_option(options: &HashMap) -> ZarrFdwResult { + let Some(raw) = options.get(OPT_COMPRESSED_CACHE_BYTES) else { + return Ok(DEFAULT_COMPRESSED_CACHE_BYTES); + }; + let value = raw + .parse::() + .map_err(|_| ZarrFdwError::InvalidOptionValue { + option: OPT_COMPRESSED_CACHE_BYTES.to_string(), + message: format!("must be an integer between 0 and {MAX_COMPRESSED_CACHE_BYTES}"), + })?; + if value > MAX_COMPRESSED_CACHE_BYTES { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_COMPRESSED_CACHE_BYTES.to_string(), + message: format!("must be between 0 and {MAX_COMPRESSED_CACHE_BYTES}"), + }); + } + Ok(value) +} + +fn scalable_execution_options( + options: &HashMap, +) -> ZarrFdwResult<(usize, usize, usize)> { + Ok(( + bounded_server_usize_option( + options, + OPT_MAX_CONCURRENT_READS, + DEFAULT_MAX_CONCURRENT_READS, + 1, + MAX_CONCURRENT_READS, + )?, + bounded_server_usize_option( + options, + OPT_MAX_INFLIGHT_BYTES, + DEFAULT_MAX_INFLIGHT_BYTES, + MIN_MAX_INFLIGHT_BYTES, + MAX_MAX_INFLIGHT_BYTES, + )?, + compressed_cache_bytes_option(options)?, + )) +} + +fn validate_time_from_attrs_options( + time_from_attrs: bool, + has_time_unit: bool, + has_time_origin: bool, +) -> ZarrFdwResult<()> { + if time_from_attrs && (has_time_unit || has_time_origin) { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_TIME_FROM_ATTRS.to_string(), + message: format!("cannot be combined with '{OPT_TIME_UNIT}' or '{OPT_TIME_ORIGIN}'"), + }); + } + Ok(()) +} + +fn option_value(options: &[Option], name: &str) -> Option { + let prefix = format!("{name}="); + options + .iter() + .flatten() + .find_map(|kv| kv.strip_prefix(&prefix).map(|v| v.to_string())) +} + +impl ForeignDataWrapper for ZarrFdw { + fn new(server: ForeignServer) -> ZarrFdwResult { + let (configured_max_concurrent_reads, max_inflight_bytes, compressed_cache_bytes) = + scalable_execution_options(&server.options)?; + let store = ZarrStore::new(&server)?; + let max_concurrent_reads = + store.effective_max_concurrent_reads(configured_max_concurrent_reads); + let prefetch = OrderedPrefetch::new( + max_concurrent_reads, + max_inflight_bytes, + INTERRUPT_POLL_INTERVAL, + ) + .map_err(|error| ZarrFdwError::InvalidOptionValue { + option: OPT_MAX_CONCURRENT_READS.to_string(), + message: error.to_string(), + })?; + stats::inc_stats(FDW_NAME, stats::Metric::CreateTimes, 1); + Ok(Self { + store, + tgt_cols: Vec::new(), + array_dir: String::new(), + axes: Vec::new(), + array_attributes: Map::new(), + selected_ome_level: None, + axis_roles: Vec::new(), + rank: 0, + axis_meta: None, + dtype: None, + codec: None, + scientific_decoder: None, + fill_bytes: None, + coords: Vec::new(), + selection: Selection::default(), + dimension_selectors: BoundDimensionSelectors::default(), + call_dimension_selectors: DimensionSelectors::default(), + bound_call_dimension_selectors: BoundDimensionSelectors::default(), + chunk_cursor: ChunkIndexCursor::default(), + current_chunk: Vec::new(), + current_object_key: String::new(), + deferred_prefetch: None, + prefetch, + compressed_cache: CompressedChunkCache::new( + compressed_cache_bytes, + MAX_COMPRESSED_CACHE_ENTRIES, + ), + shard_index_cache: ShardIndexCache::new(0, 0), + payload_cache_bytes: compressed_cache_bytes, + shard_index_cache_bytes: 0, + cache_layout_sharded: false, + max_concurrent_reads, + max_inflight_bytes, + compressed_cache_bytes, + metrics: ZarrScanMetrics::default(), + remote_data_get_calls: Arc::new(AtomicU64::new(0)), + remote_data_encoded_bytes: Arc::new(AtomicU64::new(0)), + remote_shard_payload_get_calls: Arc::new(AtomicU64::new(0)), + remote_shard_payload_encoded_bytes: Arc::new(AtomicU64::new(0)), + flushed_encoded_bytes: 0, + flushed_cells: 0, + flushed_tuples: 0, + chunk_bytes: Vec::new(), + chunk_shape: Vec::new(), + sub_lo: Vec::new(), + sub_hi: Vec::new(), + sub_idx: Vec::new(), + capture_spatial_indices: false, + last_emitted_indices: None, + pending: false, + aggregate_defs: Vec::new(), + aggregate_quals: Vec::new(), + aggregate_reducer: None, + aggregate_emitted: false, + time_spec: TimeSpec::default(), + rows_out: 0, + }) + } + + fn supported_aggregates(&self) -> Vec { + vec![ + AggregateKind::Count, + AggregateKind::CountColumn, + AggregateKind::Sum, + AggregateKind::Avg, + AggregateKind::Min, + AggregateKind::Max, + ] + } + + fn explain(&self) -> Vec { + let Some(meta) = self.axis_meta.as_ref() else { + return Vec::new(); + }; + let chunk_shape = meta + .chunks + .iter() + .map(|&extent| usize::try_from(extent).unwrap_or(usize::MAX)) + .collect::>(); + let dtype = self + .dtype + .map(|dtype| format!("{dtype:?}")) + .unwrap_or_else(|| meta.dtype.clone()); + let codec = self + .codec + .as_ref() + .map(CodecPipeline::ordered_label) + .unwrap_or_else(|| "unknown".to_string()); + let storage_layout = meta.storage_layout.ordered_label(); + let (shard_shape, index_location) = match &meta.storage_layout { + StorageLayout::Direct => (None, None), + StorageLayout::Sharded(config) => ( + Some(config.shard_shape.as_slice()), + Some(config.index_location.label()), + ), + }; + let aggregate_mode = if self.aggregate_defs.is_empty() { + "none".to_string() + } else { + self.aggregate_defs + .iter() + .map(|aggregate| aggregate.kind.sql_name()) + .collect::>() + .join(", ") + }; + let mut properties = self + .metrics_snapshot() + .explain_properties(ZarrExplainContext { + array: &self.array_dir, + dimensions: &self.axes, + shape: &meta.shape, + chunk_shape: &chunk_shape, + dtype: &dtype, + codec: &codec, + storage_backend: self.store.backend_label(), + storage_layout: &storage_layout, + shard_shape, + index_location, + aggregate_mode: &aggregate_mode, + max_concurrent_reads: self.max_concurrent_reads, + max_inflight_bytes: self.max_inflight_bytes, + compressed_cache_bytes: self.payload_cache_bytes, + cache_entries: self.compressed_cache.len(), + cache_resident_bytes: self.compressed_cache.resident_bytes(), + shard_index_cache_bytes: self.shard_index_cache_bytes, + shard_index_cache_entries: self.shard_index_cache.len(), + shard_index_cache_resident_bytes: self.shard_index_cache.resident_bytes(), + }); + if let Some(level) = &self.selected_ome_level { + properties.push(ExplainProperty::text( + "Zarr OME Group", + if level.group_path.is_empty() { + "/" + } else { + &level.group_path + }, + )); + properties.push(ExplainProperty::unsigned( + "Zarr OME Multiscale Index", + u64::try_from(level.multiscale_index).unwrap_or(u64::MAX), + )); + properties.push(ExplainProperty::unsigned( + "Zarr OME Level Index", + u64::try_from(level.level_index).unwrap_or(u64::MAX), + )); + properties.push(ExplainProperty::text( + "Zarr OME Effective Scale", + format!("{:?}", level.transform.scale), + )); + properties.push(ExplainProperty::text( + "Zarr OME Effective Translation", + format!("{:?}", level.transform.translation), + )); + } + properties + } + + #[allow(clippy::too_many_arguments)] + fn can_pushdown_aggregate( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + quals: &[Qual], + base_columns: &[Column], + all_base_quals_extracted: bool, + _options: &HashMap, + ) -> ZarrFdwResult { + if aggregates.is_empty() || !group_by.is_empty() || !all_base_quals_extracted { + return Ok(false); + } + if aggregates.iter().any(|aggregate| { + !aggregate_signature_supported(aggregate) + || aggregate.column.as_ref().is_some_and(|column| { + !base_columns + .iter() + .any(|base| base.name == column.name && base.type_oid == column.type_oid) + }) + }) { + return Ok(false); + } + if quals.iter().any(|qual| { + !qual_shape_supported(qual) + || !base_columns.iter().any(|column| column.name == qual.field) + }) { + return Ok(false); + } + Ok(true) + } + + fn get_aggregate_rel_size( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + _quals: &[Qual], + _options: &HashMap, + ) -> ZarrFdwResult<(i64, i32)> { + debug_assert!(group_by.is_empty()); + let width = aggregates.iter().fold(0_i32, |sum, aggregate| { + sum.saturating_add(estimated_pg_type_width(aggregate.type_oid)) + }); + Ok((1, width.max(DEFAULT_EMPTY_PROJECTION_WIDTH))) + } + + fn get_rel_size( + &mut self, + _quals: &[Qual], + columns: &[Column], + _sorts: &[Sort], + _limit: &Option, + _options: &HashMap, + ) -> ZarrFdwResult<(i64, i32)> { + // Do not fetch array metadata here: this callback also runs for plain + // EXPLAIN, where remote latency/auth failures would be surprising. + Ok(conservative_rel_size(columns)) + } + + fn begin_scan( + &mut self, + quals: &[Qual], + columns: &[Column], + _sorts: &[Sort], + _limit: &Option, + options: &HashMap, + ) -> ZarrFdwResult<()> { + self.aggregate_defs.clear(); + self.aggregate_quals.clear(); + self.aggregate_reducer = None; + self.aggregate_emitted = false; + self.dimension_selectors = BoundDimensionSelectors::default(); + self.bound_call_dimension_selectors = BoundDimensionSelectors::default(); + self.tgt_cols = columns.to_vec(); + let dimension_selectors = + DimensionSelectors::parse(options.get(OPT_DIMENSION_SELECTORS).map(String::as_str))?; + let time_from_attrs = boolean_table_option(options, OPT_TIME_FROM_ATTRS)?; + validate_time_from_attrs_options( + time_from_attrs, + options.contains_key(OPT_TIME_UNIT), + options.contains_key(OPT_TIME_ORIGIN), + )?; + let decode_cf = boolean_table_option(options, OPT_DECODE_CF)?; + + let multiscale_selection = multiscale_selection_options(options)?; + let (value_node, dataset, coordinate_nodes, coordinate_fill_values, selected_ome_level) = + if let Some(selection) = multiscale_selection { + let group_attributes = + read_ome_group_attributes(&self.store, &mut self.metrics, &selection.group)?; + let level = resolve_ome_05_level( + &selection.group, + &group_attributes, + selection.index, + selection.level, + )?; + validate_ome_hierarchy_versions(&self.store, &mut self.metrics, &level.array_path)?; + let value_node = + read_array_node(&self.store, &mut self.metrics, &level.array_path)?; + validate_optional_ome_05_attributes( + &level.array_path, + Some(&value_node.attributes), + )?; + value_node.meta.validate()?; + let dataset = ome_rank2_dataset(&level.array_path, &value_node, &level)?; + let rank = dataset.dimensions().len(); + ( + value_node, + dataset, + vec![None; rank], + vec![None; rank], + Some(level), + ) + } else { + // `array_group` scopes one ordinary Zarr array; the default is + // the store root. Existing behavior is unchanged outside OME + // selection mode. + let array_dir = options + .get(OPT_ARRAY_GROUP) + .map(|path| path.trim_matches('/').to_string()) + .unwrap_or_default(); + let value_node = read_array_node(&self.store, &mut self.metrics, &array_dir)?; + value_node.meta.validate()?; + let dimension_names = named_dimensions(&value_node, &array_dir)?; + let coordinate_parent = array_parent_path(&array_dir); + let mut coordinate_nodes = Vec::with_capacity(dimension_names.len()); + let mut coordinate_fill_values = Vec::with_capacity(dimension_names.len()); + let mut aligned_nodes = Vec::with_capacity(dimension_names.len()); + for (name, &length) in dimension_names.iter().zip(value_node.meta.shape.iter()) { + let (coordinate, fill_value) = read_coordinate_metadata( + &self.store, + &mut self.metrics, + coordinate_parent, + name, + length, + )?; + aligned_nodes.push(coordinate.clone()); + coordinate_nodes.push(Some(coordinate)); + coordinate_fill_values.push(fill_value); + } + let dataset = + named_array_dataset(&array_dir, &value_node, &dimension_names, &aligned_nodes)?; + ( + value_node, + dataset, + coordinate_nodes, + coordinate_fill_values, + None, + ) + }; + self.selected_ome_level = selected_ome_level; + let meta = &value_node.meta; + let value_attributes = &value_node.attributes; + let variable = dataset.variable(); + self.array_dir = variable.path().to_string(); + self.axes = dataset.axis_names(); + let bound_dimension_selectors = dimension_selectors.bind(&self.axes, &meta.shape)?; + let bound_call_dimension_selectors = self + .call_dimension_selectors + .bind(&self.axes, &meta.shape)?; + debug_assert_eq!(variable.dimensions(), self.axes.as_slice()); + self.axis_roles = dataset + .dimensions() + .iter() + .map(|dimension| dimension.semantic_role()) + .collect(); + let rank = dataset.dimensions().len(); + self.rank = rank; + let dtype = DType::parse(variable.dtype())?; + let time_axis = dataset + .dimensions() + .iter() + .position(|dimension| dimension.semantic_role() == DimensionRole::Time); + let has_manual_time_options = + options.contains_key(OPT_TIME_UNIT) || options.contains_key(OPT_TIME_ORIGIN); + if (time_from_attrs || has_manual_time_options) && time_axis.is_none() { + let option = if time_from_attrs { + OPT_TIME_FROM_ATTRS + } else if options.contains_key(OPT_TIME_UNIT) { + OPT_TIME_UNIT + } else { + OPT_TIME_ORIGIN + }; + return Err(ZarrFdwError::InvalidOptionValue { + option: option.to_string(), + message: "requires exactly one discovered Time dimension".to_string(), + }); + } + let time_spec = if time_from_attrs { + let axis = time_axis.ok_or_else(|| ZarrFdwError::InvalidOptionValue { + option: OPT_TIME_FROM_ATTRS.to_string(), + message: "requires exactly one discovered Time dimension".to_string(), + })?; + let coordinate = coordinate_nodes[axis].as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata( + "time_from_attrs requires a stored time coordinate array".to_string(), + ) + })?; + TimeSpec::from_cf_attributes(&coordinate.attributes)? + } else { + TimeSpec::from_legacy_options( + options.get(OPT_TIME_UNIT).map(String::as_str), + options.get(OPT_TIME_ORIGIN).map(String::as_str), + )? + }; + let scientific_decoder = if decode_cf { + Some(ScientificValueDecoder::from_attributes( + dtype, + value_attributes, + )?) + } else { + None + }; + validate_column_types(columns, &dataset, dtype, decode_cf)?; + + // single-array MVP: at most one non-dimension (value) column allowed + let value_cols = columns + .iter() + .filter(|column| !dataset.is_dimension(&column.name)) + .count(); + if value_cols > 1 { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_BANDS.to_string(), + message: format!( + "single-array execution supports at most one value column, got {value_cols}" + ), + }); + } + self.axis_meta = Some(meta.clone()); + self.array_attributes = value_attributes.clone(); + self.dtype = Some(dtype); + self.codec = Some(codec_pipeline_for_execution(meta)?); + self.scientific_decoder = scientific_decoder; + self.fill_bytes = fill_value_bytes(dtype, &meta.fill_value)?; + self.time_spec = time_spec; + if matches!(&meta.storage_layout, StorageLayout::Sharded(_)) + || coordinate_nodes.iter().flatten().any(|coordinate| { + matches!(&coordinate.meta.storage_layout, StorageLayout::Sharded(_)) + }) + { + self.configure_sharded_cache_budget(); + } + + // Coordinate metadata is required for every dimension, but coordinate + // chunk values are needed only for projected or restricted dimensions. + let required_coordinates = dataset + .dimensions() + .iter() + .enumerate() + .map(|(axis, dimension)| { + columns.iter().any(|column| column.name == dimension.name()) + || bound_dimension_selectors.requires_coordinate(axis) + || bound_call_dimension_selectors.requires_coordinate(axis) + }) + .collect::>(); + checked_total_coordinate_values( + dataset + .dimensions() + .iter() + .zip(required_coordinates.iter()) + .map(|(dimension, &required)| (dimension.name(), dimension.length(), required)), + MAX_TOTAL_COORDINATE_VALUES, + )?; + let mut coords = Vec::with_capacity(rank); + for (axis, dimension) in dataset.dimensions().iter().enumerate() { + if !required_coordinates[axis] { + coords.push(None); + continue; + } + let values = match dimension.coordinate_source() { + CoordinateSource::Stored(coordinate) => { + let node = coordinate_nodes[axis].as_ref().ok_or_else(|| { + ZarrFdwError::InvalidMetadata(format!( + "stored coordinate '{}' has no loaded metadata", + dimension.name() + )) + })?; + validate_coordinate_decoding_attributes(dimension.name(), &node.attributes)?; + read_coordinate_values( + self, + coordinate.parent(), + coordinate.name(), + node, + coordinate_fill_values[axis], + )? + } + CoordinateSource::Affine { scale, translation } => affine_coordinate_values( + dimension.name(), + dimension.length(), + *scale, + *translation, + )?, + }; + validate_coordinate_values(dimension.name(), &values)?; + coords.push(Some(values)); + } + self.coords = coords; + + // translate quals into per-axis value ranges and intersect them + let mut ranges: Vec = vec![(None, None); rank]; + for q in quals { + let Some(axis) = self.axes.iter().position(|a| a == &q.field) else { + continue; + }; + let is_time = self.axis_roles[axis] == DimensionRole::Time; + let Some((lo, hi)) = qual_to_range(q, is_time, self.time_spec)? else { + continue; + }; + let cur = &mut ranges[axis]; + if let Some(l) = lo { + cur.0 = Some(match cur.0 { + Some(x) => x.max(l), + None => l, + }); + } + if let Some(h) = hi { + cur.1 = Some(match cur.1 { + Some(x) => x.min(h), + None => h, + }); + } + } + + let planner = ScanPlanner::new(meta); + let qual_selection = + planner.selection_from_coordinate_ranges(&self.axes, &self.coords, &ranges)?; + let selector_selection = + bound_dimension_selectors.resolve(&meta.shape, &self.coords, || { + process_postgres_interrupts(); + Ok(()) + })?; + let call_selector_selection = + bound_call_dimension_selectors.resolve(&meta.shape, &self.coords, || { + process_postgres_interrupts(); + Ok(()) + })?; + let plan = planner.plan( + qual_selection + .intersect(selector_selection) + .intersect(call_selector_selection), + )?; + self.dimension_selectors = bound_dimension_selectors; + self.bound_call_dimension_selectors = bound_call_dimension_selectors; + self.apply_scan_plan(plan, false) + } + + fn begin_aggregate_scan_with_base_columns( + &mut self, + aggregates: &[Aggregate], + group_by: &[Column], + quals: &[Qual], + base_columns: &[Column], + options: &HashMap, + ) -> ZarrFdwResult<()> { + debug_assert!(!aggregates.is_empty()); + debug_assert!(group_by.is_empty()); + >::begin_scan( + self, + quals, + base_columns, + &[], + &None, + options, + )?; + self.aggregate_defs = aggregates.to_vec(); + self.aggregate_quals = quals.to_vec(); + self.aggregate_reducer = Some(AggregateReducer::new(aggregates)?); + self.aggregate_emitted = false; + Ok(()) + } + + fn iter_scan(&mut self, row: &mut Row) -> ZarrFdwResult> { + let result = if !self.aggregate_defs.is_empty() { + self.iter_aggregate_scan(row) + } else { + self.iter_scalar_scan(row) + }; + if result.is_err() { + // PostgreSQL may unwind immediately after receiving the error. + // Drop all owned I/O futures before leaving the callback. + self.prefetch.clear(); + self.deferred_prefetch = None; + } + result + } + + fn re_scan(&mut self) -> ZarrFdwResult<()> { + self.chunk_cursor.reset(); + self.current_chunk.clear(); + self.prefetch.clear(); + self.deferred_prefetch = None; + self.metrics.record_rescan(); + self.chunk_bytes.clear(); + self.chunk_shape.clear(); + [self.sub_lo, self.sub_hi, self.sub_idx] = zeroed_scan_cursors(self.rank); + self.last_emitted_indices = None; + self.pending = false; + self.rows_out = 0; + if !self.aggregate_defs.is_empty() { + self.aggregate_reducer = Some(AggregateReducer::new(&self.aggregate_defs)?); + self.aggregate_emitted = false; + } + Ok(()) + } + + fn end_scan(&mut self) -> ZarrFdwResult<()> { + self.prefetch.clear(); + self.deferred_prefetch = None; + self.tgt_cols.clear(); + self.array_dir.clear(); + self.axes.clear(); + self.array_attributes.clear(); + self.selected_ome_level = None; + self.chunk_bytes.clear(); + self.chunk_shape.clear(); + self.chunk_cursor = ChunkIndexCursor::default(); + self.current_chunk.clear(); + self.coords.clear(); + self.selection = Selection::default(); + self.dimension_selectors = BoundDimensionSelectors::default(); + self.call_dimension_selectors = DimensionSelectors::default(); + self.bound_call_dimension_selectors = BoundDimensionSelectors::default(); + self.axis_roles.clear(); + self.sub_lo.clear(); + self.sub_hi.clear(); + self.sub_idx.clear(); + self.capture_spatial_indices = false; + self.last_emitted_indices = None; + self.rank = 0; + self.axis_meta = None; + self.dtype = None; + self.codec = None; + self.scientific_decoder = None; + self.fill_bytes = None; + self.pending = false; + self.aggregate_defs.clear(); + self.aggregate_quals.clear(); + self.aggregate_reducer = None; + self.aggregate_emitted = false; + self.time_spec = TimeSpec::default(); + self.rows_out = 0; + Ok(()) + } + + fn validator(options: Vec>, catalog: Option) -> ZarrFdwResult<()> { + if let Some(oid) = catalog { + match oid { + FOREIGN_SERVER_RELATION_ID => { + check_options_contain(&options, "store_url")?; + let server_options = options + .iter() + .flatten() + .filter_map(|option| option.split_once('=')) + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect::>(); + let backend = validate_store_options(&server_options)?; + validate_store_definition_privilege(backend)?; + scalable_execution_options(&server_options)?; + } + FOREIGN_TABLE_RELATION_ID => { + let table_options = options + .iter() + .flatten() + .filter_map(|option| option.split_once('=')) + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect::>(); + multiscale_selection_options(&table_options)?; + if let Some(v) = option_value(&options, OPT_TIME_UNIT) { + TimeSpec::from_legacy_options(Some(&v), None)?; + } + if let Some(v) = option_value(&options, OPT_TIME_ORIGIN) { + TimeSpec::from_legacy_options(None, Some(&v))?; + } + if let Some(v) = option_value(&options, OPT_TIME_FROM_ATTRS) { + let time_from_attrs = parse_boolean_option(OPT_TIME_FROM_ATTRS, &v)?; + validate_time_from_attrs_options( + time_from_attrs, + option_value(&options, OPT_TIME_UNIT).is_some(), + option_value(&options, OPT_TIME_ORIGIN).is_some(), + )?; + } + if let Some(v) = option_value(&options, OPT_DECODE_CF) { + parse_boolean_option(OPT_DECODE_CF, &v)?; + } + DimensionSelectors::parse( + option_value(&options, OPT_DIMENSION_SELECTORS).as_deref(), + )?; + if let Some(v) = option_value(&options, OPT_ARRAY_GROUP) + && (v.trim_matches('/').is_empty() || v.contains("..")) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_ARRAY_GROUP.to_string(), + message: "must be a non-empty array path inside the store".to_string(), + }); + } + if let Some(v) = option_value(&options, OPT_BANDS) + && v.split(',').any(|b| b.trim().is_empty()) + { + return Err(ZarrFdwError::InvalidOptionValue { + option: OPT_BANDS.to_string(), + message: "must be a comma-separated list of band column names" + .to_string(), + }); + } + } + _ => {} + } + } + Ok(()) + } +} + +#[cfg(test)] +mod unit_tests { + use super::super::meta::{ChunkKeyEncoding, ZarrFormat}; + use super::*; + + fn column(name: &str, type_oid: pg_sys::Oid) -> Column { + Column { + name: name.to_string(), + num: 1, + type_oid, + } + } + + fn array_meta(shape: Vec, chunks: Vec) -> ArrayMeta { + ArrayMeta { + zarr_format: 2, + shape, + chunks, + dtype: ", + chunks: Vec, + attributes: Map, + ) -> ArrayNode { + ArrayNode { + format: ZarrFormat::V2, + meta: array_meta(shape, chunks), + attributes, + dimension_names: None, + native_dtype: " Map { + value.as_object().unwrap().clone() + } + + fn named_dataset( + names: &[&str], + coordinate_attributes: Vec>, + ) -> Dataset { + let names = names + .iter() + .map(|name| (*name).to_string()) + .collect::>(); + let value_node = array_node( + vec![2; names.len()], + vec![1; names.len()], + attributes(serde_json::json!({ "_ARRAY_DIMENSIONS": names })), + ); + let coordinate_nodes = coordinate_attributes + .into_iter() + .map(|attributes| array_node(vec![2], vec![1], attributes)) + .collect::>(); + named_array_dataset("nested/value", &value_node, &names, &coordinate_nodes).unwrap() + } + + #[test] + fn array_metadata_document_requires_exactly_one_format() { + assert!(matches!( + select_array_metadata_document("value", None, Some(vec![2])), + Ok(ArrayMetadataDocument::V2(bytes)) if bytes == vec![2] + )); + assert!(matches!( + select_array_metadata_document("value", Some(vec![3]), None), + Ok(ArrayMetadataDocument::V3(bytes)) if bytes == vec![3] + )); + assert!(matches!( + select_array_metadata_document("value", Some(vec![3]), Some(vec![2])), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("both zarr.json and .zarray") + )); + assert!(matches!( + select_array_metadata_document("value", None, None), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("neither zarr.json nor .zarray") + )); + } + + #[test] + fn array_ancestor_paths_include_root_and_every_parent_group() { + assert!(array_ancestor_paths("").is_empty()); + assert_eq!(array_ancestor_paths("value"), vec![""]); + assert_eq!( + array_ancestor_paths("outer/inner/value"), + vec!["", "outer", "outer/inner"] + ); + } + + #[test] + fn scan_cursors_are_sized_to_the_array_rank() { + let [lo, hi, idx] = zeroed_scan_cursors(3); + assert_eq!(lo, vec![0; 3]); + assert_eq!(hi, vec![0; 3]); + assert_eq!(idx, vec![0; 3]); + } + + #[test] + fn spatial_time_layout_accepts_axis_orders_and_singleton_extras() { + let time_y_x = discover_spatial_time_layout( + 3, + &[ + DimensionRole::Time, + DimensionRole::SpatialY, + DimensionRole::SpatialX, + ], + &[2, 5, 6], + ) + .unwrap(); + assert_eq!(time_y_x.time, 0); + assert_eq!((time_y_x.horizontal.x, time_y_x.horizontal.y), (2, 1)); + + let x_time_y = discover_spatial_time_layout( + 3, + &[ + DimensionRole::Longitude, + DimensionRole::Time, + DimensionRole::Latitude, + ], + &[6, 2, 5], + ) + .unwrap(); + assert_eq!(x_time_y.time, 1); + assert_eq!((x_time_y.horizontal.x, x_time_y.horizontal.y), (0, 2)); + + let with_singleton_band = discover_spatial_time_layout( + 4, + &[ + DimensionRole::Band, + DimensionRole::SpatialY, + DimensionRole::SpatialX, + DimensionRole::Time, + ], + &[1, 5, 6, 2], + ) + .unwrap(); + assert_eq!(with_singleton_band.time, 3); + assert_eq!( + ( + with_singleton_band.horizontal.x, + with_singleton_band.horizontal.y, + ), + (2, 1) + ); + } + + #[test] + fn spatial_time_layout_rejects_invalid_rank_roles_and_extras() { + assert!( + discover_spatial_time_layout( + 2, + &[DimensionRole::SpatialY, DimensionRole::SpatialX], + &[5, 6], + ) + .is_err() + ); + assert!( + discover_spatial_time_layout( + 4, + &[ + DimensionRole::Band, + DimensionRole::SpatialY, + DimensionRole::SpatialX, + DimensionRole::Time, + ], + &[2, 5, 6, 2], + ) + .is_err() + ); + assert!( + discover_spatial_time_layout( + 4, + &[ + DimensionRole::Time, + DimensionRole::SpatialY, + DimensionRole::SpatialX, + DimensionRole::Time, + ], + &[2, 5, 6, 1], + ) + .is_err() + ); + } + + #[test] + fn spatial_time_range_is_exact_for_unordered_duplicate_coordinates() { + let spec = TimeSpec::default(); + let start = spec.raw_to_pg_micros(1.0).unwrap(); + let end = spec.raw_to_pg_micros(3.0).unwrap(); + let values = [2.0, 0.0, 2.0, 1.0, 3.0]; + let selected = values + .into_iter() + .enumerate() + .filter_map(|(index, raw)| { + spatial_time_value_in_range(spec, raw, start, end) + .unwrap() + .then_some(index) + }) + .collect::>(); + assert_eq!(selected, vec![0, 2, 3]); + } + + #[test] + fn missing_chunk_repeats_typed_fill_or_rejects_null() { + assert_eq!( + filled_chunk_bytes(Some(&[0x00, 0x40]), 3, "nested/raw/0.0").unwrap(), + vec![0x00, 0x40, 0x00, 0x40, 0x00, 0x40] + ); + assert!(matches!( + filled_chunk_bytes(None, 3, "nested/raw/0.0"), + Err(ZarrFdwError::MissingChunkWithoutFillValue { key }) + if key == "nested/raw/0.0" + )); + } + + #[test] + fn missing_coordinate_chunk_uses_fill_or_rejects_null() { + assert_eq!( + filled_coordinate_values(Some(4.5), 3, "nested/x/0", "x").unwrap(), + vec![4.5; 3] + ); + let err = filled_coordinate_values(None, 3, "nested/x/0", "x").unwrap_err(); + assert_eq!( + err.to_string(), + "failed to read coordinate 'x': zarr chunk 'nested/x/0' is absent and fill_value is null, so its contents are undefined" + ); + + let non_finite = filled_coordinate_values(Some(f64::NAN), 2, "nested/x/0", "x").unwrap(); + assert!(validate_coordinate_values("x", &non_finite).is_err()); + } + + #[test] + fn chunk_layout_uses_checked_bounded_arithmetic() { + let normal = array_meta(vec![2, 3, 4], vec![2, 3, 4]); + let (shape, cells, bytes) = checked_chunk_layout(&normal, 4).unwrap(); + assert_eq!(shape, vec![2, 3, 4]); + assert_eq!(cells, 24); + assert_eq!(bytes, 96); + + let too_large = array_meta( + vec![(MAX_DECODED_CHUNK_BYTES / 8 + 1) as u64, 1], + vec![(MAX_DECODED_CHUNK_BYTES / 8 + 1) as u64, 1], + ); + assert!(matches!( + checked_chunk_layout(&too_large, 8), + Err(ZarrFdwError::InvalidMetadata(message)) + if message.contains("exceeding the safety limit") + )); + assert!(filled_chunk_bytes(Some(&[0]), MAX_DECODED_CHUNK_BYTES + 1, "oversized").is_err()); + assert!( + filled_coordinate_values(Some(0.0), MAX_COORDINATE_VALUES + 1, "x/0", "x").is_err() + ); + } + + #[test] + fn flat_chunk_offsets_and_byte_ranges_are_checked() { + assert_eq!(checked_flat_offset(&[1, 2, 3], &[2, 3, 4]).unwrap(), 23); + assert!(checked_flat_offset(&[2, 0, 0], &[2, 3, 4]).is_err()); + assert!(checked_flat_offset(&[0, 0], &[usize::MAX, 2]).is_err()); + assert_eq!(checked_chunk_byte_range(23, 4, 96).unwrap(), 92..96); + assert!(checked_chunk_byte_range(24, 4, 96).is_err()); + assert!(checked_chunk_byte_range(usize::MAX, 8, usize::MAX).is_err()); + } + + #[test] + fn decoded_chunks_must_have_exact_declared_length() { + require_exact_decoded_len("0.0", 96, 96).unwrap(); + for actual in [95, 97] { + assert!(matches!( + require_exact_decoded_len("0.0", actual, 96), + Err(ZarrFdwError::ReadError(_)) + )); + } + } + + #[test] + fn value_dtypes_map_to_exact_postgres_types() { + let cases = [ + (DType::F32, pg_sys::FLOAT4OID, "real"), + (DType::F64, pg_sys::FLOAT8OID, "double precision"), + (DType::I8, pg_sys::CHAROID, r#""char""#), + (DType::I16, pg_sys::INT2OID, "smallint"), + (DType::I32, pg_sys::INT4OID, "integer"), + (DType::I64, pg_sys::INT8OID, "bigint"), + ]; + for (dtype, oid, name) in cases { + assert_eq!(expected_value_pg_type(dtype, false), (oid, name)); + assert_eq!( + expected_value_pg_type(dtype, true), + (pg_sys::FLOAT8OID, "double precision") + ); + } + } + + #[test] + fn planner_estimate_is_positive_bounded_and_uses_projected_types() { + let columns = vec![ + column("forecast_time", pg_sys::TIMESTAMPTZOID), + column("easting", pg_sys::FLOAT8OID), + column("value", pg_sys::FLOAT4OID), + ]; + assert_eq!(conservative_rel_size(&columns), (1_000_000, 20)); + assert_eq!(conservative_rel_size(&[]), (1_000_000, 8)); + assert_eq!( + conservative_rel_size(&[column("unknown", pg_sys::TEXTOID)]), + (1_000_000, 32) + ); + } + + #[test] + fn accepts_exact_coordinate_and_value_column_types() { + let dataset = named_dataset( + &["forecast_time", "level", "band", "channel"], + vec![ + attributes(serde_json::json!({"standard_name": "time"})), + attributes(serde_json::json!({"axis": "Z"})), + Map::new(), + Map::new(), + ], + ); + let columns = vec![ + column("forecast_time", pg_sys::TIMESTAMPTZOID), + column("level", pg_sys::FLOAT8OID), + column("band", pg_sys::FLOAT8OID), + column("channel", pg_sys::FLOAT8OID), + column("value", pg_sys::FLOAT4OID), + ]; + validate_column_types(&columns, &dataset, DType::F32, false).unwrap(); + } + + #[test] + fn coordinates_do_not_need_to_be_projected() { + let dataset = named_dataset(&["latitude", "longitude"], vec![Map::new(), Map::new()]); + validate_column_types( + &[column("value", pg_sys::FLOAT4OID)], + &dataset, + DType::F32, + false, + ) + .unwrap(); + validate_column_types( + &[ + column("longitude", pg_sys::FLOAT8OID), + column("value", pg_sys::FLOAT4OID), + ], + &dataset, + DType::F32, + false, + ) + .unwrap(); + } + + #[test] + fn rejects_incompatible_discovered_coordinate_type() { + let dataset = named_dataset(&["level"], vec![Map::new()]); + assert!(matches!( + validate_column_types( + &[column("level", pg_sys::INT4OID)], + &dataset, + DType::F32, + false + ), + Err(ZarrFdwError::ColumnTypeMismatch { column, .. }) if column == "level" + )); + } + + #[test] + fn rejects_incompatible_discovered_time_type() { + let dataset = named_dataset( + &["forecast_time"], + vec![attributes(serde_json::json!({"axis": "T"}))], + ); + assert!(matches!( + validate_column_types( + &[column("forecast_time", pg_sys::TIMESTAMPOID)], + &dataset, + DType::F32, + false + ), + Err(ZarrFdwError::ColumnTypeMismatch { column, .. }) if column == "forecast_time" + )); + } + + #[test] + fn arbitrary_non_dimension_name_is_a_value_column() { + let dataset = named_dataset(&["latitude", "longitude"], vec![Map::new(), Map::new()]); + validate_column_types( + &[column("time", pg_sys::FLOAT4OID)], + &dataset, + DType::F32, + false, + ) + .unwrap(); + } + + #[test] + fn cf_decoding_requires_double_precision_value_columns() { + let dataset = named_dataset(&["latitude", "longitude"], vec![Map::new(), Map::new()]); + validate_column_types( + &[column("value", pg_sys::FLOAT8OID)], + &dataset, + DType::F32, + true, + ) + .unwrap(); + assert!(matches!( + validate_column_types( + &[column("value", pg_sys::FLOAT4OID)], + &dataset, + DType::F32, + true, + ), + Err(ZarrFdwError::ColumnTypeMismatch { column, .. }) if column == "value" + )); + } + + #[test] + fn coordinate_values_must_be_finite_but_may_be_unordered() { + for finite in [ + vec![0.0, 1.0, 2.0], + vec![2.0, 1.0, 0.0], + vec![1.0, 1.0, 2.0], + vec![1.0, 1.0, 1.0], + vec![0.0, 2.0, 1.0], + ] { + validate_coordinate_values("x", &finite).unwrap(); + } + + for invalid in [ + vec![0.0, f64::NAN], + vec![0.0, f64::INFINITY], + vec![0.0, f64::NEG_INFINITY], + ] { + assert!(validate_coordinate_values("x", &invalid).is_err()); + } + } + + #[test] + fn coordinate_nan_qualifiers_never_narrow_chunk_ranges() { + for operator in ["=", "<", "<=", ">", ">="] { + let qual = Qual { + field: "x".to_string(), + operator: operator.to_string(), + value: Value::Cell(Cell::F64(f64::NAN)), + use_or: false, + param: None, + }; + assert_eq!( + qual_to_range(&qual, false, TimeSpec::default()).unwrap(), + None + ); + } + + let finite_and_nan = Qual { + field: "x".to_string(), + operator: "=".to_string(), + value: Value::Array(vec![Cell::F64(f64::NAN), Cell::F64(5.0)]), + use_or: true, + param: None, + }; + assert_eq!( + qual_to_range(&finite_and_nan, false, TimeSpec::default()).unwrap(), + Some((Some(5.0), Some(5.0))) + ); + } + + #[test] + fn required_coordinates_have_a_cumulative_value_budget() { + assert_eq!( + checked_total_coordinate_values([("a", 3, true), ("b", 2, false), ("c", 4, true)], 7) + .unwrap(), + 7 + ); + assert!(checked_total_coordinate_values([("a", 4, true), ("b", 4, true)], 7).is_err()); + } + + #[test] + fn required_coordinates_reject_unsupported_scientific_decoding() { + for attribute in UNSUPPORTED_COORDINATE_DECODING_ATTRIBUTES { + let mut attributes = Map::new(); + attributes.insert(attribute.to_string(), serde_json::json!(1)); + assert!(validate_coordinate_decoding_attributes("level", &attributes).is_err()); + } + validate_coordinate_decoding_attributes("level", &Map::new()).unwrap(); + } + + #[test] + fn scalable_execution_options_are_bounded_and_cache_can_be_disabled() { + assert_eq!( + scalable_execution_options(&HashMap::new()).unwrap(), + ( + DEFAULT_MAX_CONCURRENT_READS, + DEFAULT_MAX_INFLIGHT_BYTES, + DEFAULT_COMPRESSED_CACHE_BYTES, + ) + ); + let valid = HashMap::from([ + (OPT_MAX_CONCURRENT_READS.to_string(), "32".to_string()), + ( + OPT_MAX_INFLIGHT_BYTES.to_string(), + MIN_MAX_INFLIGHT_BYTES.to_string(), + ), + (OPT_COMPRESSED_CACHE_BYTES.to_string(), "0".to_string()), + ]); + assert_eq!( + scalable_execution_options(&valid).unwrap(), + (32, MIN_MAX_INFLIGHT_BYTES, 0) + ); + + for invalid in [ + HashMap::from([(OPT_MAX_CONCURRENT_READS.to_string(), "0".to_string())]), + HashMap::from([( + OPT_MAX_INFLIGHT_BYTES.to_string(), + (MIN_MAX_INFLIGHT_BYTES - 1).to_string(), + )]), + HashMap::from([( + OPT_COMPRESSED_CACHE_BYTES.to_string(), + (MAX_COMPRESSED_CACHE_BYTES + 1).to_string(), + )]), + ] { + assert!(scalable_execution_options(&invalid).is_err()); + } + } + + #[test] + fn remote_io_metrics_count_polled_and_completed_work_only() { + let calls = Arc::new(AtomicU64::new(0)); + let bytes = Arc::new(AtomicU64::new(0)); + let unpolled = observe_data_fetch( + async { Ok::<_, ZarrFdwError>(Some(vec![1, 2, 3])) }, + Arc::clone(&calls), + Arc::clone(&bytes), + ); + drop(unpolled); + assert_eq!(calls.load(Ordering::Relaxed), 0); + assert_eq!(bytes.load(Ordering::Relaxed), 0); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let object = runtime + .block_on(observe_data_fetch( + async { Ok::<_, ZarrFdwError>(Some(vec![1, 2, 3, 4])) }, + Arc::clone(&calls), + Arc::clone(&bytes), + )) + .unwrap(); + assert_eq!(object, Some(vec![1, 2, 3, 4])); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(bytes.load(Ordering::Relaxed), 4); + } +}