diff --git a/Cargo.lock b/Cargo.lock index b43b24f07..3c6878b1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4044,6 +4044,7 @@ dependencies = [ "nodedb-codec", "nodedb-types", "nodedb-wal", + "rustc-hash", "serde", "serde_json", "thiserror 2.0.21", @@ -4095,6 +4096,8 @@ dependencies = [ "serde_json", "sonic-rs", "tempfile", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "tokio", "tokio-postgres", ] diff --git a/README.md b/README.md index a1e77b106..352e51771 100644 --- a/README.md +++ b/README.md @@ -118,8 +118,8 @@ cargo install nodedb Requires Linux kernel >= 5.1 (io_uring). Connect: ```bash -ndb # native CLI (connects to localhost:6433) -psql -h localhost -p 6432 # or any PostgreSQL client +ndb -w # native CLI as user nodedb on localhost:6433, prompts for the password (or set NDB_PASSWORD) +psql -h localhost -p 6432 -U nodedb -d default # password: /.superuser_password, or set NODEDB_SUPERUSER_PASSWORD ``` ```sql @@ -164,7 +164,7 @@ See [NodeDB-Lite](https://github.com/NodeDB-Lab/nodedb-lite) for platform detail **Write-time validation** -- Typeguards enforce types, required fields, CHECK constraints, and DEFAULT/VALUE expressions on schemaless collections. Graduate to strict schema with `CONVERT COLLECTION x TO document_strict`. -**Bitemporal queries** -- System time (audit trail) and valid time (temporal semantics) across all engines. Query data as it existed in the past, or as it was valid on a past date. GDPR-compliant tile purge on array engine. +**Bitemporal queries** -- System time (audit trail) and valid time (temporal semantics) on Document, Columnar, Timeseries, Graph and Array. Vector, Full-Text, Spatial and KV indexes reach time through the collection they index. Query data as it existed in the past, or as it was valid on a past date. GDPR-compliant tile purge on array engine. **Multi-dimensional arrays** -- Scientific computing with Z-order indexed tiles, tile-level compression, and bitemporal support. Combine with vector/graph/text in fused queries via cross-engine identity. @@ -214,11 +214,9 @@ cargo nextest run --all-features ## Release Status -NodeDB Origin is in **public beta**; the current release is **v0.5.0 (2026-08-04)**. All eight engines are feature-complete and covered by tests. The query-facing wire protocols (pgwire, HTTP, native MessagePack, RESP) are stable — clients written against 0.1.0 will keep working through 1.0. The CRDT sync protocol and ILP ingest are still moving: 0.5.0 changes the sync frame and requires authentication on ILP connections, so upgrade NodeDB Lite and ingest clients together with the server. +NodeDB Origin is in **public beta**; the current release is **v0.6.2 (2026-10-16)**. All eight engines are feature-complete and covered by tests. Server binaries, Docker images and the crates.io crates ship at the same version. The CRDT sync protocol and the native protocol still change between minor releases, so upgrade NodeDB Lite and `nodedb-client` together with the server. -v0.5.0 ships as server binaries and Docker images; the crates.io publication follows separately, so `cargo add` still resolves 0.4.0 for now. Library users who need the 0.5.0 fixes today can depend on the `v0.5.0` tag by git. - -**v0.5.0 — Beta (today).** Build new products on it. The public surface (SQL dialect, query wire protocols, configuration) is stable; expect internal changes (storage layout, on-disk format, replication internals) between minor releases — a 0.4.0 data directory needs a dump/reload to move to 0.5.0. Patch and minor bumps land as needed — 0.2 through 0.5 shipped as the work demanded. The beta period is deliberately for real workloads to surface edge cases we can't manufacture in-house. +**v0.6 — Beta (today).** Build new products on it. Expect internal changes (storage layout, on-disk format, replication internals) between minor releases: a 0.5.0 data directory cannot be opened by 0.6, so start 0.6 on an empty directory. Patch and minor bumps land as needed. The beta period is deliberately for real workloads to surface edge cases we can't manufacture in-house. **v1.0.0 — Production-ready (target: after the beta hardening period).** What 1.0 guarantees: diff --git a/docs/array.md b/docs/array.md index 77ff4be44..dd5777d26 100644 --- a/docs/array.md +++ b/docs/array.md @@ -50,14 +50,38 @@ CREATE ARRAY spatial_grid | Parameter | Required | Default | Description | | ----------------- | -------- | ----------- | ------------------------------------------------------------------------------------------------------------------ | | `DIMS` | Yes | — | List of dimensions. Each has a name, type (`INT64`, `FLOAT64`, `TIMESTAMP_MS`, `STRING`), and optional domain bounds `[lo..hi]`. | -| `ATTRS` | Yes | — | List of attributes (cell values). Each has a name and type (`INT64`, `FLOAT64`, `STRING`, `BYTES`). | +| `ATTRS` | Yes | — | List of attributes (cell values). Each has a name, a type (see [Attribute Types](#attribute-types)), and an optional `NOT NULL`. | | `TILE_EXTENTS` | Yes | — | Tuple of tile extent per dimension. All > 0. Determines cell locality and compression block granularity. | -| `CELL_ORDER` | No | `Z_ORDER` | `ROW_MAJOR`, `COL_MAJOR`, `Z_ORDER`, or `HILBERT`. Affects spatial cache locality. | -| `TILE_ORDER` | No | — | `ROW_MAJOR`, `COL_MAJOR`, `Z_ORDER`, or `HILBERT`. Optional tile ordering. | +| `CELL_ORDER` | No | `HILBERT` | Unquoted: `ROW_MAJOR`, `COL_MAJOR`, `HILBERT`, or `ZORDER` (`Z_ORDER` is accepted). Affects spatial cache locality. | +| `TILE_ORDER` | No | `HILBERT` | Unquoted: `ROW_MAJOR`, `COL_MAJOR`, `HILBERT`, or `ZORDER` (`Z_ORDER` is accepted). Affects tile layout. | | `prefix_bits` | No | `8` | Bits for key prefix compression (1–16). Lower = better compression, higher = faster prefix filtering. | | `audit_retain_ms` | No | `NULL` | Milliseconds. Tiles older than now - `audit_retain_ms` (by system time) are eligible for purge. `NULL` = keep all. | | `minimum_audit_retain_ms` | No | — | Minimum retention window enforced for compliance. | +### Attribute Types + +An attribute stores each cell at its declared width. A value outside the width is refused with SQLSTATE `22003` on `INSERT`, `COPY`, and the native bulk load. It is never wrapped or clamped. + +| Type | Aliases | Range | Column type (OID) | +| --------- | -------------------------- | ---------------------------------- | ------------------------- | +| `UINT8` | | 0 to 255 | `smallint` (21) | +| `INT16` | `SMALLINT`, `INT2` | -32768 to 32767 | `smallint` (21) | +| `UINT16` | | 0 to 65535 | `integer` (23) | +| `INT32` | `INTEGER`, `INT`, `INT4` | -2147483648 to 2147483647 | `integer` (23) | +| `UINT32` | | 0 to 4294967295 | `bigint` (20) | +| `INT64` | `BIGINT` | 64-bit signed | `bigint` (20) | +| `FLOAT32` | `REAL`, `FLOAT4` | 32-bit IEEE float | `real` (700) | +| `FLOAT64` | `DOUBLE`, `FLOAT` | 64-bit IEEE float | `double precision` (701) | +| `STRING` | `TEXT`, `VARCHAR` | UTF-8 text | `text` (25) | +| `BYTES` | `BLOB`, `BYTEA` | raw bytes | `bytea` (17) | + +- An unsigned type reads as the next wider signed column type, so every stored value fits. +- `INT` is the 4-byte `INT32`, as in PostgreSQL. A bare `FLOAT` is the 8-byte `FLOAT64`. +- A float written to a `FLOAT32` attribute rounds to the nearest `f32`. Only a finite value past the `f32` range is refused. `0.1` stored as `FLOAT32` reads back as `0.1`. +- An integer literal is accepted by every integer and float type. A float literal is accepted by `FLOAT32` and `FLOAT64` only. +- An integer column compresses by its value range, so a `UINT16` column takes the same bytes on disk as an `INT64` column of the same values. A `FLOAT32` column stores at most 4 bytes per value, half the raw width of `FLOAT64`. +- A `NULL` stays `NULL` through a flush and a compaction, in every type. + ## Examples ### Basic Creation and Insertion @@ -74,12 +98,18 @@ CREATE ARRAY elevation_map ) TILE_EXTENTS (256, 256); --- Insert cells (or rows of cells) --- Cells are written as objects with dim names as keys -INSERT INTO elevation_map (lon, lat, height) VALUES - (-73.5, 40.7, 10.5), - (-73.6, 40.8, 12.3), - (-73.7, 40.6, 8.9); +-- Insert cells: COORDS lists the dims in schema order, VALUES lists the attrs in schema order. +-- Values are literals: numbers, strings, or NULL for an attribute. +INSERT INTO ARRAY elevation_map + COORDS (-73.5, 40.7) VALUES (10.5), + COORDS (-73.6, 40.8) VALUES (12.3), + COORDS (-73.7, 40.6) VALUES (8.9); + +-- Remove cells by coordinate +DELETE FROM ARRAY elevation_map WHERE COORDS IN ((-73.7, 40.6)); + +-- Writing a coordinate again overwrites its cell. +-- DROP ARRAY [IF EXISTS] removes the whole array. -- Flush memtable to persistent storage SELECT ARRAY_FLUSH('elevation_map'); @@ -103,13 +133,13 @@ CREATE ARRAY climate_forecast WITH (audit_retain_ms = 7776000000); -- 90 days -- Insert forecast data at a specific valid time -INSERT INTO climate_forecast (lon, lat, altitude_m, temp_c, humidity) VALUES - (10, 20, 5000, -10.5, 65.0), - (10, 20, 10000, -20.3, 40.0); +INSERT INTO ARRAY climate_forecast + COORDS (10, 20, 5000) VALUES (-10.5, 65.0), + COORDS (10, 20, 10000) VALUES (-20.3, 40.0); -- Query at a specific moment in time (system time) SELECT lon, lat, altitude_m, temp_c -FROM ARRAY_SLICE('climate_forecast', {lon: [10, 15), lat: [20, 25), altitude_m: [5000, 15000)}, ['temp_c']) +FROM ARRAY_SLICE('climate_forecast', '{lon: [10, 14], lat: [20, 24], altitude_m: [5000, 14999]}', ['temp_c']) AS OF SYSTEM TIME 1700000000000; ``` @@ -124,7 +154,7 @@ Returns cells within a multi-dimensional range: ```sql SELECT * FROM ARRAY_SLICE( 'elevation_map', - {lon: [-74.0, -73.0), lat: [40.0, 41.0)}, + '{lon: [-74.0, -73.0], lat: [40.0, 41.0]}', ['height'], -- projecting only 'height' (optional) 1000 -- limit to 1000 cells (optional) ); @@ -132,14 +162,62 @@ SELECT * FROM ARRAY_SLICE( | Parameter | Required | Type | Description | | --------- | -------- | ------------- | --------------------------------------------------------------- | -| `array` | Yes | STRING | Array name | -| `bounds` | Yes | OBJECT | Dict of dim name → `[lo, hi)` bounds. Omitted dims = full range | -| `attrs` | No | ARRAY[STRING] | Attributes to project. `NULL` = all attributes | -| `limit` | No | INT64 | Max cells returned. `NULL` = no limit | +| `array` | Yes | STRING | Array name, as a string literal | +| `bounds` | Yes | STRING | Quoted object literal `'{ dim: [lo, hi] }'`. Each range is closed: both ends are included. Omitted dims = full range | +| `attrs` | No | ARRAY[STRING] | Attributes to project, such as `['height']`. `'*'` = all attributes. Omit it to read all attributes | +| `limit` | No | INT | Max cells returned, as a non-negative integer literal. Omit it for no limit | + +#### Result Relation + +`ARRAY_SLICE` and `ARRAY_PROJECT` return one row per cell. The row has one column per dimension, then one column per requested attribute, or every attribute when the list is empty. Each column carries its declared type: + +| Declared type | Column type | +| ------------- | ------------------ | +| `UINT8`, `INT16` (attribute only) | `smallint` | +| `UINT16`, `INT32` (attribute only) | `integer` | +| `UINT32` (attribute only) | `bigint` | +| `INT64` | `bigint` | +| `FLOAT32` (attribute only) | `real` | +| `FLOAT64` | `double precision` | +| `TIMESTAMP_MS` (dimension) | `timestamp` | +| `STRING` | `text` | +| `BYTES` | `bytea` (attribute only) | + +`SELECT *` lists the dimensions first, then the attributes, in schema order. For a 2-D image with `nir` and `red` bands: + +```sql +SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3], y: [0, 3]}'); +-- x | y | nir | red +-- ---+---+------+------ +-- 0 | 0 | 0.80 | 0.20 +-- 0 | 1 | 0.60 | 0.30 +``` + +Every SELECT clause applies to the rows: `WHERE`, computed columns, `GROUP BY`, aggregates, `HAVING`, `ORDER BY`, `DISTINCT`, `LIMIT` / `OFFSET`, and window functions. Qualify a column with the call's alias, or with the array name when the call has no alias. + +```sql +-- NDVI per cell, vegetated cells only, greenest first +SELECT x, y, (nir - red) / (nir + red) AS ndvi +FROM ARRAY_SLICE('img', '{x: [0, 3], y: [0, 3]}') AS t +WHERE t.nir > t.red +ORDER BY ndvi DESC +LIMIT 10; + +-- Mean NIR per column of pixels +SELECT x, AVG(nir) AS mean_nir +FROM ARRAY_SLICE('img', '{x: [0, 3], y: [0, 3]}') +GROUP BY x; +``` + +The 4th-argument `limit` caps the cells the slice reads before any clause applies. An outer `LIMIT` applies to the query's result, and a smaller one also lowers the read cap when no other clause reads past the cut. A `WHERE` bound on a dimension narrows the read range only when the slice has no cap, so a capped slice keeps the cells its cap selects. + +`AS OF SYSTEM TIME NULL` adds a `_ts_system` column: the version's system time in milliseconds. `SELECT *` omits it. Name it to read it. + +`ARRAY_AGG` and `ARRAY_ELEMENTWISE` return rows of their own shape, not dimensions and attributes. A query over them takes `*` or their own output columns, and `LIMIT`. Every other clause is refused with an error that names the clause. Use `ARRAY_SLICE` to filter, group, or order cells. ### ARRAY_PROJECT — Select Attributes -Returns all cells, optionally filtered to specific attributes: +Returns all cells with only the listed attributes. The attribute list is required and must not be empty. It reads the whole domain, as an `ARRAY_SLICE` with no bounds does, and honors `AS OF` the same way: ```sql SELECT * FROM ARRAY_PROJECT( @@ -148,36 +226,73 @@ SELECT * FROM ARRAY_PROJECT( ); ``` -### ARRAY_AGG — Aggregate Over Dimensions +### ARRAY_AGG — Aggregate an Attribute + +Reduces one attribute with a reducer. The optional 4th argument names a dimension to group by: + +```sql +-- One result for the whole array +SELECT * FROM ARRAY_AGG('spatial_grid', 'temperature', 'SUM'); + +-- One result per value of x +SELECT * FROM ARRAY_AGG('spatial_grid', 'temperature', 'SUM', 'x'); +``` + +The result has a `result` column. A grouped call also has a `group` column that holds the dimension value. The `group` column takes the group dimension's declared type: `TIMESTAMP_MS` gives `timestamp`, `FLOAT64` gives `float8`, `INT64` gives `int8`, and `STRING` gives `text`. A temporal call (`AS OF`) appends one summary row with the `bool` column `truncated_before_horizon`. Reducers are case-insensitive: `'SUM'`, `'COUNT'`, `'MIN'`, `'MAX'`, and `'MEAN'` (or `'AVG'`). + +The `result` type follows the reducer and the attribute type: + +| Reducer | Integer attribute | Float attribute | +| ------------- | -------------------------------------- | ------------------ | +| `COUNT` | `bigint` | `bigint` | +| `SUM` | `bigint`, exact | `double precision` | +| `MIN` / `MAX` | the attribute's own column type, exact | the attribute's own column type | +| `MEAN` | `double precision` | `double precision` | + +`SUM`, `MEAN`, `MIN`, and `MAX` accept integer and float attributes only. `COUNT` accepts every attribute type. Another reducer over a `STRING` or `BYTES` attribute is refused at plan time with a type mismatch error. + +An integer `SUM` adds exactly, so `2^53 + 1` stays `9007199254740993`. A `SUM` past the `bigint` range is refused with SQLSTATE `22003`. `MIN` and `MAX` sort `NaN` above every other float, as PostgreSQL does. -Aggregates an attribute over a dimension, reducing dimensionality: +Three more optional arguments bound the cells the reducer reads: + +- **5th, box**: a slice predicate in the `ARRAY_SLICE` form, such as `'{lon: [101, 103], lat: [2, 4]}'`. Bounds are inclusive. +- **6th and 7th, zone**: a WKT `POLYGON` or `MULTIPOLYGON`, then the two dimensions it is drawn over, x first, such as `ARRAY['lon', 'lat']`. Both dimensions must be `INT64` or `FLOAT64`. Pass both arguments or neither. + +Pass `NULL` for an argument to skip it: ```sql --- Sum temperature over the x dimension, keeping y and z +-- Mean NDVI inside a field boundary SELECT * FROM ARRAY_AGG( - 'spatial_grid', - 'temperature', - 'SUM', - 'x' -- aggregate over x; result has one less dimension + 'scene', 'ndvi', 'MEAN', NULL, NULL, + 'POLYGON((101 2, 103 2, 103 4, 101 4, 101 2))', ARRAY['lon', 'lat'] ); + +-- Sum per lon value inside a box +SELECT * FROM ARRAY_AGG('scene', 'ndvi', 'SUM', 'lon', '{lon: [101, 103]}'); ``` -Supported reducers: `'SUM'`, `'AVG'`, `'MIN'`, `'MAX'`, `'COUNT'`. +A cell is in a zone when it lies inside one member polygon and outside that polygon's holes. Edges are half-open: on a rectangle, a cell on the lower x or y edge is inside and a cell on the upper edge is outside. A region that holds no cell returns `NULL` for every reducer, `COUNT` included. Tiles whose bounding boxes miss the region are skipped before decode. ### ARRAY_ELEMENTWISE — Element-wise Operations -Applies an operation between two arrays (or array and scalar) with the same shape: +Applies an operation between two arrays of the same shape. Both arrays must have the same dimensions, attributes, tile extents, cell order, and `prefix_bits`. Arrays that differ are refused at planning: ```sql -- Subtract a baseline from all cells SELECT * FROM ARRAY_ELEMENTWISE( 'current_grid', 'baseline_grid', - 'SUBTRACT', - 'temperature' -- the attribute to operate on + 'sub', -- 'add', 'sub', 'mul', or 'div' (case-insensitive) + 'temperature' -- the attribute to operate on ); ``` +The result has `coords` and `attrs` columns, both JSON. `attrs` holds one value: the operation applied to the named attribute. A coordinate present in one array only takes that array's value. + +On integer attributes, `add`, `sub`, and `mul` are exact 64-bit integers, and a result past the `bigint` range is refused with SQLSTATE `22003`. `div`, and every operation on a float attribute, yields a 64-bit float. Division by zero yields `NULL`. + +`AS OF SYSTEM TIME` and `AS OF VALID TIME` apply to both arrays. A temporal query appends one `{"truncated_before_horizon": bool}` row, as `ARRAY_AGG` does. `AS OF SYSTEM TIME NULL` is refused. + ## Maintenance Functions ### ARRAY_FLUSH @@ -185,22 +300,22 @@ SELECT * FROM ARRAY_ELEMENTWISE( Forces in-memory cells to durable storage: ```sql -SELECT ARRAY_FLUSH('spatial_grid') AS result; --- Returns: {result: true} +SELECT ARRAY_FLUSH('spatial_grid'); +-- Returns: {"flushed": 1} ``` -Returns a single row `{result: BOOL}`. Always returns `true`; failure is fatal and raises an error. +Returns one object `{"flushed": 1}` on success. A failure raises an error. ### ARRAY_COMPACT Merges tile versions and reclaims space: ```sql -SELECT ARRAY_COMPACT('spatial_grid') AS result; --- Returns: {result: true} +SELECT ARRAY_COMPACT('spatial_grid'); +-- Returns: {"compacted": 1} ``` -Compaction is background-automatic, but can be triggered manually. +Returns one object `{"compacted": n}`. `n` counts the cores that merged segments. It is `0` when the compaction picker found nothing to merge. A failure raises an error. Compaction is background-automatic, but can be triggered manually. ## Bitemporal Support @@ -215,7 +330,7 @@ Query as of either or both: -- Read cells as they existed at a point in the past SELECT * FROM ARRAY_SLICE( 'data', - {x: [0, 100), y: [0, 100)}, + '{x: [0, 99], y: [0, 99]}', ['value'] ) AS OF SYSTEM TIME 1700000000000; @@ -223,7 +338,7 @@ AS OF SYSTEM TIME 1700000000000; -- Read cells that were valid at a specific time SELECT * FROM ARRAY_SLICE( 'forecast', - {x: [0, 100), y: [0, 100)}, + '{x: [0, 99], y: [0, 99]}', ['temp'] ) AS OF VALID TIME 1700000000000; @@ -231,23 +346,23 @@ AS OF VALID TIME 1700000000000; -- Read cells that were valid AND existed at a specific time SELECT * FROM ARRAY_SLICE( 'forecast', - {x: [0, 100), y: [0, 100)}, + '{x: [0, 99], y: [0, 99]}', ['temp'] ) AS OF SYSTEM TIME 1700000000000 AS OF VALID TIME 1700000001000; ``` +A time is an integer in milliseconds since the Unix epoch, `NOW()`, or an ISO-8601 string such as `'2024-01-15T00:00:00Z'`. + ## Cross-Engine Integration -Array cells are addressable via surrogate identity alongside other engines. Combine array queries with vector search, graph traversal, and full-text search: +Array cells are addressable via surrogate identity alongside other engines. An `ARRAY_SLICE` or `ARRAY_PROJECT` call can be a `JOIN` source. A slice row holds one field per dimension and attribute, keyed by name, so a join key reads them directly. A slice row has no `id` column. ```sql --- Find cells near a vector embedding, return array slice -SELECT * -FROM ARRAY_SLICE('spatial_data', slice_bounds, ['attr1', 'attr2']) -WHERE id IN ( - SEARCH vectors USING VECTOR(embedding, query_vec, 100) -); +-- Join slice cells to a collection on a dimension value +SELECT c.x, c.y, c.attr1, s.name +FROM ARRAY_SLICE('spatial_data', '{x: [0, 999], y: [0, 999]}', ['attr1']) AS c +JOIN sites s ON c.x = s.grid_x; ``` See [Architecture — Cross-engine identity](architecture.md#cross-engine-identity) for details. diff --git a/docs/bitemporal.md b/docs/bitemporal.md index 3bd2faf5f..8023a3526 100644 --- a/docs/bitemporal.md +++ b/docs/bitemporal.md @@ -115,11 +115,11 @@ Retrieve array cells as they existed at a point in time: -- Multidimensional scientific data with temporal audit trail CREATE ARRAY climate_grid DIMS ( - lon INT32 DOMAIN [-180, 180), - lat INT32 DOMAIN [-90, 90) + lon INT64 [-180..180], + lat INT64 [-90..90] ) ATTRS ( - temp_c FLOAT32 + temp_c FLOAT64 ) TILE_EXTENTS (64, 64) WITH (audit_retain_ms = 7776000000); -- 90 days @@ -127,7 +127,7 @@ CREATE ARRAY climate_grid -- Query cells as they were committed yesterday SELECT lon, lat, temp_c FROM ARRAY_SLICE( 'climate_grid', - {lon: [-10, 10), lat: [0, 20)}, + '{lon: [-10, 9], lat: [0, 19]}', ['temp_c'] ) AS OF SYSTEM TIME '2026-06-06T00:00:00Z'; @@ -162,7 +162,7 @@ AS OF SYSTEM TIME NULL -- returns every system-time version, ascending ORDER BY _ts_system ASC; ``` -`AS OF SYSTEM TIME NULL` returns every system-time version of each matching row, ordered ascending. Each version row carries the user columns plus three synthetic temporal columns: `_ts_system`, `_ts_valid_from`, and `_ts_valid_until`. Unbounded valid-time ends surface as the raw `i64::MIN` / `i64::MAX` sentinels (matching the Columnar and Timeseries engines). It is supported on the Document (strict and schemaless), Columnar, and Timeseries engines. It is not supported on the Graph or Array engines, nor when reading through a database clone — those return a typed error rather than collapsing to a single version. +`AS OF SYSTEM TIME NULL` returns every system-time version of each matching row, ordered ascending. Each version row carries the user columns plus three synthetic temporal columns: `_ts_system`, `_ts_valid_from`, and `_ts_valid_until`. Unbounded valid-time ends surface as the raw `i64::MIN` / `i64::MAX` sentinels (matching the Columnar and Timeseries engines). It is supported on the Document (strict and schemaless), Columnar, and Timeseries engines, and on `ARRAY_SLICE`. An `ARRAY_SLICE` version row carries the dimension and attribute columns plus `_ts_system` only. It is not supported on the Graph engine or on `ARRAY_AGG`, nor when reading through a database clone — those return a typed error rather than collapsing to a single version. `AS OF` reads run at parity with non-temporal reads: `WHERE` predicates (including on strict Binary-Tuple collections), `ORDER BY`, computed columns, and window functions all apply to temporal queries. diff --git a/docs/databases.md b/docs/databases.md index e5a17194c..f87bcaee7 100644 --- a/docs/databases.md +++ b/docs/databases.md @@ -309,7 +309,7 @@ This allows safe retry after network failures without duplicating the move. ### Per-Database Audit Filtering ```sql -SHOW AUDIT IN DATABASE staging [WHERE event_type = 'database_quota_changed']; +SHOW AUDIT IN DATABASE staging [WHERE event_type = 'database_quota_changed'] [LIMIT n]; ``` Events tagged with `database_id` include: `DatabaseCreated`, `DatabaseDropped`, `DatabaseRenamed`, `DatabaseQuotaChanged`, `DatabaseCloned`, `DatabaseMirrored`, `DatabasePromoted`, `TenantMoved`, etc. diff --git a/docs/getting-started.md b/docs/getting-started.md index b95b51d3c..f41858104 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -726,17 +726,21 @@ An env override still wins over an expanded value for the same field. ### With the `ndb` CLI ```bash -./target/release/ndb +./target/release/ndb -w ``` +`ndb` connects to `127.0.0.1:6433` as user `nodedb`. The server runs in password mode by default, and `ndb` sends no password unless given `-w`, which prompts for it. Set `NDB_PASSWORD` to skip the prompt. `NDB_USER` and `NDB_DATABASE` set the user and database. + You get a full TUI with syntax highlighting, tab completion, and history search. See the [CLI guide](cli.md) for details. ### With psql ```bash -psql -h localhost -p 6432 +psql -h localhost -p 6432 -U nodedb -d default ``` +The server starts in password mode with the superuser `nodedb`. Set `NODEDB_SUPERUSER_PASSWORD` before the first start, or read the generated password from `/.superuser_password`. `-d default` names the built-in database. Without it, psql asks for a database named after the user, and the server refuses the connection with `3D000`. + NodeDB speaks PostgreSQL's wire protocol, so standard tools like `psql`, ORMs, and BI tools work out of the box. ### With the Rust SDK or FFI diff --git a/docs/protocols.md b/docs/protocols.md index 4bc204237..5717a3bf1 100644 --- a/docs/protocols.md +++ b/docs/protocols.md @@ -21,7 +21,7 @@ NodeDB speaks six wire protocols. All SQL-capable protocols share the same query Standard PostgreSQL wire protocol. Any tool that speaks Postgres works with NodeDB. ```bash -psql -h localhost -p 6432 +psql -h localhost -p 6432 -U nodedb -d default # Or any Postgres-compatible client # JDBC, libpq, SQLAlchemy, Prisma, etc. @@ -76,11 +76,11 @@ Use SQL for complex queries, ad-hoc exploration, and rapid prototyping. Use nati **Connection:** ```bash -# ndb CLI connects automatically -./target/release/ndb +# ndb CLI: user nodedb, port 6433, -w prompts for the password (or set NDB_PASSWORD) +./target/release/ndb -w # Or specify host -./target/release/ndb --host localhost --port 6433 +./target/release/ndb --host localhost --port 6433 -w ``` ## HTTP API diff --git a/docs/query-language.md b/docs/query-language.md index 570c5db60..69118e8b4 100644 --- a/docs/query-language.md +++ b/docs/query-language.md @@ -72,6 +72,10 @@ SELECT UPPER(label), SUM(score) FROM events GROUP BY 1; GROUP BY keys follow PostgreSQL semantics: output columns appear in SELECT-list order, and a key projected as `SELECT k AS label` is emitted under its alias. +A column in the SELECT list that is neither a GROUP BY key nor inside an aggregate is refused with SQLSTATE `42803` (`grouping_error`). The message reads `column "x" must appear in the GROUP BY clause or be used in an aggregate function`. + +A planning error reaches a driver at Parse (prepare) with its SQLSTATE when the statement has no `$N` placeholders. Class `42` errors and `0A000` report at Parse. These errors defer to Execute: `42501` (insufficient privilege), data errors of class `22`, and retryable errors of class `40`. With placeholders, Parse plans each `$N` as `NULL`. It reports only these errors: `42P01` (unknown table), `42703` (unknown column), `42702` (ambiguous column), `42883` (unknown function), and `42803`. Every other error defers to Execute, which plans with the bound values. + A scalar aggregate (no GROUP BY) over zero input rows returns one identity row: `COUNT(*)` → `0`; `SUM`/`AVG`/`MIN`/`MAX` → `NULL`. A `GROUP BY` aggregate over zero rows returns zero rows. ### Joins @@ -384,7 +388,7 @@ SHOW SESSIONS [IN DATABASE prod_us]; KILL SESSION 'sid_abc123'; -- Audit events tagged by database -SHOW AUDIT IN DATABASE prod_us [WHERE event_type = 'database_quota_changed']; +SHOW AUDIT IN DATABASE prod_us [WHERE event_type = 'database_quota_changed'] [LIMIT n]; ``` All database DDL is replicated via the metadata Raft group and is consistent across the cluster. For full examples and context, see [Databases](databases.md). @@ -830,7 +834,7 @@ ORDER BY sparse_score(terms, '{3: 1.0, 7: 0.5}') DESC LIMIT 3; ``` -The inverted index is maintained automatically on document writes (insert, update, delete). Documents sharing no dimension with the query are never scanned. `sparse_score` is a standalone surface — it is not fused into `rrf_score` hybrid ranking. +The inverted index is maintained automatically on document writes (insert, update, delete). Documents sharing no dimension with the query are never scanned. A `WHERE` clause filters the hits before the top-k cut. Projected columns read from the row. `sparse_score` is a standalone surface — it is not fused into `rrf_score` hybrid ranking. ### Full-Text Search @@ -988,47 +992,73 @@ SELECT crdt_apply('collab_docs', 'doc123', ''); -- Create array with dimensions and attributes CREATE ARRAY spatial_grid DIMS ( - x INT64 DOMAIN [0, 1000), - y INT64 DOMAIN [0, 1000), - z INT64 DOMAIN [0, 1000) + x INT64 [0..1000], + y INT64 [0..1000], + z INT64 [0..1000] ) ATTRS ( - temperature FLOAT32, - pressure FLOAT32 + temperature FLOAT64, + pressure FLOAT64 ) TILE_EXTENTS (64, 64, 64); --- Insert cells -INSERT INTO spatial_grid (x, y, z, temperature, pressure) VALUES - (10, 20, 5, 23.5, 1013.2), - (10, 20, 10, 22.1, 1013.1); +-- Insert cells: COORDS lists the dims, VALUES lists the attrs +INSERT INTO ARRAY spatial_grid + COORDS (10, 20, 5) VALUES (23.5, 1013.2), + COORDS (10, 20, 10) VALUES (22.1, 1013.1); --- Slice: range query over dimensions +-- Delete cells by coordinate +DELETE FROM ARRAY spatial_grid WHERE COORDS IN ((10, 20, 10)); + +-- Slice: range query over dimensions (closed ranges, quoted object literal) SELECT x, y, z, temperature FROM ARRAY_SLICE( 'spatial_grid', - {x: [0, 100), y: [0, 100), z: [0, 50)}, + '{x: [0, 99], y: [0, 99], z: [0, 49]}', ['temperature'] ); -- Project: select specific attributes SELECT * FROM ARRAY_PROJECT('spatial_grid', ['temperature', 'pressure']); --- Aggregate: reduce dimensionality via aggregation +-- Aggregate: reduce an attribute, grouped by the dimension x SELECT * FROM ARRAY_AGG('spatial_grid', 'temperature', 'AVG', 'x'); --- Element-wise operations +-- Element-wise operations ('add', 'sub', 'mul', 'div') SELECT * FROM ARRAY_ELEMENTWISE( 'current', 'baseline', - 'SUBTRACT', + 'sub', 'temperature' ); -- Flush and compact SELECT ARRAY_FLUSH('spatial_grid'); SELECT ARRAY_COMPACT('spatial_grid'); + +-- Narrow attribute types store each cell at its declared width +CREATE ARRAY scene + DIMS (y INT64 [0..10979], x INT64 [0..10979]) + ATTRS (b04 UINT16, b08 UINT16, scl UINT8, ndvi FLOAT32) + TILE_EXTENTS (512, 512); ``` +Attribute types and the column type each reads as: + +| Attribute type | Aliases | Column type (OID) | +| -------------- | ------------------------ | ------------------------ | +| `UINT8` | | `smallint` (21) | +| `INT16` | `SMALLINT`, `INT2` | `smallint` (21) | +| `UINT16` | | `integer` (23) | +| `INT32` | `INTEGER`, `INT`, `INT4` | `integer` (23) | +| `UINT32` | | `bigint` (20) | +| `INT64` | `BIGINT` | `bigint` (20) | +| `FLOAT32` | `REAL`, `FLOAT4` | `real` (700) | +| `FLOAT64` | `DOUBLE`, `FLOAT` | `double precision` (701) | +| `STRING` | `TEXT`, `VARCHAR` | `text` (25) | +| `BYTES` | `BLOB`, `BYTEA` | `bytea` (17) | + +A value outside an attribute's width is `22003`. `ARRAY_AGG` `SUM` over an integer attribute is an exact `bigint`, and `22003` past its range. See [Array](array.md#attribute-types). + ## Temporal Queries (Bitemporal) Temporal queries use `AS OF` clauses to query data at a point in time. Both system time (when data entered) and valid time (when it represents) are supported across Document, Columnar, Timeseries, Spatial, Graph, and Array engines. @@ -1142,17 +1172,48 @@ SELECT * FROM WEIGHTED_PICK('reward_pool', weight => 'chance', count => 5, Returns rows with columns: `pick_index`, `key`, `weight`. Uses Vose's alias method (O(N) setup, O(1) per pick). Default RNG is ChaCha-based CSPRNG; `SEED` enables deterministic reproducibility. -## Bulk Import +## Bulk Import and Export -Bytes stream from the client over pgwire COPY; the database never opens -a file by a caller-named path. +A collection loads from a file on the server. An array also loads from a client stream with `COPY FROM STDIN`. ```sql --- Import from a client-side file (NDJSON, JSON array, or CSV) -COPY users FROM STDIN; -COPY users FROM STDIN WITH (FORMAT csv); +-- Import a server-side file into a collection (NDJSON, JSON array, or CSV) +COPY users FROM '/data/users.ndjson'; +COPY users FROM '/data/users.json'; +COPY orders FROM '/data/orders.csv' WITH (FORMAT csv, DELIMITER ';', HEADER false); + +-- Export a collection or a query to a server-side file +COPY users TO '/exports/users.ndjson'; +COPY (SELECT id, name FROM users WHERE status = 'active') TO '/exports/active.csv' WITH (FORMAT csv); ``` +- A path must be absolute and must not contain `..`. A relative path is `42601`. A path with `..` is `42501`. +- `FORMAT` is `ndjson` (or `jsonl`), `json`, or `csv`. Without `FORMAT`, the extension decides: `.ndjson`, `.jsonl`, `.json`, `.csv`. `DELIMITER` is one character. `HEADER` defaults to `true`. +- `COPY FROM STDIN` is refused with `0A000`. Use the server-path form. +- A collection on the `timeseries` or `spatial` engine refuses `COPY FROM` with `0A000`. +- A file over 16 GiB is `54000`. +- `COPY TO` writes `.tmp` and renames it to ``. It is not transactional. +- `COPY FROM ''` loads CSV only. Use `WITH (FORMAT csv)`. + +### Array bulk load + +```sql +-- Columns: the dims in schema order, then the attrs in schema order. +COPY scene FROM STDIN WITH (FORMAT csv); +COPY scene FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER); +COPY scene FROM STDIN; -- text format, tab-delimited +``` + +- `FORMAT` is `csv` or `text`. The default is `text`, tab-delimited. `HEADER` is off unless given. A bare `HEADER` turns it on. The file form `COPY FROM ''` takes CSV only, and its `HEADER` defaults to `true`. +- Each value parses by its column's declared type. No type is guessed. +- With `HEADER`, the first line must name the dims, then the attrs, in schema order. A mismatch is `42703`. +- A wrong column count or an unparsable value is `22P02`. A NULL in a NOT NULL attr is `23502`. A coordinate outside its dim domain is `22003`. A value outside its attr's width, such as `256` in a `UINT8` attr, is `22003`. +- The load streams. Rows are checked as they arrive and load in chunks of about 4 MiB. Each chunk is written before the next rows are read, so the server holds a few chunks whatever the input size. A repeated coordinate keeps its last value. +- Outside a transaction block each chunk commits on its own, as each row of a collection `COPY` does. A bad row or a refused chunk, such as `53200` under memory pressure, stops the load. The chunks before it stay committed. The error names the row and the committed row count: `COPY: CSV row 150000: ... (139810 rows committed)`. Those rows are in the array. +- Inside `BEGIN` every chunk stages, and an error aborts the block, so the load commits or rolls back whole. A chunk that takes the staged redo past the 64 MiB WAL record cap is refused with `54000` before it stages. Load more than that outside a block. +- A record must end within the array engine's memory headroom. A record with no line break past it is `54000`. +- The native protocol loads the same way through the `ArrayPutBatch` opcode (`NodeDb::array_put_cells` in `nodedb-client`). The same rules apply to each request. + ## Introspection ```sql @@ -1307,7 +1368,7 @@ SHOW USERS; | `WITH RECURSIVE` | Supported | Iterative fixed-point execution. For graph traversal, the native `GRAPH TRAVERSE`, `GRAPH PATH`, and algorithm functions remain more efficient. | | `UPDATE/DELETE ... JOIN` | Not supported | The Data Plane executes mutations as single-collection atomic operations through the SPSC bridge. Multi-collection mutations would require cross-engine coordination that breaks the isolation model. Rewrite as a subquery: `DELETE FROM orders WHERE user_id IN (SELECT id FROM users WHERE ...)`. | | `FOREIGN KEY` | Not enforced | In a distributed system with CRDT sync and eventual consistency at the edge, enforcing FK constraints across collections would require cross-shard coordination on every write — killing write throughput. CRDT constraint validation (UNIQUE, FK, CHECK) is enforced at delta-apply time for synced collections, but not for general SQL. | -| `COPY TO` (export) | Not supported | The Data Plane is write-optimized with io_uring for ingest, but export requires serialization across all shards and cores. Use the HTTP API (`/v1/query/stream`) for NDJSON export or query into Parquet via L2 cold storage. | +| `COPY TO` (export) | Supported | Writes a collection or a query result to an absolute server-side path as NDJSON, JSON array, or CSV. See [Bulk Import and Export](#bulk-import-and-export). | | `UPDATE/DELETE` on timeseries | Not supported | Timeseries collections use append-only columnar memtables with cascading compression (ALP + FastLanes + FSST + Gorilla + LZ4). In-place mutation would break compression chains and invalidate block statistics. Use retention policies to age out old data. | | `EXPLAIN ANALYZE` | Not yet | Requires instrumentation across the SPSC bridge to collect per-core execution stats from the Data Plane and merge them on the Control Plane. The bridge currently returns results but not timing metadata. Planned. | diff --git a/docs/security/audit.md b/docs/security/audit.md index 35478785d..210d1087b 100644 --- a/docs/security/audit.md +++ b/docs/security/audit.md @@ -10,50 +10,68 @@ SHOW AUDIT LOG; SHOW AUDIT LOG LIMIT 50; ``` -**Columns:** `seq`, `timestamp`, `event`, `tenant_id`, `source`, `detail` +**Columns:** `seq`, `timestamp_us`, `event`, `tenant_id`, `database_id`, `source`, `detail` -Each entry also records `auth_user_id`, `auth_user_name`, and `session_id` for correlation. +`SHOW AUDIT LOG`, `SHOW AUDIT WHERE` and `SHOW AUDIT IN DATABASE` return the newest entries first. Without `LIMIT`, each returns 100 rows. Only a superuser can run them. + +The `event` column prints the snake_case event name, for example `auth_success`. `WHERE event_type = '...'` takes the same snake_case name. + +A read sees every entry recorded before it starts. There is no flush delay. + +`_system.audit_log` has the columns `seq`, `timestamp_us`, `event`, `tenant_id`, `source`, `detail` and `prev_hash`. Its `event` column prints the snake_case name too. It has no `database_id` column. A superuser or a role with `monitor` can read it. It returns at most the newest 100,000 entries. + +Each stored entry also records `auth_user_id`, `auth_user_name`, and `session_id` for correlation. ## Audit Events -| Event | Level | Description | -| ---------------------------- | -------- | ---------------------------------------------- | -| `AuthSuccess` | Minimal | Successful authentication | -| `AuthFailure` | Minimal | Failed login attempt | -| `PermissionDenied` | Minimal | Permission check failed | -| `AuthzDenied` | Minimal | Legacy: permission denied | -| `PrivilegeChange` | Standard | GRANT, REVOKE, role changes | -| `SessionConnect` | Standard | New connection | -| `SessionDisconnect` | Standard | Connection closed | -| `SessionRevoked` | Standard | Admin terminated session | -| `LockoutTriggered` | Standard | Account locked after failed logins | -| `LoginRateLimited` | Standard | Too many login attempts | -| `RlsRejected` | Standard | RLS policy blocked row | -| `AdminAction` | Standard | DDL, config changes | -| `TenantCreated` | Standard | New tenant | -| `TenantDeleted` | Standard | Tenant removed | -| `DatabaseCreated` | Standard | New database | -| `DatabaseDropped` | Standard | Database deleted | -| `DatabaseRenamed` | Standard | Database renamed | -| `DatabaseQuotaChanged` | Standard | Quota limit changed | -| `DatabaseCloned` | Standard | Database cloned | -| `DatabaseMirrored` | Standard | Database mirrored | -| `DatabasePromoted` | Standard | Replica promoted to primary | -| `DatabaseMaterialized` | Standard | Materialized view materialized | -| `TenantMoved` | Standard | Tenant reassigned to database | -| `DatabaseBackedUp` | Standard | Backup created | -| `DatabaseRestored` | Standard | Backup restored | -| `DatabaseAuditDmlChanged` | Standard | DML audit setting changed | -| `DatabaseIdleTimeoutChanged` | Standard | Idle timeout changed | -| `OidcProviderChanged` | Standard | OIDC provider created/altered/dropped | -| `DmlAudit` | Forensic | Individual DML operation (opt-in per database) | -| `SnapshotBegin/End` | Standard | Backup lifecycle | -| `RestoreBegin/End` | Standard | Restore lifecycle | -| `CertRotation` | Standard | TLS cert rotated | -| `KeyRotation` | Standard | Encryption key rotated | -| `NodeJoined/Left` | Standard | Cluster membership | -| `QueryExec` | Full | Every query executed | -| `RowChange` | Forensic | Individual row mutations | +| Event | Level | Description | +| -------------------------------------- | -------- | ---------------------------------------------- | +| `auth_success` | Minimal | Successful authentication | +| `auth_failure` | Minimal | Failed login attempt | +| `permission_denied` | Minimal | Permission check failed | +| `authz_denied` | Minimal | Authorization denied | +| `rls_rejected` | Minimal | RLS write policy rejected a write | +| `lockout_triggered` | Minimal | Account locked after failed logins | +| `login_rate_limited` | Minimal | Too many login attempts | +| `session_revoked` | Minimal | Open session revoked | +| `audit_bus_lagged` | Minimal | Audit bus dropped events, rows may be missing | +| `audit_checkpoint` | Minimal | Retention prune checkpoint | +| `privilege_change` | Standard | GRANT, REVOKE, role changes | +| `session_connect` | Standard | New connection | +| `session_disconnect` | Standard | Connection closed | +| `session_handle_fingerprint_mismatch` | Standard | Session handle used from a different origin | +| `session_handle_resolve_miss_spike` | Standard | Many failed session handle lookups | +| `admin_action` | Standard | Administrative actions | +| `config_change` | Standard | Config change | +| `ddl_change` | Standard | DDL change committed to the metadata group | +| `tenant_created` | Standard | New tenant | +| `tenant_deleted` | Standard | Tenant removed | +| `tenant_moved` | Standard | Tenant reassigned to database | +| `database_created` | Standard | New database | +| `database_dropped` | Standard | Database dropped | +| `database_renamed` | Standard | Database renamed | +| `database_quota_changed` | Standard | Quota limit changed | +| `database_cloned` | Standard | Database cloned | +| `database_mirrored` | Standard | Database mirrored | +| `database_promoted` | Standard | Mirror promoted to writable primary | +| `database_materialized` | Standard | Cloned database materialized | +| `database_backed_up` | Standard | Database backup started | +| `database_restored` | Standard | Database restored from backup | +| `database_audit_dml_changed` | Standard | DML audit mode changed | +| `database_idle_timeout_changed` | Standard | Idle timeout changed | +| `oidc_provider_changed` | Standard | OIDC provider created, altered, or dropped | +| `snapshot_begin`, `snapshot_end` | Standard | Snapshot lifecycle | +| `restore_begin`, `restore_end` | Standard | Restore lifecycle | +| `cert_rotation` | Standard | TLS cert rotated | +| `cert_rotation_failed` | Standard | TLS cert rotation failed | +| `key_rotation` | Standard | Encryption key rotated | +| `node_joined`, `node_left` | Standard | Cluster membership | +| `query_exec` | Full | Every query executed | +| `rls_denied` | Full | RLS denial | +| `row_change` | Forensic | Individual row mutations | +| `dml_audit` | Forensic | Individual DML operation (opt-in per database) | + +Events at the Minimal level are recorded at every audit level. ## Audit Levels @@ -80,17 +98,39 @@ Filter audit events by database to reduce noise when operating a single database -- Show audit events for a specific database SHOW AUDIT IN DATABASE prod LIMIT 50; --- Combine with event type filter -SHOW AUDIT IN DATABASE prod WHERE event_type = 'DmlAudit' LIMIT 100; +-- Combine with an event type filter +SHOW AUDIT IN DATABASE prod WHERE event_type = 'dml_audit' LIMIT 100; + +-- Filter by event type across all databases +SHOW AUDIT WHERE event_type = 'dml_audit' LIMIT 100; ``` +`SHOW AUDIT LOG`, `SHOW AUDIT WHERE` and `SHOW AUDIT IN DATABASE ` all take `[WHERE event_type = ''] [LIMIT n]`. `SHOW AUDIT WHERE` requires the `WHERE` clause. Any other `WHERE` form is refused with SQLSTATE `42601`, and the message names the clause. A bad `LIMIT` value or a stray token is also refused with `42601`. An unknown database name returns SQLSTATE `3D000`. + Database-scoped DDL (CLONE, MIRROR, PROMOTE, etc.) include `database_id` in the audit entry, enabling per-database filtering and forensics. ## Hash Chain Integrity -Every `AuditEntry` contains `prev_hash` — the SHA-256 of the previous entry. This creates a tamper-evident chain verified on startup. If an entry is modified or deleted, the chain breaks and is flagged. +Every `AuditEntry` contains `prev_hash` — the SHA-256 of the previous entry. If an entry is modified or deleted, the chain breaks. The hash covers `database_id` when the entry has one. + +Two commands check the chain. Only a superuser can run them. + +```sql +VERIFY AUDIT CHAIN; +SELECT VERIFY_AUDIT_CHAIN(1, 5000); +``` + +- Both read the whole log: catalog rows plus the in-memory buffer. +- `VERIFY AUDIT CHAIN` checks the whole chain. It returns `status = VALID` and the entry count, or the error `XX001` `audit chain broken at sequence N`. +- `VERIFY_AUDIT_CHAIN(from, to)` checks the entries in that sequence range. Entries before `from` still feed the first link. It returns `valid`, `entries_checked`, `from_seq`, `to_seq`, `broken_at_seq`, and `last_hash`. +- `broken_at_seq` is the first sequence whose link fails. An altered entry is at that sequence or the one before it. +- The newest entry is also compared with the live tail hash. + +## Retention Pruning + +`audit_retention_days` prunes entries older than that age. `audit_max_entries` prunes the oldest entries beyond that count. Age pruning runs first. Both default to 0, which keeps everything. -Hash chain is extended with `database_id` when present (new events include this field; legacy events without database scope are preserved for compatibility). +Each prune that deletes entries first records an `audit_checkpoint` entry. The checkpoint holds `pruned_count`, `oldest_kept_seq`, and `anchor_hash`, the hash of the last deleted entry. The checkpoint is a normal chain link, so the chain still verifies after a prune. A prune that deletes nothing records no checkpoint. ## DML Audit (Opt-In) @@ -107,7 +147,7 @@ ALTER DATABASE prod SET AUDIT_DML = 'all'; ALTER DATABASE prod SET AUDIT_DML = 'none'; -- View DML audit entries -SHOW AUDIT IN DATABASE prod WHERE event_type = 'DmlAudit' LIMIT 100; +SHOW AUDIT WHERE event_type = 'dml_audit' LIMIT 100; ``` **DML audit entry includes:** diff --git a/docs/security/auth.md b/docs/security/auth.md index 953c8c134..34c3ae466 100644 --- a/docs/security/auth.md +++ b/docs/security/auth.md @@ -25,7 +25,7 @@ SHOW USERS; Connect via psql: ```bash -psql -h localhost -p 6432 -U alice +psql -h localhost -p 6432 -U alice -d default ``` ## API Keys diff --git a/docs/security/oidc.md b/docs/security/oidc.md index 377db8085..8e14a6f3c 100644 --- a/docs/security/oidc.md +++ b/docs/security/oidc.md @@ -218,24 +218,29 @@ Then Alice's session is bound to the `prod` database. She cannot query other dat OIDC authentication events appear in the audit log: ```sql -SHOW AUDIT WHERE event_type = 'AuthSuccess' AND provider = 'okta'; +SHOW AUDIT WHERE event_type = 'auth_success'; ``` -**Audit entry includes:** +`SHOW AUDIT WHERE` filters on `event_type` only, with the snake_case event name. It does not filter on other fields. -| Field | Value | -| ------------- | ----------------------------------- | -| `event_type` | `AuthSuccess` | -| `provider` | OIDC provider name | -| `jwt_subject` | JWT `sub` claim (e.g. user's email) | -| `auth_method` | `OidcBearer` | -| `timestamp` | Login time | +Every accepted bearer login records an `auth_success` entry. The entry carries `auth_user_id`, `auth_user_name` and `session_id`. Its detail reads: -**Failed OIDC attempts:** +``` +bearer login accepted: method= provider="" sub="" +``` + +**Failed attempts:** ```sql -SHOW AUDIT WHERE event_type = 'AuthFailure' AND auth_method = 'OidcBearer'; --- Returns: invalid signature, missing kid, expired token, audience mismatch, etc. +SHOW AUDIT WHERE event_type = 'auth_failure'; +``` + +Every refusal records an `auth_failure` entry. Its detail reads: + ``` +bearer login refused: method= provider= sub= reason="" +``` + +The audit log never holds the token. A bearer refusal does not count toward account lockout or the login rate limit. [Back to security](README.md) diff --git a/docs/security/sessions.md b/docs/security/sessions.md index 061ce4b66..516783bf1 100644 --- a/docs/security/sessions.md +++ b/docs/security/sessions.md @@ -171,7 +171,7 @@ login_lockout_duration_secs = 900 # Lock for 15 minutes **Audit:** ```sql -SHOW AUDIT WHERE event_type = 'LockoutTriggered'; +SHOW AUDIT WHERE event_type = 'lockout_triggered'; -- Returns: username, ip, locked_until (timestamp) ``` @@ -209,7 +209,7 @@ login_attempts_per_user_per_min = 10 **Audit:** ```sql -SHOW AUDIT WHERE event_type = 'LoginRateLimited'; +SHOW AUDIT WHERE event_type = 'login_rate_limited'; -- Returns: username, ip, timestamp ``` diff --git a/docs/security/tenants.md b/docs/security/tenants.md index 38e41cf59..f429cd051 100644 --- a/docs/security/tenants.md +++ b/docs/security/tenants.md @@ -30,7 +30,7 @@ CREATE RLS POLICY tenant_isolation ON orders FOR ALL -- Audit is per-database; RLS is per-tenant ALTER DATABASE us_west SET AUDIT_DML = 'writes'; -SHOW AUDIT IN DATABASE us_west WHERE event_type = 'DmlAudit'; +SHOW AUDIT IN DATABASE us_west WHERE event_type = 'dml_audit'; ``` ## Creating Tenants diff --git a/docs/spatial.md b/docs/spatial.md index c5c91a0b0..f4aaad6ba 100644 --- a/docs/spatial.md +++ b/docs/spatial.md @@ -19,7 +19,7 @@ Spatial is a **columnar profile**. Collections with a `SPATIAL_INDEX` column mod - **H3 hexagonal index** — Uber's H3 via h3o for uniform-area spatial binning - **OGC predicates** — `ST_Contains`, `ST_Intersects`, `ST_Within`, `ST_Disjoint`, `ST_DWithin`, `ST_IsValid` - **Geometry constructors** — `ST_Point`, `ST_MakePoint`, `ST_GeomFromText` (WKT), `ST_GeomFromWKB`, `ST_GeomFromGeoJSON`, `ST_MakeLine`, `ST_MakePolygon`, `ST_MakeEnvelope` -- **Accessors and measures** — `ST_AsText`, `ST_AsGeoJSON`, `ST_X`, `ST_Y`, `ST_GeometryType`, `ST_NPoints`, `ST_SRID`, `ST_Distance`, `ST_Length`, `ST_Perimeter`, `ST_Area` +- **Accessors and measures** — `ST_AsText`, `ST_AsGeoJSON`, `ST_X`, `ST_Y`, `ST_GeometryType`, `ST_NPoints`, `ST_SRID`, `ST_SetSRID`, `ST_Transform`, `ST_Distance`, `ST_Length`, `ST_Perimeter`, `ST_Area` - **Geometry operations** — `ST_Buffer`, `ST_Envelope`, `ST_Centroid`, `ST_Union`, `ST_Intersection` - **Format support** — WKB, WKT, GeoJSON interchange. GeoParquet v1.1.0 + GeoArrow metadata. - **Hybrid spatial-vector** — Spatial R\*-tree narrows candidates by location, then HNSW ranks by semantic similarity in one query @@ -117,12 +117,48 @@ A string literal in geometry position is read as either WKT or GeoJSON, so | `ST_X(geom)` / `ST_Y(geom)` | Ordinate of a Point; NULL for any other geometry | | `ST_GeometryType(geom)` | Type name, e.g. `Point` | | `ST_NPoints(geom)` | Total vertex count | -| `ST_SRID(geom)` | Always `4326` — geometry is stored as GeoJSON, which RFC 7946 defines as WGS 84 | +| `ST_SRID(geom)` | SRID of the geometry; `4326` unless the geometry carries another SRID. See [Spatial reference systems](#spatial-reference-systems) | | `ST_IsValid(geom)` | Whether the geometry is well-formed | +### Spatial reference systems + +Every geometry carries an SRID. The default is `4326` (WGS 84, longitude and latitude in degrees). GeoJSON reads and writes a `crs` member for any SRID other than `4326`. + +| SRID | System | Unit | +| --- | --- | --- | +| `4326` | WGS 84 geographic | degrees | +| `32647` | UTM zone 47N on WGS 84 | metres | +| `32648` | UTM zone 48N on WGS 84 | metres | +| `3375` | GDM2000 Peninsula RSO | metres | +| `3376` | GDM2000 East Malaysia (Borneo) RSO | metres | + +The set is closed. Any other SRID is refused. + +| Function | Result | +| --- | --- | +| `ST_SRID(geom)` | SRID of the geometry | +| `ST_SetSRID(geom, srid)` | Geometry relabeled with `srid`. Coordinates stay as they are. | +| `ST_Transform(geom, srid)` | Geometry with every coordinate reprojected into `srid` | + +`ST_GeomFromText` and `ST_GeomFromWKB` take an optional `srid` argument. `ST_SetSRID` and `ST_Transform` return NULL when `srid` is NULL. + +```sql +SELECT ST_Transform(ST_Point(101.7117, 3.1579), 32647); +``` + +Geometry functions read the SRID of each operand. They never transform implicitly. These cases raise a data exception with SQLSTATE `22000`: + +- Two operands in different SRIDs. The error reads `operands have different SRIDs and ; use ST_Transform`. +- An SRID outside the supported set. The error lists the supported codes. +- A coordinate outside the region the target projection is defined on. + +Call `ST_Transform` on one operand to bring both into the same SRID. + +A spatial index fixes its SRID at its first entry. It refuses geometries in any other SRID. + ### Measures -Every measure is geodesic: lengths in meters, areas in square meters. +Every measure is in meters, and square meters for area. A WGS 84 (`4326`) geometry is measured geodesically. A geometry in a projected SRID is measured in planar grid meters. `ST_Distance` needs both operands in the same SRID. | Function | Result | | --- | --- | @@ -138,7 +174,7 @@ Every measure is geodesic: lengths in meters, areas in square meters. | `ST_Contains(a, b)` | Whether `a` contains `b` | | `ST_Within(a, b)` | Whether `a` lies within `b` — the converse of `ST_Contains` | | `ST_Intersects(a, b)` / `ST_Disjoint(a, b)` | Whether the two geometries share any point | -| `ST_DWithin(a, b, meters)` | Whether the two geometries lie within a geodesic distance | +| `ST_DWithin(a, b, meters)` | Whether the two geometries lie within a distance: geodesic for `4326`, planar for a projected SRID | | `ST_Buffer(geom, meters [, segments])` | Geometry grown by a distance | | `ST_Envelope(geom)` | Bounding rectangle | | `ST_Centroid(geom)` | Centroid, weighted by the highest dimension present | diff --git a/docs/vectors.md b/docs/vectors.md index db22dcfca..634d2062f 100644 --- a/docs/vectors.md +++ b/docs/vectors.md @@ -97,8 +97,16 @@ INSERT INTO sparse_docs (id, terms) VALUES ('a', '{3: 1.0, 7: 1.0}'); SELECT id FROM sparse_docs ORDER BY sparse_score(terms, '{3: 1.0, 7: 0.5}') DESC LIMIT 3; + +-- A WHERE clause filters the hits before the top-k cut +SELECT id FROM sparse_docs +WHERE id <> 'a' +ORDER BY sparse_score(terms, '{3: 1.0, 7: 0.5}') DESC +LIMIT 3; ``` +A `WHERE` clause filters the candidates before the top-k cut, so a matching row is never dropped by the cut. Projected columns read from the row. + `SPARSEVECTOR` carries no fixed dimensionality — only non-zero `(dimension, weight)` pairs are stored and indexed. Documents sharing no dimension with the query are excluded by the inverted index rather than scored at zero. `sparse_score` ranks by descending dot product; it is a standalone surface, separate from the dense `vector_distance` path and the RRF hybrid fusion. ## Vector-Primary Collections diff --git a/docs/wasm.md b/docs/wasm.md index 4a718edfc..f481e4e8a 100644 --- a/docs/wasm.md +++ b/docs/wasm.md @@ -146,19 +146,19 @@ const results = await db.sql( // Create array await db.sql(` CREATE ARRAY spatial_grid - DIMS (x INT64 DOMAIN [0, 1000), y INT64 DOMAIN [0, 1000)) - ATTRS (temperature FLOAT32) + DIMS (x INT64 [0..1000], y INT64 [0..1000]) + ATTRS (temperature FLOAT64) TILE_EXTENTS (64, 64) `); // Insert cells await db.sql( - "INSERT INTO spatial_grid (x, y, temperature) VALUES (10, 20, 23.5)", + "INSERT INTO ARRAY spatial_grid COORDS (10, 20) VALUES (23.5)", ); // Query slice const results = await db.sql( - `SELECT * FROM ARRAY_SLICE('spatial_grid', {x: [0, 100), y: [0, 100)}, ['temperature'])`, + `SELECT * FROM ARRAY_SLICE('spatial_grid', '{x: [0, 99], y: [0, 99]}', ['temperature'])`, ); ``` diff --git a/fuzz/src/targets/strict_tuple.rs b/fuzz/src/targets/strict_tuple.rs index 304101a13..0a85987e6 100644 --- a/fuzz/src/targets/strict_tuple.rs +++ b/fuzz/src/targets/strict_tuple.rs @@ -1,6 +1,6 @@ use nodedb_strict::{TupleDecoder, TupleEncoder}; use nodedb_types::columnar::{ColumnDef, ColumnType, StrictSchema}; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::value::Value; use nodedb_types::{NdbDateTime, NdbDuration}; @@ -97,10 +97,10 @@ fn valid_value(column_type: ColumnType, seed: u8) -> Option { ])) } ColumnType::SparseVector => Some(Value::String("{1: 1.0}".to_owned())), - ColumnType::Geometry => Some(Value::Geometry(Geometry::point( + ColumnType::Geometry => Some(Value::Geometry(SridGeometry::wgs84(Geometry::point( f64::from(seed) - 128.0, f64::from(seed % 181) - 90.0, - ))), + )))), ColumnType::Json => Some(Value::Array(vec![Value::Integer(i64::from(seed))])), ColumnType::Ulid => Some(Value::Ulid("01ARZ3NDEKTSV4RRFFQ69G5FAV".to_owned())), ColumnType::Duration => Some(Value::Duration(NdbDuration::from_micros(i64::from(seed)))), diff --git a/nodedb-array/Cargo.toml b/nodedb-array/Cargo.toml index fdc96723e..05c22889f 100644 --- a/nodedb-array/Cargo.toml +++ b/nodedb-array/Cargo.toml @@ -23,6 +23,7 @@ crc32c = { workspace = true } uuid = { workspace = true } loro = { workspace = true } nodedb-wal = { workspace = true } +rustc-hash = { workspace = true } [dev-dependencies] serde_json = { workspace = true } diff --git a/nodedb-array/src/codec/attr_col/column.rs b/nodedb-array/src/codec/attr_col/column.rs new file mode 100644 index 000000000..f2b8cf7e6 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/column.rs @@ -0,0 +1,315 @@ +// SPDX-License-Identifier: Apache-2.0 + +// Attr column format: +// +// [u8 tag] the column's attr type, or NULL_ONLY +// [u32 LE count] rows, NULL rows included +// validity section see `validity.rs` +// payload the non-null values only, by tag +// +// A NULL row stores no value, so a NULL never decodes as 0 or 0.0. Every +// non-null value of a column has one type: the schema check at write time +// holds a stored column to its attr's declared type. + +use crate::codec::limits::{MAX_COLUMN_ENTRIES, checked_capacity}; +use crate::error::{ArrayError, ArrayResult}; +use crate::schema::attr_spec::AttrType; +use crate::types::cell_value::value::CellValue; + +use super::float::{decode_f32, decode_f64, encode_f32, encode_f64}; +use super::int::{decode_ints, encode_ints}; +use super::validity::{decode_validity, encode_validity}; +use super::varlen::{decode_varlen, encode_varlen}; + +const TAG_NULL_ONLY: u8 = 0; + +fn tag_of(dtype: AttrType) -> u8 { + match dtype { + AttrType::Int64 => 1, + AttrType::Float64 => 2, + AttrType::String => 3, + AttrType::Bytes => 4, + AttrType::UInt8 => 5, + AttrType::Int16 => 6, + AttrType::UInt16 => 7, + AttrType::Int32 => 8, + AttrType::UInt32 => 9, + AttrType::Float32 => 10, + } +} + +fn type_of(tag: u8) -> ArrayResult> { + Ok(Some(match tag { + TAG_NULL_ONLY => return Ok(None), + 1 => AttrType::Int64, + 2 => AttrType::Float64, + 3 => AttrType::String, + 4 => AttrType::Bytes, + 5 => AttrType::UInt8, + 6 => AttrType::Int16, + 7 => AttrType::UInt16, + 8 => AttrType::Int32, + 9 => AttrType::UInt32, + 10 => AttrType::Float32, + other => { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col: unknown tag {other:#04x}"), + }); + } + })) +} + +/// The one attr type of the column's non-null values, or `None` when every +/// row is NULL. A column that mixes types is refused. +fn column_type(values: &[CellValue]) -> ArrayResult> { + let mut found: Option = None; + for t in values.iter().filter_map(CellValue::attr_type) { + match found { + None => found = Some(t), + Some(f) if f == t => {} + Some(f) => { + return Err(ArrayError::InvalidOp { + detail: format!( + "attr column mixes {} and {} values", + f.sql_name(), + t.sql_name() + ), + }); + } + } + } + Ok(found) +} + +pub fn encode_attr_col(values: &[CellValue]) -> ArrayResult> { + let dtype = column_type(values)?; + let mut out = Vec::new(); + out.push(dtype.map_or(TAG_NULL_ONLY, tag_of)); + out.extend_from_slice(&(values.len() as u32).to_le_bytes()); + let Some(dtype) = dtype else { + return Ok(out); + }; + encode_validity(values, &mut out); + let present: Vec<&CellValue> = values.iter().filter(|v| !v.is_null()).collect(); + let payload = match dtype { + AttrType::Int64 + | AttrType::UInt8 + | AttrType::Int16 + | AttrType::UInt16 + | AttrType::Int32 + | AttrType::UInt32 => encode_ints(&present)?, + AttrType::Float64 => encode_f64(&present)?, + AttrType::Float32 => encode_f32(&present)?, + AttrType::String | AttrType::Bytes => encode_varlen(&present)?, + }; + out.extend_from_slice(&payload); + Ok(out) +} + +/// Decode an attr column of `rows` rows. A column that holds another row count +/// is corrupt and is refused before any allocation. +pub fn decode_attr_col(data: &[u8], rows: usize) -> ArrayResult> { + let Some((&tag, rest)) = data.split_first() else { + return Err(ArrayError::SegmentCorruption { + detail: "attr col: empty payload".into(), + }); + }; + let Some((count_bytes, _)) = rest.split_first_chunk::<4>() else { + return Err(ArrayError::SegmentCorruption { + detail: "attr col: truncated count".into(), + }); + }; + let stored = u32::from_le_bytes(*count_bytes) as usize; + if stored != rows { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col holds {stored} rows, its tile holds {rows}"), + }); + } + let count = checked_capacity(stored, MAX_COLUMN_ENTRIES, rows, 0, "attr_col count")?; + let mut pos = 5; + let Some(dtype) = type_of(tag)? else { + if data.len() != pos { + return Err(ArrayError::SegmentCorruption { + detail: "attr col: a NULL-only column carries a payload".into(), + }); + } + return Ok(vec![CellValue::Null; count]); + }; + let validity = decode_validity(data, &mut pos, count)?; + let present = (0..count).filter(|&row| validity.is_present(row)).count(); + if present == 0 { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col {}: no value under a typed tag", dtype.sql_name()), + }); + } + let body = &data[pos..]; + let values = match dtype { + AttrType::Int64 + | AttrType::UInt8 + | AttrType::Int16 + | AttrType::UInt16 + | AttrType::Int32 + | AttrType::UInt32 => decode_ints(body, dtype, present)?, + AttrType::Float64 => decode_f64(body, present)?, + AttrType::Float32 => decode_f32(body, present)?, + AttrType::String | AttrType::Bytes => decode_varlen(body, dtype, present)?, + }; + let mut values = values.into_iter(); + let mut out = Vec::with_capacity(count); + for row in 0..count { + if validity.is_present(row) { + out.push(values.next().ok_or_else(|| ArrayError::SegmentCorruption { + detail: "attr col: fewer values than the validity holds".into(), + })?); + } else { + out.push(CellValue::Null); + } + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn round_trip(values: &[CellValue]) -> Vec { + decode_attr_col(&encode_attr_col(values).expect("encode"), values.len()).expect("decode") + } + + fn with_nulls(values: Vec) -> Vec { + values + .into_iter() + .enumerate() + .map(|(i, v)| if i % 4 == 1 { CellValue::Null } else { v }) + .collect() + } + + #[test] + fn every_type_round_trips_with_nulls() { + let columns: Vec> = vec![ + (0..40).map(CellValue::Int64).collect(), + (0..40) + .map(|i| CellValue::Float64(i as f64 * 0.5)) + .collect(), + (0..40) + .map(|i| CellValue::String(format!("s{i}"))) + .collect(), + (0..40) + .map(|i| CellValue::Bytes(vec![i as u8; 3])) + .collect(), + (0..40).map(|i| CellValue::UInt8(i as u8 * 6)).collect(), + (0..40).map(|i| CellValue::Int16(i as i16 * -800)).collect(), + (0..40) + .map(|i| CellValue::UInt16(i as u16 * 1600)) + .collect(), + (0..40).map(|i| CellValue::Int32(i * -50_000_000)).collect(), + (0..40) + .map(|i| CellValue::UInt32(i as u32 * 100_000_000)) + .collect(), + (0..40) + .map(|i| CellValue::Float32(i as f32 * 0.1)) + .collect(), + ]; + for col in columns { + let col = with_nulls(col); + assert_eq!(round_trip(&col), col); + } + } + + /// A NULL keeps its NULL through the codec, in a numeric column too. + #[test] + fn numeric_null_stays_null() { + let ints = vec![CellValue::Int64(5), CellValue::Null, CellValue::Int64(0)]; + assert_eq!(round_trip(&ints), ints); + let floats = vec![CellValue::Null, CellValue::Float64(0.0), CellValue::Null]; + assert_eq!(round_trip(&floats), floats); + let narrow = vec![CellValue::UInt16(0), CellValue::Null]; + assert_eq!(round_trip(&narrow), narrow); + } + + #[test] + fn all_null_and_empty_round_trip() { + let nulls = vec![CellValue::Null; 17]; + let encoded = encode_attr_col(&nulls).unwrap(); + assert_eq!(encoded.len(), 5); + assert_eq!(round_trip(&nulls), nulls); + assert!(round_trip(&[]).is_empty()); + } + + #[test] + fn a_mixed_column_is_refused() { + let mixed = vec![CellValue::Int64(1), CellValue::Float64(1.0)]; + assert!(matches!( + encode_attr_col(&mixed), + Err(ArrayError::InvalidOp { .. }) + )); + } + + /// FastLanes packs by value range, so a narrow integer column costs the + /// bytes of the same values in an `INT64` column. + #[test] + fn uint16_bytes_equal_int64_bytes() { + let narrow: Vec = (0..1000u16).map(|i| CellValue::UInt16(i * 37)).collect(); + let wide: Vec = (0..1000i64).map(|i| CellValue::Int64(i * 37)).collect(); + assert_eq!( + encode_attr_col(&narrow).unwrap().len(), + encode_attr_col(&wide).unwrap().len() + ); + } + + /// A FLOAT32 column of reflectance-like values is smaller than the same + /// values stored as FLOAT64. + #[test] + fn float32_is_smaller_than_float64() { + let vals: Vec = (0..1000) + .map(|i| ((i * 7919) % 10_000) as f32 / 10_000.0) + .collect(); + let narrow: Vec = vals.iter().map(|v| CellValue::Float32(*v)).collect(); + let wide: Vec = vals + .iter() + .map(|v| CellValue::Float64(f64::from(*v))) + .collect(); + let narrow_len = encode_attr_col(&narrow).unwrap().len(); + let wide_len = encode_attr_col(&wide).unwrap().len(); + assert!(narrow_len < wide_len, "{narrow_len} >= {wide_len}"); + assert!(narrow_len <= 5 + 1 + 1 + 4000); + } + + #[test] + fn a_huge_count_in_a_tiny_payload_is_refused() { + let mut data = vec![tag_of(AttrType::String)]; + data.extend_from_slice(&(MAX_COLUMN_ENTRIES as u32).to_le_bytes()); + data.push(0); + assert!(matches!( + decode_attr_col(&data, MAX_COLUMN_ENTRIES), + Err(ArrayError::SegmentCorruption { .. }) + )); + } + + #[test] + fn an_out_of_width_stored_value_is_corruption() { + let mut data = vec![tag_of(AttrType::UInt8)]; + data.extend_from_slice(&1u32.to_le_bytes()); + data.push(0); + data.extend_from_slice(&nodedb_codec::fastlanes::encode(&[300]).unwrap()); + assert!(matches!( + decode_attr_col(&data, 1), + Err(ArrayError::SegmentCorruption { .. }) + )); + } + + #[test] + fn an_unknown_tag_is_corruption() { + assert!(decode_attr_col(&[0xEE, 0, 0, 0, 0], 0).is_err()); + } + + #[test] + fn a_row_count_that_differs_from_the_tile_is_refused() { + let mut data = vec![TAG_NULL_ONLY]; + data.extend_from_slice(&8_000_000u32.to_le_bytes()); + assert!(matches!( + decode_attr_col(&data, 3), + Err(ArrayError::SegmentCorruption { .. }) + )); + } +} diff --git a/nodedb-array/src/codec/attr_col/float.rs b/nodedb-array/src/codec/attr_col/float.rs new file mode 100644 index 000000000..f231224f6 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/float.rs @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: Apache-2.0 + +// Float attr payloads. +// +// FLOAT64: Gorilla XOR over the non-null values. +// FLOAT32: [u8 mode][body]. Mode 0 holds raw f32 LE values, mode 1 holds +// 32-bit Gorilla XOR. The encoder writes the smaller one. + +use crate::error::{ArrayError, ArrayResult}; +use crate::types::cell_value::value::CellValue; + +use super::int::codec_err; + +const F32_RAW: u8 = 0; +const F32_GORILLA: u8 = 1; + +fn not_float(v: &CellValue, width: &str) -> ArrayError { + ArrayError::InvalidOp { + detail: format!("{width} attr column holds {v:?}"), + } +} + +fn count_mismatch(width: &str, got: usize, expected: usize) -> ArrayError { + ArrayError::SegmentCorruption { + detail: format!("attr col {width}: {got} values decoded, validity holds {expected}"), + } +} + +pub(super) fn encode_f64(present: &[&CellValue]) -> ArrayResult> { + let floats = present + .iter() + .map(|v| match v { + CellValue::Float64(f) => Ok(*f), + other => Err(not_float(other, "FLOAT64")), + }) + .collect::>>()?; + Ok(nodedb_codec::gorilla::encode_f64(&floats)) +} + +pub(super) fn decode_f64(body: &[u8], expected: usize) -> ArrayResult> { + let floats = nodedb_codec::gorilla::decode_f64(body).map_err(codec_err)?; + if floats.len() != expected { + return Err(count_mismatch("FLOAT64", floats.len(), expected)); + } + Ok(floats.into_iter().map(CellValue::Float64).collect()) +} + +pub(super) fn encode_f32(present: &[&CellValue]) -> ArrayResult> { + let floats = present + .iter() + .map(|v| match v { + CellValue::Float32(f) => Ok(*f), + other => Err(not_float(other, "FLOAT32")), + }) + .collect::>>()?; + let gorilla = nodedb_codec::gorilla_f32::encode_f32(&floats); + let raw_len = floats.len() * 4; + let mut out = Vec::with_capacity(1 + gorilla.len().min(raw_len)); + if gorilla.len() < raw_len { + out.push(F32_GORILLA); + out.extend_from_slice(&gorilla); + } else { + out.push(F32_RAW); + for f in &floats { + out.extend_from_slice(&f.to_le_bytes()); + } + } + Ok(out) +} + +pub(super) fn decode_f32(body: &[u8], expected: usize) -> ArrayResult> { + let Some((&mode, rest)) = body.split_first() else { + return Err(ArrayError::SegmentCorruption { + detail: "attr col FLOAT32: missing mode byte".into(), + }); + }; + let floats = match mode { + F32_RAW => { + if rest.len() != expected * 4 { + return Err(count_mismatch("FLOAT32", rest.len() / 4, expected)); + } + rest.as_chunks::<4>() + .0 + .iter() + .map(|c| f32::from_le_bytes(*c)) + .collect::>() + } + F32_GORILLA => nodedb_codec::gorilla_f32::decode_f32(rest).map_err(codec_err)?, + other => { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col FLOAT32: unknown mode {other}"), + }); + } + }; + if floats.len() != expected { + return Err(count_mismatch("FLOAT32", floats.len(), expected)); + } + Ok(floats.into_iter().map(CellValue::Float32).collect()) +} diff --git a/nodedb-array/src/codec/attr_col/int.rs b/nodedb-array/src/codec/attr_col/int.rs new file mode 100644 index 000000000..40d00f9e5 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/int.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 + +// Integer attr payload: the non-null values widened to `i64`, FastLanes +// encoded. FOR + bit-packing packs a narrow column by its value range, so a +// `UINT16` column costs the same bytes as an `INT64` column of equal values. + +use nodedb_codec::error::CodecError; + +use crate::error::{ArrayError, ArrayResult}; +use crate::schema::attr_spec::AttrType; +use crate::types::cell_value::value::CellValue; + +pub(super) fn codec_err(e: CodecError) -> ArrayError { + ArrayError::SegmentCorruption { + detail: format!("codec error: {e}"), + } +} + +/// Encode the integer values of `present`. Every value is an integer cell. +pub(super) fn encode_ints(present: &[&CellValue]) -> ArrayResult> { + let ints = present + .iter() + .map(|v| { + v.as_i64().ok_or_else(|| ArrayError::InvalidOp { + detail: format!("integer attr column holds {v:?}"), + }) + }) + .collect::>>()?; + nodedb_codec::fastlanes::encode(&ints).map_err(codec_err) +} + +/// Decode `expected` integer cells of type `dtype`. +pub(super) fn decode_ints( + body: &[u8], + dtype: AttrType, + expected: usize, +) -> ArrayResult> { + let ints = nodedb_codec::fastlanes::decode(body).map_err(codec_err)?; + if ints.len() != expected { + return Err(ArrayError::SegmentCorruption { + detail: format!( + "attr col {}: {} values decoded, validity holds {expected}", + dtype.sql_name(), + ints.len() + ), + }); + } + ints.into_iter() + .map(|v| { + CellValue::from_int_literal(v, dtype).map_err(|e| ArrayError::SegmentCorruption { + detail: format!("attr col: {e}"), + }) + }) + .collect() +} diff --git a/nodedb-array/src/codec/attr_col/mod.rs b/nodedb-array/src/codec/attr_col/mod.rs new file mode 100644 index 000000000..c1266a4c5 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 + +pub mod column; +mod float; +mod int; +mod validity; +mod varlen; + +pub use column::{decode_attr_col, encode_attr_col}; diff --git a/nodedb-array/src/codec/attr_col/validity.rs b/nodedb-array/src/codec/attr_col/validity.rs new file mode 100644 index 000000000..8475f2da4 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/validity.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 + +// Validity section of an attr column. +// +// [u8 has_nulls] +// [ceil(count / 8) bytes bitmap] only when has_nulls == 1 +// +// Bit `i` (LSB first within each byte) is set when row `i` holds a value. +// A column with no NULL writes no bitmap. + +use crate::error::{ArrayError, ArrayResult}; +use crate::types::cell_value::value::CellValue; + +/// Write the validity section of `values` into `out`. +pub(super) fn encode_validity(values: &[CellValue], out: &mut Vec) { + if !values.iter().any(CellValue::is_null) { + out.push(0); + return; + } + out.push(1); + let mut bitmap = vec![0u8; values.len().div_ceil(8)]; + for (i, v) in values.iter().enumerate() { + if !v.is_null() { + bitmap[i / 8] |= 1 << (i % 8); + } + } + out.extend_from_slice(&bitmap); +} + +/// Which rows of a decoded column hold a value. +pub(super) enum Validity<'a> { + /// Every row holds a value. + AllPresent, + /// Bit `row` of the borrowed bitmap is set when that row holds a value. + Bitmap(&'a [u8]), +} + +impl Validity<'_> { + /// Whether `row` holds a value. A row past the bitmap holds none. + pub(super) fn is_present(&self, row: usize) -> bool { + match self { + Self::AllPresent => true, + Self::Bitmap(bits) => bits + .get(row / 8) + .is_some_and(|byte| byte & (1 << (row % 8)) != 0), + } + } +} + +/// Read the validity section of a `count`-row column at `*pos`. The bitmap is +/// borrowed from `data`, so a hostile count allocates nothing here. +pub(super) fn decode_validity<'a>( + data: &'a [u8], + pos: &mut usize, + count: usize, +) -> ArrayResult> { + let has_nulls = *data + .get(*pos) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: "attr col: truncated validity flag".into(), + })?; + *pos += 1; + match has_nulls { + 0 => Ok(Validity::AllPresent), + 1 => { + let len = count.div_ceil(8); + let end = pos + .checked_add(len) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: "attr col: validity bitmap length overflows".into(), + })?; + let bits = data + .get(*pos..end) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: format!("attr col: validity bitmap needs {len} bytes"), + })?; + *pos = end; + Ok(Validity::Bitmap(bits)) + } + other => Err(ArrayError::SegmentCorruption { + detail: format!("attr col: validity flag {other} is neither 0 nor 1"), + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn no_null_writes_no_bitmap() { + let mut out = Vec::new(); + encode_validity(&[CellValue::Int64(1), CellValue::Int64(2)], &mut out); + assert_eq!(out, vec![0]); + let mut pos = 0; + let validity = decode_validity(&out, &mut pos, 2).unwrap(); + assert!(matches!(validity, Validity::AllPresent)); + assert_eq!(pos, 1); + } + + #[test] + fn nulls_round_trip_through_the_bitmap() { + let values: Vec = (0..10) + .map(|i| { + if i % 3 == 0 { + CellValue::Null + } else { + CellValue::Int64(i) + } + }) + .collect(); + let mut out = Vec::new(); + encode_validity(&values, &mut out); + assert_eq!(out.len(), 1 + 2); + let mut pos = 0; + let validity = decode_validity(&out, &mut pos, values.len()).unwrap(); + let flags: Vec = (0..values.len()) + .map(|row| validity.is_present(row)) + .collect(); + let want: Vec = values.iter().map(|v| !v.is_null()).collect(); + assert_eq!(flags, want); + } + + #[test] + fn a_truncated_bitmap_is_corruption() { + let mut pos = 0; + assert!(decode_validity(&[1, 0], &mut pos, 20).is_err()); + let mut pos = 0; + assert!(decode_validity(&[7], &mut pos, 1).is_err()); + } +} diff --git a/nodedb-array/src/codec/attr_col/varlen.rs b/nodedb-array/src/codec/attr_col/varlen.rs new file mode 100644 index 000000000..879f20cf5 --- /dev/null +++ b/nodedb-array/src/codec/attr_col/varlen.rs @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 + +// String and Bytes attr payload: per non-null value, [u32 LE len][bytes]. + +use crate::codec::limits::{MAX_COLUMN_ENTRIES, checked_capacity}; +use crate::error::{ArrayError, ArrayResult}; +use crate::schema::attr_spec::AttrType; +use crate::types::cell_value::value::CellValue; + +pub(super) fn encode_varlen(present: &[&CellValue]) -> ArrayResult> { + let mut out = Vec::new(); + for v in present { + let bytes = match v { + CellValue::String(s) => s.as_bytes(), + CellValue::Bytes(b) => b.as_slice(), + other => { + return Err(ArrayError::InvalidOp { + detail: format!("variable-length attr column holds {other:?}"), + }); + } + }; + out.extend_from_slice(&(bytes.len() as u32).to_le_bytes()); + out.extend_from_slice(bytes); + } + Ok(out) +} + +/// Decode `expected` values of `dtype`, which is `String` or `Bytes`. +pub(super) fn decode_varlen( + body: &[u8], + dtype: AttrType, + expected: usize, +) -> ArrayResult> { + let capacity = checked_capacity( + expected, + MAX_COLUMN_ENTRIES, + body.len(), + 4, + "attr_col varlen count", + )?; + let mut values = Vec::with_capacity(capacity); + let mut pos = 0usize; + for _ in 0..expected { + let len_bytes = body + .get(pos..pos + 4) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: "attr col varlen: truncated entry len".into(), + })?; + let len = + u32::from_le_bytes([len_bytes[0], len_bytes[1], len_bytes[2], len_bytes[3]]) as usize; + pos += 4; + let bytes = body + .get(pos..pos + len) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: "attr col varlen: truncated entry bytes".into(), + })?; + pos += len; + let value = match dtype { + AttrType::String => { + let s = std::str::from_utf8(bytes).map_err(|e| ArrayError::SegmentCorruption { + detail: format!("attr col STRING: {e}"), + })?; + CellValue::String(s.to_string()) + } + AttrType::Bytes => CellValue::Bytes(bytes.to_vec()), + AttrType::Int64 + | AttrType::Float64 + | AttrType::UInt8 + | AttrType::Int16 + | AttrType::UInt16 + | AttrType::Int32 + | AttrType::UInt32 + | AttrType::Float32 => { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col varlen: tag names {}", dtype.sql_name()), + }); + } + }; + values.push(value); + } + if pos != body.len() { + return Err(ArrayError::SegmentCorruption { + detail: format!("attr col varlen: {} trailing bytes", body.len() - pos), + }); + } + Ok(values) +} diff --git a/nodedb-array/src/codec/column_codec.rs b/nodedb-array/src/codec/column_codec.rs index 66e986b00..3dc31dfeb 100644 --- a/nodedb-array/src/codec/column_codec.rs +++ b/nodedb-array/src/codec/column_codec.rs @@ -1,19 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 -// Per-column codecs for SparseTile scalar columns. +// Per-column codecs for SparseTile structural columns. // // surrogates (Vec) → fastlanes (u32 as i64) // row_kinds (Vec) → raw bytes (sentinel-dominated, not compressible) // *_ms cols (Vec) → gorilla timestamp encoding -// attr cols (Vec) → type-dispatch: fastlanes for Int64/Float64, -// raw zerompk for String/Bytes/Null +// +// Attr columns live in `attr_col`. use nodedb_codec::error::CodecError; use nodedb_types::Surrogate; -use crate::codec::limits::{MAX_COLUMN_ENTRIES, checked_capacity}; use crate::error::{ArrayError, ArrayResult}; -use crate::types::cell_value::value::CellValue; // --------------------------------------------------------------------------- // Error conversion @@ -84,146 +82,6 @@ pub fn decode_timestamps_col(data: &[u8]) -> ArrayResult> { nodedb_codec::gorilla::decode_timestamps(data).map_err(codec_err) } -// --------------------------------------------------------------------------- -// Attribute columns: type-dispatched -// --------------------------------------------------------------------------- - -/// Tag byte for each attr column format variant. -const ATTR_TAG_INT64: u8 = 0; -const ATTR_TAG_FLOAT64: u8 = 1; -const ATTR_TAG_MSGPACK: u8 = 2; // String, Bytes, Null — zerompk per-value - -pub fn encode_attr_col(values: &[CellValue]) -> ArrayResult> { - if values.is_empty() { - let mut out = vec![ATTR_TAG_MSGPACK]; - out.extend_from_slice(&0u32.to_le_bytes()); - return Ok(out); - } - - // Check if all values are Int64 or Float64 — only then use numeric codec. - let all_int = values - .iter() - .all(|v| matches!(v, CellValue::Int64(_) | CellValue::Null)); - let all_float = values - .iter() - .all(|v| matches!(v, CellValue::Float64(_) | CellValue::Null)); - - if all_int { - let ints: Vec = values - .iter() - .map(|v| match v { - CellValue::Int64(i) => *i, - _ => 0, - }) - .collect(); - let encoded = nodedb_codec::fastlanes::encode(&ints).map_err(codec_err)?; - let mut out = vec![ATTR_TAG_INT64]; - out.extend_from_slice(&encoded); - return Ok(out); - } - - if all_float { - // Gorilla XOR-encodes f64 series — exploits common-prefix bits across - // adjacent values, which dominates the size of monotonic / smooth - // numeric columns. fastlanes-on-bits used to live here but treats - // each f64 as an independent i64, missing the inter-value redundancy. - let floats: Vec = values - .iter() - .map(|v| match v { - CellValue::Float64(f) => *f, - _ => 0.0, - }) - .collect(); - let encoded = nodedb_codec::gorilla::encode_f64(&floats); - let mut out = vec![ATTR_TAG_FLOAT64]; - out.extend_from_slice(&encoded); - return Ok(out); - } - - // Generic zerompk fallback for String / Bytes / mixed / Null. - let mut out = vec![ATTR_TAG_MSGPACK]; - out.extend_from_slice(&(values.len() as u32).to_le_bytes()); - for v in values { - let bytes = zerompk::to_msgpack_vec(v).map_err(|e| ArrayError::SegmentCorruption { - detail: format!("attr col encode: {e}"), - })?; - out.extend_from_slice(&(bytes.len() as u32).to_le_bytes()); - out.extend_from_slice(&bytes); - } - Ok(out) -} - -pub fn decode_attr_col(data: &[u8]) -> ArrayResult> { - if data.is_empty() { - return Err(ArrayError::SegmentCorruption { - detail: "attr col: empty payload".into(), - }); - } - let tag = data[0]; - let body = &data[1..]; - - match tag { - ATTR_TAG_INT64 => { - let ints = nodedb_codec::fastlanes::decode(body).map_err(codec_err)?; - Ok(ints.into_iter().map(CellValue::Int64).collect()) - } - ATTR_TAG_FLOAT64 => { - let floats = nodedb_codec::gorilla::decode_f64(body).map_err(codec_err)?; - Ok(floats.into_iter().map(CellValue::Float64).collect()) - } - ATTR_TAG_MSGPACK => { - if body.len() < 4 { - return Err(ArrayError::SegmentCorruption { - detail: "attr col msgpack: truncated count".into(), - }); - } - let count = u32::from_le_bytes( - body[0..4] - .try_into() - .expect("invariant: bounds-checked above (body.len() >= 4)"), - ) as usize; - let capacity = checked_capacity( - count, - MAX_COLUMN_ENTRIES, - body.len() - 4, - 4, - "attr_col_msgpack count", - )?; - let mut pos = 4; - let mut values = Vec::with_capacity(capacity); - for _ in 0..count { - if pos + 4 > body.len() { - return Err(ArrayError::SegmentCorruption { - detail: "attr col msgpack: truncated entry len".into(), - }); - } - let len = u32::from_le_bytes( - body[pos..pos + 4] - .try_into() - .expect("invariant: bounds-checked above (pos + 4 <= body.len())"), - ) as usize; - pos += 4; - if pos + len > body.len() { - return Err(ArrayError::SegmentCorruption { - detail: "attr col msgpack: truncated entry bytes".into(), - }); - } - let v: CellValue = zerompk::from_msgpack(&body[pos..pos + len]).map_err(|e| { - ArrayError::SegmentCorruption { - detail: format!("attr col decode: {e}"), - } - })?; - pos += len; - values.push(v); - } - Ok(values) - } - other => Err(ArrayError::SegmentCorruption { - detail: format!("attr col: unknown tag {other:#04x}"), - }), - } -} - #[cfg(test)] mod tests { use super::*; @@ -275,67 +133,6 @@ mod tests { assert_eq!(out, ts); } - #[test] - fn attr_col_int64_roundtrip() { - let vals = vec![ - CellValue::Int64(10), - CellValue::Int64(-5), - CellValue::Int64(0), - ]; - let data = encode_attr_col(&vals).unwrap(); - let out = decode_attr_col(&data).unwrap(); - assert_eq!(out, vals); - } - - #[test] - fn attr_col_float64_roundtrip() { - let vals = vec![CellValue::Float64(1.5), CellValue::Float64(-2.5)]; - let data = encode_attr_col(&vals).unwrap(); - let out = decode_attr_col(&data).unwrap(); - assert_eq!(out, vals); - } - - #[test] - fn attr_col_string_roundtrip() { - let vals = vec![ - CellValue::String("hello".into()), - CellValue::String("world".into()), - ]; - let data = encode_attr_col(&vals).unwrap(); - let out = decode_attr_col(&data).unwrap(); - assert_eq!(out, vals); - } - - #[test] - fn huge_msgpack_count_in_tiny_payload_is_rejected_before_allocation() { - let mut data = vec![ATTR_TAG_MSGPACK]; - data.extend_from_slice(&(MAX_COLUMN_ENTRIES as u32).to_le_bytes()); - let result = decode_attr_col(&data); - assert!(matches!( - result, - Err(crate::error::ArrayError::SegmentCorruption { .. }) - )); - } - - #[test] - fn attr_col_empty_roundtrip() { - let data = encode_attr_col(&[]).unwrap(); - let out = decode_attr_col(&data).unwrap(); - assert!(out.is_empty()); - } - - #[test] - fn attr_col_mixed_types_roundtrip() { - let vals = vec![ - CellValue::String("x".into()), - CellValue::Null, - CellValue::Bytes(vec![1, 2, 3]), - ]; - let data = encode_attr_col(&vals).unwrap(); - let out = decode_attr_col(&data).unwrap(); - assert_eq!(out, vals); - } - #[test] fn surrogates_large_roundtrip() { let vals: Vec = (0u32..1000).map(|i| Surrogate::new(i * 7)).collect(); diff --git a/nodedb-array/src/codec/mod.rs b/nodedb-array/src/codec/mod.rs index cce535f31..6d3df2e66 100644 --- a/nodedb-array/src/codec/mod.rs +++ b/nodedb-array/src/codec/mod.rs @@ -3,6 +3,7 @@ // Structural sparse-tile codec. Dense tiles retain zerompk serialization; // this module covers sparse tiles only. See tile_encode / tile_decode for // the entry points used by the segment writer and reader. +pub mod attr_col; pub mod column_codec; pub mod coord_delta; pub mod coord_rle; @@ -11,9 +12,10 @@ pub mod tag; pub mod tile_decode; pub mod tile_encode; +pub use attr_col::{decode_attr_col, encode_attr_col}; pub use column_codec::{ - decode_attr_col, decode_row_kinds, decode_surrogates, decode_timestamps_col, encode_attr_col, - encode_row_kinds, encode_surrogates, encode_timestamps_col, + decode_row_kinds, decode_surrogates, decode_timestamps_col, encode_row_kinds, + encode_surrogates, encode_timestamps_col, }; pub use coord_delta::{decode_coord_axis, encode_coord_axis}; pub use coord_rle::{decode_coord_axis_rle, encode_coord_axis_rle}; diff --git a/nodedb-array/src/codec/tile_decode.rs b/nodedb-array/src/codec/tile_decode.rs index 8ba84ecb4..6d597b701 100644 --- a/nodedb-array/src/codec/tile_decode.rs +++ b/nodedb-array/src/codec/tile_decode.rs @@ -7,17 +7,18 @@ // NOT handled here; the segment reader peeks the tag and dispatches to // zerompk for those. -use crate::codec::column_codec::{ - decode_attr_col, decode_row_kinds, decode_surrogates, decode_timestamps_col, -}; +use crate::codec::attr_col::decode_attr_col; +use crate::codec::column_codec::{decode_row_kinds, decode_surrogates, decode_timestamps_col}; use crate::codec::coord_rle::decode_coord_axis_rle; use crate::codec::limits::{ MAX_ATTRS_PER_TILE, MAX_AXES_PER_TILE, MAX_CELLS_PER_TILE, check_decoded_size, checked_capacity, }; use crate::codec::tag::{CodecTag, peek_tag}; use crate::error::{ArrayError, ArrayResult}; -use crate::tile::mbr::TileMBR; -use crate::tile::sparse_tile::{RowKind, SparseTile}; +use crate::tile::mbr::{MbrBuilder, TileMBR}; +use crate::tile::sparse_tile::{DimDict, RowKind, SparseTile}; +use crate::types::cell_value::value::CellValue; +use crate::types::coord::value::CoordValue; const SUPPORTED_PAYLOAD_VERSION: u8 = 1; @@ -49,8 +50,23 @@ fn read_framed<'a>(data: &'a [u8], pos: &mut usize) -> ArrayResult<&'a [u8]> { /// Decode a tile payload previously written by `encode_sparse_tile`. /// /// The `payload` slice must start at the tag byte (i.e. after BlockFraming -/// has been unwrapped by the segment reader). +/// has been unwrapped by the segment reader). A structural payload stores no +/// MBR, so this rebuilds it from the rows. A segment reader holds the MBR in +/// the tile's footer entry and decodes with [`decode_sparse_tile_with_mbr`]. pub fn decode_sparse_tile(payload: &[u8]) -> ArrayResult { + decode_sparse_tile_inner(payload, None) +} + +/// Decode a tile payload whose MBR is `stored_mbr`, the one its segment +/// footer entry holds. A raw payload carries its own MBR and keeps it. +pub fn decode_sparse_tile_with_mbr(payload: &[u8], stored_mbr: TileMBR) -> ArrayResult { + decode_sparse_tile_inner(payload, Some(stored_mbr)) +} + +fn decode_sparse_tile_inner( + payload: &[u8], + stored_mbr: Option, +) -> ArrayResult { if payload.len() < 2 { return Err(ArrayError::SegmentCorruption { detail: "sparse tile payload too short".into(), @@ -77,7 +93,7 @@ pub fn decode_sparse_tile(payload: &[u8]) -> ArrayResult { match tag_result { CodecTag::Raw => decode_raw(&payload[2..]), - CodecTag::Structural => decode_structural(&payload[2..]), + CodecTag::Structural => decode_structural(&payload[2..], stored_mbr), } } @@ -90,7 +106,7 @@ fn decode_raw(body: &[u8]) -> ArrayResult { Ok(tile) } -fn decode_structural(body: &[u8]) -> ArrayResult { +fn decode_structural(body: &[u8], stored_mbr: Option) -> ArrayResult { let mut pos = 0; if pos + 8 > body.len() { @@ -167,19 +183,18 @@ fn decode_structural(body: &[u8]) -> ArrayResult { "attr_count", )?; + // An attr column holds one value per live row. + let live_rows = row_kinds + .iter() + .filter(|&&kind| kind == RowKind::Live.as_u8()) + .count(); let mut attr_cols = Vec::with_capacity(attr_capacity); for _ in 0..attr_count { let col_bytes = read_framed(body, &mut pos)?; - let col = decode_attr_col(col_bytes)?; + let col = decode_attr_col(col_bytes, live_rows)?; attr_cols.push(col); } - // Reconstruct MBR from decoded data — we don't persist the MBR inside the - // structural payload (it lives in the footer TileEntry instead). - // Use a zero-dimension MBR so the tile can be used for reads; compaction - // re-derives the full MBR if needed. - let mbr = TileMBR::new(axis_count, attr_count); - // Validate sizes match cell_count. if row_kinds.len() != cell_count { return Err(ArrayError::SegmentCorruption { @@ -190,6 +205,13 @@ fn decode_structural(body: &[u8]) -> ArrayResult { }); } let surrogates = spread_live_surrogates(&row_kinds, live_surrogates)?; + // The structural payload stores no MBR. Readers count and bound a tile by + // its MBR: the segment footer's copy when the caller has it, else one + // rebuilt from the decoded rows. + let mbr = match stored_mbr { + Some(mbr) => mbr, + None => rebuild_mbr(&dim_dicts, &attr_cols, &row_kinds)?, + }; let tile = SparseTile { dim_dicts, @@ -204,6 +226,45 @@ fn decode_structural(body: &[u8]) -> ArrayResult { Ok(tile) } +/// The MBR the tile builder folds for these rows: every live row's coords +/// and attrs, and no sentinel row. Attr columns hold one value per live row. +fn rebuild_mbr( + dim_dicts: &[DimDict], + attr_cols: &[Vec], + row_kinds: &[u8], +) -> ArrayResult { + let mut mbr = MbrBuilder::new(dim_dicts.len(), attr_cols.len()); + let mut coord: Vec = Vec::with_capacity(dim_dicts.len()); + let mut attrs: Vec = Vec::with_capacity(attr_cols.len()); + let mut live = 0usize; + for (row, &kind) in row_kinds.iter().enumerate() { + if RowKind::from_u8(kind)? != RowKind::Live { + continue; + } + coord.clear(); + for dict in dim_dicts { + let value = dict + .indices + .get(row) + .and_then(|&i| dict.values.get(i as usize)) + .ok_or_else(|| ArrayError::SegmentCorruption { + detail: format!("structural tile: row {row} has no coordinate on an axis"), + })?; + coord.push(value.clone()); + } + attrs.clear(); + for col in attr_cols { + let value = col.get(live).ok_or_else(|| ArrayError::SegmentCorruption { + detail: format!("structural tile: live row {live} has no value in an attr column"), + })?; + attrs.push(value.clone()); + } + mbr.fold(&coord, &attrs); + live += 1; + } + Ok(mbr.build()) +} + /// Place the live-row surrogate column back onto every row: a live row takes /// the next surrogate, a tombstone or erasure row holds none. The column must /// hold exactly one surrogate per live row. @@ -381,7 +442,7 @@ mod tests { let mut body = Vec::new(); body.extend_from_slice(&0u32.to_le_bytes()); body.extend_from_slice(&(MAX_AXES_PER_TILE as u32).to_le_bytes()); - let error = decode_structural(&body).unwrap_err(); + let error = decode_structural(&body, None).unwrap_err(); assert!(matches!(error, ArrayError::SegmentCorruption { .. })); } diff --git a/nodedb-array/src/codec/tile_encode.rs b/nodedb-array/src/codec/tile_encode.rs index 95abb069d..6128e47a3 100644 --- a/nodedb-array/src/codec/tile_encode.rs +++ b/nodedb-array/src/codec/tile_encode.rs @@ -16,15 +16,14 @@ // [u32 LE valid_from_ms_len][gorilla payload] // [u32 LE valid_until_ms_len][gorilla payload] // [u32 LE attr_count] -// per attr: [u32 LE col_len][column_codec payload] +// per attr: [u32 LE col_len][attr_col payload] // // For CodecTag::Raw, the full SparseTile is serialised with zerompk and // written verbatim after the two header bytes (tag + version). The decoder // mirrors this. -use crate::codec::column_codec::{ - encode_attr_col, encode_row_kinds, encode_surrogates, encode_timestamps_col, -}; +use crate::codec::attr_col::encode_attr_col; +use crate::codec::column_codec::{encode_row_kinds, encode_surrogates, encode_timestamps_col}; use crate::codec::coord_rle::encode_coord_axis_rle; use crate::codec::tag::CodecTag; use crate::error::{ArrayError, ArrayResult}; @@ -212,6 +211,10 @@ mod tests { assert_eq!(decoded.surrogates, tile.surrogates); assert_eq!(decoded.attr_cols, tile.attr_cols); assert_eq!(decoded.dim_dicts.len(), tile.dim_dicts.len()); + // Readers count and bound a tile by its MBR, so a structural tile + // decodes with the MBR it was built with. + assert_eq!(decoded.nnz(), 100); + assert_eq!(decoded.mbr, tile.mbr); } #[test] @@ -284,6 +287,38 @@ mod tests { assert_eq!(decoded.surrogates, tile.surrogates); assert_eq!(decoded.surrogates[0], None); assert_eq!(decoded.surrogates[1], Some(Surrogate::new(101))); + // The MBR folds live rows only, as the builder does. + assert_eq!(decoded.mbr, tile.mbr); + assert_eq!(decoded.nnz(), 13); + } + + /// A NULL in a numeric attr survives the structural codec as NULL. + #[test] + fn structural_tile_keeps_numeric_nulls() { + let s = schema(); + let mut b = SparseTileBuilder::new(&s); + for i in 0..20i64 { + let attr = if i % 2 == 0 { + CellValue::Null + } else { + CellValue::Int64(i) + }; + b.push_row(SparseRow::live( + &[CoordValue::Int64(i), CoordValue::Int64(i)], + &[attr], + Surrogate::new(1 + i as u32), + 0, + OPEN_UPPER, + )) + .unwrap(); + } + let tile = b.build(); + let mut buf = Vec::new(); + encode_sparse_tile(&tile, &mut buf).unwrap(); + assert_eq!(buf[0], CodecTag::Structural.as_byte()); + let decoded = decode_sparse_tile(&buf).unwrap(); + assert_eq!(decoded.attr_cols, tile.attr_cols); + assert_eq!(decoded.attr_cols[0][0], CellValue::Null); } /// A derived result tile holds no identity, so it is never written. diff --git a/nodedb-array/src/coord/normalize.rs b/nodedb-array/src/coord/normalize.rs index b1096457c..fa7c57241 100644 --- a/nodedb-array/src/coord/normalize.rs +++ b/nodedb-array/src/coord/normalize.rs @@ -51,6 +51,26 @@ pub fn normalize_coord( Ok(out) } +/// Check one cell coordinate as [`normalize_coord`] does, with no allocation: +/// the arity, each value's type, and each value against its dim's domain. +pub fn check_coord_domain( + schema: &ArraySchema, + components: &[CoordValue], + bits: u32, +) -> ArrayResult<()> { + if components.len() != schema.arity() { + return Err(ArrayError::CoordArityMismatch { + array: schema.name.clone(), + expected: schema.arity(), + got: components.len(), + }); + } + for (dim, value) in schema.dims.iter().zip(components.iter()) { + normalize_one(&schema.name, dim, value, bits)?; + } + Ok(()) +} + fn normalize_one(array: &str, dim: &DimSpec, value: &CoordValue, bits: u32) -> ArrayResult { let max = if bits >= 64 { u64::MAX diff --git a/nodedb-array/src/error.rs b/nodedb-array/src/error.rs index 705bdbfbc..5d17ce8cc 100644 --- a/nodedb-array/src/error.rs +++ b/nodedb-array/src/error.rs @@ -9,6 +9,8 @@ use nodedb_types::NodeDbError; use thiserror::Error; +use crate::types::cell_value::NarrowError; + /// Crate-internal result alias. pub type ArrayResult = std::result::Result; @@ -60,6 +62,19 @@ pub enum ArrayError { detail: String, }, + /// A literal does not fit its attribute's declared type. + #[error("attribute '{attr}' on '{array}': {source}")] + Narrow { + array: String, + attr: String, + source: NarrowError, + }, + + /// Integer arithmetic over cells left the `i64` range, such as an + /// exact `SUM` or an elementwise product. SQLSTATE `22003`. + #[error("numeric value out of range: {detail}")] + NumericOverflow { detail: String }, + #[error("segment corruption: {detail}")] SegmentCorruption { detail: String }, @@ -136,8 +151,10 @@ impl ArrayError { | ArrayError::InvalidTileExtents { array, .. } | ArrayError::CoordArityMismatch { array, .. } | ArrayError::CoordOutOfDomain { array, .. } - | ArrayError::CellTypeMismatch { array, .. } => array, - ArrayError::SegmentCorruption { .. } + | ArrayError::CellTypeMismatch { array, .. } + | ArrayError::Narrow { array, .. } => array, + ArrayError::NumericOverflow { .. } + | ArrayError::SegmentCorruption { .. } | ArrayError::UnsupportedFormat { .. } | ArrayError::UnsupportedSegmentFormat { .. } | ArrayError::InvalidReplicaId { .. } diff --git a/nodedb-array/src/lib.rs b/nodedb-array/src/lib.rs index 4a4d0daaa..08b35c6be 100644 --- a/nodedb-array/src/lib.rs +++ b/nodedb-array/src/lib.rs @@ -36,7 +36,7 @@ pub use sync::{ SnapshotHeader, SnapshotSink, TileSnapshot, }; pub use tile::{ - AttrStats, DenseTile, SparseTile, TileMBR, cell_tile_prefix, tile_id_for_cell, + AttrStats, DenseTile, SparseTile, TileMBR, TilePrefixMemo, cell_tile_prefix, tile_id_for_cell, tile_indices_for_cell, tile_prefix, tiles_per_dim, }; pub use types::{ArrayId, CellValue, Coord, Domain, TileId}; diff --git a/nodedb-array/src/query/agg_number.rs b/nodedb-array/src/query/agg_number.rs new file mode 100644 index 000000000..6f40e3211 --- /dev/null +++ b/nodedb-array/src/query/agg_number.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Running and final numbers of the array reducers. +//! +//! An integer attr accumulates exactly in `i128`, so `SUM` of `INT64` +//! cells past 2^53 keeps every unit. A float attr accumulates in `f64`. +//! `MIN` / `MAX` follow PostgreSQL: NaN sorts above every other float. + +use std::cmp::Ordering; + +use crate::error::{ArrayError, ArrayResult}; +use crate::types::cell_value::value::CellValue; + +/// A running reducer number. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum AggNumber { + Int(i128), + Float(f64), +} + +/// A finalized reducer value. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum AggScalar { + Int(i64), + Float(f64), +} + +impl AggNumber { + /// The number a numeric cell contributes. `None` for a NULL, String, or + /// Bytes cell. + pub fn of_cell(v: &CellValue) -> Option { + match v.as_i64() { + Some(i) => Some(AggNumber::Int(i128::from(i))), + None => v.as_float().map(AggNumber::Float), + } + } + + pub fn to_f64(self) -> f64 { + match self { + AggNumber::Int(v) => v as f64, + AggNumber::Float(v) => v, + } + } + + /// The sum of two numbers. An `i128` sum saturates, which is far past + /// the `i64` range a final `SUM` must fit, so it still refuses. + pub fn sum_with(self, other: AggNumber) -> AggNumber { + match (self, other) { + (AggNumber::Int(a), AggNumber::Int(b)) => AggNumber::Int(a.saturating_add(b)), + (a, b) => AggNumber::Float(a.to_f64() + b.to_f64()), + } + } + + /// The smaller number. NaN loses to every other value. + pub fn min_with(self, other: AggNumber) -> AggNumber { + match cmp_nan_high(self, other) { + Ordering::Greater => other, + Ordering::Less | Ordering::Equal => self, + } + } + + /// The larger number. NaN wins over every other value. + pub fn max_with(self, other: AggNumber) -> AggNumber { + match cmp_nan_high(self, other) { + Ordering::Less => other, + Ordering::Greater | Ordering::Equal => self, + } + } + + /// The value as a final `i64` or `f64`. An integer past the `i64` + /// range is refused as SQLSTATE `22003`. + pub fn finish(self) -> ArrayResult { + match self { + AggNumber::Int(v) => { + i64::try_from(v) + .map(AggScalar::Int) + .map_err(|_| ArrayError::NumericOverflow { + detail: format!("aggregate value {v} is out of range for bigint"), + }) + } + AggNumber::Float(v) => Ok(AggScalar::Float(v)), + } + } +} + +/// Total order with NaN above every other value, as PostgreSQL sorts floats. +fn cmp_nan_high(a: AggNumber, b: AggNumber) -> Ordering { + match (a, b) { + (AggNumber::Int(x), AggNumber::Int(y)) => x.cmp(&y), + (x, y) => { + let (x, y) = (x.to_f64(), y.to_f64()); + match (x.is_nan(), y.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => x.partial_cmp(&y).unwrap_or(Ordering::Equal), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn integer_sum_is_exact_past_2_pow_53() { + let big = AggNumber::Int(1 << 53); + let sum = big.sum_with(AggNumber::Int(1)); + assert_eq!(sum.finish().unwrap(), AggScalar::Int((1 << 53) + 1)); + } + + #[test] + fn integer_sum_past_i64_is_refused() { + let sum = AggNumber::Int(i128::from(i64::MAX)).sum_with(AggNumber::Int(1)); + assert!(matches!( + sum.finish(), + Err(ArrayError::NumericOverflow { .. }) + )); + } + + #[test] + fn nan_sorts_above_every_float() { + let nan = AggNumber::Float(f64::NAN); + let one = AggNumber::Float(1.0); + assert!(nan.max_with(one).to_f64().is_nan()); + assert!(one.max_with(nan).to_f64().is_nan()); + assert_eq!(nan.min_with(one), one); + assert_eq!(one.min_with(nan), one); + } + + #[test] + fn narrow_cells_contribute_exact_integers() { + assert_eq!( + AggNumber::of_cell(&CellValue::UInt16(65535)), + Some(AggNumber::Int(65535)) + ); + assert_eq!( + AggNumber::of_cell(&CellValue::Float32(0.5)), + Some(AggNumber::Float(0.5)) + ); + assert_eq!(AggNumber::of_cell(&CellValue::Null), None); + } +} diff --git a/nodedb-array/src/query/aggregate.rs b/nodedb-array/src/query/aggregate.rs index 1aa8adc41..e9be5ae76 100644 --- a/nodedb-array/src/query/aggregate.rs +++ b/nodedb-array/src/query/aggregate.rs @@ -15,6 +15,8 @@ use std::collections::HashMap; +use crate::error::{ArrayError, ArrayResult}; +use crate::query::agg_number::{AggNumber, AggScalar}; use crate::tile::sparse_tile::{RowKind, SparseTile}; use crate::types::cell_value::value::CellValue; use crate::types::coord::value::CoordValue; @@ -29,26 +31,27 @@ pub enum Reducer { } /// Partial result for one reducer over one (group of) cell(s). -/// `Mean` retains `(sum, count)` so partials merge exactly. +/// `Mean` retains `(sum, count)` so partials merge exactly. An integer attr +/// accumulates exactly; see [`AggNumber`]. #[derive(Debug, Clone, Copy, PartialEq)] pub enum AggregateResult { Sum { - value: f64, + value: AggNumber, count: u64, }, Count { count: u64, }, Min { - value: f64, + value: AggNumber, count: u64, }, Max { - value: f64, + value: AggNumber, count: u64, }, Mean { - sum: f64, + sum: AggNumber, count: u64, }, /// No non-null cells observed yet for this reducer. @@ -73,7 +76,7 @@ impl AggregateResult { count: cb, }, ) => Sum { - value: a + b, + value: a.sum_with(b), count: ca + cb, }, (Count { count: a }, Count { count: b }) => Count { count: a + b }, @@ -87,7 +90,7 @@ impl AggregateResult { count: cb, }, ) => Min { - value: if a <= b { a } else { b }, + value: a.min_with(b), count: ca + cb, }, ( @@ -100,25 +103,35 @@ impl AggregateResult { count: cb, }, ) => Max { - value: if a >= b { a } else { b }, + value: a.max_with(b), count: ca + cb, }, (Mean { sum: a, count: ca }, Mean { sum: b, count: cb }) => Mean { - sum: a + b, + sum: a.sum_with(b), count: ca + cb, }, (lhs, _) => lhs, } } - /// Final scalar — useful for tests and the planner's last fold. - pub fn finalize(self) -> Option { + /// Final value, or `None` when no cell contributed. `SUM` of an integer + /// attr is an exact integer and refuses past `i64` (SQLSTATE `22003`). + /// `MIN` / `MAX` keep the attr's kind. `COUNT` is an integer. `MEAN` is + /// always a float. + pub fn finalize(self) -> ArrayResult> { match self { - AggregateResult::Sum { value, .. } => Some(value), - AggregateResult::Count { count } => Some(count as f64), - AggregateResult::Min { value, .. } | AggregateResult::Max { value, .. } => Some(value), - AggregateResult::Mean { sum, count } if count > 0 => Some(sum / count as f64), - _ => None, + AggregateResult::Sum { value, .. } + | AggregateResult::Min { value, .. } + | AggregateResult::Max { value, .. } => value.finish().map(Some), + AggregateResult::Count { count } => Ok(Some(AggScalar::Int( + i64::try_from(count).map_err(|_| ArrayError::NumericOverflow { + detail: format!("count {count} is out of range for bigint"), + })?, + ))), + AggregateResult::Mean { sum, count } if count > 0 => { + Ok(Some(AggScalar::Float(sum.to_f64() / count as f64))) + } + AggregateResult::Mean { .. } | AggregateResult::Empty(_) => Ok(None), } } } @@ -145,11 +158,7 @@ pub fn aggregate_attr(tile: &SparseTile, attr_idx: usize, reducer: Reducer) -> A fn single_cell(v: &CellValue, reducer: Reducer) -> AggregateResult { use AggregateResult::*; - let n = match v { - CellValue::Int64(x) => Some(*x as f64), - CellValue::Float64(x) => Some(*x), - CellValue::String(_) | CellValue::Bytes(_) | CellValue::Null => None, - }; + let n = AggNumber::of_cell(v); match (reducer, n) { (Reducer::Count, _) if !v.is_null() => Count { count: 1 }, (Reducer::Count, _) => Empty(Reducer::Count), @@ -254,44 +263,134 @@ mod tests { b.build() } + fn fin(r: AggregateResult) -> Option { + r.finalize().expect("finalize") + } + + fn float(v: f64) -> Option { + Some(AggScalar::Float(v)) + } + + fn int_tile(dtype: AttrType, values: &[CellValue]) -> SparseTile { + let s = ArraySchemaBuilder::new("n") + .dim(DimSpec::new( + "k", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("v", dtype, true)) + .tile_extents(vec![16]) + .build() + .unwrap(); + let mut b = SparseTileBuilder::new(&s); + for (k, v) in values.iter().enumerate() { + b.push(&[CoordValue::Int64(k as i64)], std::slice::from_ref(v)) + .unwrap(); + } + b.build() + } + #[test] fn sum_skips_nulls() { let t = tile(&[(0, Some(1.0)), (1, None), (2, Some(3.0))]); let r = aggregate_attr(&t, 0, Reducer::Sum); - assert_eq!(r.finalize(), Some(4.0)); + assert_eq!(fin(r), float(4.0)); } #[test] fn count_includes_nonnull_only() { let t = tile(&[(0, Some(1.0)), (1, None), (2, Some(3.0))]); let r = aggregate_attr(&t, 0, Reducer::Count); - assert_eq!(r.finalize(), Some(2.0)); + assert_eq!(fin(r), Some(AggScalar::Int(2))); } #[test] fn min_max_mean() { let t = tile(&[(0, Some(1.0)), (1, Some(5.0)), (2, Some(3.0))]); - assert_eq!(aggregate_attr(&t, 0, Reducer::Min).finalize(), Some(1.0)); - assert_eq!(aggregate_attr(&t, 0, Reducer::Max).finalize(), Some(5.0)); - assert_eq!(aggregate_attr(&t, 0, Reducer::Mean).finalize(), Some(3.0)); + assert_eq!(fin(aggregate_attr(&t, 0, Reducer::Min)), float(1.0)); + assert_eq!(fin(aggregate_attr(&t, 0, Reducer::Max)), float(5.0)); + assert_eq!(fin(aggregate_attr(&t, 0, Reducer::Mean)), float(3.0)); } #[test] fn merge_combines_partials_exactly() { let a = AggregateResult::Mean { - sum: 10.0, + sum: AggNumber::Float(10.0), count: 4, }; - let b = AggregateResult::Mean { sum: 6.0, count: 2 }; + let b = AggregateResult::Mean { + sum: AggNumber::Float(6.0), + count: 2, + }; // (10 + 6) / (4 + 2) = 16/6 - assert_eq!(a.merge(b).finalize(), Some(16.0 / 6.0)); + assert_eq!(fin(a.merge(b)), float(16.0 / 6.0)); } #[test] fn empty_reducer_finalizes_to_none() { let t = tile(&[(0, None)]); let r = aggregate_attr(&t, 0, Reducer::Sum); - assert_eq!(r.finalize(), None); + assert_eq!(fin(r), None); + } + + /// An `INT64` SUM keeps every unit past 2^53, and MIN / MAX stay exact + /// integers. MEAN is a float. + #[test] + fn integer_reducers_are_exact() { + let big = (1i64 << 53) + 1; + let t = int_tile( + AttrType::Int64, + &[CellValue::Int64(big), CellValue::Null, CellValue::Int64(0)], + ); + assert_eq!( + fin(aggregate_attr(&t, 0, Reducer::Sum)), + Some(AggScalar::Int(big)) + ); + assert_eq!( + fin(aggregate_attr(&t, 0, Reducer::Max)), + Some(AggScalar::Int(big)) + ); + assert_eq!( + fin(aggregate_attr(&t, 0, Reducer::Min)), + Some(AggScalar::Int(0)) + ); + assert!(matches!( + fin(aggregate_attr(&t, 0, Reducer::Mean)), + Some(AggScalar::Float(_)) + )); + } + + #[test] + fn integer_sum_past_bigint_is_refused() { + let t = int_tile( + AttrType::Int64, + &[CellValue::Int64(i64::MAX), CellValue::Int64(1)], + ); + assert!(matches!( + aggregate_attr(&t, 0, Reducer::Sum).finalize(), + Err(ArrayError::NumericOverflow { .. }) + )); + } + + #[test] + fn narrow_attrs_reduce() { + let t = int_tile( + AttrType::UInt16, + &[CellValue::UInt16(65535), CellValue::UInt16(65535)], + ); + assert_eq!( + fin(aggregate_attr(&t, 0, Reducer::Sum)), + Some(AggScalar::Int(131_070)) + ); + let f = int_tile( + AttrType::Float32, + &[CellValue::Float32(0.5), CellValue::Float32(f32::NAN)], + ); + assert!(matches!( + fin(aggregate_attr(&f, 0, Reducer::Max)), + Some(AggScalar::Float(v)) if v.is_nan() + )); + assert_eq!(fin(aggregate_attr(&f, 0, Reducer::Min)), float(0.5)); } #[test] @@ -306,8 +405,8 @@ mod tests { assert_eq!(g.len(), 2); // first-seen order: 0, then 1 assert_eq!(g[0].key, CoordValue::Int64(0)); - assert_eq!(g[0].result.finalize(), Some(4.0)); + assert_eq!(fin(g[0].result), float(4.0)); assert_eq!(g[1].key, CoordValue::Int64(1)); - assert_eq!(g[1].result.finalize(), Some(6.0)); + assert_eq!(fin(g[1].result), float(6.0)); } } diff --git a/nodedb-array/src/query/elementwise.rs b/nodedb-array/src/query/elementwise.rs index 207849755..368c475a1 100644 --- a/nodedb-array/src/query/elementwise.rs +++ b/nodedb-array/src/query/elementwise.rs @@ -7,10 +7,14 @@ //! in only one operand contribute a Null on the other side, which //! propagates as Null through the op (SQL-style null semantics). //! -//! Numeric attrs (Int64, Float64) participate in the op; non-numeric -//! attrs (String, Bytes) are passed through from the left operand -//! unchanged — the op is undefined for them. Division by zero yields -//! Null rather than infinity so downstream aggregates stay finite. +//! The op reads one named attr. Every other attr passes through from the +//! left operand. A non-numeric value (String, Bytes) in the named attr also +//! passes through, since the op is undefined for it. Division by zero yields +//! Null rather than infinity, so downstream aggregates stay finite. +//! +//! Integer Add, Sub, and Mul are exact `Int64` and refuse past the `i64` +//! range with SQLSTATE `22003`. Div, and any op over a float attr, yields +//! `Float64`. use std::collections::BTreeMap; @@ -28,15 +32,25 @@ pub enum BinaryOp { Div, } -/// Apply `op` to every aligned pair of cells. The result tile carries -/// the union of both coord sets. +/// Apply `op` to attr `attr_idx` of every aligned pair of cells. The result +/// tile carries the union of both coord sets. Every other attr passes +/// through, from the left operand when it has the cell. pub fn elementwise( schema: &ArraySchema, a: &SparseTile, b: &SparseTile, op: BinaryOp, + attr_idx: usize, ) -> ArrayResult { let n_attrs = schema.attrs.len(); + if attr_idx >= n_attrs { + return Err(ArrayError::InvalidOp { + detail: format!( + "elementwise: attr index {attr_idx} is out of range for {n_attrs} attrs of '{}'", + schema.name + ), + }); + } let by_coord_a = index_rows(a)?; let by_coord_b = index_rows(b)?; let mut keys: BTreeMap, ()> = BTreeMap::new(); @@ -53,17 +67,38 @@ pub fn elementwise( for i in 0..n_attrs { let l = lhs.map(|v| v[i].clone()).unwrap_or(CellValue::Null); let r = rhs.map(|v| v[i].clone()).unwrap_or(CellValue::Null); - out.push(apply(&l, &r, op)); + out.push(if i == attr_idx { + apply(&l, &r, op)? + } else { + passthrough(&l, &r) + }); } builder.push(&coord, &out)?; } Ok(builder.build()) } -fn apply(l: &CellValue, r: &CellValue, op: BinaryOp) -> CellValue { - let (lf, rf) = match (to_f64(l), to_f64(r)) { +/// Add, Sub, and Mul of two integer cells is an exact `Int64`, refused past +/// the `i64` range. Div, and any op with a float operand, is a `Float64`. +fn apply(l: &CellValue, r: &CellValue, op: BinaryOp) -> ArrayResult { + if let (Some(a), Some(b)) = (l.as_i64(), r.as_i64()) { + let exact = match op { + BinaryOp::Add => Some(a.checked_add(b)), + BinaryOp::Sub => Some(a.checked_sub(b)), + BinaryOp::Mul => Some(a.checked_mul(b)), + BinaryOp::Div => None, + }; + if let Some(v) = exact { + return v + .map(CellValue::Int64) + .ok_or_else(|| ArrayError::NumericOverflow { + detail: format!("elementwise {op:?} of {a} and {b} is out of range for bigint"), + }); + } + } + let (lf, rf) = match (l.as_f64(), r.as_f64()) { (Some(a), Some(b)) => (a, b), - _ => return passthrough(l, r), + _ => return Ok(widen_to_result(passthrough(l, r), op)), }; let v = match op { BinaryOp::Add => lf + rf, @@ -71,20 +106,28 @@ fn apply(l: &CellValue, r: &CellValue, op: BinaryOp) -> CellValue { BinaryOp::Mul => lf * rf, BinaryOp::Div => { if rf == 0.0 { - return CellValue::Null; + return Ok(CellValue::Null); } lf / rf } }; - // Preserve Int64 when both operands were integral and result is - // exact — keeps integer attrs from drifting to floats unnecessarily. - if let (CellValue::Int64(_), CellValue::Int64(_)) = (l, r) - && v.fract() == 0.0 - && v.is_finite() - { - return CellValue::Int64(v as i64); + Ok(CellValue::Float64(v)) +} + +/// A passed-through numeric value of the op attr takes the type the op +/// yields, so the result column holds one type: `Int64` for an integer +/// Add, Sub, or Mul, and `Float64` otherwise. +fn widen_to_result(v: CellValue, op: BinaryOp) -> CellValue { + let int_op = match op { + BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul => true, + BinaryOp::Div => false, + }; + match (v.as_i64(), v.as_float()) { + (Some(i), _) if int_op => CellValue::Int64(i), + (Some(i), _) => CellValue::Float64(i as f64), + (None, Some(f)) => CellValue::Float64(f), + (None, None) => v, } - CellValue::Float64(v) } fn passthrough(l: &CellValue, r: &CellValue) -> CellValue { @@ -97,14 +140,6 @@ fn passthrough(l: &CellValue, r: &CellValue) -> CellValue { } } -fn to_f64(v: &CellValue) -> Option { - match v { - CellValue::Int64(x) => Some(*x as f64), - CellValue::Float64(x) => Some(*x), - _ => None, - } -} - /// Hashable, totally-ordered key for a coordinate tuple. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] enum CoordKey { @@ -206,7 +241,7 @@ mod tests { let s = schema(); let a = tile(&[(0, 1), (1, 2)]); let b = tile(&[(0, 10), (1, 20)]); - let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap(); + let out = elementwise(&s, &a, &b, BinaryOp::Add, 0).unwrap(); assert_eq!(out.nnz(), 2); assert_eq!(out.attr_cols[0][0], CellValue::Int64(11)); assert_eq!(out.attr_cols[0][1], CellValue::Int64(22)); @@ -217,7 +252,7 @@ mod tests { let s = schema(); let a = tile(&[(0, 1)]); let b = tile(&[(1, 2)]); - let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap(); + let out = elementwise(&s, &a, &b, BinaryOp::Add, 0).unwrap(); assert_eq!(out.nnz(), 2); // Both rows are missing one operand → both become Null+something // → passthrough gives the present operand. @@ -232,7 +267,7 @@ mod tests { let s = schema(); let a = tile(&[(0, 5)]); let b = tile(&[(0, 0)]); - let out = elementwise(&s, &a, &b, BinaryOp::Div).unwrap(); + let out = elementwise(&s, &a, &b, BinaryOp::Div, 0).unwrap(); assert_eq!(out.attr_cols[0][0], CellValue::Null); } @@ -267,7 +302,7 @@ mod tests { .unwrap(); let a = builder.build(); let b = tile(&[(1, 6)]); - let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap(); + let out = elementwise(&s, &a, &b, BinaryOp::Add, 0).unwrap(); assert_eq!(out.nnz(), 1); assert_eq!(out.attr_cols[0][0], CellValue::Int64(10)); } @@ -277,9 +312,115 @@ mod tests { let s = schema(); let a = tile(&[(0, 10)]); let b = tile(&[(0, 3)]); - let s1 = elementwise(&s, &a, &b, BinaryOp::Sub).unwrap(); - let s2 = elementwise(&s, &a, &b, BinaryOp::Mul).unwrap(); + let s1 = elementwise(&s, &a, &b, BinaryOp::Sub, 0).unwrap(); + let s2 = elementwise(&s, &a, &b, BinaryOp::Mul, 0).unwrap(); assert_eq!(s1.attr_cols[0][0], CellValue::Int64(7)); assert_eq!(s2.attr_cols[0][0], CellValue::Int64(30)); } + + fn two_attr_schema() -> ArraySchema { + ArraySchemaBuilder::new("g2") + .dim(DimSpec::new( + "k", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("a", AttrType::Int64, true)) + .attr(AttrSpec::new("b", AttrType::Int64, true)) + .tile_extents(vec![16]) + .build() + .unwrap() + } + + fn two_attr_tile(s: &ArraySchema, k: i64, a: i64, b: i64) -> SparseTile { + let mut builder = SparseTileBuilder::new(s); + builder + .push( + &[CoordValue::Int64(k)], + &[CellValue::Int64(a), CellValue::Int64(b)], + ) + .unwrap(); + builder.build() + } + + /// Only the named attr takes the op. The other attr keeps the left + /// operand's value. + #[test] + fn the_op_reads_the_named_attr_only() { + let s = two_attr_schema(); + let left = two_attr_tile(&s, 0, 1, 2); + let right = two_attr_tile(&s, 0, 10, 20); + let on_b = elementwise(&s, &left, &right, BinaryOp::Add, 1).unwrap(); + assert_eq!(on_b.attr_cols[1][0], CellValue::Int64(22)); + assert_eq!(on_b.attr_cols[0][0], CellValue::Int64(1)); + let on_a = elementwise(&s, &left, &right, BinaryOp::Add, 0).unwrap(); + assert_eq!(on_a.attr_cols[0][0], CellValue::Int64(11)); + assert_eq!(on_a.attr_cols[1][0], CellValue::Int64(2)); + } + + fn typed_tile(dtype: AttrType, rows: &[(i64, CellValue)]) -> (ArraySchema, SparseTile) { + let s = ArraySchemaBuilder::new("t") + .dim(DimSpec::new( + "k", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("v", dtype, true)) + .tile_extents(vec![16]) + .build() + .unwrap(); + let mut b = SparseTileBuilder::new(&s); + for (k, v) in rows { + b.push(&[CoordValue::Int64(*k)], std::slice::from_ref(v)) + .unwrap(); + } + let t = b.build(); + (s, t) + } + + /// Narrow integers add exactly into `Int64`, and a passed-through cell + /// takes the same type. + #[test] + fn narrow_integers_add_into_int64() { + let (s, a) = typed_tile( + AttrType::UInt16, + &[(0, CellValue::UInt16(65535)), (1, CellValue::UInt16(3))], + ); + let (_, b) = typed_tile(AttrType::UInt16, &[(0, CellValue::UInt16(65535))]); + let out = elementwise(&s, &a, &b, BinaryOp::Add, 0).unwrap(); + assert_eq!(out.attr_cols[0][0], CellValue::Int64(131_070)); + assert_eq!(out.attr_cols[0][1], CellValue::Int64(3)); + } + + #[test] + fn integer_overflow_is_refused() { + let s = schema(); + let a = tile(&[(0, i64::MAX)]); + let b = tile(&[(0, 2)]); + assert!(matches!( + elementwise(&s, &a, &b, BinaryOp::Mul, 0), + Err(ArrayError::NumericOverflow { .. }) + )); + } + + /// Div of integers, and any op over a float attr, yields `Float64`. + #[test] + fn div_and_floats_yield_float64() { + let s = schema(); + let a = tile(&[(0, 6), (1, 7)]); + let b = tile(&[(0, 3)]); + let out = elementwise(&s, &a, &b, BinaryOp::Div, 0).unwrap(); + assert_eq!(out.attr_cols[0][0], CellValue::Float64(2.0)); + assert_eq!(out.attr_cols[0][1], CellValue::Float64(7.0)); + let (fs, fa) = typed_tile(AttrType::Float32, &[(0, CellValue::Float32(0.5))]); + let out = elementwise(&fs, &fa, &fa, BinaryOp::Add, 0).unwrap(); + assert_eq!(out.attr_cols[0][0], CellValue::Float64(1.0)); + } + + #[test] + fn an_attr_index_out_of_range_is_refused() { + let s = two_attr_schema(); + let t = two_attr_tile(&s, 0, 1, 2); + assert!(elementwise(&s, &t, &t, BinaryOp::Add, 2).is_err()); + } } diff --git a/nodedb-array/src/query/mod.rs b/nodedb-array/src/query/mod.rs index 68608f76f..19e09ea18 100644 --- a/nodedb-array/src/query/mod.rs +++ b/nodedb-array/src/query/mod.rs @@ -3,24 +3,31 @@ //! Query operators over tiles. //! //! Pure, tile-local primitives: slice (coord-range filter), project -//! (attr subset), aggregate (sum/count/min/max/mean with optional -//! group-by + cross-tile merge), elementwise (binary ops between two -//! aligned tiles) and rechunk (re-bucket cells into a new tile-extent -//! layout). No engine state, no plane crossings — the executor wires these +//! (attr subset), region (coordinate box plus polygon zone), aggregate +//! (sum/count/min/max/mean with optional group-by + cross-tile merge), +//! elementwise (binary ops between two aligned tiles) and rechunk +//! (re-bucket cells into a new tile-extent layout). No engine state, no plane crossings — the executor wires these //! into the physical plan. +pub mod agg_number; pub mod aggregate; pub mod ceiling; pub mod elementwise; pub mod project; pub mod rechunk; +pub mod region; pub mod retention; pub mod slice; +pub use agg_number::{AggNumber, AggScalar}; pub use aggregate::{AggregateResult, GroupAggregate, Reducer, aggregate_attr, group_by_dim}; pub use ceiling::{CeilingParams, CeilingResult, ceiling_resolve_cell}; pub use elementwise::{BinaryOp, elementwise}; pub use project::{Projection, project_sparse}; pub use rechunk::rechunk_sparse; +pub use region::{ArrayRegion, PolygonZone, ZonePolygon, filter_region_owned}; pub use retention::{RetentionMergeResult, merge_for_retention}; -pub use slice::{DimRange, Slice, slice_sparse, tile_overlaps_slice}; +pub use slice::{ + DimRange, Slice, filter_sparse, slice_sparse, slice_sparse_owned, tile_overlaps_slice, + tile_within_slice, +}; diff --git a/nodedb-array/src/query/region.rs b/nodedb-array/src/query/region.rs new file mode 100644 index 000000000..b2d4fc39e --- /dev/null +++ b/nodedb-array/src/query/region.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Region predicate over an array: a coordinate box and an optional +//! polygon zone on two numeric dims. +//! +//! A cell lies in an [`ArrayRegion`] when it passes the box ([`Slice`]) and, +//! when a zone is set, lies inside one member polygon of the zone and inside +//! none of that polygon's holes. Polygon membership uses the ray-casting rule +//! of [`nodedb_types::geometry::point_in_polygon`]: a point on a minimum +//! edge is inside, a point on a maximum edge is outside. +//! +//! Two entry points: +//! * [`ArrayRegion::tile_may_contain`] — MBR-only check. A `false` proves +//! the tile holds no live cell in the region, so no payload is decoded. +//! * [`filter_region_owned`] — per-cell filter of one tile. + +use nodedb_types::geometry::point_in_polygon; + +use super::slice::{Slice, cell_in_slice, filter_sparse, slice_sparse_owned, tile_overlaps_slice}; +use crate::error::ArrayResult; +use crate::schema::ArraySchema; +use crate::tile::mbr::TileMBR; +use crate::tile::sparse_tile::SparseTile; +use crate::types::coord::value::CoordValue; +use crate::types::domain::DomainBound; + +/// One polygon of a zone. Each ring is a closed list of `[x, y]` positions. +#[derive( + Debug, + Clone, + PartialEq, + serde::Serialize, + serde::Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ZonePolygon { + pub exterior: Vec<[f64; 2]>, + pub holes: Vec>, +} + +impl ZonePolygon { + fn contains(&self, x: f64, y: f64) -> bool { + point_in_polygon(x, y, &self.exterior) + && !self.holes.iter().any(|hole| point_in_polygon(x, y, hole)) + } + + /// `[min_x, min_y, max_x, max_y]` of the exterior ring, or `None` for an + /// empty ring. + fn bbox(&self) -> Option<[f64; 4]> { + let first = self.exterior.first()?; + let mut b = [first[0], first[1], first[0], first[1]]; + for p in &self.exterior { + b[0] = b[0].min(p[0]); + b[1] = b[1].min(p[1]); + b[2] = b[2].max(p[0]); + b[3] = b[3].max(p[1]); + } + Some(b) + } +} + +/// A polygon zone over two numeric dims. `x_dim` and `y_dim` are schema dim +/// positions. A cell is in the zone when it lies in any member polygon. +#[derive( + Debug, + Clone, + PartialEq, + serde::Serialize, + serde::Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct PolygonZone { + pub x_dim: u32, + pub y_dim: u32, + pub polygons: Vec, +} + +impl PolygonZone { + fn contains(&self, coord: &[CoordValue]) -> bool { + let (Some(x), Some(y)) = ( + coord.get(self.x_dim as usize).and_then(coord_to_f64), + coord.get(self.y_dim as usize).and_then(coord_to_f64), + ) else { + return false; + }; + self.polygons.iter().any(|p| p.contains(x, y)) + } + + /// False only when no member polygon's bounding box meets the MBR's + /// `x_dim` × `y_dim` rectangle. A bound that is not numeric proves + /// nothing, so the tile stays in. + fn may_overlap(&self, mbr: &TileMBR) -> bool { + let bound = |dims: &[DomainBound], dim: u32| dims.get(dim as usize).and_then(bound_to_f64); + let (Some(tx0), Some(ty0), Some(tx1), Some(ty1)) = ( + bound(&mbr.dim_mins, self.x_dim), + bound(&mbr.dim_mins, self.y_dim), + bound(&mbr.dim_maxs, self.x_dim), + bound(&mbr.dim_maxs, self.y_dim), + ) else { + return true; + }; + self.polygons + .iter() + .filter_map(ZonePolygon::bbox) + .any(|[px0, py0, px1, py1]| px0 <= tx1 && tx0 <= px1 && py0 <= ty1 && ty0 <= py1) + } +} + +/// The cells an aggregate reads: a coordinate box and an optional zone. +/// `slice.dim_ranges` is empty or holds one entry per schema dim. +#[derive( + Debug, + Clone, + Default, + serde::Serialize, + serde::Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayRegion { + pub slice: Slice, + pub zone: Option, +} + +impl ArrayRegion { + pub fn new(slice: Slice, zone: Option) -> Self { + Self { slice, zone } + } + + /// False only when no live cell of a tile with this MBR lies in the + /// region. + pub fn tile_may_contain(&self, mbr: &TileMBR) -> bool { + tile_overlaps_slice(&mbr.dim_mins, &mbr.dim_maxs, &self.slice) + && self.zone.as_ref().is_none_or(|z| z.may_overlap(mbr)) + } + + /// True if the cell at `coord` lies in the region. + pub fn contains(&self, coord: &[CoordValue]) -> bool { + cell_in_slice(coord, &self.slice) && self.zone.as_ref().is_none_or(|z| z.contains(coord)) + } +} + +/// The live cells of `tile` that lie in `region`. A tile whose MBR misses +/// the region is replaced by an empty tile with no per-cell work. +pub fn filter_region_owned( + schema: &ArraySchema, + tile: SparseTile, + region: &ArrayRegion, +) -> ArrayResult { + if !region.tile_may_contain(&tile.mbr) { + return Ok(SparseTile::empty(schema)); + } + match ®ion.zone { + None => slice_sparse_owned(schema, tile, ®ion.slice), + Some(_) => filter_sparse(schema, &tile, |coord| region.contains(coord)), + } +} + +fn coord_to_f64(c: &CoordValue) -> Option { + match c { + CoordValue::Int64(v) => Some(*v as f64), + CoordValue::Float64(v) => Some(*v), + CoordValue::TimestampMs(_) | CoordValue::String(_) => None, + } +} + +fn bound_to_f64(b: &DomainBound) -> Option { + match b { + DomainBound::Int64(v) => Some(*v as f64), + DomainBound::Float64(v) => Some(*v), + DomainBound::TimestampMs(_) | DomainBound::String(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::query::slice::DimRange; + use crate::schema::ArraySchemaBuilder; + use crate::schema::attr_spec::{AttrSpec, AttrType}; + use crate::schema::dim_spec::{DimSpec, DimType}; + use crate::tile::sparse_tile::SparseTileBuilder; + use crate::types::cell_value::value::CellValue; + use crate::types::domain::Domain; + + fn schema() -> ArraySchema { + ArraySchemaBuilder::new("g") + .dim(DimSpec::new( + "x", + DimType::Float64, + Domain::new(DomainBound::Float64(0.0), DomainBound::Float64(100.0)), + )) + .dim(DimSpec::new( + "y", + DimType::Float64, + Domain::new(DomainBound::Float64(0.0), DomainBound::Float64(100.0)), + )) + .attr(AttrSpec::new("v", AttrType::Float64, true)) + .tile_extents(vec![10, 10]) + .build() + .unwrap() + } + + fn tile(cells: &[(f64, f64)]) -> SparseTile { + let s = schema(); + let mut b = SparseTileBuilder::new(&s); + for (x, y) in cells { + b.push( + &[CoordValue::Float64(*x), CoordValue::Float64(*y)], + &[CellValue::Float64(1.0)], + ) + .unwrap(); + } + b.build() + } + + fn square(lo: f64, hi: f64) -> Vec<[f64; 2]> { + vec![[lo, lo], [hi, lo], [hi, hi], [lo, hi], [lo, lo]] + } + + fn zone(polygons: Vec) -> ArrayRegion { + ArrayRegion::new( + Slice::default(), + Some(PolygonZone { + x_dim: 0, + y_dim: 1, + polygons, + }), + ) + } + + fn at(x: f64, y: f64) -> [CoordValue; 2] { + [CoordValue::Float64(x), CoordValue::Float64(y)] + } + + #[test] + fn a_minimum_edge_is_in_and_a_maximum_edge_is_out() { + let r = zone(vec![ZonePolygon { + exterior: square(0.0, 4.0), + holes: Vec::new(), + }]); + assert!(r.contains(&at(0.0, 2.0)), "min x edge"); + assert!(r.contains(&at(2.0, 0.0)), "min y edge"); + assert!(!r.contains(&at(4.0, 2.0)), "max x edge"); + assert!(!r.contains(&at(2.0, 4.0)), "max y edge"); + assert!(r.contains(&at(2.0, 2.0))); + assert!(!r.contains(&at(5.0, 2.0))); + } + + #[test] + fn a_hole_excludes_its_cells() { + let r = zone(vec![ZonePolygon { + exterior: square(0.0, 10.0), + holes: vec![square(4.0, 6.0)], + }]); + assert!(r.contains(&at(2.0, 2.0))); + assert!(!r.contains(&at(5.0, 5.0))); + } + + #[test] + fn any_member_polygon_admits_a_cell() { + let r = zone(vec![ + ZonePolygon { + exterior: square(0.0, 2.0), + holes: Vec::new(), + }, + ZonePolygon { + exterior: square(8.0, 10.0), + holes: Vec::new(), + }, + ]); + assert!(r.contains(&at(1.0, 1.0))); + assert!(r.contains(&at(9.0, 9.0))); + assert!(!r.contains(&at(5.0, 5.0))); + } + + #[test] + fn the_box_and_the_zone_both_bound_the_region() { + let mut r = zone(vec![ZonePolygon { + exterior: square(0.0, 10.0), + holes: Vec::new(), + }]); + r.slice = Slice::new(vec![ + Some(DimRange::new( + DomainBound::Float64(0.0), + DomainBound::Float64(3.0), + )), + None, + ]); + assert!(r.contains(&at(2.0, 8.0))); + assert!(!r.contains(&at(5.0, 8.0)), "outside the box"); + } + + #[test] + fn a_tile_beside_every_polygon_is_pruned() { + let r = zone(vec![ZonePolygon { + exterior: square(0.0, 4.0), + holes: Vec::new(), + }]); + assert!(r.tile_may_contain(&tile(&[(1.0, 1.0), (9.0, 9.0)]).mbr)); + assert!(!r.tile_may_contain(&tile(&[(6.0, 6.0), (9.0, 9.0)]).mbr)); + assert!(ArrayRegion::default().tile_may_contain(&tile(&[(9.0, 9.0)]).mbr)); + } + + #[test] + fn the_filter_keeps_the_cells_inside() { + let s = schema(); + let r = zone(vec![ZonePolygon { + exterior: square(0.0, 4.0), + holes: Vec::new(), + }]); + let out = filter_region_owned(&s, tile(&[(1.0, 1.0), (0.0, 3.0), (4.0, 1.0)]), &r).unwrap(); + assert_eq!(out.nnz(), 2); + let pruned = filter_region_owned(&s, tile(&[(6.0, 6.0)]), &r).unwrap(); + assert_eq!(pruned.nnz(), 0); + let whole = filter_region_owned(&s, tile(&[(6.0, 6.0)]), &ArrayRegion::default()).unwrap(); + assert_eq!(whole.nnz(), 1); + } + + #[test] + fn a_region_round_trips_through_msgpack() { + let r = zone(vec![ZonePolygon { + exterior: square(0.0, 4.0), + holes: vec![square(1.0, 2.0)], + }]); + let bytes = zerompk::to_msgpack_vec(&r).unwrap(); + let back: ArrayRegion = zerompk::from_msgpack(&bytes).unwrap(); + assert_eq!(back.zone, r.zone); + } +} diff --git a/nodedb-array/src/query/slice.rs b/nodedb-array/src/query/slice.rs index 7ee98c72b..a3ef2fd3f 100644 --- a/nodedb-array/src/query/slice.rs +++ b/nodedb-array/src/query/slice.rs @@ -86,6 +86,44 @@ pub fn tile_overlaps_slice( true } +/// True if a tile's per-dim MBR lies inside every constrained dim range, so +/// every cell of the tile passes the slice. A bound of another type than +/// its range never lies inside it. +pub fn tile_within_slice( + dim_mins: &[DomainBound], + dim_maxs: &[DomainBound], + slice: &Slice, +) -> bool { + for (i, range) in slice.dim_ranges.iter().enumerate() { + let Some(r) = range else { continue }; + let (Some(tile_min), Some(tile_max)) = (dim_mins.get(i), dim_maxs.get(i)) else { + return false; + }; + let same_type = [tile_min, tile_max, &r.hi] + .iter() + .all(|b| std::mem::discriminant(*b) == std::mem::discriminant(&r.lo)); + if !same_type || bound_lt(tile_min, &r.lo) || bound_lt(&r.hi, tile_max) { + return false; + } + } + true +} + +/// [`slice_sparse`] over an owned tile. A tile of live rows only whose MBR +/// lies inside the slice is returned as it is, with no rebuild. +pub fn slice_sparse_owned( + schema: &ArraySchema, + tile: SparseTile, + slice: &Slice, +) -> ArrayResult { + use crate::tile::sparse_tile::RowKind; + let all_live = tile.row_kinds.iter().all(|&k| k == RowKind::Live.as_u8()); + if all_live && tile_within_slice(&tile.mbr.dim_mins, &tile.mbr.dim_maxs, slice) { + return Ok(tile); + } + slice_sparse(schema, &tile, slice) +} + /// Filter cells in `tile` to those whose coords pass every dim range. /// Result is a freshly-built [`SparseTile`] — dictionaries shrink to /// surviving values, MBR/attr_stats are recomputed. @@ -93,6 +131,17 @@ pub fn slice_sparse( schema: &ArraySchema, tile: &SparseTile, slice: &Slice, +) -> ArrayResult { + filter_sparse(schema, tile, |coord| cell_in_slice(coord, slice)) +} + +/// Keep the live cells of `tile` whose coordinate passes `keep`. Result is a +/// freshly-built [`SparseTile`] — dictionaries shrink to surviving values, +/// MBR/attr_stats are recomputed. Sentinel rows never survive. +pub fn filter_sparse( + schema: &ArraySchema, + tile: &SparseTile, + keep: impl Fn(&[CoordValue]) -> bool, ) -> ArrayResult { use crate::tile::sparse_tile::RowKind; let mut b = SparseTileBuilder::new(schema); @@ -108,7 +157,7 @@ pub fn slice_sparse( .iter() .map(|d| d.values[d.indices[row] as usize].clone()) .collect(); - if !cell_in_slice(&coord, slice) { + if !keep(&coord) { live_idx += 1; continue; } @@ -139,7 +188,8 @@ pub fn slice_sparse( Ok(b.build()) } -fn cell_in_slice(coord: &[CoordValue], slice: &Slice) -> bool { +/// True if `coord` passes every constrained dim range of `slice`. +pub(crate) fn cell_in_slice(coord: &[CoordValue], slice: &Slice) -> bool { for (i, range) in slice.dim_ranges.iter().enumerate() { let Some(r) = range else { continue }; let Some(c) = coord.get(i) else { @@ -278,4 +328,29 @@ mod tests { ]); assert!(tile_overlaps_slice(&mins, &maxs, &sl)); } + + #[test] + fn mbr_within_needs_both_bounds_inside() { + let mins = vec![DomainBound::Int64(5), DomainBound::Int64(0)]; + let maxs = vec![DomainBound::Int64(9), DomainBound::Int64(9)]; + let inside = Slice::new(vec![ + Some(DimRange::new(DomainBound::Int64(5), DomainBound::Int64(9))), + None, + ]); + assert!(tile_within_slice(&mins, &maxs, &inside)); + assert!(tile_within_slice(&mins, &maxs, &Slice::default())); + let partial = Slice::new(vec![ + Some(DimRange::new(DomainBound::Int64(6), DomainBound::Int64(20))), + None, + ]); + assert!(!tile_within_slice(&mins, &maxs, &partial)); + let other_type = Slice::new(vec![ + Some(DimRange::new( + DomainBound::Float64(0.0), + DomainBound::Float64(100.0), + )), + None, + ]); + assert!(!tile_within_slice(&mins, &maxs, &other_type)); + } } diff --git a/nodedb-array/src/schema/attr_spec.rs b/nodedb-array/src/schema/attr_spec.rs index 288259f80..fdb9add65 100644 --- a/nodedb-array/src/schema/attr_spec.rs +++ b/nodedb-array/src/schema/attr_spec.rs @@ -7,6 +7,10 @@ use serde::{Deserialize, Serialize}; /// Attribute type tag. One-to-one with /// [`crate::types::cell_value::value::CellValue`] (excluding `Null`, /// which is per-cell, not per-attribute). +/// +/// The narrow integer and `Float32` types store a cell at its declared +/// width. A value outside the width is refused at write time, never +/// wrapped or saturated. #[derive( Debug, Clone, @@ -25,6 +29,62 @@ pub enum AttrType { Float64, String, Bytes, + #[serde(rename = "uint8")] + UInt8, + Int16, + #[serde(rename = "uint16")] + UInt16, + Int32, + #[serde(rename = "uint32")] + UInt32, + Float32, +} + +impl AttrType { + /// The canonical DDL name of the type. + pub fn sql_name(self) -> &'static str { + match self { + AttrType::Int64 => "INT64", + AttrType::Float64 => "FLOAT64", + AttrType::String => "STRING", + AttrType::Bytes => "BYTES", + AttrType::UInt8 => "UINT8", + AttrType::Int16 => "INT16", + AttrType::UInt16 => "UINT16", + AttrType::Int32 => "INT32", + AttrType::UInt32 => "UINT32", + AttrType::Float32 => "FLOAT32", + } + } + + /// The inclusive range an integer type holds, or `None` for a + /// non-integer type. + pub fn int_range(self) -> Option<(i64, i64)> { + match self { + AttrType::Int64 => Some((i64::MIN, i64::MAX)), + AttrType::UInt8 => Some((0, i64::from(u8::MAX))), + AttrType::Int16 => Some((i64::from(i16::MIN), i64::from(i16::MAX))), + AttrType::UInt16 => Some((0, i64::from(u16::MAX))), + AttrType::Int32 => Some((i64::from(i32::MIN), i64::from(i32::MAX))), + AttrType::UInt32 => Some((0, i64::from(u32::MAX))), + AttrType::Float64 | AttrType::Float32 | AttrType::String | AttrType::Bytes => None, + } + } + + /// Whether the type is `Float32` or `Float64`. + pub fn is_float(self) -> bool { + match self { + AttrType::Float64 | AttrType::Float32 => true, + AttrType::Int64 + | AttrType::UInt8 + | AttrType::Int16 + | AttrType::UInt16 + | AttrType::Int32 + | AttrType::UInt32 + | AttrType::String + | AttrType::Bytes => false, + } + } } /// One attribute column on an [`super::ArraySchema`]. @@ -65,4 +125,15 @@ mod tests { let b = a.clone(); assert_eq!(a, b); } + + #[test] + fn integer_ranges_match_their_widths() { + assert_eq!(AttrType::UInt8.int_range(), Some((0, 255))); + assert_eq!(AttrType::Int16.int_range(), Some((-32768, 32767))); + assert_eq!(AttrType::UInt16.int_range(), Some((0, 65535))); + assert_eq!(AttrType::UInt32.int_range(), Some((0, 4_294_967_295))); + assert_eq!(AttrType::Float32.int_range(), None); + assert!(AttrType::Float32.is_float()); + assert!(!AttrType::Int32.is_float()); + } } diff --git a/nodedb-array/src/schema/cell_check.rs b/nodedb-array/src/schema/cell_check.rs new file mode 100644 index 000000000..f7cc26907 --- /dev/null +++ b/nodedb-array/src/schema/cell_check.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Checking a live cell's attribute values against the schema. +//! +//! Every write that stores cells passes this gate, so a stored column +//! holds one value type plus `Null`, and the column codec can rely on it. + +use crate::error::{ArrayError, ArrayResult}; +use crate::schema::array_schema::ArraySchema; +use crate::types::cell_value::value::CellValue; + +impl ArraySchema { + /// Check one live cell's attrs: one value per attr, each `Null` on a + /// nullable attr or of the attr's exact declared type. + pub fn check_cell(&self, attrs: &[CellValue]) -> ArrayResult<()> { + if attrs.len() != self.attrs.len() { + return Err(ArrayError::InvalidOp { + detail: format!( + "a cell of '{}' carries {} attrs, the schema has {}", + self.name, + attrs.len(), + self.attrs.len() + ), + }); + } + for (value, spec) in attrs.iter().zip(&self.attrs) { + match value.attr_type() { + None if spec.nullable => {} + None => { + return Err(ArrayError::CellTypeMismatch { + array: self.name.clone(), + attr: spec.name.clone(), + detail: "the attribute is NOT NULL".to_string(), + }); + } + Some(got) if got == spec.dtype => {} + Some(got) => { + return Err(ArrayError::CellTypeMismatch { + array: self.name.clone(), + attr: spec.name.clone(), + detail: format!( + "declared {}, the cell holds {}", + spec.dtype.sql_name(), + got.sql_name() + ), + }); + } + } + } + Ok(()) + } + + /// Check every cell of a write with [`Self::check_cell`]. The first + /// refused cell refuses the whole write. + pub fn check_cells<'a>( + &self, + cells: impl IntoIterator, + ) -> ArrayResult<()> { + cells + .into_iter() + .try_for_each(|attrs| self.check_cell(attrs)) + } +} + +#[cfg(test)] +mod tests { + use crate::error::ArrayError; + use crate::schema::ArraySchemaBuilder; + use crate::schema::attr_spec::{AttrSpec, AttrType}; + use crate::schema::dim_spec::{DimSpec, DimType}; + use crate::types::cell_value::value::CellValue; + use crate::types::domain::{Domain, DomainBound}; + + fn schema() -> crate::schema::ArraySchema { + ArraySchemaBuilder::new("scene") + .dim(DimSpec::new( + "x", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("b1", AttrType::UInt16, false)) + .attr(AttrSpec::new("ndvi", AttrType::Float32, true)) + .tile_extents(vec![16]) + .build() + .expect("schema") + } + + #[test] + fn typed_cells_pass() { + let s = schema(); + let a = [CellValue::UInt16(7), CellValue::Float32(0.5)]; + let b = [CellValue::UInt16(8), CellValue::Null]; + assert!(s.check_cells([&a[..], &b[..]]).is_ok()); + } + + #[test] + fn a_wider_value_is_refused() { + let s = schema(); + let a = [CellValue::Int64(7), CellValue::Null]; + assert!(matches!( + s.check_cell(&a), + Err(ArrayError::CellTypeMismatch { attr, .. }) if attr == "b1" + )); + } + + #[test] + fn null_on_not_null_and_arity_are_refused() { + let s = schema(); + assert!(matches!( + s.check_cell(&[CellValue::Null, CellValue::Null]), + Err(ArrayError::CellTypeMismatch { .. }) + )); + assert!(matches!( + s.check_cell(&[CellValue::UInt16(1)]), + Err(ArrayError::InvalidOp { .. }) + )); + } +} diff --git a/nodedb-array/src/schema/mod.rs b/nodedb-array/src/schema/mod.rs index cef61dd3e..5772614bb 100644 --- a/nodedb-array/src/schema/mod.rs +++ b/nodedb-array/src/schema/mod.rs @@ -3,6 +3,7 @@ pub mod array_schema; pub mod attr_spec; pub mod builder; +pub mod cell_check; pub mod cell_order; pub mod dim_spec; pub mod validation; diff --git a/nodedb-array/src/segment/reader.rs b/nodedb-array/src/segment/reader.rs index 7d1ba3c6d..5a4ea4f06 100644 --- a/nodedb-array/src/segment/reader.rs +++ b/nodedb-array/src/segment/reader.rs @@ -11,7 +11,9 @@ use super::format::{ SegmentFooter, SegmentHeader, TileEntry, TileKind, framing::BlockFraming, header::HEADER_SIZE, }; use crate::codec::tag::peek_tag; -use crate::codec::tile_decode::decode_sparse_tile; +use std::sync::Arc; + +use crate::codec::tile_decode::decode_sparse_tile_with_mbr; use crate::error::{ArrayError, ArrayResult}; use crate::tile::cell_payload::{CELL_GDPR_ERASURE_SENTINEL, CELL_TOMBSTONE_SENTINEL, CellPayload}; use crate::tile::dense_tile::DenseTile; @@ -112,9 +114,12 @@ pub fn extract_cell_bytes(tile: &SparseTile, coord: &[CoordValue]) -> ArrayResul /// /// - Recognized tag (0 or 1): delegate to `decode_sparse_tile`. /// - Empty payload or unrecognized byte: return a corruption error. -fn read_sparse_tile(payload: &[u8]) -> ArrayResult { +/// +/// The footer entry holds the tile's MBR, so the decode takes it instead +/// of rebuilding it from the rows. +fn read_sparse_tile(payload: &[u8], entry: &TileEntry) -> ArrayResult { match peek_tag(payload) { - Some(_) => decode_sparse_tile(payload), + Some(_) => decode_sparse_tile_with_mbr(payload, entry.mbr.clone()), None => Err(ArrayError::SegmentCorruption { detail: format!( "sparse tile has unrecognized tag byte {:#04x}", @@ -134,10 +139,26 @@ pub enum TilePayload { pub struct SegmentReader<'a> { bytes: &'a [u8], header: SegmentHeader, - footer: SegmentFooter, + footer: Arc, } impl<'a> SegmentReader<'a> { + /// A reader over `bytes` whose header and footer were decoded from those + /// same bytes by an earlier [`Self::open`]. An open segment keeps them, + /// so each read does not decode the footer again. + pub fn from_parts(bytes: &'a [u8], header: SegmentHeader, footer: Arc) -> Self { + Self { + bytes, + header, + footer, + } + } + + /// The decoded footer, shared with every reader built from it. + pub fn footer(&self) -> &Arc { + &self.footer + } + pub fn open(bytes: &'a [u8]) -> ArrayResult { if bytes.len() < HEADER_SIZE { return Err(ArrayError::SegmentCorruption { @@ -157,7 +178,7 @@ impl<'a> SegmentReader<'a> { Ok(Self { bytes, header, - footer, + footer: Arc::new(footer), }) } @@ -284,7 +305,7 @@ impl<'a> SegmentReader<'a> { let (payload, _) = BlockFraming::decode(&self.bytes[off..end])?; match entry.kind { TileKind::Sparse => { - let t = read_sparse_tile(payload)?; + let t = read_sparse_tile(payload, entry)?; Ok(TilePayload::Sparse(t)) } TileKind::Dense => { @@ -308,7 +329,7 @@ pub struct OwnedSegmentReader { /// Decrypted plaintext segment bytes. plaintext: Vec, header: SegmentHeader, - footer: SegmentFooter, + footer: Arc, } impl OwnedSegmentReader { @@ -344,7 +365,7 @@ impl OwnedSegmentReader { Ok(Self { plaintext, header, - footer, + footer: Arc::new(footer), }) } diff --git a/nodedb-array/src/sync/apply.rs b/nodedb-array/src/sync/apply.rs index dcb3b32e4..f44584670 100644 --- a/nodedb-array/src/sync/apply.rs +++ b/nodedb-array/src/sync/apply.rs @@ -16,6 +16,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Mutex; use crate::error::{ArrayError, ArrayResult}; +use crate::schema::ArraySchema; use crate::sync::hlc::Hlc; use crate::sync::op::{ArrayOp, ArrayOpKind}; use crate::types::coord::value::CoordValue; @@ -57,6 +58,13 @@ pub enum ApplyRejection { /// Human-readable description of the violation. detail: String, }, + /// A `Put` op's attrs do not match the local schema: wrong arity, a + /// value of another type than the attr declares, or `Null` on a + /// NOT NULL attr. + SchemaViolation { + /// Human-readable description of the violation. + detail: String, + }, /// The engine rejected the op (e.g. a transient storage error). /// /// The rejection is wrapped here so callers can record it without @@ -81,6 +89,10 @@ pub trait ApplyEngine { /// not known to this replica. fn schema_hlc(&self, array: &str) -> ArrayResult>; + /// Return the current schema of `array`, or `None` if the array is not + /// known to this replica. + fn array_schema(&self, array: &str) -> ArrayResult>; + /// Return `true` if `hlc` has already been applied to `array`. /// /// Used for idempotent re-delivery detection. @@ -112,11 +124,13 @@ pub trait ApplyEngine { /// `op.header.schema_hlc` strictly later on the clock than `local` /// → `Rejected(SchemaTooNew)`. Same-instant schemas are admitted. /// 3. Idempotency — already seen → `Idempotent`. -/// 4. Dispatch to `apply_put`/`apply_delete`/`apply_erase`. Engine errors +/// 4. Schema check — a `Put` op's attrs are checked against the local schema +/// with [`ArraySchema::check_cell`] — on error: `Rejected(SchemaViolation)`. +/// 5. Dispatch to `apply_put`/`apply_delete`/`apply_erase`. Engine errors /// that indicate corruption (`SegmentCorruption`, `HlcLockPoisoned`) are /// propagated; all others become `Rejected(EngineRejected)`. -/// 5. Tile-cache invalidation. -/// 6. Return `Applied`. +/// 6. Tile-cache invalidation. +/// 7. Return `Applied`. pub fn apply_op(engine: &mut E, op: &ArrayOp) -> ArrayResult { // 1. Shape validation. if let Err(e) = op.validate_shape() { @@ -150,7 +164,21 @@ pub fn apply_op(engine: &mut E, op: &ArrayOp) -> ArrayResult engine.apply_put(op), ArrayOpKind::Delete => engine.apply_delete(op), @@ -169,10 +197,10 @@ pub fn apply_op(engine: &mut E, op: &ArrayOp) -> ArrayResult(engine: &mut E, op: &ArrayOp) -> ArrayResult, + schemas: HashMap, seen: HashSet<(String, [u8; 18])>, pub applied: Vec, pub invalidated: Vec<(String, Vec)>, @@ -208,9 +236,28 @@ impl MockEngine { } } - /// Register an array so that ops targeting it can be applied. + /// Register an array so that ops targeting it can be applied. Its schema + /// has one `Int64` dim `k` and one nullable `Float64` attr `v`. pub fn register_array(&mut self, array: &str, schema_hlc: Hlc) { - self.schemas.insert(array.to_string(), schema_hlc); + use crate::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; + use crate::types::domain::{Domain, DomainBound}; + let schema = ArraySchemaBuilder::new(array) + .dim(DimSpec::new( + "k", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("v", AttrType::Float64, true)) + .tile_extents(vec![16]) + .build() + .expect("invariant: the mock schema is well formed"); + self.register_array_with_schema(schema_hlc, schema); + } + + /// Register an array under its own schema. + pub fn register_array_with_schema(&mut self, schema_hlc: Hlc, schema: ArraySchema) { + self.schemas + .insert(schema.name.clone(), (schema_hlc, schema)); } /// Inject an error that will be returned by the next `apply_*` call. @@ -239,7 +286,11 @@ impl Default for MockEngine { #[cfg(test)] impl ApplyEngine for MockEngine { fn schema_hlc(&self, array: &str) -> ArrayResult> { - Ok(self.schemas.get(array).copied()) + Ok(self.schemas.get(array).map(|(hlc, _)| *hlc)) + } + + fn array_schema(&self, array: &str) -> ArrayResult> { + Ok(self.schemas.get(array).map(|(_, schema)| schema.clone())) } fn already_seen(&self, array: &str, hlc: Hlc) -> ArrayResult { @@ -423,6 +474,50 @@ mod tests { )); } + /// A Put whose attr value is of another type than the local schema + /// declares is refused before the engine sees it. + #[test] + fn apply_rejects_a_put_that_breaks_the_schema() { + use crate::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; + use crate::types::domain::{Domain, DomainBound}; + let mut engine = MockEngine::new(); + let schema = ArraySchemaBuilder::new("scene") + .dim(DimSpec::new( + "k", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("b1", AttrType::UInt16, false)) + .tile_extents(vec![16]) + .build() + .unwrap(); + engine.register_array_with_schema(hlc(100), schema); + + let mut wide = put_op("scene", 50, 100); + wide.attrs = Some(vec![CellValue::Int64(70_000)]); + let outcome = apply_op(&mut engine, &wide).unwrap(); + assert!(matches!( + outcome, + ApplyOutcome::Rejected(ApplyRejection::SchemaViolation { .. }) + )); + + let mut null = put_op("scene", 51, 100); + null.attrs = Some(vec![CellValue::Null]); + let outcome = apply_op(&mut engine, &null).unwrap(); + assert!(matches!( + outcome, + ApplyOutcome::Rejected(ApplyRejection::SchemaViolation { .. }) + )); + assert!(engine.applied.is_empty()); + + let mut typed = put_op("scene", 52, 100); + typed.attrs = Some(vec![CellValue::UInt16(65_535)]); + assert_eq!( + apply_op(&mut engine, &typed).unwrap(), + ApplyOutcome::Applied + ); + } + #[test] fn apply_wraps_engine_error_as_rejection() { let mut engine = MockEngine::new(); diff --git a/nodedb-array/src/tile/layout.rs b/nodedb-array/src/tile/layout.rs index d6c22f102..4e5ac759c 100644 --- a/nodedb-array/src/tile/layout.rs +++ b/nodedb-array/src/tile/layout.rs @@ -8,7 +8,7 @@ //! same tile share a [`TileId::hilbert_prefix`]. use crate::coord::encode::encode_tile_prefix_with_order; -use crate::coord::normalize::{bits_per_dim, normalize_coord}; +use crate::coord::normalize::{MAX_DIMS, bits_per_dim, check_coord_domain}; use crate::error::{ArrayError, ArrayResult}; use crate::schema::ArraySchema; use crate::schema::dim_spec::{DimSpec, DimType}; @@ -25,46 +25,47 @@ pub fn tile_indices_for_cell(schema: &ArraySchema, coord: &[CoordValue]) -> Arra got: coord.len(), }); } - let mut out = Vec::with_capacity(schema.arity()); - for (i, dim) in schema.dims.iter().enumerate() { - let extent = schema.tile_extents[i]; - let cell = &coord[i]; - let lo = &dim.domain.lo; - let off = match (dim.dtype, cell, lo) { - (DimType::Int64, CoordValue::Int64(v), DomainBound::Int64(lo)) - | (DimType::TimestampMs, CoordValue::TimestampMs(v), DomainBound::TimestampMs(lo)) => { - if v < lo { - return out_of_domain(schema, dim.name.as_str(), "below domain_lo"); - } - let delta = (*v as i128) - (*lo as i128); - (delta as u128 / u128::from(extent)) as u64 - } - (DimType::Float64, CoordValue::Float64(v), DomainBound::Float64(lo)) => { - if !v.is_finite() || v < lo { - return out_of_domain( - schema, - dim.name.as_str(), - "below domain_lo or non-finite", - ); - } - ((v - lo) as u64) / extent - } - (DimType::String, CoordValue::String(s), DomainBound::String(_)) => { - // Strings have no natural tile-extent semantics — bucket - // by hash modulo extent for stable grouping. - crate::coord::string_hash::hash_string_modulo(s, extent.max(1)) + coord + .iter() + .enumerate() + .map(|(i, cell)| tile_index(schema, i, cell)) + .collect() +} + +/// The tile index of `cell` along dim `i` of `schema`. +pub(crate) fn tile_index(schema: &ArraySchema, i: usize, cell: &CoordValue) -> ArrayResult { + let dim = &schema.dims[i]; + let extent = schema.tile_extents[i]; + let lo = &dim.domain.lo; + match (dim.dtype, cell, lo) { + (DimType::Int64, CoordValue::Int64(v), DomainBound::Int64(lo)) + | (DimType::TimestampMs, CoordValue::TimestampMs(v), DomainBound::TimestampMs(lo)) => { + if v < lo { + return out_of_domain(schema, dim.name.as_str(), "below domain_lo"); } - _ => { - return Err(ArrayError::CoordOutOfDomain { - array: schema.name.clone(), - dim: dim.name.clone(), - detail: "coord variant does not match dim dtype".to_string(), - }); + let delta = (*v as i128) - (*lo as i128); + Ok((delta as u128 / u128::from(extent)) as u64) + } + (DimType::Float64, CoordValue::Float64(v), DomainBound::Float64(lo)) => { + if !v.is_finite() || v < lo { + return out_of_domain(schema, dim.name.as_str(), "below domain_lo or non-finite"); } - }; - out.push(off); + Ok(((v - lo) as u64) / extent) + } + (DimType::String, CoordValue::String(s), DomainBound::String(_)) => { + // Strings have no natural tile-extent semantics — bucket + // by hash modulo extent for stable grouping. + Ok(crate::coord::string_hash::hash_string_modulo( + s, + extent.max(1), + )) + } + _ => Err(ArrayError::CoordOutOfDomain { + array: schema.name.clone(), + dim: dim.name.clone(), + detail: "coord variant does not match dim dtype".to_string(), + }), } - Ok(out) } /// Compute the [`TileId`] for one cell coordinate at a given system @@ -107,19 +108,27 @@ pub fn tile_prefix(schema: &ArraySchema, tile_indices: &[u64]) -> ArrayResult = tile_indices - .iter() - .zip(schema.dims.iter().zip(schema.tile_extents.iter())) - .map(|(index, (dim, extent))| scale_tile_index(*index, tiles_per_dim(dim, *extent), bits)) - .collect(); - encode_tile_prefix_with_order(schema, &scaled, bits) + let n = tile_indices.len(); + let mut buf = [0u64; MAX_DIMS]; + let scaled = buf.get_mut(..n).ok_or_else(|| ArrayError::InvalidSchema { + array: schema.name.clone(), + detail: format!("tile prefix: arity {n} exceeds MAX_DIMS={MAX_DIMS}"), + })?; + for (slot, (index, (dim, extent))) in scaled.iter_mut().zip( + tile_indices + .iter() + .zip(schema.dims.iter().zip(schema.tile_extents.iter())), + ) { + *slot = scale_tile_index(*index, tiles_per_dim(dim, *extent), bits); + } + encode_tile_prefix_with_order(schema, scaled, bits) } /// The key of the tile that holds the cell at `coord`: the key every /// cluster path routes the cell by. Fails when `coord` lies outside the /// schema's domain. pub fn cell_tile_prefix(schema: &ArraySchema, coord: &[CoordValue]) -> ArrayResult { - normalize_coord(schema, coord, bits_per_dim(schema.arity()))?; + check_coord_domain(schema, coord, bits_per_dim(schema.arity()))?; tile_prefix(schema, &tile_indices_for_cell(schema, coord)?) } diff --git a/nodedb-array/src/tile/mbr.rs b/nodedb-array/src/tile/mbr.rs index c23429db1..81e050e03 100644 --- a/nodedb-array/src/tile/mbr.rs +++ b/nodedb-array/src/tile/mbr.rs @@ -23,7 +23,8 @@ use crate::types::domain::DomainBound; zerompk::FromMessagePack, )] pub enum AttrStats { - /// Numeric (`Int64` / `Float64`) min + max + null count. + /// Numeric min + max + null count. Every integer and float attr type + /// folds here, widened to `f64`. Numeric { min: f64, max: f64, null_count: u32 }, /// String / bytes — only cardinality + null count, no min/max. Categorical { distinct: u32, null_count: u32 }, @@ -94,7 +95,26 @@ impl MbrBuilder { } } + /// Fold one live cell: its coordinate, each of its attrs, and its count. pub fn fold(&mut self, coord: &[CoordValue], attrs: &[CellValue]) { + self.fold_coord(coord); + for (i, a) in attrs.iter().take(self.n_attrs).enumerate() { + self.fold_attr(i, a); + } + self.nnz += 1; + } + + /// Fold one live cell whose attr `i` is `attr(i)`, for a caller that + /// holds the attrs column-wise rather than as one slice. + pub fn fold_with<'v>(&mut self, coord: &[CoordValue], attr: impl Fn(usize) -> &'v CellValue) { + self.fold_coord(coord); + for i in 0..self.n_attrs { + self.fold_attr(i, attr(i)); + } + self.nnz += 1; + } + + fn fold_coord(&mut self, coord: &[CoordValue]) { for (i, c) in coord.iter().take(self.arity).enumerate() { let lo_replaces = self.dim_mins[i].as_ref().is_none_or(|cur| coord_lt(c, cur)); if lo_replaces { @@ -105,27 +125,33 @@ impl MbrBuilder { self.dim_maxs[i] = Some(c.clone()); } } - for (i, a) in attrs.iter().take(self.n_attrs).enumerate() { - self.attr_total[i] += 1; - match a { - CellValue::Null => self.attr_nulls[i] += 1, - CellValue::Int64(v) => { - fold_numeric(&mut self.attr_min[i], &mut self.attr_max[i], *v as f64) - } - CellValue::Float64(v) => { - fold_numeric(&mut self.attr_min[i], &mut self.attr_max[i], *v) - } - CellValue::String(s) => { - self.attr_is_categorical[i] = true; - self.attr_distinct[i].insert(s.as_bytes().to_vec()); - } - CellValue::Bytes(b) => { - self.attr_is_categorical[i] = true; - self.attr_distinct[i].insert(b.clone()); + } + + fn fold_attr(&mut self, i: usize, a: &CellValue) { + self.attr_total[i] += 1; + match a { + CellValue::Null => self.attr_nulls[i] += 1, + CellValue::Int64(_) + | CellValue::Float64(_) + | CellValue::UInt8(_) + | CellValue::Int16(_) + | CellValue::UInt16(_) + | CellValue::Int32(_) + | CellValue::UInt32(_) + | CellValue::Float32(_) => { + if let Some(v) = a.as_f64() { + fold_numeric(&mut self.attr_min[i], &mut self.attr_max[i], v); } } + CellValue::String(s) => { + self.attr_is_categorical[i] = true; + self.attr_distinct[i].insert(s.as_bytes().to_vec()); + } + CellValue::Bytes(b) => { + self.attr_is_categorical[i] = true; + self.attr_distinct[i].insert(b.clone()); + } } - self.nnz += 1; } pub fn build(self) -> TileMBR { diff --git a/nodedb-array/src/tile/mod.rs b/nodedb-array/src/tile/mod.rs index f2a285abf..c8955e758 100644 --- a/nodedb-array/src/tile/mod.rs +++ b/nodedb-array/src/tile/mod.rs @@ -4,6 +4,7 @@ pub mod cell_payload; pub mod dense_tile; pub mod layout; pub mod mbr; +pub mod prefix_memo; pub mod sparse_tile; pub use cell_payload::{ @@ -15,4 +16,5 @@ pub use layout::{ cell_tile_prefix, tile_id_for_cell, tile_indices_for_cell, tile_prefix, tiles_per_dim, }; pub use mbr::{AttrStats, TileMBR}; +pub use prefix_memo::TilePrefixMemo; pub use sparse_tile::SparseTile; diff --git a/nodedb-array/src/tile/prefix_memo.rs b/nodedb-array/src/tile/prefix_memo.rs new file mode 100644 index 000000000..1d2614144 --- /dev/null +++ b/nodedb-array/src/tile/prefix_memo.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Tile keys of a batch of cells, one curve encode per run of cells in one +//! tile. +//! +//! A cell's tile key is the space-filling-curve encode of its tile indices. +//! The indices are a subtract and a divide per dim, and the encode costs far +//! more. Consecutive cells of one tile share the indices, so a batch keeps +//! the last tile's indices and key and encodes only when the tile changes. +//! A batch in row-major or tile order encodes once per tile row segment. + +use crate::coord::normalize::{MAX_DIMS, bits_per_dim, check_coord_domain}; +use crate::error::{ArrayError, ArrayResult}; +use crate::schema::ArraySchema; +use crate::types::TileId; +use crate::types::coord::value::CoordValue; + +use super::layout::{tile_index, tile_indices_for_cell, tile_prefix}; + +/// The last tile one batch keyed: its indices and its key. +/// +/// One memo serves one schema. A batch builds its own memo and drops it. +#[derive(Debug, Default)] +pub struct TilePrefixMemo { + last: Option<([u64; MAX_DIMS], u64)>, +} + +impl TilePrefixMemo { + pub fn new() -> Self { + Self::default() + } + + /// [`super::tile_id_for_cell`], encoding the curve only when the tile + /// differs from the last cell's. + pub fn tile_id( + &mut self, + schema: &ArraySchema, + coord: &[CoordValue], + system_from_ms: i64, + ) -> ArrayResult { + Ok(TileId::new(self.prefix(schema, coord)?, system_from_ms)) + } + + /// [`super::cell_tile_prefix`], encoding the curve only when the tile + /// differs from the last cell's. Fails when `coord` lies outside the + /// schema's domain. + pub fn cell_prefix(&mut self, schema: &ArraySchema, coord: &[CoordValue]) -> ArrayResult { + check_coord_domain(schema, coord, bits_per_dim(schema.arity()))?; + self.prefix(schema, coord) + } + + /// The tile key of `coord`, as [`tile_prefix`] over its tile indices. + fn prefix(&mut self, schema: &ArraySchema, coord: &[CoordValue]) -> ArrayResult { + let n = schema.arity(); + if coord.len() != n { + return Err(ArrayError::CoordArityMismatch { + array: schema.name.clone(), + expected: n, + got: coord.len(), + }); + } + if n > MAX_DIMS { + // No curve encodes this arity: the unmemoized path reports why. + return tile_prefix(schema, &tile_indices_for_cell(schema, coord)?); + } + let mut indices = [0u64; MAX_DIMS]; + for (i, cell) in coord.iter().enumerate() { + indices[i] = tile_index(schema, i, cell)?; + } + if let Some((last, key)) = &self.last + && last[..n] == indices[..n] + { + return Ok(*key); + } + let key = tile_prefix(schema, &indices[..n])?; + self.last = Some((indices, key)); + Ok(key) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema::ArraySchemaBuilder; + use crate::schema::attr_spec::{AttrSpec, AttrType}; + use crate::schema::dim_spec::{DimSpec, DimType}; + use crate::tile::{cell_tile_prefix, tile_id_for_cell}; + use crate::types::domain::{Domain, DomainBound}; + + fn schema() -> ArraySchema { + ArraySchemaBuilder::new("g") + .dim(DimSpec::new( + "y", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(250)), + )) + .dim(DimSpec::new( + "x", + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(255)), + )) + .attr(AttrSpec::new("v", AttrType::Int64, false)) + .tile_extents(vec![64, 64]) + .build() + .unwrap() + } + + /// Every cell of a row-major sweep, back and forth across tile edges, + /// keys as the unmemoized functions key it. + #[test] + fn a_memo_agrees_with_the_unmemoized_keys() { + let s = schema(); + let mut memo = TilePrefixMemo::new(); + let cells = (0..=250i64) + .step_by(7) + .flat_map(|y| (0..256i64).step_by(5).map(move |x| (y, x))) + .chain([(3, 200), (200, 3), (3, 201)]); + for (y, x) in cells { + let coord = [CoordValue::Int64(y), CoordValue::Int64(x)]; + assert_eq!( + memo.cell_prefix(&s, &coord).unwrap(), + cell_tile_prefix(&s, &coord).unwrap(), + "({y}, {x})" + ); + assert_eq!( + memo.tile_id(&s, &coord, 9).unwrap(), + tile_id_for_cell(&s, &coord, 9).unwrap() + ); + } + } + + /// A cell past the domain inside the last tile shares that tile's + /// indices, and the memo still refuses it. + #[test] + fn a_memo_hit_still_checks_the_domain() { + let s = schema(); + let mut memo = TilePrefixMemo::new(); + memo.cell_prefix(&s, &[CoordValue::Int64(250), CoordValue::Int64(0)]) + .unwrap(); + assert!( + memo.cell_prefix(&s, &[CoordValue::Int64(251), CoordValue::Int64(0)]) + .is_err() + ); + assert!(memo.cell_prefix(&s, &[CoordValue::Int64(1)]).is_err()); + } +} diff --git a/nodedb-array/src/tile/sparse_tile.rs b/nodedb-array/src/tile/sparse_tile.rs index c92ca81f7..7d8dd1b54 100644 --- a/nodedb-array/src/tile/sparse_tile.rs +++ b/nodedb-array/src/tile/sparse_tile.rs @@ -8,6 +8,7 @@ //! that have many repeats per tile. use nodedb_types::{OPEN_UPPER, Surrogate}; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use super::mbr::{MbrBuilder, TileMBR}; @@ -79,15 +80,6 @@ impl DimDict { } } - fn push(&mut self, v: &CoordValue) { - if let Some(idx) = self.values.iter().position(|x| x == v) { - self.indices.push(idx as u32); - } else { - self.indices.push(self.values.len() as u32); - self.values.push(v.clone()); - } - } - pub fn cardinality(&self) -> usize { self.values.len() } @@ -300,12 +292,16 @@ pub struct SparseTileBuilder<'a> { valid_until_ms: Vec, row_kinds: Vec, mbr: MbrBuilder, + /// Per dim: dictionary value → its index in `dim_dicts`, so a row's + /// dictionary lookup costs one hash, not a scan of the dictionary. + dim_index: Vec>, } impl<'a> SparseTileBuilder<'a> { pub fn new(schema: &'a ArraySchema) -> Self { Self { schema, + dim_index: (0..schema.arity()).map(|_| FxHashMap::default()).collect(), dim_dicts: (0..schema.arity()).map(|_| DimDict::new()).collect(), attr_cols: (0..schema.attrs.len()).map(|_| Vec::new()).collect(), surrogates: Vec::new(), @@ -371,7 +367,7 @@ impl<'a> SparseTileBuilder<'a> { }); } for (i, c) in coord.iter().enumerate() { - self.dim_dicts[i].push(c); + self.push_dim(i, c); } // For sentinel rows `attrs` is empty — no entries added to attr_cols. // For live rows, every attr col gets one entry. @@ -390,6 +386,64 @@ impl<'a> SparseTileBuilder<'a> { Ok(()) } + fn push_dim(&mut self, dim: usize, value: &CoordValue) { + let dict = &mut self.dim_dicts[dim]; + let index = *self.dim_index[dim].entry(value.clone()).or_insert_with(|| { + dict.values.push(value.clone()); + (dict.values.len() - 1) as u32 + }); + dict.indices.push(index); + } + + /// Push live row `row` of the stored tile `tile`, whose coordinate is + /// `coord`, reading its attrs, surrogate and valid time from the tile's + /// own columns. `live_row` is the row's index in the attr columns. + pub fn push_tile_row( + &mut self, + tile: &SparseTile, + row: usize, + live_row: usize, + coord: &[CoordValue], + ) -> ArrayResult<()> { + let corrupt = |detail: String| ArrayError::SegmentCorruption { detail }; + if coord.len() != self.schema.arity() { + return Err(ArrayError::CoordArityMismatch { + array: self.schema.name.clone(), + expected: self.schema.arity(), + got: coord.len(), + }); + } + if tile.attr_cols.len() != self.schema.attrs.len() + || tile.attr_cols.iter().any(|col| col.len() <= live_row) + { + return Err(corrupt(format!( + "tile live row {live_row} is missing from an attr column of array '{}'", + self.schema.name + ))); + } + let surrogate = tile.live_surrogate(row)?; + let valid_from_ms = *tile + .valid_from_ms + .get(row) + .ok_or_else(|| corrupt(format!("tile row {row} has no valid_from_ms")))?; + let valid_until_ms = *tile + .valid_until_ms + .get(row) + .ok_or_else(|| corrupt(format!("tile row {row} has no valid_until_ms")))?; + for (i, c) in coord.iter().enumerate() { + self.push_dim(i, c); + } + for (out, col) in self.attr_cols.iter_mut().zip(&tile.attr_cols) { + out.push(col[live_row].clone()); + } + self.surrogates.push(Some(surrogate)); + self.valid_from_ms.push(valid_from_ms); + self.valid_until_ms.push(valid_until_ms); + self.row_kinds.push(RowKind::Live.as_u8()); + self.mbr.fold_with(coord, |i| &tile.attr_cols[i][live_row]); + Ok(()) + } + /// Push a derived live row, such as an elementwise result. It holds no /// surrogate, so the tile it lands in is a query result and never stored. /// A stored cell goes through [`Self::push_row`] with its bound surrogate. diff --git a/nodedb-array/src/types/cell_value/mod.rs b/nodedb-array/src/types/cell_value/mod.rs index 173c7f69f..3829274f7 100644 --- a/nodedb-array/src/types/cell_value/mod.rs +++ b/nodedb-array/src/types/cell_value/mod.rs @@ -1,5 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 +pub mod narrow; pub mod value; +pub mod widen; +pub use narrow::NarrowError; pub use value::CellValue; +pub use widen::widen_f32; diff --git a/nodedb-array/src/types/cell_value/narrow.rs b/nodedb-array/src/types/cell_value/narrow.rs new file mode 100644 index 000000000..2744b24c0 --- /dev/null +++ b/nodedb-array/src/types/cell_value/narrow.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Narrowing a numeric literal to an attribute's declared type. +//! +//! An integer outside the declared width is refused. A float narrows to +//! `f32` by rounding, and only a finite value that rounds to infinity is +//! refused. Both refusals are SQLSTATE `22003` at the SQL surface. + +use thiserror::Error; + +use crate::schema::attr_spec::AttrType; +use crate::types::cell_value::value::CellValue; + +/// Why a literal does not fit an attribute type. +#[derive(Debug, Clone, PartialEq, Error)] +pub enum NarrowError { + /// The value is outside the type's range. + #[error("value {value} is out of range for type {dtype}")] + OutOfRange { value: String, dtype: &'static str }, + /// The literal's kind cannot be stored in the type at all. + #[error("a {got} value cannot be stored in a {dtype} attribute")] + WrongKind { + got: &'static str, + dtype: &'static str, + }, +} + +impl NarrowError { + /// The array error for this refusal on attribute `attr` of `array`. + pub fn for_attr(self, array: &str, attr: &str) -> crate::error::ArrayError { + crate::error::ArrayError::Narrow { + array: array.to_string(), + attr: attr.to_string(), + source: self, + } + } + + /// Whether the value is out of range, SQLSTATE `22003`, as opposed to a + /// kind the type never holds. + pub fn is_out_of_range(&self) -> bool { + match self { + NarrowError::OutOfRange { .. } => true, + NarrowError::WrongKind { .. } => false, + } + } +} + +impl CellValue { + /// The cell an integer literal stores as in a `dtype` attribute. + pub fn from_int_literal(v: i64, dtype: AttrType) -> Result { + let out_of_range = || NarrowError::OutOfRange { + value: v.to_string(), + dtype: dtype.sql_name(), + }; + match dtype { + AttrType::Int64 => Ok(CellValue::Int64(v)), + AttrType::UInt8 => u8::try_from(v) + .map(CellValue::UInt8) + .map_err(|_| out_of_range()), + AttrType::Int16 => i16::try_from(v) + .map(CellValue::Int16) + .map_err(|_| out_of_range()), + AttrType::UInt16 => u16::try_from(v) + .map(CellValue::UInt16) + .map_err(|_| out_of_range()), + AttrType::Int32 => i32::try_from(v) + .map(CellValue::Int32) + .map_err(|_| out_of_range()), + AttrType::UInt32 => u32::try_from(v) + .map(CellValue::UInt32) + .map_err(|_| out_of_range()), + AttrType::Float64 => Ok(CellValue::Float64(v as f64)), + AttrType::Float32 => Ok(CellValue::Float32(v as f32)), + AttrType::String | AttrType::Bytes => Err(NarrowError::WrongKind { + got: "integer", + dtype: dtype.sql_name(), + }), + } + } + + /// The cell a float literal stores as in a `dtype` attribute. + pub fn from_float_literal(v: f64, dtype: AttrType) -> Result { + match dtype { + AttrType::Float64 => Ok(CellValue::Float64(v)), + AttrType::Float32 => { + let narrowed = v as f32; + if v.is_finite() && narrowed.is_infinite() { + return Err(NarrowError::OutOfRange { + value: v.to_string(), + dtype: dtype.sql_name(), + }); + } + Ok(CellValue::Float32(narrowed)) + } + AttrType::Int64 + | AttrType::UInt8 + | AttrType::Int16 + | AttrType::UInt16 + | AttrType::Int32 + | AttrType::UInt32 + | AttrType::String + | AttrType::Bytes => Err(NarrowError::WrongKind { + got: "float", + dtype: dtype.sql_name(), + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn integers_store_at_their_width() { + assert_eq!( + CellValue::from_int_literal(255, AttrType::UInt8), + Ok(CellValue::UInt8(255)) + ); + assert_eq!( + CellValue::from_int_literal(-32768, AttrType::Int16), + Ok(CellValue::Int16(-32768)) + ); + assert_eq!( + CellValue::from_int_literal(65535, AttrType::UInt16), + Ok(CellValue::UInt16(65535)) + ); + assert_eq!( + CellValue::from_int_literal(4_294_967_295, AttrType::UInt32), + Ok(CellValue::UInt32(u32::MAX)) + ); + assert_eq!( + CellValue::from_int_literal(3, AttrType::Float32), + Ok(CellValue::Float32(3.0)) + ); + } + + #[test] + fn integers_past_the_width_are_refused() { + let cases = [ + (256, AttrType::UInt8), + (-1, AttrType::UInt8), + (32768, AttrType::Int16), + (-32769, AttrType::Int16), + (65536, AttrType::UInt16), + (-1, AttrType::UInt16), + (2_147_483_648, AttrType::Int32), + (4_294_967_296, AttrType::UInt32), + (-1, AttrType::UInt32), + ]; + for (v, dtype) in cases { + assert!( + matches!( + CellValue::from_int_literal(v, dtype), + Err(NarrowError::OutOfRange { .. }) + ), + "{v} into {dtype:?}" + ); + } + } + + #[test] + fn floats_round_into_float32() { + assert_eq!( + CellValue::from_float_literal(0.1, AttrType::Float32), + Ok(CellValue::Float32(0.1)) + ); + assert!(matches!( + CellValue::from_float_literal(1e39, AttrType::Float32), + Err(NarrowError::OutOfRange { .. }) + )); + assert!(matches!( + CellValue::from_float_literal(f64::INFINITY, AttrType::Float32), + Ok(CellValue::Float32(v)) if v.is_infinite() + )); + } + + #[test] + fn a_float_is_never_an_integer() { + assert!(matches!( + CellValue::from_float_literal(1.0, AttrType::Int16), + Err(NarrowError::WrongKind { .. }) + )); + assert!(matches!( + CellValue::from_int_literal(1, AttrType::String), + Err(NarrowError::WrongKind { .. }) + )); + } +} diff --git a/nodedb-array/src/types/cell_value/value.rs b/nodedb-array/src/types/cell_value/value.rs index d9cc9c565..659fc1295 100644 --- a/nodedb-array/src/types/cell_value/value.rs +++ b/nodedb-array/src/types/cell_value/value.rs @@ -4,11 +4,14 @@ //! //! A cell's value is one [`CellValue`] per attribute defined on the //! schema. The schema's `attrs.len()` and ordering determine the -//! payload layout — validation lives in -//! [`crate::schema::validation::attrs`]. +//! payload layout. A write checks a cell against the schema with +//! `ArraySchema::check_cell`. use serde::{Deserialize, Serialize}; +use crate::schema::attr_spec::AttrType; +use crate::types::cell_value::widen::widen_f32; + /// Typed attribute payload. /// /// Keep this enum strictly in sync with [`crate::schema::AttrType`]: @@ -31,12 +34,80 @@ pub enum CellValue { Bytes(Vec), /// Cell present but value absent — distinct from "not written". Null, + UInt8(u8), + Int16(i16), + UInt16(u16), + Int32(i32), + UInt32(u32), + Float32(f32), } impl CellValue { pub fn is_null(&self) -> bool { matches!(self, CellValue::Null) } + + /// The attribute type this value belongs to. `Null` belongs to none. + pub fn attr_type(&self) -> Option { + match self { + CellValue::Int64(_) => Some(AttrType::Int64), + CellValue::Float64(_) => Some(AttrType::Float64), + CellValue::String(_) => Some(AttrType::String), + CellValue::Bytes(_) => Some(AttrType::Bytes), + CellValue::Null => None, + CellValue::UInt8(_) => Some(AttrType::UInt8), + CellValue::Int16(_) => Some(AttrType::Int16), + CellValue::UInt16(_) => Some(AttrType::UInt16), + CellValue::Int32(_) => Some(AttrType::Int32), + CellValue::UInt32(_) => Some(AttrType::UInt32), + CellValue::Float32(_) => Some(AttrType::Float32), + } + } + + /// The value of an integer cell, widened to `i64`. `None` for every + /// other cell. + pub fn as_i64(&self) -> Option { + match self { + CellValue::Int64(v) => Some(*v), + CellValue::UInt8(v) => Some(i64::from(*v)), + CellValue::Int16(v) => Some(i64::from(*v)), + CellValue::UInt16(v) => Some(i64::from(*v)), + CellValue::Int32(v) => Some(i64::from(*v)), + CellValue::UInt32(v) => Some(i64::from(*v)), + CellValue::Float64(_) + | CellValue::Float32(_) + | CellValue::String(_) + | CellValue::Bytes(_) + | CellValue::Null => None, + } + } + + /// The value of a float cell, widened to `f64`. A `Float32` widens + /// through [`widen_f32`]. `None` for every other cell, integer cells + /// included. + pub fn as_float(&self) -> Option { + match self { + CellValue::Float64(v) => Some(*v), + CellValue::Float32(v) => Some(widen_f32(*v)), + CellValue::Int64(_) + | CellValue::UInt8(_) + | CellValue::Int16(_) + | CellValue::UInt16(_) + | CellValue::Int32(_) + | CellValue::UInt32(_) + | CellValue::String(_) + | CellValue::Bytes(_) + | CellValue::Null => None, + } + } + + /// The value of any numeric cell as `f64`. An `Int64` past 2^53 rounds. + pub fn as_f64(&self) -> Option { + match self.as_i64() { + Some(v) => Some(v as f64), + None => self.as_float(), + } + } } #[cfg(test)] @@ -48,4 +119,14 @@ mod tests { assert!(CellValue::Null.is_null()); assert!(!CellValue::Int64(0).is_null()); } + + #[test] + fn numeric_cells_widen() { + assert_eq!(CellValue::UInt16(65535).as_i64(), Some(65535)); + assert_eq!(CellValue::Int16(-3).as_f64(), Some(-3.0)); + assert_eq!(CellValue::Float32(0.5).as_float(), Some(0.5)); + assert_eq!(CellValue::Float32(0.5).as_i64(), None); + assert_eq!(CellValue::UInt8(1).attr_type(), Some(AttrType::UInt8)); + assert_eq!(CellValue::Null.attr_type(), None); + } } diff --git a/nodedb-array/src/types/cell_value/widen.rs b/nodedb-array/src/types/cell_value/widen.rs new file mode 100644 index 000000000..1f074c6cd --- /dev/null +++ b/nodedb-array/src/types/cell_value/widen.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Widening a stored `f32` to the `f64` every output path carries. +//! +//! The `f64` is the nearest double to the shortest decimal that names the +//! `f32`, so `0.1` stored as `FLOAT32` reads as `0.1` over JSON, the native +//! protocol, and pgwire alike. Narrowing that `f64` back gives the stored +//! `f32` bit for bit. + +/// The `f64` an `f32` cell reads as. +pub fn widen_f32(v: f32) -> f64 { + let exact = f64::from(v); + // The shortest decimal of a finite `f32` round-trips to it. A value that + // does not, such as a NaN payload, widens exactly. + v.to_string() + .parse::() + .ok() + .filter(|d| (*d as f32).to_bits() == v.to_bits()) + .unwrap_or(exact) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn point_one_widens_to_point_one() { + assert_eq!(widen_f32(0.1), 0.1_f64); + assert_ne!(f64::from(0.1_f32), 0.1_f64); + } + + #[test] + fn f32_max_widens_to_its_shortest_decimal() { + assert_eq!(widen_f32(f32::MAX), 3.4028235e38_f64); + assert_eq!(widen_f32(f32::MAX) as f32, f32::MAX); + } + + #[test] + fn every_widened_value_narrows_back_bit_exact() { + for v in [ + 0.0_f32, + -0.0, + f32::from_bits(1), + f32::MIN_POSITIVE, + 0.3, + -123.456, + 16_777_217.0, + f32::INFINITY, + f32::NEG_INFINITY, + ] { + assert_eq!((widen_f32(v) as f32).to_bits(), v.to_bits(), "{v}"); + } + assert!(widen_f32(f32::NAN).is_nan()); + } +} diff --git a/nodedb-client-tests/Cargo.toml b/nodedb-client-tests/Cargo.toml index 22cc6923d..d022de19f 100644 --- a/nodedb-client-tests/Cargo.toml +++ b/nodedb-client-tests/Cargo.toml @@ -23,3 +23,6 @@ tokio-postgres = { workspace = true } tempfile = { workspace = true } serde_json = { workspace = true } sonic-rs = { workspace = true } +# The array load measurement installs jemalloc and reads its counters. +tikv-jemallocator = { workspace = true } +tikv-jemalloc-ctl = { workspace = true } diff --git a/nodedb-client-tests/tests/array_load_measure.rs b/nodedb-client-tests/tests/array_load_measure.rs new file mode 100644 index 000000000..0a78a10f7 --- /dev/null +++ b/nodedb-client-tests/tests/array_load_measure.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Measurement of one 1024 x 1024 x 6 array load through +//! `NativeClient::array_put_cells` and the server's bulk-load core. +//! +//! The server runs in this process, so this binary installs jemalloc as its +//! global allocator and the counters cover the server's allocations. The +//! test prints the jemalloc allocation count and bytes of the load alone. +//! +//! Judge a change by instruction and allocation counts, never wall time. +//! Run it ignored, then under `perf stat`, once with the load and once +//! without, and subtract: +//! +//! ```text +//! cargo nextest run -p nodedb-client-tests --test array_load_measure \ +//! --run-ignored only --no-capture +//! BIN=$(ls -t target/debug/deps/array_load_measure-* | grep -v '\.d$' | head -1) +//! perf stat -e instructions:u -- "$BIN" --ignored --exact \ +//! array_load_1024x1024x6 --nocapture +//! NODEDB_MEASURE_SKIP_LOAD=1 perf stat -e instructions:u -- "$BIN" --ignored \ +//! --exact array_load_1024x1024x6 --nocapture +//! ``` +//! +//! The first `perf stat` minus the second is the load's instruction count. +//! Rebuild both sides of an A/B before measuring, and never measure while a +//! build runs. + +use nodedb_client::native::pool::PoolConfig; +use nodedb_client::{NativeClient, NodeDb, Value}; +use nodedb_test_support::pgwire_harness::TestServer; +use nodedb_types::protocol::ArrayCellWrite; + +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +const SIDE: i64 = 1024; +const BANDS: usize = 6; + +/// jemalloc's totals over every arena: allocation calls and bytes allocated. +fn alloc_counters() -> (u64, u64) { + tikv_jemalloc_ctl::epoch::advance().expect("advance the jemalloc epoch"); + // SAFETY: each name is a NUL-terminated jemalloc ctl key of type u64. + let calls = unsafe { + tikv_jemalloc_ctl::raw::read::(b"stats.arenas.4096.small.nmalloc\0") + .expect("small nmalloc") + + tikv_jemalloc_ctl::raw::read::(b"stats.arenas.4096.large.nmalloc\0") + .expect("large nmalloc") + }; + let bytes = tikv_jemalloc_ctl::stats::allocated::read().expect("stats.allocated") as u64; + (calls, bytes) +} + +fn scene() -> Vec { + (0..SIDE) + .flat_map(|y| { + (0..SIDE).map(move |x| ArrayCellWrite { + coords: vec![Value::Integer(y), Value::Integer(x)], + attrs: (1..=BANDS) + .map(|k| Value::Float(((y * SIDE + x) * k as i64) as f64 + 0.5)) + .collect(), + }) + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "measurement: run with --run-ignored only under perf stat"] +async fn array_load_1024x1024x6() { + let server = TestServer::start().await; + let client = NativeClient::new(PoolConfig::new( + format!("127.0.0.1:{}", server.native_port), + nodedb_types::protocol::AuthMethod::Trust { + username: "nodedb".into(), + }, + )); + client + .execute_sql( + "CREATE ARRAY scene DIMS (y INT64 [0..1023], x INT64 [0..1023]) \ + ATTRS (b1 FLOAT64, b2 FLOAT64, b3 FLOAT64, b4 FLOAT64, b5 FLOAT64, b6 FLOAT64) \ + TILE_EXTENTS (128, 128)", + &[], + ) + .await + .expect("CREATE ARRAY scene"); + let cells = scene(); + if std::env::var_os("NODEDB_MEASURE_SKIP_LOAD").is_some() { + return; + } + + let (calls_before, bytes_before) = alloc_counters(); + let loaded = client + .array_put_cells("scene", cells) + .await + .expect("array_put_cells"); + let (calls_after, bytes_after) = alloc_counters(); + + assert_eq!(loaded, (SIDE * SIDE) as u64); + println!( + "array_load_1024x1024x6: {loaded} cells, {} jemalloc allocations, \ + {} bytes allocated at the end versus the start", + calls_after - calls_before, + bytes_after as i64 - bytes_before as i64 + ); +} diff --git a/nodedb-client-tests/tests/array_put_cells_round_trip.rs b/nodedb-client-tests/tests/array_put_cells_round_trip.rs new file mode 100644 index 000000000..d7e91fecc --- /dev/null +++ b/nodedb-client-tests/tests/array_put_cells_round_trip.rs @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `NativeClient::array_put_cells` loads cells through `ArrayPutBatch` +//! requests and the server's array bulk-load core. The pgwire client +//! answers the call with the trait's unsupported error. + +use std::collections::HashMap; + +use nodedb_client::native::pool::PoolConfig; +use nodedb_client::{NativeClient, NodeDb, NodeDbRemote, Value}; +use nodedb_test_support::pgwire_harness::TestServer; +use nodedb_types::protocol::ArrayCellWrite; + +const SIDE: i64 = 64; + +fn native(server: &TestServer) -> NativeClient { + NativeClient::new(PoolConfig::new( + format!("127.0.0.1:{}", server.native_port), + nodedb_types::protocol::AuthMethod::Trust { + username: "nodedb".into(), + }, + )) +} + +async fn create_tile(client: &NativeClient, name: &str) { + client + .execute_sql( + &format!( + "CREATE ARRAY {name} DIMS (y INT64 [0..63], x INT64 [0..63]) \ + ATTRS (v FLOAT64) TILE_EXTENTS (16, 16)" + ), + &[], + ) + .await + .unwrap_or_else(|e| panic!("CREATE ARRAY {name}: {e}")); +} + +fn value_of(y: i64, x: i64) -> f64 { + (y * SIDE + x) as f64 + 0.25 +} + +fn cells() -> Vec { + (0..SIDE) + .flat_map(|y| { + (0..SIDE).map(move |x| ArrayCellWrite { + coords: vec![Value::Integer(y), Value::Integer(x)], + attrs: vec![Value::Float(value_of(y, x))], + }) + }) + .collect() +} + +fn number(value: &Value) -> f64 { + match value { + Value::Float(f) => *f, + Value::Integer(i) => *i as f64, + Value::String(s) => s.parse().unwrap_or_else(|_| panic!("not a number: {s}")), + other => panic!("not a number: {other:?}"), + } +} + +/// Every cell of `name`, keyed by `(y, x)`, read by column name. +async fn read_cells(client: &NativeClient, name: &str) -> HashMap<(i64, i64), f64> { + let result = client + .execute_sql( + &format!( + "SELECT * FROM ARRAY_SLICE('{name}', '{{\"y\":[0,63],\"x\":[0,63]}}', '*', 100000)" + ), + &[], + ) + .await + .unwrap_or_else(|e| panic!("ARRAY_SLICE {name}: {e}")); + // An empty result carries no columns. + if result.rows.is_empty() { + return HashMap::new(); + } + let column = |wanted: &str| { + result + .columns + .iter() + .position(|c| c == wanted) + .unwrap_or_else(|| panic!("column {wanted} in {:?}", result.columns)) + }; + let (y, x, v) = (column("y"), column("x"), column("v")); + result + .rows + .iter() + .map(|row| { + ( + (number(&row[y]) as i64, number(&row[x]) as i64), + number(&row[v]), + ) + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn native_array_put_cells_loads_every_cell() { + let server = TestServer::start().await; + let client = native(&server); + create_tile(&client, "native_tile").await; + + let loaded = client + .array_put_cells("native_tile", cells()) + .await + .expect("array_put_cells"); + assert_eq!(loaded, (SIDE * SIDE) as u64); + + let read = read_cells(&client, "native_tile").await; + assert_eq!(read.len(), (SIDE * SIDE) as usize); + for y in 0..SIDE { + for x in 0..SIDE { + assert_eq!(read.get(&(y, x)), Some(&value_of(y, x)), "cell ({y}, {x})"); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn a_cell_of_the_wrong_type_stops_the_load_with_the_committed_count() { + let server = TestServer::start().await; + let client = native(&server); + create_tile(&client, "native_bad").await; + + let mut batch = cells(); + batch[100].attrs[0] = Value::String("not a float".into()); + let message = client + .array_put_cells("native_bad", batch) + .await + .expect_err("a string in a FLOAT64 attr is refused") + .to_string(); + assert!(message.contains("batch row 101"), "{message}"); + let head = message + .rfind(" rows committed)") + .map(|end| &message[..end]) + .unwrap_or_else(|| panic!("no committed count in {message}")); + let committed: usize = head[head.rfind('(').expect("open paren") + 1..] + .parse() + .unwrap_or_else(|_| panic!("committed count in {message}")); + assert!(committed < 100, "{message}"); + assert_eq!( + read_cells(&client, "native_bad").await.len(), + committed, + "exactly the committed cells are in the array" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn the_pgwire_client_answers_unsupported() { + let server = TestServer::start().await; + let conn_str = format!( + "host=127.0.0.1 port={} user=nodedb dbname=default", + server.pg_port + ); + let remote = NodeDbRemote::connect(&conn_str) + .await + .expect("pgwire connect to harness must succeed"); + let err = remote + .array_put_cells("anything", cells()) + .await + .expect_err("the pgwire client has no array_put_cells"); + assert!(err.to_string().contains("array_put_cells"), "{err}"); +} diff --git a/nodedb-client/src/native/client/array.rs b/nodedb-client/src/native/client/array.rs new file mode 100644 index 000000000..734e6944b --- /dev/null +++ b/nodedb-client/src/native/client/array.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Array bulk-load implementation for `NativeClient`: `ArrayPutBatch` +//! requests over one connection. + +use nodedb_types::error::{NodeDbError, NodeDbResult}; +use nodedb_types::protocol::{ArrayCellWrite, ArrayPutBatch, MAX_FRAME_SIZE, OpCode, TextFields}; + +use super::core::NativeClient; +use crate::native::connection::check_error; + +/// The payload size one request targets: three quarters of the frame cap, +/// leaving room for the request envelope. +const TARGET_PAYLOAD_BYTES: usize = (MAX_FRAME_SIZE as usize) * 3 / 4; + +/// How many cells the first request encodes to measure the cell size. +const PROBE_CELLS: usize = 1024; + +impl NativeClient { + /// Send `cells` as `ArrayPutBatch` requests on one connection, each + /// under one frame. Outside a transaction block each chunk of a request + /// commits on its own. A bad cell stops the load: the requests and + /// chunks before it stay committed, and the error names the cell's row + /// in its request and that request's committed row count. Inside a + /// block a refusal aborts the block. Returns the cells loaded. + pub(super) async fn array_put_cells_impl( + &self, + array: &str, + cells: Vec, + ) -> NodeDbResult { + if cells.is_empty() { + return Ok(0); + } + let mut conn = self.pool.acquire().await?; + let per_request = cells_per_request(&cells)?; + let mut loaded = 0u64; + let mut pending: Vec<&[ArrayCellWrite]> = cells.chunks(per_request).rev().collect(); + while let Some(slice) = pending.pop() { + let payload = + zerompk::to_msgpack_vec(&ArrayPutBatch::from_cells(slice)?).map_err(|e| { + NodeDbError::serialization("msgpack", format!("ArrayPutBatch: {e}")) + })?; + // Cells vary in size. A slice over the target splits in two, and + // the halves go first, in order. + if payload.len() > TARGET_PAYLOAD_BYTES && slice.len() > 1 { + let (head, tail) = slice.split_at(slice.len() / 2); + pending.push(tail); + pending.push(head); + continue; + } + let request = TextFields { + collection: Some(array.to_string()), + payload: Some(payload), + ..Default::default() + }; + let resp = conn.send(OpCode::ArrayPutBatch, request).await?; + check_error(&resp)?; + loaded += resp.rows_affected.ok_or_else(|| { + NodeDbError::internal("ArrayPutBatch: the server answered with no row count") + })?; + } + Ok(loaded) + } +} + +/// How many cells one request carries, from the encoded size of a probe. +fn cells_per_request(cells: &[ArrayCellWrite]) -> NodeDbResult { + let probe = &cells[..cells.len().min(PROBE_CELLS)]; + let bytes = zerompk::to_msgpack_vec(&ArrayPutBatch::from_cells(probe)?) + .map_err(|e| NodeDbError::serialization("msgpack", format!("ArrayPutBatch: {e}")))? + .len(); + let per_cell = bytes.div_ceil(probe.len()).max(1); + Ok((TARGET_PAYLOAD_BYTES / per_cell).max(1)) +} diff --git a/nodedb-client/src/native/client/dispatch.rs b/nodedb-client/src/native/client/dispatch.rs index a5509d29f..147583f20 100644 --- a/nodedb-client/src/native/client/dispatch.rs +++ b/nodedb-client/src/native/client/dispatch.rs @@ -207,6 +207,14 @@ impl NodeDb for NativeClient { async fn execute_sql(&self, query: &str, params: &[Value]) -> NodeDbResult { self.execute_sql_impl(query, params).await } + + async fn array_put_cells( + &self, + array: &str, + cells: Vec, + ) -> NodeDbResult { + self.array_put_cells_impl(array, cells).await + } } #[cfg(test)] diff --git a/nodedb-client/src/native/client/mod.rs b/nodedb-client/src/native/client/mod.rs index e232e9dfa..b024f4649 100644 --- a/nodedb-client/src/native/client/mod.rs +++ b/nodedb-client/src/native/client/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 +mod array; pub mod core; mod crdt_list; mod dispatch; diff --git a/nodedb-client/src/traits/core/default_impls/batch.rs b/nodedb-client/src/traits/core/default_impls/batch.rs index 718fa8a20..83c9e8958 100644 --- a/nodedb-client/src/traits/core/default_impls/batch.rs +++ b/nodedb-client/src/traits/core/default_impls/batch.rs @@ -2,9 +2,21 @@ use nodedb_types::error::{NodeDbError, NodeDbResult}; use nodedb_types::id::NodeId; +use nodedb_types::protocol::ArrayCellWrite; use super::super::trait_def::NodeDb; +pub(in crate::traits::core) fn array_put_cells_default( + array: &str, + cells: Vec, +) -> NodeDbResult { + let _ = cells; + Err(NodeDbError::storage(format!( + "array_put_cells is not implemented on this client; load array '{array}' \ + through the native client, or with COPY {array} FROM STDIN" + ))) +} + pub(in crate::traits::core) async fn batch_vector_insert_default( this: &T, collection: &str, diff --git a/nodedb-client/src/traits/core/default_impls/mod.rs b/nodedb-client/src/traits/core/default_impls/mod.rs index ac901aa3a..ba34fe33d 100644 --- a/nodedb-client/src/traits/core/default_impls/mod.rs +++ b/nodedb-client/src/traits/core/default_impls/mod.rs @@ -9,7 +9,9 @@ mod collections; mod dispatch; mod graph; -pub(super) use batch::{batch_graph_insert_edges_default, batch_vector_insert_default}; +pub(super) use batch::{ + array_put_cells_default, batch_graph_insert_edges_default, batch_vector_insert_default, +}; pub(super) use collections::{ drop_collection_purge_default, list_dropped_collections_default, on_collection_purged_default, undrop_collection_default, diff --git a/nodedb-client/src/traits/core/trait_def.rs b/nodedb-client/src/traits/core/trait_def.rs index b82e912e7..482576203 100644 --- a/nodedb-client/src/traits/core/trait_def.rs +++ b/nodedb-client/src/traits/core/trait_def.rs @@ -19,16 +19,12 @@ use std::collections::HashSet; use async_trait::async_trait; -use nodedb_types::document::Document; -use nodedb_types::dropped_collection::DroppedCollection; -use nodedb_types::error::NodeDbResult; -use nodedb_types::filter::{EdgeFilter, MetadataFilter}; -use nodedb_types::graph::GraphStats; -use nodedb_types::id::{EdgeId, NodeId}; -use nodedb_types::protocol::Limits; -use nodedb_types::result::{QueryResult, SearchResult, SubGraph}; -use nodedb_types::text_search::TextSearchParams; -use nodedb_types::value::Value; +use nodedb_types::{ + document::Document, dropped_collection::DroppedCollection, error::NodeDbResult, + filter::EdgeFilter, filter::MetadataFilter, graph::GraphStats, id::EdgeId, id::NodeId, + protocol::ArrayCellWrite, protocol::Limits, result::QueryResult, result::SearchResult, + result::SubGraph, text_search::TextSearchParams, value::Value, +}; use super::default_impls; use super::marker::NodeDbMarker; @@ -422,6 +418,15 @@ pub trait NodeDb: NodeDbMarker { default_impls::batch_graph_insert_edges_default(self, collection, edges).await } + /// Load `cells` into `array` through the bulk path. Returns the count. + /// A cell holds one coordinate per dim and one value per attr, in schema + /// order. Outside a block a bad cell stops the load, earlier chunks kept. + /// On Native: `OpCode::ArrayPutBatch` requests, each under one frame. + /// The default returns `Err`: the pgwire client does not implement it. + async fn array_put_cells(&self, array: &str, cells: Vec) -> NodeDbResult { + default_impls::array_put_cells_default(array, cells) + } + /// The protocol version negotiated during the connection handshake. /// `0` for an implementation that performs no handshake. fn proto_version(&self) -> u16 { diff --git a/nodedb-cluster-tests/tests/common_suite/cases/bitemporal_array_cluster.rs b/nodedb-cluster-tests/tests/common_suite/cases/bitemporal_array_cluster.rs index f174d3465..aad7fbfe1 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/bitemporal_array_cluster.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/bitemporal_array_cluster.rs @@ -53,31 +53,11 @@ async fn query_col0(client: &tokio_postgres::Client, sql: &str) -> Vec { .collect() } -/// Parse the `attrs[0]` int64 from a single ARRAY_SLICE result row. -/// -/// `SELECT *` over ARRAY_SLICE expands to one pgwire field per declared -/// column. This test uses the column-0 helper, which lands on the -/// `attrs` JSON-array column for the slice TVF (e.g. `[99]`). Older -/// codecs returned the full `{"coords": [...], "attrs": [v]}` envelope -/// in that slot — accept either shape so the test stays robust to the -/// projection tightening. +/// Parse the int64 attr value of a single-column `SELECT v FROM +/// ARRAY_SLICE(...)` row. A slice row names each attr as its own column. fn parse_int_attr(row: &str) -> i64 { - let v: serde_json::Value = - serde_json::from_str(row).unwrap_or_else(|e| panic!("row is not JSON: {row}: {e}")); - let arr = match v { - // Bare attrs array — `[99]`. - serde_json::Value::Array(a) => a, - // Envelope with `attrs` field. - serde_json::Value::Object(map) => map - .get("attrs") - .and_then(|a| a.as_array()) - .cloned() - .unwrap_or_else(|| panic!("envelope missing attrs array: {row}")), - _ => panic!("row is neither array nor object: {row}"), - }; - arr.first() - .and_then(|n| n.as_i64()) - .unwrap_or_else(|| panic!("cannot extract attrs[0] as i64 from: {row}")) + row.parse() + .unwrap_or_else(|e| panic!("attr `v` is not an int64: {row}: {e}")) } /// Wall-clock milliseconds since the Unix epoch. @@ -251,7 +231,7 @@ async fn cluster_array_as_of_returns_consistent_version_across_shards() { // Live read (no AS OF) must return v2 = 99. let live_rows = query_col0( client, - "SELECT * FROM ARRAY_SLICE('bt2', '{x: [0, 0]}', ['v'], 10)", + "SELECT v FROM ARRAY_SLICE('bt2', '{x: [0, 0]}', ['v'], 10)", ) .await; assert_eq!( @@ -269,7 +249,7 @@ async fn cluster_array_as_of_returns_consistent_version_across_shards() { let as_of_rows = query_col0( client, &format!( - "SELECT * FROM ARRAY_SLICE('bt2', '{{x: [0, 0]}}', ['v'], 10) \ + "SELECT v FROM ARRAY_SLICE('bt2', '{{x: [0, 0]}}', ['v'], 10) \ AS OF SYSTEM TIME {cutoff_ms}" ), ) @@ -291,7 +271,7 @@ async fn cluster_array_as_of_returns_consistent_version_across_shards() { // before the earliest stamp. let before_rows = query_col0( client, - "SELECT * FROM ARRAY_SLICE('bt2', '{x: [0, 0]}', ['v'], 10) \ + "SELECT v FROM ARRAY_SLICE('bt2', '{x: [0, 0]}', ['v'], 10) \ AS OF SYSTEM TIME 1", ) .await; diff --git a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array.rs b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array.rs index af2c15c38..e2364386e 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array.rs @@ -143,9 +143,8 @@ async fn cluster_array_slice_spans_multiple_shards() { async fn assert_chr1_slice(client: &tokio_postgres::Client, node_idx: usize) { // Slice chr=1, pos 0..99 — should return exactly the 3 cells for chr=1. - // ARRAY_SLICE projects one pgwire field per declared column: a - // `coords` JSON-array column followed by an `attrs` JSON-array column - // carrying the requested attribute values in projection order. + // ARRAY_SLICE projects one pgwire field per dim, then one per requested + // attr, each named after it. let rows = query_named_rows( client, "SELECT * FROM ARRAY_SLICE('genome', '{chr: [1, 1], pos: [0, 99]}', ['qual'], 100)", @@ -161,16 +160,17 @@ async fn assert_chr1_slice(client: &tokio_postgres::Client, node_idx: usize) { let mut quals: Vec = rows .iter() .map(|row| { - let attrs_text = row - .get("attrs") - .unwrap_or_else(|| panic!("missing attrs column in {row:?}")); - let attrs: serde_json::Value = serde_json::from_str(attrs_text) - .unwrap_or_else(|e| panic!("attrs not JSON: {attrs_text}: {e}")); - attrs - .as_array() - .and_then(|a| a.first()) - .and_then(|v| v.as_f64()) - .unwrap_or_else(|| panic!("missing qual in row: {row:?}")) + assert_eq!( + row.get("chr").map(String::as_str), + Some("1"), + "node {node_idx}: every row is on chr=1: {row:?}" + ); + assert!(row.contains_key("pos"), "missing pos column in {row:?}"); + let qual = row + .get("qual") + .unwrap_or_else(|| panic!("missing qual column in {row:?}")); + qual.parse::() + .unwrap_or_else(|e| panic!("qual not a float: {qual}: {e}")) }) .collect(); quals.sort_by(|a, b| a.partial_cmp(b).unwrap()); diff --git a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_cell_raft_replication.rs b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_cell_raft_replication.rs index 7cf464eb8..6a06c8867 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_cell_raft_replication.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_cell_raft_replication.rs @@ -79,7 +79,8 @@ fn array_surrogate( .map(|s| s.as_u32()) } -/// Run `ARRAY_SLICE` and return each row's `attrs` JSON text. +/// Run `ARRAY_SLICE` and return each row's `v` attr text. A slice row names +/// each attr as its own column. async fn slice_attrs(client: &tokio_postgres::Client, sql: &str) -> Vec { let msgs = client .simple_query(sql) @@ -91,7 +92,7 @@ async fn slice_attrs(client: &tokio_postgres::Client, sql: &str) -> Vec let names: Vec = r.columns().iter().map(|c| c.name().to_string()).collect(); names .iter() - .position(|n| n == "attrs") + .position(|n| n == "v") .and_then(|i| r.get(i)) .map(|s| s.to_string()) } else { @@ -187,13 +188,9 @@ async fn cluster_array_cell_write_replicates_to_all_replicas() { 1, "ARRAY_SLICE for x=5 must return exactly the one cell; got {attrs:?}" ); - let parsed: serde_json::Value = serde_json::from_str(&attrs[0]) - .unwrap_or_else(|e| panic!("attrs not JSON: {}: {e}", attrs[0])); - let v = parsed - .as_array() - .and_then(|a| a.first()) - .and_then(|x| x.as_f64()) - .unwrap_or_else(|| panic!("missing v in slice row: {}", attrs[0])); + let v: f64 = attrs[0] + .parse() + .unwrap_or_else(|e| panic!("attr v is not a float: {}: {e}", attrs[0])); assert!((v - 42.0).abs() < 1e-6, "expected v=42.0, got {v}"); cluster.shutdown().await; diff --git a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_project_elementwise.rs b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_project_elementwise.rs new file mode 100644 index 000000000..133c3a749 --- /dev/null +++ b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_project_elementwise.rs @@ -0,0 +1,200 @@ +// SPDX-License-Identifier: BUSL-1.1 +//! 3-node cluster tests for `ARRAY_PROJECT`, `ARRAY_ELEMENTWISE`, +//! `ARRAY_FLUSH` and `ARRAY_COMPACT`. +//! +//! Every test matches `binary(/cluster/)` in nextest.toml, so it runs in the +//! serial `cluster` test group. The arrays are created on one node and +//! written through another. Their 32 one-cell tiles spread over the shards of +//! every node, and every node answers each query with every cell. + +use std::collections::BTreeMap; + +use crate::common; + +use common::cluster_harness::TestCluster; + +const CELLS: i64 = 32; + +/// The DDL of a 1-D array `name` with one FLOAT64 attr `v`. +fn ddl(name: &str) -> String { + format!( + "CREATE ARRAY {name} \ + DIMS (k INT64 [0..31]) \ + ATTRS (v FLOAT64) \ + TILE_EXTENTS (1)" + ) +} + +/// The INSERT that writes `v = scale * k` for every cell of `name`. +fn insert(name: &str, scale: f64) -> String { + let cells: Vec = (0..CELLS) + .map(|k| format!("COORDS ({k}) VALUES ({})", scale * k as f64)) + .collect(); + format!("INSERT INTO ARRAY {name} {}", cells.join(", ")) +} + +/// Spawn a 3-node cluster holding every `(name, scale)` array, written and +/// flushed through a node other than the one that ran the DDL. +async fn spawn_with_arrays(arrays: &[(&str, f64)]) -> TestCluster { + let cluster = TestCluster::spawn_three() + .await + .expect("3-node cluster spawn"); + for &(name, scale) in arrays { + let ddl_idx = cluster + .exec_ddl_on_any_leader(&ddl(name)) + .await + .unwrap_or_else(|e| panic!("CREATE ARRAY {name}: {e}")); + let writer = &cluster.nodes[(ddl_idx + 1) % cluster.nodes.len()]; + writer + .exec(&insert(name, scale)) + .await + .unwrap_or_else(|e| panic!("INSERT INTO ARRAY {name}: {e}")); + writer + .exec(&format!("SELECT ARRAY_FLUSH('{name}')")) + .await + .unwrap_or_else(|e| panic!("ARRAY_FLUSH {name}: {e}")); + } + cluster + .wait_for_full_apply_convergence(std::time::Duration::from_secs(10)) + .await; + cluster +} + +fn int(text: &str) -> i64 { + text.parse() + .unwrap_or_else(|e| panic!("{text:?} is not an integer: {e}")) +} + +fn number(text: &str) -> f64 { + text.parse() + .unwrap_or_else(|e| panic!("{text:?} is not a number: {e}")) +} + +/// The first element of a JSON array column. +fn first(column: &str) -> serde_json::Value { + let value: serde_json::Value = serde_json::from_str(column) + .unwrap_or_else(|e| panic!("column {column:?} is not JSON: {e}")); + value + .as_array() + .and_then(|items| items.first().cloned()) + .unwrap_or_else(|| panic!("column {column:?} is not a non-empty JSON array")) +} + +/// The `(k, v)` cells of `sql` on node `idx`. +async fn cells_on(cluster: &TestCluster, idx: usize, sql: &str) -> BTreeMap { + let rows = cluster.nodes[idx] + .query_named_rows(sql) + .await + .unwrap_or_else(|e| panic!("node {idx}: {sql}: {e}")); + let map: BTreeMap = rows + .iter() + .map(|row| (int(&row["k"]), number(&row["v"]))) + .collect(); + assert_eq!( + map.len(), + rows.len(), + "node {idx}: {sql}: a cell came twice" + ); + map +} + +fn scaled(range: std::ops::Range, scale: f64) -> BTreeMap { + range.map(|k| (k, scale * k as f64)).collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_project_returns_every_cell_on_every_node() { + let cluster = spawn_with_arrays(&[("g", 1.0)]).await; + for idx in 0..cluster.nodes.len() { + let got = cells_on(&cluster, idx, "SELECT k, v FROM ARRAY_PROJECT('g', ['v'])").await; + assert_eq!(got, scaled(0..CELLS, 1.0), "node {idx}"); + } + cluster.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_project_feeds_a_relational_tail() { + let cluster = spawn_with_arrays(&[("g", 1.0)]).await; + for idx in 0..cluster.nodes.len() { + let got = cells_on( + &cluster, + idx, + "SELECT k, v FROM ARRAY_PROJECT('g', ['v']) WHERE v >= 20 ORDER BY k", + ) + .await; + assert_eq!(got, scaled(20..CELLS, 1.0), "node {idx}"); + + let sql = "SELECT count(*) AS n FROM ARRAY_PROJECT('g', ['v'])"; + let rows = cluster.nodes[idx] + .query_named_rows(sql) + .await + .unwrap_or_else(|e| panic!("node {idx}: {sql}: {e}")); + assert_eq!(rows.len(), 1, "node {idx}: one count row: {rows:?}"); + assert_eq!(int(&rows[0]["n"]), CELLS, "node {idx}"); + } + cluster.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_elementwise_pairs_every_cell_on_every_node() { + let cluster = spawn_with_arrays(&[("a", 1.0), ("b", 100.0)]).await; + let sql = "SELECT coords, attrs FROM ARRAY_ELEMENTWISE('a', 'b', 'add', 'v')"; + for (idx, node) in cluster.nodes.iter().enumerate() { + let rows = node + .query_named_rows(sql) + .await + .unwrap_or_else(|e| panic!("node {idx}: {sql}: {e}")); + let got: BTreeMap = rows + .iter() + .map(|row| { + let k = first(&row["coords"]) + .as_i64() + .unwrap_or_else(|| panic!("coords {:?} holds no integer", row["coords"])); + let v = first(&row["attrs"]) + .as_f64() + .unwrap_or_else(|| panic!("attrs {:?} holds no number", row["attrs"])); + (k, v) + }) + .collect(); + assert_eq!(got.len(), rows.len(), "node {idx}: a cell paired twice"); + assert_eq!(got, scaled(0..CELLS, 101.0), "node {idx}"); + } + cluster.shutdown().await; +} + +/// Three more flushed rounds give every shard four segments. The compaction +/// merges them on every shard that holds cells, and every node reads the +/// newest values afterwards. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_flush_and_compact_reach_every_shard() { + let cluster = spawn_with_arrays(&[("m", 1.0)]).await; + for round in 2..=4 { + let node = &cluster.nodes[round % cluster.nodes.len()]; + node.exec(&insert("m", round as f64)) + .await + .unwrap_or_else(|e| panic!("INSERT round {round}: {e}")); + let rows = node + .query_named_rows("SELECT ARRAY_FLUSH('m')") + .await + .unwrap_or_else(|e| panic!("ARRAY_FLUSH round {round}: {e}")); + assert_eq!(int(&rows[0]["flushed"]), 1, "round {round}"); + } + + let rows = cluster.nodes[0] + .query_named_rows("SELECT ARRAY_COMPACT('m')") + .await + .expect("ARRAY_COMPACT"); + assert!( + int(&rows[0]["compacted"]) >= 1, + "the shards that hold cells merge: {rows:?}" + ); + cluster + .wait_for_full_apply_convergence(std::time::Duration::from_secs(10)) + .await; + + for idx in 0..cluster.nodes.len() { + let got = cells_on(&cluster, idx, "SELECT k, v FROM ARRAY_PROJECT('m', ['v'])").await; + assert_eq!(got, scaled(0..CELLS, 4.0), "node {idx}"); + } + cluster.shutdown().await; +} diff --git a/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_region_agg.rs b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_region_agg.rs new file mode 100644 index 000000000..419f3b6de --- /dev/null +++ b/nodedb-cluster-tests/tests/common_suite/cases/cluster_array_region_agg.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: BUSL-1.1 +//! 3-node cluster tests for region-bounded `ARRAY_AGG`. +//! +//! Every test matches `binary(/cluster/)` in nextest.toml, so it runs in the +//! serial `cluster` test group. The arrays are created on one node and +//! written through another. Every node answers each query: the coordinator +//! sends the region to every shard it fans out to and merges the partials. + +use crate::common; + +use common::cluster_harness::TestCluster; + +/// Spawn a 3-node cluster holding `ddl` with `insert` written and flushed +/// through a node other than the one that ran the DDL. +async fn spawn_with_array(ddl: &str, insert: &str, name: &str) -> TestCluster { + let cluster = TestCluster::spawn_three() + .await + .expect("3-node cluster spawn"); + let ddl_idx = cluster + .exec_ddl_on_any_leader(ddl) + .await + .expect("CREATE ARRAY"); + let writer = &cluster.nodes[(ddl_idx + 1) % cluster.nodes.len()]; + writer.exec(insert).await.expect("INSERT INTO ARRAY"); + writer + .exec(&format!("SELECT ARRAY_FLUSH('{name}')")) + .await + .expect("ARRAY_FLUSH"); + cluster + .wait_for_full_apply_convergence(std::time::Duration::from_secs(10)) + .await; + cluster +} + +/// The genome fixture of `cluster_array.rs`: qual 1/2/3 on chr 0, +/// 10/20/30 on chr 1, 100/200/300 on chr 2, each at pos 10/20/30. +async fn spawn_genome() -> TestCluster { + spawn_with_array( + "CREATE ARRAY genome \ + DIMS (chr INT64 [0..9], pos INT64 [0..99]) \ + ATTRS (qual FLOAT64) \ + TILE_EXTENTS (1, 100) \ + CELL_ORDER HILBERT", + "INSERT INTO ARRAY genome \ + COORDS (0, 10) VALUES (1.0), COORDS (0, 20) VALUES (2.0), \ + COORDS (0, 30) VALUES (3.0), COORDS (1, 10) VALUES (10.0), \ + COORDS (1, 20) VALUES (20.0), COORDS (1, 30) VALUES (30.0), \ + COORDS (2, 10) VALUES (100.0), COORDS (2, 20) VALUES (200.0), \ + COORDS (2, 30) VALUES (300.0)", + "genome", + ) + .await +} + +fn number(text: &str) -> f64 { + text.parse() + .unwrap_or_else(|e| panic!("{text:?} is not a number: {e}")) +} + +/// The scalar `result` every node answers for `sql`. +async fn scalar_on_every_node(cluster: &TestCluster, sql: &str) -> Vec { + let mut out = Vec::with_capacity(cluster.nodes.len()); + for (idx, node) in cluster.nodes.iter().enumerate() { + let rows = node + .query_named_rows(sql) + .await + .unwrap_or_else(|e| panic!("node {idx}: {sql}: {e}")); + assert_eq!(rows.len(), 1, "node {idx}: one scalar row: {rows:?}"); + out.push(number(&rows[0]["result"])); + } + out +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_agg_box_sums_only_its_rows() { + let cluster = spawn_genome().await; + let sums = scalar_on_every_node( + &cluster, + "SELECT * FROM ARRAY_AGG('genome', 'qual', 'sum', NULL, '{chr: [1, 2]}')", + ) + .await; + for (idx, sum) in sums.iter().enumerate() { + assert!((sum - 660.0).abs() < 1e-9, "node {idx}: sum {sum}"); + } + cluster.shutdown().await; +} + +/// A zone drawn over `pos` then `chr`: pos in `[15, 35)`, chr in `[0, 1.5)`. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_agg_zone_sums_the_cells_inside_it() { + let cluster = spawn_genome().await; + let sums = scalar_on_every_node( + &cluster, + "SELECT * FROM ARRAY_AGG('genome', 'qual', 'sum', NULL, NULL, \ + 'POLYGON((15 0, 35 0, 35 1.5, 15 1.5, 15 0))', ARRAY['pos', 'chr'])", + ) + .await; + for (idx, sum) in sums.iter().enumerate() { + assert!((sum - 55.0).abs() < 1e-9, "node {idx}: sum {sum}"); + } + cluster.shutdown().await; +} + +/// A group-by on a FLOAT64 dim reports each group under its dim value on +/// every node, across shards. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cluster_array_agg_groups_by_a_float_dim_value() { + let cluster = spawn_with_array( + "CREATE ARRAY scene \ + DIMS (lon FLOAT64 [100..106], lat FLOAT64 [0..8]) \ + ATTRS (ndvi FLOAT64) \ + TILE_EXTENTS (1, 1) \ + CELL_ORDER HILBERT", + "INSERT INTO ARRAY scene \ + COORDS (101.5, 0.5) VALUES (1.0), COORDS (101.5, 5.5) VALUES (2.0), \ + COORDS (103.25, 2.5) VALUES (4.0), COORDS (103.25, 7.5) VALUES (8.0), \ + COORDS (105.5, 1.5) VALUES (16.0)", + "scene", + ) + .await; + let sql = "SELECT * FROM ARRAY_AGG('scene', 'ndvi', 'sum', 'lon', '{lon: [100, 104]}')"; + for (idx, node) in cluster.nodes.iter().enumerate() { + let rows = node + .query_named_rows(sql) + .await + .unwrap_or_else(|e| panic!("node {idx}: {sql}: {e}")); + let mut groups: Vec<(f64, f64)> = rows + .iter() + .map(|row| (number(&row["group"]), number(&row["result"]))) + .collect(); + groups.sort_by(|a, b| a.0.total_cmp(&b.0)); + assert_eq!( + groups, + vec![(101.5, 3.0), (103.25, 12.0)], + "node {idx}: one group per lon value" + ); + } + cluster.shutdown().await; +} diff --git a/nodedb-cluster-tests/tests/common_suite/cases/mod.rs b/nodedb-cluster-tests/tests/common_suite/cases/mod.rs index bfb83a2c5..eaf7c5454 100644 --- a/nodedb-cluster-tests/tests/common_suite/cases/mod.rs +++ b/nodedb-cluster-tests/tests/common_suite/cases/mod.rs @@ -51,6 +51,8 @@ mod clone_materialize_cross_node; mod cluster_array; mod cluster_array_body_read; mod cluster_array_cell_raft_replication; +mod cluster_array_project_elementwise; +mod cluster_array_region_agg; mod cluster_array_ws_rpc; mod cluster_backup_remote_cut; mod cluster_backup_restore; diff --git a/nodedb-cluster/src/calvin/completion.rs b/nodedb-cluster/src/calvin/completion.rs index af3b6f708..90ca42062 100644 --- a/nodedb-cluster/src/calvin/completion.rs +++ b/nodedb-cluster/src/calvin/completion.rs @@ -44,29 +44,63 @@ impl TxnId { pub enum AttemptOutcome { /// The global verdict was COMMIT. Completed, - /// The global verdict was ABORT. `reason` names the cause. - Aborted { reason: AbortReason }, + /// The global verdict was ABORT, with its cause. + Aborted(AbortCause), +} + +/// The cause of an abort vote or verdict. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AbortCause { + /// The class of the abort. + pub reason: AbortReason, + /// The typed refusal a participant's stage raised, encoded by the host + /// crate and opaque to the sequencer. Set only with + /// [`AbortReason::ParticipantError`]. + /// + /// Every replica stages the same slice and raises the same refusal. A + /// participant that never staged (an expired request, a drained core) + /// carries none. + pub refusal: Option>, +} + +impl AbortCause { + /// A participant error that carries the typed refusal `refusal`. + pub fn refused(refusal: Vec) -> Self { + Self { + reason: AbortReason::ParticipantError, + refusal: Some(refusal), + } + } +} + +impl From for AbortCause { + fn from(reason: AbortReason) -> Self { + Self { + reason, + refusal: None, + } + } } /// One participant vShard's durable commit vote. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum ParticipantVote { Commit, - Abort(AbortReason), + Abort(AbortCause), } /// A commit/abort decision: the tally aggregated from votes, and the /// authoritative verdict stored on a `PendingCompletion`. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum VerdictOutcome { Commit, - Abort(AbortReason), + Abort(AbortCause), } impl VerdictOutcome { /// `true` for a commit decision. The scheduler's flush/drop gate needs only - /// this bit; the reason travels to the coordinator instead. - pub fn is_commit(self) -> bool { + /// this bit. The cause travels to the coordinator instead. + pub fn is_commit(&self) -> bool { matches!(self, Self::Commit) } } @@ -120,7 +154,7 @@ impl Inner { pub(crate) fn settle(&mut self, txn: TxnId) { if let Some(entry) = self.completions.get_mut(&txn) && entry.is_terminal() - && let Some(verdict) = entry.verdict + && let Some(verdict) = entry.verdict.clone() && let Some(waiter) = entry.completion_tx.take() { let results = entry.take_ack_results(); @@ -141,7 +175,7 @@ impl Inner { fn outcome_for(verdict: VerdictOutcome) -> AttemptOutcome { match verdict { VerdictOutcome::Commit => AttemptOutcome::Completed, - VerdictOutcome::Abort(reason) => AttemptOutcome::Aborted { reason }, + VerdictOutcome::Abort(cause) => AttemptOutcome::Aborted(cause), } } @@ -538,7 +572,7 @@ mod tests { reg.seed_expected(txn, 2); reg.note_verdict( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict), + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()), ); let rx = reg.register_completion(txn, 2); reg.note_completion_ack(txn, 10); @@ -546,9 +580,7 @@ mod tests { let outcome = rx.await.expect("completion fires"); assert_eq!( outcome, - AttemptOutcome::Aborted { - reason: AbortReason::SerializationConflict - }, + AttemptOutcome::Aborted(AbortReason::SerializationConflict.into()), "a stored ABORT verdict must surface as Aborted, not Completed" ); assert_eq!(reg.pending_completions_len(), 0); @@ -562,15 +594,13 @@ mod tests { let reg = CalvinCompletionRegistry::new_detached(); let txn = TxnId::new(32, 1); reg.seed_expected(txn, 2); - reg.note_verdict(txn, VerdictOutcome::Abort(AbortReason::PlanRejected)); + reg.note_verdict(txn, VerdictOutcome::Abort(AbortReason::PlanRejected.into())); reg.note_completion_ack(txn, 10); reg.note_completion_ack(txn, 20); let mut rx = reg.register_completion(txn, 2); assert_eq!( rx.try_recv(), - Ok(AttemptOutcome::Aborted { - reason: AbortReason::PlanRejected - }), + Ok(AttemptOutcome::Aborted(AbortReason::PlanRejected.into())), "registration on a terminal entry fires at once" ); assert_eq!( @@ -683,22 +713,24 @@ mod tests { let txn = TxnId::new(40, 2); reg.seed_expected(txn, 1); let rx = reg.register_completion(txn, 1); - reg.note_vote(txn, 1, ParticipantVote::Abort(AbortReason::PlanRejected)); - reg.note_verdict(txn, VerdictOutcome::Abort(AbortReason::PlanRejected)); + reg.note_vote( + txn, + 1, + ParticipantVote::Abort(AbortReason::PlanRejected.into()), + ); + reg.note_verdict(txn, VerdictOutcome::Abort(AbortReason::PlanRejected.into())); reg.note_completion_ack(txn, 1); assert_eq!( rx.await.expect("completion fires"), - AttemptOutcome::Aborted { - reason: AbortReason::PlanRejected - } + AttemptOutcome::Aborted(AbortReason::PlanRejected.into()) ); let progress = reg.participant_progress(txn, 1).expect("entry stays"); assert!(progress.voted && progress.acked && progress.has_verdict); assert_eq!(reg.verdict(txn), Some(false)); assert_eq!( - reg.vote_tally(txn).and_then(|tally| tally.get(&1).copied()), - Some(ParticipantVote::Abort(AbortReason::PlanRejected)) + reg.vote_tally(txn).and_then(|tally| tally.get(&1).cloned()), + Some(ParticipantVote::Abort(AbortReason::PlanRejected.into())) ); reg.set_waiterless_ttl(Duration::ZERO); diff --git a/nodedb-cluster/src/calvin/completion_parts.rs b/nodedb-cluster/src/calvin/completion_parts.rs index 06ba9390e..32bf4265b 100644 --- a/nodedb-cluster/src/calvin/completion_parts.rs +++ b/nodedb-cluster/src/calvin/completion_parts.rs @@ -21,17 +21,16 @@ impl CalvinCompletionRegistry { /// Idempotent. A later ack of `txn` lands on the terminal entry, which /// the waiterless sweep removes. pub fn note_parts_abandoned(&self, txn: TxnId) { - let verdict = VerdictOutcome::Abort(AbortReason::PartsLost); let mut inner = self.inner.lock().unwrap_or_else(|p| p.into_inner()); { let entry = inner.entry_mut(txn); entry.abandoned = true; - entry.verdict = Some(verdict); + entry.verdict = Some(VerdictOutcome::Abort(AbortReason::PartsLost.into())); } let signal = VerdictSignal { epoch: txn.epoch, position: txn.position, - verdict, + commit: false, }; for tx in inner.verdict_signal_senders.values() { let _ = tx.try_send(signal); @@ -55,9 +54,7 @@ mod tests { reg.note_parts_abandoned(txn); assert_eq!( rx.await.expect("outcome"), - AttemptOutcome::Aborted { - reason: AbortReason::PartsLost - } + AttemptOutcome::Aborted(AbortReason::PartsLost.into()) ); assert_eq!(reg.verdict(txn), Some(false)); } @@ -70,9 +67,7 @@ mod tests { let rx = reg.register_completion(txn, 3); assert_eq!( rx.await.expect("outcome"), - AttemptOutcome::Aborted { - reason: AbortReason::PartsLost - } + AttemptOutcome::Aborted(AbortReason::PartsLost.into()) ); } } diff --git a/nodedb-cluster/src/calvin/completion_verdict.rs b/nodedb-cluster/src/calvin/completion_verdict.rs index 8edb5ad14..456836334 100644 --- a/nodedb-cluster/src/calvin/completion_verdict.rs +++ b/nodedb-cluster/src/calvin/completion_verdict.rs @@ -11,7 +11,7 @@ use std::collections::BTreeMap; use tokio::sync::mpsc; use super::TxnId; -use super::completion::{CalvinCompletionRegistry, ParticipantVote, VerdictOutcome}; +use super::completion::{AbortCause, CalvinCompletionRegistry, ParticipantVote, VerdictOutcome}; use super::sequencer::AbortReason; /// Push notification that a staged cross-shard txn's authoritative global @@ -28,8 +28,10 @@ use super::sequencer::AbortReason; pub struct VerdictSignal { pub epoch: u64, pub position: u32, - /// The decision, carrying the abort reason when it aborts. - pub verdict: VerdictOutcome, + /// `true` for a commit decision. A participant needs only this bit to + /// flush or drop. The abort cause reaches the coordinator through + /// `AttemptOutcome`. + pub commit: bool, } impl CalvinCompletionRegistry { @@ -155,22 +157,27 @@ impl CalvinCompletionRegistry { /// channel) is a non-fatal drop — the scheduler's stall re-probe backstops /// it — never a block or panic, mirroring the apply fan-out drop discipline. pub fn note_verdict(&self, txn: TxnId, verdict: VerdictOutcome) { + let commit = verdict.is_commit(); let mut inner = self.inner.lock().unwrap_or_else(|p| p.into_inner()); { let entry = inner.entry_mut(txn); - match entry.verdict { - Some(existing) if existing.is_commit() != verdict.is_commit() => { + let differs = entry + .verdict + .as_ref() + .map(|existing| existing.is_commit() != commit); + match differs { + Some(true) => { tracing::warn!( epoch = txn.epoch, position = txn.position, - existing = ?existing, + existing = ?entry.verdict, proposed = ?verdict, "calvin verdict differs from a previously applied one; \ determinism bug — overwriting with the latest" ); entry.verdict = Some(verdict); } - Some(_) => {} + Some(false) => {} None => entry.verdict = Some(verdict), } } @@ -181,7 +188,7 @@ impl CalvinCompletionRegistry { let signal = VerdictSignal { epoch: txn.epoch, position: txn.position, - verdict, + commit, }; for tx in inner.verdict_signal_senders.values() { let _ = tx.try_send(signal); @@ -192,7 +199,7 @@ impl CalvinCompletionRegistry { /// The parked scheduler's flush/drop gate: `Some(true)` commit, `Some(false)` /// abort, `None` if no verdict has been applied yet. Deliberately drops the - /// abort reason — the reason reaches the coordinator via `AttemptOutcome`, + /// abort cause — the cause reaches the coordinator via `AttemptOutcome`, /// and a participant only needs to know whether to flush. pub fn verdict(&self, txn_id: TxnId) -> Option { self.inner @@ -200,7 +207,7 @@ impl CalvinCompletionRegistry { .unwrap_or_else(|p| p.into_inner()) .completions .get(&txn_id) - .and_then(|entry| entry.verdict) + .and_then(|entry| entry.verdict.as_ref()) .map(VerdictOutcome::is_commit) } @@ -217,37 +224,45 @@ impl CalvinCompletionRegistry { } /// Aggregate a complete vote tally: commit only when every participant voted -/// commit, otherwise abort with the highest-precedence reason. +/// commit, otherwise abort with the highest-precedence cause. +/// +/// Precedence, highest first: +/// +/// - `PartsLost`: no participant holds the whole transaction. +/// - `PlanRejected`: the plans fail the same way on every resubmit. +/// - `ParticipantError` with a refusal: the stage refused the slice, and a +/// resubmit refuses it the same way. The client sees that refusal. +/// - `ParticipantError` without a refusal: a peer that never staged makes +/// every other verdict unverifiable. +/// - `CollectionSuperseded`: a retry against the same read-set is pointless. +/// - `PredictionDrift`: the coordinator retries with a fresh reconnaissance. +/// - `SerializationConflict`. /// -/// Precedence: `PartsLost` outranks everything, since no participant holds -/// the whole transaction. `PlanRejected` outranks `ParticipantError`, which -/// outranks `CollectionSuperseded`, which outranks `PredictionDrift`, which -/// outranks `SerializationConflict`. Rejected plans fail the same way on every -/// resubmit, a peer that never staged makes every other verdict unverifiable, -/// a superseded collection makes a retry against the same read-set pointless, -/// and a drifted prediction retries with a fresh reconnaissance. The tally is -/// order-independent. +/// Among causes of equal rank the lowest vShard wins. The votes iterate in +/// vShard order, so every replica picks the same cause, whatever the order +/// the votes applied in. fn tally_verdict(votes: &BTreeMap) -> VerdictOutcome { - fn rank(reason: AbortReason) -> u8 { - match reason { - AbortReason::PartsLost => 5, - AbortReason::PlanRejected => 4, + fn rank(cause: &AbortCause) -> u8 { + match cause.reason { + AbortReason::PartsLost => 6, + AbortReason::PlanRejected => 5, + AbortReason::ParticipantError if cause.refusal.is_some() => 4, AbortReason::ParticipantError => 3, AbortReason::CollectionSuperseded => 2, AbortReason::PredictionDrift => 1, AbortReason::SerializationConflict => 0, } } - let mut reason: Option = None; + let mut held: Option<&AbortCause> = None; for vote in votes.values() { - if let ParticipantVote::Abort(vote_reason) = *vote - && reason.is_none_or(|held| rank(vote_reason) > rank(held)) + if let ParticipantVote::Abort(cause) = vote + && held.is_none_or(|held| rank(cause) > rank(held)) { - reason = Some(vote_reason); + held = Some(cause); } } - match reason { - Some(reason) => VerdictOutcome::Abort(reason), + match held { + Some(cause) => VerdictOutcome::Abort(cause.clone()), None => VerdictOutcome::Commit, } } @@ -264,14 +279,16 @@ mod tests { reg.note_vote( txn, 2, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); let tally = reg.vote_tally(txn).expect("entry created by note_vote"); assert_eq!(tally.len(), 2); assert_eq!(tally.get(&1), Some(&ParticipantVote::Commit)); assert_eq!( tally.get(&2), - Some(&ParticipantVote::Abort(AbortReason::SerializationConflict)) + Some(&ParticipantVote::Abort( + AbortReason::SerializationConflict.into() + )) ); } @@ -313,14 +330,14 @@ mod tests { reg.note_vote( txn, 2, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); // Any abort vote makes the aggregated verdict an abort. assert_eq!( rx.try_recv().expect("verdict emitted"), ( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict) + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()) ) ); assert!(rx.try_recv().is_err()); @@ -337,16 +354,19 @@ mod tests { reg.note_vote( txn, 1, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); reg.note_vote( txn, 2, - ParticipantVote::Abort(AbortReason::ParticipantError), + ParticipantVote::Abort(AbortReason::ParticipantError.into()), ); assert_eq!( rx.try_recv().expect("verdict emitted"), - (txn, VerdictOutcome::Abort(AbortReason::ParticipantError)), + ( + txn, + VerdictOutcome::Abort(AbortReason::ParticipantError.into()) + ), "a peer that never staged makes the stale read-set unverifiable" ); } @@ -362,33 +382,40 @@ mod tests { reg.note_vote( txn, 1, - ParticipantVote::Abort(AbortReason::ParticipantError), + ParticipantVote::Abort(AbortReason::ParticipantError.into()), ); reg.note_vote( txn, 2, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); assert_eq!( rx.try_recv().expect("verdict emitted"), - (txn, VerdictOutcome::Abort(AbortReason::ParticipantError)) + ( + txn, + VerdictOutcome::Abort(AbortReason::ParticipantError.into()) + ) ); } + /// Tally `votes`, the vote of vShard `1`, `2`, ... in turn. + fn tally(votes: [&ParticipantVote; N]) -> VerdictOutcome { + let votes: BTreeMap = (1..).zip(votes.into_iter().cloned()).collect(); + tally_verdict(&votes) + } + + fn aborted(reason: AbortReason) -> VerdictOutcome { + VerdictOutcome::Abort(reason.into()) + } + /// A superseded collection aborts a transaction other slices voted to /// commit, outranks a conflict, and yields to a participant error, in any /// vote order. #[test] fn a_superseded_collection_aborts_the_whole_transaction() { - let superseded = ParticipantVote::Abort(AbortReason::CollectionSuperseded); - let conflict = ParticipantVote::Abort(AbortReason::SerializationConflict); - let errored = ParticipantVote::Abort(AbortReason::ParticipantError); - let tally = |votes: [&ParticipantVote; 2]| { - let votes: BTreeMap = - (1..).zip(votes.into_iter().copied()).collect(); - tally_verdict(&votes) - }; - let aborted = VerdictOutcome::Abort; + let superseded = ParticipantVote::Abort(AbortReason::CollectionSuperseded.into()); + let conflict = ParticipantVote::Abort(AbortReason::SerializationConflict.into()); + let errored = ParticipantVote::Abort(AbortReason::ParticipantError.into()); assert_eq!( tally([&ParticipantVote::Commit, &superseded]), @@ -405,30 +432,80 @@ mod tests { /// A rejected plan outranks every staged abort cause, in any vote order. #[test] fn plan_rejected_outranks_staged_abort_causes() { - let rejected = ParticipantVote::Abort(AbortReason::PlanRejected); - let tally = |votes: [&ParticipantVote; 2]| { - let votes: BTreeMap = - (1..).zip(votes.into_iter().copied()).collect(); - tally_verdict(&votes) - }; + let rejected = ParticipantVote::Abort(AbortReason::PlanRejected.into()); for other in [ - AbortReason::SerializationConflict, - AbortReason::PredictionDrift, - AbortReason::CollectionSuperseded, - AbortReason::ParticipantError, + AbortCause::from(AbortReason::SerializationConflict), + AbortCause::from(AbortReason::PredictionDrift), + AbortCause::from(AbortReason::CollectionSuperseded), + AbortCause::from(AbortReason::ParticipantError), + AbortCause::refused(vec![1, 2, 3]), ] { let other = ParticipantVote::Abort(other); for votes in [[&rejected, &other], [&other, &rejected]] { assert_eq!( tally(votes), - VerdictOutcome::Abort(AbortReason::PlanRejected), + aborted(AbortReason::PlanRejected), "{other:?}" ); } } - let lost = ParticipantVote::Abort(AbortReason::PartsLost); + let lost = ParticipantVote::Abort(AbortReason::PartsLost.into()); for votes in [[&rejected, &lost], [&lost, &rejected]] { - assert_eq!(tally(votes), VerdictOutcome::Abort(AbortReason::PartsLost)); + assert_eq!(tally(votes), aborted(AbortReason::PartsLost)); + } + } + + /// A refusal outranks a participant that never staged, a superseded + /// collection, a drifted prediction and a conflict, in any vote order. + #[test] + fn a_refusal_outranks_every_unrefused_participant_cause() { + let refused = ParticipantVote::Abort(AbortCause::refused(vec![7])); + for other in [ + AbortReason::ParticipantError, + AbortReason::CollectionSuperseded, + AbortReason::PredictionDrift, + AbortReason::SerializationConflict, + ] { + let other = ParticipantVote::Abort(other.into()); + for votes in [[&refused, &other], [&other, &refused]] { + assert_eq!( + tally(votes), + VerdictOutcome::Abort(AbortCause::refused(vec![7])), + "{other:?}" + ); + } + } + } + + /// Two participants refuse with different refusals. Every replica picks + /// the lowest vShard's refusal, whatever order the votes applied in. + #[tokio::test] + async fn two_refusals_choose_the_lowest_vshard_on_every_replica() { + let txn = TxnId::new(32, 0); + let low = ParticipantVote::Abort(AbortCause::refused(b"vshard 3".to_vec())); + let high = ParticipantVote::Abort(AbortCause::refused(b"vshard 9".to_vec())); + let orders: [[(u32, &ParticipantVote); 3]; 2] = [ + [(3, &low), (5, &ParticipantVote::Commit), (9, &high)], + [(9, &high), (5, &ParticipantVote::Commit), (3, &low)], + ]; + for order in orders { + let (tx, mut rx) = mpsc::channel(8); + let replica = CalvinCompletionRegistry::new(tx); + replica.seed_expected(txn, 3); + for (vshard, vote) in order { + replica.note_vote(txn, vshard, vote.clone()); + } + let expected = VerdictOutcome::Abort(AbortCause::refused(b"vshard 3".to_vec())); + assert_eq!( + rx.try_recv().expect("verdict emitted"), + (txn, expected.clone()), + "{order:?}" + ); + assert_eq!( + replica.drain_unproposed_verdicts(), + vec![(txn, expected)], + "{order:?}" + ); } } @@ -515,7 +592,7 @@ mod tests { VerdictSignal { epoch: 50, position: 2, - verdict: VerdictOutcome::Commit, + commit: true, } ); } @@ -531,23 +608,13 @@ mod tests { let txn = TxnId::new(51, 0); reg.note_verdict( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict), + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()), ); // Both locally registered vShard schedulers receive the broadcast; each // filters by its own parked (epoch, position). - assert!( - !rx1.try_recv() - .expect("push to vshard 1") - .verdict - .is_commit() - ); - assert!( - !rx2.try_recv() - .expect("push to vshard 2") - .verdict - .is_commit() - ); + assert!(!rx1.try_recv().expect("push to vshard 1").commit); + assert!(!rx2.try_recv().expect("push to vshard 2").commit); } #[tokio::test] @@ -581,7 +648,7 @@ mod tests { // Must not panic despite the closed channel; the verdict still stores. reg.note_verdict( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict), + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()), ); assert_eq!(reg.verdict(txn), Some(false)); } @@ -628,14 +695,14 @@ mod tests { reg.note_vote( txn, 2, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); assert_eq!(reg.verdict(txn), None); assert_eq!( reg.drain_unproposed_verdicts(), vec![( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict) + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()) )], "any abort vote makes the re-proposed verdict an abort" ); diff --git a/nodedb-cluster/src/calvin/mod.rs b/nodedb-cluster/src/calvin/mod.rs index c66504177..3f2e007dc 100644 --- a/nodedb-cluster/src/calvin/mod.rs +++ b/nodedb-cluster/src/calvin/mod.rs @@ -12,8 +12,8 @@ pub mod types; pub use applied_acks::{AppliedAckLog, AppliedCompletionAck}; pub use completion::{ - Assignment, AssignmentReceiver, AttemptOutcome, CalvinCompletionRegistry, ParticipantProgress, - ParticipantVote, TxnId, VerdictOutcome, + AbortCause, Assignment, AssignmentReceiver, AttemptOutcome, CalvinCompletionRegistry, + ParticipantProgress, ParticipantVote, TxnId, VerdictOutcome, }; pub use completion_verdict::VerdictSignal; pub use completion_waiter::CompletionReport; diff --git a/nodedb-cluster/src/calvin/sequencer/entry.rs b/nodedb-cluster/src/calvin/sequencer/entry.rs index 9774f2232..c73e36c1d 100644 --- a/nodedb-cluster/src/calvin/sequencer/entry.rs +++ b/nodedb-cluster/src/calvin/sequencer/entry.rs @@ -116,18 +116,22 @@ pub enum SequencerEntry { reason: ReleaseReason, }, /// One participant vShard's durable ABORT vote, carrying why it aborted. + /// `refusal` is the typed refusal its stage raised, encoded by the host + /// crate. It is set only with `ParticipantError`. AbortVote { epoch: u64, position: u32, vshard: u32, reason: AbortReason, + refusal: Option>, }, /// The global ABORT verdict for a staged cross-shard txn, carrying the - /// winning participant reason. Participants drop their staged buffer. + /// winning participant cause. Participants drop their staged buffer. AbortVerdict { epoch: u64, position: u32, reason: AbortReason, + refusal: Option>, }, /// A backup's consistent-cut marker, carrying the backup's watermark /// `hlc`. Every replica fans it out to each of its vShard schedulers in @@ -304,6 +308,7 @@ mod tests { position: 2, vshard: 9, reason: AbortReason::ParticipantError, + refusal: Some(vec![0x23, 0x00, 0xff]), }; let bytes = zerompk::to_msgpack_vec(&entry).expect("encode"); let decoded: SequencerEntry = zerompk::from_msgpack(&bytes).expect("decode"); @@ -312,14 +317,17 @@ mod tests { #[test] fn abort_verdict_msgpack_roundtrip() { - let entry = SequencerEntry::AbortVerdict { - epoch: 5, - position: 2, - reason: AbortReason::SerializationConflict, - }; - let bytes = zerompk::to_msgpack_vec(&entry).expect("encode"); - let decoded: SequencerEntry = zerompk::from_msgpack(&bytes).expect("decode"); - assert_eq!(entry, decoded); + for refusal in [None, Some(vec![0x23, 0x50, 0x02])] { + let entry = SequencerEntry::AbortVerdict { + epoch: 5, + position: 2, + reason: AbortReason::SerializationConflict, + refusal, + }; + let bytes = zerompk::to_msgpack_vec(&entry).expect("encode"); + let decoded: SequencerEntry = zerompk::from_msgpack(&bytes).expect("decode"); + assert_eq!(entry, decoded); + } } #[test] diff --git a/nodedb-cluster/src/calvin/sequencer/service/verdict_entry.rs b/nodedb-cluster/src/calvin/sequencer/service/verdict_entry.rs index ebfe7e006..cd8e75652 100644 --- a/nodedb-cluster/src/calvin/sequencer/service/verdict_entry.rs +++ b/nodedb-cluster/src/calvin/sequencer/service/verdict_entry.rs @@ -4,21 +4,22 @@ //! the leader proposes for it. use crate::calvin::sequencer::entry::SequencerEntry; -use crate::calvin::{TxnId, VerdictOutcome}; +use crate::calvin::{AbortCause, TxnId, VerdictOutcome}; /// Encode a decision as the entry the leader proposes. A commit is `Verdict`. -/// An abort is `AbortVerdict`, which carries the reason the coordinator -/// reports. +/// An abort is `AbortVerdict`, which carries the reason and the refusal the +/// coordinator reports. pub(crate) fn verdict_entry(txn: TxnId, outcome: VerdictOutcome) -> SequencerEntry { match outcome { VerdictOutcome::Commit => SequencerEntry::Verdict { epoch: txn.epoch, position: txn.position, }, - VerdictOutcome::Abort(reason) => SequencerEntry::AbortVerdict { + VerdictOutcome::Abort(AbortCause { reason, refusal }) => SequencerEntry::AbortVerdict { epoch: txn.epoch, position: txn.position, reason, + refusal, }, } } @@ -32,7 +33,7 @@ mod tests { fn abort_proposes_abort_verdict_with_its_reason() { let entry = verdict_entry( TxnId::new(4, 1), - VerdictOutcome::Abort(AbortReason::ParticipantError), + VerdictOutcome::Abort(AbortReason::ParticipantError.into()), ); assert_eq!( entry, @@ -40,6 +41,24 @@ mod tests { epoch: 4, position: 1, reason: AbortReason::ParticipantError, + refusal: None, + } + ); + } + + #[test] + fn abort_proposes_abort_verdict_with_its_refusal() { + let entry = verdict_entry( + TxnId::new(4, 3), + VerdictOutcome::Abort(AbortCause::refused(vec![5, 6])), + ); + assert_eq!( + entry, + SequencerEntry::AbortVerdict { + epoch: 4, + position: 3, + reason: AbortReason::ParticipantError, + refusal: Some(vec![5, 6]), } ); } diff --git a/nodedb-cluster/src/calvin/sequencer/state_machine/apply.rs b/nodedb-cluster/src/calvin/sequencer/state_machine/apply.rs index 0015a141c..5125273ba 100644 --- a/nodedb-cluster/src/calvin/sequencer/state_machine/apply.rs +++ b/nodedb-cluster/src/calvin/sequencer/state_machine/apply.rs @@ -17,7 +17,7 @@ use tracing::{error, warn}; use crate::calvin::sequencer::entry::SequencerEntry; use crate::calvin::sequencer::epoch_guard::{EpochCheck, SequencerHalt, classify}; use crate::calvin::types::{EpochBatch, SchedulerInput, TxnIdWire}; -use crate::calvin::{ParticipantVote, TxnId, VerdictOutcome}; +use crate::calvin::{AbortCause, ParticipantVote, TxnId, VerdictOutcome}; use super::core::{Delivery, SequencerStateMachine}; @@ -117,10 +117,11 @@ impl SequencerStateMachine { position, vshard, reason, + refusal, } => self.completion_registry.note_vote( TxnId::new(epoch, position), vshard, - ParticipantVote::Abort(reason), + ParticipantVote::Abort(AbortCause { reason, refusal }), ), // Authoritative verdict for a staged cross-shard txn, proposed by // the leader once every participant voted. Applied on ALL replicas @@ -134,9 +135,11 @@ impl SequencerStateMachine { epoch, position, reason, - } => self - .completion_registry - .note_verdict(TxnId::new(epoch, position), VerdictOutcome::Abort(reason)), + refusal, + } => self.completion_registry.note_verdict( + TxnId::new(epoch, position), + VerdictOutcome::Abort(AbortCause { reason, refusal }), + ), // The owning vShard's scheduler installs the SHARED lock. SequencerEntry::ReserveRead { owner, vshard, key } => self.forward_reservation( index, @@ -954,19 +957,19 @@ mod tests { epoch: 9, position: 5, reason: crate::calvin::AbortReason::ParticipantError, + refusal: Some(vec![4, 2]), }); sm.apply(1, &data); // The participant gate still reads a plain abort... assert_eq!(registry.verdict(txn), Some(false)); - // ...and the coordinator gets the cause, not a blanket conflict. + // ...and the coordinator gets the cause and its refusal, not a + // blanket conflict. let rx = registry.register_completion(txn, 1); registry.note_completion_ack(txn, 3); assert_eq!( rx.await.expect("completion fires"), - crate::calvin::AttemptOutcome::Aborted { - reason: crate::calvin::AbortReason::ParticipantError - } + crate::calvin::AttemptOutcome::Aborted(crate::calvin::AbortCause::refused(vec![4, 2])) ); assert_eq!(sm.last_applied_epoch(), None); } @@ -983,13 +986,14 @@ mod tests { position: 6, vshard: 3, reason: crate::calvin::AbortReason::ParticipantError, + refusal: Some(vec![9]), }); sm.apply(1, &data); assert_eq!( - registry.vote_tally(txn).and_then(|t| t.get(&3).copied()), + registry.vote_tally(txn).and_then(|t| t.get(&3).cloned()), Some(crate::calvin::ParticipantVote::Abort( - crate::calvin::AbortReason::ParticipantError + crate::calvin::AbortCause::refused(vec![9]) )) ); } @@ -1008,7 +1012,7 @@ mod tests { sm.apply(1, &data); assert_eq!( - registry.vote_tally(txn).and_then(|t| t.get(&3).copied()), + registry.vote_tally(txn).and_then(|t| t.get(&3).cloned()), Some(crate::calvin::ParticipantVote::Commit) ); } diff --git a/nodedb-cluster/src/distributed_array/agg_number.rs b/nodedb-cluster/src/distributed_array/agg_number.rs new file mode 100644 index 000000000..142009218 --- /dev/null +++ b/nodedb-cluster/src/distributed_array/agg_number.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The running number a distributed array aggregate partial carries. +//! +//! An integer attr's partial carries an exact `i128`, so a cluster `SUM` +//! equals the single-node `SUM` to the unit. A float attr's partial carries +//! an `f64`. `MIN` / `MAX` sort NaN above every other float, as PostgreSQL +//! does. + +use std::cmp::Ordering; + +use serde::{Deserialize, Serialize}; + +/// A partial's sum, minimum, or maximum. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub enum ArrayAggNumber { + /// An exact integer: the high and low 64 bits of an `i128`. + Int { + hi: i64, + lo: u64, + }, + Float(f64), +} + +impl ArrayAggNumber { + pub fn int(v: i128) -> Self { + ArrayAggNumber::Int { + hi: (v >> 64) as i64, + lo: v as u64, + } + } + + /// The exact integer, or `None` for a float. + pub fn as_i128(self) -> Option { + match self { + ArrayAggNumber::Int { hi, lo } => Some(join(hi, lo)), + ArrayAggNumber::Float(_) => None, + } + } + + pub fn as_f64(self) -> f64 { + match self { + ArrayAggNumber::Int { hi, lo } => join(hi, lo) as f64, + ArrayAggNumber::Float(v) => v, + } + } + + /// The sum of two numbers. An integer sum saturates at the `i128` + /// bounds, far past the `i64` range a final `SUM` must fit. + pub fn sum_with(self, other: Self) -> Self { + match (self.as_i128(), other.as_i128()) { + (Some(a), Some(b)) => Self::int(a.saturating_add(b)), + _ => ArrayAggNumber::Float(self.as_f64() + other.as_f64()), + } + } + + /// The smaller number. NaN loses to every other value. + pub fn min_with(self, other: Self) -> Self { + match cmp_nan_high(self, other) { + Ordering::Greater => other, + Ordering::Less | Ordering::Equal => self, + } + } + + /// The larger number. NaN wins over every other value. + pub fn max_with(self, other: Self) -> Self { + match cmp_nan_high(self, other) { + Ordering::Less => other, + Ordering::Greater | Ordering::Equal => self, + } + } +} + +fn join(hi: i64, lo: u64) -> i128 { + (i128::from(hi) << 64) | i128::from(lo) +} + +fn cmp_nan_high(a: ArrayAggNumber, b: ArrayAggNumber) -> Ordering { + if let (Some(x), Some(y)) = (a.as_i128(), b.as_i128()) { + return x.cmp(&y); + } + let (x, y) = (a.as_f64(), b.as_f64()); + match (x.is_nan(), y.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => x.partial_cmp(&y).unwrap_or(Ordering::Equal), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_split_i128_round_trips() { + for v in [0i128, -1, 1 << 70, -(1 << 70), i128::from(i64::MIN) - 5] { + assert_eq!(ArrayAggNumber::int(v).as_i128(), Some(v)); + } + } + + #[test] + fn integer_sums_stay_exact_past_2_pow_53() { + let a = ArrayAggNumber::int(1 << 53); + let b = ArrayAggNumber::int(1); + assert_eq!(a.sum_with(b).as_i128(), Some((1 << 53) + 1)); + } + + #[test] + fn nan_sorts_high() { + let nan = ArrayAggNumber::Float(f64::NAN); + let one = ArrayAggNumber::Float(1.0); + assert!(one.max_with(nan).as_f64().is_nan()); + assert_eq!(nan.min_with(one), one); + } + + #[test] + fn the_number_survives_msgpack() { + let n = ArrayAggNumber::int(-(1 << 80)); + let bytes = zerompk::to_msgpack_vec(&n).expect("encode"); + let back: ArrayAggNumber = zerompk::from_msgpack(&bytes).expect("decode"); + assert_eq!(back, n); + } +} diff --git a/nodedb-cluster/src/distributed_array/coordinator/elementwise.rs b/nodedb-cluster/src/distributed_array/coordinator/elementwise.rs new file mode 100644 index 000000000..0bfbbca74 --- /dev/null +++ b/nodedb-cluster/src/distributed_array/coordinator/elementwise.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Distributed `ARRAY_ELEMENTWISE`: each shard pairs the cells it owns. +//! +//! Two arrays with one schema hash and one `prefix_bits` route every +//! coordinate to the same shard, so a pair never spans shards. The +//! coordinator sends each primary vShard the request with that shard's +//! bucket, and concatenates the shards' rows. The binary op runs on each +//! shard's Data Plane. + +use crate::error::{ClusterError, Result}; + +use super::super::merge::merge_read_versions; +use super::super::scatter::{ArrayCoordError, FanOutPartitionedParams, fan_out_partitioned}; +use super::super::wire::{ArrayShardElementwiseReq, ArrayShardElementwiseResp}; +use super::read::{ArrayCoordinator, decode_resps, shard_bucket}; + +/// Result of a coordinated elementwise fan-out. +#[derive(Debug, Clone, Default)] +pub struct CoordElementwiseResult { + /// Every shard's rows, in shard-arrival order. + pub rows: Vec>, + /// OR of every shard's below-horizon signal. + pub truncated_before_horizon: bool, + /// The write versions every shard leg observed, one per vShard. + pub read_versions: Vec, +} + +impl ArrayCoordinator { + /// Fan out an elementwise request to every target shard and concatenate + /// the rows. Each shard pairs only the cells in its own bucket. Any shard + /// error fails the whole operation. + pub async fn coord_elementwise( + &self, + req: ArrayShardElementwiseReq, + ) -> std::result::Result { + let prefix_bits = self.params.prefix_bits; + let per_shard: Vec<(u32, Vec)> = self + .params + .shard_ids + .iter() + .map(|&shard_id| { + let per_shard_req = ArrayShardElementwiseReq { + shard_bucket: shard_bucket(shard_id, prefix_bits)?, + ..req.clone() + }; + let bytes = + zerompk::to_msgpack_vec(&per_shard_req).map_err(|e| ClusterError::Codec { + detail: format!("ArrayShardElementwiseReq serialise: {e}"), + })?; + Ok((shard_id, bytes)) + }) + .collect::>>()?; + + let fo_params = FanOutPartitionedParams { + source_node: self.params.source_node, + timeout_ms: self.params.timeout_ms, + }; + let raw = fan_out_partitioned( + &fo_params, + super::super::opcodes::ARRAY_SHARD_ELEMENTWISE_REQ, + &per_shard, + &self.dispatch, + &self.circuit_breaker, + ) + .await?; + let resps = decode_resps::(&raw)?; + Ok(CoordElementwiseResult { + truncated_before_horizon: resps.iter().any(|r| r.truncated_before_horizon), + read_versions: merge_read_versions(resps.iter().map(|r| r.read_versions.as_slice())), + rows: resps.into_iter().flat_map(|r| r.rows_msgpack).collect(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use async_trait::async_trait; + + use super::*; + use crate::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; + use crate::distributed_array::coordinator::read::ArrayCoordParams; + use crate::distributed_array::rpc::ShardRpcDispatch; + use crate::wire::{VShardEnvelope, VShardMessageType}; + + /// Answers each shard with one row naming the bucket it was sent. + struct BucketEcho; + + #[async_trait] + impl ShardRpcDispatch for BucketEcho { + async fn call(&self, req: VShardEnvelope, _timeout_ms: u64) -> Result { + assert_eq!(req.msg_type, VShardMessageType::ArrayShardElementwiseReq); + let shard_req: ArrayShardElementwiseReq = + zerompk::from_msgpack(&req.payload).expect("request decodes"); + let bucket = shard_req.shard_bucket.expect("each shard gets its bucket"); + let resp = ArrayShardElementwiseResp { + shard_id: req.vshard_id, + rows_msgpack: vec![bucket.bucket.to_be_bytes().to_vec()], + // One shard reports its cutoff below the horizon. + truncated_before_horizon: bucket.bucket == 2, + read_versions: Vec::new(), + }; + Ok(VShardEnvelope::new( + VShardMessageType::ArrayShardElementwiseResp, + req.target_node, + req.source_node, + req.vshard_id, + zerompk::to_msgpack_vec(&resp).expect("encode"), + )) + } + } + + fn request() -> ArrayShardElementwiseReq { + ArrayShardElementwiseReq { + left_id_msgpack: Vec::new(), + right_id_msgpack: Vec::new(), + op_msgpack: vec![0x00], + attr_idx: 0, + shard_bucket: None, + system_as_of: None, + valid_at_ms: None, + txn_id: None, + } + } + + /// Every primary shard gets its own bucket, and every shard's rows reach + /// the merged result. + #[tokio::test] + async fn every_shard_pairs_its_own_bucket() { + let bits = 2u8; + let shard_ids = + crate::distributed_array::routing::array_vshards_for_slice(&[], bits, 1024).unwrap(); + assert_eq!(shard_ids.len(), 4); + let coord = ArrayCoordinator::new( + ArrayCoordParams { + source_node: 1, + shard_ids, + timeout_ms: 1000, + prefix_bits: bits, + slice_tile_keys: Vec::new(), + }, + Arc::new(BucketEcho), + Arc::new(CircuitBreaker::new(CircuitBreakerConfig::default())), + ); + let result = coord + .coord_elementwise(request()) + .await + .expect("fan-out succeeds"); + let mut buckets: Vec = result + .rows + .iter() + .map(|row| u32::from_be_bytes(row.as_slice().try_into().expect("4 bytes"))) + .collect(); + buckets.sort_unstable(); + assert_eq!(buckets, vec![0, 1, 2, 3]); + assert!( + result.truncated_before_horizon, + "one shard's below-horizon signal reaches the result" + ); + } +} diff --git a/nodedb-cluster/src/distributed_array/coordinator/maint.rs b/nodedb-cluster/src/distributed_array/coordinator/maint.rs new file mode 100644 index 000000000..99a56a439 --- /dev/null +++ b/nodedb-cluster/src/distributed_array/coordinator/maint.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Distributed `ARRAY_FLUSH` / `ARRAY_COMPACT`: every primary vShard runs the +//! maintenance on the core that holds its tiles. + +use crate::error::ClusterError; + +use super::super::scatter::{ArrayCoordError, fan_out}; +use super::super::wire::{ArrayShardMaintReq, ArrayShardMaintResp}; +use super::read::{ArrayCoordinator, decode_resps}; + +impl ArrayCoordinator { + /// Send `req` to every target shard and return each shard's + /// acknowledgement. Any shard error fails the whole operation. + pub async fn coord_maint( + &self, + req: ArrayShardMaintReq, + ) -> std::result::Result, ArrayCoordError> { + let req_bytes = zerompk::to_msgpack_vec(&req).map_err(|e| ClusterError::Codec { + detail: format!("ArrayShardMaintReq serialise: {e}"), + })?; + let raw = fan_out( + &self.fan_out_params(), + super::super::opcodes::ARRAY_SHARD_MAINT_REQ, + &req_bytes, + &self.dispatch, + &self.circuit_breaker, + ) + .await?; + Ok(decode_resps::(&raw)?) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use async_trait::async_trait; + + use super::*; + use crate::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; + use crate::distributed_array::coordinator::read::ArrayCoordParams; + use crate::distributed_array::rpc::ShardRpcDispatch; + use crate::distributed_array::wire::ArrayShardMaintOp; + use crate::error::Result; + use crate::wire::{VShardEnvelope, VShardMessageType}; + + /// Answers each shard with its own vShard id as the count. + struct ShardEcho; + + #[async_trait] + impl ShardRpcDispatch for ShardEcho { + async fn call(&self, req: VShardEnvelope, _timeout_ms: u64) -> Result { + assert_eq!(req.msg_type, VShardMessageType::ArrayShardMaintReq); + let shard_req: ArrayShardMaintReq = + zerompk::from_msgpack(&req.payload).expect("request decodes"); + assert_eq!(shard_req.op, ArrayShardMaintOp::Flush); + let resp = ArrayShardMaintResp { + shard_id: req.vshard_id, + affected: u64::from(req.vshard_id), + }; + Ok(VShardEnvelope::new( + VShardMessageType::ArrayShardMaintResp, + req.target_node, + req.source_node, + req.vshard_id, + zerompk::to_msgpack_vec(&resp).expect("encode"), + )) + } + } + + /// Every primary shard runs the maintenance and answers. + #[tokio::test] + async fn every_primary_shard_runs_the_maintenance() { + let bits = 2u8; + let shard_ids = + crate::distributed_array::routing::array_vshards_for_slice(&[], bits, 1024).unwrap(); + let expected: Vec = shard_ids.clone(); + let coord = ArrayCoordinator::new( + ArrayCoordParams { + source_node: 1, + shard_ids, + timeout_ms: 1000, + prefix_bits: bits, + slice_tile_keys: Vec::new(), + }, + Arc::new(ShardEcho), + Arc::new(CircuitBreaker::new(CircuitBreakerConfig::default())), + ); + let resps = coord + .coord_maint(ArrayShardMaintReq { + array_id_msgpack: Vec::new(), + op: ArrayShardMaintOp::Flush, + }) + .await + .expect("fan-out succeeds"); + let mut seen: Vec = resps.iter().map(|r| r.shard_id).collect(); + seen.sort_unstable(); + assert_eq!(seen, expected); + assert!(resps.iter().all(|r| r.affected == u64::from(r.shard_id))); + } +} diff --git a/nodedb-cluster/src/distributed_array/coordinator/mod.rs b/nodedb-cluster/src/distributed_array/coordinator/mod.rs index cc0de6bdf..433beb0fe 100644 --- a/nodedb-cluster/src/distributed_array/coordinator/mod.rs +++ b/nodedb-cluster/src/distributed_array/coordinator/mod.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: BUSL-1.1 +pub mod elementwise; +pub mod maint; pub mod read; pub mod write; +pub use elementwise::CoordElementwiseResult; pub use read::{ArrayCoordParams, ArrayCoordinator, CoordAggResult, CoordSliceResult}; pub use write::{ArrayWriteCoordParams, coord_delete, coord_put, coord_put_partitioned}; diff --git a/nodedb-cluster/src/distributed_array/coordinator/read.rs b/nodedb-cluster/src/distributed_array/coordinator/read.rs index 09ea665fe..2b0d09b76 100644 --- a/nodedb-cluster/src/distributed_array/coordinator/read.rs +++ b/nodedb-cluster/src/distributed_array/coordinator/read.rs @@ -73,7 +73,10 @@ pub struct CoordAggResult { /// The bucket filter for a read sent to vShard `shard_id`. `None` when /// `prefix_bits` is 0, which reads every tile. -fn shard_bucket(shard_id: u32, prefix_bits: u8) -> Result> { +pub(super) fn shard_bucket( + shard_id: u32, + prefix_bits: u8, +) -> Result> { if prefix_bits == 0 { return Ok(None); } @@ -323,8 +326,13 @@ mod tests { use async_trait::async_trait; use crate::circuit_breaker::CircuitBreakerConfig; + use crate::distributed_array::group_key::ArrayAggGroupKey; use crate::wire::{VShardEnvelope, VShardMessageType}; + fn g(k: i64) -> ArrayAggGroupKey { + ArrayAggGroupKey::Int64(k) + } + use super::*; /// The version a mock shard leg reports for its vShard. @@ -488,13 +496,14 @@ mod tests { system_as_of: None, valid_at_ms: None, txn_id: None, + region_msgpack: vec![], } } #[tokio::test] async fn coord_agg_merges_scalar_partials_from_shards() { let dispatch: Arc = Arc::new(AggEchoDispatch { - partials: vec![ArrayAggPartial::from_single(0, 10.0)], + partials: vec![ArrayAggPartial::from_single(ArrayAggGroupKey::Scalar, 10.0)], }); // 3 shards each returning a partial with sum=10 → merged sum=30. let coord = make_coordinator(vec![0, 1, 2], dispatch); @@ -505,7 +514,7 @@ mod tests { assert_eq!(merged.partials.len(), 1); assert_eq!(merged.partials[0].count, 3); - assert!((merged.partials[0].sum - 30.0).abs() < f64::EPSILON); + assert!((merged.partials[0].sum.as_f64() - 30.0).abs() < f64::EPSILON); assert!(!merged.truncated_before_horizon); } @@ -554,10 +563,10 @@ mod tests { } let dispatch: Arc = Arc::new(GroupedDispatch { - shard0_partials: vec![ArrayAggPartial::from_single(0, 5.0)], + shard0_partials: vec![ArrayAggPartial::from_single(g(0), 5.0)], shard1_partials: vec![ - ArrayAggPartial::from_single(0, 15.0), - ArrayAggPartial::from_single(1, 20.0), + ArrayAggPartial::from_single(g(0), 15.0), + ArrayAggPartial::from_single(g(1), 20.0), ], }); let coord = make_coordinator(vec![0, 1], dispatch); @@ -571,16 +580,16 @@ mod tests { let g0 = merged .partials .iter() - .find(|p| p.group_key == 0) + .find(|p| p.group_key == g(0)) .expect("group 0"); let g1 = merged .partials .iter() - .find(|p| p.group_key == 1) + .find(|p| p.group_key == g(1)) .expect("group 1"); - assert!((g0.sum - 20.0).abs() < f64::EPSILON); + assert!((g0.sum.as_f64() - 20.0).abs() < f64::EPSILON); assert_eq!(g0.count, 2); - assert!((g1.sum - 20.0).abs() < f64::EPSILON); + assert!((g1.sum.as_f64() - 20.0).abs() < f64::EPSILON); assert_eq!(g1.count, 1); } @@ -596,7 +605,10 @@ mod tests { async fn call(&self, req: VShardEnvelope, _timeout_ms: u64) -> Result { // Shard 1 is below horizon (zero partials); shard 0 has data. let (partials, below) = if req.vshard_id == 0 { - (vec![ArrayAggPartial::from_single(0, 10.0)], false) + ( + vec![ArrayAggPartial::from_single(ArrayAggGroupKey::Scalar, 10.0)], + false, + ) } else { (vec![], true) }; diff --git a/nodedb-cluster/src/distributed_array/coordinator/write.rs b/nodedb-cluster/src/distributed_array/coordinator/write.rs index 17c284ff3..50b96aff1 100644 --- a/nodedb-cluster/src/distributed_array/coordinator/write.rs +++ b/nodedb-cluster/src/distributed_array/coordinator/write.rs @@ -196,7 +196,7 @@ mod tests { applied_lsn: shard_req.wal_lsn, // The shard's cell count, as the Data Plane's `{"inserted": n}` // would report for a batch with no conflicts. - affected: zerompk::from_msgpack::>>(&shard_req.cells_msgpack) + affected: zerompk::from_msgpack::>(&shard_req.cells_msgpack) .unwrap() .len() as u64, }; @@ -236,7 +236,7 @@ mod tests { applied_lsn: shard_req.wal_lsn, // The shard's coordinate count, as the Data Plane's // `{"deleted": n}` would report when every cell exists. - affected: zerompk::from_msgpack::>>(&shard_req.coords_msgpack) + affected: zerompk::from_msgpack::>(&shard_req.coords_msgpack) .unwrap() .len() as u64, }; @@ -346,7 +346,8 @@ mod tests { let p0 = 0x0000_0000_0000_0000u64; let p1 = 0x0040_0000_0000_0000u64; - let coords = vec![(p0, vec![0xAAu8]), (p1, vec![0xBBu8]), (p0, vec![0xCCu8])]; + // Each coord is one msgpack value: a positive fixint. + let coords = vec![(p0, vec![0x0Au8]), (p1, vec![0x0Bu8]), (p0, vec![0x0Cu8])]; let dispatch: Arc = Arc::new(DeleteEchoDispatch); let mut resps = coord_delete(&write_params(), vec![], 10, 55, &coords, &dispatch, &cb()) diff --git a/nodedb-cluster/src/distributed_array/group_key.rs b/nodedb-cluster/src/distributed_array/group_key.rs new file mode 100644 index 000000000..afe5e6c2e --- /dev/null +++ b/nodedb-cluster/src/distributed_array/group_key.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Typed group key of a distributed array aggregate partial. +//! +//! The key mirrors the array coordinate variants without a dependency on +//! `nodedb-array`. Two shards report one group only when their keys are the +//! same typed value, so a Float64 or String dim groups by its value. + +use std::cmp::Ordering; + +use serde::{Deserialize, Serialize}; + +/// The group a partial aggregate belongs to. +/// +/// `Scalar` keys the one partial of an aggregate with no group-by. Float64 +/// keys compare by `f64::total_cmp`, so equal keys have equal bits. +#[derive( + Debug, Clone, Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack, +)] +pub enum ArrayAggGroupKey { + Scalar, + Int64(i64), + Float64(f64), + TimestampMs(i64), + String(String), +} + +impl ArrayAggGroupKey { + fn rank(&self) -> u8 { + match self { + ArrayAggGroupKey::Scalar => 0, + ArrayAggGroupKey::Int64(_) => 1, + ArrayAggGroupKey::Float64(_) => 2, + ArrayAggGroupKey::TimestampMs(_) => 3, + ArrayAggGroupKey::String(_) => 4, + } + } +} + +impl Ord for ArrayAggGroupKey { + fn cmp(&self, other: &Self) -> Ordering { + use ArrayAggGroupKey::*; + match (self, other) { + (Scalar, Scalar) => Ordering::Equal, + (Int64(a), Int64(b)) | (TimestampMs(a), TimestampMs(b)) => a.cmp(b), + (Float64(a), Float64(b)) => a.total_cmp(b), + (String(a), String(b)) => a.cmp(b), + _ => self.rank().cmp(&other.rank()), + } + } +} + +impl PartialOrd for ArrayAggGroupKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for ArrayAggGroupKey { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for ArrayAggGroupKey {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn float_keys_group_by_value() { + assert_eq!( + ArrayAggGroupKey::Float64(1.5), + ArrayAggGroupKey::Float64(1.5) + ); + assert!(ArrayAggGroupKey::Float64(1.5) < ArrayAggGroupKey::Float64(2.5)); + assert_ne!( + ArrayAggGroupKey::Float64(1.5), + ArrayAggGroupKey::Float64(2.5) + ); + } + + #[test] + fn keys_of_different_kinds_never_match() { + assert_ne!(ArrayAggGroupKey::Int64(3), ArrayAggGroupKey::TimestampMs(3)); + assert_ne!(ArrayAggGroupKey::Scalar, ArrayAggGroupKey::Int64(0)); + } + + #[test] + fn a_string_key_round_trips_through_msgpack() { + let key = ArrayAggGroupKey::String("chrX".into()); + let bytes = zerompk::to_msgpack_vec(&key).expect("encode"); + let back: ArrayAggGroupKey = zerompk::from_msgpack(&bytes).expect("decode"); + assert_eq!(back, key); + } +} diff --git a/nodedb-cluster/src/distributed_array/handler.rs b/nodedb-cluster/src/distributed_array/handler.rs index 9ac775379..b574a89b2 100644 --- a/nodedb-cluster/src/distributed_array/handler.rs +++ b/nodedb-cluster/src/distributed_array/handler.rs @@ -19,12 +19,14 @@ use crate::error::{ClusterError, Result}; use super::local_executor::ArrayLocalExecutor; use super::opcodes::{ - ARRAY_SHARD_AGG_REQ, ARRAY_SHARD_DELETE_REQ, ARRAY_SHARD_PUT_REQ, ARRAY_SHARD_SLICE_REQ, + ARRAY_SHARD_AGG_REQ, ARRAY_SHARD_DELETE_REQ, ARRAY_SHARD_ELEMENTWISE_REQ, + ARRAY_SHARD_MAINT_REQ, ARRAY_SHARD_PUT_REQ, ARRAY_SHARD_SLICE_REQ, ARRAY_SHARD_SURROGATE_BITMAP_REQ, }; use super::routing::{array_vshard_for_tile, array_vshards_for_slice}; use super::wire::{ ArrayShardAggReq, ArrayShardAggResp, ArrayShardDeleteReq, ArrayShardDeleteResp, + ArrayShardElementwiseReq, ArrayShardElementwiseResp, ArrayShardMaintReq, ArrayShardMaintResp, ArrayShardPutReq, ArrayShardPutResp, ArrayShardSliceReq, ArrayShardSliceResp, ArrayShardSurrogateBitmapReq, ArrayShardSurrogateBitmapResp, }; @@ -52,6 +54,8 @@ pub async fn handle_array_shard_rpc( ARRAY_SHARD_SURROGATE_BITMAP_REQ => { handle_surrogate_bitmap(local_vshard_id, payload, executor).await } + ARRAY_SHARD_ELEMENTWISE_REQ => handle_elementwise(local_vshard_id, payload, executor).await, + ARRAY_SHARD_MAINT_REQ => handle_maint(local_vshard_id, payload, executor).await, other => Err(ClusterError::Codec { detail: format!("handle_array_shard_rpc: unknown opcode {other}"), }), @@ -104,6 +108,43 @@ async fn handle_agg( serialise(resp) } +async fn handle_elementwise( + local_vshard_id: u32, + payload: &[u8], + executor: &Arc, +) -> Result> { + let req: ArrayShardElementwiseReq = + zerompk::from_msgpack(payload).map_err(|e| ClusterError::Codec { + detail: format!("ArrayShardElementwiseReq decode: {e}"), + })?; + + let exec = executor.exec_elementwise(local_vshard_id, &req).await?; + + let resp = ArrayShardElementwiseResp { + shard_id: local_vshard_id, + rows_msgpack: exec.rows, + truncated_before_horizon: exec.truncated_before_horizon, + read_versions: exec.read_versions, + }; + serialise(resp) +} + +async fn handle_maint( + local_vshard_id: u32, + payload: &[u8], + executor: &Arc, +) -> Result> { + let req: ArrayShardMaintReq = + zerompk::from_msgpack(payload).map_err(|e| ClusterError::Codec { + detail: format!("ArrayShardMaintReq decode: {e}"), + })?; + let affected = executor.exec_maint(local_vshard_id, &req).await?; + serialise(ArrayShardMaintResp { + shard_id: local_vshard_id, + affected, + }) +} + async fn handle_put( local_vshard_id: u32, payload: &[u8], @@ -246,6 +287,7 @@ mod tests { use async_trait::async_trait; + use crate::distributed_array::group_key::ArrayAggGroupKey; use crate::distributed_array::merge::ArrayAggPartial; use crate::distributed_array::wire::{ArrayShardAggReq, ArrayShardDeleteReq, ArrayShardPutReq}; use crate::error::Result; @@ -317,6 +359,26 @@ mod tests { }) } + async fn exec_elementwise( + &self, + local_vshard_id: u32, + _req: &super::super::wire::ArrayShardElementwiseReq, + ) -> Result { + Ok(super::super::local_executor::ArrayElementwiseExec { + rows: self.rows.clone(), + truncated_before_horizon: self.truncated_before_horizon, + read_versions: stub_versions(local_vshard_id), + }) + } + + async fn exec_maint( + &self, + _local_vshard_id: u32, + _req: &super::super::wire::ArrayShardMaintReq, + ) -> Result { + Ok(self.affected) + } + async fn exec_put( &self, _local_vshard_id: u32, @@ -374,6 +436,7 @@ mod tests { system_as_of: None, valid_at_ms: None, txn_id: None, + region_msgpack: vec![], }; zerompk::to_msgpack_vec(&req).unwrap() } @@ -386,6 +449,73 @@ mod tests { zerompk::to_msgpack_vec(&req).unwrap() } + #[tokio::test] + async fn handle_elementwise_returns_executor_rows_and_versions() { + let row = b"pair".to_vec(); + let executor: Arc = Arc::new(StubExecutor { + rows: vec![row.clone()], + bitmap: vec![], + partials: vec![], + truncated_before_horizon: true, + affected: 0, + }); + let req = super::super::wire::ArrayShardElementwiseReq { + left_id_msgpack: vec![], + right_id_msgpack: vec![], + op_msgpack: vec![0x00], + attr_idx: 0, + shard_bucket: None, + system_as_of: None, + valid_at_ms: None, + txn_id: None, + }; + let payload = zerompk::to_msgpack_vec(&req).unwrap(); + let resp_bytes = handle_array_shard_rpc( + super::super::opcodes::ARRAY_SHARD_ELEMENTWISE_REQ, + 3, + &payload, + &executor, + ) + .await + .expect("elementwise handler should succeed"); + + let resp: super::super::wire::ArrayShardElementwiseResp = + zerompk::from_msgpack(&resp_bytes).expect("response should deserialise"); + assert_eq!(resp.shard_id, 3); + assert_eq!(resp.rows_msgpack, vec![row]); + assert!(resp.truncated_before_horizon); + assert_eq!(resp.read_versions, stub_versions(3)); + } + + #[tokio::test] + async fn handle_maint_returns_the_executor_count() { + let executor: Arc = Arc::new(StubExecutor { + rows: vec![], + bitmap: vec![], + partials: vec![], + truncated_before_horizon: false, + affected: 1, + }); + let req = super::super::wire::ArrayShardMaintReq { + array_id_msgpack: vec![], + op: super::super::wire::ArrayShardMaintOp::Compact { + audit_retain_ms: Some(60_000), + }, + }; + let payload = zerompk::to_msgpack_vec(&req).unwrap(); + let resp_bytes = handle_array_shard_rpc( + super::super::opcodes::ARRAY_SHARD_MAINT_REQ, + 5, + &payload, + &executor, + ) + .await + .expect("maint handler should succeed"); + let resp: super::super::wire::ArrayShardMaintResp = + zerompk::from_msgpack(&resp_bytes).expect("response should deserialise"); + assert_eq!((resp.shard_id, resp.affected), (5, 1)); + } + #[tokio::test] async fn handle_slice_returns_executor_rows() { let row = b"row-data".to_vec(); @@ -505,7 +635,7 @@ mod tests { #[tokio::test] async fn handle_agg_returns_executor_partials() { - let partial = ArrayAggPartial::from_single(0, 42.0); + let partial = ArrayAggPartial::from_single(ArrayAggGroupKey::Scalar, 42.0); let executor: Arc = Arc::new(StubExecutor { rows: vec![], bitmap: vec![], @@ -523,7 +653,7 @@ mod tests { assert_eq!(resp.shard_id, 3); assert_eq!(resp.partials.len(), 1); assert_eq!(resp.partials[0].count, 1); - assert!((resp.partials[0].sum - 42.0).abs() < f64::EPSILON); + assert!((resp.partials[0].sum.as_f64() - 42.0).abs() < f64::EPSILON); } fn make_put_req_bytes(prefix_bits: u8) -> Vec { diff --git a/nodedb-cluster/src/distributed_array/local_executor.rs b/nodedb-cluster/src/distributed_array/local_executor.rs index 971737412..48cd7dfb4 100644 --- a/nodedb-cluster/src/distributed_array/local_executor.rs +++ b/nodedb-cluster/src/distributed_array/local_executor.rs @@ -14,7 +14,8 @@ use crate::distributed_array::merge::ArrayAggPartial; use crate::distributed_array::wire::{ - ArrayShardAggReq, ArrayShardDeleteReq, ArrayShardPutReq, ArrayShardSliceReq, + ArrayShardAggReq, ArrayShardDeleteReq, ArrayShardElementwiseReq, ArrayShardMaintReq, + ArrayShardPutReq, ArrayShardSliceReq, }; use crate::error::Result; @@ -43,6 +44,15 @@ pub struct ArrayAggExec { pub read_versions: Vec, } +/// Result of a local shard elementwise op: one row per output cell, plus the +/// bitemporal below-horizon signal (see [`ArraySliceExec`]). +pub struct ArrayElementwiseExec { + pub rows: Vec>, + pub truncated_before_horizon: bool, + /// The write versions the shard read observed (see [`ArraySliceExec`]). + pub read_versions: Vec, +} + /// Result of a local shard write (put or delete): the LSN the coordinator /// acks with, plus the real cell count the Data Plane handler reported. /// @@ -111,6 +121,24 @@ pub trait ArrayLocalExecutor: Send + Sync + 'static { /// this shard. async fn exec_agg(&self, local_vshard_id: u32, req: &ArrayShardAggReq) -> Result; + /// Pair the cells of two co-located arrays on this shard and apply the + /// request's binary op on the Data Plane. Only cells in + /// `req.shard_bucket` take part. `req.txn_id` folds in that transaction's + /// staged cells on this shard. + /// + /// Returns one native-msgpack row per output cell. + async fn exec_elementwise( + &self, + local_vshard_id: u32, + req: &ArrayShardElementwiseReq, + ) -> Result; + + /// Flush or compact the array on the core that owns `local_vshard_id`. + /// + /// Returns the count the Data Plane handler answered with: `1` for a + /// flush, the number of segments merged for a compaction. + async fn exec_maint(&self, local_vshard_id: u32, req: &ArrayShardMaintReq) -> Result; + /// Apply a cell-batch write to the local array engine. /// /// `req.cells_msgpack` is a zerompk encoding of diff --git a/nodedb-cluster/src/distributed_array/merge.rs b/nodedb-cluster/src/distributed_array/merge.rs index 9d11396e7..d29c039a1 100644 --- a/nodedb-cluster/src/distributed_array/merge.rs +++ b/nodedb-cluster/src/distributed_array/merge.rs @@ -20,6 +20,8 @@ use serde::{Deserialize, Serialize}; +use super::agg_number::ArrayAggNumber; +use super::group_key::ArrayAggGroupKey; use super::wire::{ArrayShardAggResp, ArrayShardSliceResp}; use crate::error::{ClusterError, Result}; @@ -31,27 +33,50 @@ use crate::error::{ClusterError, Result}; Debug, Clone, Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack, )] pub struct ArrayAggPartial { - /// Group-by dimension value, or 0 when group_by_dim < 0 (scalar aggregate). - pub group_key: i64, + /// Group-by dimension value, or `Scalar` when group_by_dim < 0. + pub group_key: ArrayAggGroupKey, pub count: u64, - pub sum: f64, - pub min: f64, - pub max: f64, + /// Exact for an integer attr. + pub sum: ArrayAggNumber, + pub min: ArrayAggNumber, + pub max: ArrayAggNumber, /// Welford mean — enables MEAN without a second pass. pub welford_mean: f64, pub welford_m2: f64, } impl ArrayAggPartial { - /// Create from a single cell value. - pub fn from_single(group_key: i64, val: f64) -> Self { + /// Create from a single float cell value. + pub fn from_single(group_key: ArrayAggGroupKey, val: f64) -> Self { + Self::from_number(group_key, ArrayAggNumber::Float(val)) + } + + /// Create from a single integer cell value. + pub fn from_single_int(group_key: ArrayAggGroupKey, val: i64) -> Self { + Self::from_number(group_key, ArrayAggNumber::int(i128::from(val))) + } + + fn from_number(group_key: ArrayAggGroupKey, val: ArrayAggNumber) -> Self { Self { group_key, count: 1, sum: val, min: val, max: val, - welford_mean: val, + welford_mean: val.as_f64(), + welford_m2: 0.0, + } + } + + /// The partial of a group no cell contributed to. Merging it is a no-op. + pub fn empty(group_key: ArrayAggGroupKey) -> Self { + Self { + group_key, + count: 0, + sum: ArrayAggNumber::Float(0.0), + min: ArrayAggNumber::Float(f64::INFINITY), + max: ArrayAggNumber::Float(f64::NEG_INFINITY), + welford_mean: 0.0, welford_m2: 0.0, } } @@ -65,13 +90,9 @@ impl ArrayAggPartial { *self = other.clone(); return; } - self.sum += other.sum; - if other.min < self.min { - self.min = other.min; - } - if other.max > self.max { - self.max = other.max; - } + self.sum = self.sum.sum_with(other.sum); + self.min = self.min.min_with(other.min); + self.max = self.max.max_with(other.max); let new_count = self.count + other.count; let delta = other.welford_mean - self.welford_mean; let combined_mean = (self.welford_mean * self.count as f64 @@ -156,7 +177,8 @@ pub fn merge_slice_rows( /// coord proxy) before applying the limit. /// /// Each row is the native (untagged) msgpack map emitted by the array slice -/// handler — `{"coords": [...], "attrs": [...], "_ts_system": }`. In +/// handler — one field per dim and per attr, then `"_ts_system": `, +/// written in schema order so one cell's rows encode identically. In /// audit-log mode every row carries `_ts_system`; a row that fails to decode /// or is missing that column is a wire-contract violation (corruption or a /// handler bug), so this function returns an error rather than silently @@ -226,11 +248,11 @@ fn decode_system_time(bytes: &[u8]) -> Result { /// Groups by `group_key`; uses `ArrayAggPartial::merge` for each group. pub fn reduce_agg_partials(shard_resps: &[ArrayShardAggResp]) -> Vec { use std::collections::BTreeMap; - let mut buckets: BTreeMap = BTreeMap::new(); + let mut buckets: BTreeMap = BTreeMap::new(); for resp in shard_resps { for partial in &resp.partials { buckets - .entry(partial.group_key) + .entry(partial.group_key.clone()) .and_modify(|existing| existing.merge(partial)) .or_insert_with(|| partial.clone()); } @@ -242,24 +264,43 @@ pub fn reduce_agg_partials(shard_resps: &[ArrayShardAggResp]) -> Vec ArrayAggGroupKey { + ArrayAggGroupKey::Int64(k) + } + #[test] fn reduce_sum_across_shards() { let resp_a = ArrayShardAggResp { shard_id: 0, - partials: vec![ArrayAggPartial::from_single(0, 10.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 10.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let resp_b = ArrayShardAggResp { shard_id: 1, - partials: vec![ArrayAggPartial::from_single(0, 20.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 20.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let merged = reduce_agg_partials(&[resp_a, resp_b]); assert_eq!(merged.len(), 1); assert_eq!(merged[0].count, 2); - assert!((merged[0].sum - 30.0).abs() < f64::EPSILON); + assert!((merged[0].sum.as_f64() - 30.0).abs() < f64::EPSILON); + } + + /// An integer SUM across shards keeps every unit past 2^53. + #[test] + fn reduce_integer_sum_is_exact_across_shards() { + let shard = |shard_id, v| ArrayShardAggResp { + shard_id, + partials: vec![ArrayAggPartial::from_single_int(g(0), v)], + truncated_before_horizon: false, + read_versions: Vec::new(), + }; + let merged = reduce_agg_partials(&[shard(0, 1 << 53), shard(1, 1)]); + assert_eq!(merged[0].sum.as_i128(), Some((1 << 53) + 1)); + assert_eq!(merged[0].max.as_i128(), Some(1 << 53)); + assert_eq!(merged[0].min.as_i128(), Some(1)); } #[test] @@ -267,8 +308,8 @@ mod tests { let resp = ArrayShardAggResp { shard_id: 0, partials: vec![ - ArrayAggPartial::from_single(0, 5.0), - ArrayAggPartial::from_single(1, 15.0), + ArrayAggPartial::from_single(g(0), 5.0), + ArrayAggPartial::from_single(g(1), 15.0), ], truncated_before_horizon: false, read_versions: Vec::new(), @@ -279,14 +320,14 @@ mod tests { #[test] fn merge_empty_partial_is_noop() { - let mut a = ArrayAggPartial::from_single(0, 42.0); + let mut a = ArrayAggPartial::from_single(g(0), 42.0); let empty = ArrayAggPartial { count: 0, - ..ArrayAggPartial::from_single(0, 0.0) + ..ArrayAggPartial::from_single(g(0), 0.0) }; a.merge(&empty); assert_eq!(a.count, 1); - assert!((a.sum - 42.0).abs() < f64::EPSILON); + assert!((a.sum.as_f64() - 42.0).abs() < f64::EPSILON); } #[test] @@ -328,51 +369,51 @@ mod tests { fn reduce_min_across_shards() { let resp_a = ArrayShardAggResp { shard_id: 0, - partials: vec![ArrayAggPartial::from_single(0, 5.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 5.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let resp_b = ArrayShardAggResp { shard_id: 1, - partials: vec![ArrayAggPartial::from_single(0, 3.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 3.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let merged = reduce_agg_partials(&[resp_a, resp_b]); assert_eq!(merged.len(), 1); - assert!((merged[0].min - 3.0).abs() < f64::EPSILON); + assert!((merged[0].min.as_f64() - 3.0).abs() < f64::EPSILON); } #[test] fn reduce_max_across_shards() { let resp_a = ArrayShardAggResp { shard_id: 0, - partials: vec![ArrayAggPartial::from_single(0, 5.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 5.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let resp_b = ArrayShardAggResp { shard_id: 1, - partials: vec![ArrayAggPartial::from_single(0, 99.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 99.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let merged = reduce_agg_partials(&[resp_a, resp_b]); assert_eq!(merged.len(), 1); - assert!((merged[0].max - 99.0).abs() < f64::EPSILON); + assert!((merged[0].max.as_f64() - 99.0).abs() < f64::EPSILON); } #[test] fn reduce_avg_welford_merge_exact() { // Two shards: shard A has one value of 10, shard B has one value of 20. // Combined mean should be exactly 15. - let mut a = ArrayAggPartial::from_single(0, 10.0); - let b = ArrayAggPartial::from_single(0, 20.0); + let mut a = ArrayAggPartial::from_single(g(0), 10.0); + let b = ArrayAggPartial::from_single(g(0), 20.0); a.merge(&b); // welford_mean after merge = 15.0 assert!((a.welford_mean - 15.0).abs() < 1e-9); assert_eq!(a.count, 2); - assert!((a.sum - 30.0).abs() < f64::EPSILON); + assert!((a.sum.as_f64() - 30.0).abs() < f64::EPSILON); } #[test] @@ -381,8 +422,8 @@ mod tests { let resp_a = ArrayShardAggResp { shard_id: 0, partials: vec![ - ArrayAggPartial::from_single(0, 5.0), - ArrayAggPartial::from_single(1, 10.0), + ArrayAggPartial::from_single(g(0), 5.0), + ArrayAggPartial::from_single(g(1), 10.0), ], truncated_before_horizon: false, read_versions: Vec::new(), @@ -390,20 +431,20 @@ mod tests { let resp_b = ArrayShardAggResp { shard_id: 1, partials: vec![ - ArrayAggPartial::from_single(1, 20.0), - ArrayAggPartial::from_single(2, 30.0), + ArrayAggPartial::from_single(g(1), 20.0), + ArrayAggPartial::from_single(g(2), 30.0), ], truncated_before_horizon: false, read_versions: Vec::new(), }; let merged = reduce_agg_partials(&[resp_a, resp_b]); assert_eq!(merged.len(), 3); - let g0 = merged.iter().find(|p| p.group_key == 0).unwrap(); - let g1 = merged.iter().find(|p| p.group_key == 1).unwrap(); - let g2 = merged.iter().find(|p| p.group_key == 2).unwrap(); - assert!((g0.sum - 5.0).abs() < f64::EPSILON); - assert!((g1.sum - 30.0).abs() < f64::EPSILON); - assert!((g2.sum - 30.0).abs() < f64::EPSILON); + let g0 = merged.iter().find(|p| p.group_key == g(0)).unwrap(); + let g1 = merged.iter().find(|p| p.group_key == g(1)).unwrap(); + let g2 = merged.iter().find(|p| p.group_key == g(2)).unwrap(); + assert!((g0.sum.as_f64() - 5.0).abs() < f64::EPSILON); + assert!((g1.sum.as_f64() - 30.0).abs() < f64::EPSILON); + assert!((g2.sum.as_f64() - 30.0).abs() < f64::EPSILON); } #[test] @@ -452,24 +493,66 @@ mod tests { // Shard A has only group 0; shard B has only group 1 — no overlap. let resp_a = ArrayShardAggResp { shard_id: 0, - partials: vec![ArrayAggPartial::from_single(0, 7.0)], + partials: vec![ArrayAggPartial::from_single(g(0), 7.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let resp_b = ArrayShardAggResp { shard_id: 1, - partials: vec![ArrayAggPartial::from_single(1, 13.0)], + partials: vec![ArrayAggPartial::from_single(g(1), 13.0)], truncated_before_horizon: false, read_versions: Vec::new(), }; let merged = reduce_agg_partials(&[resp_a, resp_b]); assert_eq!(merged.len(), 2); - let g0 = merged.iter().find(|p| p.group_key == 0).unwrap(); - let g1 = merged.iter().find(|p| p.group_key == 1).unwrap(); + let g0 = merged.iter().find(|p| p.group_key == g(0)).unwrap(); + let g1 = merged.iter().find(|p| p.group_key == g(1)).unwrap(); assert_eq!(g0.count, 1); assert_eq!(g1.count, 1); } + /// Float64 and String group keys merge by value across shards, and + /// distinct values stay distinct groups. + #[test] + fn reduce_merges_typed_keys_by_value() { + let shard = |shard_id, partials| ArrayShardAggResp { + shard_id, + partials, + truncated_before_horizon: false, + read_versions: Vec::new(), + }; + let f = ArrayAggGroupKey::Float64; + let s = |v: &str| ArrayAggGroupKey::String(v.into()); + let merged = reduce_agg_partials(&[ + shard( + 0, + vec![ + ArrayAggPartial::from_single(f(101.5), 1.0), + ArrayAggPartial::from_single(s("chrX"), 4.0), + ], + ), + shard( + 1, + vec![ + ArrayAggPartial::from_single(f(101.5), 2.0), + ArrayAggPartial::from_single(f(101.25), 8.0), + ArrayAggPartial::from_single(s("chrX"), 16.0), + ], + ), + ]); + assert_eq!(merged.len(), 3); + let sum_of = |key: ArrayAggGroupKey| { + merged + .iter() + .find(|p| p.group_key == key) + .map(|p| p.sum.as_f64()) + .expect("group present") + }; + assert!((sum_of(f(101.5)) - 3.0).abs() < f64::EPSILON); + assert!((sum_of(f(101.25)) - 8.0).abs() < f64::EPSILON); + assert!((sum_of(s("chrX")) - 20.0).abs() < f64::EPSILON); + } + #[test] fn merge_slice_rows_limit_across_shards() { let r0 = ArrayShardSliceResp { @@ -578,8 +661,8 @@ mod tests { // also drops the below-horizon signal. This pins the contract so the // producer and consumer can't drift apart again. let partials = vec![ - ArrayAggPartial::from_single(0, 10.0), - ArrayAggPartial::from_single(1, 20.0), + ArrayAggPartial::from_single(g(0), 10.0), + ArrayAggPartial::from_single(g(1), 20.0), ]; let bytes = zerompk::to_msgpack_vec(&(&partials, true)).expect("encode tuple"); diff --git a/nodedb-cluster/src/distributed_array/mod.rs b/nodedb-cluster/src/distributed_array/mod.rs index acf71a1cd..5c39847bc 100644 --- a/nodedb-cluster/src/distributed_array/mod.rs +++ b/nodedb-cluster/src/distributed_array/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: BUSL-1.1 +pub mod agg_number; pub mod coordinator; +pub mod group_key; pub mod handler; pub mod local_executor; pub mod merge; @@ -11,13 +13,15 @@ pub mod rpc; pub mod scatter; pub mod wire; +pub use agg_number::ArrayAggNumber; pub use coordinator::{ - ArrayCoordinator, ArrayWriteCoordParams, CoordAggResult, CoordSliceResult, coord_delete, - coord_put, coord_put_partitioned, + ArrayCoordinator, ArrayWriteCoordParams, CoordAggResult, CoordElementwiseResult, + CoordSliceResult, coord_delete, coord_put, coord_put_partitioned, }; +pub use group_key::ArrayAggGroupKey; pub use handler::handle_array_shard_rpc; pub use local_executor::{ - ArrayAggExec, ArrayLocalExecutor, ArrayShardWriteOutcome, ArraySliceExec, + ArrayAggExec, ArrayElementwiseExec, ArrayLocalExecutor, ArrayShardWriteOutcome, ArraySliceExec, }; pub use merge::{ ArrayAggPartial, any_truncated_before_horizon_agg, any_truncated_before_horizon_slice, @@ -27,6 +31,7 @@ pub use routing::{array_vshard_for_tile, array_vshards_for_slice}; pub use rpc::ShardRpcDispatch; pub use wire::{ ArrayShardAggReq, ArrayShardAggResp, ArrayShardDeleteReq, ArrayShardDeleteResp, - ArrayShardPutReq, ArrayShardPutResp, ArrayShardSliceReq, ArrayShardSliceResp, - ArrayShardSurrogateBitmapReq, ArrayShardSurrogateBitmapResp, + ArrayShardElementwiseReq, ArrayShardElementwiseResp, ArrayShardMaintOp, ArrayShardMaintReq, + ArrayShardMaintResp, ArrayShardPutReq, ArrayShardPutResp, ArrayShardSliceReq, + ArrayShardSliceResp, ArrayShardSurrogateBitmapReq, ArrayShardSurrogateBitmapResp, }; diff --git a/nodedb-cluster/src/distributed_array/opcodes.rs b/nodedb-cluster/src/distributed_array/opcodes.rs index c1a0d2324..14a90242e 100644 --- a/nodedb-cluster/src/distributed_array/opcodes.rs +++ b/nodedb-cluster/src/distributed_array/opcodes.rs @@ -26,3 +26,11 @@ pub const ARRAY_SHARD_DELETE_RESP: u32 = 87; pub const ARRAY_SHARD_SURROGATE_BITMAP_REQ: u32 = 88; /// Shard → coordinator: surrogate bitmap result. pub const ARRAY_SHARD_SURROGATE_BITMAP_RESP: u32 = 89; +/// Coordinator → shard: pair two arrays' cells and apply a binary op. +pub const ARRAY_SHARD_ELEMENTWISE_REQ: u32 = 90; +/// Shard → coordinator: paired result rows. +pub const ARRAY_SHARD_ELEMENTWISE_RESP: u32 = 91; +/// Coordinator → shard: flush or compact an array. +pub const ARRAY_SHARD_MAINT_REQ: u32 = 92; +/// Shard → coordinator: maintenance acknowledgement. +pub const ARRAY_SHARD_MAINT_RESP: u32 = 93; diff --git a/nodedb-cluster/src/distributed_array/partition.rs b/nodedb-cluster/src/distributed_array/partition.rs index 89ce0b32f..fddba1ceb 100644 --- a/nodedb-cluster/src/distributed_array/partition.rs +++ b/nodedb-cluster/src/distributed_array/partition.rs @@ -15,14 +15,18 @@ //! //! ## Output shape //! +//! Each element arrives as one msgpack value. A bucket is a msgpack array +//! header followed by its elements' bytes, in input order. The bucket is then +//! the msgpack encoding of `Vec`, the shape the shard's Data Plane +//! decodes, so the shard passes it through without a decode. +//! //! `partition_put_cells` returns one `PutBucket` per vShard: its -//! `representative_tile_key` and `cells_msgpack`, a zerompk-encoded -//! `Vec` (the same bytes the shard handler passes to -//! `exec_put`). +//! `representative_tile_key` and `cells_msgpack`, the encoding of +//! `Vec`. //! //! `partition_delete_coords` returns one `DeleteBucket` per vShard: its -//! `representative_tile_key` and `coords_msgpack`, a zerompk-encoded -//! `Vec`. +//! `representative_tile_key` and `coords_msgpack`, the encoding of +//! `Vec`. use std::collections::HashMap; @@ -36,7 +40,7 @@ pub struct PutBucket { /// Tile key of the first cell added to this bucket. The shard handler /// routes it to check ownership (`validate_put_routing`). pub representative_tile_key: u64, - /// zerompk-encoded `Vec` for all cells in this bucket. + /// The msgpack array of every cell in this bucket, in input order. pub cells_msgpack: Vec, } @@ -46,7 +50,7 @@ pub struct DeleteBucket { /// Tile key of the first coord added to this bucket. The shard handler /// routes it to check ownership (`validate_delete_routing`). pub representative_tile_key: u64, - /// zerompk-encoded `Vec` for all coords in this bucket. + /// The msgpack array of every coord in this bucket, in input order. pub coords_msgpack: Vec, } @@ -55,33 +59,17 @@ pub struct DeleteBucket { /// `cells` — each element is `(tile_key, zerompk-encoded single cell)`. /// `prefix_bits` — routing granularity (1–16). /// -/// Returns one `PutBucket` per unique owning vShard. The order of buckets -/// in the output is unspecified but deterministic within a call. +/// Returns one `PutBucket` per unique owning vShard, in the order of each +/// vShard's first cell. pub fn partition_put_cells(cells: &[(u64, Vec)], prefix_bits: u8) -> Result> { - if cells.is_empty() { - return Ok(Vec::new()); - } - - // vshard_id → (representative tile key, per-cell byte blobs) - let mut buckets: HashMap>)> = HashMap::new(); - - for (tile_key, cell_bytes) in cells { - let vshard_id = array_vshard_for_tile(*tile_key, prefix_bits)?; - let entry = buckets.entry(vshard_id).or_insert((*tile_key, Vec::new())); - entry.1.push(cell_bytes.clone()); - } - - buckets + Ok(partition_flat(cells, prefix_bits)? .into_iter() - .map(|(vshard_id, (representative, cell_blobs))| { - let cells_msgpack = encode_blob_vec(&cell_blobs)?; - Ok(PutBucket { - vshard_id, - representative_tile_key: representative, - cells_msgpack, - }) + .map(|bucket| PutBucket { + vshard_id: bucket.vshard_id, + representative_tile_key: bucket.representative_tile_key, + cells_msgpack: bucket.msgpack, }) - .collect() + .collect()) } /// Group `(tile_key, coord_msgpack)` pairs by owning vShard. @@ -89,64 +77,119 @@ pub fn partition_put_cells(cells: &[(u64, Vec)], prefix_bits: u8) -> Result< /// `coords` — each element is `(tile_key, zerompk-encoded single coord)`. /// `prefix_bits` — routing granularity (1–16). /// -/// Returns one `DeleteBucket` per unique owning vShard. +/// Returns one `DeleteBucket` per unique owning vShard, in the order of each +/// vShard's first coord. pub fn partition_delete_coords( coords: &[(u64, Vec)], prefix_bits: u8, ) -> Result> { - if coords.is_empty() { - return Ok(Vec::new()); - } + Ok(partition_flat(coords, prefix_bits)? + .into_iter() + .map(|bucket| DeleteBucket { + vshard_id: bucket.vshard_id, + representative_tile_key: bucket.representative_tile_key, + coords_msgpack: bucket.msgpack, + }) + .collect()) +} - // vshard_id → (representative tile key, coord byte blobs) - let mut buckets: HashMap>)> = HashMap::new(); +/// One vShard's elements as one msgpack array. +struct FlatBucket { + vshard_id: u32, + representative_tile_key: u64, + msgpack: Vec, +} - for (tile_key, coord_bytes) in coords { +/// Group `(tile_key, element_msgpack)` pairs by owning vShard. Each bucket +/// is sized once and written once: no element is copied twice. +fn partition_flat(items: &[(u64, Vec)], prefix_bits: u8) -> Result> { + struct Slot { + vshard_id: u32, + representative_tile_key: u64, + count: usize, + bytes: usize, + } + let mut slot_of_vshard: HashMap = HashMap::new(); + let mut slots: Vec = Vec::new(); + let mut slot_of: Vec = Vec::with_capacity(items.len()); + // Consecutive elements mostly share a vShard, so the last one answers + // before the map. + let mut last: Option<(u32, usize)> = None; + for (tile_key, element) in items { let vshard_id = array_vshard_for_tile(*tile_key, prefix_bits)?; - let entry = buckets.entry(vshard_id).or_insert((*tile_key, Vec::new())); - entry.1.push(coord_bytes.clone()); + let slot = match last { + Some((vshard, slot)) if vshard == vshard_id => slot, + _ => *slot_of_vshard.entry(vshard_id).or_insert_with(|| { + slots.push(Slot { + vshard_id, + representative_tile_key: *tile_key, + count: 0, + bytes: 0, + }); + slots.len() - 1 + }), + }; + last = Some((vshard_id, slot)); + slots[slot].count += 1; + slots[slot].bytes += element.len(); + slot_of.push(slot); } - - buckets + let mut buffers: Vec> = slots + .iter() + .map(|slot| { + let mut buf = Vec::with_capacity(ARRAY_HEADER_MAX + slot.bytes); + write_array_header(&mut buf, slot.count)?; + Ok(buf) + }) + .collect::>()?; + for ((_, element), &slot) in items.iter().zip(&slot_of) { + buffers[slot].extend_from_slice(element); + } + Ok(slots .into_iter() - .map(|(vshard_id, (representative, coord_blobs))| { - let coords_msgpack = encode_blob_vec(&coord_blobs)?; - Ok(DeleteBucket { - vshard_id, - representative_tile_key: representative, - coords_msgpack, - }) + .zip(buffers) + .map(|(slot, msgpack)| FlatBucket { + vshard_id: slot.vshard_id, + representative_tile_key: slot.representative_tile_key, + msgpack, }) - .collect() + .collect()) } -/// Encode a `Vec>` (collection of opaque blobs) as a zerompk array. -/// -/// Each blob is emitted as raw msgpack bytes (already encoded by the caller) -/// wrapped in a msgpack bin or str value. We re-encode the entire collection -/// as a msgpack array of byte-arrays so the shard handler can deserialise it -/// as `Vec` and forward to the engine. -fn encode_blob_vec(blobs: &[Vec]) -> Result> { - // zerompk::to_msgpack_vec requires Sized, so pass as a Vec reference. - let owned: Vec> = blobs.to_vec(); - zerompk::to_msgpack_vec(&owned).map_err(|e| ClusterError::Codec { - detail: format!("partition encode blob vec: {e}"), - }) +/// The longest msgpack array header: the `array 32` marker and a `u32`. +const ARRAY_HEADER_MAX: usize = 5; + +/// Write the msgpack array header for `len` elements, in the shortest form, +/// as zerompk writes it. +fn write_array_header(buf: &mut Vec, len: usize) -> Result<()> { + if len < 16 { + buf.push(0x90 | len as u8); + } else if let Ok(len) = u16::try_from(len) { + buf.push(0xdc); + buf.extend_from_slice(&len.to_be_bytes()); + } else { + let len = u32::try_from(len).map_err(|_| ClusterError::Codec { + detail: format!("partition: a bucket of {len} elements exceeds a msgpack array"), + })?; + buf.push(0xdd); + buf.extend_from_slice(&len.to_be_bytes()); + } + Ok(()) } #[cfg(test)] mod tests { use super::*; - // Build a fake (tile_key, cell_bytes) pair. The cell_bytes are - // opaque for the partitioner — any non-empty Vec works. + // Build a fake (tile_key, cell_bytes) pair. The cell is one msgpack + // value: `tag` as a positive fixint. fn cell(prefix: u64, tag: u8) -> (u64, Vec) { - (prefix, vec![tag]) + (prefix, vec![tag & 0x7f]) } // Build a fake coord pair. fn coord(prefix: u64, tag: u8) -> (u64, Vec) { - (prefix, vec![tag]) + (prefix, vec![tag & 0x7f]) } // Helper: compute expected vshard for a given tile key and prefix_bits. @@ -202,9 +245,34 @@ mod tests { .iter() .find(|b| b.vshard_id == vshard) .expect("bucket for key"); - let blobs: Vec> = zerompk::from_msgpack(&bucket.cells_msgpack).unwrap(); - assert_eq!(blobs.len(), want, "key={key:#x}"); + let cells: Vec = zerompk::from_msgpack(&bucket.cells_msgpack).unwrap(); + assert_eq!(cells.len(), want, "key={key:#x}"); } + // A bucket holds its cells in input order. + let first = buckets + .iter() + .find(|b| b.vshard_id == expected_vshard(p0, 10)) + .expect("bucket for p0"); + let cells: Vec = zerompk::from_msgpack(&first.cells_msgpack).unwrap(); + assert_eq!(cells, vec![0x01, 0x03]); + } + + /// A bucket of more than 15 elements takes the `array 16` header, and + /// still decodes as the flat array of its elements. + #[test] + fn a_large_bucket_is_one_flat_msgpack_array() { + let cells: Vec<(u64, Vec)> = (0..40u64) + .map(|i| (0, zerompk::to_msgpack_vec(&(i * 1_000)).unwrap())) + .collect(); + let buckets = partition_put_cells(&cells, 10).unwrap(); + assert_eq!(buckets.len(), 1); + let decoded: Vec = zerompk::from_msgpack(&buckets[0].cells_msgpack).unwrap(); + assert_eq!(decoded, (0..40u64).map(|i| i * 1_000).collect::>()); + assert_eq!( + buckets[0].cells_msgpack, + zerompk::to_msgpack_vec(&decoded).unwrap(), + "the bucket is byte-identical to zerompk's own encoding" + ); } #[test] @@ -222,8 +290,8 @@ mod tests { .iter() .find(|b| b.vshard_id == vshard) .expect("bucket for key"); - let blobs: Vec> = zerompk::from_msgpack(&bucket.coords_msgpack).unwrap(); - assert_eq!(blobs.len(), want, "key={key:#x}"); + let coords: Vec = zerompk::from_msgpack(&bucket.coords_msgpack).unwrap(); + assert_eq!(coords.len(), want, "key={key:#x}"); } } diff --git a/nodedb-cluster/src/distributed_array/routing.rs b/nodedb-cluster/src/distributed_array/routing.rs index 29125d5aa..9b64379ff 100644 --- a/nodedb-cluster/src/distributed_array/routing.rs +++ b/nodedb-cluster/src/distributed_array/routing.rs @@ -100,10 +100,11 @@ mod tests { #[test] fn slice_routing_empty_keys_returns_every_primary() { - // prefix_bits=8 (stride=4): primaries 0, 4, …, 1020. + // prefix_bits=8 (stride=4): one primary in each stride of four. let shards = array_vshards_for_slice(&[], 8, VSHARD_COUNT).unwrap(); assert_eq!(shards.len(), 256); - assert!(shards.iter().all(|v| v % 4 == 0)); + let buckets: Vec = shards.iter().map(|v| v >> 2).collect(); + assert_eq!(buckets, (0u32..256).collect::>()); // total_shards=4 keeps only primary 0. assert_eq!(array_vshards_for_slice(&[], 8, 4).unwrap(), vec![0]); } diff --git a/nodedb-cluster/src/distributed_array/scatter.rs b/nodedb-cluster/src/distributed_array/scatter.rs index 9dd0fc845..d88e25c83 100644 --- a/nodedb-cluster/src/distributed_array/scatter.rs +++ b/nodedb-cluster/src/distributed_array/scatter.rs @@ -314,7 +314,7 @@ pub async fn fan_out_partitioned( /// Map an opcode constant to a `VShardMessageType`. /// -/// All array opcodes are in the range 80–89. Any other value is a +/// All array opcodes are in the range 80–93. Any other value is a /// programming error in the coordinator, not a runtime condition. fn msg_type_from_opcode(opcode: u32) -> Result { match opcode { @@ -328,6 +328,10 @@ fn msg_type_from_opcode(opcode: u32) -> Result { 87 => Ok(VShardMessageType::ArrayShardDeleteResp), 88 => Ok(VShardMessageType::ArrayShardSurrogateBitmapReq), 89 => Ok(VShardMessageType::ArrayShardSurrogateBitmapResp), + 90 => Ok(VShardMessageType::ArrayShardElementwiseReq), + 91 => Ok(VShardMessageType::ArrayShardElementwiseResp), + 92 => Ok(VShardMessageType::ArrayShardMaintReq), + 93 => Ok(VShardMessageType::ArrayShardMaintResp), other => Err(ClusterError::Codec { detail: format!("msg_type_from_opcode: unknown array opcode {other}"), }), diff --git a/nodedb-cluster/src/distributed_array/wire.rs b/nodedb-cluster/src/distributed_array/wire.rs index 683f796d5..3b3a7616f 100644 --- a/nodedb-cluster/src/distributed_array/wire.rs +++ b/nodedb-cluster/src/distributed_array/wire.rs @@ -105,6 +105,9 @@ pub struct ArrayShardAggReq { /// Each shard folds only its own overlay cells into its partial, so the /// coordinator's merge counts every staged cell exactly once. pub txn_id: Option, + /// Zerompk encoding of `nodedb_array::query::ArrayRegion`, forwarded from + /// `ArrayOp::Aggregate::region_msgpack`. Empty = the whole array. + pub region_msgpack: Vec, } /// Gather response: shard returns partial aggregate(s) for merge. @@ -121,6 +124,81 @@ pub struct ArrayShardAggResp { pub read_versions: Vec, } +/// Scatter request: coordinator asks a shard to pair the cells of two arrays +/// it owns and apply a binary op to each pair. +/// +/// A cell routes by its tile key, and the tile key is a function of the +/// schema's dims and tile extents alone. Two arrays with one schema hash and +/// one `prefix_bits` therefore place every coordinate on the same shard, so +/// each shard pairs its own cells and no cell crosses shards. +#[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)] +pub struct ArrayShardElementwiseReq { + /// Zerompk encoding of the left `nodedb_array::types::ArrayId`. + pub left_id_msgpack: Vec, + /// Zerompk encoding of the right `nodedb_array::types::ArrayId`. + pub right_id_msgpack: Vec, + /// Zerompk encoding of `ArrayBinaryOp` (c_enum, 1 byte on wire). + pub op_msgpack: Vec, + pub attr_idx: u32, + /// The bucket this shard owns. The shard pairs only cells whose tile key + /// falls in it, so vShards that share one Data Plane emit no row twice. + /// `None` pairs every cell. + pub shard_bucket: Option, + /// Bitemporal system-time cutoff for both operands. `None` = live read. + pub system_as_of: Option, + /// Bitemporal valid-time point for both operands. `None` = no + /// valid-time filter. + pub valid_at_ms: Option, + /// The reading transaction's id, for read-your-own-writes against the + /// shard's staging overlay (see [`ArrayShardSliceReq::txn_id`]). + pub txn_id: Option, +} + +/// Gather response: shard returns its paired rows. +#[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)] +pub struct ArrayShardElementwiseResp { + pub shard_id: u32, + /// One element per output cell, each the native-msgpack encoding of one + /// `{"coords", "attrs"}` row. + pub rows_msgpack: Vec>, + /// Whether the system-time cutoff predates this shard's retained history + /// (see [`ArrayShardSliceResp::truncated_before_horizon`]). + pub truncated_before_horizon: bool, + /// The write versions the shard read observed (see + /// [`ArrayShardSliceResp::read_versions`]). + pub read_versions: Vec, +} + +/// The maintenance a shard runs on its local copy of an array. +#[derive(Debug, Clone, PartialEq, Eq, zerompk::ToMessagePack, zerompk::FromMessagePack)] +pub enum ArrayShardMaintOp { + /// Flush the memtable to a segment. + Flush, + /// Merge segments. `audit_retain_ms` is the array's audit-retention + /// horizon from the catalog; `None` keeps every version. + Compact { audit_retain_ms: Option }, +} + +/// Scatter request: coordinator asks a shard to flush or compact an array. +/// +/// A core holds the tiles of every bucket its vShards own, so the request +/// goes to every primary vShard. A core that two vShards share runs it +/// twice; the second run finds nothing left to do. +#[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)] +pub struct ArrayShardMaintReq { + pub array_id_msgpack: Vec, + pub op: ArrayShardMaintOp, +} + +/// Acknowledgement: shard ran the maintenance. +#[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)] +pub struct ArrayShardMaintResp { + pub shard_id: u32, + /// The count the Data Plane handler answered with: `1` for a flush, the + /// number of segments merged for a compaction. + pub affected: u64, +} + /// Scatter request: coordinator forwards a cell write to the owning shard. /// /// `cells_msgpack` is a zerompk encoding of @@ -152,7 +230,8 @@ pub struct ArrayShardPutResp { /// Scatter request: coordinator asks a shard to delete cells by exact coords. /// -/// `coords_msgpack` is a zerompk encoding of `Vec>`. +/// `coords_msgpack` is a zerompk encoding of +/// `Vec`. #[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)] pub struct ArrayShardDeleteReq { pub array_id_msgpack: Vec, diff --git a/nodedb-cluster/src/distributed_spatial/coordinator.rs b/nodedb-cluster/src/distributed_spatial/coordinator.rs index 36faefa68..dc47e9760 100644 --- a/nodedb-cluster/src/distributed_spatial/coordinator.rs +++ b/nodedb-cluster/src/distributed_spatial/coordinator.rs @@ -23,8 +23,9 @@ pub struct SpatialScatterPayload { pub collection: String, pub field: String, pub predicate: String, - /// Typed query geometry, parsed and validated on the originating CP. - pub query_geometry: nodedb_types::geometry::Geometry, + /// Typed query geometry and its SRID, parsed and validated on the + /// originating CP. + pub query_geometry: nodedb_types::geometry::SridGeometry, pub distance_meters: f64, pub limit: u32, } @@ -70,7 +71,7 @@ impl SpatialScatterGather { collection: &str, field: &str, predicate: &str, - query_geometry: nodedb_types::geometry::Geometry, + query_geometry: nodedb_types::geometry::SridGeometry, distance_meters: f64, limit: usize, ) -> Result> { @@ -210,7 +211,9 @@ mod tests { "buildings", "geom", "st_dwithin", - nodedb_types::geometry::Geometry::point(0.0, 0.0), + nodedb_types::geometry::SridGeometry::wgs84( + nodedb_types::geometry::Geometry::point(0.0, 0.0), + ), 1000.0, 100, ) diff --git a/nodedb-cluster/src/rpc_codec/discriminants.rs b/nodedb-cluster/src/rpc_codec/discriminants.rs index 4fcf1567a..dc52abb6b 100644 --- a/nodedb-cluster/src/rpc_codec/discriminants.rs +++ b/nodedb-cluster/src/rpc_codec/discriminants.rs @@ -167,7 +167,7 @@ pub const RPC_FRAME_REFUSAL: u8 = 58; /// reason, so the sender learns the peer is up and refused this request. pub const RPC_REQUEST_REFUSAL: u8 = 59; -// VShardMessageType discriminants for distributed array ops (u16, range 80-89). +// VShardMessageType discriminants for distributed array ops (u16, range 80-93). // These mirror `crate::wire::VShardMessageType` repr values and are declared // here so external code can reference them without importing the full enum. pub const VSHARD_ARRAY_SHARD_SLICE_REQ: u16 = 80; @@ -180,3 +180,7 @@ pub const VSHARD_ARRAY_SHARD_DELETE_REQ: u16 = 86; pub const VSHARD_ARRAY_SHARD_DELETE_RESP: u16 = 87; pub const VSHARD_ARRAY_SHARD_SURROGATE_BITMAP_REQ: u16 = 88; pub const VSHARD_ARRAY_SHARD_SURROGATE_BITMAP_RESP: u16 = 89; +pub const VSHARD_ARRAY_SHARD_ELEMENTWISE_REQ: u16 = 90; +pub const VSHARD_ARRAY_SHARD_ELEMENTWISE_RESP: u16 = 91; +pub const VSHARD_ARRAY_SHARD_MAINT_REQ: u16 = 92; +pub const VSHARD_ARRAY_SHARD_MAINT_RESP: u16 = 93; diff --git a/nodedb-cluster/src/vshard_handler.rs b/nodedb-cluster/src/vshard_handler.rs index d11694861..17e7402a3 100644 --- a/nodedb-cluster/src/vshard_handler.rs +++ b/nodedb-cluster/src/vshard_handler.rs @@ -102,6 +102,10 @@ pub fn dispatch_by_type(envelope: &VShardEnvelope) -> DispatchTarget { VShardMessageType::ArrayShardDeleteResp => DispatchTarget::ArrayCoordinator, VShardMessageType::ArrayShardSurrogateBitmapReq => DispatchTarget::ArrayShard, VShardMessageType::ArrayShardSurrogateBitmapResp => DispatchTarget::ArrayCoordinator, + VShardMessageType::ArrayShardElementwiseReq => DispatchTarget::ArrayShard, + VShardMessageType::ArrayShardElementwiseResp => DispatchTarget::ArrayCoordinator, + VShardMessageType::ArrayShardMaintReq => DispatchTarget::ArrayShard, + VShardMessageType::ArrayShardMaintResp => DispatchTarget::ArrayCoordinator, } } @@ -248,6 +252,10 @@ mod tests { VShardMessageType::ArrayShardDeleteResp, VShardMessageType::ArrayShardSurrogateBitmapReq, VShardMessageType::ArrayShardSurrogateBitmapResp, + VShardMessageType::ArrayShardElementwiseReq, + VShardMessageType::ArrayShardElementwiseResp, + VShardMessageType::ArrayShardMaintReq, + VShardMessageType::ArrayShardMaintResp, ]; for msg_type in all_types { let env = VShardEnvelope::new(msg_type, 1, 2, 0, vec![]); diff --git a/nodedb-cluster/src/wire.rs b/nodedb-cluster/src/wire.rs index 90a9f7e3a..f236659bd 100644 --- a/nodedb-cluster/src/wire.rs +++ b/nodedb-cluster/src/wire.rs @@ -143,6 +143,14 @@ pub enum VShardMessageType { ArrayShardSurrogateBitmapReq = 88, /// Gather: shard returns the surrogate bitmap for matching cells. ArrayShardSurrogateBitmapResp = 89, + /// Scatter: coordinator asks a shard to pair two arrays' cells. + ArrayShardElementwiseReq = 90, + /// Gather: shard returns its paired rows. + ArrayShardElementwiseResp = 91, + /// Scatter: coordinator asks a shard to flush or compact an array. + ArrayShardMaintReq = 92, + /// Gather: shard acknowledges the maintenance. + ArrayShardMaintResp = 93, } /// Current wire protocol version. @@ -263,6 +271,10 @@ impl VShardMessageType { 87 => VShardMessageType::ArrayShardDeleteResp, 88 => VShardMessageType::ArrayShardSurrogateBitmapReq, 89 => VShardMessageType::ArrayShardSurrogateBitmapResp, + 90 => VShardMessageType::ArrayShardElementwiseReq, + 91 => VShardMessageType::ArrayShardElementwiseResp, + 92 => VShardMessageType::ArrayShardMaintReq, + 93 => VShardMessageType::ArrayShardMaintResp, _ => return None, }; Some(msg_type) @@ -334,6 +346,10 @@ mod tests { VShardMessageType::ArrayShardDeleteResp, VShardMessageType::ArrayShardSurrogateBitmapReq, VShardMessageType::ArrayShardSurrogateBitmapResp, + VShardMessageType::ArrayShardElementwiseReq, + VShardMessageType::ArrayShardElementwiseResp, + VShardMessageType::ArrayShardMaintReq, + VShardMessageType::ArrayShardMaintResp, ]; for msg_type in types { diff --git a/nodedb-codec/src/double_delta.rs b/nodedb-codec/src/double_delta.rs index e42cec162..14a3a788a 100644 --- a/nodedb-codec/src/double_delta.rs +++ b/nodedb-codec/src/double_delta.rs @@ -119,11 +119,38 @@ impl<'a> BitReader<'a> { Ok(bit) } + /// Read `num_bits` (at most 64) bits, most significant first. Whole + /// byte runs are taken at once, so a wide field costs one step per byte, + /// not one per bit. A read past the buffer fails and consumes nothing. pub(crate) fn read_bits(&mut self, num_bits: usize) -> Result { + if num_bits == 0 { + return Ok(0); + } + if num_bits > 64 { + return Err(CodecError::Corrupt { + detail: format!("bit field of {num_bits} bits exceeds 64"), + }); + } + let end = checked_add(self.bit_pos, num_bits, "DoubleDelta bit offset")?; + let bytes_needed = end.div_ceil(8); + if bytes_needed > self.buf.len() { + return Err(CodecError::Truncated { + expected: bytes_needed, + actual: self.buf.len(), + }); + } let mut value = 0u64; - for _ in 0..num_bits { - value = (value << 1) | u64::from(self.read_bit()?); + let mut pos = self.bit_pos; + while pos < end { + let byte = self.buf[pos / 8]; + let offset = pos % 8; + let available = 8 - offset; + let take = available.min(end - pos); + let bits = (u64::from(byte) >> (available - take)) & ((1u64 << take) - 1); + value = (value << take) | bits; + pos += take; } + self.bit_pos = end; Ok(value) } } @@ -596,4 +623,43 @@ mod tests { assert!(decode(&[1, 0, 0, 0]).is_err()); assert!(decode(&[2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]).is_err()); } + + /// Fields of every width at every bit offset read back as written, and + /// a read past the end fails without consuming any bit. + #[test] + fn read_bits_matches_written_fields_at_any_offset() { + let fields: Vec<(u64, usize)> = (0..=64usize) + .map(|width| { + let value = match width { + 0 => 0, + 64 => 0xA5C3_0F1E_9B7D_2468, + w => 0x9E37_79B9_7F4A_7C15 & ((1u64 << w) - 1), + }; + (value, width) + }) + .collect(); + for lead in 0..8usize { + let mut w = BitWriter::new(); + w.write_bits(0, lead); + for &(value, width) in &fields { + w.write_bits(value, width); + } + let mut r = BitReader::new(w.as_bytes()); + assert_eq!(r.read_bits(lead).unwrap(), 0); + for &(value, width) in &fields { + assert_eq!( + r.read_bits(width).unwrap(), + value, + "width {width}, lead {lead}" + ); + } + } + + let bytes = [0xFFu8]; + let mut r = BitReader::new(&bytes); + assert_eq!(r.read_bits(3).unwrap(), 0b111); + assert!(matches!(r.read_bits(6), Err(CodecError::Truncated { .. }))); + assert_eq!(r.read_bits(5).unwrap(), 0b11111); + assert!(matches!(r.read_bits(65), Err(CodecError::Corrupt { .. }))); + } } diff --git a/nodedb-codec/src/gorilla_f32.rs b/nodedb-codec/src/gorilla_f32.rs new file mode 100644 index 000000000..a26998181 --- /dev/null +++ b/nodedb-codec/src/gorilla_f32.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Gorilla XOR encoding for `f32` value columns. +//! +//! The 32-bit form of the Gorilla value channel: each value XORs against +//! the previous one, and only the meaningful bits of a non-zero XOR are +//! written. No timestamp channel. +//! +//! Wire format: +//! ```text +//! [4 bytes] value count (LE u32) +//! [N bytes] bitstream: first value raw (32 bits), then per value: +//! 0 XOR is zero +//! 1 0 XOR fits the previous window +//! 1 1 <5 lead> <5 len-1> new window +//! ``` + +use crate::double_delta::{BitReader, BitWriter}; +use crate::error::CodecError; + +/// Encode `values` with 32-bit Gorilla XOR compression. +pub fn encode_f32(values: &[f32]) -> Vec { + let mut buf = BitWriter::new(); + let mut prev = 0u32; + let mut prev_leading = u8::MAX; + let mut prev_trailing = 0u8; + for (i, v) in values.iter().enumerate() { + let bits = v.to_bits(); + if i == 0 { + buf.write_bits(u64::from(bits), 32); + prev = bits; + continue; + } + let xor = prev ^ bits; + prev = bits; + if xor == 0 { + buf.write_bit(false); + continue; + } + buf.write_bit(true); + let leading = xor.leading_zeros() as u8; + let trailing = xor.trailing_zeros() as u8; + if prev_leading != u8::MAX && leading >= prev_leading && trailing >= prev_trailing { + buf.write_bit(false); + let meaningful = 32 - prev_leading - prev_trailing; + buf.write_bits(u64::from(xor >> prev_trailing), meaningful as usize); + } else { + buf.write_bit(true); + let meaningful = 32 - leading - trailing; + buf.write_bits(u64::from(leading), 5); + buf.write_bits(u64::from(meaningful - 1), 5); + buf.write_bits(u64::from(xor >> trailing), meaningful as usize); + prev_leading = leading; + prev_trailing = trailing; + } + } + let bitstream = buf.as_bytes(); + let mut out = Vec::with_capacity(4 + bitstream.len()); + out.extend_from_slice(&(values.len() as u32).to_le_bytes()); + out.extend_from_slice(bitstream); + out +} + +/// Decode values written by [`encode_f32`]. +pub fn decode_f32(data: &[u8]) -> Result, CodecError> { + let Some((header, body)) = data.split_first_chunk::<4>() else { + return Err(CodecError::Truncated { + expected: 4, + actual: data.len(), + }); + }; + let count = u32::from_le_bytes(*header) as usize; + if count == 0 { + return Ok(Vec::new()); + } + // The first value takes 32 bits and every later value at least one, so a + // count past that bound is corrupt and is refused before allocation. + let max_count = (body.len() * 8).saturating_sub(31); + if count > max_count { + return Err(CodecError::Corrupt { + detail: format!( + "gorilla_f32 count {count} exceeds the {max_count} values a {}-byte stream holds", + body.len() + ), + }); + } + let count = crate::bounds::checked_capacity(count, size_of::(), "gorilla_f32 count")?; + let mut reader = BitReader::new(body); + let mut out = Vec::with_capacity(count); + let mut prev = reader.read_bits(32)? as u32; + out.push(f32::from_bits(prev)); + let mut prev_leading = 0u8; + let mut prev_trailing = 0u8; + for _ in 1..count { + if !reader.read_bit()? { + out.push(f32::from_bits(prev)); + continue; + } + let xor = if !reader.read_bit()? { + let meaningful = 32 - prev_leading - prev_trailing; + (reader.read_bits(meaningful as usize)? as u32) << prev_trailing + } else { + let leading = reader.read_bits(5)? as u8; + let meaningful = reader.read_bits(5)? as u8 + 1; + if leading + meaningful > 32 { + return Err(CodecError::Corrupt { + detail: format!( + "invalid gorilla_f32 window: {leading} leading + {meaningful} meaningful bits" + ), + }); + } + let trailing = 32 - leading - meaningful; + prev_leading = leading; + prev_trailing = trailing; + (reader.read_bits(meaningful as usize)? as u32) << trailing + }; + prev ^= xor; + out.push(f32::from_bits(prev)); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn round_trip(values: &[f32]) -> Vec { + decode_f32(&encode_f32(values)).expect("decode") + } + + #[test] + fn empty_round_trips() { + assert!(round_trip(&[]).is_empty()); + } + + #[test] + fn values_round_trip_bit_exact() { + let values = [ + 0.1f32, + 0.1, + -2.5, + f32::MAX, + f32::MIN_POSITIVE, + 0.0, + -0.0, + f32::INFINITY, + 1.0e-40, + ]; + let out = round_trip(&values); + let bits: Vec = out.iter().map(|v| v.to_bits()).collect(); + let want: Vec = values.iter().map(|v| v.to_bits()).collect(); + assert_eq!(bits, want); + } + + #[test] + fn nan_keeps_its_bits() { + let out = round_trip(&[f32::NAN, 1.0, f32::NAN]); + assert!(out[0].is_nan() && out[2].is_nan()); + assert_eq!(out[1], 1.0); + } + + #[test] + fn smooth_series_beats_raw_width() { + let values: Vec = (0..1024).map(|i| 20.0 + (i % 16) as f32 * 0.25).collect(); + let encoded = encode_f32(&values); + assert!(encoded.len() < values.len() * 4, "{} bytes", encoded.len()); + assert_eq!(round_trip(&values), values); + } + + #[test] + fn a_count_past_the_stream_is_refused() { + let mut data = (1_000_000u32).to_le_bytes().to_vec(); + data.extend_from_slice(&[0u8; 8]); + assert!(matches!(decode_f32(&data), Err(CodecError::Corrupt { .. }))); + } + + #[test] + fn a_truncated_header_is_refused() { + assert!(matches!( + decode_f32(&[1, 0]), + Err(CodecError::Truncated { .. }) + )); + } +} diff --git a/nodedb-codec/src/lib.rs b/nodedb-codec/src/lib.rs index 4485697d5..c41e084f8 100644 --- a/nodedb-codec/src/lib.rs +++ b/nodedb-codec/src/lib.rs @@ -31,6 +31,7 @@ pub mod error; pub mod fastlanes; pub mod fsst; pub mod gorilla; +pub mod gorilla_f32; pub mod lz4; pub mod pcodec; pub mod pipeline; diff --git a/nodedb-physical/src/physical_plan/array.rs b/nodedb-physical/src/physical_plan/array.rs index b0e3606dd..153b76a86 100644 --- a/nodedb-physical/src/physical_plan/array.rs +++ b/nodedb-physical/src/physical_plan/array.rs @@ -3,7 +3,7 @@ //! Array engine operations dispatched to the Data Plane. //! //! `ArrayOp` is the wire type for array query operators (slice, -//! project, aggregate, elementwise) plus the write-side ops (put, +//! aggregate, elementwise) plus the write-side ops (put, //! delete) and engine maintenance (flush, compact). Complex nested //! payloads — schemas, slice predicates, cell batches — ride as //! opaque msgpack bytes (`*_msgpack`) decoded on the Data Plane via @@ -91,8 +91,8 @@ pub enum ArrayOp { }, /// Insert one or more cells. `cells_msgpack` is an zerompk - /// encoding of `Vec`. - /// Each cell carries a Control-Plane-allocated `Surrogate` parallel + /// encoding of `Vec`. The WAL + /// record splices these bytes in without decoding them. Each cell carries a Control-Plane-allocated `Surrogate` parallel /// to its `coord` tuple — cross-engine bitmap joins read the /// surrogate column on read paths instead of translating coords /// back to user-visible PKs. @@ -113,8 +113,9 @@ pub enum ArrayOp { vshard_id: u32, }, - /// Delete by exact coordinates. `coords_msgpack` is an zerompk - /// encoding of `Vec>`. + /// Delete by exact coordinates. `coords_msgpack` is a zerompk + /// encoding of `Vec`. + /// The WAL record splices these bytes in without decoding them. /// /// `provenance` is `Some` only on the inbound sync path (Lite → Origin). Delete { @@ -168,12 +169,6 @@ pub enum ArrayOp { valid_at_ms: Option, }, - /// Attribute projection (no coord filter). - Project { - array_id: ArrayId, - attr_indices: Vec, - }, - /// Reduce one attribute column. `group_by_dim < 0` means no /// group-by (one scalar partial per tile, merged across tiles by /// the executor). @@ -203,6 +198,9 @@ pub enum ArrayOp { /// `valid_from_ms <= vt < valid_until_ms` contribute. `None` = no /// valid-time filter. valid_at_ms: Option, + /// zerompk-encoded `nodedb_array::query::ArrayRegion`. Only cells + /// inside it contribute. Empty = the whole array. + region_msgpack: Vec, }, /// Pairwise op between two coord-aligned arrays. Both must share @@ -217,6 +215,16 @@ pub enum ArrayOp { /// fallthroughs from either side are also excluded. /// `None` = all cells. cell_filter: Option, + /// Optional shard bucket. When set, only cells whose tile key routes + /// to this bucket take part on either side, so vShards that share + /// one Data Plane emit no row twice. `None` includes every cell. + shard_bucket: Option, + /// Bitemporal system-time cutoff for both operands. `None` = live + /// read. + system_as_of: Option, + /// Bitemporal valid-time point for both operands. `None` = no + /// valid-time filter. + valid_at_ms: Option, }, /// Force a memtable flush. Returns the new segment ref's id + @@ -265,9 +273,9 @@ pub enum ArrayOp { } impl ArrayOp { - /// The array this op targets. For `Elementwise`, returns the - /// left operand — vshard routing pins to the left array's - /// shard and the right is fetched cross-shard. + /// The array this op targets. For `Elementwise`, returns the left + /// operand. Both operands share one tile layout and routing, so a + /// shard holds the matching cells of both. pub fn primary_array(&self) -> &ArrayId { match self { ArrayOp::OpenArray { array_id, .. } @@ -275,7 +283,6 @@ impl ArrayOp { | ArrayOp::Delete { array_id, .. } | ArrayOp::Slice { array_id, .. } | ArrayOp::SurrogateBitmapScan { array_id, .. } - | ArrayOp::Project { array_id, .. } | ArrayOp::Aggregate { array_id, .. } | ArrayOp::Flush { array_id, .. } | ArrayOp::Compact { array_id, .. } @@ -308,6 +315,7 @@ mod tests { shard_bucket: None, system_as_of: None, valid_at_ms: None, + region_msgpack: vec![1, 2, 3], }; let bytes = zerompk::to_msgpack_vec(&op).unwrap(); let back: ArrayOp = zerompk::from_msgpack(&bytes).unwrap(); @@ -368,6 +376,9 @@ mod tests { op: ArrayBinaryOp::Add, attr_idx: 0, cell_filter: None, + shard_bucket: None, + system_as_of: None, + valid_at_ms: None, }; assert_eq!(op.primary_array().name, "L"); } diff --git a/nodedb-physical/src/physical_plan/cluster_array.rs b/nodedb-physical/src/physical_plan/cluster_array.rs index acf6487c2..a4e21fcbd 100644 --- a/nodedb-physical/src/physical_plan/cluster_array.rs +++ b/nodedb-physical/src/physical_plan/cluster_array.rs @@ -16,6 +16,8 @@ use nodedb_array::types::ArrayId; use nodedb_types::SystemTimeScope; +use super::array::ArrayBinaryOp; + /// Cluster-mode array operations executed by the coordinator on the /// Control Plane. #[derive( @@ -76,6 +78,32 @@ pub enum ClusterArrayOp { system_as_of: Option, /// Bitemporal valid-time point. `None` = no valid-time filter. valid_at_ms: Option, + /// zerompk-encoded `nodedb_array::query::ArrayRegion`, forwarded to + /// every shard. Empty = the whole array. + region_msgpack: Vec, + }, + + /// Fan-out an elementwise op to every primary vShard and concatenate the + /// rows. + /// + /// The planner admits only two arrays with one schema hash and one + /// `prefix_bits`. Their tile keys then match per coordinate, so each + /// shard holds both sides of every pair it owns and pairs them on its + /// Data Plane. + Elementwise { + left: ArrayId, + right: ArrayId, + op: ArrayBinaryOp, + /// Index of the attribute the op reads. + attr_idx: u32, + /// Routing granularity, shared by both arrays. + prefix_bits: u8, + /// Bitemporal system-time cutoff for both operands. `None` = live + /// read. + system_as_of: Option, + /// Bitemporal valid-time point for both operands. `None` = no + /// valid-time filter. + valid_at_ms: Option, }, /// Fan-out a cell write batch, partitioned by routing bucket. @@ -109,18 +137,40 @@ pub enum ClusterArrayOp { /// Routing granularity. prefix_bits: u8, }, + + /// Fan-out a memtable flush to every primary vShard. Each shard's core + /// stamps the flush with its own WAL frontier. + Flush { + array_id: ArrayId, + /// Routing granularity. + prefix_bits: u8, + }, + + /// Fan-out a compaction to every primary vShard. + /// + /// `audit_retain_ms` is the array's retention window from the catalog. + /// `None` keeps every version. + Compact { + array_id: ArrayId, + audit_retain_ms: Option, + /// Routing granularity. + prefix_bits: u8, + }, } impl ClusterArrayOp { /// The array this op targets. Total over every variant, so read-set /// tracking and the commit validator key on the same name a shard-local - /// `ArrayOp` reports. + /// `ArrayOp` reports. For `Elementwise`, the left operand. pub fn array_id(&self) -> &ArrayId { match self { ClusterArrayOp::Slice { array_id, .. } | ClusterArrayOp::Agg { array_id, .. } | ClusterArrayOp::Put { array_id, .. } - | ClusterArrayOp::Delete { array_id, .. } => array_id, + | ClusterArrayOp::Delete { array_id, .. } + | ClusterArrayOp::Flush { array_id, .. } + | ClusterArrayOp::Compact { array_id, .. } => array_id, + ClusterArrayOp::Elementwise { left, .. } => left, } } } diff --git a/nodedb-physical/src/physical_plan/collection.rs b/nodedb-physical/src/physical_plan/collection.rs index 220213162..0215eb2ed 100644 --- a/nodedb-physical/src/physical_plan/collection.rs +++ b/nodedb-physical/src/physical_plan/collection.rs @@ -14,10 +14,10 @@ impl PhysicalPlan { /// Primary read/target collection this plan touches, if it maps to /// exactly one user collection. /// - /// Plane-neutral twin of the Control-Plane - /// `crate::control::server::shared::plan_util::extract_collection` — the - /// two MUST stay in sync (logic replicated, not called, since the core - /// crate depends on `nodedb-physical`, not the reverse). + /// The Control Plane's `extract_collection` answers a wider question: + /// which collection a plan reads or writes, for redaction, clone + /// interception, and read-set scoping. It reports a collection for more + /// variants than this accessor does. pub fn collection(&self) -> Option<&str> { match self { PhysicalPlan::Document(DocumentOp::PointGet { collection, .. }) diff --git a/nodedb-physical/src/physical_plan/query.rs b/nodedb-physical/src/physical_plan/query.rs index 317c90be5..134415cac 100644 --- a/nodedb-physical/src/physical_plan/query.rs +++ b/nodedb-physical/src/physical_plan/query.rs @@ -209,6 +209,13 @@ pub enum QueryOp { /// groups are finalized and HAVING is filtered. #[serde(default)] sort_keys: Vec, + /// SELECT-list output alias of each GROUP BY key, parallel to + /// `group_by`. A `Some` entry renames the key's `output_name` in every + /// result row, after HAVING, sort, and limit. HAVING and the sort keys + /// read the `output_name`. Empty, or a `None` entry, keeps the + /// `output_name`. + #[serde(default)] + group_key_aliases: Vec>, }, /// Partial aggregate: each core computes locally, Control Plane merges. diff --git a/nodedb-physical/src/physical_plan/spatial.rs b/nodedb-physical/src/physical_plan/spatial.rs index eca3968d5..735779650 100644 --- a/nodedb-physical/src/physical_plan/spatial.rs +++ b/nodedb-physical/src/physical_plan/spatial.rs @@ -2,7 +2,7 @@ //! Spatial engine operations dispatched to the Data Plane. -use nodedb_types::{QualifiedCollection, Surrogate, SurrogateBitmap, geometry::Geometry}; +use nodedb_types::{QualifiedCollection, Surrogate, SurrogateBitmap, geometry::SridGeometry}; /// Spatial predicate type for R-tree index scan. #[derive( @@ -50,8 +50,9 @@ pub enum SpatialOp { field: String, /// Stable global surrogate for the row, assigned on the Control Plane. surrogate: Surrogate, - /// Typed geometry, deserialised on the Control Plane from wire bytes. - geometry: Geometry, + /// Typed geometry and its SRID, deserialised on the Control Plane + /// from wire bytes. + geometry: SridGeometry, /// Sync provenance: identifies the originating peer and sequence for idempotency. #[serde(default)] provenance: Option, @@ -76,8 +77,9 @@ pub enum SpatialOp { collection: QualifiedCollection, field: String, predicate: SpatialPredicate, - /// Typed query geometry, parsed and validated on the Control Plane. - query_geometry: Geometry, + /// Typed query geometry and its SRID, parsed and validated on the + /// Control Plane. An SRID other than the index's fails the scan. + query_geometry: SridGeometry, /// Distance threshold in meters (for ST_DWithin). 0 for non-distance predicates. distance_meters: f64, /// Additional attribute filters applied after spatial candidates. diff --git a/nodedb-physical/src/physical_plan/timeseries.rs b/nodedb-physical/src/physical_plan/timeseries.rs index fceb106a9..dfcc8e981 100644 --- a/nodedb-physical/src/physical_plan/timeseries.rs +++ b/nodedb-physical/src/physical_plan/timeseries.rs @@ -71,6 +71,13 @@ pub enum TimeseriesOp { /// raw rows or to a downsample tier's buckets. #[serde(default)] tiering: Option, + /// The output columns of an aggregating scan, in SELECT-list order. + /// Each pair is `(encoded key, output name)`: the encoded key is + /// `bucket`, a GROUP BY column, or an aggregate's canonical key. Each + /// group row carries exactly these columns, renamed after HAVING, + /// sort, and limit. Empty keeps every encoded key under its own name. + #[serde(default)] + output_columns: Vec<(String, String)>, }, /// Write a batch of samples to the columnar memtable. diff --git a/nodedb-physical/src/physical_plan/vector/op.rs b/nodedb-physical/src/physical_plan/vector/op.rs index abd5df3c2..40f600f1c 100644 --- a/nodedb-physical/src/physical_plan/vector/op.rs +++ b/nodedb-physical/src/physical_plan/vector/op.rs @@ -228,6 +228,11 @@ pub enum VectorOp { query_entries: Vec<(u32, f32)>, /// Maximum results to return. top_k: usize, + /// The statement's residual `WHERE` (serialized `Vec`). + /// Empty when none applies. With a filter the Data Plane widens the + /// candidate window and attaches each body. The Control Plane keeps + /// the `top_k` highest-scoring hits that pass. + filters: Vec, }, /// Delete a document from the sparse inverted index. diff --git a/nodedb-query/src/expr/eval.rs b/nodedb-query/src/expr/eval.rs index fec2d9141..6d7ebe01d 100644 --- a/nodedb-query/src/expr/eval.rs +++ b/nodedb-query/src/expr/eval.rs @@ -26,7 +26,9 @@ use super::types::{SqlExpr, WHOLE_ROW_COLUMN}; /// - a function argument it cannot compute on: vectors of different /// dimensions, an argument of the wrong type, a malformed JSONPath. /// SQLSTATE `22000`. A `NULL` argument stays `NULL` instead; -/// - an exact integer SUM past the `Decimal` range, SQLSTATE `22000`. +/// - an exact integer SUM past the `Decimal` range, SQLSTATE `22000`; +/// - a geometry SRID fault: an unsupported SRID, operands in different +/// SRIDs, or a coordinate outside a projection's domain. SQLSTATE `22000`. #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] pub enum EvalError { #[error("division by zero")] @@ -66,6 +68,25 @@ pub enum EvalError { /// An exact integer aggregate total lies outside the `Decimal` range. #[error("{function}(): integer total out of range")] NumericOverflow { function: &'static str }, + /// An SRID argument, or an SRID a geometry declares, is outside the + /// supported set. SQLSTATE `22000`. + #[error( + "{function}(): unsupported SRID {srid}; supported: {}", + nodedb_types::geometry::SUPPORTED_SRID_CODES + )] + UnsupportedSrid { function: &'static str, srid: i64 }, + /// Two geometry operands are in different SRIDs. No implicit transform + /// runs. SQLSTATE `22000`. + #[error("{function}(): operands have different SRIDs {left} and {right}; use ST_Transform")] + MixedSrid { + function: &'static str, + left: u32, + right: u32, + }, + /// A coordinate lies outside the region the target projection is + /// defined on. SQLSTATE `22000`. + #[error("{function}(): coordinate is outside the domain of SRID {srid}")] + ProjectionDomain { function: &'static str, srid: u32 }, } /// Row scope for `SqlExpr::eval_scope`: how `Column(..)` and `OldColumn(..)` diff --git a/nodedb-query/src/functions/eval.rs b/nodedb-query/src/functions/eval.rs index cd6e0a920..601f6af6e 100644 --- a/nodedb-query/src/functions/eval.rs +++ b/nodedb-query/src/functions/eval.rs @@ -22,7 +22,7 @@ use super::{ /// - a name no family module implements returns /// `Err(EvalError::UnknownFunction)`, never a silent `NULL`. /// -/// The fallible families (`math`, `vector`, `json`) return +/// The fallible families (`math`, `vector`, `json`, geo) return /// `Option>`. The rest stay `Option`-shaped /// and are wrapped in `Ok` at this dispatch boundary. pub fn eval_function(name: &str, args: &[Value]) -> Result { @@ -62,10 +62,14 @@ pub fn eval_function(name: &str, args: &[Value]) -> Result { if let Some(v) = text_chunk::try_eval(name, args) { return Ok(v); } - // Geo / Spatial functions — delegated to geo_functions module. - crate::geo_functions::eval_geo_function(name, args).ok_or_else(|| EvalError::UnknownFunction { - name: name.to_owned(), - }) + // Geo / Spatial functions — delegated to geo_functions module. A geo + // fault (unsupported or mixed SRID, projection domain) fails the call. + match crate::geo_functions::eval_geo_function(name, args) { + Some(r) => r, + None => Err(EvalError::UnknownFunction { + name: name.to_owned(), + }), + } } #[cfg(test)] @@ -321,7 +325,7 @@ mod tests { fn point_value(lng: f64, lat: f64) -> Value { let geom = nodedb_types::geometry::Geometry::point(lng, lat); - Value::Geometry(geom) + Value::Geometry(nodedb_types::geometry::SridGeometry::wgs84(geom)) } fn square_value() -> Value { @@ -332,7 +336,7 @@ mod tests { [0.0, 10.0], [0.0, 0.0], ]]); - Value::Geometry(geom) + Value::Geometry(nodedb_types::geometry::SridGeometry::wgs84(geom)) } #[test] @@ -392,10 +396,9 @@ mod tests { #[test] fn geo_length_sql() { - let line = Value::Geometry(nodedb_types::geometry::Geometry::line_string(vec![ - [0.0, 0.0], - [0.0, 1.0], - ])); + let line = Value::Geometry(nodedb_types::geometry::SridGeometry::wgs84( + nodedb_types::geometry::Geometry::line_string(vec![[0.0, 0.0], [0.0, 1.0]]), + )); let result = eval_fn("geo_length", vec![line]); let d = result.as_f64().unwrap(); assert!((d - 111_195.0).abs() < 500.0, "got {d}"); diff --git a/nodedb-query/src/geo_functions/accessors.rs b/nodedb-query/src/geo_functions/accessors.rs index 73163da97..7c4f8f07a 100644 --- a/nodedb-query/src/geo_functions/accessors.rs +++ b/nodedb-query/src/geo_functions/accessors.rs @@ -1,37 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 //! Readers over a geometry's own contents: ordinates, type name, vertex count, -//! reference system, and the WKT / GeoJSON renderings. +//! and the WKT / GeoJSON renderings. The SRID reader lives in `crs`. use nodedb_types::Value; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; -use super::helpers::geom_arg; +use super::helpers::{GeoEval, geom_arg}; use crate::value_ops::to_value_number; -/// Geometry in NodeDB is stored as GeoJSON, whose coordinate reference system -/// is defined by RFC 7946 to be WGS 84 — EPSG:4326. There is no per-geometry -/// SRID to carry, so every stored geometry reports the same one. -const WGS84_SRID: i64 = 4326; - -pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { let result = match canonical { "st_x" => ordinate(args, 0), "st_y" => ordinate(args, 1), - "st_astext" => unary(args, |g| Value::String(nodedb_spatial::geometry_to_wkt(g))), + // WKT stays plain: it has no SRID slot. + "st_astext" => unary(args, |g| { + Value::String(nodedb_spatial::geometry_to_wkt(&g.geometry)) + }), + // GeoJSON carries the `crs` member for any SRID other than 4326. "st_asgeojson" => unary(args, |g| match sonic_rs::to_string(g) { Ok(json) => Value::String(json), Err(_) => Value::Null, }), - "st_geometrytype" => unary(args, |g| Value::String(g.geometry_type().to_string())), - "st_npoints" => unary(args, |g| Value::Integer(count_points(g) as i64)), - "st_srid" => unary(args, |_| Value::Integer(WGS84_SRID)), + "st_geometrytype" => unary(args, |g| { + Value::String(g.geometry.geometry_type().to_string()) + }), + "st_npoints" => unary(args, |g| Value::Integer(count_points(&g.geometry) as i64)), _ => return None, }; - Some(result) + Some(Ok(result)) } -fn unary(args: &[Value], f: fn(&Geometry) -> Value) -> Value { +fn unary(args: &[Value], f: fn(&SridGeometry) -> Value) -> Value { match geom_arg(args, 0) { Some(geom) => f(&geom), None => Value::Null, @@ -41,7 +41,7 @@ fn unary(args: &[Value], f: fn(&Geometry) -> Value) -> Value { /// Read one ordinate of a point. PostGIS defines `ST_X`/`ST_Y` on points only; /// any other geometry yields NULL rather than an arbitrary vertex. fn ordinate(args: &[Value], index: usize) -> Value { - match geom_arg(args, 0) { + match geom_arg(args, 0).map(|g| g.geometry) { Some(Geometry::Point { coordinates }) => to_value_number(coordinates[index]), _ => Value::Null, } @@ -72,8 +72,13 @@ fn count_points(geom: &Geometry) -> usize { mod tests { use super::*; + /// The evaluator with its infallible accessors unwrapped. + fn eval(canonical: &'static str, args: &[Value]) -> Option { + super::eval(canonical, args).map(|r| r.expect("accessors never fail")) + } + fn geom(v: Geometry) -> Vec { - vec![Value::Geometry(v)] + vec![Value::Geometry(SridGeometry::wgs84(v))] } #[test] @@ -125,14 +130,33 @@ mod tests { }; assert_eq!( super::super::helpers::geometry_from_text(&wkt), - Some(original.clone()) + Some(SridGeometry::wgs84(original.clone())) ); assert_eq!( super::super::helpers::geometry_from_text(&json), - Some(original) + Some(SridGeometry::wgs84(original)) ); } + /// GeoJSON keeps a projected SRID in its `crs` member; WKT stays plain. + #[test] + fn projected_renderings() { + let g = SridGeometry::new( + nodedb_types::geometry::Srid::Utm48N, + Geometry::point(1.0, 2.0), + ); + let args = vec![Value::Geometry(g.clone())]; + let Some(Value::String(json)) = eval("st_asgeojson", &args) else { + panic!("expected GeoJSON"); + }; + assert!(json.contains("EPSG:32648"), "{json}"); + assert_eq!(super::super::helpers::geometry_from_text(&json), Some(g)); + let Some(Value::String(wkt)) = eval("st_astext", &args) else { + panic!("expected WKT"); + }; + assert!(!wkt.contains("SRID"), "{wkt}"); + } + #[test] fn geometrytype_names_the_variant() { assert_eq!( @@ -151,20 +175,11 @@ mod tests { assert_eq!(eval("st_npoints", &line), Some(Value::Integer(3))); } - #[test] - fn srid_reports_wgs84() { - assert_eq!( - eval("st_srid", &geom(Geometry::point(0.0, 0.0))), - Some(Value::Integer(4326)) - ); - } - #[test] fn unreadable_geometry_yields_null() { let bad = vec![Value::String("not a geometry".into())]; assert_eq!(eval("st_astext", &bad), Some(Value::Null)); assert_eq!(eval("st_npoints", &bad), Some(Value::Null)); - assert_eq!(eval("st_srid", &bad), Some(Value::Null)); } /// A WKT string argument is a geometry, so accessors read it directly. diff --git a/nodedb-query/src/geo_functions/catalog/spec.rs b/nodedb-query/src/geo_functions/catalog/spec.rs index 0977a5e5a..5378abec4 100644 --- a/nodedb-query/src/geo_functions/catalog/spec.rs +++ b/nodedb-query/src/geo_functions/catalog/spec.rs @@ -16,6 +16,8 @@ pub enum GeoArgShape { Geometry2, /// `(geom1, geom2, distance_m)` Geometry2Distance, + /// `(geom, srid)` + GeometrySrid, /// `(geom, distance_m [, segments])` GeometryBuffer, /// `(lng, lat)` @@ -51,7 +53,7 @@ impl GeoArgShape { match self { Self::Geometry1 | Self::Text1 => (1, 1), Self::TextSrid | Self::BytesSrid => (1, 2), - Self::Geometry2 | Self::LngLat => (2, 2), + Self::Geometry2 | Self::LngLat | Self::GeometrySrid => (2, 2), Self::GeometryBuffer | Self::LngLatZ | Self::LngLatPrecision => (2, 3), Self::Geometry2Distance | Self::LatLngResolution => (3, 3), Self::Circle => (3, 4), diff --git a/nodedb-query/src/geo_functions/catalog/table.rs b/nodedb-query/src/geo_functions/catalog/table.rs index e75df2ff3..36a5bbcd5 100644 --- a/nodedb-query/src/geo_functions/catalog/table.rs +++ b/nodedb-query/src/geo_functions/catalog/table.rs @@ -42,7 +42,10 @@ pub static GEO_FUNCTIONS: &[GeoFunctionSpec] = &[ f("st_asgeojson", &["geo_as_geojson"], Geometry1, Text), f("st_geometrytype", &["geo_type"], Geometry1, Text), f("st_npoints", &["geo_num_points"], Geometry1, Int), + // ── Spatial reference systems ── f("st_srid", &[], Geometry1, Int), + f("st_setsrid", &[], GeometrySrid, Geometry), + f("st_transform", &[], GeometrySrid, Geometry), // ── Constructors ── f("st_point", &["geo_point"], LngLat, Geometry), f("st_makepoint", &[], LngLatZ, Geometry), diff --git a/nodedb-query/src/geo_functions/constructors.rs b/nodedb-query/src/geo_functions/constructors.rs index 36a37436e..7556088a0 100644 --- a/nodedb-query/src/geo_functions/constructors.rs +++ b/nodedb-query/src/geo_functions/constructors.rs @@ -1,69 +1,139 @@ // SPDX-License-Identifier: Apache-2.0 //! Geometry constructors and geometry-returning operations. +//! +//! A constructor from bare numbers builds a WGS 84 geometry. An operation on +//! geometries keeps its input SRID. A binary operation needs both inputs in +//! one SRID. use nodedb_types::Value; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, Srid, SridGeometry}; -use super::helpers::{geom_arg, geometry_from_text, num_arg, str_arg}; +use super::helpers::{ + GeoEval, SridArg, geom_arg, num_arg, parse_text_geometry, same_srid, srid_arg, str_arg, + text_fault, +}; +use crate::expr::EvalError; /// Default segment count for the circular approximations `ST_Buffer` and /// `geo_circle` produce when the caller does not specify one. const DEFAULT_BUFFER_SEGMENTS: usize = 32; -pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { let result = match canonical { - "st_point" => point(args), - "st_makepoint" => point(args), - "st_geomfromtext" | "st_geomfromgeojson" => match str_arg(args, 0) { - Some(text) => match geometry_from_text(&text) { - Some(geom) => Value::Geometry(geom), - None => Value::Null, - }, - None => Value::Null, - }, - "st_geomfromwkb" => match wkb_bytes(args) { - Some(bytes) => match nodedb_spatial::geometry_from_wkb(&bytes) { - Some(geom) => Value::Geometry(geom), - None => Value::Null, - }, - None => Value::Null, - }, - "st_makeline" => make_line(args), - "st_makepolygon" => make_polygon(args), - "st_makeenvelope" => make_envelope(args), - "geo_circle" => circle(args), - "st_buffer" => { - let Some(geom) = geom_arg(args, 0) else { - return Some(Value::Null); - }; - let Some(distance) = num_arg(args, 1) else { - return Some(Value::Null); - }; - let segments = segments_arg(args, 2); - Value::Geometry(nodedb_spatial::st_buffer(&geom, distance, segments)) - } - "st_envelope" => unary(args, |g| Some(nodedb_spatial::st_envelope(g))), - "st_centroid" => unary(args, nodedb_spatial::st_centroid), - "st_union" => binary(args, nodedb_spatial::st_union), - "st_intersection" => binary(args, nodedb_spatial::st_intersection), + "st_point" | "st_makepoint" => Ok(point(args)), + "st_geomfromtext" => from_text(canonical, args, true), + "st_geomfromgeojson" => from_text(canonical, args, false), + "st_geomfromwkb" => from_wkb(canonical, args), + "st_makeline" => make_line(canonical, args), + "st_makepolygon" => Ok(make_polygon(args)), + "st_makeenvelope" => Ok(make_envelope(args)), + "geo_circle" => Ok(circle(args)), + "st_buffer" => Ok(buffer(args)), + "st_envelope" => Ok(unary(args, |g| { + Some(nodedb_spatial::st_envelope_in( + &g.geometry, + g.srid.bbox_space(), + )) + })), + "st_centroid" => Ok(unary(args, |g| { + nodedb_spatial::st_centroid_in(&g.geometry, g.srid.bbox_space()) + })), + "st_union" => binary(canonical, args, nodedb_spatial::st_union), + "st_intersection" => binary(canonical, args, nodedb_spatial::st_intersection), _ => return None, }; Some(result) } -fn unary(args: &[Value], f: fn(&Geometry) -> Option) -> Value { - match geom_arg(args, 0).as_ref().and_then(f) { - Some(geom) => Value::Geometry(geom), +/// A unary geometry operation whose result keeps the input SRID. +fn unary(args: &[Value], f: fn(&SridGeometry) -> Option) -> Value { + match geom_arg(args, 0) { + Some(g) => match f(&g) { + Some(out) => Value::Geometry(g.with_geometry(out)), + None => Value::Null, + }, None => Value::Null, } } -fn binary(args: &[Value], f: fn(&Geometry, &Geometry) -> Geometry) -> Value { +/// A binary geometry operation over two inputs in one SRID. +fn binary( + function: &'static str, + args: &[Value], + f: fn(&Geometry, &Geometry) -> Geometry, +) -> Result { let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else { + return Ok(Value::Null); + }; + let srid = same_srid(function, &a, &b)?; + Ok(Value::Geometry(SridGeometry::new( + srid, + f(&a.geometry, &b.geometry), + ))) +} + +/// `ST_GeomFromText(text [, srid])` and `ST_GeomFromGeoJSON(text)`. A GeoJSON +/// text keeps the SRID its `crs` member names. An explicit SRID argument +/// relabels the result. +fn from_text(function: &'static str, args: &[Value], takes_srid: bool) -> Result { + let explicit = if takes_srid { + srid_arg(args, 1, function)? + } else { + SridArg::Absent + }; + let Some(text) = str_arg(args, 0) else { + return Ok(Value::Null); + }; + let mut geom = match parse_text_geometry(&text) { + Ok(g) => g, + Err(e) => return text_fault(function, e), + }; + match explicit { + SridArg::Absent => {} + SridArg::Null => return Ok(Value::Null), + SridArg::Srid(srid) => geom.srid = srid, + } + Ok(Value::Geometry(geom)) +} + +/// `ST_GeomFromWKB(wkb [, srid])`. An explicit SRID argument wins over the +/// SRID an EWKB header carries. ISO WKB with no SRID argument is WGS 84. +fn from_wkb(function: &'static str, args: &[Value]) -> Result { + let explicit = srid_arg(args, 1, function)?; + let Some(bytes) = wkb_bytes(args) else { + return Ok(Value::Null); + }; + let Some((geometry, header_srid)) = nodedb_spatial::geometry_from_wkb(&bytes) else { + return Ok(Value::Null); + }; + let srid = match (explicit, header_srid) { + (SridArg::Null, _) => return Ok(Value::Null), + (SridArg::Srid(srid), _) => srid, + (SridArg::Absent, Some(code)) => { + Srid::from_code(i64::from(code)).map_err(|e| EvalError::UnsupportedSrid { + function, + srid: e.code, + })? + } + (SridArg::Absent, None) => Srid::Wgs84, + }; + Ok(Value::Geometry(SridGeometry::new(srid, geometry))) +} + +/// `ST_Buffer(geom, distance [, segments])`. Distance is metres in either +/// space: geodesic-approximate for 4326, planar for a projected SRID. +fn buffer(args: &[Value]) -> Value { + let (Some(g), Some(distance)) = (geom_arg(args, 0), num_arg(args, 1)) else { return Value::Null; }; - Value::Geometry(f(&a, &b)) + let segments = segments_arg(args, 2); + let out = if g.srid.is_geographic() { + nodedb_spatial::st_buffer(&g.geometry, distance, segments) + } else { + nodedb_spatial::planar_buffer(&g.geometry, distance, segments) + }; + Value::Geometry(g.with_geometry(out)) } /// Segment count for a circular approximation, floored at 3 — fewer than three @@ -100,7 +170,7 @@ fn point(args: &[Value]) -> Value { let (Some(lng), Some(lat)) = (num_arg(args, 0), num_arg(args, 1)) else { return Value::Null; }; - Value::Geometry(Geometry::point(lng, lat)) + Value::Geometry(SridGeometry::wgs84(Geometry::point(lng, lat))) } fn circle(args: &[Value]) -> Value { @@ -110,28 +180,41 @@ fn circle(args: &[Value]) -> Value { return Value::Null; }; let segments = segments_arg(args, 3); - Value::Geometry(nodedb_spatial::st_buffer( + Value::Geometry(SridGeometry::wgs84(nodedb_spatial::st_buffer( &Geometry::point(lng, lat), radius, segments, - )) + ))) } /// `ST_MakeLine(point, point, ...)` — every argument must be a readable point; /// a single unreadable one makes the whole line NULL rather than silently -/// dropping a vertex and shortening the line. -fn make_line(args: &[Value]) -> Value { +/// dropping a vertex and shortening the line. Every point must share one +/// SRID, which the line keeps. +fn make_line(function: &'static str, args: &[Value]) -> Result { let mut coords = Vec::with_capacity(args.len()); + let mut first: Option = None; for idx in 0..args.len() { - match geom_arg(args, idx) { - Some(Geometry::Point { coordinates }) => coords.push(coordinates), - _ => return Value::Null, + let Some(g) = geom_arg(args, idx) else { + return Ok(Value::Null); + }; + let Geometry::Point { coordinates } = g.geometry else { + return Ok(Value::Null); + }; + if let Some(head) = &first { + same_srid(function, head, &g)?; } + coords.push(coordinates); + first.get_or_insert(g); } if coords.len() < 2 { - return Value::Null; + return Ok(Value::Null); } - Value::Geometry(Geometry::line_string(coords)) + let srid = first.map_or(Srid::Wgs84, |g| g.srid); + Ok(Value::Geometry(SridGeometry::new( + srid, + Geometry::line_string(coords), + ))) } /// `ST_MakePolygon(ring, ...)` — each argument is an array of coordinate @@ -163,7 +246,7 @@ fn make_polygon(args: &[Value]) -> Value { if rings.is_empty() { return Value::Null; } - Value::Geometry(Geometry::polygon(rings)) + Value::Geometry(SridGeometry::wgs84(Geometry::polygon(rings))) } fn make_envelope(args: &[Value]) -> Value { @@ -175,25 +258,38 @@ fn make_envelope(args: &[Value]) -> Value { ) else { return Value::Null; }; - Value::Geometry(Geometry::polygon(vec![vec![ + Value::Geometry(SridGeometry::wgs84(Geometry::polygon(vec![vec![ [min_lng, min_lat], [max_lng, min_lat], [max_lng, max_lat], [min_lng, max_lat], [min_lng, min_lat], - ]])) + ]]))) } #[cfg(test)] mod tests { use super::*; + /// The evaluator, with a fault turned into a test failure. + fn run(canonical: &'static str, args: &[Value]) -> Option { + eval(canonical, args).map(|r| r.expect("no fault expected")) + } + + fn wgs84(g: Geometry) -> Value { + Value::Geometry(SridGeometry::wgs84(g)) + } + + fn projected(srid: Srid, g: Geometry) -> Value { + Value::Geometry(SridGeometry::new(srid, g)) + } + #[test] fn point_constructors_agree() { let args = vec![Value::Float(1.0), Value::Float(2.0)]; - let expected = Some(Value::Geometry(Geometry::point(1.0, 2.0))); - assert_eq!(eval("st_point", &args), expected); - assert_eq!(eval("st_makepoint", &args), expected); + let expected = Some(wgs84(Geometry::point(1.0, 2.0))); + assert_eq!(run("st_point", &args), expected); + assert_eq!(run("st_makepoint", &args), expected); } #[test] @@ -202,9 +298,30 @@ mod tests { let json = vec![Value::String( r#"{"type":"Point","coordinates":[1.0,2.0]}"#.into(), )]; - let expected = Some(Value::Geometry(Geometry::point(1.0, 2.0))); - assert_eq!(eval("st_geomfromtext", &wkt), expected); - assert_eq!(eval("st_geomfromgeojson", &json), expected); + let expected = Some(wgs84(Geometry::point(1.0, 2.0))); + assert_eq!(run("st_geomfromtext", &wkt), expected); + assert_eq!(run("st_geomfromgeojson", &json), expected); + } + + #[test] + fn geomfromtext_srid_argument_sets_srid() { + let args = vec![Value::String("POINT(1 2)".into()), Value::Integer(3375)]; + assert_eq!( + run("st_geomfromtext", &args), + Some(projected(Srid::PeninsularRso, Geometry::point(1.0, 2.0))) + ); + } + + #[test] + fn geomfromtext_unsupported_srid_fails() { + let args = vec![Value::String("POINT(1 2)".into()), Value::Integer(2154)]; + assert_eq!( + eval("st_geomfromtext", &args), + Some(Err(EvalError::UnsupportedSrid { + function: "st_geomfromtext", + srid: 2154 + })) + ); } /// A byte literal and its hex-string spelling must decode identically. @@ -213,20 +330,37 @@ mod tests { // WKB for POINT(1 2): little-endian, type 1, x=1.0, y=2.0. let hex = "0101000000000000000000F03F0000000000000040"; let bytes = decode_hex(hex).expect("test vector must decode"); - let from_bytes = eval("st_geomfromwkb", &[Value::Bytes(bytes)]); - let from_hex = eval("st_geomfromwkb", &[Value::String(hex.into())]); + let from_bytes = run("st_geomfromwkb", &[Value::Bytes(bytes)]); + let from_hex = run("st_geomfromwkb", &[Value::String(hex.into())]); assert_eq!(from_bytes, from_hex); - assert_eq!(from_bytes, Some(Value::Geometry(Geometry::point(1.0, 2.0)))); + assert_eq!(from_bytes, Some(wgs84(Geometry::point(1.0, 2.0)))); + } + + #[test] + fn geomfromwkb_reads_ewkb_srid_and_explicit_srid_wins() { + let g = SridGeometry::new(Srid::Utm48N, Geometry::point(1.0, 2.0)); + let ewkb = nodedb_spatial::geometry_to_ewkb(&g).expect("encode"); + assert_eq!( + run("st_geomfromwkb", &[Value::Bytes(ewkb.clone())]), + Some(Value::Geometry(g)) + ); + assert_eq!( + run( + "st_geomfromwkb", + &[Value::Bytes(ewkb), Value::Integer(3376)] + ), + Some(projected(Srid::BorneoRso, Geometry::point(1.0, 2.0))) + ); } #[test] fn malformed_hex_yields_null() { assert_eq!( - eval("st_geomfromwkb", &[Value::String("zz".into())]), + run("st_geomfromwkb", &[Value::String("zz".into())]), Some(Value::Null) ); assert_eq!( - eval("st_geomfromwkb", &[Value::String("010".into())]), + run("st_geomfromwkb", &[Value::String("010".into())]), Some(Value::Null) ); } @@ -234,7 +368,7 @@ mod tests { #[test] fn malformed_text_yields_null() { let bad = vec![Value::String("POINT(".into())]; - assert_eq!(eval("st_geomfromtext", &bad), Some(Value::Null)); + assert_eq!(run("st_geomfromtext", &bad), Some(Value::Null)); } /// A dropped vertex would silently shorten the line; the call must fail @@ -242,16 +376,32 @@ mod tests { #[test] fn makeline_rejects_a_non_point_argument() { let args = vec![ - Value::Geometry(Geometry::point(0.0, 0.0)), + wgs84(Geometry::point(0.0, 0.0)), Value::String("not a geometry".into()), ]; - assert_eq!(eval("st_makeline", &args), Some(Value::Null)); + assert_eq!(run("st_makeline", &args), Some(Value::Null)); } #[test] fn makeline_needs_at_least_two_points() { - let one = vec![Value::Geometry(Geometry::point(0.0, 0.0))]; - assert_eq!(eval("st_makeline", &one), Some(Value::Null)); + let one = vec![wgs84(Geometry::point(0.0, 0.0))]; + assert_eq!(run("st_makeline", &one), Some(Value::Null)); + } + + #[test] + fn makeline_rejects_mixed_srids() { + let args = vec![ + wgs84(Geometry::point(0.0, 0.0)), + projected(Srid::Utm48N, Geometry::point(1.0, 1.0)), + ]; + assert_eq!( + eval("st_makeline", &args), + Some(Err(EvalError::MixedSrid { + function: "st_makeline", + left: 4326, + right: 32648 + })) + ); } #[test] @@ -262,8 +412,10 @@ mod tests { Value::Float(1.0), Value::Float(1.0), ]; - let Some(Value::Geometry(Geometry::Polygon { coordinates })) = - eval("st_makeenvelope", &args) + let Some(Value::Geometry(SridGeometry { + geometry: Geometry::Polygon { coordinates }, + .. + })) = run("st_makeenvelope", &args) else { panic!("expected a polygon"); }; @@ -276,11 +428,14 @@ mod tests { #[test] fn buffer_segment_count_is_floored() { let args = vec![ - Value::Geometry(Geometry::point(0.0, 0.0)), + wgs84(Geometry::point(0.0, 0.0)), Value::Float(1000.0), Value::Integer(0), ]; - let Some(Value::Geometry(Geometry::Polygon { coordinates })) = eval("st_buffer", &args) + let Some(Value::Geometry(SridGeometry { + geometry: Geometry::Polygon { coordinates }, + .. + })) = run("st_buffer", &args) else { panic!("expected a buffered polygon"); }; @@ -290,10 +445,48 @@ mod tests { ); } + #[test] + fn projected_buffer_is_planar_and_keeps_srid() { + let args = vec![ + projected(Srid::Utm48N, Geometry::point(500_000.0, 100_000.0)), + Value::Float(100.0), + ]; + let Some(Value::Geometry(g)) = run("st_buffer", &args) else { + panic!("expected a geometry"); + }; + assert_eq!(g.srid, Srid::Utm48N); + let Geometry::Polygon { coordinates } = &g.geometry else { + panic!("expected a polygon"); + }; + let [x, y] = coordinates[0][0]; + assert!(((x - 500_000.0).hypot(y - 100_000.0) - 100.0).abs() < 1e-6); + } + + #[test] + fn union_rejects_mixed_srids() { + let args = vec![ + wgs84(Geometry::point(0.0, 0.0)), + projected(Srid::PeninsularRso, Geometry::point(1.0, 1.0)), + ]; + assert!(matches!( + eval("st_union", &args), + Some(Err(EvalError::MixedSrid { .. })) + )); + } + + #[test] + fn centroid_keeps_srid() { + let args = vec![projected(Srid::BorneoRso, Geometry::point(5.0, 6.0))]; + assert_eq!( + run("st_centroid", &args), + Some(projected(Srid::BorneoRso, Geometry::point(5.0, 6.0))) + ); + } + #[test] fn centroid_of_unreadable_geometry_is_null() { let bad = vec![Value::String("not a geometry".into())]; - assert_eq!(eval("st_centroid", &bad), Some(Value::Null)); + assert_eq!(run("st_centroid", &bad), Some(Value::Null)); } #[test] diff --git a/nodedb-query/src/geo_functions/crs.rs b/nodedb-query/src/geo_functions/crs.rs new file mode 100644 index 000000000..825c4417a --- /dev/null +++ b/nodedb-query/src/geo_functions/crs.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Spatial reference system functions: read, set, and transform a +//! geometry's SRID. + +use nodedb_spatial::CrsError; +use nodedb_types::Value; + +use super::helpers::{GeoEval, SridArg, geom_arg, srid_arg}; +use crate::expr::EvalError; + +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { + let result = match canonical { + "st_srid" => Ok(match geom_arg(args, 0) { + Some(g) => Value::Integer(i64::from(g.srid.code())), + None => Value::Null, + }), + "st_setsrid" => set_srid(canonical, args), + "st_transform" => transform(canonical, args), + _ => return None, + }; + Some(result) +} + +/// `ST_SetSRID(geom, srid)`: relabel the geometry's SRID. Coordinates are +/// kept as they are. +fn set_srid(function: &'static str, args: &[Value]) -> Result { + let srid = match srid_arg(args, 1, function)? { + SridArg::Srid(srid) => srid, + SridArg::Absent | SridArg::Null => return Ok(Value::Null), + }; + Ok(match geom_arg(args, 0) { + Some(mut g) => { + g.srid = srid; + Value::Geometry(g) + } + None => Value::Null, + }) +} + +/// `ST_Transform(geom, srid)`: reproject every coordinate into `srid`. +fn transform(function: &'static str, args: &[Value]) -> Result { + let target = match srid_arg(args, 1, function)? { + SridArg::Srid(srid) => srid, + SridArg::Absent | SridArg::Null => return Ok(Value::Null), + }; + let Some(g) = geom_arg(args, 0) else { + return Ok(Value::Null); + }; + nodedb_spatial::transform_geometry(&g, target) + .map(Value::Geometry) + .map_err(|e| projection_fault(function, e)) +} + +/// A projection that cannot map a coordinate fails the statement. +pub(super) fn projection_fault(function: &'static str, e: CrsError) -> EvalError { + EvalError::ProjectionDomain { + function, + srid: e.srid().code(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_types::geometry::{Geometry, Srid, SridGeometry}; + + fn point(srid: Srid, x: f64, y: f64) -> Value { + Value::Geometry(SridGeometry::new(srid, Geometry::point(x, y))) + } + + #[test] + fn srid_reads_the_geometry_srid() { + assert_eq!( + eval("st_srid", &[point(Srid::Wgs84, 0.0, 0.0)]), + Some(Ok(Value::Integer(4326))) + ); + assert_eq!( + eval("st_srid", &[point(Srid::BorneoRso, 0.0, 0.0)]), + Some(Ok(Value::Integer(3376))) + ); + assert_eq!( + eval("st_srid", &[Value::String("not a geometry".into())]), + Some(Ok(Value::Null)) + ); + } + + #[test] + fn setsrid_relabels_without_moving() { + let got = eval( + "st_setsrid", + &[point(Srid::Wgs84, 1.0, 2.0), Value::Integer(32648)], + ); + assert_eq!(got, Some(Ok(point(Srid::Utm48N, 1.0, 2.0)))); + } + + #[test] + fn transform_projects_klcc_into_utm47n() { + let Some(Ok(Value::Geometry(g))) = eval( + "st_transform", + &[point(Srid::Wgs84, 101.7117, 3.1579), Value::Integer(32647)], + ) else { + panic!("expected a geometry"); + }; + assert_eq!(g.srid, Srid::Utm47N); + let Geometry::Point { coordinates } = g.geometry else { + panic!("expected a point"); + }; + assert!((coordinates[0] - 801_401.794_9).abs() < 1e-3); + assert!((coordinates[1] - 349_439.510_6).abs() < 1e-3); + } + + #[test] + fn transform_to_unsupported_srid_names_every_supported_code() { + let Some(Err(err)) = eval( + "st_transform", + &[point(Srid::Wgs84, 101.7, 3.1), Value::Integer(2154)], + ) else { + panic!("expected an error"); + }; + assert_eq!( + err, + EvalError::UnsupportedSrid { + function: "st_transform", + srid: 2154 + } + ); + let text = err.to_string(); + for code in ["4326", "32647", "32648", "3375", "3376"] { + assert!(text.contains(code), "{text}"); + } + } + + #[test] + fn transform_out_of_domain_is_a_projection_fault() { + let got = eval( + "st_transform", + &[point(Srid::Wgs84, 179.0, 3.0), Value::Integer(32648)], + ); + assert_eq!( + got, + Some(Err(EvalError::ProjectionDomain { + function: "st_transform", + srid: 32648 + })) + ); + } + + #[test] + fn null_srid_yields_null() { + assert_eq!( + eval("st_setsrid", &[point(Srid::Wgs84, 1.0, 2.0), Value::Null]), + Some(Ok(Value::Null)) + ); + } +} diff --git a/nodedb-query/src/geo_functions/dispatch.rs b/nodedb-query/src/geo_functions/dispatch.rs index 08a830ec8..a1ee8935c 100644 --- a/nodedb-query/src/geo_functions/dispatch.rs +++ b/nodedb-query/src/geo_functions/dispatch.rs @@ -5,7 +5,8 @@ use nodedb_types::Value; use super::catalog; -use super::{accessors, constructors, indexing, measures, predicates}; +use super::{accessors, constructors, crs, indexing, measures, predicates}; +use crate::expr::EvalError; /// Evaluate a geospatial function. /// @@ -14,8 +15,10 @@ use super::{accessors, constructors, indexing, measures, predicates}; /// before dispatch, so every spelling of a capability behaves identically. /// /// Returns `None` only when `name` is not a geospatial function at all, so the -/// caller can fall through to its own function table. -pub fn eval_geo_function(name: &str, args: &[Value]) -> Option { +/// caller can fall through to its own function table. `Some(Err(_))` is a +/// fault that fails the statement: an unsupported SRID, operands in different +/// SRIDs, or a coordinate outside a projection's domain. +pub fn eval_geo_function(name: &str, args: &[Value]) -> Option> { let spec = catalog::lookup(name)?; eval_canonical(spec.canonical, args) } @@ -25,10 +28,11 @@ pub fn eval_geo_function(name: &str, args: &[Value]) -> Option { /// A `None` here means the catalog declares a function the evaluators do not /// implement. That drift is caught by `every_catalog_entry_is_evaluable` /// below rather than being allowed to surface as a silent NULL at runtime. -fn eval_canonical(canonical: &str, args: &[Value]) -> Option { +fn eval_canonical(canonical: &'static str, args: &[Value]) -> Option> { predicates::eval(canonical, args) .or_else(|| measures::eval(canonical, args)) .or_else(|| accessors::eval(canonical, args)) + .or_else(|| crs::eval(canonical, args)) .or_else(|| constructors::eval(canonical, args)) .or_else(|| indexing::eval(canonical, args)) } @@ -36,14 +40,14 @@ fn eval_canonical(canonical: &str, args: &[Value]) -> Option { #[cfg(test)] mod tests { use super::*; - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, SridGeometry}; /// Arguments that satisfy any catalog shape: enough geometry-shaped and /// number-shaped values that no evaluator indexes past the end. fn probe_args() -> Vec { vec![ - Value::Geometry(Geometry::point(1.0, 2.0)), - Value::Geometry(Geometry::point(3.0, 4.0)), + Value::Geometry(SridGeometry::wgs84(Geometry::point(1.0, 2.0))), + Value::Geometry(SridGeometry::wgs84(Geometry::point(3.0, 4.0))), Value::Float(5.0), Value::Float(6.0), ] @@ -101,18 +105,38 @@ mod tests { #[test] fn standard_accessor_names_read_stored_geometry() { use crate::value_ops::to_value_number; - let point = vec![Value::Geometry(Geometry::point(1.0, 2.0))]; + let point = vec![Value::Geometry(SridGeometry::wgs84(Geometry::point( + 1.0, 2.0, + )))]; assert_eq!( eval_geo_function("st_x", &point), - Some(to_value_number(1.0)) + Some(Ok(to_value_number(1.0))) ); assert_eq!( eval_geo_function("st_y", &point), - Some(to_value_number(2.0)) + Some(Ok(to_value_number(2.0))) ); assert!(matches!( eval_geo_function("st_astext", &point), - Some(Value::String(_)) + Some(Ok(Value::String(_))) )); } + + /// `ST_SRID(ST_SetSRID(ST_Point(1, 2), 32648))` is 32648. + #[test] + fn setsrid_then_srid_round_trips() { + let Some(Ok(point)) = + eval_geo_function("st_point", &[Value::Float(1.0), Value::Float(2.0)]) + else { + panic!("expected a point"); + }; + let Some(Ok(tagged)) = eval_geo_function("ST_SetSRID", &[point, Value::Integer(32648)]) + else { + panic!("expected a geometry"); + }; + assert_eq!( + eval_geo_function("ST_SRID", &[tagged]), + Some(Ok(Value::Integer(32648))) + ); + } } diff --git a/nodedb-query/src/geo_functions/helpers.rs b/nodedb-query/src/geo_functions/helpers.rs index 31e6febd9..06c7687e1 100644 --- a/nodedb-query/src/geo_functions/helpers.rs +++ b/nodedb-query/src/geo_functions/helpers.rs @@ -2,9 +2,14 @@ //! Argument coercion shared by the geo function evaluators. +use crate::expr::EvalError; use crate::value_ops::value_to_f64; use nodedb_types::Value; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{GeoJsonError, Srid, SridGeometry, parse_geojson}; + +/// What a geo family evaluator returns: `None` when the canonical name is +/// not its own, else the call's value or the fault that fails the statement. +pub(super) type GeoEval = Option>; pub(super) fn str_arg(args: &[Value], idx: usize) -> Option { args.get(idx)?.as_str().map(|s| s.to_string()) @@ -22,15 +27,13 @@ pub(super) fn num_arg(args: &[Value], idx: usize) -> Option { /// as), and a WKT string. Accepting WKT here is what makes a bare /// `'POINT(1 2)'` literal usable wherever a geometry is expected, matching /// PostGIS's treatment of an unknown-typed literal in geometry position. -pub(super) fn geom_arg(args: &[Value], idx: usize) -> Option { +/// GeoJSON keeps the SRID its `crs` member names. WKT is WGS 84. +pub(super) fn geom_arg(args: &[Value], idx: usize) -> Option { match args.get(idx)? { Value::Geometry(g) => Some(g.clone()), Value::String(s) => geometry_from_text(s), - other => { - // GeoJSON object: convert Value → JSON → Geometry. - let json = serde_json::Value::from(other.clone()); - serde_json::from_value(json).ok() - } + obj @ Value::Object(_) => SridGeometry::from_value(obj), + _ => None, } } @@ -38,25 +41,123 @@ pub(super) fn geom_arg(args: &[Value], idx: usize) -> Option { /// /// The two are unambiguous — GeoJSON always starts with `{`, WKT with a type /// keyword — so trying them in order cannot misread one as the other. -pub fn geometry_from_text(text: &str) -> Option { +pub fn geometry_from_text(text: &str) -> Option { + parse_text_geometry(text).ok() +} + +/// [`geometry_from_text`] that reports why a text is not a geometry. A +/// GeoJSON `crs` naming an unsupported SRID gives +/// [`GeoJsonError::UnsupportedSrid`]. Any other unreadable text gives +/// [`GeoJsonError::Malformed`]. +pub(super) fn parse_text_geometry(text: &str) -> Result { let trimmed = text.trim(); if trimmed.starts_with('{') { - nodedb_types::geometry::from_geojson_str(trimmed) + parse_geojson(trimmed) } else { nodedb_spatial::geometry_from_wkt(trimmed) + .map(SridGeometry::wgs84) + .ok_or(GeoJsonError::Malformed) + } +} + +/// An optional SRID argument. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum SridArg { + /// The call has no argument at this position. + Absent, + /// The argument is SQL NULL. + Null, + Srid(Srid), +} + +/// Read the SRID argument at `idx`. +/// +/// An integer, an integral float, or integer text resolves through +/// [`Srid::from_code`]. A code outside the supported set fails with +/// [`EvalError::UnsupportedSrid`]. A value of any other kind fails with +/// [`EvalError::ArgumentType`]. +pub(super) fn srid_arg( + args: &[Value], + idx: usize, + function: &'static str, +) -> Result { + let Some(value) = args.get(idx) else { + return Ok(SridArg::Absent); + }; + let code = match value { + Value::Null => return Ok(SridArg::Null), + Value::Integer(i) => Some(*i), + Value::Float(f) if f.is_finite() && f.fract() == 0.0 => Some(*f as i64), + Value::String(s) => s.trim().parse::().ok(), + _ => None, + }; + let Some(code) = code else { + return Err(EvalError::ArgumentType { + function, + position: idx + 1, + expected: "an integer SRID", + got: value.type_name(), + }); + }; + Srid::from_code(code) + .map(SridArg::Srid) + .map_err(|e| EvalError::UnsupportedSrid { + function, + srid: e.code, + }) +} + +/// The shared SRID of two operands, or [`EvalError::MixedSrid`]. +pub(super) fn same_srid( + function: &'static str, + a: &SridGeometry, + b: &SridGeometry, +) -> Result { + if a.srid == b.srid { + Ok(a.srid) + } else { + Err(EvalError::MixedSrid { + function, + left: a.srid.code(), + right: b.srid.code(), + }) + } +} + +/// Map a GeoJSON parse fault: an unsupported SRID fails the statement, +/// anything else is an unreadable argument and yields NULL. +pub(super) fn text_fault(function: &'static str, e: GeoJsonError) -> Result { + match e { + GeoJsonError::UnsupportedSrid(u) => Err(EvalError::UnsupportedSrid { + function, + srid: u.code, + }), + GeoJsonError::Malformed | GeoJsonError::UnknownCrsName { .. } => Ok(Value::Null), } } #[cfg(test)] mod tests { use super::*; + use nodedb_types::geometry::Geometry; #[test] fn geojson_and_wkt_strings_both_parse() { let geojson = geometry_from_text(r#"{"type":"Point","coordinates":[1.0,2.0]}"#); let wkt = geometry_from_text("POINT(1 2)"); assert_eq!(geojson, wkt); - assert_eq!(geojson, Some(Geometry::point(1.0, 2.0))); + assert_eq!( + geojson, + Some(SridGeometry::wgs84(Geometry::point(1.0, 2.0))) + ); + } + + #[test] + fn geojson_crs_member_sets_srid() { + let g = geometry_from_text( + r#"{"type":"Point","coordinates":[1.0,2.0],"crs":{"type":"name","properties":{"name":"EPSG:3375"}}}"#, + ); + assert_eq!(g.map(|g| g.srid), Some(Srid::PeninsularRso)); } #[test] @@ -70,4 +171,34 @@ mod tests { assert!(geometry_from_text("not a geometry").is_none()); assert!(geometry_from_text("").is_none()); } + + #[test] + fn srid_arg_forms() { + let f = "st_setsrid"; + let args = [ + Value::Integer(32648), + Value::Float(3375.0), + Value::String("3376".into()), + Value::Null, + ]; + assert_eq!(srid_arg(&args, 0, f), Ok(SridArg::Srid(Srid::Utm48N))); + assert_eq!( + srid_arg(&args, 1, f), + Ok(SridArg::Srid(Srid::PeninsularRso)) + ); + assert_eq!(srid_arg(&args, 2, f), Ok(SridArg::Srid(Srid::BorneoRso))); + assert_eq!(srid_arg(&args, 3, f), Ok(SridArg::Null)); + assert_eq!(srid_arg(&args, 4, f), Ok(SridArg::Absent)); + assert_eq!( + srid_arg(&[Value::Integer(2154)], 0, f), + Err(EvalError::UnsupportedSrid { + function: f, + srid: 2154 + }) + ); + assert!(matches!( + srid_arg(&[Value::Bool(true)], 0, f), + Err(EvalError::ArgumentType { .. }) + )); + } } diff --git a/nodedb-query/src/geo_functions/indexing.rs b/nodedb-query/src/geo_functions/indexing.rs index b40c7ce25..d18915470 100644 --- a/nodedb-query/src/geo_functions/indexing.rs +++ b/nodedb-query/src/geo_functions/indexing.rs @@ -5,8 +5,9 @@ use std::collections::HashMap; use nodedb_types::Value; +use nodedb_types::geometry::SridGeometry; -use super::helpers::{num_arg, str_arg}; +use super::helpers::{GeoEval, num_arg, str_arg}; /// Geohash precision used when the caller does not specify one — ~1.2 km /// cells, the conventional default for point-of-interest bucketing. @@ -15,7 +16,13 @@ const DEFAULT_GEOHASH_PRECISION: u8 = 6; /// H3 resolution used when the caller does not specify one — ~5 km hexes. const DEFAULT_H3_RESOLUTION: u8 = 7; -pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { +/// Geohash and H3 functions never fail the statement: an unreadable input +/// yields NULL. +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { + eval_infallible(canonical, args).map(Ok) +} + +fn eval_infallible(canonical: &str, args: &[Value]) -> Option { let result = match canonical { "st_geohash" | "geo_h3" => { let (Some(lng), Some(lat)) = (num_arg(args, 0), num_arg(args, 1)) else { @@ -83,7 +90,7 @@ pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { }, "geo_h3_to_boundary" => match h3_index(args) { Some(index) => match nodedb_spatial::h3::h3_to_boundary(index) { - Some(geom) => Value::Geometry(geom), + Some(geom) => Value::Geometry(SridGeometry::wgs84(geom)), None => Value::Null, }, None => Value::Null, @@ -128,6 +135,10 @@ fn precision_arg(args: &[Value], idx: usize, default: u8) -> u8 { mod tests { use super::*; + fn eval(canonical: &str, args: &[Value]) -> Option { + eval_infallible(canonical, args) + } + #[test] fn geohash_encodes_and_decodes() { let args = vec![Value::Float(-122.4), Value::Float(37.8), Value::Integer(6)]; diff --git a/nodedb-query/src/geo_functions/measures.rs b/nodedb-query/src/geo_functions/measures.rs index 1c5c16403..148ef8bfc 100644 --- a/nodedb-query/src/geo_functions/measures.rs +++ b/nodedb-query/src/geo_functions/measures.rs @@ -2,36 +2,55 @@ //! Distance, length, perimeter, area, and bearing. //! -//! Every measure is geodesic and reported in meters (or square meters), the -//! same basis as `nodedb_spatial::st_distance`. +//! Every measure is reported in metres (or square metres). A WGS 84 geometry +//! is measured geodesically, the same basis as `nodedb_spatial::st_distance`. +//! A projected geometry is measured in planar grid metres. use nodedb_types::Value; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; -use super::helpers::{geom_arg, num_arg}; +use super::helpers::{GeoEval, geom_arg, num_arg, same_srid}; use crate::value_ops::to_value_number; -pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { let result = match canonical { "st_distance" => { let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else { - return Some(Value::Null); + return Some(Ok(Value::Null)); }; - to_value_number(nodedb_spatial::st_distance(&a, &b)) + let srid = match same_srid(canonical, &a, &b) { + Ok(srid) => srid, + Err(e) => return Some(Err(e)), + }; + let d = if srid.is_geographic() { + nodedb_spatial::st_distance(&a.geometry, &b.geometry) + } else { + nodedb_spatial::planar_distance(&a.geometry, &b.geometry) + }; + to_value_number(d) } - "st_length" => unary(args, linear_length), - "st_perimeter" => unary(args, perimeter), - "st_area" => unary(args, nodedb_spatial::st_area), + "st_length" => unary(args, linear_length, nodedb_spatial::planar_length), + "st_perimeter" => unary(args, perimeter, nodedb_spatial::planar_perimeter), + "st_area" => unary(args, nodedb_spatial::st_area, nodedb_spatial::planar_area), "geo_distance" => haversine(args, nodedb_types::geometry::haversine_distance), "geo_bearing" => haversine(args, nodedb_types::geometry::haversine_bearing), _ => return None, }; - Some(result) + Some(Ok(result)) } -fn unary(args: &[Value], f: fn(&Geometry) -> f64) -> Value { +/// A unary measure: `geodesic` for a WGS 84 geometry, `planar` for a +/// projected one. +fn unary(args: &[Value], geodesic: fn(&Geometry) -> f64, planar: fn(&Geometry) -> f64) -> Value { match geom_arg(args, 0) { - Some(geom) => to_value_number(f(&geom)), + Some(SridGeometry { srid, geometry }) => { + let f = if srid.is_geographic() { + geodesic + } else { + planar + }; + to_value_number(f(&geometry)) + } None => Value::Null, } } @@ -106,9 +125,55 @@ fn closed_ring_length(ring: &[[f64; 2]]) -> f64 { mod tests { use super::*; use crate::value_ops::value_to_f64; + use nodedb_types::geometry::Srid; + + /// The evaluator with a fault turned into a test failure. + fn eval(canonical: &'static str, args: &[Value]) -> Option { + super::eval(canonical, args).map(|r| r.expect("no fault expected")) + } fn geom(v: Geometry) -> Vec { - vec![Value::Geometry(v)] + vec![Value::Geometry(SridGeometry::wgs84(v))] + } + + fn utm48(v: Geometry) -> Value { + Value::Geometry(SridGeometry::new(Srid::Utm48N, v)) + } + + #[test] + fn projected_area_of_hundred_metre_square() { + let square = utm48(Geometry::polygon(vec![vec![ + [500_000.0, 100_000.0], + [500_100.0, 100_000.0], + [500_100.0, 100_100.0], + [500_000.0, 100_100.0], + [500_000.0, 100_000.0], + ]])); + let area = eval("st_area", &[square]).and_then(|v| value_to_f64(&v, true)); + assert_eq!(area, Some(10_000.0)); + } + + #[test] + fn projected_distance_and_length_are_planar() { + let a = utm48(Geometry::point(500_000.0, 0.0)); + let b = utm48(Geometry::point(500_300.0, 400.0)); + let d = eval("st_distance", &[a, b]).and_then(|v| value_to_f64(&v, true)); + assert_eq!(d, Some(500.0)); + let line = utm48(Geometry::line_string(vec![[0.0, 0.0], [300.0, 400.0]])); + let l = eval("st_length", &[line]).and_then(|v| value_to_f64(&v, true)); + assert_eq!(l, Some(500.0)); + } + + #[test] + fn mixed_srid_distance_fails() { + let args = [ + Value::Geometry(SridGeometry::wgs84(Geometry::point(0.0, 0.0))), + utm48(Geometry::point(0.0, 0.0)), + ]; + assert!(matches!( + super::eval("st_distance", &args), + Some(Err(crate::expr::EvalError::MixedSrid { .. })) + )); } #[test] diff --git a/nodedb-query/src/geo_functions/mod.rs b/nodedb-query/src/geo_functions/mod.rs index 40a09f6ea..cf34c7fab 100644 --- a/nodedb-query/src/geo_functions/mod.rs +++ b/nodedb-query/src/geo_functions/mod.rs @@ -5,6 +5,7 @@ pub mod accessors; pub mod catalog; pub mod constructors; +pub mod crs; pub mod dispatch; pub mod helpers; pub mod indexing; diff --git a/nodedb-query/src/geo_functions/predicates.rs b/nodedb-query/src/geo_functions/predicates.rs index 0ed55b43a..940908b03 100644 --- a/nodedb-query/src/geo_functions/predicates.rs +++ b/nodedb-query/src/geo_functions/predicates.rs @@ -2,42 +2,118 @@ //! Topological predicate evaluation. Every predicate is NULL-propagating: an //! argument that is not a readable geometry yields NULL rather than `false`, -//! so "no geometry" stays distinguishable from "does not match". +//! so "no geometry" stays distinguishable from "does not match". Two +//! operands in different SRIDs fail with `EvalError::MixedSrid`. use nodedb_types::Value; use nodedb_types::geometry::Geometry; -use super::helpers::{geom_arg, num_arg}; +use super::helpers::{GeoEval, geom_arg, num_arg, same_srid}; +use crate::expr::EvalError; -pub(super) fn eval(canonical: &str, args: &[Value]) -> Option { +pub(super) fn eval(canonical: &'static str, args: &[Value]) -> GeoEval { let result = match canonical { - "st_contains" => binary(args, nodedb_spatial::st_contains), - "st_intersects" => binary(args, nodedb_spatial::st_intersects), - "st_within" => binary(args, nodedb_spatial::st_within), - "st_disjoint" => binary(args, nodedb_spatial::st_disjoint), - "st_dwithin" => { - let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else { - return Some(Value::Null); - }; - let Some(distance) = num_arg(args, 2) else { - return Some(Value::Null); - }; - Value::Bool(nodedb_spatial::st_dwithin(&a, &b, distance)) - } - "st_isvalid" => { - let Some(geom) = geom_arg(args, 0) else { - return Some(Value::Null); - }; - Value::Bool(nodedb_spatial::is_valid(&geom)) - } + "st_contains" => binary(canonical, args, nodedb_spatial::st_contains), + "st_intersects" => binary(canonical, args, nodedb_spatial::st_intersects), + "st_within" => binary(canonical, args, nodedb_spatial::st_within), + "st_disjoint" => binary(canonical, args, nodedb_spatial::st_disjoint), + "st_dwithin" => dwithin(canonical, args), + "st_isvalid" => Ok(match geom_arg(args, 0) { + Some(g) => Value::Bool(nodedb_spatial::is_valid(&g.geometry)), + None => Value::Null, + }), _ => return None, }; Some(result) } -fn binary(args: &[Value], f: fn(&Geometry, &Geometry) -> bool) -> Value { +fn binary( + function: &'static str, + args: &[Value], + f: fn(&Geometry, &Geometry) -> bool, +) -> Result { let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else { - return Value::Null; + return Ok(Value::Null); }; - Value::Bool(f(&a, &b)) + same_srid(function, &a, &b)?; + Ok(Value::Bool(f(&a.geometry, &b.geometry))) +} + +/// `ST_DWithin(a, b, distance)`, distance in metres: geodesic-approximate for +/// 4326, planar for a projected SRID. +fn dwithin(function: &'static str, args: &[Value]) -> Result { + let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else { + return Ok(Value::Null); + }; + let Some(distance) = num_arg(args, 2) else { + return Ok(Value::Null); + }; + let srid = same_srid(function, &a, &b)?; + let within = if srid.is_geographic() { + nodedb_spatial::st_dwithin(&a.geometry, &b.geometry, distance) + } else { + nodedb_spatial::planar_dwithin(&a.geometry, &b.geometry, distance) + }; + Ok(Value::Bool(within)) +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_types::geometry::{Srid, SridGeometry}; + + fn geom(srid: Srid, g: Geometry) -> Value { + Value::Geometry(SridGeometry::new(srid, g)) + } + + #[test] + fn mixed_srid_intersects_fails() { + let args = vec![ + geom(Srid::Wgs84, Geometry::point(1.0, 2.0)), + geom(Srid::Utm48N, Geometry::point(1.0, 2.0)), + ]; + assert_eq!( + eval("st_intersects", &args), + Some(Err(EvalError::MixedSrid { + function: "st_intersects", + left: 4326, + right: 32648 + })) + ); + } + + #[test] + fn projected_dwithin_is_planar_metres() { + let args = |d: f64| { + vec![ + geom(Srid::Utm48N, Geometry::point(500_000.0, 0.0)), + geom(Srid::Utm48N, Geometry::point(500_300.0, 400.0)), + Value::Float(d), + ] + }; + assert_eq!( + eval("st_dwithin", &args(500.0)), + Some(Ok(Value::Bool(true))) + ); + assert_eq!( + eval("st_dwithin", &args(499.0)), + Some(Ok(Value::Bool(false))) + ); + } + + #[test] + fn same_srid_predicate_evaluates() { + let square = Geometry::polygon(vec![vec![ + [0.0, 0.0], + [10.0, 0.0], + [10.0, 10.0], + [0.0, 10.0], + [0.0, 0.0], + ]]); + let args = vec![ + geom(Srid::PeninsularRso, square), + geom(Srid::PeninsularRso, Geometry::point(5.0, 5.0)), + ]; + assert_eq!(eval("st_contains", &args), Some(Ok(Value::Bool(true)))); + } } diff --git a/nodedb-spatial/src/crs/ellipsoid.rs b/nodedb-spatial/src/crs/ellipsoid.rs new file mode 100644 index 000000000..240623529 --- /dev/null +++ b/nodedb-spatial/src/crs/ellipsoid.rs @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Reference ellipsoids. + +/// An ellipsoid of revolution, defined by its semi-major axis and inverse +/// flattening. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Ellipsoid { + /// Semi-major axis in metres. + pub a: f64, + /// Inverse flattening, 1/f. + pub inv_f: f64, +} + +/// GRS 1980 (EPSG:7019), the ellipsoid of GDM2000. +pub const GRS80: Ellipsoid = Ellipsoid { + a: 6_378_137.0, + inv_f: 298.257_222_101, +}; + +/// WGS 84 (EPSG:7030). +pub const WGS84: Ellipsoid = Ellipsoid { + a: 6_378_137.0, + inv_f: 298.257_223_563, +}; + +impl Ellipsoid { + /// Flattening f. + pub fn f(&self) -> f64 { + 1.0 / self.inv_f + } + + /// First eccentricity squared, e² = f(2 − f). + pub fn e2(&self) -> f64 { + let f = self.f(); + f * (2.0 - f) + } + + /// First eccentricity e. + pub fn e(&self) -> f64 { + self.e2().sqrt() + } + + /// Third flattening n = f / (2 − f). + pub fn n(&self) -> f64 { + let f = self.f(); + f / (2.0 - f) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn grs80_and_wgs84_differ_only_in_flattening() { + assert_eq!(GRS80.a, WGS84.a); + assert!((GRS80.e2() - 0.006_694_380_022_90).abs() < 1e-14); + assert!((WGS84.e2() - 0.006_694_379_990_14).abs() < 1e-14); + } + + #[test] + fn third_flattening_matches_definition() { + let n = WGS84.n(); + let f = WGS84.f(); + assert!((n - f / (2.0 - f)).abs() < 1e-18); + assert!((WGS84.e() * WGS84.e() - WGS84.e2()).abs() < 1e-18); + } +} diff --git a/nodedb-spatial/src/crs/error.rs b/nodedb-spatial/src/crs/error.rs new file mode 100644 index 000000000..4d55bd78d --- /dev/null +++ b/nodedb-spatial/src/crs/error.rs @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Coordinate transformation errors. + +use nodedb_types::geometry::Srid; + +/// A coordinate a projection cannot map. +#[derive(Debug, Clone, Copy, PartialEq, thiserror::Error)] +pub enum CrsError { + /// The coordinate lies outside the region where the projection is + /// defined or accurate. + #[error("coordinate ({x}, {y}) is outside the domain of SRID {srid}: {reason}")] + OutOfDomain { + srid: Srid, + x: f64, + y: f64, + reason: DomainReason, + }, +} + +/// Why a coordinate is outside a projection's domain. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DomainReason { + /// An ordinate is NaN or infinite. + NonFinite, + /// Latitude is beyond ±90°. + LatitudeRange, + /// Longitude is too far from a Transverse Mercator central meridian for + /// the Krüger series to hold. + CentralMeridianOffset, + /// The Hotine Oblique Mercator U term reaches ±1, where the projection + /// diverges. + ObliqueSingularity, +} + +impl std::fmt::Display for DomainReason { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + DomainReason::NonFinite => "non-finite ordinate", + DomainReason::LatitudeRange => "latitude beyond ±90°", + DomainReason::CentralMeridianOffset => "too far from the central meridian", + DomainReason::ObliqueSingularity => "on the oblique projection's singular line", + }) + } +} + +/// Largest absolute latitude, in degrees. +const MAX_LATITUDE_DEG: f64 = 90.0; + +/// Fail unless `[lon, lat]` is a finite geographic coordinate with latitude +/// inside ±90°. +pub(crate) fn check_lon_lat(srid: Srid, coord: [f64; 2]) -> Result<(), CrsError> { + let [lon, lat] = coord; + if !lon.is_finite() || !lat.is_finite() { + return Err(CrsError::out_of_domain( + srid, + coord, + DomainReason::NonFinite, + )); + } + if lat.abs() > MAX_LATITUDE_DEG { + return Err(CrsError::out_of_domain( + srid, + coord, + DomainReason::LatitudeRange, + )); + } + Ok(()) +} + +/// Fail unless both ordinates are finite. +pub(crate) fn check_finite(srid: Srid, coord: [f64; 2]) -> Result<(), CrsError> { + if coord[0].is_finite() && coord[1].is_finite() { + Ok(()) + } else { + Err(CrsError::out_of_domain( + srid, + coord, + DomainReason::NonFinite, + )) + } +} + +impl CrsError { + pub(crate) fn out_of_domain(srid: Srid, [x, y]: [f64; 2], reason: DomainReason) -> Self { + CrsError::OutOfDomain { srid, x, y, reason } + } + + /// The SRID whose projection rejected the coordinate. + pub fn srid(&self) -> Srid { + match self { + CrsError::OutOfDomain { srid, .. } => *srid, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lon_lat_checks() { + assert!(check_lon_lat(Srid::Utm48N, [105.0, 90.0]).is_ok()); + assert_eq!( + check_lon_lat(Srid::Utm48N, [105.0, 90.5]), + Err(CrsError::out_of_domain( + Srid::Utm48N, + [105.0, 90.5], + DomainReason::LatitudeRange + )) + ); + assert!(check_lon_lat(Srid::Utm48N, [f64::NAN, 0.0]).is_err()); + assert!(check_finite(Srid::BorneoRso, [0.0, f64::INFINITY]).is_err()); + } + + #[test] + fn message_names_srid_and_reason() { + let e = CrsError::out_of_domain( + Srid::Utm47N, + [150.0, 3.0], + DomainReason::CentralMeridianOffset, + ); + assert_eq!(e.srid(), Srid::Utm47N); + let text = e.to_string(); + assert!( + text.contains("32647") && text.contains("central meridian"), + "{text}" + ); + } +} diff --git a/nodedb-spatial/src/crs/hom.rs b/nodedb-spatial/src/crs/hom.rs new file mode 100644 index 000000000..920e085b6 --- /dev/null +++ b/nodedb-spatial/src/crs/hom.rs @@ -0,0 +1,284 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Hotine Oblique Mercator, by IOGP Guidance Note 7-2 §3.2.4. +//! +//! Variant A (EPSG method 9812) measures u from the natural origin, the PROJ +//! `+no_uoff` form. Variant B (EPSG method 9815) measures u from the +//! projection centre. + +use std::f64::consts::FRAC_PI_4; + +use nodedb_types::geometry::Srid; + +use super::ellipsoid::Ellipsoid; +use super::error::{CrsError, DomainReason, check_finite, check_lon_lat}; + +/// |cos αc| below which the initial-line azimuth counts as 90°, where +/// GN7-2 replaces the uc formula with uc = A (λc − λ0). +const AZIMUTH_RIGHT_ANGLE_EPS: f64 = 1e-12; + +/// Where u is measured from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HomVariant { + /// EPSG 9812: from the natural origin. False easting and northing apply + /// at the natural origin. + A, + /// EPSG 9815: from the projection centre. False easting and northing + /// apply at the projection centre. + B, +} + +/// Defining parameters of one Hotine Oblique Mercator projection. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct HomParams { + pub variant: HomVariant, + /// Latitude of the projection centre φc, degrees. + pub lat_c_deg: f64, + /// Longitude of the projection centre λc, degrees. + pub lon_c_deg: f64, + /// Azimuth of the initial line αc, degrees. + pub alpha_c_deg: f64, + /// Angle from the rectified grid to the skew grid γc, degrees. + pub gamma_c_deg: f64, + /// Scale factor on the initial line. + pub k_c: f64, + /// False easting, metres. + pub false_easting: f64, + /// False northing, metres. + pub false_northing: f64, + pub ellipsoid: Ellipsoid, +} + +/// A Hotine Oblique Mercator projection with its constants computed. +#[derive(Debug, Clone)] +pub struct HotineObliqueMercator { + srid: Srid, + params: HomParams, + e: f64, + e2: f64, + b: f64, + a: f64, + h: f64, + sin_g0: f64, + cos_g0: f64, + sin_gc: f64, + cos_gc: f64, + /// λ0 in radians. + lon0: f64, + /// Offset of u for variant B: |uc| · sign(φc). Zero for variant A. + u_offset: f64, +} + +impl HotineObliqueMercator { + pub fn new(srid: Srid, params: HomParams) -> Self { + let el = params.ellipsoid; + let (e, e2) = (el.e(), el.e2()); + let phi_c = params.lat_c_deg.to_radians(); + let lam_c = params.lon_c_deg.to_radians(); + let alpha_c = params.alpha_c_deg.to_radians(); + let gamma_c = params.gamma_c_deg.to_radians(); + let (sin_pc, cos_pc) = phi_c.sin_cos(); + let sign_pc = if phi_c < 0.0 { -1.0 } else { 1.0 }; + + let b = (1.0 + e2 * cos_pc.powi(4) / (1.0 - e2)).sqrt(); + let a = el.a * b * params.k_c * (1.0 - e2).sqrt() / (1.0 - e2 * sin_pc * sin_pc); + let t0 = isometric_t(phi_c, e); + let d = b * (1.0 - e2).sqrt() / (cos_pc * (1.0 - e2 * sin_pc * sin_pc).sqrt()); + // GN7-2: when D < 1, take D² = 1 so F stays real. + let d2 = (d * d).max(1.0); + let f = d + (d2 - 1.0).sqrt() * sign_pc; + let h = f * t0.powf(b); + let g = (f - 1.0 / f) / 2.0; + let gamma0 = (alpha_c.sin() / d).asin(); + let lon0 = lam_c - (g * gamma0.tan()).asin() / b; + let u_offset = match params.variant { + HomVariant::A => 0.0, + HomVariant::B => { + let uc = if alpha_c.cos().abs() < AZIMUTH_RIGHT_ANGLE_EPS { + a * (lam_c - lon0) + } else { + a / b * ((d2 - 1.0).sqrt() / alpha_c.cos()).atan() * sign_pc + }; + uc.abs() * sign_pc + } + }; + let (sin_g0, cos_g0) = gamma0.sin_cos(); + let (sin_gc, cos_gc) = gamma_c.sin_cos(); + Self { + srid, + params, + e, + e2, + b, + a, + h, + sin_g0, + cos_g0, + sin_gc, + cos_gc, + lon0, + u_offset, + } + } + + /// Project `[lon, lat]` degrees to `[easting, northing]` metres. + pub fn forward(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + check_lon_lat(self.srid, coord)?; + let phi = coord[1].to_radians(); + let lam = coord[0].to_radians(); + let t = isometric_t(phi, self.e); + let q = self.h / t.powf(self.b); + let s = (q - 1.0 / q) / 2.0; + let tt = (q + 1.0 / q) / 2.0; + let (v_sin, b_dlam_cos) = (self.b * (lam - self.lon0)).sin_cos(); + let u_big = (-v_sin * self.cos_g0 + s * self.sin_g0) / tt; + if !u_big.is_finite() || u_big.abs() >= 1.0 { + return Err(CrsError::out_of_domain( + self.srid, + coord, + DomainReason::ObliqueSingularity, + )); + } + let v = self.a * ((1.0 - u_big) / (1.0 + u_big)).ln() / (2.0 * self.b); + let u = self.a * (s * self.cos_g0 + v_sin * self.sin_g0).atan2(b_dlam_cos) / self.b + - self.u_offset; + let out = [ + v * self.cos_gc + u * self.sin_gc + self.params.false_easting, + u * self.cos_gc - v * self.sin_gc + self.params.false_northing, + ]; + check_finite(self.srid, out)?; + Ok(out) + } + + /// Unproject `[easting, northing]` metres to `[lon, lat]` degrees. + pub fn inverse(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + check_finite(self.srid, coord)?; + let de = coord[0] - self.params.false_easting; + let dn = coord[1] - self.params.false_northing; + let v = de * self.cos_gc - dn * self.sin_gc; + let u = dn * self.cos_gc + de * self.sin_gc + self.u_offset; + let q = (-self.b * v / self.a).exp(); + let s = (q - 1.0 / q) / 2.0; + let tt = (q + 1.0 / q) / 2.0; + let (v_sin, bu_cos) = (self.b * u / self.a).sin_cos(); + let u_big = (v_sin * self.cos_g0 + s * self.sin_g0) / tt; + if !u_big.is_finite() || u_big.abs() >= 1.0 { + return Err(CrsError::out_of_domain( + self.srid, + coord, + DomainReason::ObliqueSingularity, + )); + } + let t = (self.h / ((1.0 + u_big) / (1.0 - u_big)).sqrt()).powf(1.0 / self.b); + let chi = std::f64::consts::FRAC_PI_2 - 2.0 * t.atan(); + let phi = latitude_from_conformal(chi, self.e2); + let lam = self.lon0 - (s * self.cos_g0 - v_sin * self.sin_g0).atan2(bu_cos) / self.b; + let out = [lam.to_degrees(), phi.to_degrees()]; + check_finite(self.srid, out)?; + Ok(out) + } +} + +/// The GN7-2 t function: tan(π/4 − φ/2) / [(1 − e sin φ)/(1 + e sin φ)]^(e/2). +fn isometric_t(phi: f64, e: f64) -> f64 { + let es = e * phi.sin(); + (FRAC_PI_4 - phi / 2.0).tan() / ((1.0 - es) / (1.0 + es)).powf(e / 2.0) +} + +/// Geodetic latitude from conformal latitude χ, the GN7-2 series in e². +/// The first omitted term is of order e¹⁰ ≈ 1.4e-11 rad, under 0.1 mm. +fn latitude_from_conformal(chi: f64, e2: f64) -> f64 { + let e4 = e2 * e2; + let e6 = e4 * e2; + let e8 = e6 * e2; + chi + (2.0 * chi).sin() * (e2 / 2.0 + 5.0 * e4 / 24.0 + e6 / 12.0 + 13.0 * e8 / 360.0) + + (4.0 * chi).sin() * (7.0 * e4 / 48.0 + 29.0 * e6 / 240.0 + 811.0 * e8 / 11520.0) + + (6.0 * chi).sin() * (7.0 * e6 / 120.0 + 81.0 * e8 / 1120.0) + + (8.0 * chi).sin() * (4279.0 * e8 / 161_280.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::crs::params::{BORNEO_RSO, PENINSULAR_RSO}; + + /// One millimetre, the tolerance for every PROJ reference value. + const MM: f64 = 1e-3; + + fn rso_peninsular() -> HotineObliqueMercator { + HotineObliqueMercator::new(Srid::PeninsularRso, PENINSULAR_RSO) + } + + fn rso_borneo() -> HotineObliqueMercator { + HotineObliqueMercator::new(Srid::BorneoRso, BORNEO_RSO) + } + + fn assert_close(got: [f64; 2], want: [f64; 2]) { + assert!( + (got[0] - want[0]).abs() < MM && (got[1] - want[1]).abs() < MM, + "got {got:?}, want {want:?}" + ); + } + + // Reference values: PROJ 9.8.1 via pyproj, EPSG:4326 → target, always_xy. + + #[test] + fn peninsular_rso_klcc() { + let got = rso_peninsular().forward([101.7117, 3.1579]).unwrap(); + assert_close(got, [412_834.468_1, 349_472.917_7]); + } + + #[test] + fn peninsular_rso_105_5() { + let got = rso_peninsular().forward([105.0, 5.0]).unwrap(); + assert_close(got, [778_104.319_5, 553_022.237_0]); + } + + #[test] + fn borneo_rso_kota_kinabalu() { + let got = rso_borneo().forward([116.0735, 5.9804]).unwrap(); + assert_close(got, [708_697.150_6, 662_326.916_8]); + } + + #[test] + fn borneo_rso_105_5() { + let got = rso_borneo().forward([105.0, 5.0]).unwrap(); + assert_close(got, [-521_596.204_7, 554_604.058_6]); + } + + #[test] + fn inverse_recovers_input_to_a_millimetre() { + for hom in [rso_peninsular(), rso_borneo()] { + for input in [[101.7117, 3.1579], [116.0735, 5.9804], [105.0, 5.0]] { + let projected = hom.forward(input).unwrap(); + let back = hom.inverse(projected).unwrap(); + let again = hom.forward(back).unwrap(); + assert_close(again, projected); + } + } + } + + /// Variant B moves u to the projection centre: the centre itself maps to + /// the false origin. + #[test] + fn variant_b_maps_centre_to_false_origin() { + let params = HomParams { + variant: HomVariant::B, + false_easting: 0.0, + false_northing: 0.0, + ..PENINSULAR_RSO + }; + let hom = HotineObliqueMercator::new(Srid::PeninsularRso, params); + let centre = hom.forward([params.lon_c_deg, params.lat_c_deg]).unwrap(); + assert_close(centre, [0.0, 0.0]); + let back = hom.inverse(centre).unwrap(); + assert!((back[0] - params.lon_c_deg).abs() < 1e-9); + assert!((back[1] - params.lat_c_deg).abs() < 1e-9); + } + + #[test] + fn pole_is_out_of_domain() { + assert!(rso_peninsular().forward([102.0, 90.0]).is_err()); + assert!(rso_peninsular().forward([102.0, 90.5]).is_err()); + } +} diff --git a/nodedb-spatial/src/crs/mod.rs b/nodedb-spatial/src/crs/mod.rs new file mode 100644 index 000000000..5734313fc --- /dev/null +++ b/nodedb-spatial/src/crs/mod.rs @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Coordinate reference system transforms for the supported SRIDs: +//! WGS 84, UTM 47N / 48N, and the GDM2000 Malaysian RSO grids. + +pub mod ellipsoid; +pub mod error; +pub mod hom; +pub mod params; +pub mod tm; +pub mod transform; + +pub use error::{CrsError, DomainReason}; +pub use hom::{HomParams, HomVariant, HotineObliqueMercator}; +pub use tm::{ + TAU_NEWTON_MAX_ITERS, TAU_TOLERANCE, TM_MAX_LON_OFFSET_DEG, TmParams, TransverseMercator, +}; +pub use transform::{Projection, Transformer, transform_geometry}; diff --git a/nodedb-spatial/src/crs/params.rs b/nodedb-spatial/src/crs/params.rs new file mode 100644 index 000000000..ff75d39cb --- /dev/null +++ b/nodedb-spatial/src/crs/params.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Defining parameters of every supported projected SRID. +//! +//! Values are the EPSG definitions as PROJ publishes them. +//! +//! Datum bound: GDM2000 (EPSG:3375, 3376) is ITRF2000 at epoch 2000.0. Its +//! published `towgs84` is all zeros, so WGS 84 → GDM2000 is taken as the +//! identity. GDM2000 is fixed to the plate while WGS 84 tracks ITRF, so the +//! identity drifts with Sundaland plate motion: about 1 m by 2026. Results in +//! 3375 / 3376 are exact to the projection, not to that datum offset. + +use super::ellipsoid::{GRS80, WGS84}; +use super::hom::{HomParams, HomVariant}; +use super::tm::TmParams; + +/// UTM scale factor on the central meridian. +const UTM_K0: f64 = 0.9996; + +/// UTM false easting, metres. +const UTM_FALSE_EASTING: f64 = 500_000.0; + +/// UTM false northing in the northern hemisphere, metres. +const UTM_FALSE_NORTHING_NORTH: f64 = 0.0; + +/// RSO scale factor on the initial line, both Malaysian RSO grids. +const RSO_K: f64 = 0.99984; + +/// EPSG:32647, UTM zone 47N on WGS 84. +pub const UTM_47N: TmParams = TmParams { + lon0_deg: 99.0, + k0: UTM_K0, + false_easting: UTM_FALSE_EASTING, + false_northing: UTM_FALSE_NORTHING_NORTH, + ellipsoid: WGS84, +}; + +/// EPSG:32648, UTM zone 48N on WGS 84. +pub const UTM_48N: TmParams = TmParams { + lon0_deg: 105.0, + k0: UTM_K0, + false_easting: UTM_FALSE_EASTING, + false_northing: UTM_FALSE_NORTHING_NORTH, + ellipsoid: WGS84, +}; + +/// EPSG:3375, GDM2000 Peninsula RSO. +/// `+proj=omerc +no_uoff +lat_0=4 +lonc=102.25 +alpha=323.025796466667 +/// +gamma=323.130102361111 +k=0.99984 +x_0=804671 +y_0=0 +ellps=GRS80`. +pub const PENINSULAR_RSO: HomParams = HomParams { + variant: HomVariant::A, + lat_c_deg: 4.0, + lon_c_deg: 102.25, + alpha_c_deg: 323.025_796_466_667, + gamma_c_deg: 323.130_102_361_111, + k_c: RSO_K, + false_easting: 804_671.0, + false_northing: 0.0, + ellipsoid: GRS80, +}; + +/// EPSG:3376, GDM2000 East Malaysia (Borneo) RSO. +/// `+proj=omerc +no_uoff +lat_0=4 +lonc=115 +alpha=53.31580995 +/// +gamma=53.1301023611111 +k=0.99984 +x_0=0 +y_0=0 +ellps=GRS80`. +pub const BORNEO_RSO: HomParams = HomParams { + variant: HomVariant::A, + lat_c_deg: 4.0, + lon_c_deg: 115.0, + alpha_c_deg: 53.315_809_95, + gamma_c_deg: 53.130_102_361_111_1, + k_c: RSO_K, + false_easting: 0.0, + false_northing: 0.0, + ellipsoid: GRS80, +}; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn utm_zones_differ_only_in_central_meridian() { + assert_eq!(UTM_47N.lon0_deg, 99.0); + assert_eq!(UTM_48N.lon0_deg, 105.0); + assert_eq!( + TmParams { + lon0_deg: UTM_48N.lon0_deg, + ..UTM_47N + }, + UTM_48N + ); + } + + #[test] + fn rso_grids_are_variant_a_on_grs80() { + for p in [PENINSULAR_RSO, BORNEO_RSO] { + assert_eq!(p.variant, HomVariant::A); + assert_eq!(p.ellipsoid, GRS80); + assert_eq!(p.k_c, RSO_K); + assert_eq!(p.lat_c_deg, 4.0); + } + } +} diff --git a/nodedb-spatial/src/crs/tm.rs b/nodedb-spatial/src/crs/tm.rs new file mode 100644 index 000000000..b8fd47c3d --- /dev/null +++ b/nodedb-spatial/src/crs/tm.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Transverse Mercator, by the Krüger series to n⁶. +//! +//! Follows C. F. F. Karney, "Transverse Mercator with an accuracy of a few +//! nanometers", J. Geodesy 85 (2011), eqs. 7–9 and 35–36. Within 35° of the +//! central meridian the series errors stay far below a millimetre. + +use nodedb_types::geometry::Srid; + +use super::ellipsoid::Ellipsoid; +use super::error::{CrsError, DomainReason, check_finite, check_lon_lat}; + +/// Largest |λ − λ0| in degrees that the projection accepts. Karney 2011 +/// bounds the n⁶ series error below 5 nm inside 35° of the central meridian. +pub const TM_MAX_LON_OFFSET_DEG: f64 = 35.0; + +/// Newton iterations allowed when inverting τ' → τ. Karney 2011 §2.2 reports +/// convergence to double precision in two; the cap only bounds bad input. +pub const TAU_NEWTON_MAX_ITERS: usize = 10; + +/// Newton step on τ below which the inverse stops. A step this small moves +/// latitude by under 1e-14 rad, about 64 nm on the ground. +pub const TAU_TOLERANCE: f64 = 1e-14; + +/// Half a turn in degrees, for wrapping a longitude difference. +const HALF_TURN_DEG: f64 = 180.0; + +/// A full turn in degrees. +const FULL_TURN_DEG: f64 = 360.0; + +/// Defining parameters of one Transverse Mercator projection. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct TmParams { + /// Central meridian λ0, degrees. + pub lon0_deg: f64, + /// Scale factor on the central meridian. + pub k0: f64, + /// False easting, metres. + pub false_easting: f64, + /// False northing, metres. + pub false_northing: f64, + pub ellipsoid: Ellipsoid, +} + +/// A Transverse Mercator projection with its series coefficients computed. +#[derive(Debug, Clone)] +pub struct TransverseMercator { + srid: Srid, + params: TmParams, + e: f64, + e2: f64, + /// k0 · A, where A is the rectifying radius. + k0_a: f64, + /// Krüger α coefficients (forward). + alpha: [f64; 6], + /// Krüger β coefficients (inverse). + beta: [f64; 6], +} + +impl TransverseMercator { + pub fn new(srid: Srid, params: TmParams) -> Self { + let n = params.ellipsoid.n(); + let n2 = n * n; + let n3 = n2 * n; + let n4 = n3 * n; + let n5 = n4 * n; + let n6 = n5 * n; + // Rectifying radius A, Karney 2011 eq. 14. + let a_rect = params.ellipsoid.a / (1.0 + n) * (1.0 + n2 / 4.0 + n4 / 64.0 + n6 / 256.0); + // Karney 2011 eq. 35. + let alpha = [ + n / 2.0 - 2.0 / 3.0 * n2 + 5.0 / 16.0 * n3 + 41.0 / 180.0 * n4 - 127.0 / 288.0 * n5 + + 7891.0 / 37800.0 * n6, + 13.0 / 48.0 * n2 - 3.0 / 5.0 * n3 + 557.0 / 1440.0 * n4 + 281.0 / 630.0 * n5 + - 1_983_433.0 / 1_935_360.0 * n6, + 61.0 / 240.0 * n3 - 103.0 / 140.0 * n4 + + 15061.0 / 26880.0 * n5 + + 167_603.0 / 181_440.0 * n6, + 49561.0 / 161_280.0 * n4 - 179.0 / 168.0 * n5 + 6_601_661.0 / 7_257_600.0 * n6, + 34729.0 / 80640.0 * n5 - 3_418_889.0 / 1_995_840.0 * n6, + 212_378_941.0 / 319_334_400.0 * n6, + ]; + // Karney 2011 eq. 36. + let beta = [ + n / 2.0 - 2.0 / 3.0 * n2 + 37.0 / 96.0 * n3 - 1.0 / 360.0 * n4 - 81.0 / 512.0 * n5 + + 96199.0 / 604_800.0 * n6, + 1.0 / 48.0 * n2 + 1.0 / 15.0 * n3 - 437.0 / 1440.0 * n4 + 46.0 / 105.0 * n5 + - 1_118_711.0 / 3_870_720.0 * n6, + 17.0 / 480.0 * n3 - 37.0 / 840.0 * n4 - 209.0 / 4480.0 * n5 + 5569.0 / 90720.0 * n6, + 4397.0 / 161_280.0 * n4 - 11.0 / 504.0 * n5 - 830_251.0 / 7_257_600.0 * n6, + 4583.0 / 161_280.0 * n5 - 108_847.0 / 3_991_680.0 * n6, + 20_648_693.0 / 638_668_800.0 * n6, + ]; + Self { + srid, + params, + e: params.ellipsoid.e(), + e2: params.ellipsoid.e2(), + k0_a: params.k0 * a_rect, + alpha, + beta, + } + } + + /// Project `[lon, lat]` degrees to `[easting, northing]` metres. + pub fn forward(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + check_lon_lat(self.srid, coord)?; + let [lon, lat] = coord; + let dlon_deg = wrap_lon_offset(lon - self.params.lon0_deg); + if dlon_deg.abs() > TM_MAX_LON_OFFSET_DEG { + return Err(CrsError::out_of_domain( + self.srid, + coord, + DomainReason::CentralMeridianOffset, + )); + } + let dlon = dlon_deg.to_radians(); + let tau_p = conformal_tau(lat.to_radians().tan(), self.e); + let xi_p = tau_p.atan2(dlon.cos()); + let eta_p = (dlon.sin() / tau_p.hypot(dlon.cos())).asinh(); + let mut xi = xi_p; + let mut eta = eta_p; + for (j, a) in self.alpha.iter().enumerate() { + let k = 2.0 * (j + 1) as f64; + xi += a * (k * xi_p).sin() * (k * eta_p).cosh(); + eta += a * (k * xi_p).cos() * (k * eta_p).sinh(); + } + let out = [ + self.params.false_easting + self.k0_a * eta, + self.params.false_northing + self.k0_a * xi, + ]; + check_finite(self.srid, out)?; + Ok(out) + } + + /// Unproject `[easting, northing]` metres to `[lon, lat]` degrees. + pub fn inverse(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + check_finite(self.srid, coord)?; + let xi = (coord[1] - self.params.false_northing) / self.k0_a; + let eta = (coord[0] - self.params.false_easting) / self.k0_a; + let mut xi_p = xi; + let mut eta_p = eta; + for (j, b) in self.beta.iter().enumerate() { + let k = 2.0 * (j + 1) as f64; + xi_p -= b * (k * xi).sin() * (k * eta).cosh(); + eta_p -= b * (k * xi).cos() * (k * eta).sinh(); + } + let tau_p = xi_p.sin() / eta_p.sinh().hypot(xi_p.cos()); + let dlon = eta_p.sinh().atan2(xi_p.cos()); + let tau = self.tau_from_conformal(tau_p); + let dlon_deg = dlon.to_degrees(); + if dlon_deg.abs() > TM_MAX_LON_OFFSET_DEG { + return Err(CrsError::out_of_domain( + self.srid, + coord, + DomainReason::CentralMeridianOffset, + )); + } + let out = [self.params.lon0_deg + dlon_deg, tau.atan().to_degrees()]; + check_finite(self.srid, out)?; + Ok(out) + } + + /// Solve τ' = τ'(τ) for τ by Newton's method, Karney 2011 eqs. 19–21. + fn tau_from_conformal(&self, tau_p: f64) -> f64 { + let one_minus_e2 = 1.0 - self.e2; + let mut tau = tau_p; + for _ in 0..TAU_NEWTON_MAX_ITERS { + let tau_i_p = conformal_tau(tau, self.e); + let step = (tau_p - tau_i_p) / (1.0 + tau_i_p * tau_i_p).sqrt() + * (1.0 + one_minus_e2 * tau * tau) + / (one_minus_e2 * (1.0 + tau * tau).sqrt()); + tau += step; + if step.abs() < TAU_TOLERANCE { + break; + } + } + tau + } +} + +/// Conformal-latitude tangent τ' from geodetic τ = tan φ, Karney 2011 eq. 7. +fn conformal_tau(tau: f64, e: f64) -> f64 { + let tau1 = (1.0 + tau * tau).sqrt(); + let sigma = (e * (e * tau / tau1).atanh()).sinh(); + tau * (1.0 + sigma * sigma).sqrt() - sigma * tau1 +} + +/// Wrap a longitude difference into [-180°, 180°]. +fn wrap_lon_offset(dlon_deg: f64) -> f64 { + let d = dlon_deg % FULL_TURN_DEG; + if d > HALF_TURN_DEG { + d - FULL_TURN_DEG + } else if d < -HALF_TURN_DEG { + d + FULL_TURN_DEG + } else { + d + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::crs::params::{UTM_47N, UTM_48N}; + + /// One millimetre, the tolerance for every PROJ reference value. + const MM: f64 = 1e-3; + + fn utm47() -> TransverseMercator { + TransverseMercator::new(Srid::Utm47N, UTM_47N) + } + + fn utm48() -> TransverseMercator { + TransverseMercator::new(Srid::Utm48N, UTM_48N) + } + + fn assert_close(got: [f64; 2], want: [f64; 2]) { + assert!( + (got[0] - want[0]).abs() < MM && (got[1] - want[1]).abs() < MM, + "got {got:?}, want {want:?}" + ); + } + + // Reference values: PROJ 9.8.1 via pyproj, EPSG:4326 → target, always_xy. + + #[test] + fn utm47n_klcc() { + let got = utm47().forward([101.7117, 3.1579]).unwrap(); + assert_close(got, [801_401.794_9, 349_439.510_6]); + } + + #[test] + fn utm48n_klcc() { + let got = utm48().forward([101.7117, 3.1579]).unwrap(); + assert_close(got, [134_445.592_4, 349_624.727_8]); + } + + #[test] + fn utm48n_origin_on_central_meridian() { + assert_close(utm48().forward([105.0, 0.0]).unwrap(), [500_000.0, 0.0]); + } + + #[test] + fn utm48n_central_meridian_five_north() { + let got = utm48().forward([105.0, 5.0]).unwrap(); + assert_close(got, [500_000.0, 552_664.296_9]); + } + + #[test] + fn inverse_recovers_input() { + let tm = utm48(); + for input in [ + [101.7117, 3.1579], + [105.0, 0.0], + [109.5, 6.2], + [100.0, -2.0], + ] { + let back = tm.inverse(tm.forward(input).unwrap()).unwrap(); + assert!( + (back[0] - input[0]).abs() < 1e-9 && (back[1] - input[1]).abs() < 1e-9, + "{input:?} -> {back:?}" + ); + } + } + + #[test] + fn far_from_central_meridian_is_out_of_domain() { + let err = utm48().forward([150.0, 3.0]).unwrap_err(); + assert!(matches!( + err, + CrsError::OutOfDomain { + reason: DomainReason::CentralMeridianOffset, + .. + } + )); + } + + #[test] + fn bad_latitude_and_non_finite_are_out_of_domain() { + assert!(utm48().forward([105.0, 91.0]).is_err()); + assert!(utm48().forward([f64::NAN, 1.0]).is_err()); + assert!(utm48().inverse([f64::INFINITY, 1.0]).is_err()); + } + + #[test] + fn wrap_keeps_offset_in_half_turn() { + assert_eq!(wrap_lon_offset(200.0), -160.0); + assert_eq!(wrap_lon_offset(-200.0), 160.0); + assert_eq!(wrap_lon_offset(10.0), 10.0); + } +} diff --git a/nodedb-spatial/src/crs/transform.rs b/nodedb-spatial/src/crs/transform.rs new file mode 100644 index 000000000..d328727f3 --- /dev/null +++ b/nodedb-spatial/src/crs/transform.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Coordinate transformation between supported SRIDs. +//! +//! Every path runs through WGS 84 geographic coordinates: the inverse of the +//! source projection (none for 4326), then the forward of the target. The +//! WGS 84 → GDM2000 datum step is the identity (see `params`). + +use nodedb_types::geometry::{Srid, SridGeometry}; + +use super::error::CrsError; +use super::hom::HotineObliqueMercator; +use super::params::{BORNEO_RSO, PENINSULAR_RSO, UTM_47N, UTM_48N}; +use super::tm::TransverseMercator; + +/// The projection behind one projected SRID. +#[derive(Debug, Clone)] +pub enum Projection { + TransverseMercator(TransverseMercator), + HotineObliqueMercator(HotineObliqueMercator), +} + +impl Projection { + /// The projection of `srid`, or `None` for the geographic WGS 84. + pub fn for_srid(srid: Srid) -> Option { + match srid { + Srid::Wgs84 => None, + Srid::Utm47N => Some(Projection::TransverseMercator(TransverseMercator::new( + srid, UTM_47N, + ))), + Srid::Utm48N => Some(Projection::TransverseMercator(TransverseMercator::new( + srid, UTM_48N, + ))), + Srid::PeninsularRso => Some(Projection::HotineObliqueMercator( + HotineObliqueMercator::new(srid, PENINSULAR_RSO), + )), + Srid::BorneoRso => Some(Projection::HotineObliqueMercator( + HotineObliqueMercator::new(srid, BORNEO_RSO), + )), + } + } + + /// `[lon, lat]` degrees to `[easting, northing]` metres. + pub fn forward(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + match self { + Projection::TransverseMercator(p) => p.forward(coord), + Projection::HotineObliqueMercator(p) => p.forward(coord), + } + } + + /// `[easting, northing]` metres to `[lon, lat]` degrees. + pub fn inverse(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + match self { + Projection::TransverseMercator(p) => p.inverse(coord), + Projection::HotineObliqueMercator(p) => p.inverse(coord), + } + } +} + +/// A prepared `from → to` coordinate transform. +#[derive(Debug, Clone)] +pub struct Transformer { + from: Srid, + source: Option, + target: Option, +} + +impl Transformer { + pub fn new(from: Srid, to: Srid) -> Self { + // A same-SRID transform maps nothing. + if from == to { + return Self { + from, + source: None, + target: None, + }; + } + Self { + from, + source: Projection::for_srid(from), + target: Projection::for_srid(to), + } + } + + /// Map one coordinate. + pub fn apply(&self, coord: [f64; 2]) -> Result<[f64; 2], CrsError> { + let geographic = match &self.source { + Some(p) => p.inverse(coord)?, + None => { + super::error::check_finite(self.from, coord)?; + coord + } + }; + match &self.target { + Some(p) => p.forward(geographic), + None => Ok(geographic), + } + } +} + +/// Transform every coordinate of `geom` into `target`. +/// +/// A geometry already in `target` comes back unchanged. Any coordinate a +/// projection cannot map fails the whole call with +/// [`CrsError::OutOfDomain`]. +pub fn transform_geometry(geom: &SridGeometry, target: Srid) -> Result { + if geom.srid == target { + return Ok(geom.clone()); + } + let transformer = Transformer::new(geom.srid, target); + let geometry = geom + .geometry + .try_map_coords(&mut |c| transformer.apply(c))?; + Ok(SridGeometry::new(target, geometry)) +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_types::geometry::Geometry; + + /// Largest round-trip error accepted, metres. + const ROUND_TRIP_TOLERANCE_M: f64 = 1e-3; + + /// Metres per degree of latitude, rounded up. Converts a metre tolerance + /// to a conservative degree tolerance. + const METRES_PER_DEGREE: f64 = 111_320.0; + + fn point(srid: Srid, x: f64, y: f64) -> SridGeometry { + SridGeometry::new(srid, Geometry::point(x, y)) + } + + fn coords(g: &SridGeometry) -> [f64; 2] { + match g.geometry { + Geometry::Point { coordinates } => coordinates, + _ => panic!("expected a point"), + } + } + + #[test] + fn wgs84_round_trip_through_every_srid() { + let inputs = [[101.7117, 3.1579], [116.0735, 5.9804], [103.5, 1.5]]; + for srid in &Srid::SUPPORTED[1..] { + for input in inputs { + let src = point(Srid::Wgs84, input[0], input[1]); + let projected = transform_geometry(&src, *srid).unwrap(); + assert_eq!(projected.srid, *srid); + let back = transform_geometry(&projected, Srid::Wgs84).unwrap(); + let [lon, lat] = coords(&back); + let tol = ROUND_TRIP_TOLERANCE_M / METRES_PER_DEGREE; + assert!( + (lon - input[0]).abs() < tol && (lat - input[1]).abs() < tol, + "{srid:?}: {input:?} -> {:?}", + [lon, lat] + ); + } + } + } + + #[test] + fn peninsular_rso_through_utm47n_round_trips() { + let src = point(Srid::PeninsularRso, 412_834.468_1, 349_472.917_7); + let utm = transform_geometry(&src, Srid::Utm47N).unwrap(); + let back = transform_geometry(&utm, Srid::PeninsularRso).unwrap(); + let [x, y] = coords(&back); + assert!( + (x - 412_834.468_1).abs() < ROUND_TRIP_TOLERANCE_M + && (y - 349_472.917_7).abs() < ROUND_TRIP_TOLERANCE_M, + "{x} {y}" + ); + let [ux, uy] = coords(&utm); + assert!((ux - 801_401.794_9).abs() < ROUND_TRIP_TOLERANCE_M, "{ux}"); + assert!((uy - 349_439.510_6).abs() < ROUND_TRIP_TOLERANCE_M, "{uy}"); + } + + #[test] + fn same_srid_is_identity() { + let src = point(Srid::Utm48N, 1.0, 2.0); + assert_eq!(transform_geometry(&src, Srid::Utm48N), Ok(src)); + } + + #[test] + fn every_coordinate_is_mapped() { + let src = SridGeometry::wgs84(Geometry::polygon(vec![vec![ + [101.0, 3.0], + [102.0, 3.0], + [102.0, 4.0], + [101.0, 3.0], + ]])); + let out = transform_geometry(&src, Srid::Utm47N).unwrap(); + let Geometry::Polygon { coordinates } = out.geometry else { + panic!("shape changed"); + }; + assert_eq!(coordinates[0].len(), 4); + assert!(coordinates[0].iter().all(|c| c[0] > 100_000.0)); + } + + #[test] + fn out_of_domain_fails_the_call() { + let src = point(Srid::Wgs84, 179.0, 3.0); + assert!(matches!( + transform_geometry(&src, Srid::Utm48N), + Err(CrsError::OutOfDomain { .. }) + )); + } +} diff --git a/nodedb-spatial/src/hybrid.rs b/nodedb-spatial/src/hybrid.rs index eef11e6b3..d63f1932d 100644 --- a/nodedb-spatial/src/hybrid.rs +++ b/nodedb-spatial/src/hybrid.rs @@ -23,7 +23,7 @@ use nodedb_mem::ScopedMemory; #[cfg(test)] use nodedb_types::BoundingBox; use nodedb_types::geometry::Geometry; -use nodedb_types::geometry_bbox; +use nodedb_types::{BboxSpace, geometry_bbox_in}; use crate::predicates; use crate::rtree::RTree; @@ -61,11 +61,13 @@ pub fn spatial_prefilter( exact_geometries: &dyn Fn(u64) -> Option, memory: &ScopedMemory, ) -> SpatialPreFilterResult { - // Step 1: Compute search bbox. - let search_bbox = if let Some(dist) = distance_meters { - geometry_bbox(query_geometry).expand_meters(dist) - } else { - geometry_bbox(query_geometry) + // Compute the search bbox in the tree's space. The query geometry is in + // the tree's SRID: the caller rejects a mismatch before this point. + let space = rtree.space(); + let query_bbox = geometry_bbox_in(query_geometry, space); + let search_bbox = match distance_meters { + Some(dist) => query_bbox.expand_meters_in(dist, space), + None => query_bbox, }; // Step 2: R-tree range search. @@ -78,10 +80,14 @@ pub fn spatial_prefilter( let mut candidate_ids = Vec::with_capacity(rtree_candidates); for entry in &rtree_results { if let Some(doc_geom) = exact_geometries(entry.id) { - let passes = if let Some(dist) = distance_meters { - predicates::st_dwithin(&doc_geom, query_geometry, dist) - } else { - predicates::st_intersects(&doc_geom, query_geometry) + let passes = match (distance_meters, space) { + (Some(dist), BboxSpace::Geographic) => { + predicates::st_dwithin(&doc_geom, query_geometry, dist) + } + (Some(dist), BboxSpace::Planar) => { + crate::planar::planar_dwithin(&doc_geom, query_geometry, dist) + } + (None, _) => predicates::st_intersects(&doc_geom, query_geometry), }; if passes { candidate_ids.push(entry.id); @@ -135,7 +141,7 @@ mod tests { use crate::test_support::test_memory; fn make_tree() -> RTree { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), nodedb_types::geometry::Srid::Wgs84); // 10 points in a grid from (0,0) to (9,9). for i in 0..10 { tree.insert(RTreeEntry { @@ -197,6 +203,26 @@ mod tests { assert!(!result.candidate_ids.contains(&9)); } + /// A projected tree expands the search box by metres and checks the + /// threshold in Euclidean metres. + #[test] + fn projected_prefilter_uses_planar_metres() { + let mut tree = RTree::new(test_memory(), nodedb_types::geometry::Srid::Utm48N); + let points = [[500_000.0, 0.0], [500_300.0, 400.0], [510_000.0, 0.0]]; + for (id, p) in points.iter().enumerate() { + tree.insert(RTreeEntry { + id: id as u64, + bbox: BoundingBox::from_point(p[0], p[1]), + }); + } + let get_geom = |id: u64| points.get(id as usize).map(|p| Geometry::point(p[0], p[1])); + let query = Geometry::point(500_000.0, 0.0); + let result = spatial_prefilter(&tree, &query, Some(500.0), &get_geom, &test_memory()); + let mut ids = result.candidate_ids.clone(); + ids.sort_unstable(); + assert_eq!(ids, vec![0, 1]); + } + #[test] fn bitmap_roundtrip() { let ids = vec![0, 5, 7, 63, 100]; diff --git a/nodedb-spatial/src/lib.rs b/nodedb-spatial/src/lib.rs index d6d24479a..bd7726cac 100644 --- a/nodedb-spatial/src/lib.rs +++ b/nodedb-spatial/src/lib.rs @@ -10,6 +10,7 @@ //! and use the columnar storage core; this crate provides the index //! structures and predicate evaluators only. +pub mod crs; pub mod geo_meta; pub mod geohash; pub mod geohash_index; @@ -18,6 +19,7 @@ pub mod hybrid; pub mod measures; pub mod operations; pub mod persist; +pub mod planar; pub mod predicates; pub mod rtree; pub mod spatial_join; @@ -27,6 +29,7 @@ pub mod validate; pub mod wkb; pub mod wkt; +pub use crs::{CrsError, transform_geometry}; pub use geohash::{geohash_decode, geohash_encode, geohash_neighbors}; pub use geohash_index::GeohashIndex; pub use h3::{ @@ -34,16 +37,19 @@ pub use h3::{ h3_to_boundary, h3_to_center, }; pub use hybrid::{SpatialPreFilterResult, bitmap_contains, ids_to_bitmap, spatial_prefilter}; -pub use measures::{st_area, st_centroid}; -pub use operations::{st_buffer, st_envelope, st_union}; +pub use measures::{st_area, st_centroid, st_centroid_in}; +pub use operations::{st_buffer, st_buffer_in, st_envelope, st_envelope_in, st_union}; pub use persist::{ RTreeCheckpointError, SpatialIndexMeta, SpatialIndexType, deserialize_meta, meta_storage_key, rtree_storage_key, serialize_meta, }; +pub use planar::{ + planar_area, planar_buffer, planar_distance, planar_dwithin, planar_length, planar_perimeter, +}; pub use predicates::{ st_contains, st_disjoint, st_distance, st_dwithin, st_intersection, st_intersects, st_within, }; -pub use rtree::{RTree, RTreeEntry}; +pub use rtree::{RTree, RTreeEntry, SridMismatch}; pub use validate::{is_valid, validate_geometry}; -pub use wkb::{geometry_from_wkb, geometry_to_wkb}; +pub use wkb::{geometry_from_wkb, geometry_to_ewkb, geometry_to_wkb}; pub use wkt::{geometry_from_wkt, geometry_to_wkt}; diff --git a/nodedb-spatial/src/measures.rs b/nodedb-spatial/src/measures.rs index 97aa18622..cb94c2ff8 100644 --- a/nodedb-spatial/src/measures.rs +++ b/nodedb-spatial/src/measures.rs @@ -9,6 +9,7 @@ //! present, so a `GeometryCollection` of a polygon and a stray point returns //! the polygon's centroid, not a blend of the two. +use nodedb_types::BboxSpace; use nodedb_types::geometry::Geometry; /// Mean Earth radius in meters, matching `nodedb_types::geometry`'s haversine. @@ -79,11 +80,38 @@ fn ring_area(ring: &[[f64; 2]]) -> f64 { /// vertices at all, or when every component of its winning dimension is /// degenerate (zero area / zero length) — those fall back to the next /// dimension down rather than returning nothing. +/// +/// Weights are geodesic (spherical area, haversine length). Same as +/// `st_centroid_in(geom, BboxSpace::Geographic)`. pub fn st_centroid(geom: &Geometry) -> Option { - let [lng, lat] = centroid_coords(geom)?; + st_centroid_in(geom, BboxSpace::Geographic) +} + +/// Centroid of a geometry in `space`, as a `Point`. +/// +/// Planar space weights components by planar area and Euclidean length, so +/// a projected geometry's centroid is in the same metres as its input. +pub fn st_centroid_in(geom: &Geometry, space: BboxSpace) -> Option { + let [lng, lat] = centroid_coords(geom, space)?; Some(Geometry::point(lng, lat)) } +/// Unsigned area of a ring in `space`: square metres in either space. +fn ring_weight(ring: &[[f64; 2]], space: BboxSpace) -> f64 { + match space { + BboxSpace::Geographic => ring_area(ring), + BboxSpace::Planar => crate::planar::area::planar_ring_area(ring), + } +} + +/// Length of one segment in `space`, metres. +fn segment_weight(a: [f64; 2], b: [f64; 2], space: BboxSpace) -> f64 { + match space { + BboxSpace::Geographic => nodedb_types::geometry::haversine_distance(a[0], a[1], b[0], b[1]), + BboxSpace::Planar => (b[0] - a[0]).hypot(b[1] - a[1]), + } +} + /// A weighted point accumulator: `(Σ w·lng, Σ w·lat, Σ w)`. #[derive(Default, Clone, Copy)] struct Weighted { @@ -114,26 +142,26 @@ impl Weighted { } } -fn centroid_coords(geom: &Geometry) -> Option<[f64; 2]> { - areal_centroid(geom) +fn centroid_coords(geom: &Geometry, space: BboxSpace) -> Option<[f64; 2]> { + areal_centroid(geom, space) .and_then(Weighted::resolve) - .or_else(|| linear_centroid(geom).and_then(Weighted::resolve)) + .or_else(|| linear_centroid(geom, space).and_then(Weighted::resolve)) .or_else(|| punctual_centroid(geom).and_then(Weighted::resolve)) } /// Area-weighted centroid of every areal component, or `None` if there are no /// areal components at all. -fn areal_centroid(geom: &Geometry) -> Option { +fn areal_centroid(geom: &Geometry, space: BboxSpace) -> Option { match geom { - Geometry::Polygon { coordinates } => Some(polygon_centroid(coordinates)), + Geometry::Polygon { coordinates } => Some(polygon_centroid(coordinates, space)), Geometry::MultiPolygon { coordinates } => { let mut acc = Weighted::default(); for rings in coordinates { - acc.merge(polygon_centroid(rings)); + acc.merge(polygon_centroid(rings, space)); } Some(acc) } - Geometry::GeometryCollection { geometries } => collect(geometries, areal_centroid), + Geometry::GeometryCollection { geometries } => collect(geometries, space, areal_centroid), Geometry::Point { .. } | Geometry::LineString { .. } | Geometry::MultiPoint { .. } @@ -145,17 +173,17 @@ fn areal_centroid(geom: &Geometry) -> Option { } /// Length-weighted centroid of every linear component. -fn linear_centroid(geom: &Geometry) -> Option { +fn linear_centroid(geom: &Geometry, space: BboxSpace) -> Option { match geom { - Geometry::LineString { coordinates } => Some(line_centroid(coordinates)), + Geometry::LineString { coordinates } => Some(line_centroid(coordinates, space)), Geometry::MultiLineString { coordinates } => { let mut acc = Weighted::default(); for line in coordinates { - acc.merge(line_centroid(line)); + acc.merge(line_centroid(line, space)); } Some(acc) } - Geometry::GeometryCollection { geometries } => collect(geometries, linear_centroid), + Geometry::GeometryCollection { geometries } => collect(geometries, space, linear_centroid), Geometry::Point { .. } | Geometry::MultiPoint { .. } | Geometry::Polygon { .. } @@ -178,11 +206,15 @@ fn punctual_centroid(geom: &Geometry) -> Option { /// Merge the per-component results of `f` across a collection, yielding `None` /// when no member contributed at that dimension. -fn collect(geometries: &[Geometry], f: fn(&Geometry) -> Option) -> Option { +fn collect( + geometries: &[Geometry], + space: BboxSpace, + f: fn(&Geometry, BboxSpace) -> Option, +) -> Option { let mut acc = Weighted::default(); let mut any = false; for member in geometries { - if let Some(part) = f(member) { + if let Some(part) = f(member, space) { acc.merge(part); any = true; } @@ -192,18 +224,18 @@ fn collect(geometries: &[Geometry], f: fn(&Geometry) -> Option) -> Opt /// Centroid of one polygon, weighting the exterior ring positively and each /// hole negatively so the hole's mass is removed from the result. -fn polygon_centroid(rings: &[Vec<[f64; 2]>]) -> Weighted { +fn polygon_centroid(rings: &[Vec<[f64; 2]>], space: BboxSpace) -> Weighted { let mut acc = Weighted::default(); let mut iter = rings.iter(); let Some(exterior) = iter.next() else { return acc; }; if let Some(point) = ring_centroid(exterior) { - acc.add(point, ring_area(exterior)); + acc.add(point, ring_weight(exterior, space)); } for hole in iter { if let Some(point) = ring_centroid(hole) { - acc.add(point, -ring_area(hole)); + acc.add(point, -ring_weight(hole, space)); } } // Holes can only ever cancel the exterior down to nothing; a negative @@ -243,12 +275,12 @@ fn ring_centroid(ring: &[[f64; 2]]) -> Option<[f64; 2]> { } /// Length-weighted centroid of a polyline: each segment contributes its -/// midpoint weighted by its geodesic length. -fn line_centroid(coords: &[[f64; 2]]) -> Weighted { +/// midpoint weighted by its length in `space`. +fn line_centroid(coords: &[[f64; 2]], space: BboxSpace) -> Weighted { let mut acc = Weighted::default(); for pair in coords.windows(2) { let (a, b) = (pair[0], pair[1]); - let length = nodedb_types::geometry::haversine_distance(a[0], a[1], b[0], b[1]); + let length = segment_weight(a, b, space); acc.add([(a[0] + b[0]) / 2.0, (a[1] + b[1]) / 2.0], length); } acc diff --git a/nodedb-spatial/src/operations.rs b/nodedb-spatial/src/operations.rs index 9fb9855c6..d910c8cf2 100644 --- a/nodedb-spatial/src/operations.rs +++ b/nodedb-spatial/src/operations.rs @@ -5,35 +5,65 @@ //! ST_Buffer, ST_Envelope, ST_Union — produce new geometries from existing ones. use nodedb_types::geometry::Geometry; -use nodedb_types::{BoundingBox, geometry_bbox}; +use nodedb_types::{BboxSpace, BoundingBox, geometry_bbox_in}; -/// ST_Buffer(geom, distance_meters) → Polygon. +/// Metres per degree of latitude (WGS 84 meridian average). +const METRES_PER_DEGREE_LAT: f64 = 110_540.0; + +/// Metres per degree of longitude at the equator (WGS 84). +const METRES_PER_DEGREE_LNG_EQUATOR: f64 = 111_320.0; + +/// Floor on cos(latitude) so a buffer near a pole stays finite. +const MIN_COS_LAT: f64 = 0.001; + +/// Floor on a vector length before normalising it, so a degenerate edge +/// does not divide by zero. +const MIN_VECTOR_LEN: f64 = 1e-12; + +/// ST_Buffer(geom, distance_meters) → Polygon, in geographic space. /// -/// Expands geometry by distance using equirectangular approximation. -/// - Point → 32-sided circle polygon +/// Same as `st_buffer_in(geom, distance_meters, segments, BboxSpace::Geographic)`. +pub fn st_buffer(geom: &Geometry, distance_meters: f64, segments: usize) -> Geometry { + st_buffer_in(geom, distance_meters, segments, BboxSpace::Geographic) +} + +/// ST_Buffer(geom, distance_meters) → Polygon, in `space`. +/// +/// - Point → `segments`-sided circle polygon /// - LineString → buffered corridor (parallel offset + end caps) /// - Polygon → expanded polygon (vertex offset outward) /// -/// Meter-to-degree conversion: `Δlng = meters / (111320 * cos(lat))`, -/// `Δlat = meters / 110540`. Accurate to <0.5% for ±70° latitude. -pub fn st_buffer(geom: &Geometry, distance_meters: f64, segments: usize) -> Geometry { +/// Geographic space converts metres to degrees by equirectangular +/// approximation: `Δlng = meters / (111320 * cos(lat))`, +/// `Δlat = meters / 110540`, accurate to <0.5% within ±70° latitude. +/// Planar space adds metres directly. +pub fn st_buffer_in( + geom: &Geometry, + distance_meters: f64, + segments: usize, + space: BboxSpace, +) -> Geometry { let segments = segments.max(4); match geom { - Geometry::Point { coordinates } => { - buffer_point(coordinates[0], coordinates[1], distance_meters, segments) - } + Geometry::Point { coordinates } => buffer_point( + coordinates[0], + coordinates[1], + distance_meters, + segments, + space, + ), Geometry::LineString { coordinates } => { - buffer_linestring(coordinates, distance_meters, segments) + buffer_linestring(coordinates, distance_meters, segments, space) } - Geometry::Polygon { coordinates } => buffer_polygon(coordinates, distance_meters, segments), + Geometry::Polygon { coordinates } => buffer_polygon(coordinates, distance_meters, space), // Multi* types: buffer each component and collect. Geometry::MultiPoint { coordinates } => { let polys: Vec>> = coordinates .iter() .map(|pt| { if let Geometry::Polygon { coordinates: rings } = - buffer_point(pt[0], pt[1], distance_meters, segments) + buffer_point(pt[0], pt[1], distance_meters, segments, space) { rings } else { @@ -45,15 +75,32 @@ pub fn st_buffer(geom: &Geometry, distance_meters: f64, segments: usize) -> Geom } _ => { // For unsupported types, fall back to buffering the bbox. - let bb = geometry_bbox(geom).expand_meters(distance_meters); + let bb = geometry_bbox_in(geom, space).expand_meters_in(distance_meters, space); bbox_to_polygon(&bb) } } } -/// ST_Envelope(geom) → Polygon — bounding box as a polygon. +/// Coordinate deltas `(dlng, dlat)` that `meters` spans at latitude `lat`. +fn metre_deltas(meters: f64, lat: f64, space: BboxSpace) -> (f64, f64) { + match space { + BboxSpace::Geographic => ( + meters / (METRES_PER_DEGREE_LNG_EQUATOR * lat.to_radians().cos().max(MIN_COS_LAT)), + meters / METRES_PER_DEGREE_LAT, + ), + BboxSpace::Planar => (meters, meters), + } +} + +/// ST_Envelope(geom) → Polygon — bounding box as a polygon, in geographic +/// space. pub fn st_envelope(geom: &Geometry) -> Geometry { - let bb = geometry_bbox(geom); + st_envelope_in(geom, BboxSpace::Geographic) +} + +/// ST_Envelope(geom) → Polygon — bounding box as a polygon, in `space`. +pub fn st_envelope_in(geom: &Geometry, space: BboxSpace) -> Geometry { + let bb = geometry_bbox_in(geom, space); bbox_to_polygon(&bb) } @@ -107,9 +154,8 @@ pub fn st_union(a: &Geometry, b: &Geometry) -> Geometry { // ── Buffer helpers ── /// Buffer a point into an approximate circle polygon. -fn buffer_point(lng: f64, lat: f64, meters: f64, segments: usize) -> Geometry { - let dlat = meters / 110_540.0; - let dlng = meters / (111_320.0 * lat.to_radians().cos().max(0.001)); +fn buffer_point(lng: f64, lat: f64, meters: f64, segments: usize, space: BboxSpace) -> Geometry { + let (dlng, dlat) = metre_deltas(meters, lat, space); let mut ring = Vec::with_capacity(segments + 1); for i in 0..segments { @@ -125,14 +171,19 @@ fn buffer_point(lng: f64, lat: f64, meters: f64, segments: usize) -> Geometry { } /// Buffer a linestring by creating a corridor (parallel offset on both sides + end caps). -fn buffer_linestring(line: &[[f64; 2]], meters: f64, segments: usize) -> Geometry { +fn buffer_linestring( + line: &[[f64; 2]], + meters: f64, + segments: usize, + space: BboxSpace, +) -> Geometry { if line.is_empty() { return Geometry::Polygon { coordinates: vec![vec![]], }; } if line.len() == 1 { - return buffer_point(line[0][0], line[0][1], meters, segments); + return buffer_point(line[0][0], line[0][1], meters, segments, space); } // Build the offset polygon: left side forward, right side backward, with @@ -141,7 +192,7 @@ fn buffer_linestring(line: &[[f64; 2]], meters: f64, segments: usize) -> Geometr let mut right_side = Vec::new(); for i in 0..line.len() - 1 { - let (nl, nr) = offset_segment(line[i], line[i + 1], meters); + let (nl, nr) = offset_segment(line[i], line[i + 1], meters, space); left_side.push(nl.0); left_side.push(nl.1); right_side.push(nr.0); @@ -159,8 +210,7 @@ fn buffer_linestring(line: &[[f64; 2]], meters: f64, segments: usize) -> Geometr for i in 0..=cap_segments { let angle = end_dir - std::f64::consts::FRAC_PI_2 + std::f64::consts::PI * (i as f64) / (cap_segments as f64); - let dlat = meters / 110_540.0; - let dlng = meters / (111_320.0 * last[1].to_radians().cos().max(0.001)); + let (dlng, dlat) = metre_deltas(meters, last[1], space); ring.push([last[0] + dlng * angle.cos(), last[1] + dlat * angle.sin()]); } @@ -176,8 +226,7 @@ fn buffer_linestring(line: &[[f64; 2]], meters: f64, segments: usize) -> Geometr let angle = start_dir + std::f64::consts::FRAC_PI_2 + std::f64::consts::PI * (i as f64) / (cap_segments as f64); - let dlat = meters / 110_540.0; - let dlng = meters / (111_320.0 * first[1].to_radians().cos().max(0.001)); + let (dlng, dlat) = metre_deltas(meters, first[1], space); ring.push([first[0] + dlng * angle.cos(), first[1] + dlat * angle.sin()]); } @@ -192,7 +241,7 @@ fn buffer_linestring(line: &[[f64; 2]], meters: f64, segments: usize) -> Geometr } /// Buffer a polygon by offsetting each vertex outward. -fn buffer_polygon(rings: &[Vec<[f64; 2]>], meters: f64, _segments: usize) -> Geometry { +fn buffer_polygon(rings: &[Vec<[f64; 2]>], meters: f64, space: BboxSpace) -> Geometry { let mut new_rings = Vec::with_capacity(rings.len()); for (ring_idx, ring) in rings.iter().enumerate() { @@ -223,11 +272,10 @@ fn buffer_polygon(rings: &[Vec<[f64; 2]>], meters: f64, _segments: usize) -> Geo let bisect = [(n1[0] + n2[0]) / 2.0, (n1[1] + n2[1]) / 2.0]; let len = (bisect[0] * bisect[0] + bisect[1] * bisect[1]) .sqrt() - .max(1e-12); + .max(MIN_VECTOR_LEN); let unit = [bisect[0] / len, bisect[1] / len]; - let dlat = offset_m / 110_540.0; - let dlng = offset_m / (111_320.0 * curr[1].to_radians().cos().max(0.001)); + let (dlng, dlat) = metre_deltas(offset_m, curr[1], space); new_ring.push([curr[0] + unit[0] * dlng, curr[1] + unit[1] * dlat]); } @@ -248,11 +296,9 @@ fn buffer_polygon(rings: &[Vec<[f64; 2]>], meters: f64, _segments: usize) -> Geo /// A pair of offset segments: (left_segment, right_segment). type OffsetPair = (([f64; 2], [f64; 2]), ([f64; 2], [f64; 2])); -fn offset_segment(a: [f64; 2], b: [f64; 2], meters: f64) -> OffsetPair { +fn offset_segment(a: [f64; 2], b: [f64; 2], meters: f64, space: BboxSpace) -> OffsetPair { let normal = edge_outward_normal(a, b); - let dlat = meters / 110_540.0; - let mid_lat = ((a[1] + b[1]) / 2.0).to_radians().cos().max(0.001); - let dlng = meters / (111_320.0 * mid_lat); + let (dlng, dlat) = metre_deltas(meters, (a[1] + b[1]) / 2.0, space); let left = ( [a[0] + normal[0] * dlng, a[1] + normal[1] * dlat], @@ -273,7 +319,7 @@ fn offset_segment(a: [f64; 2], b: [f64; 2], meters: f64) -> OffsetPair { fn edge_outward_normal(a: [f64; 2], b: [f64; 2]) -> [f64; 2] { let dx = b[0] - a[0]; let dy = b[1] - a[1]; - let len = (dx * dx + dy * dy).sqrt().max(1e-12); + let len = (dx * dx + dy * dy).sqrt().max(MIN_VECTOR_LEN); // Right normal: (dy, -dx) normalized — outward for CCW rings. [dy / len, -dx / len] } @@ -299,6 +345,24 @@ fn bbox_to_polygon(bb: &BoundingBox) -> Geometry { #[cfg(test)] mod tests { use super::*; + use nodedb_types::geometry_bbox; + + #[test] + fn planar_buffer_point_is_a_metre_circle() { + let circle = st_buffer_in( + &Geometry::point(500_000.0, 100_000.0), + 250.0, + 16, + BboxSpace::Planar, + ); + let Geometry::Polygon { coordinates } = &circle else { + panic!("expected Polygon"); + }; + for pt in &coordinates[0][..16] { + let d = (pt[0] - 500_000.0).hypot(pt[1] - 100_000.0); + assert!((d - 250.0).abs() < 1e-6, "d={d}"); + } + } #[test] fn buffer_point_circle() { diff --git a/nodedb-spatial/src/persist.rs b/nodedb-spatial/src/persist.rs index 73e2900bc..af7cfcb82 100644 --- a/nodedb-spatial/src/persist.rs +++ b/nodedb-spatial/src/persist.rs @@ -24,6 +24,7 @@ use nodedb_mem::ScopedMemory; use nodedb_types::BoundingBox; +use nodedb_types::geometry::{Srid, UnsupportedSrid}; use serde::{Deserialize, Serialize}; use zerompk::{FromMessagePack, ToMessagePack}; @@ -43,8 +44,9 @@ const SEGV_MAGIC: [u8; 4] = *b"SEGV"; /// Magic header for rkyv-serialized R-tree snapshots (6 bytes). const RTREE_RKYV_MAGIC: &[u8; 6] = b"RKSPT\0"; -/// Current format version for rkyv-serialized R-tree snapshots. -pub const RTREE_FORMAT_VERSION: u8 = 1; +/// Current format version for rkyv-serialized R-tree snapshots. Version 2 +/// carries the tree's SRID. +pub const RTREE_FORMAT_VERSION: u8 = 2; // ── SEGV framing helpers ─────────────────────────────────────────────────── @@ -95,6 +97,8 @@ pub struct SpatialIndexMeta { pub entry_count: u64, /// Bounding box of all indexed geometries (spatial extent). pub extent: Option, + /// SRID of every indexed geometry. + pub srid: Srid, } /// Type of spatial index. @@ -129,6 +133,8 @@ impl std::fmt::Display for SpatialIndexType { /// rkyv-serialized R-tree snapshot. #[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] struct RTreeSnapshotRkyv { + /// EPSG code of the tree's SRID. + srid: u32, entries: Vec, } @@ -145,6 +151,7 @@ impl RTree { kek: Option<&nodedb_wal::crypto::WalEncryptionKey>, ) -> Result, RTreeCheckpointError> { let snapshot = RTreeSnapshotRkyv { + srid: self.srid().code(), entries: self.entries().into_iter().cloned().collect(), }; let rkyv_bytes = rkyv::to_bytes::(&snapshot) @@ -219,7 +226,8 @@ impl RTree { let snapshot: RTreeSnapshotRkyv = rkyv::from_bytes::(&aligned) .map_err(|e| RTreeCheckpointError::RkyvDeserialize(e.to_string()))?; - Ok(RTree::bulk_load(snapshot.entries, memory)) + let srid = Srid::try_from(snapshot.srid)?; + Ok(RTree::bulk_load(snapshot.entries, memory, srid)) } } @@ -297,6 +305,9 @@ pub enum RTreeCheckpointError { /// AES-256-GCM decryption failed. #[error("spatial checkpoint decryption failed: {0}")] DecryptionFailed(String), + /// The checkpoint names an SRID outside the supported set. + #[error("R-tree checkpoint: {0}")] + UnsupportedSrid(#[from] UnsupportedSrid), } // ── Tests ────────────────────────────────────────────────────────────────── @@ -315,7 +326,7 @@ mod tests { #[test] fn checkpoint_roundtrip_empty() { - let tree = RTree::new(test_memory()); + let tree = RTree::new(test_memory(), Srid::Wgs84); let bytes = tree.checkpoint_to_bytes(None).unwrap(); let restored = RTree::from_checkpoint(&bytes, None, test_memory()).unwrap(); assert_eq!(restored.len(), 0); @@ -323,7 +334,7 @@ mod tests { #[test] fn checkpoint_roundtrip_entries() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..100 { tree.insert(make_entry(i, (i as f64) * 0.5, (i as f64) * 0.3)); } @@ -340,7 +351,7 @@ mod tests { #[test] fn checkpoint_preserves_ids() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(42, 10.0, 20.0)); tree.insert(make_entry(99, 30.0, 40.0)); @@ -368,12 +379,24 @@ mod tests { index_type: SpatialIndexType::RTree, entry_count: 1000, extent: Some(BoundingBox::new(-180.0, -90.0, 180.0, 90.0)), + srid: Srid::PeninsularRso, }; let bytes = serialize_meta(&meta).unwrap(); let restored = deserialize_meta(&bytes).unwrap(); assert_eq!(restored.collection, "buildings"); assert_eq!(restored.entry_count, 1000); assert_eq!(restored.index_type, SpatialIndexType::RTree); + assert_eq!(restored.srid, Srid::PeninsularRso); + } + + #[test] + fn checkpoint_preserves_srid() { + let mut tree = RTree::new(test_memory(), Srid::BorneoRso); + tree.insert(make_entry(7, 708_697.0, 662_326.0)); + let bytes = tree.checkpoint_to_bytes(None).unwrap(); + let restored = RTree::from_checkpoint(&bytes, None, test_memory()).unwrap(); + assert_eq!(restored.srid(), Srid::BorneoRso); + assert_eq!(restored.len(), 1); } #[test] @@ -387,7 +410,7 @@ mod tests { #[test] fn checkpoint_size_reasonable() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..1000 { tree.insert(make_entry(i, (i as f64) * 0.01, (i as f64) * 0.01)); } @@ -408,7 +431,7 @@ mod tests { #[test] fn golden_header_layout() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); let bytes = tree.checkpoint_to_bytes(None).unwrap(); // Magic at bytes[0..6]. @@ -421,7 +444,7 @@ mod tests { #[test] fn version_mismatch_returns_error() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); let mut bytes = tree.checkpoint_to_bytes(None).unwrap(); // Corrupt the version byte to an unsupported value. @@ -443,7 +466,7 @@ mod tests { #[test] fn spatial_rtree_checkpoint_encrypted_at_rest() { let kek = make_test_kek(); - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..50 { tree.insert(make_entry(i, i as f64, i as f64 * 0.5)); } @@ -463,7 +486,7 @@ mod tests { #[test] fn spatial_rtree_refuses_plaintext_when_kek_required() { let kek = make_test_kek(); - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); // Write plaintext checkpoint. @@ -479,7 +502,7 @@ mod tests { #[test] fn spatial_rtree_refuses_encrypted_without_kek() { let kek = make_test_kek(); - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); let enc_bytes = tree.checkpoint_to_bytes(Some(&kek)).unwrap(); @@ -494,7 +517,7 @@ mod tests { #[test] fn spatial_rtree_tampered_ciphertext_rejected() { let kek = make_test_kek(); - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); let mut enc_bytes = tree.checkpoint_to_bytes(Some(&kek)).unwrap(); diff --git a/nodedb-spatial/src/planar/area.rs b/nodedb-spatial/src/planar/area.rs new file mode 100644 index 000000000..286c8b8f4 --- /dev/null +++ b/nodedb-spatial/src/planar/area.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Planar area of projected geometries, in square metres. + +use nodedb_types::geometry::Geometry; + +/// Unsigned area of one ring by the shoelace formula. A ring with fewer than +/// three vertices encloses nothing. An unclosed ring is closed implicitly. +/// +/// Coordinates are taken relative to the first vertex. Projected eastings and +/// northings run to hundreds of kilometres, and summing their raw cross +/// products cancels away the digits a small ring's area lives in. +pub fn planar_ring_area(ring: &[[f64; 2]]) -> f64 { + if ring.len() < 3 { + return 0.0; + } + let origin = ring[0]; + let rel = |p: [f64; 2]| [p[0] - origin[0], p[1] - origin[1]]; + let mut twice = 0.0; + for i in 0..ring.len() { + let a = rel(ring[i]); + let b = rel(ring[(i + 1) % ring.len()]); + twice += a[0] * b[1] - b[0] * a[1]; + } + (twice / 2.0).abs() +} + +/// Area of one polygon: exterior ring less every hole, never negative. +fn polygon_area(rings: &[Vec<[f64; 2]>]) -> f64 { + let mut iter = rings.iter(); + let Some(exterior) = iter.next() else { + return 0.0; + }; + let holes: f64 = iter.map(|ring| planar_ring_area(ring)).sum(); + (planar_ring_area(exterior) - holes).max(0.0) +} + +/// Planar area of a geometry in square coordinate units (square metres for a +/// projected SRID). +/// +/// Points and lines have zero area, matching PostGIS. A GeometryCollection +/// sums its members. +pub fn planar_area(geom: &Geometry) -> f64 { + match geom { + Geometry::Polygon { coordinates } => polygon_area(coordinates), + Geometry::MultiPolygon { coordinates } => { + coordinates.iter().map(|rings| polygon_area(rings)).sum() + } + Geometry::GeometryCollection { geometries } => geometries.iter().map(planar_area).sum(), + Geometry::Point { .. } + | Geometry::LineString { .. } + | Geometry::MultiPoint { .. } + | Geometry::MultiLineString { .. } => 0.0, + // `Geometry` is `#[non_exhaustive]`. A kind added upstream is not + // known to be areal. + _ => 0.0, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hundred_metre_square_is_ten_thousand_square_metres() { + let square = Geometry::polygon(vec![vec![ + [500_000.0, 100_000.0], + [500_100.0, 100_000.0], + [500_100.0, 100_100.0], + [500_000.0, 100_100.0], + [500_000.0, 100_000.0], + ]]); + assert!((planar_area(&square) - 10_000.0).abs() < 1e-6); + } + + #[test] + fn hole_is_subtracted_and_orientation_ignored() { + let cw = vec![ + [0.0, 0.0], + [0.0, 10.0], + [10.0, 10.0], + [10.0, 0.0], + [0.0, 0.0], + ]; + let hole = vec![[2.0, 2.0], [4.0, 2.0], [4.0, 4.0], [2.0, 4.0], [2.0, 2.0]]; + let poly = Geometry::polygon(vec![cw, hole]); + assert!((planar_area(&poly) - 96.0).abs() < 1e-9); + } + + #[test] + fn lines_and_points_have_no_area() { + assert_eq!(planar_area(&Geometry::point(1.0, 2.0)), 0.0); + let line = Geometry::line_string(vec![[0.0, 0.0], [5.0, 5.0]]); + assert_eq!(planar_area(&line), 0.0); + } +} diff --git a/nodedb-spatial/src/planar/buffer.rs b/nodedb-spatial/src/planar/buffer.rs new file mode 100644 index 000000000..d3f2e2474 --- /dev/null +++ b/nodedb-spatial/src/planar/buffer.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Planar buffer of projected geometries, by a distance in metres. + +use nodedb_types::BboxSpace; +use nodedb_types::geometry::Geometry; + +use crate::operations::st_buffer_in; + +/// Buffer `geom` by `distance` coordinate units (metres for a projected +/// SRID), approximating curves with `segments` sides. +/// +/// Uses the same construction as the geographic `st_buffer`, with metres +/// added to coordinates directly instead of converted to degrees. +pub fn planar_buffer(geom: &Geometry, distance: f64, segments: usize) -> Geometry { + st_buffer_in(geom, distance, segments, BboxSpace::Planar) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::planar::area::planar_area; + + #[test] + fn buffered_point_area_approaches_circle() { + let segments = 64; + let circle = planar_buffer(&Geometry::point(500_000.0, 200_000.0), 100.0, segments); + let area = planar_area(&circle); + // A regular n-gon inscribed in radius r has area n r² sin(2π/n) / 2. + let n = segments as f64; + let expected = n * 100.0 * 100.0 * (std::f64::consts::TAU / n).sin() / 2.0; + assert!( + (area - expected).abs() < 1e-6, + "area={area} expected={expected}" + ); + } +} diff --git a/nodedb-spatial/src/planar/distance.rs b/nodedb-spatial/src/planar/distance.rs new file mode 100644 index 000000000..e3ea7cfda --- /dev/null +++ b/nodedb-spatial/src/planar/distance.rs @@ -0,0 +1,265 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Planar distance and length of projected geometries, in metres. +//! +//! Every measure is Euclidean in coordinate units, which for a projected SRID +//! are metres on the grid. + +use nodedb_types::geometry::{Geometry, point_in_polygon}; +use nodedb_types::{BboxSpace, geometry_bbox_in}; + +use crate::predicates::edge::{point_on_ring_boundary, ring_edges, segment_to_segment_dist_sq}; + +/// The pieces of a geometry a distance is measured between. +#[derive(Default)] +struct Parts<'a> { + /// Every edge, with a lone point as a zero-length edge. + segments: Vec<([f64; 2], [f64; 2])>, + /// Every vertex. + vertices: Vec<[f64; 2]>, + /// Every polygon, as its rings. + polygons: Vec<&'a [Vec<[f64; 2]>]>, +} + +impl<'a> Parts<'a> { + fn of(geom: &'a Geometry) -> Self { + let mut parts = Parts::default(); + parts.add(geom); + parts + } + + fn add(&mut self, geom: &'a Geometry) { + match geom { + Geometry::Point { coordinates } => self.add_point(*coordinates), + Geometry::MultiPoint { coordinates } => { + coordinates.iter().for_each(|c| self.add_point(*c)); + } + Geometry::LineString { coordinates } => self.add_path(coordinates), + Geometry::MultiLineString { coordinates } => { + coordinates.iter().for_each(|line| self.add_path(line)); + } + Geometry::Polygon { coordinates } => self.add_polygon(coordinates), + Geometry::MultiPolygon { coordinates } => { + coordinates.iter().for_each(|rings| self.add_polygon(rings)); + } + Geometry::GeometryCollection { geometries } => { + geometries.iter().for_each(|g| self.add(g)); + } + // `Geometry` is `#[non_exhaustive]`; a kind added upstream exposes + // no coordinates this crate can read. + _ => {} + } + } + + fn add_point(&mut self, c: [f64; 2]) { + self.segments.push((c, c)); + self.vertices.push(c); + } + + fn add_path(&mut self, path: &[[f64; 2]]) { + match path { + [] => {} + [only] => self.add_point(*only), + _ => { + self.segments.extend(path.windows(2).map(|w| (w[0], w[1]))); + self.vertices.extend_from_slice(path); + } + } + } + + fn add_polygon(&mut self, rings: &'a [Vec<[f64; 2]>]) { + for ring in rings { + self.segments.extend(ring_edges(ring)); + self.vertices.extend_from_slice(ring); + } + self.polygons.push(rings); + } + + /// Whether any polygon of `self` covers any vertex of `other`. + fn covers_any_vertex_of(&self, other: &Parts<'_>) -> bool { + self.polygons + .iter() + .any(|rings| other.vertices.iter().any(|v| polygon_covers(rings, *v))) + } +} + +/// Whether `pt` lies inside or on a polygon, outside every hole interior. +fn polygon_covers(rings: &[Vec<[f64; 2]>], pt: [f64; 2]) -> bool { + let Some(exterior) = rings.first() else { + return false; + }; + if !point_in_polygon(pt[0], pt[1], exterior) && !point_on_ring_boundary(pt, exterior) { + return false; + } + !rings[1..] + .iter() + .any(|hole| point_in_polygon(pt[0], pt[1], hole) && !point_on_ring_boundary(pt, hole)) +} + +/// Minimum Euclidean distance between two geometries, in coordinate units. +/// +/// Zero when one geometry's polygon covers a vertex of the other, or when +/// any edges cross. Otherwise the minimum edge-to-edge distance. An empty +/// geometry is infinitely far from everything. +pub fn planar_distance(a: &Geometry, b: &Geometry) -> f64 { + let pa = Parts::of(a); + let pb = Parts::of(b); + if pa.covers_any_vertex_of(&pb) || pb.covers_any_vertex_of(&pa) { + return 0.0; + } + let mut min_sq = f64::INFINITY; + for &(a1, a2) in &pa.segments { + for &(b1, b2) in &pb.segments { + min_sq = min_sq.min(segment_to_segment_dist_sq(a1, a2, b1, b2)); + } + } + min_sq.sqrt() +} + +/// Whether two geometries lie within `distance` coordinate units. +/// +/// A planar bbox pre-filter rejects clearly distant pairs before the exact +/// distance runs. +pub fn planar_dwithin(a: &Geometry, b: &Geometry, distance: f64) -> bool { + let a_bb = geometry_bbox_in(a, BboxSpace::Planar).expand_meters_in(distance, BboxSpace::Planar); + let b_bb = geometry_bbox_in(b, BboxSpace::Planar); + if !a_bb.intersects(&b_bb) { + return false; + } + planar_distance(a, b) <= distance +} + +fn path_length(path: &[[f64; 2]]) -> f64 { + path.windows(2) + .map(|w| (w[1][0] - w[0][0]).hypot(w[1][1] - w[0][1])) + .sum() +} + +fn closed_ring_length(ring: &[[f64; 2]]) -> f64 { + ring_edges(ring) + .iter() + .map(|(a, b)| (b[0] - a[0]).hypot(b[1] - a[1])) + .sum() +} + +/// Total Euclidean length of every linear component. Areal geometries and +/// points have no length, matching PostGIS. +pub fn planar_length(geom: &Geometry) -> f64 { + match geom { + Geometry::LineString { coordinates } => path_length(coordinates), + Geometry::MultiLineString { coordinates } => { + coordinates.iter().map(|line| path_length(line)).sum() + } + Geometry::GeometryCollection { geometries } => geometries.iter().map(planar_length).sum(), + Geometry::Point { .. } + | Geometry::MultiPoint { .. } + | Geometry::Polygon { .. } + | Geometry::MultiPolygon { .. } => 0.0, + // `Geometry` is `#[non_exhaustive]`; a kind added upstream is not + // known to be linear. + _ => 0.0, + } +} + +/// Total Euclidean length of every areal component's rings, holes included. +/// An unclosed ring is closed implicitly. +pub fn planar_perimeter(geom: &Geometry) -> f64 { + match geom { + Geometry::Polygon { coordinates } => coordinates + .iter() + .map(|ring| closed_ring_length(ring)) + .sum(), + Geometry::MultiPolygon { coordinates } => coordinates + .iter() + .flat_map(|rings| rings.iter()) + .map(|ring| closed_ring_length(ring)) + .sum(), + Geometry::GeometryCollection { geometries } => { + geometries.iter().map(planar_perimeter).sum() + } + Geometry::Point { .. } + | Geometry::MultiPoint { .. } + | Geometry::LineString { .. } + | Geometry::MultiLineString { .. } => 0.0, + // `Geometry` is `#[non_exhaustive]`; a kind added upstream is not + // known to be areal. + _ => 0.0, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn square(x0: f64, y0: f64, side: f64) -> Geometry { + Geometry::polygon(vec![vec![ + [x0, y0], + [x0 + side, y0], + [x0 + side, y0 + side], + [x0, y0 + side], + [x0, y0], + ]]) + } + + #[test] + fn point_to_point_is_euclidean() { + let d = planar_distance(&Geometry::point(0.0, 0.0), &Geometry::point(300.0, 400.0)); + assert!((d - 500.0).abs() < 1e-9); + } + + #[test] + fn point_inside_polygon_is_zero() { + let d = planar_distance(&square(0.0, 0.0, 100.0), &Geometry::point(50.0, 50.0)); + assert_eq!(d, 0.0); + } + + #[test] + fn point_to_polygon_edge() { + let d = planar_distance(&square(0.0, 0.0, 100.0), &Geometry::point(150.0, 50.0)); + assert!((d - 50.0).abs() < 1e-9); + } + + #[test] + fn crossing_line_is_zero() { + let line = Geometry::line_string(vec![[-10.0, 50.0], [110.0, 50.0]]); + assert_eq!(planar_distance(&square(0.0, 0.0, 100.0), &line), 0.0); + } + + #[test] + fn point_in_hole_measures_to_hole_edge() { + let poly = Geometry::polygon(vec![ + vec![ + [0.0, 0.0], + [100.0, 0.0], + [100.0, 100.0], + [0.0, 100.0], + [0.0, 0.0], + ], + vec![ + [40.0, 40.0], + [60.0, 40.0], + [60.0, 60.0], + [40.0, 60.0], + [40.0, 40.0], + ], + ]); + let d = planar_distance(&poly, &Geometry::point(50.0, 50.0)); + assert!((d - 10.0).abs() < 1e-9); + } + + #[test] + fn dwithin_respects_threshold() { + let a = Geometry::point(500_000.0, 0.0); + let b = Geometry::point(500_300.0, 400.0); + assert!(planar_dwithin(&a, &b, 500.0)); + assert!(!planar_dwithin(&a, &b, 499.0)); + } + + #[test] + fn length_and_perimeter() { + let line = Geometry::line_string(vec![[0.0, 0.0], [3.0, 4.0], [3.0, 10.0]]); + assert!((planar_length(&line) - 11.0).abs() < 1e-9); + assert!((planar_perimeter(&square(0.0, 0.0, 100.0)) - 400.0).abs() < 1e-9); + assert_eq!(planar_length(&square(0.0, 0.0, 1.0)), 0.0); + } +} diff --git a/nodedb-spatial/src/planar/mod.rs b/nodedb-spatial/src/planar/mod.rs new file mode 100644 index 000000000..13d0d1848 --- /dev/null +++ b/nodedb-spatial/src/planar/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Euclidean measures for geometries in a projected SRID, in metres. + +pub mod area; +pub mod buffer; +pub mod distance; + +pub use area::{planar_area, planar_ring_area}; +pub use buffer::planar_buffer; +pub use distance::{planar_distance, planar_dwithin, planar_length, planar_perimeter}; diff --git a/nodedb-spatial/src/predicates/contains.rs b/nodedb-spatial/src/predicates/contains.rs index bd9fc741a..61ab25407 100644 --- a/nodedb-spatial/src/predicates/contains.rs +++ b/nodedb-spatial/src/predicates/contains.rs @@ -11,18 +11,22 @@ use nodedb_types::geometry::{Geometry, point_in_polygon}; use nodedb_types::geometry_bbox; -use super::edge::{point_on_ring_boundary, ring_edges}; +use super::edge::{is_collection, point_on_ring_boundary, ring_edges}; /// ST_Contains(a, b) — does geometry A fully contain geometry B? /// /// Returns false if any part of B lies outside A, or if B only touches /// A's boundary without entering its interior. pub fn st_contains(a: &Geometry, b: &Geometry) -> bool { - // Fast path: bbox pre-filter. - let a_bb = geometry_bbox(a); - let b_bb = geometry_bbox(b); - if !a_bb.contains_bbox(&b_bb) { - return false; + // Fast path: bbox pre-filter. A GeometryCollection skips it: its bbox is + // a union whose wrap rule depends on the SRID, and the arms below test + // each member on its own. + if !is_collection(a) && !is_collection(b) { + let a_bb = geometry_bbox(a); + let b_bb = geometry_bbox(b); + if !a_bb.contains_bbox(&b_bb) { + return false; + } } match (a, b) { diff --git a/nodedb-spatial/src/predicates/edge.rs b/nodedb-spatial/src/predicates/edge.rs index f485115ff..763360af6 100644 --- a/nodedb-spatial/src/predicates/edge.rs +++ b/nodedb-spatial/src/predicates/edge.rs @@ -16,6 +16,14 @@ pub enum Orientation { CounterClockwise, } +/// Whether `geom` is a GeometryCollection. +pub fn is_collection(geom: &nodedb_types::geometry::Geometry) -> bool { + matches!( + geom, + nodedb_types::geometry::Geometry::GeometryCollection { .. } + ) +} + /// Compute orientation of ordered triplet (p, q, r). pub fn orientation(p: [f64; 2], q: [f64; 2], r: [f64; 2]) -> Orientation { let val = (q[1] - p[1]) * (r[0] - q[0]) - (q[0] - p[0]) * (r[1] - q[1]); diff --git a/nodedb-spatial/src/predicates/intersects.rs b/nodedb-spatial/src/predicates/intersects.rs index 9582e9e2b..9bae4d4e9 100644 --- a/nodedb-spatial/src/predicates/intersects.rs +++ b/nodedb-spatial/src/predicates/intersects.rs @@ -11,15 +11,19 @@ use nodedb_types::geometry::{Geometry, point_in_polygon}; use nodedb_types::geometry_bbox; -use super::edge::{point_on_ring_boundary, ring_edges, segments_intersect}; +use super::edge::{is_collection, point_on_ring_boundary, ring_edges, segments_intersect}; /// ST_Intersects(a, b) — do geometries A and B share any space? pub fn st_intersects(a: &Geometry, b: &Geometry) -> bool { - // Bbox pre-filter. - let a_bb = geometry_bbox(a); - let b_bb = geometry_bbox(b); - if !a_bb.intersects(&b_bb) { - return false; + // Bbox pre-filter. A GeometryCollection skips it: its bbox is a union + // whose wrap rule depends on the SRID, and the recursion below filters + // each member on its own. + if !is_collection(a) && !is_collection(b) { + let a_bb = geometry_bbox(a); + let b_bb = geometry_bbox(b); + if !a_bb.intersects(&b_bb) { + return false; + } } match (a, b) { diff --git a/nodedb-spatial/src/rtree/bulk_load.rs b/nodedb-spatial/src/rtree/bulk_load.rs index 8d17d7b9b..1d669546f 100644 --- a/nodedb-spatial/src/rtree/bulk_load.rs +++ b/nodedb-spatial/src/rtree/bulk_load.rs @@ -3,19 +3,23 @@ //! Sort-Tile-Recursive (STR) bulk loading for R-tree. use nodedb_mem::ScopedMemory; +use nodedb_types::BboxSpace; +use nodedb_types::geometry::Srid; use super::node::{ChildRef, INTERNAL_CAPACITY, LEAF_CAPACITY, Node, NodeKind, RTreeEntry}; use super::tree::RTree; impl RTree { - /// Bulk load entries using Sort-Tile-Recursive packing, tracked against `memory`. + /// Bulk load entries in `srid` using Sort-Tile-Recursive packing, tracked + /// against `memory`. /// /// More efficient than repeated single inserts for large datasets. /// Produces better packing (less overlap between nodes). - pub fn bulk_load(entries: Vec, memory: ScopedMemory) -> Self { + pub fn bulk_load(entries: Vec, memory: ScopedMemory, srid: Srid) -> Self { if entries.is_empty() { - return Self::new(memory); + return Self::new(memory, srid); } + let space = srid.bbox_space(); let len = entries.len(); // Reserve budget for nodes vec (approximately len / LEAF_CAPACITY nodes, @@ -30,15 +34,17 @@ impl RTree { nodes: Vec::new(), root: 0, len, + srid, + space, memory, }; - tree.root = str_pack(&mut tree.nodes, entries); + tree.root = str_pack(&mut tree.nodes, entries, space); tree } } /// Recursively pack entries into leaf nodes, then group into internal nodes. -fn str_pack(nodes: &mut Vec, mut entries: Vec) -> usize { +fn str_pack(nodes: &mut Vec, mut entries: Vec, space: BboxSpace) -> usize { if entries.len() <= LEAF_CAPACITY { let mut node = Node::new_leaf(); if let NodeKind::Leaf { @@ -47,7 +53,7 @@ fn str_pack(nodes: &mut Vec, mut entries: Vec) -> usize { { *leaf = entries; } - node.recompute_bbox(); + node.recompute_bbox(space); let idx = nodes.len(); nodes.push(node); return idx; @@ -76,16 +82,16 @@ fn str_pack(nodes: &mut Vec, mut entries: Vec) -> usize { }); for chunk in slice_vec.chunks(LEAF_CAPACITY) { - let child_idx = str_pack(nodes, chunk.to_vec()); + let child_idx = str_pack(nodes, chunk.to_vec(), space); child_nodes.push(child_idx); } } - pack_internal(nodes, child_nodes) + pack_internal(nodes, child_nodes, space) } /// Recursively group child node indices into internal nodes. -fn pack_internal(nodes: &mut Vec, child_indices: Vec) -> usize { +fn pack_internal(nodes: &mut Vec, child_indices: Vec, space: BboxSpace) -> usize { if child_indices.len() <= INTERNAL_CAPACITY { let level = nodes[child_indices[0]].level + 1; let mut node = Node::new_internal(level); @@ -97,7 +103,7 @@ fn pack_internal(nodes: &mut Vec, child_indices: Vec) -> usize { }); } } - node.recompute_bbox(); + node.recompute_bbox(space); let idx = nodes.len(); nodes.push(node); return idx; @@ -115,13 +121,13 @@ fn pack_internal(nodes: &mut Vec, child_indices: Vec) -> usize { }); } } - node.recompute_bbox(); + node.recompute_bbox(space); let idx = nodes.len(); nodes.push(node); new_children.push(idx); } - pack_internal(nodes, new_children) + pack_internal(nodes, new_children, space) } #[cfg(test)] @@ -142,7 +148,7 @@ mod tests { let entries: Vec = (0..500) .map(|i| make_entry(i, (i as f64) * 0.1 - 25.0, (i as f64) * 0.06 - 15.0)) .collect(); - let tree = RTree::bulk_load(entries, test_memory()); + let tree = RTree::bulk_load(entries, test_memory(), Srid::Wgs84); assert_eq!(tree.len(), 500); let all = tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0)); @@ -155,13 +161,22 @@ mod tests { #[test] fn bulk_load_small() { let entries = vec![make_entry(1, 0.0, 0.0), make_entry(2, 1.0, 1.0)]; - let tree = RTree::bulk_load(entries, test_memory()); + let tree = RTree::bulk_load(entries, test_memory(), Srid::Wgs84); assert_eq!(tree.len(), 2); } #[test] fn bulk_load_empty() { - let tree = RTree::bulk_load(Vec::new(), test_memory()); + let tree = RTree::bulk_load(Vec::new(), test_memory(), Srid::Utm47N); assert!(tree.is_empty()); + assert_eq!(tree.srid(), Srid::Utm47N); + } + + #[test] + fn bulk_load_keeps_srid_and_planar_space() { + let entries = vec![make_entry(1, 1e5, 0.0), make_entry(2, 8e5, 0.0)]; + let tree = RTree::bulk_load(entries, test_memory(), Srid::Utm48N); + assert_eq!(tree.srid(), Srid::Utm48N); + assert_eq!(tree.search(&BoundingBox::new(7e5, -1.0, 9e5, 1.0)).len(), 1); } } diff --git a/nodedb-spatial/src/rtree/delete.rs b/nodedb-spatial/src/rtree/delete.rs index 49bdc5a94..e707fb4d4 100644 --- a/nodedb-spatial/src/rtree/delete.rs +++ b/nodedb-spatial/src/rtree/delete.rs @@ -2,13 +2,15 @@ //! R-tree entry deletion with underflow handling. +use nodedb_types::BboxSpace; + use super::node::{EntryId, NodeKind, RTreeEntry}; use super::tree::{RTree, collect_entries_owned}; impl RTree { /// Delete an entry by ID. Returns true if found and removed. pub fn delete(&mut self, id: EntryId) -> bool { - let removed = delete_entry(&mut self.nodes, self.root, id); + let removed = delete_entry(&mut self.nodes, self.root, id, self.space); if let Some(orphans) = removed { self.len -= 1; for entry in orphans { @@ -27,6 +29,7 @@ fn delete_entry( nodes: &mut Vec, node_idx: usize, id: EntryId, + space: BboxSpace, ) -> Option> { let is_leaf = nodes[node_idx].is_leaf(); @@ -37,7 +40,7 @@ fn delete_entry( if let NodeKind::Leaf { entries } = &mut nodes[node_idx].kind { entries.remove(pos); } - nodes[node_idx].recompute_bbox(); + nodes[node_idx].recompute_bbox(space); return Some(Vec::new()); } } @@ -52,7 +55,7 @@ fn delete_entry( }; for child_idx in child_indices { - if let Some(mut orphans) = delete_entry(nodes, child_idx, id) { + if let Some(mut orphans) = delete_entry(nodes, child_idx, id, space) { // Update child bbox in parent. let child_bbox = nodes[child_idx].bbox; if let NodeKind::Internal { children } = &mut nodes[node_idx].kind { @@ -63,7 +66,7 @@ fn delete_entry( } } } - nodes[node_idx].recompute_bbox(); + nodes[node_idx].recompute_bbox(space); // Check underflow. if nodes[child_idx].is_underflow() { @@ -72,7 +75,7 @@ fn delete_entry( if let NodeKind::Internal { children } = &mut nodes[node_idx].kind { children.retain(|c| c.node_idx != child_idx); } - nodes[node_idx].recompute_bbox(); + nodes[node_idx].recompute_bbox(space); orphans.extend(collected); } diff --git a/nodedb-spatial/src/rtree/insert.rs b/nodedb-spatial/src/rtree/insert.rs index eeecda036..26321f023 100644 --- a/nodedb-spatial/src/rtree/insert.rs +++ b/nodedb-spatial/src/rtree/insert.rs @@ -2,7 +2,7 @@ //! R*-tree insertion with overflow treatment (forced reinsert). -use nodedb_types::BoundingBox; +use nodedb_types::{BboxSpace, BoundingBox}; use super::node::{ChildRef, NodeKind, REINSERT_COUNT_LEAF, RTreeEntry}; use super::split::split_node; @@ -15,7 +15,7 @@ impl RTree { let mut reinserted_levels = Vec::new(); insert_entry(self, entry, 0, &mut reinserted_levels); self.len += 1; - self.nodes[self.root].bbox = self.nodes[self.root].bbox.union(&bbox); + self.nodes[self.root].bbox = self.nodes[self.root].bbox.union_in(&bbox, self.space); } /// Insert with reinsert tracking (used internally and by delete). @@ -43,7 +43,7 @@ fn insert_entry( return; } } - tree.nodes[leaf_idx].recompute_bbox(); + tree.nodes[leaf_idx].recompute_bbox(tree.space); if tree.nodes[leaf_idx].is_overflow() { treat_overflow(tree, leaf_idx, reinserted_levels); @@ -69,21 +69,25 @@ fn choose_subtree( return node_idx; } let best = if tree.nodes[children[0].node_idx].is_leaf() { - choose_least_overlap(children, entry_bbox) + choose_least_overlap(children, entry_bbox, tree.space) } else { - choose_least_enlargement(children, entry_bbox) + choose_least_enlargement(children, entry_bbox, tree.space) }; choose_subtree(tree, children[best].node_idx, entry_bbox, target_level) } } } -fn choose_least_enlargement(children: &[ChildRef], entry_bbox: &BoundingBox) -> usize { +fn choose_least_enlargement( + children: &[ChildRef], + entry_bbox: &BoundingBox, + space: BboxSpace, +) -> usize { let mut best = 0; let mut best_enlarge = f64::INFINITY; let mut best_area = f64::INFINITY; for (i, child) in children.iter().enumerate() { - let enlarge = child.bbox.enlargement(entry_bbox); + let enlarge = child.bbox.enlargement_in(entry_bbox, space); let area = child.bbox.area(); if enlarge < best_enlarge || (enlarge == best_enlarge && area < best_area) { best = i; @@ -94,13 +98,17 @@ fn choose_least_enlargement(children: &[ChildRef], entry_bbox: &BoundingBox) -> best } -fn choose_least_overlap(children: &[ChildRef], entry_bbox: &BoundingBox) -> usize { +fn choose_least_overlap( + children: &[ChildRef], + entry_bbox: &BoundingBox, + space: BboxSpace, +) -> usize { let mut best = 0; let mut best_oi = f64::INFINITY; let mut best_enlarge = f64::INFINITY; let mut best_area = f64::INFINITY; for (i, child) in children.iter().enumerate() { - let enlarged = child.bbox.union(entry_bbox); + let enlarged = child.bbox.union_in(entry_bbox, space); let mut before = 0.0_f64; let mut after = 0.0_f64; for (j, other) in children.iter().enumerate() { @@ -110,7 +118,7 @@ fn choose_least_overlap(children: &[ChildRef], entry_bbox: &BoundingBox) -> usiz } } let oi = after - before; - let enlarge = child.bbox.enlargement(entry_bbox); + let enlarge = child.bbox.enlargement_in(entry_bbox, space); let area = child.bbox.area(); if oi < best_oi || (oi == best_oi && enlarge < best_enlarge) @@ -160,7 +168,7 @@ fn forced_reinsert(tree: &mut RTree, node_idx: usize) -> Vec { db.partial_cmp(&da).unwrap_or(std::cmp::Ordering::Equal) }); let removed: Vec = entries.drain(..reinsert_count).collect(); - tree.nodes[node_idx].recompute_bbox(); + tree.nodes[node_idx].recompute_bbox(tree.space); removed } else { Vec::new() diff --git a/nodedb-spatial/src/rtree/mod.rs b/nodedb-spatial/src/rtree/mod.rs index 14b7a57fc..774159861 100644 --- a/nodedb-spatial/src/rtree/mod.rs +++ b/nodedb-spatial/src/rtree/mod.rs @@ -10,4 +10,4 @@ pub mod tree; pub use node::{EntryId, RTreeEntry}; pub use search::NnResult; -pub use tree::RTree; +pub use tree::{RTree, SridMismatch}; diff --git a/nodedb-spatial/src/rtree/node.rs b/nodedb-spatial/src/rtree/node.rs index cef1da77d..ace19b3ce 100644 --- a/nodedb-spatial/src/rtree/node.rs +++ b/nodedb-spatial/src/rtree/node.rs @@ -2,7 +2,7 @@ //! R-tree node types and constants. -use nodedb_types::BoundingBox; +use nodedb_types::{BboxSpace, BoundingBox}; use serde::{Deserialize, Serialize}; use zerompk::{FromMessagePack, ToMessagePack}; @@ -119,7 +119,8 @@ impl Node { self.len() < self.min_fill() } - pub fn recompute_bbox(&mut self) { + /// Recompute this node's bbox from its children, unioned in `space`. + pub fn recompute_bbox(&mut self, space: BboxSpace) { match &self.kind { NodeKind::Internal { children } => { if children.is_empty() { @@ -128,7 +129,7 @@ impl Node { } let mut bb = children[0].bbox; for c in &children[1..] { - bb = bb.union(&c.bbox); + bb = bb.union_in(&c.bbox, space); } self.bbox = bb; } @@ -139,7 +140,7 @@ impl Node { } let mut bb = entries[0].bbox; for e in &entries[1..] { - bb = bb.union(&e.bbox); + bb = bb.union_in(&e.bbox, space); } self.bbox = bb; } diff --git a/nodedb-spatial/src/rtree/search.rs b/nodedb-spatial/src/rtree/search.rs index 496f01869..768d61848 100644 --- a/nodedb-spatial/src/rtree/search.rs +++ b/nodedb-spatial/src/rtree/search.rs @@ -2,7 +2,7 @@ //! R-tree range search and nearest-neighbor queries. -use nodedb_types::BoundingBox; +use nodedb_types::{BboxSpace, BoundingBox}; use std::collections::BinaryHeap; use super::node::{EntryId, Node, NodeKind, RTreeEntry}; @@ -12,7 +12,8 @@ use super::node::{EntryId, Node, NodeKind, RTreeEntry}; pub struct NnResult { pub entry_id: EntryId, pub bbox: BoundingBox, - /// Minimum distance in degrees (approximate). + /// Minimum distance in the tree's coordinate units: approximate degrees + /// for a geographic tree, metres for a projected one. pub distance: f64, } @@ -45,15 +46,27 @@ pub(crate) fn search_node<'a>( } } +/// A nearest-neighbor query point and the space its distances are taken in. +pub(crate) struct NnQuery { + pub lng: f64, + pub lat: f64, + pub k: usize, + pub space: BboxSpace, +} + /// Nearest-neighbor search via priority queue (min-heap). pub(crate) fn nearest( nodes: &[Node], root: usize, - query_lng: f64, - query_lat: f64, - k: usize, + query: NnQuery, is_empty: bool, ) -> Vec { + let NnQuery { + lng: query_lng, + lat: query_lat, + k, + space, + } = query; if k == 0 || is_empty { return Vec::new(); } @@ -62,7 +75,7 @@ pub(crate) fn nearest( let mut results: Vec = Vec::with_capacity(k); heap.push(HeapItem { - dist: min_dist_point_bbox(query_lng, query_lat, &nodes[root].bbox), + dist: min_dist_point_bbox(query_lng, query_lat, &nodes[root].bbox, space), node_idx: root, }); @@ -74,7 +87,7 @@ pub(crate) fn nearest( match &node.kind { NodeKind::Internal { children } => { for child in children { - let d = min_dist_point_bbox(query_lng, query_lat, &child.bbox); + let d = min_dist_point_bbox(query_lng, query_lat, &child.bbox, space); if results.len() < k || d <= results[results.len() - 1].distance { heap.push(HeapItem { dist: d, @@ -85,7 +98,7 @@ pub(crate) fn nearest( } NodeKind::Leaf { entries } => { for entry in entries { - let d = min_dist_point_bbox(query_lng, query_lat, &entry.bbox); + let d = min_dist_point_bbox(query_lng, query_lat, &entry.bbox, space); if results.len() < k || d < results[results.len() - 1].distance { let nn = NnResult { entry_id: entry.id, @@ -132,8 +145,10 @@ impl Ord for HeapItem { } } -/// Minimum distance from point to bbox (in degrees, approximate). -fn min_dist_point_bbox(lng: f64, lat: f64, bbox: &BoundingBox) -> f64 { +/// Minimum Euclidean distance from point to bbox in coordinate units: +/// approximate degrees in geographic space, metres in planar space. Only +/// geographic space honours an anti-meridian-wrapping box. +fn min_dist_point_bbox(lng: f64, lat: f64, bbox: &BoundingBox, space: BboxSpace) -> f64 { let dlat = if lat < bbox.min_lat { bbox.min_lat - lat } else if lat > bbox.max_lat { @@ -142,7 +157,11 @@ fn min_dist_point_bbox(lng: f64, lat: f64, bbox: &BoundingBox) -> f64 { 0.0 }; - let dlng = if bbox.crosses_antimeridian() { + let wraps = match space { + BboxSpace::Geographic => bbox.crosses_antimeridian(), + BboxSpace::Planar => false, + }; + let dlng = if wraps { if lng >= bbox.min_lng || lng <= bbox.max_lng { 0.0 } else { diff --git a/nodedb-spatial/src/rtree/split.rs b/nodedb-spatial/src/rtree/split.rs index cd710e566..38cba4b3b 100644 --- a/nodedb-spatial/src/rtree/split.rs +++ b/nodedb-spatial/src/rtree/split.rs @@ -2,6 +2,8 @@ //! R*-tree node splitting (axis split strategy). +use nodedb_types::BboxSpace; + use super::node::{ChildRef, MIN_FILL_INTERNAL, MIN_FILL_LEAF, Node, NodeKind, RTreeEntry}; use super::tree::RTree; @@ -9,21 +11,22 @@ use super::tree::RTree; pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) { let is_root = node_idx == tree.root; let level = tree.nodes[node_idx].level; + let space = tree.space; let (sibling_idx, sibling_bbox) = match &mut tree.nodes[node_idx].kind { NodeKind::Leaf { entries } => { let all = std::mem::take(entries); - let (keep, split_off) = split_leaf_entries(all); + let (keep, split_off) = split_leaf_entries(all, space); if let NodeKind::Leaf { entries } = &mut tree.nodes[node_idx].kind { *entries = keep; } - tree.nodes[node_idx].recompute_bbox(); + tree.nodes[node_idx].recompute_bbox(space); let mut sibling = Node::new_leaf(); if let NodeKind::Leaf { entries } = &mut sibling.kind { *entries = split_off; } - sibling.recompute_bbox(); + sibling.recompute_bbox(space); let bbox = sibling.bbox; let idx = tree.nodes.len(); tree.nodes.push(sibling); @@ -31,17 +34,17 @@ pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) { } NodeKind::Internal { children } => { let all = std::mem::take(children); - let (keep, split_off) = split_internal_children(all); + let (keep, split_off) = split_internal_children(all, space); if let NodeKind::Internal { children } = &mut tree.nodes[node_idx].kind { *children = keep; } - tree.nodes[node_idx].recompute_bbox(); + tree.nodes[node_idx].recompute_bbox(space); let mut sibling = Node::new_internal(level); if let NodeKind::Internal { children } = &mut sibling.kind { *children = split_off; } - sibling.recompute_bbox(); + sibling.recompute_bbox(space); let bbox = sibling.bbox; let idx = tree.nodes.len(); tree.nodes.push(sibling); @@ -62,7 +65,7 @@ pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) { node_idx: sibling_idx, }); } - new_root.recompute_bbox(); + new_root.recompute_bbox(space); let new_root_idx = tree.nodes.len(); tree.nodes.push(new_root); tree.root = new_root_idx; @@ -82,7 +85,7 @@ pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) { } } } - tree.nodes[pidx].recompute_bbox(); + tree.nodes[pidx].recompute_bbox(space); if tree.nodes[pidx].is_overflow() { split_node(tree, pidx); } @@ -91,29 +94,35 @@ pub(crate) fn split_node(tree: &mut RTree, node_idx: usize) { } /// Choose split axis (min margin sum), then split index (min overlap). -fn split_leaf_entries(mut entries: Vec) -> (Vec, Vec) { +fn split_leaf_entries( + mut entries: Vec, + space: BboxSpace, +) -> (Vec, Vec) { let min_fill = MIN_FILL_LEAF; - let best_axis = choose_best_axis_leaf(&mut entries, min_fill); + let best_axis = choose_best_axis_leaf(&mut entries, min_fill, space); sort_entries_by_axis(&mut entries, best_axis); - let best_k = choose_best_split_leaf(&entries, min_fill); + let best_k = choose_best_split_leaf(&entries, min_fill, space); let split_off = entries.split_off(best_k); (entries, split_off) } -fn split_internal_children(mut children: Vec) -> (Vec, Vec) { +fn split_internal_children( + mut children: Vec, + space: BboxSpace, +) -> (Vec, Vec) { let min_fill = MIN_FILL_INTERNAL; - let best_axis = choose_best_axis_internal(&mut children, min_fill); + let best_axis = choose_best_axis_internal(&mut children, min_fill, space); sort_children_by_axis(&mut children, best_axis); - let best_k = choose_best_split_internal(&children, min_fill); + let best_k = choose_best_split_internal(&children, min_fill, space); let split_off = children.split_off(best_k); (children, split_off) } -fn choose_best_axis_leaf(entries: &mut [RTreeEntry], min_fill: usize) -> usize { +fn choose_best_axis_leaf(entries: &mut [RTreeEntry], min_fill: usize, space: BboxSpace) -> usize { let mut best_axis = 0; let mut best_margin = f64::INFINITY; for axis in 0..2 { @@ -122,10 +131,10 @@ fn choose_best_axis_leaf(entries: &mut [RTreeEntry], min_fill: usize) -> usize { for k in min_fill..=(entries.len() - min_fill) { let left = entries[..k] .iter() - .fold(entries[0].bbox, |a, e| a.union(&e.bbox)); + .fold(entries[0].bbox, |a, e| a.union_in(&e.bbox, space)); let right = entries[k..] .iter() - .fold(entries[k].bbox, |a, e| a.union(&e.bbox)); + .fold(entries[k].bbox, |a, e| a.union_in(&e.bbox, space)); margin_sum += left.margin() + right.margin(); } if margin_sum < best_margin { @@ -136,7 +145,11 @@ fn choose_best_axis_leaf(entries: &mut [RTreeEntry], min_fill: usize) -> usize { best_axis } -fn choose_best_axis_internal(children: &mut [ChildRef], min_fill: usize) -> usize { +fn choose_best_axis_internal( + children: &mut [ChildRef], + min_fill: usize, + space: BboxSpace, +) -> usize { let mut best_axis = 0; let mut best_margin = f64::INFINITY; for axis in 0..2 { @@ -145,10 +158,10 @@ fn choose_best_axis_internal(children: &mut [ChildRef], min_fill: usize) -> usiz for k in min_fill..=(children.len() - min_fill) { let left = children[..k] .iter() - .fold(children[0].bbox, |a, c| a.union(&c.bbox)); + .fold(children[0].bbox, |a, c| a.union_in(&c.bbox, space)); let right = children[k..] .iter() - .fold(children[k].bbox, |a, c| a.union(&c.bbox)); + .fold(children[k].bbox, |a, c| a.union_in(&c.bbox, space)); margin_sum += left.margin() + right.margin(); } if margin_sum < best_margin { @@ -159,17 +172,17 @@ fn choose_best_axis_internal(children: &mut [ChildRef], min_fill: usize) -> usiz best_axis } -fn choose_best_split_leaf(entries: &[RTreeEntry], min_fill: usize) -> usize { +fn choose_best_split_leaf(entries: &[RTreeEntry], min_fill: usize, space: BboxSpace) -> usize { let mut best_k = min_fill; let mut best_overlap = f64::INFINITY; let mut best_area = f64::INFINITY; for k in min_fill..=(entries.len() - min_fill) { let left = entries[..k] .iter() - .fold(entries[0].bbox, |a, e| a.union(&e.bbox)); + .fold(entries[0].bbox, |a, e| a.union_in(&e.bbox, space)); let right = entries[k..] .iter() - .fold(entries[k].bbox, |a, e| a.union(&e.bbox)); + .fold(entries[k].bbox, |a, e| a.union_in(&e.bbox, space)); let overlap = left.overlap_area(&right); let area = left.area() + right.area(); if overlap < best_overlap || (overlap == best_overlap && area < best_area) { @@ -181,17 +194,17 @@ fn choose_best_split_leaf(entries: &[RTreeEntry], min_fill: usize) -> usize { best_k } -fn choose_best_split_internal(children: &[ChildRef], min_fill: usize) -> usize { +fn choose_best_split_internal(children: &[ChildRef], min_fill: usize, space: BboxSpace) -> usize { let mut best_k = min_fill; let mut best_overlap = f64::INFINITY; let mut best_area = f64::INFINITY; for k in min_fill..=(children.len() - min_fill) { let left = children[..k] .iter() - .fold(children[0].bbox, |a, c| a.union(&c.bbox)); + .fold(children[0].bbox, |a, c| a.union_in(&c.bbox, space)); let right = children[k..] .iter() - .fold(children[k].bbox, |a, c| a.union(&c.bbox)); + .fold(children[k].bbox, |a, c| a.union_in(&c.bbox, space)); let overlap = left.overlap_area(&right); let area = left.area() + right.area(); if overlap < best_overlap || (overlap == best_overlap && area < best_area) { diff --git a/nodedb-spatial/src/rtree/tree.rs b/nodedb-spatial/src/rtree/tree.rs index 6723c48ec..f607c46c0 100644 --- a/nodedb-spatial/src/rtree/tree.rs +++ b/nodedb-spatial/src/rtree/tree.rs @@ -2,7 +2,8 @@ //! R*-tree public API and core structure. -use nodedb_types::BoundingBox; +use nodedb_types::geometry::Srid; +use nodedb_types::{BboxSpace, BoundingBox}; use super::node::{Node, NodeKind, RTreeEntry}; use super::search::NnResult; @@ -14,6 +15,9 @@ use nodedb_mem::ScopedMemory; /// incremental nearest-neighbor queries. Array-backed nodes for cache /// friendliness. /// +/// Every entry of one tree is in one SRID. The tree takes it from its first +/// entry, and it fixes the bbox space every union and enlargement uses. +/// /// References: /// - Beckmann et al., "The R*-tree" (1990) /// - Hjaltason & Samet, "Distance Browsing in Spatial Databases" (1999) @@ -21,17 +25,32 @@ pub struct RTree { pub(crate) nodes: Vec, pub(crate) root: usize, pub(crate) len: usize, + /// SRID of every entry. + pub(crate) srid: Srid, + /// Bbox space derived from `srid`, cached for the hot insert path. + pub(crate) space: BboxSpace, /// Enforces database, tenant, and engine memory budgets on large batch /// allocations (bulk load, full-scan serialization, range search results). pub(crate) memory: ScopedMemory, } +/// A geometry in one SRID offered to an R-tree holding another. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("R-tree holds SRID {index}, entry is SRID {entry}")] +pub struct SridMismatch { + pub index: Srid, + pub entry: Srid, +} + impl RTree { - pub fn new(memory: ScopedMemory) -> Self { + /// An empty tree for entries in `srid`. + pub fn new(memory: ScopedMemory, srid: Srid) -> Self { Self { nodes: vec![Node::new_leaf()], root: 0, len: 0, + srid, + space: srid.bbox_space(), memory, } } @@ -41,6 +60,42 @@ impl RTree { &self.memory } + /// The SRID every entry is in. + pub fn srid(&self) -> Srid { + self.srid + } + + /// The bbox space unions and enlargements use. + pub fn space(&self) -> BboxSpace { + self.space + } + + /// Whether an entry in `srid` can join this tree. An empty tree accepts + /// any SRID, since its first entry decides. + pub fn accepts_srid(&self, srid: Srid) -> bool { + self.is_empty() || self.srid == srid + } + + /// Fix the tree's SRID for an entry about to be inserted. + /// + /// An empty tree adopts `srid`. A non-empty tree in another SRID refuses + /// with [`SridMismatch`] and is left unchanged. + pub fn bind_srid(&mut self, srid: Srid) -> Result<(), SridMismatch> { + if self.is_empty() { + self.srid = srid; + self.space = srid.bbox_space(); + return Ok(()); + } + if self.srid == srid { + Ok(()) + } else { + Err(SridMismatch { + index: self.srid, + entry: srid, + }) + } + } + pub fn len(&self) -> usize { self.len } @@ -62,13 +117,19 @@ impl RTree { } /// Nearest-neighbor search using incremental distance ordering. + /// + /// Distance is in the tree's coordinate units: degrees for a geographic + /// tree, metres (Euclidean) for a projected one. pub fn nearest(&self, query_lng: f64, query_lat: f64, k: usize) -> Vec { super::search::nearest( &self.nodes, self.root, - query_lng, - query_lat, - k, + super::search::NnQuery { + lng: query_lng, + lat: query_lat, + k, + space: self.space, + }, self.is_empty(), ) } @@ -151,7 +212,7 @@ mod tests { #[test] fn empty_tree() { - let tree = RTree::new(test_memory()); + let tree = RTree::new(test_memory(), Srid::Wgs84); assert!(tree.is_empty()); assert!( tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0)) @@ -161,7 +222,7 @@ mod tests { #[test] fn insert_and_search_single() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 10.0, 20.0)); assert_eq!(tree.len(), 1); assert_eq!( @@ -176,7 +237,7 @@ mod tests { #[test] fn insert_many_and_search() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..200 { tree.insert(make_entry( i, @@ -191,7 +252,7 @@ mod tests { #[test] fn delete_entry() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..50 { tree.insert(make_entry(i, i as f64, i as f64)); } @@ -204,7 +265,7 @@ mod tests { #[test] fn nearest_neighbor() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_entry(1, 0.0, 0.0)); tree.insert(make_entry(2, 10.0, 10.0)); tree.insert(make_entry(3, 5.0, 5.0)); @@ -216,12 +277,16 @@ mod tests { #[test] fn nearest_empty() { - assert!(RTree::new(test_memory()).nearest(0.0, 0.0, 5).is_empty()); + assert!( + RTree::new(test_memory(), Srid::Wgs84) + .nearest(0.0, 0.0, 5) + .is_empty() + ); } #[test] fn rect_overlap_search() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); tree.insert(make_rect(1, 0.0, 0.0, 10.0, 10.0)); tree.insert(make_rect(2, 5.0, 5.0, 15.0, 15.0)); tree.insert(make_rect(3, 20.0, 20.0, 30.0, 30.0)); @@ -231,7 +296,7 @@ mod tests { #[test] fn stress_insert_delete() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..100_u64 { tree.insert(make_entry(i, (i as f64) * 0.5, (i as f64) * 0.3)); } @@ -245,7 +310,7 @@ mod tests { #[test] fn triggers_node_split() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); let count = LEAF_CAPACITY * 3; for i in 0..count as u64 { tree.insert(make_entry(i, (i as f64) * 0.1, (i as f64) * 0.1)); @@ -255,9 +320,59 @@ mod tests { assert_eq!(all.len(), count); } + /// Disjoint metre boxes in a projected tree must union planar: the root + /// spans them directly instead of wrapping as if x were longitude. + #[test] + fn projected_tree_unions_planar() { + let mut tree = RTree::new(test_memory(), Srid::Utm48N); + tree.insert(make_rect(1, 1e5, 0.0, 2e5, 10.0)); + tree.insert(make_rect(2, 7e5, 0.0, 8e5, 10.0)); + let root = tree.nodes[tree.root].bbox; + assert_eq!((root.min_lng, root.max_lng), (1e5, 8e5)); + assert!(!root.crosses_antimeridian()); + assert_eq!( + tree.search(&BoundingBox::new(1.5e5, 0.0, 1.6e5, 10.0)) + .len(), + 1 + ); + assert!( + tree.search(&BoundingBox::new(4e5, 0.0, 5e5, 10.0)) + .is_empty() + ); + } + + #[test] + fn projected_nearest_is_euclidean_metres() { + let mut tree = RTree::new(test_memory(), Srid::Utm48N); + tree.insert(make_entry(1, 500_000.0, 0.0)); + tree.insert(make_entry(2, 500_300.0, 400.0)); + let nn = tree.nearest(500_000.0, 0.0, 2); + assert_eq!(nn[0].entry_id, 1); + assert!((nn[1].distance - 500.0).abs() < 1e-9); + } + + #[test] + fn empty_tree_adopts_first_srid_then_refuses_another() { + let mut tree = RTree::new(test_memory(), Srid::Wgs84); + assert!(tree.accepts_srid(Srid::Utm48N)); + assert_eq!(tree.bind_srid(Srid::Utm48N), Ok(())); + tree.insert(make_entry(1, 500_000.0, 0.0)); + assert_eq!(tree.srid(), Srid::Utm48N); + assert_eq!(tree.space(), BboxSpace::Planar); + assert!(!tree.accepts_srid(Srid::Wgs84)); + assert_eq!( + tree.bind_srid(Srid::Wgs84), + Err(SridMismatch { + index: Srid::Utm48N, + entry: Srid::Wgs84 + }) + ); + assert_eq!(tree.len(), 1); + } + #[test] fn entries_enumeration() { - let mut tree = RTree::new(test_memory()); + let mut tree = RTree::new(test_memory(), Srid::Wgs84); for i in 0..10 { tree.insert(make_entry(i, i as f64, i as f64)); } diff --git a/nodedb-spatial/src/spatial_join.rs b/nodedb-spatial/src/spatial_join.rs index bcbcec8f6..cd999c9ed 100644 --- a/nodedb-spatial/src/spatial_join.rs +++ b/nodedb-spatial/src/spatial_join.rs @@ -14,8 +14,8 @@ //! to index based on collection cardinality. use nodedb_mem::ScopedMemory; -use nodedb_types::geometry::Geometry; -use nodedb_types::{BoundingBox, geometry_bbox}; +use nodedb_types::geometry::{Geometry, Srid}; +use nodedb_types::{BboxSpace, BoundingBox, geometry_bbox_in}; use crate::predicates; use crate::rtree::{RTree, RTreeEntry}; @@ -71,9 +71,14 @@ pub fn spatial_join( predicates::st_contains(&probe_geom, &indexed_geom) } SpatialJoinPredicate::Within => predicates::st_within(&probe_geom, &indexed_geom), - SpatialJoinPredicate::DWithin(dist) => { - predicates::st_dwithin(&probe_geom, &indexed_geom, dist) - } + SpatialJoinPredicate::DWithin(dist) => match indexed_side.space() { + BboxSpace::Geographic => { + predicates::st_dwithin(&probe_geom, &indexed_geom, dist) + } + BboxSpace::Planar => { + crate::planar::planar_dwithin(&probe_geom, &indexed_geom, dist) + } + }, }; if matches { @@ -89,16 +94,18 @@ pub fn spatial_join( } } -/// Build an R-tree from a list of (entry_id, geometry) pairs for join. -pub fn build_join_index(entries: &[(u64, Geometry)], memory: ScopedMemory) -> RTree { +/// Build an R-tree from a list of (entry_id, geometry) pairs, all in +/// `srid`, for join. +pub fn build_join_index(entries: &[(u64, Geometry)], memory: ScopedMemory, srid: Srid) -> RTree { + let space = srid.bbox_space(); let rtree_entries: Vec = entries .iter() .map(|(id, geom)| RTreeEntry { id: *id, - bbox: geometry_bbox(geom), + bbox: geometry_bbox_in(geom, space), }) .collect(); - RTree::bulk_load(rtree_entries, memory) + RTree::bulk_load(rtree_entries, memory, srid) } /// Which spatial predicate to use for the join. @@ -115,6 +122,7 @@ pub enum SpatialJoinPredicate { mod tests { use super::*; use crate::test_support::test_memory; + use nodedb_types::geometry_bbox; #[test] fn join_overlapping_squares() { @@ -153,7 +161,7 @@ mod tests { .collect(); // Build index on left side. - let index = build_join_index(&left, test_memory()); + let index = build_join_index(&left, test_memory(), Srid::Wgs84); // Probe with right side. let probe_entries: Vec<(u64, BoundingBox)> = right @@ -200,7 +208,7 @@ mod tests { ]]), )]; - let index = build_join_index(&left, test_memory()); + let index = build_join_index(&left, test_memory(), Srid::Wgs84); let probes: Vec<(u64, BoundingBox)> = right .iter() .map(|(id, g)| (*id, geometry_bbox(g))) @@ -225,7 +233,7 @@ mod tests { let left = vec![(1u64, Geometry::point(0.0, 0.0))]; let right = vec![(100u64, Geometry::point(0.001, 0.0))]; // ~111m away - let index = build_join_index(&left, test_memory()); + let index = build_join_index(&left, test_memory(), Srid::Wgs84); let probes: Vec<(u64, BoundingBox)> = right .iter() .map(|(id, g)| { diff --git a/nodedb-spatial/src/wkb.rs b/nodedb-spatial/src/wkb.rs deleted file mode 100644 index 06777fccc..000000000 --- a/nodedb-spatial/src/wkb.rs +++ /dev/null @@ -1,634 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 - -//! Well-Known Binary (WKB) serialization for Geometry types. -//! -//! WKB is the standard binary format for geometry interchange (ISO 13249). -//! Used as the Arrow `DataType::Binary` backing for spatial columns — -//! avoids JSON parse overhead during DataFusion query execution. -//! -//! Format (little-endian): -//! ```text -//! [byte_order: u8] [type: u32] [coordinates...] -//! ``` -//! -//! Byte order: 1 = little-endian (NDR), 0 = big-endian (XDR). We always -//! write little-endian and accept both on read. - -use nodedb_types::geometry::Geometry; - -// WKB geometry type codes (ISO 13249 / OGC SFA). -const WKB_POINT: u32 = 1; -const WKB_LINESTRING: u32 = 2; -const WKB_POLYGON: u32 = 3; -const WKB_MULTIPOINT: u32 = 4; -const WKB_MULTILINESTRING: u32 = 5; -const WKB_MULTIPOLYGON: u32 = 6; -const WKB_GEOMETRYCOLLECTION: u32 = 7; - -const BYTE_ORDER_LE: u8 = 1; - -// WKB count fields are attacker-controlled. The bound caps allocation and -// parsing work even when the encoded input itself is very large. -const MAX_WKB_ELEMENTS: usize = 1_000_000; -const MAX_WKB_BYTES: usize = 64 * 1024 * 1024; -// Nested geometry collections are recursive on the stack. Keep the budget -// explicit so hostile nesting cannot exhaust the decoder thread's stack. -const MAX_WKB_DEPTH: usize = 64; - -/// Serialize a Geometry to WKB (little-endian). -/// -/// Returns `None` when the geometry exceeds the decoder's element or nesting -/// limits, ensuring every emitted WKB value remains decodable. -pub fn geometry_to_wkb(geom: &Geometry) -> Option> { - let encoded_size = checked_encoded_size(geom, MAX_WKB_DEPTH)?; - let mut buf = Vec::with_capacity(encoded_size); - write_geometry(&mut buf, geom, MAX_WKB_DEPTH)?; - debug_assert_eq!(buf.len(), encoded_size); - Some(buf) -} - -/// Deserialize a Geometry from WKB bytes. -/// -/// Returns `None` if the bytes are malformed or truncated. -pub fn geometry_from_wkb(data: &[u8]) -> Option { - if data.len() > MAX_WKB_BYTES { - return None; - } - let mut cursor = 0; - read_geometry(data, &mut cursor, MAX_WKB_DEPTH) -} - -/// Extract bounding box from WKB without full deserialization. -/// -/// Scans coordinate values to find min/max. Faster than full deserialize -/// when only the bbox is needed (e.g., R-tree insertion from Arrow batch). -pub fn wkb_bbox(data: &[u8]) -> Option { - let geom = geometry_from_wkb(data)?; - Some(nodedb_types::geometry_bbox(&geom)) -} - -// ── Write helpers ── - -fn checked_encoded_size(geom: &Geometry, depth_remaining: usize) -> Option { - if depth_remaining == 0 { - return None; - } - let size = match geom { - Geometry::Point { .. } => 21, - Geometry::LineString { coordinates } => { - checked_count(coordinates.len())?; - bounded_add(9, coordinates.len().checked_mul(16)?)? - } - Geometry::Polygon { coordinates } => { - checked_count(coordinates.len())?; - let mut size = 9usize; - for ring in coordinates { - checked_count(ring.len())?; - size = bounded_add(size, bounded_add(4, ring.len().checked_mul(16)?)?)?; - } - size - } - Geometry::MultiPoint { coordinates } => { - checked_count(coordinates.len())?; - bounded_add(9, coordinates.len().checked_mul(21)?)? - } - Geometry::MultiLineString { coordinates } => { - if depth_remaining == 1 { - return None; - } - checked_count(coordinates.len())?; - let mut size = 9usize; - for line in coordinates { - checked_count(line.len())?; - size = bounded_add(size, bounded_add(9, line.len().checked_mul(16)?)?)?; - } - size - } - Geometry::MultiPolygon { coordinates } => { - if depth_remaining == 1 { - return None; - } - checked_count(coordinates.len())?; - let mut size = 9usize; - for polygon in coordinates { - checked_count(polygon.len())?; - let mut polygon_size = 9usize; - for ring in polygon { - checked_count(ring.len())?; - polygon_size = - bounded_add(polygon_size, bounded_add(4, ring.len().checked_mul(16)?)?)?; - } - size = bounded_add(size, polygon_size)?; - } - size - } - Geometry::GeometryCollection { geometries } => { - checked_count(geometries.len())?; - let mut size = 9usize; - for geometry in geometries { - size = bounded_add(size, checked_encoded_size(geometry, depth_remaining - 1)?)?; - } - size - } - _ => 9, - }; - (size <= MAX_WKB_BYTES).then_some(size) -} - -fn checked_count(count: usize) -> Option<()> { - (count <= MAX_WKB_ELEMENTS && u32::try_from(count).is_ok()).then_some(()) -} - -fn bounded_add(left: usize, right: usize) -> Option { - left.checked_add(right) - .filter(|&size| size <= MAX_WKB_BYTES) -} - -fn write_geometry(buf: &mut Vec, geom: &Geometry, depth_remaining: usize) -> Option<()> { - if depth_remaining == 0 { - return None; - } - match geom { - Geometry::Point { coordinates } => { - write_header(buf, WKB_POINT); - write_f64(buf, coordinates[0]); - write_f64(buf, coordinates[1]); - } - Geometry::LineString { coordinates } => { - write_header(buf, WKB_LINESTRING); - write_count(buf, coordinates.len())?; - for c in coordinates { - write_f64(buf, c[0]); - write_f64(buf, c[1]); - } - } - Geometry::Polygon { coordinates } => { - write_header(buf, WKB_POLYGON); - write_count(buf, coordinates.len())?; - for ring in coordinates { - write_count(buf, ring.len())?; - for c in ring { - write_f64(buf, c[0]); - write_f64(buf, c[1]); - } - } - } - Geometry::MultiPoint { coordinates } => { - write_header(buf, WKB_MULTIPOINT); - write_count(buf, coordinates.len())?; - for c in coordinates { - write_header(buf, WKB_POINT); - write_f64(buf, c[0]); - write_f64(buf, c[1]); - } - } - Geometry::MultiLineString { coordinates } => { - write_header(buf, WKB_MULTILINESTRING); - write_count(buf, coordinates.len())?; - for ls in coordinates { - write_geometry( - buf, - &Geometry::LineString { - coordinates: ls.clone(), - }, - depth_remaining - 1, - )?; - } - } - Geometry::MultiPolygon { coordinates } => { - write_header(buf, WKB_MULTIPOLYGON); - write_count(buf, coordinates.len())?; - for poly in coordinates { - write_geometry( - buf, - &Geometry::Polygon { - coordinates: poly.clone(), - }, - depth_remaining - 1, - )?; - } - } - Geometry::GeometryCollection { geometries } => { - write_header(buf, WKB_GEOMETRYCOLLECTION); - write_count(buf, geometries.len())?; - for geometry in geometries { - write_geometry(buf, geometry, depth_remaining - 1)?; - } - } - _ => { - write_header(buf, WKB_GEOMETRYCOLLECTION); - write_u32(buf, 0); - } - } - Some(()) -} - -fn write_count(buf: &mut Vec, count: usize) -> Option<()> { - if count > MAX_WKB_ELEMENTS { - return None; - } - write_u32(buf, u32::try_from(count).ok()?); - Some(()) -} - -fn write_header(buf: &mut Vec, wkb_type: u32) { - buf.push(BYTE_ORDER_LE); - write_u32(buf, wkb_type); -} - -fn write_u32(buf: &mut Vec, val: u32) { - buf.extend_from_slice(&val.to_le_bytes()); -} - -fn write_f64(buf: &mut Vec, val: f64) { - buf.extend_from_slice(&val.to_le_bytes()); -} - -// ── Read helpers ── - -fn read_geometry(data: &[u8], cursor: &mut usize, depth_remaining: usize) -> Option { - if depth_remaining == 0 { - return None; - } - let byte_order = read_u8(data, cursor)?; - let is_le = byte_order == 1; - let wkb_type = read_u32(data, cursor, is_le)?; - - match wkb_type { - WKB_POINT => { - let x = read_f64(data, cursor, is_le)?; - let y = read_f64(data, cursor, is_le)?; - Some(Geometry::Point { - coordinates: [x, y], - }) - } - WKB_LINESTRING => { - let count = read_u32(data, cursor, is_le)? as usize; - let coords = read_coords(data, cursor, count, is_le)?; - Some(Geometry::LineString { - coordinates: coords, - }) - } - WKB_POLYGON => { - let count = read_u32(data, cursor, is_le)? as usize; - let capacity = checked_wkb_capacity(data, cursor, count, 4)?; - let mut rings = Vec::with_capacity(capacity); - for _ in 0..capacity { - let ring_count = read_u32(data, cursor, is_le)? as usize; - let ring = read_coords(data, cursor, ring_count, is_le)?; - rings.push(ring); - } - Some(Geometry::Polygon { coordinates: rings }) - } - WKB_MULTIPOINT => { - let count = read_u32(data, cursor, is_le)? as usize; - let capacity = checked_wkb_capacity(data, cursor, count, 21)?; - let mut coords = Vec::with_capacity(capacity); - for _ in 0..capacity { - let inner = read_geometry(data, cursor, depth_remaining - 1)?; - if let Geometry::Point { coordinates } = inner { - coords.push(coordinates); - } else { - return None; - } - } - Some(Geometry::MultiPoint { - coordinates: coords, - }) - } - WKB_MULTILINESTRING => { - let count = read_u32(data, cursor, is_le)? as usize; - let capacity = checked_wkb_capacity(data, cursor, count, 5)?; - let mut lines = Vec::with_capacity(capacity); - for _ in 0..capacity { - let inner = read_geometry(data, cursor, depth_remaining - 1)?; - if let Geometry::LineString { coordinates } = inner { - lines.push(coordinates); - } else { - return None; - } - } - Some(Geometry::MultiLineString { coordinates: lines }) - } - WKB_MULTIPOLYGON => { - let count = read_u32(data, cursor, is_le)? as usize; - let capacity = checked_wkb_capacity(data, cursor, count, 5)?; - let mut polys = Vec::with_capacity(capacity); - for _ in 0..capacity { - let inner = read_geometry(data, cursor, depth_remaining - 1)?; - if let Geometry::Polygon { coordinates } = inner { - polys.push(coordinates); - } else { - return None; - } - } - Some(Geometry::MultiPolygon { coordinates: polys }) - } - WKB_GEOMETRYCOLLECTION => { - let count = read_u32(data, cursor, is_le)? as usize; - let capacity = checked_wkb_capacity(data, cursor, count, 5)?; - let mut geoms = Vec::with_capacity(capacity); - for _ in 0..capacity { - geoms.push(read_geometry(data, cursor, depth_remaining - 1)?); - } - Some(Geometry::GeometryCollection { geometries: geoms }) - } - _ => None, - } -} - -/// Validate an untrusted count before using it as a vector capacity. -/// -/// Every encoded element must occupy at least `min_encoded_bytes` after the -/// current cursor. This makes a truncated count fail before allocation, while -/// the independent element cap bounds work on large valid inputs. -fn checked_wkb_capacity( - data: &[u8], - cursor: &usize, - count: usize, - min_encoded_bytes: usize, -) -> Option { - if count > MAX_WKB_ELEMENTS { - return None; - } - let required = count.checked_mul(min_encoded_bytes)?; - let remaining = data.len().checked_sub(*cursor)?; - (required <= remaining).then_some(count) -} - -fn read_u8(data: &[u8], cursor: &mut usize) -> Option { - if *cursor >= data.len() { - return None; - } - let val = data[*cursor]; - *cursor += 1; - Some(val) -} - -fn read_u32(data: &[u8], cursor: &mut usize, is_le: bool) -> Option { - if *cursor + 4 > data.len() { - return None; - } - let bytes: [u8; 4] = [ - data[*cursor], - data[*cursor + 1], - data[*cursor + 2], - data[*cursor + 3], - ]; - *cursor += 4; - Some(if is_le { - u32::from_le_bytes(bytes) - } else { - u32::from_be_bytes(bytes) - }) -} - -fn read_f64(data: &[u8], cursor: &mut usize, is_le: bool) -> Option { - if *cursor + 8 > data.len() { - return None; - } - let bytes: [u8; 8] = [ - data[*cursor], - data[*cursor + 1], - data[*cursor + 2], - data[*cursor + 3], - data[*cursor + 4], - data[*cursor + 5], - data[*cursor + 6], - data[*cursor + 7], - ]; - *cursor += 8; - Some(if is_le { - f64::from_le_bytes(bytes) - } else { - f64::from_be_bytes(bytes) - }) -} - -fn read_coords( - data: &[u8], - cursor: &mut usize, - count: usize, - is_le: bool, -) -> Option> { - let capacity = checked_wkb_capacity(data, cursor, count, 16)?; - let mut coords = Vec::with_capacity(capacity); - for _ in 0..capacity { - let x = read_f64(data, cursor, is_le)?; - let y = read_f64(data, cursor, is_le)?; - coords.push([x, y]); - } - Some(coords) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn geometry_to_wkb(geometry: &Geometry) -> Vec { - super::geometry_to_wkb(geometry).expect("test geometry must fit WKB limits") - } - - #[test] - fn point_roundtrip() { - let geom = Geometry::point(-73.9857, 40.7484); - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn linestring_roundtrip() { - let geom = Geometry::line_string(vec![[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]]); - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn polygon_roundtrip() { - let geom = Geometry::polygon(vec![ - vec![ - [0.0, 0.0], - [10.0, 0.0], - [10.0, 10.0], - [0.0, 10.0], - [0.0, 0.0], - ], - vec![[2.0, 2.0], [3.0, 2.0], [3.0, 3.0], [2.0, 3.0], [2.0, 2.0]], // hole - ]); - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn multipoint_roundtrip() { - let geom = Geometry::MultiPoint { - coordinates: vec![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], - }; - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn multilinestring_roundtrip() { - let geom = Geometry::MultiLineString { - coordinates: vec![ - vec![[0.0, 0.0], [1.0, 1.0]], - vec![[2.0, 2.0], [3.0, 3.0], [4.0, 2.0]], - ], - }; - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn multipolygon_roundtrip() { - let geom = Geometry::MultiPolygon { - coordinates: vec![ - vec![vec![ - [0.0, 0.0], - [1.0, 0.0], - [1.0, 1.0], - [0.0, 1.0], - [0.0, 0.0], - ]], - vec![vec![ - [5.0, 5.0], - [6.0, 5.0], - [6.0, 6.0], - [5.0, 6.0], - [5.0, 5.0], - ]], - ], - }; - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn geometry_collection_roundtrip() { - let geom = Geometry::GeometryCollection { - geometries: vec![ - Geometry::point(1.0, 2.0), - Geometry::line_string(vec![[0.0, 0.0], [1.0, 1.0]]), - ], - }; - let wkb = geometry_to_wkb(&geom); - let decoded = geometry_from_wkb(&wkb).unwrap(); - assert_eq!(geom, decoded); - } - - #[test] - fn parsed_max_counts_are_rejected_before_allocation() { - for geometry_type in [ - WKB_LINESTRING, - WKB_POLYGON, - WKB_MULTIPOINT, - WKB_MULTILINESTRING, - WKB_MULTIPOLYGON, - WKB_GEOMETRYCOLLECTION, - ] { - let mut wkb = Vec::new(); - write_header(&mut wkb, geometry_type); - write_u32(&mut wkb, u32::MAX); - assert!(geometry_from_wkb(&wkb).is_none(), "type={geometry_type}"); - } - } - - #[test] - fn encoder_rejects_total_output_above_byte_limit_before_writing() { - let geometry = Geometry::MultiLineString { - coordinates: (0..5).map(|_| vec![[0.0, 0.0]; 840_000]).collect(), - }; - assert!(checked_encoded_size(&geometry, MAX_WKB_DEPTH).is_none()); - assert!(super::geometry_to_wkb(&geometry).is_none()); - } - - #[test] - fn encoder_rejects_element_count_above_decoder_limit() { - let geometry = Geometry::LineString { - coordinates: vec![[0.0, 0.0]; MAX_WKB_ELEMENTS + 1], - }; - assert!(super::geometry_to_wkb(&geometry).is_none()); - } - - #[test] - fn truncated_count_bodies_are_rejected_before_allocation() { - for geometry_type in [ - WKB_LINESTRING, - WKB_POLYGON, - WKB_MULTIPOINT, - WKB_MULTILINESTRING, - WKB_MULTIPOLYGON, - WKB_GEOMETRYCOLLECTION, - ] { - let mut wkb = Vec::new(); - write_header(&mut wkb, geometry_type); - write_u32(&mut wkb, 1); - assert!(geometry_from_wkb(&wkb).is_none(), "type={geometry_type}"); - } - } - - #[test] - fn geometry_collection_depth_is_bounded() { - let mut valid = Geometry::point(1.0, 2.0); - for _ in 0..MAX_WKB_DEPTH - 1 { - valid = Geometry::GeometryCollection { - geometries: vec![valid], - }; - } - let valid_wkb = geometry_to_wkb(&valid); - assert_eq!(geometry_from_wkb(&valid_wkb), Some(valid)); - - let mut too_deep = Geometry::point(1.0, 2.0); - for _ in 0..MAX_WKB_DEPTH { - too_deep = Geometry::GeometryCollection { - geometries: vec![too_deep], - }; - } - assert!(super::geometry_to_wkb(&too_deep).is_none()); - } - - #[test] - fn truncated_data_returns_none() { - let wkb = geometry_to_wkb(&Geometry::point(1.0, 2.0)); - assert!(geometry_from_wkb(&wkb[..3]).is_none()); - assert!(geometry_from_wkb(&[]).is_none()); - } - - #[test] - fn invalid_type_returns_none() { - let mut wkb = geometry_to_wkb(&Geometry::point(1.0, 2.0)); - wkb[1] = 99; // invalid WKB type - assert!(geometry_from_wkb(&wkb).is_none()); - } - - #[test] - fn wkb_bbox_extraction() { - let geom = Geometry::polygon(vec![vec![ - [-10.0, -5.0], - [10.0, -5.0], - [10.0, 5.0], - [-10.0, 5.0], - [-10.0, -5.0], - ]]); - let wkb = geometry_to_wkb(&geom); - let bb = wkb_bbox(&wkb).unwrap(); - assert_eq!(bb.min_lng, -10.0); - assert_eq!(bb.max_lng, 10.0); - assert_eq!(bb.min_lat, -5.0); - assert_eq!(bb.max_lat, 5.0); - } - - #[test] - fn point_wkb_size() { - let wkb = geometry_to_wkb(&Geometry::point(0.0, 0.0)); - // 1 (byte order) + 4 (type) + 8 (x) + 8 (y) = 21 bytes - assert_eq!(wkb.len(), 21); - } -} diff --git a/nodedb-spatial/src/wkb/ewkb.rs b/nodedb-spatial/src/wkb/ewkb.rs new file mode 100644 index 000000000..92b652af9 --- /dev/null +++ b/nodedb-spatial/src/wkb/ewkb.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! PostGIS EWKB writer: ISO WKB with the SRID in the top-level header. + +use nodedb_types::geometry::SridGeometry; + +use super::format::{EWKB_SRID_FLAG, MAX_WKB_BYTES}; +use super::write::{geometry_to_wkb, write_u32}; + +/// Length of the WKB byte-order marker plus type word. +const HEADER_LEN: usize = 5; + +/// Serialize a geometry to little-endian EWKB with its SRID. +/// +/// The top-level type word gains [`EWKB_SRID_FLAG`] and the SRID code +/// follows it. Nested geometries stay plain WKB, as PostGIS writes them. +/// Returns `None` under the same limits as `geometry_to_wkb`. +pub fn geometry_to_ewkb(geom: &SridGeometry) -> Option> { + let iso = geometry_to_wkb(&geom.geometry)?; + if iso.len() + 4 > MAX_WKB_BYTES || iso.len() < HEADER_LEN { + return None; + } + let type_word = u32::from_le_bytes(iso[1..HEADER_LEN].try_into().ok()?); + let mut out = Vec::with_capacity(iso.len() + 4); + out.push(iso[0]); + write_u32(&mut out, type_word | EWKB_SRID_FLAG); + write_u32(&mut out, geom.srid.code()); + out.extend_from_slice(&iso[HEADER_LEN..]); + Some(out) +} + +#[cfg(test)] +mod tests { + use super::super::read::geometry_from_wkb; + use super::*; + use nodedb_types::geometry::{Geometry, Srid}; + + #[test] + fn ewkb_round_trip_keeps_srid() { + for srid in Srid::SUPPORTED { + let g = SridGeometry::new( + srid, + Geometry::polygon(vec![vec![ + [0.0, 0.0], + [100.0, 0.0], + [100.0, 100.0], + [0.0, 0.0], + ]]), + ); + let bytes = geometry_to_ewkb(&g).unwrap(); + let (geometry, code) = geometry_from_wkb(&bytes).unwrap(); + assert_eq!(geometry, g.geometry); + assert_eq!(code, Some(srid.code())); + } + } + + #[test] + fn iso_wkb_carries_no_srid() { + let iso = geometry_to_wkb(&Geometry::point(1.0, 2.0)).unwrap(); + assert_eq!(geometry_from_wkb(&iso).map(|(_, s)| s), Some(None)); + } +} diff --git a/nodedb-spatial/src/wkb/format.rs b/nodedb-spatial/src/wkb/format.rs new file mode 100644 index 000000000..9d3d9f75a --- /dev/null +++ b/nodedb-spatial/src/wkb/format.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! WKB type codes and decoder limits shared by the reader and writers. + +// WKB geometry type codes (ISO 13249 / OGC SFA). +pub(super) const WKB_POINT: u32 = 1; +pub(super) const WKB_LINESTRING: u32 = 2; +pub(super) const WKB_POLYGON: u32 = 3; +pub(super) const WKB_MULTIPOINT: u32 = 4; +pub(super) const WKB_MULTILINESTRING: u32 = 5; +pub(super) const WKB_MULTIPOLYGON: u32 = 6; +pub(super) const WKB_GEOMETRYCOLLECTION: u32 = 7; + +pub(super) const BYTE_ORDER_LE: u8 = 1; + +/// PostGIS EWKB flag: a 4-byte SRID follows the type word. +pub const EWKB_SRID_FLAG: u32 = 0x2000_0000; + +// WKB count fields are attacker-controlled. The bound caps allocation and +// parsing work even when the encoded input itself is very large. +pub(super) const MAX_WKB_ELEMENTS: usize = 1_000_000; +pub(super) const MAX_WKB_BYTES: usize = 64 * 1024 * 1024; +// Nested geometry collections are recursive on the stack. Keep the budget +// explicit so hostile nesting cannot exhaust the decoder thread's stack. +pub(super) const MAX_WKB_DEPTH: usize = 64; diff --git a/nodedb-spatial/src/wkb/mod.rs b/nodedb-spatial/src/wkb/mod.rs new file mode 100644 index 000000000..678c08189 --- /dev/null +++ b/nodedb-spatial/src/wkb/mod.rs @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Well-Known Binary (WKB) serialization for Geometry types. +//! +//! WKB is the standard binary format for geometry interchange (ISO 13249). +//! Used as the Arrow `DataType::Binary` backing for spatial columns, which +//! avoids JSON parse overhead during query execution. +//! +//! Format: `[byte_order: u8] [type: u32] [coordinates...]`. Byte order 1 is +//! little-endian (NDR), 0 is big-endian (XDR). Writers always emit +//! little-endian. The reader accepts both, plus PostGIS EWKB. + +pub mod ewkb; +pub mod format; +pub mod read; +pub mod write; + +pub use ewkb::geometry_to_ewkb; +pub use format::EWKB_SRID_FLAG; +pub use read::{geometry_from_wkb, wkb_bbox}; +pub use write::geometry_to_wkb; diff --git a/nodedb-spatial/src/wkb/read.rs b/nodedb-spatial/src/wkb/read.rs new file mode 100644 index 000000000..cd236f4a4 --- /dev/null +++ b/nodedb-spatial/src/wkb/read.rs @@ -0,0 +1,387 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! WKB / EWKB reader. +//! +//! Accepts ISO WKB and PostGIS EWKB in either byte order. An EWKB type word +//! with [`EWKB_SRID_FLAG`] set carries a 4-byte SRID after it. The reader +//! masks the flag, and reports the top-level SRID to the caller. + +use nodedb_types::geometry::Geometry; + +use super::format::{ + EWKB_SRID_FLAG, MAX_WKB_BYTES, MAX_WKB_DEPTH, MAX_WKB_ELEMENTS, WKB_GEOMETRYCOLLECTION, + WKB_LINESTRING, WKB_MULTILINESTRING, WKB_MULTIPOINT, WKB_MULTIPOLYGON, WKB_POINT, WKB_POLYGON, +}; + +/// Deserialize a Geometry from WKB or EWKB bytes. +/// +/// Returns the geometry and the SRID an EWKB header declares, `None` for +/// ISO WKB. The SRID is the raw code: the caller decides whether it is +/// supported. Returns `None` overall if the bytes are malformed or truncated. +pub fn geometry_from_wkb(data: &[u8]) -> Option<(Geometry, Option)> { + if data.len() > MAX_WKB_BYTES { + return None; + } + let mut cursor = 0; + let (geometry, srid) = read_geometry(data, &mut cursor, MAX_WKB_DEPTH)?; + Some((geometry, srid)) +} + +/// Extract bounding box from WKB, in geographic space. +pub fn wkb_bbox(data: &[u8]) -> Option { + let (geom, _) = geometry_from_wkb(data)?; + Some(nodedb_types::geometry_bbox(&geom)) +} + +/// One geometry and the SRID its own header declared. +fn read_geometry( + data: &[u8], + cursor: &mut usize, + depth_remaining: usize, +) -> Option<(Geometry, Option)> { + if depth_remaining == 0 { + return None; + } + let byte_order = read_u8(data, cursor)?; + let is_le = byte_order == 1; + let raw_type = read_u32(data, cursor, is_le)?; + let srid = if raw_type & EWKB_SRID_FLAG != 0 { + Some(read_u32(data, cursor, is_le)?) + } else { + None + }; + let wkb_type = raw_type & !EWKB_SRID_FLAG; + let child = |data: &[u8], cursor: &mut usize| { + read_geometry(data, cursor, depth_remaining - 1).map(|(g, _)| g) + }; + + let geometry = match wkb_type { + WKB_POINT => { + let x = read_f64(data, cursor, is_le)?; + let y = read_f64(data, cursor, is_le)?; + Geometry::Point { + coordinates: [x, y], + } + } + WKB_LINESTRING => { + let count = read_u32(data, cursor, is_le)? as usize; + Geometry::LineString { + coordinates: read_coords(data, cursor, count, is_le)?, + } + } + WKB_POLYGON => { + let count = read_u32(data, cursor, is_le)? as usize; + let capacity = checked_wkb_capacity(data, cursor, count, 4)?; + let mut rings = Vec::with_capacity(capacity); + for _ in 0..capacity { + let ring_count = read_u32(data, cursor, is_le)? as usize; + rings.push(read_coords(data, cursor, ring_count, is_le)?); + } + Geometry::Polygon { coordinates: rings } + } + WKB_MULTIPOINT => { + let count = read_u32(data, cursor, is_le)? as usize; + let capacity = checked_wkb_capacity(data, cursor, count, 21)?; + let mut coords = Vec::with_capacity(capacity); + for _ in 0..capacity { + let Geometry::Point { coordinates } = child(data, cursor)? else { + return None; + }; + coords.push(coordinates); + } + Geometry::MultiPoint { + coordinates: coords, + } + } + WKB_MULTILINESTRING => { + let count = read_u32(data, cursor, is_le)? as usize; + let capacity = checked_wkb_capacity(data, cursor, count, 5)?; + let mut lines = Vec::with_capacity(capacity); + for _ in 0..capacity { + let Geometry::LineString { coordinates } = child(data, cursor)? else { + return None; + }; + lines.push(coordinates); + } + Geometry::MultiLineString { coordinates: lines } + } + WKB_MULTIPOLYGON => { + let count = read_u32(data, cursor, is_le)? as usize; + let capacity = checked_wkb_capacity(data, cursor, count, 5)?; + let mut polys = Vec::with_capacity(capacity); + for _ in 0..capacity { + let Geometry::Polygon { coordinates } = child(data, cursor)? else { + return None; + }; + polys.push(coordinates); + } + Geometry::MultiPolygon { coordinates: polys } + } + WKB_GEOMETRYCOLLECTION => { + let count = read_u32(data, cursor, is_le)? as usize; + let capacity = checked_wkb_capacity(data, cursor, count, 5)?; + let mut geoms = Vec::with_capacity(capacity); + for _ in 0..capacity { + geoms.push(child(data, cursor)?); + } + Geometry::GeometryCollection { geometries: geoms } + } + _ => return None, + }; + Some((geometry, srid)) +} + +/// Validate an untrusted count before using it as a vector capacity. +/// +/// Every encoded element must occupy at least `min_encoded_bytes` after the +/// current cursor. This makes a truncated count fail before allocation, while +/// the independent element cap bounds work on large valid inputs. +fn checked_wkb_capacity( + data: &[u8], + cursor: &usize, + count: usize, + min_encoded_bytes: usize, +) -> Option { + if count > MAX_WKB_ELEMENTS { + return None; + } + let required = count.checked_mul(min_encoded_bytes)?; + let remaining = data.len().checked_sub(*cursor)?; + (required <= remaining).then_some(count) +} + +fn read_u8(data: &[u8], cursor: &mut usize) -> Option { + let val = *data.get(*cursor)?; + *cursor += 1; + Some(val) +} + +fn read_u32(data: &[u8], cursor: &mut usize, is_le: bool) -> Option { + let bytes: [u8; 4] = data.get(*cursor..*cursor + 4)?.try_into().ok()?; + *cursor += 4; + Some(if is_le { + u32::from_le_bytes(bytes) + } else { + u32::from_be_bytes(bytes) + }) +} + +fn read_f64(data: &[u8], cursor: &mut usize, is_le: bool) -> Option { + let bytes: [u8; 8] = data.get(*cursor..*cursor + 8)?.try_into().ok()?; + *cursor += 8; + Some(if is_le { + f64::from_le_bytes(bytes) + } else { + f64::from_be_bytes(bytes) + }) +} + +fn read_coords( + data: &[u8], + cursor: &mut usize, + count: usize, + is_le: bool, +) -> Option> { + let capacity = checked_wkb_capacity(data, cursor, count, 16)?; + let mut coords = Vec::with_capacity(capacity); + for _ in 0..capacity { + let x = read_f64(data, cursor, is_le)?; + let y = read_f64(data, cursor, is_le)?; + coords.push([x, y]); + } + Some(coords) +} + +#[cfg(test)] +mod tests { + use super::super::write::{geometry_to_wkb as encode, write_header, write_u32}; + use super::*; + + fn geometry_to_wkb(geometry: &Geometry) -> Vec { + encode(geometry).expect("test geometry must fit WKB limits") + } + + fn decode(wkb: &[u8]) -> Option { + geometry_from_wkb(wkb).map(|(g, _)| g) + } + + #[test] + fn point_roundtrip() { + let geom = Geometry::point(-73.9857, 40.7484); + let wkb = geometry_to_wkb(&geom); + assert_eq!(geometry_from_wkb(&wkb), Some((geom, None))); + } + + #[test] + fn linestring_roundtrip() { + let geom = Geometry::line_string(vec![[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]]); + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn polygon_roundtrip() { + let geom = Geometry::polygon(vec![ + vec![ + [0.0, 0.0], + [10.0, 0.0], + [10.0, 10.0], + [0.0, 10.0], + [0.0, 0.0], + ], + vec![[2.0, 2.0], [3.0, 2.0], [3.0, 3.0], [2.0, 3.0], [2.0, 2.0]], // hole + ]); + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn multipoint_roundtrip() { + let geom = Geometry::MultiPoint { + coordinates: vec![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + }; + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn multilinestring_roundtrip() { + let geom = Geometry::MultiLineString { + coordinates: vec![ + vec![[0.0, 0.0], [1.0, 1.0]], + vec![[2.0, 2.0], [3.0, 3.0], [4.0, 2.0]], + ], + }; + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn multipolygon_roundtrip() { + let geom = Geometry::MultiPolygon { + coordinates: vec![ + vec![vec![ + [0.0, 0.0], + [1.0, 0.0], + [1.0, 1.0], + [0.0, 1.0], + [0.0, 0.0], + ]], + vec![vec![ + [5.0, 5.0], + [6.0, 5.0], + [6.0, 6.0], + [5.0, 6.0], + [5.0, 5.0], + ]], + ], + }; + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn geometry_collection_roundtrip() { + let geom = Geometry::GeometryCollection { + geometries: vec![ + Geometry::point(1.0, 2.0), + Geometry::line_string(vec![[0.0, 0.0], [1.0, 1.0]]), + ], + }; + assert_eq!(decode(&geometry_to_wkb(&geom)), Some(geom)); + } + + #[test] + fn parsed_max_counts_are_rejected_before_allocation() { + for geometry_type in [ + WKB_LINESTRING, + WKB_POLYGON, + WKB_MULTIPOINT, + WKB_MULTILINESTRING, + WKB_MULTIPOLYGON, + WKB_GEOMETRYCOLLECTION, + ] { + let mut wkb = Vec::new(); + write_header(&mut wkb, geometry_type); + write_u32(&mut wkb, u32::MAX); + assert!(decode(&wkb).is_none(), "type={geometry_type}"); + } + } + + #[test] + fn truncated_count_bodies_are_rejected_before_allocation() { + for geometry_type in [ + WKB_LINESTRING, + WKB_POLYGON, + WKB_MULTIPOINT, + WKB_MULTILINESTRING, + WKB_MULTIPOLYGON, + WKB_GEOMETRYCOLLECTION, + ] { + let mut wkb = Vec::new(); + write_header(&mut wkb, geometry_type); + write_u32(&mut wkb, 1); + assert!(decode(&wkb).is_none(), "type={geometry_type}"); + } + } + + #[test] + fn geometry_collection_depth_is_bounded() { + let mut valid = Geometry::point(1.0, 2.0); + for _ in 0..MAX_WKB_DEPTH - 1 { + valid = Geometry::GeometryCollection { + geometries: vec![valid], + }; + } + let valid_wkb = geometry_to_wkb(&valid); + assert_eq!(decode(&valid_wkb), Some(valid)); + + let mut too_deep = Geometry::point(1.0, 2.0); + for _ in 0..MAX_WKB_DEPTH { + too_deep = Geometry::GeometryCollection { + geometries: vec![too_deep], + }; + } + assert!(encode(&too_deep).is_none()); + } + + #[test] + fn truncated_data_returns_none() { + let wkb = geometry_to_wkb(&Geometry::point(1.0, 2.0)); + assert!(decode(&wkb[..3]).is_none()); + assert!(decode(&[]).is_none()); + } + + #[test] + fn invalid_type_returns_none() { + let mut wkb = geometry_to_wkb(&Geometry::point(1.0, 2.0)); + wkb[1] = 99; // invalid WKB type + assert!(decode(&wkb).is_none()); + } + + #[test] + fn ewkb_srid_flag_is_masked_and_reported() { + // EWKB POINT(1 2) SRID=32648, little-endian. + let mut wkb = Vec::new(); + write_header(&mut wkb, WKB_POINT | EWKB_SRID_FLAG); + write_u32(&mut wkb, 32648); + wkb.extend_from_slice(&1.0f64.to_le_bytes()); + wkb.extend_from_slice(&2.0f64.to_le_bytes()); + assert_eq!( + geometry_from_wkb(&wkb), + Some((Geometry::point(1.0, 2.0), Some(32648))) + ); + } + + #[test] + fn wkb_bbox_extraction() { + let geom = Geometry::polygon(vec![vec![ + [-10.0, -5.0], + [10.0, -5.0], + [10.0, 5.0], + [-10.0, 5.0], + [-10.0, -5.0], + ]]); + let wkb = geometry_to_wkb(&geom); + let bb = wkb_bbox(&wkb).unwrap(); + assert_eq!(bb.min_lng, -10.0); + assert_eq!(bb.max_lng, 10.0); + assert_eq!(bb.min_lat, -5.0); + assert_eq!(bb.max_lat, 5.0); + } +} diff --git a/nodedb-spatial/src/wkb/write.rs b/nodedb-spatial/src/wkb/write.rs new file mode 100644 index 000000000..4c188102f --- /dev/null +++ b/nodedb-spatial/src/wkb/write.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! ISO WKB writer (little-endian, no SRID). + +use nodedb_types::geometry::Geometry; + +use super::format::{ + BYTE_ORDER_LE, MAX_WKB_BYTES, MAX_WKB_DEPTH, MAX_WKB_ELEMENTS, WKB_GEOMETRYCOLLECTION, + WKB_LINESTRING, WKB_MULTILINESTRING, WKB_MULTIPOINT, WKB_MULTIPOLYGON, WKB_POINT, WKB_POLYGON, +}; + +/// Serialize a Geometry to ISO WKB (little-endian). Carries no SRID. +/// +/// Returns `None` when the geometry exceeds the decoder's element or nesting +/// limits, ensuring every emitted WKB value remains decodable. +pub fn geometry_to_wkb(geom: &Geometry) -> Option> { + let encoded_size = checked_encoded_size(geom, MAX_WKB_DEPTH)?; + let mut buf = Vec::with_capacity(encoded_size); + write_geometry(&mut buf, geom, MAX_WKB_DEPTH)?; + debug_assert_eq!(buf.len(), encoded_size); + Some(buf) +} + +pub(super) fn checked_encoded_size(geom: &Geometry, depth_remaining: usize) -> Option { + if depth_remaining == 0 { + return None; + } + let size = match geom { + Geometry::Point { .. } => 21, + Geometry::LineString { coordinates } => { + checked_count(coordinates.len())?; + bounded_add(9, coordinates.len().checked_mul(16)?)? + } + Geometry::Polygon { coordinates } => { + checked_count(coordinates.len())?; + let mut size = 9usize; + for ring in coordinates { + checked_count(ring.len())?; + size = bounded_add(size, bounded_add(4, ring.len().checked_mul(16)?)?)?; + } + size + } + Geometry::MultiPoint { coordinates } => { + checked_count(coordinates.len())?; + bounded_add(9, coordinates.len().checked_mul(21)?)? + } + Geometry::MultiLineString { coordinates } => { + if depth_remaining == 1 { + return None; + } + checked_count(coordinates.len())?; + let mut size = 9usize; + for line in coordinates { + checked_count(line.len())?; + size = bounded_add(size, bounded_add(9, line.len().checked_mul(16)?)?)?; + } + size + } + Geometry::MultiPolygon { coordinates } => { + if depth_remaining == 1 { + return None; + } + checked_count(coordinates.len())?; + let mut size = 9usize; + for polygon in coordinates { + checked_count(polygon.len())?; + let mut polygon_size = 9usize; + for ring in polygon { + checked_count(ring.len())?; + polygon_size = + bounded_add(polygon_size, bounded_add(4, ring.len().checked_mul(16)?)?)?; + } + size = bounded_add(size, polygon_size)?; + } + size + } + Geometry::GeometryCollection { geometries } => { + checked_count(geometries.len())?; + let mut size = 9usize; + for geometry in geometries { + size = bounded_add(size, checked_encoded_size(geometry, depth_remaining - 1)?)?; + } + size + } + _ => 9, + }; + (size <= MAX_WKB_BYTES).then_some(size) +} + +fn checked_count(count: usize) -> Option<()> { + (count <= MAX_WKB_ELEMENTS && u32::try_from(count).is_ok()).then_some(()) +} + +fn bounded_add(left: usize, right: usize) -> Option { + left.checked_add(right) + .filter(|&size| size <= MAX_WKB_BYTES) +} + +fn write_geometry(buf: &mut Vec, geom: &Geometry, depth_remaining: usize) -> Option<()> { + if depth_remaining == 0 { + return None; + } + match geom { + Geometry::Point { coordinates } => { + write_header(buf, WKB_POINT); + write_f64(buf, coordinates[0]); + write_f64(buf, coordinates[1]); + } + Geometry::LineString { coordinates } => { + write_header(buf, WKB_LINESTRING); + write_count(buf, coordinates.len())?; + for c in coordinates { + write_f64(buf, c[0]); + write_f64(buf, c[1]); + } + } + Geometry::Polygon { coordinates } => { + write_header(buf, WKB_POLYGON); + write_count(buf, coordinates.len())?; + for ring in coordinates { + write_count(buf, ring.len())?; + for c in ring { + write_f64(buf, c[0]); + write_f64(buf, c[1]); + } + } + } + Geometry::MultiPoint { coordinates } => { + write_header(buf, WKB_MULTIPOINT); + write_count(buf, coordinates.len())?; + for c in coordinates { + write_header(buf, WKB_POINT); + write_f64(buf, c[0]); + write_f64(buf, c[1]); + } + } + Geometry::MultiLineString { coordinates } => { + write_header(buf, WKB_MULTILINESTRING); + write_count(buf, coordinates.len())?; + for ls in coordinates { + write_geometry( + buf, + &Geometry::LineString { + coordinates: ls.clone(), + }, + depth_remaining - 1, + )?; + } + } + Geometry::MultiPolygon { coordinates } => { + write_header(buf, WKB_MULTIPOLYGON); + write_count(buf, coordinates.len())?; + for poly in coordinates { + write_geometry( + buf, + &Geometry::Polygon { + coordinates: poly.clone(), + }, + depth_remaining - 1, + )?; + } + } + Geometry::GeometryCollection { geometries } => { + write_header(buf, WKB_GEOMETRYCOLLECTION); + write_count(buf, geometries.len())?; + for geometry in geometries { + write_geometry(buf, geometry, depth_remaining - 1)?; + } + } + _ => { + write_header(buf, WKB_GEOMETRYCOLLECTION); + write_u32(buf, 0); + } + } + Some(()) +} + +fn write_count(buf: &mut Vec, count: usize) -> Option<()> { + if count > MAX_WKB_ELEMENTS { + return None; + } + write_u32(buf, u32::try_from(count).ok()?); + Some(()) +} + +pub(super) fn write_header(buf: &mut Vec, wkb_type: u32) { + buf.push(BYTE_ORDER_LE); + write_u32(buf, wkb_type); +} + +pub(super) fn write_u32(buf: &mut Vec, val: u32) { + buf.extend_from_slice(&val.to_le_bytes()); +} + +fn write_f64(buf: &mut Vec, val: f64) { + buf.extend_from_slice(&val.to_le_bytes()); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encoder_rejects_total_output_above_byte_limit_before_writing() { + let geometry = Geometry::MultiLineString { + coordinates: (0..5).map(|_| vec![[0.0, 0.0]; 840_000]).collect(), + }; + assert!(checked_encoded_size(&geometry, MAX_WKB_DEPTH).is_none()); + assert!(geometry_to_wkb(&geometry).is_none()); + } + + #[test] + fn encoder_rejects_element_count_above_decoder_limit() { + let geometry = Geometry::LineString { + coordinates: vec![[0.0, 0.0]; MAX_WKB_ELEMENTS + 1], + }; + assert!(geometry_to_wkb(&geometry).is_none()); + } + + #[test] + fn point_wkb_size() { + let wkb = geometry_to_wkb(&Geometry::point(0.0, 0.0)).unwrap(); + // 1 (byte order) + 4 (type) + 8 (x) + 8 (y) = 21 bytes + assert_eq!(wkb.len(), 21); + } +} diff --git a/nodedb-sql/src/aggregate_walk.rs b/nodedb-sql/src/aggregate_walk.rs index b517425c8..8de8c8306 100644 --- a/nodedb-sql/src/aggregate_walk.rs +++ b/nodedb-sql/src/aggregate_walk.rs @@ -46,6 +46,23 @@ pub fn contains_aggregate(expr: &Expr, functions: &FunctionRegistry) -> bool { detector.found } +/// Return `true` if `expr` is one aggregate call and nothing else, past any +/// parentheses: `SUM(x)`, not `ROUND(SUM(x), 1)` or `SUM(x) / 2`. +pub fn is_bare_aggregate_call(expr: &Expr, functions: &FunctionRegistry) -> bool { + match expr { + Expr::Nested(inner) => is_bare_aggregate_call(inner, functions), + Expr::Function(f) => functions.is_aggregate(&function_name(f)) && f.over.is_none(), + _ => false, + } +} + +/// Return `true` if `expr` computes over an aggregate: it holds an aggregate +/// call and is not one bare aggregate call. A grouped plan emits the +/// aggregate's value, so such an expression is evaluated over the group row. +pub fn computes_over_aggregate(expr: &Expr, functions: &FunctionRegistry) -> bool { + contains_aggregate(expr, functions) && !is_bare_aggregate_call(expr, functions) +} + /// Extract every aggregate function call from `expr`, binding each to /// the given output `alias`. Nested aggregates (e.g. `SUM(AVG(x))`, /// which is illegal SQL in Postgres and most other systems) are diff --git a/nodedb-sql/src/catalog.rs b/nodedb-sql/src/catalog.rs index f5de86a1a..ff79ef02d 100644 --- a/nodedb-sql/src/catalog.rs +++ b/nodedb-sql/src/catalog.rs @@ -165,6 +165,23 @@ pub trait SqlCatalog { self.get_collection(database_id, name) } + /// The designated time-key column of timeseries collection `name`. + /// + /// The native grouped timeseries scan buckets by this column only. The + /// default reads it from the declared columns by the DDL rule: the one + /// `TIME_KEY` column, else the one timestamp column. A catalog that + /// stores the designation overrides this. `None` when the collection is + /// unknown or no column is designated. + fn timeseries_time_key( + &self, + database_id: DatabaseId, + name: &str, + ) -> Result, SqlCatalogError> { + Ok(self.get_collection(database_id, name)?.and_then(|info| { + crate::ddl_ast::collection_type::designated::designated_time_key(&info.columns) + })) + } + /// Look up an array by name. Returns `None` if no array with that /// name is registered. The default implementation returns `None` so /// that catalog adapters predating array support compile without diff --git a/nodedb-sql/src/ddl_ast/collection_type/designated.rs b/nodedb-sql/src/ddl_ast/collection_type/designated.rs index d28b9002e..e8cac76ea 100644 --- a/nodedb-sql/src/ddl_ast/collection_type/designated.rs +++ b/nodedb-sql/src/ddl_ast/collection_type/designated.rs @@ -86,6 +86,27 @@ pub(crate) fn resolve_time_key(columns: &[(String, String)]) -> Result Option { + use crate::types_expr::SqlDataType; + let declared: Vec<(String, String)> = columns + .iter() + .map(|c| { + let text = match (&c.raw_type, &c.data_type) { + (Some(raw), _) => raw.clone(), + (None, SqlDataType::Timestamp | SqlDataType::Timestamptz) => "TIMESTAMP".into(), + (None, _) => String::new(), + }; + (c.name.clone(), text) + }) + .collect(); + resolve_time_key(&declared).ok() +} + /// Find the geometry column from the column list. /// /// Matches columns whose type_str contains `SPATIAL_INDEX` modifier or whose diff --git a/nodedb-sql/src/ddl_ast/parse/copy_from.rs b/nodedb-sql/src/ddl_ast/parse/copy_from.rs index 50fed8f05..9926d0395 100644 --- a/nodedb-sql/src/ddl_ast/parse/copy_from.rs +++ b/nodedb-sql/src/ddl_ast/parse/copy_from.rs @@ -2,10 +2,11 @@ //! Parse `COPY FROM '' [WITH (FORMAT ..., DELIMITER ..., HEADER ...)]`. //! -//! Only the server-side file-path form is handled here. The STDIN streaming -//! shape (`COPY ... FROM STDIN`) is handled elsewhere (backup/restore); -//! returning `None` lets dispatch fall through to that handler. `COPY ... TO` -//! and `COPY (SELECT ...) TO` belong to the `copy_to` family. +//! The server-side file-path form parses here. The STDIN streaming shape +//! (`COPY FROM STDIN`) parses in `copy_from_stdin`. The restore shape +//! `COPY tenant_restore() FROM STDIN` returns `None`, so dispatch falls +//! through to the backup/restore handler. `COPY ... TO` and +//! `COPY (SELECT ...) TO` belong to the `copy_to` family. //! //! The direction is the keyword right after the copied relation, never a //! `FROM` or `TO` inside a path, a name, or a query. @@ -81,8 +82,9 @@ fn relation_end(toks: &[Token<'_>]) -> Option { /// Try to parse a COPY statement. /// -/// Returns `None` unless the statement is `COPY FROM ''`. -/// Returns `Some(Err)` for a malformed statement of that shape. +/// Returns `None` unless the statement is `COPY FROM ''` +/// or `COPY FROM STDIN`. Returns `Some(Err)` for a malformed +/// statement of either shape. pub(super) fn try_parse( upper: &str, _parts: &[&str], @@ -91,10 +93,11 @@ pub(super) fn try_parse( if !upper.starts_with("COPY ") { return None; } - if copy_direction(trimmed) != Some(CopyDirection::FromFile) { - return None; + match copy_direction(trimmed)? { + CopyDirection::FromFile => Some(parse_copy_from(trimmed)), + CopyDirection::FromStdin => super::copy_from_stdin::try_parse(trimmed), + CopyDirection::To => None, } - Some(parse_copy_from(trimmed)) } fn parse_copy_from(sql: &str) -> Result { @@ -181,7 +184,7 @@ fn parse_format_value(s: &str) -> Result { } } -fn parse_delimiter_value(s: &str) -> Result { +pub(super) fn parse_delimiter_value(s: &str) -> Result { let mut chars = s.chars(); let ch = chars.next().ok_or_else(|| SqlError::Parse { detail: "COPY: DELIMITER must be a single character".to_string(), @@ -194,7 +197,7 @@ fn parse_delimiter_value(s: &str) -> Result { Ok(ch) } -fn parse_bool_value(s: &str) -> Option { +pub(super) fn parse_bool_value(s: &str) -> Option { match s.trim().to_lowercase().as_str() { "true" | "1" | "yes" | "on" => Some(true), "false" | "0" | "no" | "off" => Some(false), diff --git a/nodedb-sql/src/ddl_ast/parse/copy_from_stdin.rs b/nodedb-sql/src/ddl_ast/parse/copy_from_stdin.rs new file mode 100644 index 000000000..3ad96d259 --- /dev/null +++ b/nodedb-sql/src/ddl_ast/parse/copy_from_stdin.rs @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Parse `COPY FROM STDIN [WITH (FORMAT csv|text, DELIMITER '', +//! HEADER [bool])]`. +//! +//! Only a one-token relation parses here. `COPY tenant_restore(1) FROM +//! STDIN` is the restore shape, which the backup detector owns, so it +//! returns `None`. + +use super::clause::{ClauseCursor, Tok, tokenize}; +use super::copy_from::{copy_collection_name, parse_bool_value, parse_delimiter_value}; +use crate::ddl_ast::statement::{CopyStdinFormat, MiscStmt, NodedbStatement}; +use crate::error::SqlError; + +/// Parse a `COPY ... FROM STDIN` statement. `None` when the copied relation +/// is more than one name token, or the options are not this grammar's (a +/// `FORMAT json`, say): such a statement keeps the refusal the bulk-import +/// router gives it. +pub(super) fn try_parse(sql: &str) -> Option> { + let toks = tokenize(sql); + // `COPY`, the name, then `FROM`. + let one_name = matches!( + toks.get(1).map(|t| &t.tok), + Some(Tok::Word(_) | Tok::Ident(_)) + ) && matches!(toks.get(2).map(|t| &t.tok), Some(Tok::Word(w)) if w.eq_ignore_ascii_case("FROM")); + if !one_name { + return None; + } + let mut c = ClauseCursor::new(sql, &toks, "COPY"); + parse_stdin(&mut c).ok().map(Ok) +} + +fn parse_stdin(c: &mut ClauseCursor<'_, '_>) -> Result { + c.expect_keyword("COPY")?; + let collection = copy_collection_name(c)?; + c.expect_keyword("FROM")?; + c.expect_keyword("STDIN")?; + let mut format = CopyStdinFormat::Text; + let mut delimiter = None; + // PostgreSQL reads no header line unless HEADER is given. + let mut header = false; + if c.eat_keyword("WITH") { + let options = c.paren_list("WITH options", |c| { + let key = c.word("option name")?.to_ascii_uppercase(); + c.eat_punct('='); + // A bare `HEADER` turns the header line on. + if key == "HEADER" && (c.peek_punct(',') || c.peek_punct(')')) { + return Ok((key, "true".to_string())); + } + let value = match c.string(&key) { + Ok(value) => value, + Err(_) => c.word(&key)?.to_string(), + }; + Ok((key, value)) + })?; + for (key, value) in options { + match key.as_str() { + "FORMAT" => format = parse_format(&value)?, + "DELIMITER" => delimiter = Some(parse_delimiter_value(&value)?), + "HEADER" => { + header = parse_bool_value(&value).ok_or_else(|| { + c.error(format!("HEADER must be true or false — found '{value}'")) + })?; + } + other => return Err(c.error(format!("WITH: unknown option '{other}'"))), + } + } + } + if !c.at_end() { + return Err(c.error(format!("unexpected trailing content {}", c.describe_next()))); + } + Ok(NodedbStatement::Misc(MiscStmt::CopyFromStdin { + collection, + format, + delimiter, + header, + })) +} + +fn parse_format(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "text" => Ok(CopyStdinFormat::Text), + "csv" => Ok(CopyStdinFormat::Csv), + other => Err(SqlError::Parse { + detail: format!("COPY FROM STDIN: unknown FORMAT '{other}'; expected csv or text"), + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ddl_ast::parse::dispatch::parse; + + fn stdin(sql: &str) -> (String, CopyStdinFormat, Option, bool) { + match parse(sql).expect("expected Some").expect("expected Ok") { + NodedbStatement::Misc(MiscStmt::CopyFromStdin { + collection, + format, + delimiter, + header, + }) => (collection, format, delimiter, header), + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn bare_stdin_defaults_to_text_without_header() { + let (name, format, delimiter, header) = stdin("COPY scene FROM STDIN"); + assert_eq!(name, "scene"); + assert_eq!(format, CopyStdinFormat::Text); + assert_eq!(delimiter, None); + assert!(!header); + } + + #[test] + fn csv_with_bare_header_and_delimiter() { + let (_, format, delimiter, header) = + stdin("COPY scene FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER)"); + assert_eq!(format, CopyStdinFormat::Csv); + assert_eq!(delimiter, Some(';')); + assert!(header); + } + + #[test] + fn header_takes_an_explicit_boolean() { + let (_, _, _, header) = stdin("COPY scene FROM STDIN WITH (FORMAT csv, HEADER false)"); + assert!(!header); + } + + #[test] + fn quoted_name_keeps_no_quotes() { + let (name, ..) = stdin("COPY \"Scene\" FROM STDIN"); + assert_eq!(name, "scene"); + } + + #[test] + fn another_grammars_options_are_not_parsed() { + assert!(parse("COPY scene FROM STDIN WITH (FORMAT json)").is_none()); + assert!(parse("COPY scene FROM STDIN WITH (BOGUS 1)").is_none()); + } + + #[test] + fn restore_shape_is_not_parsed() { + assert!(parse("COPY tenant_restore(1) FROM STDIN").is_none()); + } +} diff --git a/nodedb-sql/src/ddl_ast/parse/mod.rs b/nodedb-sql/src/ddl_ast/parse/mod.rs index 93d3793d2..ef43801ff 100644 --- a/nodedb-sql/src/ddl_ast/parse/mod.rs +++ b/nodedb-sql/src/ddl_ast/parse/mod.rs @@ -8,6 +8,7 @@ mod cluster_admin; mod collection; mod conflict_policy; mod copy_from; +mod copy_from_stdin; mod copy_to; mod custom_type; pub mod database; diff --git a/nodedb-sql/src/ddl_ast/statement/maintenance.rs b/nodedb-sql/src/ddl_ast/statement/maintenance.rs index b7b78301c..31b687367 100644 --- a/nodedb-sql/src/ddl_ast/statement/maintenance.rs +++ b/nodedb-sql/src/ddl_ast/statement/maintenance.rs @@ -9,6 +9,17 @@ pub enum CopyToSource { Query(String), } +/// Row format of `COPY ... FROM STDIN`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CopyStdinFormat { + /// PostgreSQL text format: tab-delimited by default, `\N` is NULL, and a + /// backslash escapes the next character. + Text, + /// CSV: comma-delimited by default, `"` quotes a field, and an unquoted + /// empty field is NULL. + Csv, +} + /// Format for `COPY ... FROM` bulk import. #[derive(Debug, Clone, PartialEq, Eq)] pub enum CopyFormat { diff --git a/nodedb-sql/src/ddl_ast/statement/types/misc.rs b/nodedb-sql/src/ddl_ast/statement/types/misc.rs index fa3268682..00543d715 100644 --- a/nodedb-sql/src/ddl_ast/statement/types/misc.rs +++ b/nodedb-sql/src/ddl_ast/statement/types/misc.rs @@ -2,7 +2,7 @@ //! Miscellaneous DDL/DML statements. -use crate::ddl_ast::statement::maintenance::{CopyFormat, CopyToSource}; +use crate::ddl_ast::statement::maintenance::{CopyFormat, CopyStdinFormat, CopyToSource}; #[derive(Debug, Clone, PartialEq)] pub enum MiscStmt { @@ -18,8 +18,8 @@ pub enum MiscStmt { // ── Bulk import ────────────────────────────────────────────── /// `COPY FROM '' [WITH (FORMAT ..., DELIMITER ..., HEADER ...)]` /// - /// Server-side file-path bulk import. Does not handle STDIN streaming - /// (that is a different protocol path) or COPY ... TO. + /// Server-side file-path bulk import. `CopyFromStdin` is the streaming + /// form, and `CopyToFile` the export. CopyFromFile { collection: String, path: String, @@ -28,6 +28,17 @@ pub enum MiscStmt { header: bool, }, + /// `COPY FROM STDIN [WITH (FORMAT csv|text, DELIMITER '', HEADER [bool])]` + /// + /// Streaming bulk import: the client sends the rows over the pgwire + /// COPY protocol. `header` defaults to false, as in PostgreSQL. + CopyFromStdin { + collection: String, + format: CopyStdinFormat, + delimiter: Option, + header: bool, + }, + // ── Bulk export ────────────────────────────────────────────── /// `COPY TO '' [WITH (FORMAT ..., DELIMITER ..., HEADER ...)]` /// `COPY (SELECT ...) TO '' [WITH (...)]` diff --git a/nodedb-sql/src/engine_rules/array.rs b/nodedb-sql/src/engine_rules/array.rs index 2c96010e1..4bc321643 100644 --- a/nodedb-sql/src/engine_rules/array.rs +++ b/nodedb-sql/src/engine_rules/array.rs @@ -76,6 +76,21 @@ impl EngineRules for ArrayRules { )) } + /// A bulk load decodes its rows against the array's schema and loads them + /// in chunks. One chunk is one `InsertArray` plan, so each chunk is one + /// cell batch and one WAL record. + fn plan_bulk_load(&self, p: BulkLoadParams) -> Result { + if p.cells.is_empty() { + return Ok(BulkLoadPlan::CellChunks(Vec::new())); + } + Ok(BulkLoadPlan::CellChunks(vec![SqlPlan::InsertArray( + InsertArrayPlan { + name: p.collection, + rows: p.cells, + }, + )])) + } + fn plan_aggregate(&self, _p: AggregateParams) -> Result { Err(unsupported( "GROUP BY", @@ -118,6 +133,38 @@ mod tests { } } + #[test] + fn bulk_load_plans_one_insert_array_per_chunk() { + let probe = ArrayRules + .plan_bulk_load(BulkLoadParams { + collection: "g".into(), + cells: Vec::new(), + }) + .unwrap(); + assert!(matches!(probe, BulkLoadPlan::CellChunks(ref p) if p.is_empty())); + + let row = crate::types_array::ArrayInsertRow { + coords: vec![crate::types_array::ArrayCoordLiteral::Int64(1)], + attrs: vec![crate::types_array::ArrayAttrLiteral::Float64(0.5)], + }; + let chunk = ArrayRules + .plan_bulk_load(BulkLoadParams { + collection: "g".into(), + cells: vec![row.clone(), row], + }) + .unwrap(); + match chunk { + BulkLoadPlan::CellChunks(plans) => { + assert_eq!(plans.len(), 1); + assert!(matches!( + &plans[0], + SqlPlan::InsertArray(InsertArrayPlan { name, rows }) if name == "g" && rows.len() == 2 + )); + } + BulkLoadPlan::RowInserts => panic!("an array loads in cell chunks"), + } + } + #[test] fn every_arm_is_unsupported() { let r = ArrayRules; diff --git a/nodedb-sql/src/engine_rules/bulk_load.rs b/nodedb-sql/src/engine_rules/bulk_load.rs new file mode 100644 index 000000000..6b36ca845 --- /dev/null +++ b/nodedb-sql/src/engine_rules/bulk_load.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Parameters and answer of [`EngineRules::plan_bulk_load`](super::EngineRules::plan_bulk_load). + +use crate::types::SqlPlan; +use crate::types_array::ArrayInsertRow; + +/// Parameters for planning one chunk of a bulk load (`COPY ... FROM`). +pub struct BulkLoadParams { + /// The target collection or array. + pub collection: String, + /// The chunk's cells, typed against the array's schema. Empty when the + /// caller asks only how the engine loads, before it decodes any row. + pub cells: Vec, +} + +/// How an engine loads the rows of a bulk load. +#[derive(Debug)] +pub enum BulkLoadPlan { + /// Each row loads as one `INSERT` statement through the engine's insert + /// path, so the engine's triggers, defaults and checks apply per row. + RowInserts, + /// The rows decode against the array's schema and load in chunks. The + /// plans load the chunk the parameters carry, in order. They are empty + /// when the parameters carry no cells. + CellChunks(Vec), +} diff --git a/nodedb-sql/src/engine_rules/columnar.rs b/nodedb-sql/src/engine_rules/columnar.rs index 2a4de2768..7b1d9b6bc 100644 --- a/nodedb-sql/src/engine_rules/columnar.rs +++ b/nodedb-sql/src/engine_rules/columnar.rs @@ -125,6 +125,11 @@ impl EngineRules for ColumnarRules { }]) } + /// Each row loads as one INSERT. + fn plan_bulk_load(&self, _p: BulkLoadParams) -> Result { + Ok(BulkLoadPlan::RowInserts) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { if p.temporal.is_temporal() && !p.bitemporal { return Err(SqlError::Unsupported { diff --git a/nodedb-sql/src/engine_rules/continuous_aggregate.rs b/nodedb-sql/src/engine_rules/continuous_aggregate.rs index 087e35c0e..36cfe1c33 100644 --- a/nodedb-sql/src/engine_rules/continuous_aggregate.rs +++ b/nodedb-sql/src/engine_rules/continuous_aggregate.rs @@ -90,6 +90,11 @@ impl EngineRules for ContinuousAggregateRules { Err(read_only("TRUNCATE", &p.collection)) } + /// Each row loads as one INSERT, which `plan_insert` refuses. + fn plan_bulk_load(&self, _p: BulkLoadParams) -> Result { + Ok(BulkLoadPlan::RowInserts) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { refuse_temporal(&p.collection, &p.temporal)?; let base_scan = SqlPlan::Scan { diff --git a/nodedb-sql/src/engine_rules/document_schemaless.rs b/nodedb-sql/src/engine_rules/document_schemaless.rs index 4057aa12c..9ed77ce9b 100644 --- a/nodedb-sql/src/engine_rules/document_schemaless.rs +++ b/nodedb-sql/src/engine_rules/document_schemaless.rs @@ -140,6 +140,11 @@ impl EngineRules for SchemalessRules { }]) } + /// Each row loads as one INSERT. + fn plan_bulk_load(&self, _p: BulkLoadParams) -> Result { + Ok(BulkLoadPlan::RowInserts) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { let base_scan = SqlPlan::Scan { collection: p.collection, diff --git a/nodedb-sql/src/engine_rules/document_strict.rs b/nodedb-sql/src/engine_rules/document_strict.rs index f5eecaa14..08fdd6a64 100644 --- a/nodedb-sql/src/engine_rules/document_strict.rs +++ b/nodedb-sql/src/engine_rules/document_strict.rs @@ -140,6 +140,11 @@ impl EngineRules for StrictRules { }]) } + /// Each row loads as one INSERT. + fn plan_bulk_load(&self, _p: BulkLoadParams) -> Result { + Ok(BulkLoadPlan::RowInserts) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { let base_scan = SqlPlan::Scan { collection: p.collection, diff --git a/nodedb-sql/src/engine_rules/index_lookup.rs b/nodedb-sql/src/engine_rules/index_lookup.rs index 77c93987d..937f26337 100644 --- a/nodedb-sql/src/engine_rules/index_lookup.rs +++ b/nodedb-sql/src/engine_rules/index_lookup.rs @@ -2,6 +2,7 @@ //! Shared equality-to-index rewrite for the document engines. +use crate::planner::computed_tail; use crate::types::*; use super::params::ScanParams; @@ -90,22 +91,36 @@ pub(crate) fn try_document_index_lookup( value }; - return Some(SqlPlan::DocumentIndexLookup(DocumentIndexLookupPlan { + // The indexed fetch evaluates no expression. A computed SELECT item + // runs in a tail over the fetched rows, and the tail takes the + // DISTINCT and row bound, which apply to the computed output. + let tail = computed_tail::needs_computed_tail(¶ms.projection); + let bounds = computed_tail::TailBounds { + distinct: params.distinct, + limit: params.limit, + offset: params.offset, + }; + let lookup = SqlPlan::DocumentIndexLookup(DocumentIndexLookupPlan { collection: params.collection.clone(), alias: params.alias.clone(), engine, field: idx.field.clone(), value: lookup_value, filters: remaining, - projection: params.projection.clone(), + projection: computed_tail::body_projection(¶ms.projection), sort_keys: params.sort_keys.clone(), - limit: params.limit, - offset: params.offset, - distinct: params.distinct, + limit: if tail { None } else { params.limit }, + offset: if tail { 0 } else { params.offset }, + distinct: !tail && params.distinct, window_functions: params.window_functions.clone(), case_insensitive: idx.case_insensitive, temporal: params.temporal, - })); + }); + return Some(computed_tail::with_computed_tail( + lookup, + ¶ms.projection, + bounds, + )); } None } diff --git a/nodedb-sql/src/engine_rules/kv.rs b/nodedb-sql/src/engine_rules/kv.rs index 9f56c7458..a439128da 100644 --- a/nodedb-sql/src/engine_rules/kv.rs +++ b/nodedb-sql/src/engine_rules/kv.rs @@ -115,6 +115,11 @@ impl EngineRules for KvRules { }]) } + /// Each row loads as one INSERT. + fn plan_bulk_load(&self, _p: BulkLoadParams) -> Result { + Ok(BulkLoadPlan::RowInserts) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { let base_scan = SqlPlan::Scan { collection: p.collection, diff --git a/nodedb-sql/src/engine_rules/mod.rs b/nodedb-sql/src/engine_rules/mod.rs index a5815ae83..504233fa7 100644 --- a/nodedb-sql/src/engine_rules/mod.rs +++ b/nodedb-sql/src/engine_rules/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 pub mod array; +mod bulk_load; pub mod columnar; pub mod continuous_aggregate; pub mod document_schemaless; @@ -14,6 +15,7 @@ pub mod spatial; pub mod timeseries; mod timeseries_native; +pub use bulk_load::{BulkLoadParams, BulkLoadPlan}; pub(crate) use index_lookup::try_document_index_lookup; pub use params::{ AggregateParams, DeleteParams, InsertParams, InsertSelectParams, MergeParams, PointGetParams, diff --git a/nodedb-sql/src/engine_rules/params.rs b/nodedb-sql/src/engine_rules/params.rs index aa80ef1c7..fe16a2f35 100644 --- a/nodedb-sql/src/engine_rules/params.rs +++ b/nodedb-sql/src/engine_rules/params.rs @@ -189,6 +189,9 @@ pub struct AggregateParams { /// scan so bitemporal aggregate queries project an as-of snapshot /// before grouping. pub temporal: crate::temporal::TemporalScope, + /// The designated time-key column of a timeseries collection. The native + /// grouped timeseries scan buckets by this column only. + pub time_key: Option, } /// A `time_bucket(interval, column)` GROUP BY key. diff --git a/nodedb-sql/src/engine_rules/rules.rs b/nodedb-sql/src/engine_rules/rules.rs index 2bc28e1a1..5a4d1f88a 100644 --- a/nodedb-sql/src/engine_rules/rules.rs +++ b/nodedb-sql/src/engine_rules/rules.rs @@ -5,6 +5,7 @@ use crate::error::Result; use crate::types::SqlPlan; +use super::bulk_load::{BulkLoadParams, BulkLoadPlan}; use super::params::{ AggregateParams, DeleteParams, InsertParams, InsertSelectParams, MergeParams, PointGetParams, ScanParams, TruncateParams, UpdateFromParams, UpdateParams, UpsertParams, @@ -45,6 +46,10 @@ pub trait EngineRules { /// Plan a TRUNCATE. Returns `Err(SqlError::Unsupported)` for engines /// that carry no whole-collection clear on the SQL surface (array). fn plan_truncate(&self, params: TruncateParams) -> Result>; + /// Plan one chunk of a bulk load (`COPY ... FROM`), or say how the + /// engine loads rows when the parameters carry no cells. Returns `Err` + /// for an engine that refuses a bulk load (timeseries, spatial). + fn plan_bulk_load(&self, params: BulkLoadParams) -> Result; /// Plan a GROUP BY / aggregate query. fn plan_aggregate(&self, params: AggregateParams) -> Result; /// Plan a MERGE statement. diff --git a/nodedb-sql/src/engine_rules/spatial.rs b/nodedb-sql/src/engine_rules/spatial.rs index f0a39a331..dafa717fa 100644 --- a/nodedb-sql/src/engine_rules/spatial.rs +++ b/nodedb-sql/src/engine_rules/spatial.rs @@ -126,6 +126,16 @@ impl EngineRules for SpatialRules { }]) } + fn plan_bulk_load(&self, p: BulkLoadParams) -> Result { + Err(SqlError::Unsupported { + detail: format!( + "COPY: collection '{}' uses the spatial engine; \ + use INSERT with a WKT/GeoJSON geometry column instead", + p.collection + ), + }) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { let base_scan = SqlPlan::Scan { collection: p.collection, diff --git a/nodedb-sql/src/engine_rules/timeseries.rs b/nodedb-sql/src/engine_rules/timeseries.rs index 036858747..bf5b41338 100644 --- a/nodedb-sql/src/engine_rules/timeseries.rs +++ b/nodedb-sql/src/engine_rules/timeseries.rs @@ -67,6 +67,8 @@ impl EngineRules for TimeseriesRules { sort_keys: p.sort_keys, tiered: p.has_auto_tier, temporal: p.temporal, + output: Vec::new(), + group_keys: Vec::new(), })) } @@ -115,6 +117,16 @@ impl EngineRules for TimeseriesRules { }]) } + fn plan_bulk_load(&self, p: BulkLoadParams) -> Result { + Err(SqlError::Unsupported { + detail: format!( + "COPY: collection '{}' uses the timeseries engine; \ + use ILP or INSERT with explicit time column instead", + p.collection + ), + }) + } + fn plan_aggregate(&self, p: AggregateParams) -> Result { if p.temporal.is_temporal() && !p.bitemporal { return Err(SqlError::Unsupported { @@ -144,6 +156,8 @@ impl EngineRules for TimeseriesRules { sort_keys: Vec::new(), tiered: p.has_auto_tier, temporal: p.temporal, + output: Vec::new(), + group_keys: Vec::new(), }); return Ok(SqlPlan::Aggregate { input: Box::new(input), @@ -173,6 +187,10 @@ impl EngineRules for TimeseriesRules { sort_keys: Vec::new(), tiered: p.has_auto_tier, temporal: p.temporal, + // The SELECT-list order and aliases are attached by the caller, + // which holds the projection. + output: Vec::new(), + group_keys: p.group_by, })) } diff --git a/nodedb-sql/src/engine_rules/timeseries_native.rs b/nodedb-sql/src/engine_rules/timeseries_native.rs index 215e9c008..3629fa457 100644 --- a/nodedb-sql/src/engine_rules/timeseries_native.rs +++ b/nodedb-sql/src/engine_rules/timeseries_native.rs @@ -47,7 +47,11 @@ pub(super) fn native_shape(p: &AggregateParams) -> Option { | SqlExpr::Literal(_) | SqlExpr::Subquery(_) | SqlExpr::Wildcard => match p.time_bucket { - Some(tb) if tb.group_index == index && tb.interval_ms > 0 => { + Some(tb) + if tb.group_index == index + && tb.interval_ms > 0 + && buckets_time_key(key, p.time_key.as_deref()) => + { bucket_interval_ms = tb.interval_ms; } _ => return None, @@ -68,6 +72,29 @@ pub(super) fn native_shape(p: &AggregateParams) -> Option { }) } +/// Whether the `time_bucket` call `key` buckets the collection's time key. +/// +/// The grouped scan buckets by the time key only. A call over any other +/// column, an expression, or a collection whose time key is unknown needs +/// the generic aggregate, which buckets by the column the call names. +fn buckets_time_key(key: &SqlExpr, time_key: Option<&str>) -> bool { + let (SqlExpr::Function { args, .. }, Some(time_key)) = (key, time_key) else { + return false; + }; + let mut columns = args.iter().filter_map(|arg| match arg { + SqlExpr::Column { name, .. } => Some(name.as_str()), + _ => None, + }); + let bucketed = columns.next(); + let only_column = columns.next().is_none(); + let other_args_literal = args + .iter() + .all(|arg| matches!(arg, SqlExpr::Column { .. } | SqlExpr::Literal(_))); + only_column + && other_args_literal + && bucketed.is_some_and(|name| name.eq_ignore_ascii_case(time_key)) +} + /// Whether the grouped scan accumulates `a` exactly. fn native_aggregate(a: &AggregateExpr, columns: &[ColumnInfo]) -> bool { if a.distinct @@ -94,3 +121,31 @@ fn native_aggregate(a: &AggregateExpr, columns: &[ColumnInfo]) -> bool { _ => false, } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::SqlValue; + + fn bucket(column: &str) -> SqlExpr { + SqlExpr::Function { + name: "time_bucket".into(), + args: vec![ + SqlExpr::Literal(SqlValue::String("1 hour".into())), + SqlExpr::Column { + table: None, + name: column.into(), + }, + ], + distinct: false, + } + } + + /// Only a bucket of the time key takes the grouped scan. + #[test] + fn only_the_time_key_buckets_natively() { + assert!(buckets_time_key(&bucket("ts"), Some("ts"))); + assert!(!buckets_time_key(&bucket("recorded_at"), Some("ts"))); + assert!(!buckets_time_key(&bucket("ts"), None)); + } +} diff --git a/nodedb-sql/src/error.rs b/nodedb-sql/src/error.rs index 9ab89ba18..e1442fbf6 100644 --- a/nodedb-sql/src/error.rs +++ b/nodedb-sql/src/error.rs @@ -85,6 +85,16 @@ pub enum SqlError { #[error("ambiguous column '{column}' — qualify with table name")] AmbiguousColumn { column: String }, + /// A grouped SELECT list reads a column that is neither a GROUP BY key + /// nor inside an aggregate. The finalized group row does not carry it. + /// PostgreSQL rejects the same input with SQLSTATE `42803` + /// (`grouping_error`). + #[error( + "column \"{column}\" must appear in the GROUP BY clause or be used in an \ + aggregate function" + )] + UngroupedColumn { column: String }, + #[error("type mismatch: {detail}")] TypeMismatch { detail: String }, diff --git a/nodedb-sql/src/functions/builtins/scalars/spatial/arg_shape.rs b/nodedb-sql/src/functions/builtins/scalars/spatial/arg_shape.rs index 978be6ce0..ac8b8e9c8 100644 --- a/nodedb-sql/src/functions/builtins/scalars/spatial/arg_shape.rs +++ b/nodedb-sql/src/functions/builtins/scalars/spatial/arg_shape.rs @@ -43,6 +43,7 @@ static GEOMETRY_2_DISTANCE: &[ArgTypeSpec] = &[ arg("geom2", GEOMETRY), arg("distance_m", FLOAT), ]; +static GEOMETRY_SRID: &[ArgTypeSpec] = &[arg("geom", GEOMETRY), arg("srid", INT)]; static GEOMETRY_BUFFER: &[ArgTypeSpec] = &[ arg("geom", GEOMETRY), arg("distance_m", FLOAT), @@ -84,6 +85,7 @@ pub(super) fn arg_types(shape: GeoArgShape) -> &'static [ArgTypeSpec] { GeoArgShape::Geometry1 => GEOMETRY_1, GeoArgShape::Geometry2 => GEOMETRY_2, GeoArgShape::Geometry2Distance => GEOMETRY_2_DISTANCE, + GeoArgShape::GeometrySrid => GEOMETRY_SRID, GeoArgShape::GeometryBuffer => GEOMETRY_BUFFER, GeoArgShape::LngLat => LNG_LAT, GeoArgShape::LngLatZ => LNG_LAT_Z, diff --git a/nodedb-sql/src/parser/array_stmt/parse.rs b/nodedb-sql/src/parser/array_stmt/parse.rs index a35080449..d51d9ed23 100644 --- a/nodedb-sql/src/parser/array_stmt/parse.rs +++ b/nodedb-sql/src/parser/array_stmt/parse.rs @@ -154,6 +154,65 @@ mod tests { } } + #[test] + fn parse_create_array_narrow_attr_types() { + use crate::types_array::ArrayAttrType as T; + let sql = "CREATE ARRAY scene DIMS (x INT64 [0..255]) \ + ATTRS (a UINT8, b INT16, c SMALLINT, d INT2, e UINT16, f INT32, \ + g INTEGER, h INT, i INT4, j UINT32, k FLOAT32, l REAL, m FLOAT4, \ + n FLOAT, o BIGINT) \ + TILE_EXTENTS (64)"; + let stmt = try_parse_array_statement(sql).unwrap().unwrap(); + let ArrayStatement::Create(c) = stmt else { + panic!("wrong variant"); + }; + let types: Vec = c.attrs.iter().map(|a| a.dtype).collect(); + assert_eq!( + types, + vec![ + T::UInt8, + T::Int16, + T::Int16, + T::Int16, + T::UInt16, + T::Int32, + T::Int32, + T::Int32, + T::Int32, + T::UInt32, + T::Float32, + T::Float32, + T::Float32, + T::Float64, + T::Int64, + ] + ); + } + + #[test] + fn parse_create_array_not_null_attr() { + let sql = "CREATE ARRAY g DIMS (x INT64 [0..10]) \ + ATTRS (a UINT16 NOT NULL, b FLOAT32, c INT64 not null) TILE_EXTENTS (1)"; + let stmt = try_parse_array_statement(sql).unwrap().unwrap(); + let ArrayStatement::Create(c) = stmt else { + panic!("wrong variant"); + }; + let nullable: Vec = c.attrs.iter().map(|a| a.nullable).collect(); + assert_eq!(nullable, vec![false, true, false]); + } + + #[test] + fn create_rejects_not_without_null() { + let sql = "CREATE ARRAY g DIMS (x INT64 [0..10]) ATTRS (a INT64 NOT) TILE_EXTENTS (1)"; + assert!(try_parse_array_statement(sql).is_err()); + } + + #[test] + fn create_rejects_unknown_attr_type() { + let sql = "CREATE ARRAY g DIMS (x INT64 [0..10]) ATTRS (v UINT64) TILE_EXTENTS (1)"; + assert!(try_parse_array_statement(sql).is_err()); + } + #[test] fn create_rejects_unknown_dim_type() { let sql = "CREATE ARRAY g DIMS (x BOGUS [0..10]) ATTRS (v INT64) TILE_EXTENTS (1)"; diff --git a/nodedb-sql/src/parser/array_stmt/parse_impl/create.rs b/nodedb-sql/src/parser/array_stmt/parse_impl/create.rs index 28372f573..b9e0cb713 100644 --- a/nodedb-sql/src/parser/array_stmt/parse_impl/create.rs +++ b/nodedb-sql/src/parser/array_stmt/parse_impl/create.rs @@ -171,7 +171,8 @@ impl<'a> Parser<'a> { let dtype = parse_attr_type(&type_name) .ok_or_else(|| self.err(format!("unknown attr type `{type_name}`")))?; let nullable = if self.match_kw("NOT") { - self.expect_kw("NULL")?; + // The lexer reads `NULL` as its own token, never as an identifier. + self.expect(&Tok::Null)?; false } else { // Default: nullable. @@ -238,12 +239,20 @@ fn parse_dim_type(s: &str) -> Option { } } +/// An attr type name. `INT` is the 4-byte `INT32`, as in PostgreSQL. A +/// bare `FLOAT` is the 8-byte `FLOAT64`. fn parse_attr_type(s: &str) -> Option { match s.to_ascii_uppercase().as_str() { - "INT64" | "INT" | "BIGINT" => Some(ArrayAttrType::Int64), + "INT64" | "BIGINT" => Some(ArrayAttrType::Int64), "FLOAT64" | "DOUBLE" | "FLOAT" => Some(ArrayAttrType::Float64), "STRING" | "TEXT" | "VARCHAR" => Some(ArrayAttrType::String), "BYTES" | "BLOB" | "BYTEA" => Some(ArrayAttrType::Bytes), + "UINT8" => Some(ArrayAttrType::UInt8), + "INT16" | "SMALLINT" | "INT2" => Some(ArrayAttrType::Int16), + "UINT16" => Some(ArrayAttrType::UInt16), + "INT32" | "INTEGER" | "INT" | "INT4" => Some(ArrayAttrType::Int32), + "UINT32" => Some(ArrayAttrType::UInt32), + "FLOAT32" | "REAL" | "FLOAT4" => Some(ArrayAttrType::Float32), _ => None, } } diff --git a/nodedb-sql/src/planner/agg_bind.rs b/nodedb-sql/src/planner/agg_bind.rs index c32b15808..5db6e081d 100644 --- a/nodedb-sql/src/planner/agg_bind.rs +++ b/nodedb-sql/src/planner/agg_bind.rs @@ -1,22 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 -//! Binding aggregate calls in HAVING and post-aggregate ORDER BY. +//! Binding aggregate calls in HAVING, post-aggregate ORDER BY, and computed +//! grouped SELECT items. //! -//! Both clauses are evaluated against finalized group rows, so an aggregate -//! either mentions must actually be computed and must be addressed by the -//! column name it lands under. Left as a literal `sum(...)` call, it reaches -//! the executor as a scalar function over a row that has no such column — the -//! predicate matches nothing, or the sort key resolves to nothing. +//! These clauses are evaluated against finalized group rows. An aggregate one +//! mentions must be computed and addressed by the column it lands under. Left +//! as a literal `sum(...)` call, it reaches the executor as a scalar function +//! over a row that has no such column. The predicate matches nothing, or the +//! value resolves to nothing. //! -//! The two clauses differ in *which* name they need, because they run on +//! The clauses differ in *which* name they need, because they run on //! opposite sides of the rename to user aliases: //! //! - HAVING filters before the rename, so it addresses the canonical key //! (`sum(amount)`). -//! - ORDER BY sorts after the rename, so it addresses the output name -//! (`total`, when the projection said `SUM(amount) AS total`). +//! - ORDER BY and a computed SELECT item run after the rename, so they +//! address the output name (`total`, when the projection said +//! `SUM(amount) AS total`). -use sqlparser::ast; +use std::ops::ControlFlow; + +use sqlparser::ast::{self, VisitMut, VisitorMut}; use crate::aggregate_walk::extract_aggregates; use crate::error::{Result, SqlError}; @@ -90,6 +94,10 @@ fn resolve_select_aliases(expr: &ast::Expr, projection: &[ast::SelectItem]) -> a } } +/// Rewrite every aggregate call in `expr` at any depth: inside a scalar +/// function (`ROUND(AVG(v), 1)`), a CASE, a CAST, an IN list, and every other +/// expression position. An aggregate inside a subquery belongs to that +/// subquery, so it stays as written. fn bind( expr: &ast::Expr, aggregates: &mut Vec, @@ -97,41 +105,69 @@ fn bind( name: BindName, scope: &TableScope, ) -> Result { - match expr { - ast::Expr::Function(func) if is_aggregate_call(func, functions) => { - let column = register_aggregate(expr, aggregates, functions, name, scope)?; - Ok(ast::Expr::Identifier(ast::Ident::new(column))) + let mut bound = expr.clone(); + let mut binder = AggregateBinder { + aggregates, + functions, + name, + scope, + query_depth: 0, + error: None, + }; + let _ = bound.visit(&mut binder); + match binder.error { + Some(error) => Err(error), + None => Ok(bound), + } +} + +/// Replaces each aggregate call outside a subquery with a reference to its +/// computed column. The first registration error stops the walk. +struct AggregateBinder<'a> { + aggregates: &'a mut Vec, + functions: &'a FunctionRegistry, + name: BindName, + scope: &'a TableScope, + /// How many subqueries enclose the expression being visited. + query_depth: usize, + error: Option, +} + +impl VisitorMut for AggregateBinder<'_> { + type Break = (); + + fn pre_visit_query(&mut self, _query: &mut ast::Query) -> ControlFlow { + self.query_depth += 1; + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &mut ast::Query) -> ControlFlow { + self.query_depth -= 1; + ControlFlow::Continue(()) + } + + /// The walk visits the children of the replacement, an identifier, so it + /// does not descend into the aggregate's arguments. + fn pre_visit_expr(&mut self, expr: &mut ast::Expr) -> ControlFlow { + if self.query_depth > 0 { + return ControlFlow::Continue(()); + } + let ast::Expr::Function(func) = &*expr else { + return ControlFlow::Continue(()); + }; + if !is_aggregate_call(func, self.functions) { + return ControlFlow::Continue(()); + } + match register_aggregate(expr, self.aggregates, self.functions, self.name, self.scope) { + Ok(column) => { + *expr = ast::Expr::Identifier(ast::Ident::new(column)); + ControlFlow::Continue(()) + } + Err(error) => { + self.error = Some(error); + ControlFlow::Break(()) + } } - ast::Expr::BinaryOp { left, op, right } => Ok(ast::Expr::BinaryOp { - left: Box::new(bind(left, aggregates, functions, name, scope)?), - op: op.clone(), - right: Box::new(bind(right, aggregates, functions, name, scope)?), - }), - ast::Expr::UnaryOp { op, expr } => Ok(ast::Expr::UnaryOp { - op: *op, - expr: Box::new(bind(expr, aggregates, functions, name, scope)?), - }), - ast::Expr::Nested(inner) => Ok(ast::Expr::Nested(Box::new(bind( - inner, aggregates, functions, name, scope, - )?))), - ast::Expr::IsNull(inner) => Ok(ast::Expr::IsNull(Box::new(bind( - inner, aggregates, functions, name, scope, - )?))), - ast::Expr::IsNotNull(inner) => Ok(ast::Expr::IsNotNull(Box::new(bind( - inner, aggregates, functions, name, scope, - )?))), - ast::Expr::Between { - expr, - negated, - low, - high, - } => Ok(ast::Expr::Between { - expr: Box::new(bind(expr, aggregates, functions, name, scope)?), - negated: *negated, - low: Box::new(bind(low, aggregates, functions, name, scope)?), - high: Box::new(bind(high, aggregates, functions, name, scope)?), - }), - other => Ok(other.clone()), } } diff --git a/nodedb-sql/src/planner/aggregate.rs b/nodedb-sql/src/planner/aggregate.rs index 16d62cf00..cf8b12e2d 100644 --- a/nodedb-sql/src/planner/aggregate.rs +++ b/nodedb-sql/src/planner/aggregate.rs @@ -29,6 +29,8 @@ pub struct AggregateQuery<'a> { /// the aggregate. The engine then plans an aggregate over an input plan /// the caller can wrap. pub input_joins: bool, + /// The designated time-key column of a timeseries collection. + pub time_key: Option, } /// The aggregate's row limit without a `LIMIT` clause: every group. An @@ -45,9 +47,11 @@ pub fn plan_aggregate(query: AggregateQuery<'_>, functions: &FunctionRegistry) - scope, temporal, input_joins, + time_key, } = query; let clauses = plan_aggregate_clauses(select, order_by, &table.info.columns, scope, functions)?; let time_bucket = time_bucket_key(&clauses.group_by, functions); + let group_keys = clauses.group_by.clone(); let rules = engine_rules::resolve_engine_rules(table.info.engine); let mut plan = rules.plan_aggregate(AggregateParams { @@ -66,6 +70,7 @@ pub fn plan_aggregate(query: AggregateQuery<'_>, functions: &FunctionRegistry) - temporal: *temporal, hidden_aggregates: clauses.hidden_aggregates, input_joins, + time_key, })?; // The engine rules build the Aggregate node without a projection in scope, @@ -80,10 +85,75 @@ pub fn plan_aggregate(query: AggregateQuery<'_>, functions: &FunctionRegistry) - { *group_by_aliases = clauses.group_by_aliases; *output_order = clauses.output_order; + } else if let SqlPlan::TimeseriesScan(scan) = &mut plan { + // A native grouped timeseries scan emits its group rows under fixed + // keys. The SELECT-list order and aliases are recorded here, so the + // response shaper reads each key under the name the client asked for. + // + // An item the scan does not emit as a column (`ROUND(AVG(v), 1)`) is + // evaluated over the group rows by a tail, which can read any key or + // aggregate. The scan then emits every key and every aggregate, and + // the tail restates the SELECT list. + let order = if crate::planner::aggregate_cp_wrap::items_need_group_row_tail( + &group_keys, + &select.projection, + functions, + scope, + )? { + (0..group_keys.len()) + .map(AggOutputSlot::GroupKey) + .chain((0..scan.aggregates.len()).map(AggOutputSlot::Aggregate)) + .collect() + } else { + clauses.output_order + }; + scan.output = timeseries_agg_output(&group_keys, &clauses.group_by_aliases, &order); } Ok(plan) } +/// The SELECT-list output columns of a native grouped timeseries scan. +/// +/// `group_by` holds the GROUP BY keys the slots in `output_order` index. +/// The native scan accepts only column keys and one `time_bucket` key, so a +/// non-column key is its time bucket. +fn timeseries_agg_output( + group_by: &[SqlExpr], + group_by_aliases: &[Option], + output_order: &[AggOutputSlot], +) -> Vec { + output_order + .iter() + .filter_map(|slot| match slot { + AggOutputSlot::GroupKey(index) => { + let alias = group_by_aliases.get(*index).cloned().flatten(); + group_by.get(*index).map(|key| match key { + SqlExpr::Column { name, .. } => TimeseriesAggOutput::GroupColumn { + name: name.clone(), + alias, + }, + other => TimeseriesAggOutput::Bucket { + alias, + source: bucket_source_column(other), + }, + }) + } + AggOutputSlot::Aggregate(index) => Some(TimeseriesAggOutput::Aggregate(*index)), + }) + .collect() +} + +/// The column a `time_bucket(...)` call buckets: its one column argument. +fn bucket_source_column(key: &SqlExpr) -> Option { + match key { + SqlExpr::Function { args, .. } => args.iter().find_map(|arg| match arg { + SqlExpr::Column { name, .. } => Some(name.clone()), + _ => None, + }), + _ => None, + } +} + /// The `time_bucket(...)` key among the GROUP BY keys, with its interval. fn time_bucket_key(group_by: &[SqlExpr], functions: &FunctionRegistry) -> Option { group_by @@ -246,6 +316,16 @@ pub fn extract_aggregates_from_projection( }; let mut extracted = crate::aggregate_walk::extract_aggregates(expr, &alias, functions, scope)?; + // An item that computes over its aggregates (`ROUND(AVG(v), 1) AS m`, + // `SUM(a) / COUNT(*)`) is evaluated over the group row by a tail. Its + // aggregates are not output columns, so each takes its own canonical + // key: two aggregates of one item never share a name, and the item's + // alias stays free for the tail's result. + if crate::aggregate_walk::computes_over_aggregate(expr, functions) { + for agg in &mut extracted { + agg.alias = super::agg_naming::aggregate_output_key(agg); + } + } aggregates.append(&mut extracted); } Ok(aggregates) @@ -317,6 +397,39 @@ mod tests { assert_eq!(time_bucket_key(&[column("qtype")], &functions), None); } + /// The native scan's output follows the SELECT list: an aggregate, the + /// bucket under its alias with the column it buckets, then a tag column. + #[test] + fn timeseries_output_follows_the_select_list() { + let keys = vec![ + column("host"), + time_bucket(SqlValue::String("1 hour".into())), + ]; + let output = timeseries_agg_output( + &keys, + &[None, Some("b".into())], + &[ + AggOutputSlot::Aggregate(0), + AggOutputSlot::GroupKey(1), + AggOutputSlot::GroupKey(0), + ], + ); + assert_eq!( + output, + vec![ + TimeseriesAggOutput::Aggregate(0), + TimeseriesAggOutput::Bucket { + alias: Some("b".into()), + source: Some("ts".into()), + }, + TimeseriesAggOutput::GroupColumn { + name: "host".into(), + alias: None, + }, + ] + ); + } + #[test] fn unreadable_interval_is_zero() { let functions = FunctionRegistry::new(); diff --git a/nodedb-sql/src/planner/aggregate_clauses.rs b/nodedb-sql/src/planner/aggregate_clauses.rs index b0a8a9635..5ad4d140a 100644 --- a/nodedb-sql/src/planner/aggregate_clauses.rs +++ b/nodedb-sql/src/planner/aggregate_clauses.rs @@ -85,7 +85,7 @@ pub fn plan_aggregate_clauses( functions, scope, )?; - let group_by_aliases = group_by_output_aliases(&select.projection, &group_by); + let group_by_aliases = group_by_output_aliases(&select.projection, &group_by, scope); Ok(AggregateClauses { group_by, diff --git a/nodedb-sql/src/planner/aggregate_cp_wrap/expression.rs b/nodedb-sql/src/planner/aggregate_cp_wrap/expression.rs index 8105fb41f..a96eaf9a0 100644 --- a/nodedb-sql/src/planner/aggregate_cp_wrap/expression.rs +++ b/nodedb-sql/src/planner/aggregate_cp_wrap/expression.rs @@ -4,35 +4,69 @@ use crate::types::SqlExpr; -/// `expr` with every column reference stripped of its table qualifier. A -/// finalized group row keys its columns by bare name. -pub(super) fn unqualify_columns(expr: SqlExpr) -> SqlExpr { +/// How an expression over the input relations reads a finalized group row. +pub(super) struct GroupRowNames { + /// `(bare name, row name)` of each GROUP BY key the row carries under + /// another name: an aliased key carries its alias. + pub renames: Vec<(String, String)>, + /// `(Debug rendering, row name)` of each computed GROUP BY key. The row + /// carries the key's value, not the input columns it reads. + pub computed_keys: Vec<(String, String)>, +} + +/// `expr` read from a finalized group row. +/// +/// A sub-expression equal to a computed GROUP BY key becomes a reference to +/// the key's row column. `SqlExpr` has no `PartialEq` (it can carry a subquery +/// plan), so the match compares the `Debug` rendering. Both sides pass through +/// `convert_expr`, so equal expressions render identically. Every other column +/// reference is stripped of its table qualifier and renamed to the name the +/// row carries it under. +pub(super) fn group_row_columns(expr: SqlExpr, names: &GroupRowNames) -> SqlExpr { + if !names.computed_keys.is_empty() && !matches!(expr, SqlExpr::Column { .. }) { + let rendered = format!("{expr:?}"); + if let Some((_, name)) = names.computed_keys.iter().find(|(key, _)| *key == rendered) { + return SqlExpr::Column { + table: None, + name: name.clone(), + }; + } + } + let go = |e: SqlExpr| group_row_columns(e, names); + let go_box = |e: Box| Box::new(group_row_columns(*e, names)); match expr { - SqlExpr::Column { name, .. } => SqlExpr::Column { table: None, name }, + SqlExpr::Column { name, .. } => { + let name = names + .renames + .iter() + .find(|(from, _)| *from == name) + .map_or(name, |(_, to)| to.clone()); + SqlExpr::Column { table: None, name } + } SqlExpr::Function { name, args, distinct, } => SqlExpr::Function { name, - args: args.into_iter().map(unqualify_columns).collect(), + args: args.into_iter().map(go).collect(), distinct, }, SqlExpr::BinaryOp { left, op, right } => SqlExpr::BinaryOp { - left: Box::new(unqualify_columns(*left)), + left: go_box(left), op, - right: Box::new(unqualify_columns(*right)), + right: go_box(right), }, SqlExpr::UnaryOp { op, expr } => SqlExpr::UnaryOp { op, - expr: Box::new(unqualify_columns(*expr)), + expr: go_box(expr), }, SqlExpr::Cast { expr, to_type } => SqlExpr::Cast { - expr: Box::new(unqualify_columns(*expr)), + expr: go_box(expr), to_type, }, SqlExpr::IsNull { expr, negated } => SqlExpr::IsNull { - expr: Box::new(unqualify_columns(*expr)), + expr: go_box(expr), negated, }, SqlExpr::Case { @@ -40,20 +74,20 @@ pub(super) fn unqualify_columns(expr: SqlExpr) -> SqlExpr { when_then, else_expr, } => SqlExpr::Case { - operand: operand.map(|e| Box::new(unqualify_columns(*e))), + operand: operand.map(go_box), when_then: when_then .into_iter() - .map(|(when, then)| (unqualify_columns(when), unqualify_columns(then))) + .map(|(when, then)| (go(when), go(then))) .collect(), - else_expr: else_expr.map(|e| Box::new(unqualify_columns(*e))), + else_expr: else_expr.map(go_box), }, SqlExpr::InList { expr, list, negated, } => SqlExpr::InList { - expr: Box::new(unqualify_columns(*expr)), - list: list.into_iter().map(unqualify_columns).collect(), + expr: go_box(expr), + list: list.into_iter().map(go).collect(), negated, }, SqlExpr::Between { @@ -62,9 +96,9 @@ pub(super) fn unqualify_columns(expr: SqlExpr) -> SqlExpr { high, negated, } => SqlExpr::Between { - expr: Box::new(unqualify_columns(*expr)), - low: Box::new(unqualify_columns(*low)), - high: Box::new(unqualify_columns(*high)), + expr: go_box(expr), + low: go_box(low), + high: go_box(high), negated, }, SqlExpr::Like { @@ -73,14 +107,12 @@ pub(super) fn unqualify_columns(expr: SqlExpr) -> SqlExpr { negated, case_insensitive, } => SqlExpr::Like { - expr: Box::new(unqualify_columns(*expr)), - pattern: Box::new(unqualify_columns(*pattern)), + expr: go_box(expr), + pattern: go_box(pattern), negated, case_insensitive, }, - SqlExpr::ArrayLiteral(items) => { - SqlExpr::ArrayLiteral(items.into_iter().map(unqualify_columns).collect()) - } + SqlExpr::ArrayLiteral(items) => SqlExpr::ArrayLiteral(items.into_iter().map(go).collect()), SqlExpr::Literal(_) | SqlExpr::Subquery(_) | SqlExpr::Wildcard => expr, } } diff --git a/nodedb-sql/src/planner/aggregate_cp_wrap/mod.rs b/nodedb-sql/src/planner/aggregate_cp_wrap/mod.rs index a5372ecfd..75ec7399c 100644 --- a/nodedb-sql/src/planner/aggregate_cp_wrap/mod.rs +++ b/nodedb-sql/src/planner/aggregate_cp_wrap/mod.rs @@ -6,4 +6,5 @@ mod expression; mod projection; mod wrap; +pub(crate) use projection::items_need_group_row_tail; pub use wrap::wrap_aggregate_cp_items; diff --git a/nodedb-sql/src/planner/aggregate_cp_wrap/projection.rs b/nodedb-sql/src/planner/aggregate_cp_wrap/projection.rs index 043424ee6..52b9e9ad6 100644 --- a/nodedb-sql/src/planner/aggregate_cp_wrap/projection.rs +++ b/nodedb-sql/src/planner/aggregate_cp_wrap/projection.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 -//! Grouped output projection and sequence-accessor checks. +//! Grouped output projection: the SELECT list restated over the group rows. -use super::expression::unqualify_columns; +use super::expression::{GroupRowNames, group_row_columns}; use crate::aggregate_walk::contains_aggregate; use crate::error::{Result, SqlError}; use crate::functions::registry::FunctionRegistry; use crate::parser::normalize::normalize_ident; +use crate::planner::agg_bind::{BindName, bind_aggregate_calls}; use crate::planner::agg_naming::group_key_row_name; use crate::planner::aggregate_order::compute_output_order_by_item; use crate::planner::cp_projection::ast_calls_sequence_accessor; @@ -14,44 +15,166 @@ use crate::resolver::ColumnScope; use crate::resolver::columns::TableScope; use crate::resolver::expr::convert_expr; use crate::types::plan::{first_sequence_accessor, referenced_columns}; -use crate::types::{AggOutputSlot, AggregateExpr, Projection, SqlExpr, SqlPlan}; +use crate::types::{ + AggOutputSlot, AggregateExpr, Projection, SqlExpr, SqlPlan, TimeseriesAggOutput, + TimeseriesScanPlan, +}; use sqlparser::ast; -/// The projection restating `plan`'s output columns in SELECT-list order, -/// with a [`Projection::CpComputed`] entry at each accessor item's position. -/// `plan` is the `Aggregate` the items were planned into. +/// The names a grouped plan's finalized rows carry. +pub(super) struct GroupRow<'a> { + /// The GROUP BY keys as planned. + group_by: &'a [SqlExpr], + /// Each key's name before any alias rename: the bare column, or + /// `group_{index}` for a computed key. + key_names: Vec, + /// Each key's name in the finalized row. + row_names: Vec, + /// The computed aggregates. Each one's alias is its name in the row. + aggregates: &'a [AggregateExpr], +} + +impl<'a> GroupRow<'a> { + /// The row names of `plan`, a grouped `Aggregate` or an aggregating + /// timeseries scan. + pub(super) fn of(plan: &'a SqlPlan) -> Result { + match plan { + SqlPlan::Aggregate { + group_by, + group_by_aliases, + aggregates, + .. + } => { + let key_names = key_names(group_by); + // The aggregate renames each aliased key to its alias. + let row_names = key_names + .iter() + .enumerate() + .map(|(index, name)| { + group_by_aliases + .get(index) + .cloned() + .flatten() + .unwrap_or_else(|| name.clone()) + }) + .collect(); + Ok(Self { + group_by, + key_names, + row_names, + aggregates, + }) + } + SqlPlan::TimeseriesScan(scan) => Self::of_timeseries(scan), + other => Err(SqlError::Unsupported { + detail: format!("a grouped SELECT list over a {} plan", other.variant_name()), + }), + } + } + + /// The row names of an aggregating timeseries scan. Its rows carry every + /// key under the output name `scan.output` gives it. + fn of_timeseries(scan: &'a TimeseriesScanPlan) -> Result { + let key_names = key_names(&scan.group_keys); + let row_names = scan + .group_keys + .iter() + .map(|key| { + scan.output + .iter() + .find(|column| match (column, key) { + ( + TimeseriesAggOutput::GroupColumn { name, .. }, + SqlExpr::Column { name: key_name, .. }, + ) => name == key_name, + (TimeseriesAggOutput::Bucket { .. }, SqlExpr::Column { .. }) => false, + (TimeseriesAggOutput::Bucket { .. }, _) => true, + _ => false, + }) + .and_then(|column| column.key_and_name(&scan.aggregates)) + .map(|(_, name)| name) + .ok_or_else(|| SqlError::Unsupported { + detail: format!( + "the grouped timeseries scan of '{}' emits no column for its \ + GROUP BY key {key:?}", + scan.collection + ), + }) + }) + .collect::>()?; + Ok(Self { + group_by: &scan.group_keys, + key_names, + row_names, + aggregates: &scan.aggregates, + }) + } + + /// The renames and computed-key matches that read an expression over the + /// input relations from the finalized row. + fn names(&self) -> GroupRowNames { + GroupRowNames { + renames: self + .key_names + .iter() + .zip(&self.row_names) + .filter(|(from, to)| from != to) + .map(|(from, to)| (from.clone(), to.clone())) + .collect(), + computed_keys: self + .group_by + .iter() + .zip(&self.row_names) + .filter(|(key, _)| !matches!(key, SqlExpr::Column { .. })) + .map(|(key, name)| (format!("{key:?}"), name.clone())) + .collect(), + } + } +} + +/// The pre-rename row name of each GROUP BY key. +fn key_names(group_by: &[SqlExpr]) -> Vec { + group_by + .iter() + .enumerate() + .map(|(index, key)| group_key_row_name(key, index)) + .collect() +} + +/// The projection restating `plan`'s output columns in SELECT-list order. +/// +/// A bare GROUP BY key or a bare aggregate reads its row column. A sequence +/// accessor item becomes a [`Projection::CpComputed`] entry. Any other item +/// becomes a [`Projection::Computed`] entry evaluated over the group row: +/// `ROUND(AVG(v), 1)`, `SUM(a) / COUNT(*)`, `CASE WHEN COUNT(*) > 1 ...`, +/// `UPPER(host)` over the key `host`, or a literal. pub(super) fn aggregate_cp_projection( plan: &SqlPlan, items: &[ast::SelectItem], functions: &FunctionRegistry, scope: &TableScope, ) -> Result> { - let (group_by, aggregates) = match plan { - SqlPlan::Aggregate { - group_by, - aggregates, - .. - } => (group_by, aggregates), - other => { - return Err(SqlError::Unsupported { - detail: format!("aggregate wrap over a {} plan", other.variant_name()), - }); - } - }; - let key_names: Vec = group_by + let row = GroupRow::of(plan)?; + let names = row.names(); + let by_item = compute_output_order_by_item(items, row.group_by, functions, scope)?; + // An item resolves against the input relations, so an unknown column is + // the usual resolve error, and against the group row's names, so a + // computed key (`group_0`) and an aggregate (`avg(v)`) are addressable. + // The reference check below then narrows to those names: the finalized + // group row carries nothing else. + let row_scope = scope.with_output_names( + row.key_names + .iter() + .chain(&row.row_names) + .cloned() + .chain(row.aggregates.iter().map(|a| a.alias.clone())), + ); + let readable: Vec<&str> = row + .row_names .iter() - .enumerate() - .map(|(index, key)| group_key_row_name(key, index)) + .map(String::as_str) + .chain(row.aggregates.iter().map(|a| a.alias.as_str())) .collect(); - let by_item = compute_output_order_by_item(items, group_by, functions, scope)?; - // The item resolves against the input relations, so an unknown column is - // the usual resolve error, and against the group-key row names, so a - // computed key (`group_0`) is addressable. The reference check below - // then narrows to the keys: the finalized group row carries nothing - // else the Control Plane can read. - let cp_scope = scope - .with_output_names(key_names.iter().cloned()) - .allowing_cp_functions(); let mut projection = Vec::with_capacity(items.len()); for (item, slots) in items.iter().zip(by_item) { @@ -63,67 +186,125 @@ pub(super) fn aggregate_cp_projection( | ast::SelectItem::QualifiedWildcard(..) => continue, }; if ast_calls_sequence_accessor(expr) { + let cp_scope = row_scope.clone().allowing_cp_functions(); let converted = convert_expr(expr, &ColumnScope::Relations(&cp_scope))?; if let Some(name) = first_sequence_accessor(&converted) { - let name = name.to_string(); - projection.push(cp_item( - converted, name, alias, expr, &key_names, functions, - )?); + if contains_aggregate(expr, functions) { + return Err(SqlError::SequencePerRowUnsupported { + name: name.to_string(), + }); + } + let expr = group_row_columns(converted, &names); + require_group_row_columns(&expr, &readable)?; + projection.push(Projection::CpComputed { expr, alias }); continue; } } - for slot in slots { - projection.push(Projection::Column(slot_row_name( - slot, &key_names, aggregates, - )?)); + if item_reads_its_slots(expr, &slots, functions) { + for slot in slots { + projection.push(Projection::Column(slot_row_name(slot, &row)?)); + } + continue; } + let expr = computed_over_group_row(expr, &row, &row_scope, functions)?; + let expr = group_row_columns(expr, &names); + require_group_row_columns(&expr, &readable)?; + projection.push(Projection::Computed { expr, alias }); } Ok(projection) } -/// The Control-Plane entry for one accessor item over a grouped result. +/// Whether a SELECT list over a plan grouped by `group_by` holds an item the +/// grouped plan does not emit as one of its columns, so a tail must evaluate +/// it over the group rows. +pub(crate) fn items_need_group_row_tail( + group_by: &[SqlExpr], + items: &[ast::SelectItem], + functions: &FunctionRegistry, + scope: &TableScope, +) -> Result { + let by_item = compute_output_order_by_item(items, group_by, functions, scope)?; + Ok(items.iter().zip(&by_item).any(|(item, slots)| match item { + ast::SelectItem::UnnamedExpr(expr) | ast::SelectItem::ExprWithAlias { expr, .. } => { + !item_reads_its_slots(expr, slots, functions) + } + ast::SelectItem::ExprWithAliases { .. } + | ast::SelectItem::Wildcard(_) + | ast::SelectItem::QualifiedWildcard(..) => false, + })) +} + +/// Whether the item is exactly its output slots: one bare GROUP BY key or one +/// bare aggregate. `GROUPING(col)` items are slots too. +fn item_reads_its_slots( + expr: &ast::Expr, + slots: &[AggOutputSlot], + functions: &FunctionRegistry, +) -> bool { + match slots { + [AggOutputSlot::GroupKey(_)] => true, + [AggOutputSlot::Aggregate(_)] => { + !crate::aggregate_walk::computes_over_aggregate(expr, functions) + } + [] => false, + // Several slots from one item: `GROUPING(a) + GROUPING(b)` or an + // expression over several aggregates. An expression over aggregates + // is evaluated; a GROUPING pair keeps its slots. + _ => !contains_aggregate(expr, functions), + } +} + +/// `expr` with each aggregate call bound to the column the group row +/// carries it under, converted against the group row's names. /// -/// The item holds no aggregate: the Control Plane evaluates it over the -/// finalized group row, which carries aggregate values under their output -/// names but no per-row aggregate state. Every column it references is a -/// group key, unqualified so the reference matches the row's bare key. -fn cp_item( - converted: SqlExpr, - accessor: String, - alias: String, - raw: &ast::Expr, - key_names: &[String], +/// A SELECT item names input columns, never another item's output alias, so +/// no alias substitution applies. +fn computed_over_group_row( + expr: &ast::Expr, + row: &GroupRow<'_>, + row_scope: &TableScope, functions: &FunctionRegistry, -) -> Result { - if contains_aggregate(raw, functions) { - return Err(SqlError::SequencePerRowUnsupported { name: accessor }); +) -> Result { + let mut known = row.aggregates.to_vec(); + let bound = bind_aggregate_calls( + expr, + &[], + &mut known, + functions, + BindName::Output, + row_scope, + )?; + if known.len() != row.aggregates.len() { + return Err(SqlError::Unsupported { + detail: format!("the grouped plan computes no aggregate for `{expr}`"), + }); } - for column in referenced_columns(&converted) { - let bare = column.rsplit('.').next().unwrap_or(&column); - if !key_names.iter().any(|key| key.eq_ignore_ascii_case(bare)) { - return Err(SqlError::Unsupported { - detail: format!( - "column '{column}' beside a sequence accessor in a grouped SELECT list \ - must be a GROUP BY key" - ), - }); - } + convert_expr(&bound, &ColumnScope::Relations(row_scope)) +} + +/// Refuse an item that reads a column the finalized group row does not +/// carry: a column that is neither a GROUP BY key nor inside an aggregate. +/// +/// A column inside a GROUP BY expression is not a key by itself. With +/// `GROUP BY x / 0`, the item `x` reads no key, as PostgreSQL rules. The item +/// `x / 0` reads the key, and `group_row_columns` has already rewritten it to +/// the key's row column. +fn require_group_row_columns(expr: &SqlExpr, readable: &[&str]) -> Result<()> { + match referenced_columns(expr).into_iter().find(|column| { + !readable + .iter() + .any(|name| name.eq_ignore_ascii_case(column)) + }) { + Some(column) => Err(SqlError::UngroupedColumn { column }), + None => Ok(()), } - Ok(Projection::CpComputed { - expr: unqualify_columns(converted), - alias, - }) } /// The key a finalized group row carries one output slot under. -fn slot_row_name( - slot: AggOutputSlot, - key_names: &[String], - aggregates: &[AggregateExpr], -) -> Result { +fn slot_row_name(slot: AggOutputSlot, row: &GroupRow<'_>) -> Result { match slot { - AggOutputSlot::GroupKey(index) => key_names.get(index).cloned(), - AggOutputSlot::Aggregate(index) => aggregates.get(index).map(|a| a.alias.clone()), + AggOutputSlot::GroupKey(index) => row.row_names.get(index).cloned(), + AggOutputSlot::Aggregate(index) => row.aggregates.get(index).map(|a| a.alias.clone()), } .ok_or_else(|| SqlError::Unsupported { detail: format!("aggregate output slot {slot:?} names no output column"), diff --git a/nodedb-sql/src/planner/aggregate_cp_wrap/wrap.rs b/nodedb-sql/src/planner/aggregate_cp_wrap/wrap.rs index 7a71c603e..f852dae2b 100644 --- a/nodedb-sql/src/planner/aggregate_cp_wrap/wrap.rs +++ b/nodedb-sql/src/planner/aggregate_cp_wrap/wrap.rs @@ -1,35 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 -//! Wrapping a grouped plan so a Control-Plane-computed SELECT-list item keeps -//! its position. +//! Wrapping a grouped plan so every SELECT-list item keeps its position and +//! its value. //! -//! An `Aggregate` plan emits group keys and aggregate values only. A SELECT -//! item that calls a sequence accessor is neither, so the aggregate cannot -//! carry it. The planner wraps the finished aggregate in a `Subquery` whose -//! projection restates every output column in SELECT-list order and places a -//! [`crate::types::Projection::CpComputed`] entry at the item's position. The wrap runs -//! after ORDER BY and LIMIT are attached, so those stay on the aggregate. +//! A grouped plan (`Aggregate`, or an aggregating timeseries scan) emits +//! group keys and aggregate values only. A SELECT item that is neither, such +//! as a sequence accessor call, `ROUND(AVG(v), 1)`, `SUM(a) / COUNT(*)`, or +//! `UPPER(host)` over the key `host`, is evaluated over the finalized group +//! rows. The planner wraps the finished plan in a `Subquery` whose projection +//! restates every output column in SELECT-list order and evaluates each such +//! item at its position. The wrap runs after ORDER BY and LIMIT are attached, +//! so those stay on the grouped plan. use crate::types::CtePlan; use sqlparser::ast; -use super::projection::aggregate_cp_projection; +use super::projection::{aggregate_cp_projection, items_need_group_row_tail}; use crate::error::{Result, SqlError}; use crate::functions::registry::FunctionRegistry; use crate::planner::cp_projection::ast_sequence_accessor; use crate::resolver::columns::TableScope; -use crate::types::SqlPlan; +use crate::types::{SqlExpr, SqlPlan, TimeseriesScanPlan}; -/// Wrap `plan` when `items` hold a sequence accessor the aggregate cannot -/// carry. A plan with no such item, or one that is not grouped, returns -/// unchanged: `convert_projection` marks the item on a row plan itself. +/// Wrap `plan` when `items` hold an item the grouped plan cannot carry. A +/// plan with no such item, or one that is not grouped, returns unchanged: +/// `convert_projection` marks a sequence accessor on a row plan itself. /// /// `grouped` says whether the SELECT aggregates (GROUP BY, or an aggregate /// call in its list). `scope` is the relation set the SELECT resolved -/// against. A grouped SELECT that does not produce the statement's output -/// rows (a subquery, CTE body, UNION branch, derived table) refuses the -/// item: the aggregate planner never converts a bare non-key item, so -/// nothing else reports it. +/// against. A sequence accessor in a grouped SELECT that does not produce +/// the statement's output rows (a subquery, CTE body, UNION branch, derived +/// table) is refused: no session evaluates it there. pub fn wrap_aggregate_cp_items( plan: SqlPlan, items: &[ast::SelectItem], @@ -40,11 +41,24 @@ pub fn wrap_aggregate_cp_items( if !grouped || !scope.is_row_scope() { return Ok(plan); } - let Some(accessor) = items.iter().find_map(item_accessor) else { - return Ok(plan); - }; - if !scope.is_statement_output() { - return Err(SqlError::SequencePerRowUnsupported { name: accessor }); + let accessor = items.iter().find_map(item_accessor); + if let Some(name) = &accessor + && !scope.is_statement_output() + { + return Err(SqlError::SequencePerRowUnsupported { name: name.clone() }); + } + if accessor.is_none() { + let needs_tail = match grouped_keys(&plan) { + Some(group_by) => items_need_group_row_tail(group_by, items, functions, scope)?, + // A shape this wrap cannot restate is refused below when an item + // computes over an aggregate. Every other item it carries. + None => items + .iter() + .any(|item| item_computes_over_aggregate(item, functions)), + }; + if !needs_tail { + return Ok(plan); + } } match plan { SqlPlan::Cte(CtePlan { definitions, outer }) => Ok(SqlPlan::Cte(CtePlan { @@ -53,7 +67,7 @@ pub fn wrap_aggregate_cp_items( *outer, items, grouped, functions, scope, )?), })), - SqlPlan::Aggregate { .. } => { + SqlPlan::Aggregate { .. } | SqlPlan::TimeseriesScan(_) => { let projection = aggregate_cp_projection(&plan, items, functions, scope)?; Ok(SqlPlan::Subquery { input: Box::new(plan), @@ -79,7 +93,11 @@ pub fn wrap_aggregate_cp_items( offset, distinct, limit, - } if matches!(*input, SqlPlan::Aggregate { .. }) && projection.is_empty() => { + } if matches!( + *input, + SqlPlan::Aggregate { .. } | SqlPlan::TimeseriesScan(_) + ) && projection.is_empty() => + { let projection = aggregate_cp_projection(&input, items, functions, scope)?; Ok(SqlPlan::Subquery { input, @@ -93,19 +111,51 @@ pub fn wrap_aggregate_cp_items( }) } // A grouped query the engine rules planned as another shape (a - // timeseries bucket scan, a grouped query joined to a scalar - // subquery) names its output columns by rules this wrap does not - // restate. Refusing beats dropping the item from the result. + // grouped query joined to a scalar subquery) names its output columns + // by rules this wrap does not restate. Refusing beats dropping the + // item from the result, or returning it unevaluated. other => Err(SqlError::Unsupported { detail: format!( - "a sequence accessor in the SELECT list of a grouped query planned as {}; \ - call it in a FROM-less SELECT instead", + "a computed item in the SELECT list of a grouped query planned as {}", other.variant_name() ), }), } } +/// The GROUP BY keys of the grouped plan `plan` wraps, when it is one. +/// +/// A raw timeseries scan is not grouped. Neither is any other shape: the wrap +/// leaves it as planned. +fn grouped_keys(plan: &SqlPlan) -> Option<&[SqlExpr]> { + match plan { + SqlPlan::Aggregate { group_by, .. } => Some(group_by), + SqlPlan::TimeseriesScan(scan) if is_grouped_scan(scan) => Some(&scan.group_keys), + SqlPlan::Cte(CtePlan { outer, .. }) => grouped_keys(outer), + SqlPlan::Subquery { + input, projection, .. + } if projection.is_empty() => grouped_keys(input), + _ => None, + } +} + +/// Whether a timeseries scan groups or aggregates, rather than reading rows. +fn is_grouped_scan(scan: &TimeseriesScanPlan) -> bool { + !scan.aggregates.is_empty() || scan.bucket_interval_ms > 0 || !scan.group_by.is_empty() +} + +/// Whether a SELECT item computes over an aggregate: `ROUND(AVG(v), 1)`. +fn item_computes_over_aggregate(item: &ast::SelectItem, functions: &FunctionRegistry) -> bool { + match item { + ast::SelectItem::UnnamedExpr(expr) | ast::SelectItem::ExprWithAlias { expr, .. } => { + crate::aggregate_walk::computes_over_aggregate(expr, functions) + } + ast::SelectItem::ExprWithAliases { .. } + | ast::SelectItem::Wildcard(_) + | ast::SelectItem::QualifiedWildcard(..) => false, + } +} + /// The sequence accessor a SELECT item calls, when it calls one. fn item_accessor(item: &ast::SelectItem) -> Option { match item { @@ -315,9 +365,47 @@ mod tests { fn a_grouped_accessor_item_may_reference_group_keys_only() { let err = plan_err("SELECT grp, COUNT(*), nextval('s') + other FROM t GROUP BY grp"); assert!( - matches!(err, SqlError::Unsupported { ref detail } if detail.contains("'other'")), + matches!(err, SqlError::UngroupedColumn { ref column } if column == "other"), + "got {err:?}" + ); + } + + /// A column inside a GROUP BY expression is not a key by itself, as + /// PostgreSQL rules. The key expression itself reads the key. + #[test] + fn a_column_inside_a_group_by_expression_is_ungrouped() { + let err = plan_err("SELECT other, COUNT(*) FROM t GROUP BY other / 2"); + assert!( + matches!(err, SqlError::UngroupedColumn { ref column } if column == "other"), "got {err:?}" ); + plan("SELECT other / 2 AS half, COUNT(*) FROM t GROUP BY other / 2"); + } + + /// An aggregate inside a scalar function, a CASE, or arithmetic binds to + /// its group-row column. The item is computed over that row. + #[test] + fn an_aggregate_inside_any_expression_binds_to_its_row_column() { + for sql in [ + "SELECT ROUND(AVG(other), 1) AS m FROM t", + "SELECT grp, CASE WHEN COUNT(*) > 1 THEN 'many' ELSE 'one' END AS s FROM t GROUP BY grp", + "SELECT grp, ROUND(ABS(MIN(other) - MAX(other)), 2) AS w FROM t GROUP BY grp", + ] { + let plan = plan(sql); + let SqlPlan::Subquery { projection, .. } = plan else { + panic!("{sql}: expected a group-row tail, got {plan:?}"); + }; + let Some(Projection::Computed { expr, .. }) = projection.last() else { + panic!("{sql}: expected a computed item, got {projection:?}"); + }; + let rendered = format!("{expr:?}"); + assert!( + !rendered.contains("name: \"avg\"") + && !rendered.contains("name: \"count\"") + && !rendered.contains("name: \"min\""), + "{sql}: an aggregate call is left unbound: {rendered}" + ); + } } #[test] diff --git a/nodedb-sql/src/planner/array_dml.rs b/nodedb-sql/src/planner/array_dml.rs index e812b76fc..df6547d4d 100644 --- a/nodedb-sql/src/planner/array_dml.rs +++ b/nodedb-sql/src/planner/array_dml.rs @@ -136,15 +136,17 @@ fn coord_compatible(c: &ArrayCoordLiteral, dtype: ArrayDimType) -> bool { ) } +/// Whether a literal's kind fits the attr type. An integer fits every +/// integer and float type, a float fits the float types. Range is checked +/// where the literal narrows to the engine cell. fn attr_compatible(a: &ArrayAttrLiteral, dtype: ArrayAttrType) -> bool { - matches!( - (a, dtype), - (ArrayAttrLiteral::Int64(_), ArrayAttrType::Int64) - | (ArrayAttrLiteral::Int64(_), ArrayAttrType::Float64) - | (ArrayAttrLiteral::Float64(_), ArrayAttrType::Float64) - | (ArrayAttrLiteral::String(_), ArrayAttrType::String) - | (ArrayAttrLiteral::Bytes(_), ArrayAttrType::Bytes) - ) + match a { + ArrayAttrLiteral::Int64(_) => dtype.is_integer() || dtype.is_float(), + ArrayAttrLiteral::Float64(_) => dtype.is_float(), + ArrayAttrLiteral::String(_) => dtype == ArrayAttrType::String, + ArrayAttrLiteral::Bytes(_) => dtype == ArrayAttrType::Bytes, + ArrayAttrLiteral::Null => true, + } } #[cfg(test)] diff --git a/nodedb-sql/src/planner/array_fn/helpers.rs b/nodedb-sql/src/planner/array_fn/helpers.rs index ae063543f..d4a2a9874 100644 --- a/nodedb-sql/src/planner/array_fn/helpers.rs +++ b/nodedb-sql/src/planner/array_fn/helpers.rs @@ -7,10 +7,12 @@ use sqlparser::ast; use nodedb_types::Value; +use crate::catalog::ArrayCatalogView; use crate::error::{Result, SqlError}; use crate::parser::normalize::{SCHEMA_QUALIFIED_MSG, normalize_ident}; +use crate::parser::object_literal::parse_object_literal_complete; use crate::types::SqlCatalog; -use crate::types_array::ArrayCoordLiteral; +use crate::types_array::{ArrayCoordLiteral, ArraySliceAst, NamedDimRange}; pub(super) fn collect_args(args: &[ast::FunctionArg]) -> Vec { args.iter() @@ -120,6 +122,49 @@ pub(super) fn is_null_literal(expr: &ast::Expr) -> bool { ) } +/// Parse a slice-predicate argument of `fn_name` over array `name`. +/// +/// The argument is a quoted string holding the brace-form object literal +/// `'{dim: [lo, hi], ...}'`: the PostgreSQL dialect does not accept a bare +/// `{...}` in expression position. Each named dim must exist on the array. +pub(super) fn parse_slice_arg( + expr: &ast::Expr, + fn_name: &str, + name: &str, + view: &ArrayCatalogView, +) -> Result { + let slice_str = expect_string_literal(expr, &format!("{fn_name} slice predicate"))?; + let parsed = + parse_object_literal_complete(&slice_str).ok_or_else(|| SqlError::Unsupported { + detail: format!("{fn_name}: slice predicate must be an object literal: {slice_str}"), + })?; + let map = parsed.map_err(|detail| SqlError::Unsupported { + detail: format!("{fn_name}: slice parse: {detail}"), + })?; + let mut dim_ranges: Vec = Vec::with_capacity(map.len()); + for (dim, val) in map { + if !view.dims.iter().any(|d| d.name == dim) { + return Err(SqlError::Unsupported { + detail: format!("{fn_name}: array '{name}' has no dim '{dim}'"), + }); + } + let arr = match val { + Value::Array(a) if a.len() == 2 => a, + _ => { + return Err(SqlError::Unsupported { + detail: format!( + "{fn_name}: dim '{dim}' range must be a 2-element array [lo, hi]" + ), + }); + } + }; + let lo = value_to_coord_literal(&arr[0], &dim)?; + let hi = value_to_coord_literal(&arr[1], &dim)?; + dim_ranges.push(NamedDimRange { dim, lo, hi }); + } + Ok(ArraySliceAst { dim_ranges }) +} + pub(super) fn value_to_coord_literal(v: &Value, dim: &str) -> Result { match v { Value::Integer(i) => Ok(ArrayCoordLiteral::Int64(*i)), diff --git a/nodedb-sql/src/planner/array_fn/mod.rs b/nodedb-sql/src/planner/array_fn/mod.rs index c815f2e8b..d8cb5b395 100644 --- a/nodedb-sql/src/planner/array_fn/mod.rs +++ b/nodedb-sql/src/planner/array_fn/mod.rs @@ -16,7 +16,10 @@ mod helpers; mod maint_fn; +mod pushdown; mod table_fn; +mod zone; pub use maint_fn::try_plan_array_maint_fn; -pub use table_fn::try_plan_array_table_fn; +pub use pushdown::push_down_into_slice; +pub use table_fn::{ArrayTableFn, try_plan_array_table_fn}; diff --git a/nodedb-sql/src/planner/array_fn/pushdown.rs b/nodedb-sql/src/planner/array_fn/pushdown.rs new file mode 100644 index 000000000..1ceb437da --- /dev/null +++ b/nodedb-sql/src/planner/array_fn/pushdown.rs @@ -0,0 +1,484 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Pushdown of an outer query's constraints into the slice it reads. +//! +//! The outer query over an `ARRAY_SLICE` relation keeps every clause in its +//! own tail, so results never depend on what this module folds. Folding only +//! narrows what the Data Plane reads: +//! +//! - An outer `LIMIT n` caps the slice when nothing else in the outer query +//! reads past the cut. The cap becomes the smaller nonzero of the two. +//! - A top-level `AND` conjunct that bounds a dim narrows that dim's range. +//! This runs only when the slice has no row cap of its own: under a cap, a +//! narrower range changes which cells the cap keeps. + +use crate::types::{ + ArrayCatalogView, ArraySlicePlan, BinaryOp, CompareOp, Filter, FilterExpr, JoinType, SqlExpr, + SqlPlan, SqlValue, +}; +use crate::types_array::{ArrayCoordLiteral, ArrayDimType, ArrayDomainBound, NamedDimRange}; + +/// Narrow `slice` by the constraints `outer` places on `relation`, the name +/// the outer query reads the slice's rows by. +pub fn push_down_into_slice( + slice: &mut ArraySlicePlan, + view: &ArrayCatalogView, + outer: &SqlPlan, + relation: &str, +) { + if slice.limit == 0 + && let Some(filters) = relation_scan_filters(outer, relation) + { + for filter in filters { + narrow_by_filter(slice, view, &filter.expr); + } + } + if let Some(limit) = foldable_limit(outer, relation) { + slice.limit = folded_cap(slice.limit, limit); + } +} + +/// The filters of the scan that reads `relation` as the outer query's row +/// source. The scan sits at the root, under an aggregate, or as the probe +/// side of a join a WHERE subquery adds. Each filter there is a conjunct +/// every output row satisfies. +fn relation_scan_filters<'a>(plan: &'a SqlPlan, relation: &str) -> Option<&'a [Filter]> { + if let SqlPlan::Scan { + collection, + filters, + .. + } = plan + && collection == relation + { + return Some(filters); + } + if let SqlPlan::Aggregate { input, .. } = plan { + return relation_scan_filters(input, relation); + } + if let SqlPlan::Join { + left, join_type, .. + } = plan + && matches!(join_type, JoinType::Semi | JoinType::Anti | JoinType::Cross) + { + return relation_scan_filters(left, relation); + } + None +} + +/// The outer `LIMIT`, when the outer query is a plain read of `relation` +/// that no filter, sort, offset, DISTINCT or window reads past. +fn foldable_limit(plan: &SqlPlan, relation: &str) -> Option { + let SqlPlan::Scan { + collection, + filters, + sort_keys, + limit: Some(limit), + offset: 0, + distinct: false, + window_functions, + .. + } = plan + else { + return None; + }; + (collection == relation + && filters.is_empty() + && sort_keys.is_empty() + && window_functions.is_empty() + && *limit > 0) + .then_some(*limit) +} + +/// The smaller nonzero of the slice cap and the outer limit. A limit past +/// the `u32` cap range leaves an uncapped slice uncapped. +fn folded_cap(cap: u32, limit: usize) -> u32 { + match u32::try_from(limit) { + Ok(limit) if cap == 0 || limit < cap => limit, + Ok(_) | Err(_) => cap, + } +} + +/// Apply every top-level conjunct of `expr` that bounds a dim. +fn narrow_by_filter(slice: &mut ArraySlicePlan, view: &ArrayCatalogView, expr: &FilterExpr) { + match expr { + FilterExpr::Expr(sql) => narrow_by_expr(slice, view, sql), + FilterExpr::And(filters) => { + for filter in filters { + narrow_by_filter(slice, view, &filter.expr); + } + } + FilterExpr::Comparison { field, op, value } => { + let op = match op { + CompareOp::Eq => Some(BinaryOp::Eq), + CompareOp::Gt => Some(BinaryOp::Gt), + CompareOp::Ge => Some(BinaryOp::Ge), + CompareOp::Lt => Some(BinaryOp::Lt), + CompareOp::Le => Some(BinaryOp::Le), + CompareOp::Ne => None, + }; + if let Some(op) = op { + narrow_dim(slice, view, field, op, value); + } + } + FilterExpr::Between { field, low, high } => { + narrow_dim(slice, view, field, BinaryOp::Ge, low); + narrow_dim(slice, view, field, BinaryOp::Le, high); + } + FilterExpr::InList { .. } + | FilterExpr::IsNull { .. } + | FilterExpr::IsNotNull { .. } + | FilterExpr::Or(_) + | FilterExpr::Not(_) => {} + } +} + +fn narrow_by_expr(slice: &mut ArraySlicePlan, view: &ArrayCatalogView, expr: &SqlExpr) { + match expr { + SqlExpr::BinaryOp { + left, + op: BinaryOp::And, + right, + } => { + narrow_by_expr(slice, view, left); + narrow_by_expr(slice, view, right); + } + SqlExpr::BinaryOp { left, op, right } => { + if let (SqlExpr::Column { name, .. }, SqlExpr::Literal(value)) = + (left.as_ref(), right.as_ref()) + { + narrow_dim(slice, view, name, *op, value); + } + } + SqlExpr::Between { + expr, + low, + high, + negated: false, + } => { + if let (SqlExpr::Column { name, .. }, SqlExpr::Literal(lo), SqlExpr::Literal(hi)) = + (expr.as_ref(), low.as_ref(), high.as_ref()) + { + narrow_dim(slice, view, name, BinaryOp::Ge, lo); + narrow_dim(slice, view, name, BinaryOp::Le, hi); + } + } + _ => {} + } +} + +/// A dim bound on one numeric axis. Integer and timestamp dims compare as +/// `i64`, float dims as `f64`. +#[derive(Clone, Copy, PartialEq, PartialOrd)] +enum Num { + Int(i64), + Float(f64), +} + +/// Tighten the range of dim `name` by `name value`. +/// +/// The new range holds every coordinate the predicate admits, and may hold +/// more: a strict bound on a float dim stays inclusive. A predicate that +/// would empty the range leaves it unchanged; the outer filter answers it. +fn narrow_dim( + slice: &mut ArraySlicePlan, + view: &ArrayCatalogView, + name: &str, + op: BinaryOp, + value: &SqlValue, +) { + let Some(dim) = view.dims.iter().find(|d| d.name == name) else { + return; + }; + let Some((lo_bound, hi_bound)) = predicate_bounds(dim.dtype, op, value) else { + return; + }; + let existing = slice.slice.dim_ranges.iter().position(|r| r.dim == name); + let (lo, hi) = match existing { + Some(i) => { + let range = &slice.slice.dim_ranges[i]; + ( + literal_num(dim.dtype, &range.lo), + literal_num(dim.dtype, &range.hi), + ) + } + None => ( + domain_num(dim.dtype, &dim.lo), + domain_num(dim.dtype, &dim.hi), + ), + }; + let (Some(lo), Some(hi)) = (lo, hi) else { + return; + }; + let lo = match lo_bound { + Some(bound) if bound > lo => bound, + Some(_) | None => lo, + }; + let hi = match hi_bound { + Some(bound) if bound < hi => bound, + Some(_) | None => hi, + }; + if lo > hi { + return; + } + let range = NamedDimRange { + dim: name.to_string(), + lo: num_literal(lo), + hi: num_literal(hi), + }; + match existing { + Some(i) => slice.slice.dim_ranges[i] = range, + None => slice.slice.dim_ranges.push(range), + } +} + +/// The `(lower, upper)` bounds `dim value` places on a dim of type +/// `dtype`, each `None` when the predicate leaves that side open. `None` +/// when the predicate bounds nothing this module can express. +fn predicate_bounds( + dtype: ArrayDimType, + op: BinaryOp, + value: &SqlValue, +) -> Option<(Option, Option)> { + let (floor, ceil, strict_step) = match (dtype, value) { + (ArrayDimType::Int64, SqlValue::Int(v)) => (Num::Int(*v), Num::Int(*v), true), + (ArrayDimType::Int64, SqlValue::Float(f)) if f.is_finite() => { + (Num::Int(f.floor() as i64), Num::Int(f.ceil() as i64), false) + } + (ArrayDimType::Float64, SqlValue::Int(v)) => { + (Num::Float(*v as f64), Num::Float(*v as f64), false) + } + (ArrayDimType::Float64, SqlValue::Float(f)) if f.is_finite() => { + (Num::Float(*f), Num::Float(*f), false) + } + (ArrayDimType::TimestampMs, SqlValue::Timestamp(at) | SqlValue::Timestamptz(at)) => ( + Num::Int(at.micros.div_euclid(1000)), + Num::Int(-(-at.micros).div_euclid(1000)), + false, + ), + _ => return None, + }; + // A strict bound on an integer literal moves one step inward. Every + // other strict bound stays inclusive, which only widens the range. + let above = |n: Num| match (n, strict_step) { + (Num::Int(v), true) => v.checked_add(1).map(Num::Int), + (other, _) => Some(other), + }; + let below = |n: Num| match (n, strict_step) { + (Num::Int(v), true) => v.checked_sub(1).map(Num::Int), + (other, _) => Some(other), + }; + match op { + BinaryOp::Eq => Some((Some(floor), Some(ceil))), + BinaryOp::Ge => Some((Some(floor), None)), + BinaryOp::Gt => Some((Some(above(floor)?), None)), + BinaryOp::Le => Some((None, Some(ceil))), + BinaryOp::Lt => Some((None, Some(below(ceil)?))), + BinaryOp::Add + | BinaryOp::Sub + | BinaryOp::Mul + | BinaryOp::Div + | BinaryOp::Mod + | BinaryOp::Ne + | BinaryOp::And + | BinaryOp::Or + | BinaryOp::Concat => None, + } +} + +fn literal_num(dtype: ArrayDimType, lit: &ArrayCoordLiteral) -> Option { + match (dtype, lit) { + (ArrayDimType::Int64 | ArrayDimType::TimestampMs, ArrayCoordLiteral::Int64(v)) => { + Some(Num::Int(*v)) + } + (ArrayDimType::Float64, ArrayCoordLiteral::Int64(v)) => Some(Num::Float(*v as f64)), + (ArrayDimType::Float64, ArrayCoordLiteral::Float64(v)) => Some(Num::Float(*v)), + _ => None, + } +} + +fn domain_num(dtype: ArrayDimType, bound: &ArrayDomainBound) -> Option { + match (dtype, bound) { + (ArrayDimType::Int64, ArrayDomainBound::Int64(v)) + | (ArrayDimType::TimestampMs, ArrayDomainBound::TimestampMs(v)) => Some(Num::Int(*v)), + (ArrayDimType::Float64, ArrayDomainBound::Float64(v)) => Some(Num::Float(*v)), + _ => None, + } +} + +fn num_literal(n: Num) -> ArrayCoordLiteral { + match n { + Num::Int(v) => ArrayCoordLiteral::Int64(v), + Num::Float(v) => ArrayCoordLiteral::Float64(v), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::temporal::TemporalScope; + use crate::types::EngineType; + use crate::types_array::{ArrayAttrAst, ArrayAttrType, ArrayDimAst, ArraySliceAst}; + + fn view() -> ArrayCatalogView { + ArrayCatalogView { + name: "img".into(), + dims: vec![ + ArrayDimAst { + name: "x".into(), + dtype: ArrayDimType::Int64, + lo: ArrayDomainBound::Int64(0), + hi: ArrayDomainBound::Int64(3), + }, + ArrayDimAst { + name: "y".into(), + dtype: ArrayDimType::Int64, + lo: ArrayDomainBound::Int64(0), + hi: ArrayDomainBound::Int64(3), + }, + ], + attrs: vec![ArrayAttrAst { + name: "nir".into(), + dtype: ArrayAttrType::Float64, + nullable: true, + }], + tile_extents: vec![2, 2], + } + } + + fn slice(limit: u32) -> ArraySlicePlan { + ArraySlicePlan { + name: "img".into(), + slice: ArraySliceAst::default(), + attr_projection: Vec::new(), + limit, + temporal: TemporalScope::default(), + } + } + + fn outer(filters: Vec, limit: Option) -> SqlPlan { + SqlPlan::Scan { + collection: "rel".into(), + alias: None, + engine: EngineType::DocumentSchemaless, + filters, + projection: Vec::new(), + sort_keys: Vec::new(), + limit, + offset: 0, + distinct: false, + window_functions: Vec::new(), + temporal: TemporalScope::default(), + } + } + + fn cmp(column: &str, op: BinaryOp, value: i64) -> SqlExpr { + SqlExpr::BinaryOp { + left: Box::new(SqlExpr::Column { + table: None, + name: column.into(), + }), + op, + right: Box::new(SqlExpr::Literal(SqlValue::Int(value))), + } + } + + fn and(left: SqlExpr, right: SqlExpr) -> Filter { + Filter { + expr: FilterExpr::Expr(SqlExpr::BinaryOp { + left: Box::new(left), + op: BinaryOp::And, + right: Box::new(right), + }), + } + } + + fn range(s: &ArraySlicePlan, dim: &str) -> Option<(ArrayCoordLiteral, ArrayCoordLiteral)> { + s.slice + .dim_ranges + .iter() + .find(|r| r.dim == dim) + .map(|r| (r.lo.clone(), r.hi.clone())) + } + + #[test] + fn dim_conjuncts_narrow_an_uncapped_slice() { + let mut s = slice(0); + let filters = vec![and(cmp("x", BinaryOp::Gt, 0), cmp("y", BinaryOp::Le, 1))]; + push_down_into_slice(&mut s, &view(), &outer(filters, None), "rel"); + assert_eq!( + range(&s, "x"), + Some((ArrayCoordLiteral::Int64(1), ArrayCoordLiteral::Int64(3))) + ); + assert_eq!( + range(&s, "y"), + Some((ArrayCoordLiteral::Int64(0), ArrayCoordLiteral::Int64(1))) + ); + } + + #[test] + fn a_capped_slice_keeps_its_ranges() { + let mut s = slice(2); + let filters = vec![Filter { + expr: FilterExpr::Expr(cmp("x", BinaryOp::Eq, 3)), + }]; + push_down_into_slice(&mut s, &view(), &outer(filters, None), "rel"); + assert!(s.slice.dim_ranges.is_empty()); + assert_eq!(s.limit, 2); + } + + #[test] + fn an_attr_conjunct_and_an_or_leave_the_ranges() { + let mut s = slice(0); + let or = Filter { + expr: FilterExpr::Expr(SqlExpr::BinaryOp { + left: Box::new(cmp("x", BinaryOp::Eq, 1)), + op: BinaryOp::Or, + right: Box::new(cmp("x", BinaryOp::Eq, 2)), + }), + }; + let attr = Filter { + expr: FilterExpr::Expr(cmp("nir", BinaryOp::Gt, 0)), + }; + push_down_into_slice(&mut s, &view(), &outer(vec![or, attr], None), "rel"); + assert!(s.slice.dim_ranges.is_empty()); + } + + #[test] + fn an_empty_intersection_leaves_the_range() { + let mut s = slice(0); + let filters = vec![and(cmp("x", BinaryOp::Ge, 2), cmp("x", BinaryOp::Lt, 1))]; + push_down_into_slice(&mut s, &view(), &outer(filters, None), "rel"); + assert_eq!( + range(&s, "x"), + Some((ArrayCoordLiteral::Int64(2), ArrayCoordLiteral::Int64(3))) + ); + } + + #[test] + fn a_plain_limit_folds_to_the_smaller_cap() { + let mut s = slice(0); + push_down_into_slice(&mut s, &view(), &outer(Vec::new(), Some(5)), "rel"); + assert_eq!(s.limit, 5); + + let mut s = slice(2); + push_down_into_slice(&mut s, &view(), &outer(Vec::new(), Some(100)), "rel"); + assert_eq!(s.limit, 2); + } + + #[test] + fn a_limit_under_a_filter_does_not_fold() { + let mut s = slice(0); + let filters = vec![Filter { + expr: FilterExpr::Expr(cmp("nir", BinaryOp::Gt, 0)), + }]; + push_down_into_slice(&mut s, &view(), &outer(filters, Some(5)), "rel"); + assert_eq!(s.limit, 0); + } + + #[test] + fn limit_zero_does_not_fold() { + let mut s = slice(0); + push_down_into_slice(&mut s, &view(), &outer(Vec::new(), Some(0)), "rel"); + assert_eq!(s.limit, 0); + } +} diff --git a/nodedb-sql/src/planner/array_fn/table_fn.rs b/nodedb-sql/src/planner/array_fn/table_fn.rs index bd947ca70..d8ca61c55 100644 --- a/nodedb-sql/src/planner/array_fn/table_fn.rs +++ b/nodedb-sql/src/planner/array_fn/table_fn.rs @@ -3,37 +3,73 @@ //! `SELECT * FROM ARRAY_*(...)` table-valued function planning: //! slice / project / aggregate / elementwise. -use crate::types::{ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan}; +use crate::types::{ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan}; use sqlparser::ast; -use nodedb_types::Value; - use super::helpers::{ collect_args, expect_string_array, expect_string_literal, expect_u32, is_null_literal, - require_array_name, value_to_coord_literal, + parse_slice_arg, require_array_name, }; +use super::zone::parse_zone_args; use crate::error::{Result, SqlError}; -use crate::parser::object_literal::parse_object_literal_complete; +use crate::resolver::array_tvf::check_agg_reducer; use crate::temporal::TemporalScope; use crate::types::{SqlCatalog, SqlPlan}; -use crate::types_array::{ArrayBinaryOpAst, ArrayReducerAst, ArraySliceAst, NamedDimRange}; +use crate::types_array::{ArrayBinaryOpAst, ArrayReducerAst, ArraySliceAst}; + +/// One planned array table-valued function call. +/// +/// `Slice` and `Project` read cells: their rows carry the array's dims and +/// the requested attrs, so a SELECT reads them as a relation. `Agg` and +/// `Elementwise` emit rows of their own shape and form no such relation. +/// +/// `ARRAY_PROJECT` is a slice over the whole domain with a required attr +/// list, so it plans to the slice read and keeps only its own name. +#[derive(Debug, Clone)] +pub enum ArrayTableFn { + Slice(ArraySlicePlan), + Project(ArraySlicePlan), + Agg(ArrayAggPlan), + Elementwise(ArrayElementwisePlan), +} + +impl ArrayTableFn { + /// The plan that reads this call's rows. + pub fn into_plan(self) -> SqlPlan { + match self { + ArrayTableFn::Slice(plan) | ArrayTableFn::Project(plan) => SqlPlan::ArraySlice(plan), + ArrayTableFn::Agg(plan) => SqlPlan::ArrayAgg(plan), + ArrayTableFn::Elementwise(plan) => SqlPlan::ArrayElementwise(plan), + } + } + + /// The SQL name of the function, as an error message names it. + pub fn sql_name(&self) -> &'static str { + match self { + ArrayTableFn::Slice(_) => "ARRAY_SLICE", + ArrayTableFn::Project(_) => "ARRAY_PROJECT", + ArrayTableFn::Agg(_) => "ARRAY_AGG", + ArrayTableFn::Elementwise(_) => "ARRAY_ELEMENTWISE", + } + } +} /// Try to intercept a `SELECT * FROM array_xxx(...)` table-valued -/// function call. Returns `Ok(Some(plan))` on a match, `Ok(None)` if +/// function call. Returns `Ok(Some(call))` on a match, `Ok(None)` if /// the FROM is not an array function (caller falls through to normal /// catalog resolution). /// /// `temporal` carries any `AS OF SYSTEM TIME` / `AS OF VALID TIME` qualifiers -/// extracted by the pre-processor. It is propagated verbatim into -/// `SqlPlan::ArraySlice` and `SqlPlan::ArrayAgg` so the planner-to-physical -/// conversion can populate the corresponding `ArrayOp` fields. When neither -/// clause was present, `temporal` is `TemporalScope::default()`, which maps to -/// the live-state fast path in the Data Plane handler. +/// extracted by the pre-processor. It is propagated verbatim into the slice +/// and aggregate plans so the planner-to-physical conversion can populate +/// the corresponding `ArrayOp` fields. When neither clause was present, +/// `temporal` is `TemporalScope::default()`, which maps to the live-state +/// fast path in the Data Plane handler. pub fn try_plan_array_table_fn( from: &[ast::TableWithJoins], catalog: &dyn SqlCatalog, temporal: TemporalScope, -) -> Result> { +) -> Result> { if from.len() != 1 { return Ok(None); } @@ -58,10 +94,18 @@ pub fn try_plan_array_table_fn( let fn_name = crate::parser::normalize::normalize_object_name_checked(name)?; let arg_exprs = collect_args(&args.args); match fn_name.as_str() { - "array_slice" => Ok(Some(plan_slice(&arg_exprs, catalog, temporal)?)), - "array_project" => Ok(Some(plan_project(&arg_exprs, catalog)?)), - "array_agg" => Ok(Some(plan_agg(&arg_exprs, catalog, temporal)?)), - "array_elementwise" => Ok(Some(plan_elementwise(&arg_exprs, catalog)?)), + "array_slice" => Ok(Some(ArrayTableFn::Slice(plan_slice( + &arg_exprs, catalog, temporal, + )?))), + "array_project" => Ok(Some(ArrayTableFn::Project(plan_project( + &arg_exprs, catalog, temporal, + )?))), + "array_agg" => Ok(Some(ArrayTableFn::Agg(plan_agg( + &arg_exprs, catalog, temporal, + )?))), + "array_elementwise" => Ok(Some(ArrayTableFn::Elementwise(plan_elementwise( + &arg_exprs, catalog, temporal, + )?))), _ => Ok(None), } } @@ -70,7 +114,7 @@ fn plan_slice( args: &[ast::Expr], catalog: &dyn SqlCatalog, temporal: TemporalScope, -) -> Result { +) -> Result { if args.len() < 2 || args.len() > 4 { return Err(SqlError::Unsupported { detail: format!( @@ -86,40 +130,7 @@ fn plan_slice( detail: format!("ARRAY_SLICE: array '{name}' not found"), })?; - // Slice-predicate literal: encoded as a quoted string carrying the - // brace-form object literal. The PostgreSQL dialect does not accept - // bare `{...}` in expression position, so we decode the string - // contents here. - let slice_str = expect_string_literal(&args[1], "ARRAY_SLICE slice predicate")?; - let parsed = - parse_object_literal_complete(&slice_str).ok_or_else(|| SqlError::Unsupported { - detail: format!("ARRAY_SLICE: slice predicate must be an object literal: {slice_str}"), - })?; - let map = parsed.map_err(|detail| SqlError::Unsupported { - detail: format!("ARRAY_SLICE: slice parse: {detail}"), - })?; - let mut dim_ranges: Vec = Vec::with_capacity(map.len()); - for (dim, val) in map { - // Verify the dim exists on the array. - if !view.dims.iter().any(|d| d.name == dim) { - return Err(SqlError::Unsupported { - detail: format!("ARRAY_SLICE: array '{name}' has no dim '{dim}'"), - }); - } - let arr = match val { - Value::Array(a) if a.len() == 2 => a, - _ => { - return Err(SqlError::Unsupported { - detail: format!( - "ARRAY_SLICE: dim '{dim}' range must be a 2-element array [lo, hi]" - ), - }); - } - }; - let lo = value_to_coord_literal(&arr[0], &dim)?; - let hi = value_to_coord_literal(&arr[1], &dim)?; - dim_ranges.push(NamedDimRange { dim, lo, hi }); - } + let slice = parse_slice_arg(&args[1], "ARRAY_SLICE", &name, &view)?; let attr_projection = if args.len() >= 3 { match &args[2] { @@ -146,16 +157,20 @@ fn plan_slice( 0 }; - Ok(SqlPlan::ArraySlice(ArraySlicePlan { + Ok(ArraySlicePlan { name, - slice: ArraySliceAst { dim_ranges }, + slice, attr_projection, limit, temporal, - })) + }) } -fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { +fn plan_project( + args: &[ast::Expr], + catalog: &dyn SqlCatalog, + temporal: TemporalScope, +) -> Result { if args.len() != 2 { return Err(SqlError::Unsupported { detail: format!( @@ -183,25 +198,34 @@ fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result }); } } - Ok(SqlPlan::ArrayProject(ArrayProjectPlan { + Ok(ArraySlicePlan { name, + slice: ArraySliceAst::default(), attr_projection, - })) + limit: 0, + temporal, + }) } fn plan_agg( args: &[ast::Expr], catalog: &dyn SqlCatalog, temporal: TemporalScope, -) -> Result { - if args.len() < 3 || args.len() > 4 { +) -> Result { + if !(3..=7).contains(&args.len()) { return Err(SqlError::Unsupported { detail: format!( - "ARRAY_AGG expects 3..=4 args (name, attr, reducer, [group_by_dim]); got {}", + "ARRAY_AGG expects 3..=7 args (name, attr, reducer, [group_by_dim], [slice], \ + [zone_wkt, zone_dims]); got {}", args.len() ), }); } + if args.len() == 6 { + return Err(SqlError::Unsupported { + detail: "ARRAY_AGG: zone_wkt needs zone_dims; pass both, or neither".into(), + }); + } let name = require_array_name(args, 0, "ARRAY_AGG", catalog)?; let view = catalog .lookup_array(&name) @@ -210,39 +234,60 @@ fn plan_agg( })?; let attr = expect_string_literal(&args[1], "ARRAY_AGG attr")?; - if !view.attrs.iter().any(|a| a.name == attr) { - return Err(SqlError::Unsupported { + let attr_type = view + .attrs + .iter() + .find(|a| a.name == attr) + .map(|a| a.dtype) + .ok_or_else(|| SqlError::Unsupported { detail: format!("ARRAY_AGG: array '{name}' has no attr '{attr}'"), - }); - } + })?; let reducer_str = expect_string_literal(&args[2], "ARRAY_AGG reducer")?; let reducer = ArrayReducerAst::parse(&reducer_str).ok_or_else(|| SqlError::Unsupported { detail: format!("ARRAY_AGG: unknown reducer '{reducer_str}' (want sum/count/min/max/mean)"), })?; + check_agg_reducer(reducer, &attr, attr_type)?; - let group_by_dim = if args.len() == 4 && !is_null_literal(&args[3]) { - let dim = expect_string_literal(&args[3], "ARRAY_AGG group_by_dim")?; - if !view.dims.iter().any(|d| d.name == dim) { - return Err(SqlError::Unsupported { - detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"), - }); + let group_by_dim = match args.get(3) { + Some(arg) if !is_null_literal(arg) => { + let dim = expect_string_literal(arg, "ARRAY_AGG group_by_dim")?; + if !view.dims.iter().any(|d| d.name == dim) { + return Err(SqlError::Unsupported { + detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"), + }); + } + Some(dim) } - Some(dim) - } else { - None + _ => None, + }; + + let slice = match args.get(4) { + Some(arg) if !is_null_literal(arg) => parse_slice_arg(arg, "ARRAY_AGG", &name, &view)?, + _ => ArraySliceAst::default(), }; - Ok(SqlPlan::ArrayAgg(ArrayAggPlan { + let zone = match (args.get(5), args.get(6)) { + (Some(wkt), Some(dims)) => Some(parse_zone_args(wkt, dims, &name, &view)?), + _ => None, + }; + + Ok(ArrayAggPlan { name, attr, reducer, group_by_dim, + slice, + zone, temporal, - })) + }) } -fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { +fn plan_elementwise( + args: &[ast::Expr], + catalog: &dyn SqlCatalog, + temporal: TemporalScope, +) -> Result { if args.len() != 4 { return Err(SqlError::Unsupported { detail: format!( @@ -280,21 +325,22 @@ fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result SqlPlan { + match plan { + SqlPlan::Cte(crate::types::CtePlan { + mut definitions, .. + }) => { + assert_eq!(definitions.len(), 1, "one array relation definition"); + definitions.remove(0).1 + } + other => panic!("expected Cte over the array relation, got {other:?}"), + } + } + #[test] fn slice_happy() { let p = plan_one( "SELECT * FROM ARRAY_SLICE('g', '{chrom: [1,1], pos: [0, 100]}', ['qual'], 50)", ) .unwrap(); - match p { + match relation_body(p) { SqlPlan::ArraySlice(ArraySlicePlan { name, slice, @@ -411,15 +471,45 @@ mod tests { #[test] fn project_happy() { let p = plan_one("SELECT * FROM ARRAY_PROJECT('g', ['qual', 'variant'])").unwrap(); - match p { - SqlPlan::ArrayProject(ArrayProjectPlan { + match relation_body(p) { + SqlPlan::ArraySlice(ArraySlicePlan { name, + slice, attr_projection, + limit, + .. }) => { assert_eq!(name, "g"); + assert!( + slice.dim_ranges.is_empty(), + "a project reads the whole domain" + ); assert_eq!(attr_projection, vec!["qual".to_string(), "variant".into()]); + assert_eq!(limit, 0); + } + other => panic!("expected a whole-domain ArraySlice, got {other:?}"), + } + } + + /// A project reads the cells visible at its `AS OF` scope, as a slice + /// does. + #[test] + fn project_carries_the_temporal_scope() { + let stmts = parse_sql("SELECT * FROM ARRAY_PROJECT('g', ['qual'])").unwrap(); + let sqlparser::ast::Statement::Query(q) = &stmts[0] else { + panic!("not a query"); + }; + let temporal = crate::TemporalScope { + system_time: nodedb_types::SystemTimeScope::AsOf(1_234), + valid_time: crate::temporal::ValidTime::At(99), + }; + let p = crate::planner::select::plan_query(q, &cat(), &FunctionRegistry::new(), temporal) + .unwrap(); + match relation_body(p) { + SqlPlan::ArraySlice(ArraySlicePlan { temporal: got, .. }) => { + assert_eq!(got, temporal); } - other => panic!("expected ArrayProject, got {other:?}"), + other => panic!("expected a whole-domain ArraySlice, got {other:?}"), } } @@ -469,6 +559,131 @@ mod tests { assert!(plan_one("SELECT * FROM ARRAY_AGG('g', 'qual', 'bogus')").is_err()); } + #[test] + fn agg_carries_its_box_and_zone() { + let p = plan_one( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'mean', NULL, '{chrom: [1, 2]}', \ + 'POLYGON((0 1, 100 1, 100 3, 0 3, 0 1))', ARRAY['pos', 'chrom'])", + ) + .unwrap(); + let SqlPlan::ArrayAgg(ArrayAggPlan { + group_by_dim, + slice, + zone, + .. + }) = p + else { + panic!("expected ArrayAgg, got {p:?}"); + }; + assert!(group_by_dim.is_none()); + assert_eq!(slice.dim_ranges.len(), 1); + assert_eq!(slice.dim_ranges[0].dim, "chrom"); + let zone = zone.expect("zone"); + assert_eq!(zone.x_dim, "pos"); + assert_eq!(zone.y_dim, "chrom"); + assert_eq!(zone.polygons.len(), 1); + assert_eq!(zone.polygons[0][0].len(), 5); + } + + #[test] + fn agg_null_box_reads_the_whole_array() { + let p = plan_one( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', 'chrom', NULL, \ + 'MULTIPOLYGON(((0 0, 9 0, 9 9, 0 9, 0 0)), ((20 20, 29 20, 29 29, 20 20)))', \ + ARRAY['pos', 'chrom'])", + ) + .unwrap(); + let SqlPlan::ArrayAgg(ArrayAggPlan { + group_by_dim, + slice, + zone, + .. + }) = p + else { + panic!("expected ArrayAgg, got {p:?}"); + }; + assert_eq!(group_by_dim.as_deref(), Some("chrom")); + assert!(slice.dim_ranges.is_empty()); + assert_eq!(zone.expect("zone").polygons.len(), 2); + } + + fn agg_error(sql: &str) -> String { + format!("{}", plan_one(sql).expect_err("refused")) + } + + #[test] + fn agg_numeric_reducer_over_string_attr_is_refused() { + for reducer in ["sum", "mean", "min", "max"] { + let sql = format!("SELECT * FROM ARRAY_AGG('g', 'variant', '{reducer}')"); + match plan_one(&sql) { + Err(SqlError::TypeMismatch { detail }) => { + assert!(detail.contains("STRING attr 'variant'"), "{detail}"); + } + other => panic!("{sql} must be refused at plan time: {other:?}"), + } + } + plan_one("SELECT * FROM ARRAY_AGG('g', 'variant', 'count')") + .expect("COUNT over a STRING attr plans"); + } + + #[test] + fn agg_zone_needs_its_dims() { + let err = agg_error( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, \ + 'POLYGON((0 0, 1 0, 1 1, 0 0))')", + ); + assert!(err.contains("zone_dims"), "{err}"); + } + + #[test] + fn agg_box_on_an_unknown_dim_is_refused() { + let err = agg_error("SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, '{nope: [1, 2]}')"); + assert!( + err.contains("ARRAY_AGG: array 'g' has no dim 'nope'"), + "{err}" + ); + } + + #[test] + fn agg_zone_must_be_a_polygon() { + let err = agg_error( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, 'POINT(1 2)', \ + ARRAY['pos', 'chrom'])", + ); + assert!(err.contains("POLYGON or MULTIPOLYGON"), "{err}"); + let err = agg_error( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, 'POLYGON((0 0, 1', \ + ARRAY['pos', 'chrom'])", + ); + assert!(err.contains("not valid WKT"), "{err}"); + } + + #[test] + fn agg_zone_dims_name_two_distinct_dims() { + let zone = "'POLYGON((0 0, 1 0, 1 1, 0 0))'"; + let err = agg_error(&format!( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, {zone}, ARRAY['pos'])" + )); + assert!(err.contains("two dims"), "{err}"); + let err = agg_error(&format!( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, {zone}, ARRAY['pos', 'pos'])" + )); + assert!(err.contains("two different dims"), "{err}"); + let err = agg_error(&format!( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, {zone}, ARRAY['pos', 'lat'])" + )); + assert!(err.contains("no dim 'lat'"), "{err}"); + } + + #[test] + fn agg_more_than_seven_args_is_refused() { + let err = agg_error( + "SELECT * FROM ARRAY_AGG('g', 'qual', 'sum', NULL, NULL, \ + 'POLYGON((0 0, 1 0, 1 1, 0 0))', ARRAY['pos', 'chrom'], 1)", + ); + assert!(err.contains("3..=7 args"), "{err}"); + } + #[test] fn elementwise_happy() { let p = @@ -476,6 +691,28 @@ mod tests { assert!(matches!(p, SqlPlan::ArrayElementwise { .. })); } + /// An elementwise op reads both operands at its `AS OF` scope. + #[test] + fn elementwise_carries_the_temporal_scope() { + let stmts = + parse_sql("SELECT * FROM ARRAY_ELEMENTWISE('left', 'right', 'add', 'qual')").unwrap(); + let sqlparser::ast::Statement::Query(q) = &stmts[0] else { + panic!("not a query"); + }; + let temporal = crate::TemporalScope { + system_time: nodedb_types::SystemTimeScope::AsOf(5_000), + valid_time: crate::temporal::ValidTime::At(7), + }; + let p = crate::planner::select::plan_query(q, &cat(), &FunctionRegistry::new(), temporal) + .unwrap(); + match p { + SqlPlan::ArrayElementwise(ArrayElementwisePlan { temporal: got, .. }) => { + assert_eq!(got, temporal); + } + other => panic!("expected ArrayElementwise, got {other:?}"), + } + } + #[test] fn elementwise_unknown_op_rejected() { assert!( diff --git a/nodedb-sql/src/planner/array_fn/zone.rs b/nodedb-sql/src/planner/array_fn/zone.rs new file mode 100644 index 000000000..311fa1270 --- /dev/null +++ b/nodedb-sql/src/planner/array_fn/zone.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! The polygon zone argument pair of `ARRAY_AGG`: a WKT `POLYGON` or +//! `MULTIPOLYGON` and the two numeric dims it is drawn over, x then y. + +use sqlparser::ast; + +use nodedb_types::geometry::Geometry; + +use super::helpers::{expect_string_array, expect_string_literal}; +use crate::catalog::ArrayCatalogView; +use crate::error::{Result, SqlError}; +use crate::types_array::{ArrayDimType, ArrayZoneAst}; + +/// Parse `zone_wkt` and `zone_dims` of an `ARRAY_AGG` over array `name`. +pub(super) fn parse_zone_args( + wkt_expr: &ast::Expr, + dims_expr: &ast::Expr, + name: &str, + view: &ArrayCatalogView, +) -> Result { + let wkt = expect_string_literal(wkt_expr, "ARRAY_AGG zone")?; + let geometry = + nodedb_spatial::geometry_from_wkt(&wkt).ok_or_else(|| SqlError::InvalidFunction { + detail: format!("ARRAY_AGG: zone is not valid WKT: {wkt}"), + })?; + let issues = nodedb_spatial::validate::validate_geometry(&geometry); + if !issues.is_empty() { + return Err(SqlError::InvalidFunction { + detail: format!("ARRAY_AGG: invalid zone: {}", issues.join("; ")), + }); + } + let polygons = match geometry { + Geometry::Polygon { coordinates } => vec![coordinates], + Geometry::MultiPolygon { coordinates } => coordinates, + other => { + return Err(SqlError::InvalidFunction { + detail: format!( + "ARRAY_AGG: zone must be a POLYGON or MULTIPOLYGON, got {}", + other.geometry_type() + ), + }); + } + }; + + let dims = expect_string_array(dims_expr, "ARRAY_AGG zone dims")?; + let [x_dim, y_dim] = <[String; 2]>::try_from(dims).map_err(|dims| SqlError::Unsupported { + detail: format!( + "ARRAY_AGG: zone dims must name two dims, x then y; got {}", + dims.len() + ), + })?; + if x_dim == y_dim { + return Err(SqlError::Unsupported { + detail: format!("ARRAY_AGG: zone dims must be two different dims; got '{x_dim}' twice"), + }); + } + for dim in [&x_dim, &y_dim] { + let spec = + view.dims + .iter() + .find(|d| &d.name == dim) + .ok_or_else(|| SqlError::Unsupported { + detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"), + })?; + match spec.dtype { + ArrayDimType::Int64 | ArrayDimType::Float64 => {} + ArrayDimType::TimestampMs | ArrayDimType::String => { + return Err(SqlError::TypeMismatch { + detail: format!( + "ARRAY_AGG: zone dim '{dim}' is {:?}; a zone needs an INT64 or FLOAT64 dim", + spec.dtype + ), + }); + } + } + } + Ok(ArrayZoneAst { + polygons, + x_dim, + y_dim, + }) +} diff --git a/nodedb-sql/src/planner/catalog_fold/leaf.rs b/nodedb-sql/src/planner/catalog_fold/leaf.rs index 743657f1e..5d5fbc8b5 100644 --- a/nodedb-sql/src/planner/catalog_fold/leaf.rs +++ b/nodedb-sql/src/planner/catalog_fold/leaf.rs @@ -149,7 +149,6 @@ pub(super) fn fold_leaf( | SqlPlan::InsertArray(_) | SqlPlan::DeleteArray(_) | SqlPlan::ArraySlice(_) - | SqlPlan::ArrayProject(_) | SqlPlan::ArrayAgg(_) | SqlPlan::ArrayElementwise(_) | SqlPlan::ArrayFlush { .. } diff --git a/nodedb-sql/src/planner/catalog_plan_validate.rs b/nodedb-sql/src/planner/catalog_plan_validate.rs index 3667ae690..8f456c002 100644 --- a/nodedb-sql/src/planner/catalog_plan_validate.rs +++ b/nodedb-sql/src/planner/catalog_plan_validate.rs @@ -252,7 +252,6 @@ pub(super) fn validate_catalog_exprs( | SqlPlan::InsertArray(_) | SqlPlan::DeleteArray(_) | SqlPlan::ArraySlice(_) - | SqlPlan::ArrayProject(_) | SqlPlan::ArrayAgg(_) | SqlPlan::ArrayElementwise(_) | SqlPlan::ArrayFlush { .. } diff --git a/nodedb-sql/src/planner/computed_tail.rs b/nodedb-sql/src/planner/computed_tail.rs new file mode 100644 index 000000000..a838e7f42 --- /dev/null +++ b/nodedb-sql/src/planner/computed_tail.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Computed SELECT-list items over an access path that evaluates none. +//! +//! A spatial predicate scan and a secondary-index fetch return stored rows +//! and run no expression. A computed item over one of them is evaluated by a +//! [`SqlPlan::Subquery`] tail instead. The body emits every stored column the +//! items read. The tail evaluates each item over those rows, then projects +//! the SELECT list. An `ORDER BY` on a computed alias sorts the tail's rows, +//! which carry the alias. + +use crate::types::plan::referenced_columns; +use crate::types::{Projection, SqlPlan}; + +/// The row bound and dedup a tail applies over its body's rows. +#[derive(Debug, Clone, Copy, Default)] +pub struct TailBounds { + pub distinct: bool, + pub limit: Option, + pub offset: usize, +} + +/// Whether `projection` holds an item the body cannot evaluate. +/// +/// A renamed column (`name AS n`) counts: the body emits stored names only. +pub fn needs_computed_tail(projection: &[Projection]) -> bool { + projection + .iter() + .any(|p| matches!(p, Projection::Computed { .. })) +} + +/// The projection the body emits for `projection`. +/// +/// Without a computed item it is `projection` itself. With one, it is every +/// plain column the list names plus every column an item reads, each once. +/// A star keeps every column, so the result is empty. +pub fn body_projection(projection: &[Projection]) -> Vec { + if !needs_computed_tail(projection) { + return projection.to_vec(); + } + if projection + .iter() + .any(|p| matches!(p, Projection::Star | Projection::QualifiedStar(_))) + { + return Vec::new(); + } + let mut columns: Vec = Vec::with_capacity(projection.len()); + let mut keep = |name: String| { + if !columns.contains(&name) { + columns.push(name); + } + }; + for item in projection { + match item { + Projection::Column(name) => keep(name.clone()), + Projection::Computed { expr, .. } | Projection::CpComputed { expr, .. } => { + referenced_columns(expr).into_iter().for_each(&mut keep); + } + Projection::Star | Projection::QualifiedStar(_) => {} + } + } + columns.into_iter().map(Projection::Column).collect() +} + +/// Wrap `body` in a tail that evaluates `projection` over its rows. +/// +/// `body` must already carry [`body_projection`] of `projection`. Without a +/// computed item `body` returns unchanged, and `bounds` stays on the body. +pub fn with_computed_tail(body: SqlPlan, projection: &[Projection], bounds: TailBounds) -> SqlPlan { + if !needs_computed_tail(projection) { + return body; + } + SqlPlan::Subquery { + input: Box::new(body), + filters: Vec::new(), + projection: projection.to_vec(), + window_functions: Vec::new(), + sort_keys: Vec::new(), + offset: bounds.offset, + distinct: bounds.distinct, + limit: bounds.limit, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::SqlExpr; + + fn column(name: &str) -> SqlExpr { + SqlExpr::Column { + table: None, + name: name.into(), + } + } + + fn computed(alias: &str, expr: SqlExpr) -> Projection { + Projection::Computed { + expr, + alias: alias.into(), + } + } + + #[test] + fn a_plain_projection_needs_no_tail() { + let projection = vec![Projection::Column("name".into())]; + assert!(!needs_computed_tail(&projection)); + assert_eq!(body_projection(&projection).len(), 1); + } + + #[test] + fn the_body_emits_every_column_a_computed_item_reads() { + let projection = vec![ + Projection::Column("name".into()), + computed( + "km", + SqlExpr::Function { + name: "st_x".into(), + args: vec![column("location")], + distinct: false, + }, + ), + computed("n", column("name")), + ]; + let names: Vec = body_projection(&projection) + .into_iter() + .map(|p| match p { + Projection::Column(name) => name, + other => panic!("the body projects plain columns only, got {other:?}"), + }) + .collect(); + assert_eq!(names, ["name", "location"]); + } + + #[test] + fn a_star_beside_a_computed_item_keeps_every_column() { + let projection = vec![Projection::Star, computed("n", column("name"))]; + assert!(body_projection(&projection).is_empty()); + } +} diff --git a/nodedb-sql/src/planner/const_fold.rs b/nodedb-sql/src/planner/const_fold.rs index f8f526a8f..511037175 100644 --- a/nodedb-sql/src/planner/const_fold.rs +++ b/nodedb-sql/src/planner/const_fold.rs @@ -401,7 +401,10 @@ pub fn fold_function_call_scoped( e @ (nodedb_query::EvalError::VectorDimensionMismatch { .. } | nodedb_query::EvalError::ArgumentType { .. } | nodedb_query::EvalError::InvalidJsonPath { .. } - | nodedb_query::EvalError::NumericOverflow { .. }), + | nodedb_query::EvalError::NumericOverflow { .. } + | nodedb_query::EvalError::UnsupportedSrid { .. } + | nodedb_query::EvalError::MixedSrid { .. } + | nodedb_query::EvalError::ProjectionDomain { .. }), ) => Err(SqlError::DataException { detail: e.to_string(), }), @@ -445,7 +448,8 @@ fn ndb_to_sql_value(v: Value) -> SqlValue { // Geometry and Object values are serialized to JSON strings so that // nested function calls like ST_Distance(ST_Point(...), ST_Point(...)) // survive the SqlValue round-trip. The geo evaluator's geom_arg helper - // recovers Geometry from a GeoJSON string; Object results (e.g. from + // recovers Geometry from a GeoJSON string, SRID included through the + // `crs` member; Object results (e.g. from // ST_GeoHashDecode) reach the client as a JSON-encoded string column. Value::Geometry(g) => sonic_rs::to_string(&g) .map(SqlValue::String) diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs b/nodedb-sql/src/planner/cte/inline.rs similarity index 92% rename from nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs rename to nodedb-sql/src/planner/cte/inline.rs index 4cc6ee722..3fd6f8784 100644 --- a/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs +++ b/nodedb-sql/src/planner/cte/inline.rs @@ -1,6 +1,10 @@ -// SPDX-License-Identifier: BUSL-1.1 +// SPDX-License-Identifier: Apache-2.0 -use nodedb_sql::types::{CtePlan, Filter, Projection, SortKey, SqlPlan, WindowSpec}; +//! CTE inlining: a plan that reads a CTE by name, rewritten to read the +//! CTE's body. Pure `SqlPlan` → `SqlPlan`, shared by every executor that +//! lowers a `SqlPlan::Cte`. + +use crate::types::{CtePlan, Filter, Projection, SortKey, SqlPlan, WindowSpec}; /// Replace scans on `cte_name` with the CTE's actual subquery plan. /// @@ -18,11 +22,7 @@ use nodedb_sql::types::{CtePlan, Filter, Projection, SortKey, SqlPlan, WindowSpe /// body's materialized rows. Computed projection entries and window functions /// always take that path: the body does not evaluate them, and a name lookup /// at the response boundary yields NULL. -pub(in crate::control::planner::sql_plan_convert) fn inline_cte( - plan: &SqlPlan, - cte_name: &str, - cte_plan: &SqlPlan, -) -> SqlPlan { +pub fn inline_cte(plan: &SqlPlan, cte_name: &str, cte_plan: &SqlPlan) -> SqlPlan { match plan { // Direct scan on CTE name → replace with CTE plan. SqlPlan::Scan { @@ -266,6 +266,38 @@ impl ScanRef<'_> { } } +/// Whether [`inline_cte`] replaces `plan` with the body of `cte_name` +/// unchanged. +/// +/// That is an unconstrained `Scan` of `cte_name`. Its columns then read the +/// body's own rows, so each one resolves through the body's output schema. +pub fn reads_cte_body_unchanged(plan: &SqlPlan, cte_name: &str) -> bool { + match plan { + SqlPlan::Scan { + collection, + filters, + projection, + sort_keys, + limit, + offset, + distinct, + window_functions, + .. + } if collection == cte_name => ScanRef { + filters, + projection, + sort_keys, + limit: *limit, + offset: *offset, + distinct: *distinct, + window_functions, + has_computed: has_computed_projection(projection), + } + .is_unconstrained(), + _ => false, + } +} + /// Wrap `input` in a `Subquery` post-processor carrying the outer /// constraints that `input` has no slot for. fn wrap_in_subquery(input: SqlPlan, filters: Vec, outer: ScanRef<'_>) -> SqlPlan { @@ -416,7 +448,7 @@ fn merge_into_vector_search_body(outer: ScanRef<'_>, cte_plan: &SqlPlan) -> Opti #[cfg(test)] mod tests { use super::*; - use nodedb_sql::types::{CompareOp, EngineType, Filter, FilterExpr, SortKey, SqlValue}; + use crate::types::{CompareOp, EngineType, Filter, FilterExpr, SortKey, SqlValue}; fn vector_search_body() -> SqlPlan { SqlPlan::VectorSearch { @@ -425,10 +457,10 @@ mod tests { query_vector: vec![0.1, 0.2], top_k: 3, ef_search: 64, - metric: nodedb_sql::types::DistanceMetric::L2, + metric: crate::types::DistanceMetric::L2, filters: Vec::new(), array_prefilter: None, - ann_options: nodedb_sql::types::VectorAnnOptions::default(), + ann_options: crate::types::VectorAnnOptions::default(), skip_payload_fetch: false, payload_filters: Vec::new(), pk_prefilter: None, @@ -448,7 +480,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -528,13 +560,13 @@ mod tests { offset, distinct, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } fn id_sort_key() -> SortKey { SortKey { - expr: nodedb_sql::types::SqlExpr::Column { + expr: crate::types::SqlExpr::Column { table: Some("s".to_string()), name: "id".to_string(), }, @@ -603,20 +635,20 @@ mod tests { fn doubled_x() -> Projection { Projection::Computed { - expr: nodedb_sql::types::SqlExpr::BinaryOp { - left: Box::new(nodedb_sql::types::SqlExpr::Column { + expr: crate::types::SqlExpr::BinaryOp { + left: Box::new(crate::types::SqlExpr::Column { table: None, name: "x".to_string(), }), - op: nodedb_sql::types::BinaryOp::Mul, - right: Box::new(nodedb_sql::types::SqlExpr::Literal(SqlValue::Int(2))), + op: crate::types::BinaryOp::Mul, + right: Box::new(crate::types::SqlExpr::Literal(SqlValue::Int(2))), }, alias: "y".to_string(), } } - fn row_number_spec() -> nodedb_sql::types::WindowSpec { - nodedb_sql::types::WindowSpec { + fn row_number_spec() -> crate::types::WindowSpec { + crate::types::WindowSpec { function: "row_number".to_string(), args: Vec::new(), partition_by: Vec::new(), @@ -629,7 +661,7 @@ mod tests { /// A CTE-referencing scan carrying an outer projection and window list. fn scan_on_cte_projected( projection: Vec, - window_functions: Vec, + window_functions: Vec, ) -> SqlPlan { SqlPlan::Scan { collection: "knn".to_string(), @@ -642,7 +674,7 @@ mod tests { offset: 0, distinct: false, window_functions, - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -735,7 +767,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -769,7 +801,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), }; match inline_cte(&scan_on_cte(vec![tag_filter()], Some(2)), "knn", &body) { SqlPlan::Scan { @@ -798,7 +830,7 @@ mod tests { engine: EngineType::DocumentSchemaless, filters: Vec::new(), projection: vec![Projection::Computed { - expr: nodedb_sql::types::SqlExpr::Column { + expr: crate::types::SqlExpr::Column { table: None, name: "qty".to_string(), }, @@ -809,7 +841,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), }; match inline_cte( &scan_on_cte_projected(vec![doubled_x()], Vec::new()), @@ -842,7 +874,7 @@ mod tests { left: Box::new(other_side), right: Box::new(scan_on_cte(Vec::new(), None)), on: vec![("orders.id".to_string(), "knn.id".to_string())], - join_type: nodedb_sql::types::query::JoinType::Inner, + join_type: crate::types::query::JoinType::Inner, condition: None, limit: None, projection: Vec::new(), diff --git a/nodedb-sql/src/planner/cte/mod.rs b/nodedb-sql/src/planner/cte/mod.rs index 37abd05e5..ff98beca6 100644 --- a/nodedb-sql/src/planner/cte/mod.rs +++ b/nodedb-sql/src/planner/cte/mod.rs @@ -2,9 +2,11 @@ //! CTE (WITH clause) and WITH RECURSIVE planning. +pub mod inline; mod join_link; mod recursive_scan; mod recursive_value; mod validate; +pub use inline::{inline_cte, reads_cte_body_unchanged}; pub use recursive_scan::{DEFAULT_MAX_RECURSION_DEPTH, plan_recursive_cte}; diff --git a/nodedb-sql/src/planner/dml/bulk_load.rs b/nodedb-sql/src/planner/dml/bulk_load.rs new file mode 100644 index 000000000..cc162cbdd --- /dev/null +++ b/nodedb-sql/src/planner/dml/bulk_load.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Bulk-load (`COPY ... FROM`) planning: the target resolves to its engine +//! rules, and the rules answer how the engine loads rows. + +use nodedb_types::DatabaseId; + +use crate::catalog::SqlCatalog; +use crate::engine_rules::{self, BulkLoadParams, BulkLoadPlan}; +use crate::error::Result; +use crate::types::EngineType; + +/// Plan one chunk of a bulk load into `params.collection`, or say how its +/// engine loads rows when `params.cells` is empty. +/// +/// An array lives in its own catalog namespace, so it resolves before the +/// collection lookup can miss it, as `TRUNCATE` resolves it. An unknown +/// target loads row by row: each row's `INSERT` reports the unknown table. +pub fn plan_bulk_load(params: BulkLoadParams, catalog: &dyn SqlCatalog) -> Result { + if catalog.array_exists(¶ms.collection) { + return engine_rules::resolve_engine_rules(EngineType::Array).plan_bulk_load(params); + } + match catalog.get_collection(DatabaseId::DEFAULT, ¶ms.collection)? { + Some(info) => engine_rules::resolve_engine_rules(info.engine).plan_bulk_load(params), + None => Ok(BulkLoadPlan::RowInserts), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::{ArrayCatalogView, SqlCatalogError}; + use crate::error::SqlError; + use crate::types::CollectionInfo; + + struct ArrayOnly; + + impl SqlCatalog for ArrayOnly { + fn get_collection( + &self, + _: DatabaseId, + _: &str, + ) -> std::result::Result, SqlCatalogError> { + Ok(None) + } + fn lookup_array(&self, name: &str) -> Option { + (name == "scene").then(|| ArrayCatalogView { + name: name.into(), + dims: Vec::new(), + attrs: Vec::new(), + tile_extents: Vec::new(), + }) + } + } + + fn probe(name: &str) -> Result { + plan_bulk_load( + BulkLoadParams { + collection: name.into(), + cells: Vec::new(), + }, + &ArrayOnly, + ) + } + + #[test] + fn an_array_loads_in_cell_chunks() { + assert!(matches!(probe("scene"), Ok(BulkLoadPlan::CellChunks(_)))); + } + + #[test] + fn an_unknown_target_loads_row_by_row() { + assert!(matches!(probe("missing"), Ok(BulkLoadPlan::RowInserts))); + } + + #[test] + fn timeseries_and_spatial_refuse_a_bulk_load() { + for engine in [EngineType::Timeseries, EngineType::Spatial] { + let refused = + engine_rules::resolve_engine_rules(engine).plan_bulk_load(BulkLoadParams { + collection: "c".into(), + cells: Vec::new(), + }); + assert!(matches!(refused, Err(SqlError::Unsupported { .. }))); + } + } + + #[test] + fn row_engines_load_row_by_row() { + for engine in [ + EngineType::DocumentSchemaless, + EngineType::DocumentStrict, + EngineType::KeyValue, + EngineType::Columnar, + EngineType::ContinuousAggregate, + ] { + let plan = engine_rules::resolve_engine_rules(engine).plan_bulk_load(BulkLoadParams { + collection: "c".into(), + cells: Vec::new(), + }); + assert!(matches!(plan, Ok(BulkLoadPlan::RowInserts))); + } + } +} diff --git a/nodedb-sql/src/planner/dml/mod.rs b/nodedb-sql/src/planner/dml/mod.rs index 66450928e..5d3e2804a 100644 --- a/nodedb-sql/src/planner/dml/mod.rs +++ b/nodedb-sql/src/planner/dml/mod.rs @@ -2,12 +2,14 @@ //! INSERT, UPSERT, UPDATE, DELETE, and TRUNCATE planning. +mod bulk_load; mod insert; mod target; mod update_delete; mod update_from; mod upsert; +pub use bulk_load::plan_bulk_load; pub use insert::plan_insert; pub use update_delete::{plan_delete, plan_truncate_stmt, plan_update}; pub use upsert::plan_upsert; diff --git a/nodedb-sql/src/planner/geometry_expr/resolve.rs b/nodedb-sql/src/planner/geometry_expr/resolve.rs index fc77861f4..62f75387a 100644 --- a/nodedb-sql/src/planner/geometry_expr/resolve.rs +++ b/nodedb-sql/src/planner/geometry_expr/resolve.rs @@ -12,7 +12,7 @@ use sqlparser::ast; use nodedb_query::geo_functions; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use crate::error::{Result, SqlError}; use crate::parser::normalize::normalize_ident; @@ -24,7 +24,9 @@ use crate::types::*; /// the caller falls through to the shared constant-folding pipeline. /// /// The stored form is a GeoJSON string, which is what the spatial read path -/// parses back. A call whose arguments do not describe a valid geometry is an +/// parses back. A geometry in any SRID other than 4326 carries the `crs` +/// member, so the SRID survives storage. A call whose arguments do not +/// describe a valid geometry is an /// error rather than a NULL: silently storing NULL for a malformed /// `ST_GeomFromText('POINT(')` would lose the row's geometry with no signal. pub(crate) fn fold_geometry_function(func: &ast::Function) -> Option> { @@ -47,7 +49,7 @@ pub(crate) fn fold_geometry_function(func: &ast::Function) -> Option Result { +pub(crate) fn resolve_geometry_expr(expr: &ast::Expr) -> Result { match resolve(expr)? { Some(geom) => Ok(geom), None => Err(SqlError::InvalidFunction { @@ -62,7 +64,7 @@ pub(crate) fn resolve_geometry_expr(expr: &ast::Expr) -> Result { /// /// `Ok(None)` means the expression folded but is not a geometry, or could not /// be folded at plan time at all. -fn resolve(expr: &ast::Expr) -> Result> { +fn resolve(expr: &ast::Expr) -> Result> { let sql_expr = crate::resolver::expr::convert_expr(expr, &crate::resolver::ColumnScope::Unchecked)?; let Some(value) = crate::planner::const_fold::fold_constant_default(&sql_expr)? else { @@ -74,14 +76,15 @@ fn resolve(expr: &ast::Expr) -> Result> { /// Geometry travels through `SqlValue` as its GeoJSON string form, which is /// also how it is stored; WKT literals are accepted here for the same reason /// they are accepted by the evaluator. -fn geometry_from_value(value: &SqlValue) -> Option { +fn geometry_from_value(value: &SqlValue) -> Option { match value { SqlValue::String(text) => geo_functions::geometry_from_text(text), _ => None, } } -fn serialize(geom: &Geometry, name: &str) -> Result { +/// GeoJSON text of `geom`, with the `crs` member for a non-4326 SRID. +fn serialize(geom: &SridGeometry, name: &str) -> Result { sonic_rs::to_string(geom) .map(SqlValue::String) .map_err(|e| SqlError::InvalidFunction { @@ -111,7 +114,7 @@ fn function_name(func: &ast::Function) -> String { #[cfg(test)] mod tests { - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, Srid, SridGeometry}; use super::{fold_geometry_function, resolve_geometry_expr}; use crate::types::SqlValue; @@ -137,6 +140,7 @@ mod tests { fn resolved(sql: &str) -> Geometry { resolve_geometry_expr(&expr(sql)) + .map(|g| g.geometry) .unwrap_or_else(|e| panic!("`{sql}` must resolve to a geometry: {e}")) } @@ -219,9 +223,41 @@ mod tests { }; assert_eq!( nodedb_types::geometry::from_geojson_str(&geojson), - Some(Geometry::point(1.0, 2.0)), + Some(SridGeometry::wgs84(Geometry::point(1.0, 2.0))), "stored form must parse back through the spatial read path" ); + assert!( + !geojson.contains("crs"), + "4326 writes no crs member: {geojson}" + ); + } + + /// A projected geometry folds with its `crs` member, so the stored form + /// keeps the SRID. + #[test] + fn projected_geometry_folds_with_crs_member() { + let sqlparser::ast::Expr::Function(func) = expr("ST_GeomFromText('POINT(1 2)', 3375)") + else { + panic!("expected a function"); + }; + let Some(Ok(SqlValue::String(geojson))) = fold_geometry_function(&func) else { + panic!("a geometry constructor must fold to its stored form"); + }; + assert!(geojson.contains("EPSG:3375"), "{geojson}"); + assert_eq!( + nodedb_types::geometry::from_geojson_str(&geojson).map(|g| g.srid), + Some(Srid::PeninsularRso) + ); + let tagged = resolve_geometry_expr(&expr("ST_SetSRID(ST_Point(1, 2), 32648)")) + .expect("must resolve"); + assert_eq!(tagged.srid, Srid::Utm48N); + } + + /// A mixed-SRID operation in geometry position is an error, not a fold. + #[test] + fn mixed_srid_operation_is_an_error() { + let sql = "ST_Union(ST_Point(1, 2), ST_SetSRID(ST_Point(1, 2), 32648))"; + assert!(resolve_geometry_expr(&expr(sql)).is_err()); } /// A non-geometry call falls through so the generic constant folder handles diff --git a/nodedb-sql/src/planner/group_by.rs b/nodedb-sql/src/planner/group_by.rs index 0be6f071f..f19be5767 100644 --- a/nodedb-sql/src/planner/group_by.rs +++ b/nodedb-sql/src/planner/group_by.rs @@ -149,23 +149,56 @@ fn resolve_output_alias<'a>( } /// Derive the SELECT-list output alias for each GROUP BY key by correlating -/// each key with the projection item that references the same column. Returns -/// a `Vec` parallel to `group_by`: `Some(alias)` when a projection item -/// explicitly aliases that grouped column (`SELECT k AS label ... GROUP BY k`), -/// otherwise `None` (the output column name falls back to the raw grouped -/// column name). +/// each key with the projection item that produces it. Returns a `Vec` +/// parallel to `group_by`: `Some(alias)` when a projection item explicitly +/// aliases that key, otherwise `None` (the output column name falls back to +/// the key's own name). +/// +/// A column key matches an aliased item that references the same column +/// (`SELECT k AS label ... GROUP BY k`). A computed key matches an aliased +/// item whose expression converts to the same key +/// (`SELECT time_bucket('1 hour', ts) AS b ... GROUP BY b`). pub fn group_by_output_aliases( projection: &[ast::SelectItem], group_by: &[SqlExpr], + scope: &TableScope, ) -> Vec> { group_by .iter() - .map(|key| { - key_column_name(key).and_then(|name| projection_alias_for_column(projection, name)) + .map(|key| match key_column_name(key) { + Some(name) => projection_alias_for_column(projection, name), + None => projection_alias_for_expr(projection, key, scope), }) .collect() } +/// The explicit alias of the projection item whose expression converts to +/// the computed GROUP BY key `key`. +/// +/// `SqlExpr` has no `PartialEq` (it can carry a subquery plan), so the match +/// compares the `Debug` rendering. Both sides pass through `convert_expr`, so +/// equal expressions render identically. An item that does not convert names +/// no key. +fn projection_alias_for_expr( + projection: &[ast::SelectItem], + key: &SqlExpr, + scope: &TableScope, +) -> Option { + let rendered_key = format!("{key:?}"); + projection.iter().find_map(|item| match item { + ast::SelectItem::ExprWithAlias { expr, alias } => { + convert_expr(expr, &ColumnScope::Relations(scope)) + .ok() + .filter(|converted| format!("{converted:?}") == rendered_key) + .map(|_| normalize_ident(alias)) + } + ast::SelectItem::UnnamedExpr(_) + | ast::SelectItem::ExprWithAliases { .. } + | ast::SelectItem::QualifiedWildcard(..) + | ast::SelectItem::Wildcard(_) => None, + }) +} + /// The bare column name of a GROUP BY key, when the key is a column reference. pub fn key_column_name(key: &SqlExpr) -> Option<&str> { match key { diff --git a/nodedb-sql/src/planner/join/array_arm.rs b/nodedb-sql/src/planner/join/array_arm.rs index bbf836c0d..c20648c76 100644 --- a/nodedb-sql/src/planner/join/array_arm.rs +++ b/nodedb-sql/src/planner/join/array_arm.rs @@ -6,9 +6,8 @@ //! sibling read TVFs), the planner lowers it to the corresponding //! `SqlPlan::Array*` variant directly — bypassing catalog lookup. //! The TVF's output rows participate in the join key resolution like -//! any other relation; the dispatch handler emits `(coords, attrs)` -//! tuples that include the array's surrogate-equivalent identity -//! columns so equi-join keys work naturally. +//! any other relation. A slice or project row holds one field per dim +//! and per attr, keyed by name, so equi-join keys read them directly. use sqlparser::ast; @@ -46,7 +45,12 @@ pub(super) fn try_plan_relation( relation: rel.clone(), joins: Vec::new(), }; - super::super::array_fn::try_plan_array_table_fn(std::slice::from_ref(&twj), catalog, temporal) + Ok(super::super::array_fn::try_plan_array_table_fn( + std::slice::from_ref(&twj), + catalog, + temporal, + )? + .map(super::super::array_fn::ArrayTableFn::into_plan)) } fn is_array_tvf(name: &str) -> bool { diff --git a/nodedb-sql/src/planner/mod.rs b/nodedb-sql/src/planner/mod.rs index 81e90c463..df8c92932 100644 --- a/nodedb-sql/src/planner/mod.rs +++ b/nodedb-sql/src/planner/mod.rs @@ -16,6 +16,7 @@ pub mod catalog_fold; pub mod catalog_plan_shapes; pub mod catalog_plan_validate; pub mod collection_predicate; +pub mod computed_tail; pub mod const_fold; pub mod cp_projection; pub mod cte; @@ -35,6 +36,7 @@ pub mod merge; pub mod predicate_coerce; pub mod returning; pub mod search_scope; +pub mod search_tail; pub mod select; pub use returning::resolve_returning_items; diff --git a/nodedb-sql/src/planner/search_scope/lookup.rs b/nodedb-sql/src/planner/search_scope/lookup.rs index e57e91261..f69929cb5 100644 --- a/nodedb-sql/src/planner/search_scope/lookup.rs +++ b/nodedb-sql/src/planner/search_scope/lookup.rs @@ -124,7 +124,6 @@ pub(super) fn has_search_plan(plan: &SqlPlan) -> bool { | SqlPlan::InsertArray { .. } | SqlPlan::DeleteArray { .. } | SqlPlan::ArraySlice { .. } - | SqlPlan::ArrayProject { .. } | SqlPlan::ArrayAgg { .. } | SqlPlan::ArrayElementwise { .. } | SqlPlan::ArrayFlush { .. } diff --git a/nodedb-sql/src/planner/search_scope/scope.rs b/nodedb-sql/src/planner/search_scope/scope.rs index 892449dff..a3ba33bdc 100644 --- a/nodedb-sql/src/planner/search_scope/scope.rs +++ b/nodedb-sql/src/planner/search_scope/scope.rs @@ -150,6 +150,7 @@ impl Scope<'_> { // A search plan serves its own score call as a column; only its // residual filters run on the row evaluator. SqlPlan::VectorSearch { filters, .. } + | SqlPlan::SparseSearch { filters, .. } | SqlPlan::TextSearch(TextSearchPlan { filters, .. }) | SqlPlan::HybridSearch(HybridSearchPlan { filters, .. }) | SqlPlan::HybridSearchTriple(HybridSearchTriplePlan { filters, .. }) => { @@ -244,7 +245,6 @@ impl Scope<'_> { | SqlPlan::Truncate { .. } | SqlPlan::TimeseriesIngest { .. } | SqlPlan::MultiVectorSearch { .. } - | SqlPlan::SparseSearch { .. } | SqlPlan::RecursiveValue { .. } | SqlPlan::CreateArray { .. } | SqlPlan::DropArray { .. } @@ -252,7 +252,6 @@ impl Scope<'_> { | SqlPlan::InsertArray { .. } | SqlPlan::DeleteArray { .. } | SqlPlan::ArraySlice { .. } - | SqlPlan::ArrayProject { .. } | SqlPlan::ArrayAgg { .. } | SqlPlan::ArrayElementwise { .. } | SqlPlan::ArrayFlush { .. } diff --git a/nodedb-sql/src/planner/search_tail.rs b/nodedb-sql/src/planner/search_tail.rs new file mode 100644 index 000000000..434f3127e --- /dev/null +++ b/nodedb-sql/src/planner/search_tail.rs @@ -0,0 +1,311 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Computed SELECT-list items over a search plan. +//! +//! A text, vector, sparse, multi-vector, or hybrid search returns hit rows +//! and evaluates no expression. Its own score reaches each row under a fixed +//! key: a text score under its alias, an RRF score under the hybrid score +//! alias, a vector or sparse score under `distance`. Every other computed +//! item runs in a [`SqlPlan::Subquery`] tail over the hit rows, and a score +//! item is read from its key there. +//! +//! The pass runs once ORDER BY and LIMIT are applied, so the search keeps +//! its own top-k bound and a sort on a computed alias reads the tail's rows. + +use crate::functions::registry::{FunctionRegistry, SearchTrigger}; +use crate::planner::computed_tail::{TailBounds, body_projection, with_computed_tail}; +use crate::types::{Projection, SqlExpr, SqlPlan}; + +/// The row key a vector-family search writes its score under. +const DISTANCE_KEY: &str = "distance"; + +/// The row key a hybrid search writes its RRF score under without an alias. +const DEFAULT_RRF_KEY: &str = "rrf_score"; + +/// Lift the computed items of the search plan at the bottom of `plan` into +/// a tail. A plan that is not a search, or holds no computed item the +/// search cannot answer, returns unchanged. +pub fn lift_search_computed(plan: SqlPlan, functions: &FunctionRegistry) -> SqlPlan { + match plan { + SqlPlan::Subquery { + input, + filters, + projection, + window_functions, + sort_keys, + offset, + distinct, + limit, + } => SqlPlan::Subquery { + input: Box::new(lift_search_computed(*input, functions)), + filters, + projection, + window_functions, + sort_keys, + offset, + distinct, + limit, + }, + leaf => lift_leaf(leaf, functions), + } +} + +/// How a search plan delivers its own score. +enum ScoreSource { + /// Text scores, each under its alias. + Text(Vec), + /// The RRF score under one key. + Hybrid(String), + /// The vector or sparse score under `distance`. `field` is the column + /// the search ranks by, when the plan names one. + Vector(Option), +} + +impl ScoreSource { + fn of(plan: &SqlPlan) -> Option { + match plan { + SqlPlan::TextSearch(search) => Some(Self::Text( + search.scores.iter().map(|s| s.alias.clone()).collect(), + )), + SqlPlan::HybridSearch(search) => Some(Self::hybrid(search.score_alias.as_deref())), + SqlPlan::HybridSearchTriple(search) => { + Some(Self::hybrid(search.score_alias.as_deref())) + } + SqlPlan::VectorSearch { field, .. } | SqlPlan::SparseSearch { field, .. } => { + Some(Self::Vector(Some(field.clone()))) + } + SqlPlan::MultiVectorSearch { .. } => Some(Self::Vector(None)), + _ => None, + } + } + + fn hybrid(alias: Option<&str>) -> Self { + Self::Hybrid(alias.unwrap_or(DEFAULT_RRF_KEY).to_string()) + } + + /// The tail item that reads `item` as the search's own score, if it is + /// one. + fn score_item(&self, item: &Projection, functions: &FunctionRegistry) -> Option { + let Projection::Computed { expr, alias } = item else { + return None; + }; + match self { + Self::Text(aliases) => aliases + .contains(alias) + .then(|| Projection::Column(alias.clone())), + Self::Hybrid(key) => (trigger(expr, functions) == Some(SearchTrigger::HybridSearch)) + .then(|| read_key(key, alias)), + Self::Vector(field) => { + let SqlExpr::Function { args, .. } = expr else { + return None; + }; + let ranks = matches!( + trigger(expr, functions), + Some( + SearchTrigger::VectorSearch + | SearchTrigger::MultiVectorSearch + | SearchTrigger::SparseSearch + ) + ); + // A one-argument call (`vector_distance(ARRAY[...])`) ranks + // by the collection's own vector column. + let same_column = match (field, args.as_slice()) { + (None, _) | (Some(_), [_]) => true, + (Some(field), [SqlExpr::Column { name, .. }, ..]) => name == field, + (Some(_), _) => false, + }; + (ranks && same_column).then(|| read_key(DISTANCE_KEY, alias)) + } + } + } + + /// Whether the search already answers a score item without a tail. + fn answers_scores(&self) -> bool { + match self { + Self::Text(_) | Self::Hybrid(_) => true, + Self::Vector(_) => false, + } + } +} + +fn trigger(expr: &SqlExpr, functions: &FunctionRegistry) -> Option { + match expr { + SqlExpr::Function { name, .. } => Some(functions.search_trigger(name)), + _ => None, + } +} + +/// A tail item that emits row key `key` under `alias`. +fn read_key(key: &str, alias: &str) -> Projection { + Projection::Computed { + expr: SqlExpr::Column { + table: None, + name: key.to_string(), + }, + alias: alias.to_string(), + } +} + +/// Whether `read` emits a row key under that same name. +fn reads_own_key(read: &Projection) -> bool { + match read { + Projection::Column(_) => true, + Projection::Computed { + expr: SqlExpr::Column { table: None, name }, + alias, + } => name == alias, + Projection::Computed { .. } + | Projection::CpComputed { .. } + | Projection::Star + | Projection::QualifiedStar(_) => false, + } +} + +fn lift_leaf(mut plan: SqlPlan, functions: &FunctionRegistry) -> SqlPlan { + let Some(source) = ScoreSource::of(&plan) else { + return plan; + }; + let Some(original) = projection_mut(&mut plan).cloned() else { + return plan; + }; + + let mut tail = Vec::with_capacity(original.len()); + let mut others = Vec::with_capacity(original.len()); + let mut scores = Vec::new(); + let mut needs_tail = false; + for item in &original { + match source.score_item(item, functions) { + Some(read) => { + // A score read under its own key needs no tail: the hit row + // already carries it under the name the schema announces. + needs_tail |= !source.answers_scores() && !reads_own_key(&read); + scores.push(item.clone()); + tail.push(read); + } + None => { + needs_tail |= matches!(item, Projection::Computed { .. }); + others.push(item.clone()); + tail.push(item.clone()); + } + } + } + if !needs_tail { + return plan; + } + + let has_star = original + .iter() + .any(|p| matches!(p, Projection::Star | Projection::QualifiedStar(_))); + let mut body = if has_star { + vec![Projection::Star] + } else { + body_projection(&others) + }; + match &source { + // The body announces each score under its own key, typed as the + // search types it. + ScoreSource::Text(_) | ScoreSource::Hybrid(_) => body.extend(scores), + ScoreSource::Vector(_) if !has_star => body.push(Projection::Column(DISTANCE_KEY.into())), + ScoreSource::Vector(_) => {} + } + if let Some(projection) = projection_mut(&mut plan) { + *projection = body; + } + // `tail` holds a computed item here, so the wrap always applies. + with_computed_tail(plan, &tail, TailBounds::default()) +} + +/// The SELECT list a search plan carries. +fn projection_mut(plan: &mut SqlPlan) -> Option<&mut Vec> { + match plan { + SqlPlan::TextSearch(search) => Some(&mut search.projection), + SqlPlan::HybridSearch(search) => Some(&mut search.projection), + SqlPlan::HybridSearchTriple(search) => Some(&mut search.projection), + SqlPlan::VectorSearch { projection, .. } + | SqlPlan::SparseSearch { projection, .. } + | SqlPlan::MultiVectorSearch { projection, .. } => Some(projection), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::{TextScoreColumn, TextSearchPlan, TextSearchShape}; + use nodedb_types::text_search::QueryMode; + + fn column(name: &str) -> SqlExpr { + SqlExpr::Column { + table: None, + name: name.into(), + } + } + + fn text_search(projection: Vec) -> SqlPlan { + SqlPlan::TextSearch(TextSearchPlan { + collection: "docs".into(), + shape: TextSearchShape::ScoreScan, + filters: Vec::new(), + scores: vec![TextScoreColumn { + field: Some("body".into()), + query: "x".into(), + mode: QueryMode::Or, + fuzzy: false, + alias: "s".into(), + }], + projection, + }) + } + + fn score_item() -> Projection { + Projection::Computed { + expr: SqlExpr::Function { + name: "bm25_score".into(), + args: vec![column("body")], + distinct: false, + }, + alias: "s".into(), + } + } + + #[test] + fn a_search_with_only_its_score_stays_whole() { + let plan = text_search(vec![Projection::Column("id".into()), score_item()]); + let lifted = lift_search_computed(plan, &FunctionRegistry::new()); + assert!(matches!(lifted, SqlPlan::TextSearch(_)), "{lifted:?}"); + } + + #[test] + fn a_computed_item_moves_to_a_tail_that_reads_the_score() { + let upper = Projection::Computed { + expr: SqlExpr::Function { + name: "upper".into(), + args: vec![column("title")], + distinct: false, + }, + alias: "u".into(), + }; + let plan = text_search(vec![Projection::Column("id".into()), score_item(), upper]); + let SqlPlan::Subquery { + input, projection, .. + } = lift_search_computed(plan, &FunctionRegistry::new()) + else { + panic!("expected a tail"); + }; + assert!(matches!(&projection[1], Projection::Column(name) if name == "s")); + assert!(matches!(&projection[2], Projection::Computed { alias, .. } if alias == "u")); + let SqlPlan::TextSearch(search) = *input else { + panic!("expected the search as the body"); + }; + let body: Vec = search + .projection + .iter() + .map(|p| match p { + Projection::Column(name) => name.clone(), + Projection::Computed { alias, .. } => format!("score:{alias}"), + other => panic!("unexpected body item {other:?}"), + }) + .collect(); + assert_eq!(body, ["id", "title", "score:s"]); + } +} diff --git a/nodedb-sql/src/planner/select/array_from.rs b/nodedb-sql/src/planner/select/array_from.rs new file mode 100644 index 000000000..d4ddc0dda --- /dev/null +++ b/nodedb-sql/src/planner/select/array_from.rs @@ -0,0 +1,630 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! `FROM ARRAY_*(...)`: an array table-function call as a SELECT's source. +//! +//! `ARRAY_SLICE` and `ARRAY_PROJECT` read cells. Their rows form a relation +//! whose columns are the array's dims and the requested attrs, so the SELECT +//! plans over it like over a derived table: the call becomes the one +//! definition of a CTE, and the SELECT reads that definition by name through +//! the generic relational tail (WHERE, projection, GROUP BY, HAVING, ORDER +//! BY, DISTINCT, LIMIT / OFFSET, windows). +//! +//! `ARRAY_AGG` and `ARRAY_ELEMENTWISE` emit rows of their own shape. A +//! SELECT over them takes `*` or plain output columns and a row bound; +//! every other clause is refused by name. + +use sqlparser::ast::{self, Select}; + +use super::cte_catalog::CteCatalog; +use super::helpers::{convert_projection, plain_table_factor}; +use super::query_tail::QueryTail; +use super::select_stmt::{PlannedSelect, has_aggregation, plan_select}; +use crate::error::{Result, SqlError}; +use crate::functions::registry::FunctionRegistry; +use crate::planner::array_fn::{ArrayTableFn, push_down_into_slice, try_plan_array_table_fn}; +use crate::planner::ast_helpers::{single_table_qualifiers, strip_single_table_qualifiers}; +use crate::resolver::array_tvf::{ARRAY_SYSTEM_TIME_COLUMN, array_slice_columns}; +use crate::resolver::columns::TableScope; +use crate::temporal::TemporalScope; +use crate::types::*; + +/// Plan a SELECT whose FROM is one array table-function call. Returns +/// `Ok(None)` when the FROM is anything else. +pub(in crate::planner::select) fn try_plan_array_from( + select: &Select, + catalog: &dyn SqlCatalog, + functions: &FunctionRegistry, + temporal: TemporalScope, + tail: &QueryTail<'_>, + statement_output: bool, +) -> Result> { + let Some(call) = try_plan_array_table_fn(&select.from, catalog, temporal)? else { + return Ok(None); + }; + let ctx = OuterContext { + select, + catalog, + functions, + tail, + statement_output, + }; + let planned = match call { + ArrayTableFn::Slice(slice) => plan_cell_read(&ctx, "ARRAY_SLICE", slice)?, + ArrayTableFn::Project(slice) => plan_cell_read(&ctx, "ARRAY_PROJECT", slice)?, + call @ (ArrayTableFn::Agg(_) | ArrayTableFn::Elementwise(_)) => { + plan_non_relational(&ctx, temporal, call)? + } + }; + Ok(Some(planned)) +} + +/// Plan the SELECT over the cells `slice` reads, as the relation `fn_name` +/// names. +fn plan_cell_read( + ctx: &OuterContext<'_>, + fn_name: &str, + mut slice: ArraySlicePlan, +) -> Result { + let view = array_view(ctx.catalog, &slice.name)?; + let read = RelationRead::new(ctx, fn_name, &view.name)?; + let columns = array_slice_columns(&view, &slice); + let outer = read.plan_outer(ctx, columns)?; + push_down_into_slice(&mut slice, &view, &outer.plan, &read.name); + Ok(read.into_cte(SqlPlan::ArraySlice(slice), outer)) +} + +/// The parts of the enclosing SELECT every arm reads. +struct OuterContext<'a> { + select: &'a Select, + catalog: &'a dyn SqlCatalog, + functions: &'a FunctionRegistry, + tail: &'a QueryTail<'a>, + statement_output: bool, +} + +/// The relation one cell-reading call exposes to the SELECT over it. +struct RelationRead { + /// The name the outer query reads the rows by. No stored relation or + /// enclosing CTE holds it, so a subquery that names a real collection + /// never resolves to the array. + name: String, + /// The qualifier column references use: the FROM alias, or the array + /// name when the call has none. + qualifier: ast::Ident, +} + +impl RelationRead { + fn new(ctx: &OuterContext<'_>, fn_name: &str, array: &str) -> Result { + let alias = match &ctx.select.from[0].relation { + ast::TableFactor::Table { alias, .. } => alias.as_ref(), + _ => None, + }; + if let Some(alias) = alias + && !alias.columns.is_empty() + { + return Err(SqlError::Unsupported { + detail: format!( + "{fn_name}: a column alias list on '{}' is not supported; \ + alias each column in the SELECT list instead", + alias.name + ), + }); + } + let qualifier = match alias { + Some(alias) => alias.name.clone(), + None => ast::Ident::with_quote('"', array), + }; + Ok(Self { + name: unused_relation_name(ctx.catalog, &format!("{fn_name}({array})"))?, + qualifier, + }) + } + + /// Plan the SELECT as a read of this relation, whose rows carry `columns`. + fn plan_outer( + &self, + ctx: &OuterContext<'_>, + columns: Vec, + ) -> Result { + let star_columns: Vec = columns + .iter() + .map(|c| c.name.clone()) + .filter(|name| name != ARRAY_SYSTEM_TIME_COLUMN) + .collect(); + let info = CollectionInfo { + name: self.name.clone(), + engine: EngineType::DocumentSchemaless, + columns, + primary_key: None, + has_auto_tier: false, + indexes: Vec::new(), + bitemporal: false, + primary: nodedb_types::PrimaryEngine::Document, + vector_primary: None, + partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed, + open_schema: false, + }; + let relation_catalog = CteCatalog { + inner: ctx.catalog, + relations: vec![(self.name.clone(), info)], + }; + let qualifier = crate::reserved::check_ast_identifier(&self.qualifier)?; + let mut outer_select = ctx.select.clone(); + outer_select.from[0].relation = plain_table_factor( + ast::Ident::with_quote('"', self.name.clone()), + Some(ast::TableAlias { + explicit: true, + name: self.qualifier.clone(), + columns: Vec::new(), + at: None, + }), + ); + outer_select.projection = + expand_wildcards(&ctx.select.projection, &qualifier, &star_columns)?; + // The temporal qualifier reads the array's versions and rides on the + // slice plan. The relation the SELECT reads is those versions' + // rows, with no history of its own. + plan_select( + &outer_select, + &relation_catalog, + ctx.functions, + TemporalScope::default(), + ctx.tail, + ctx.statement_output, + ) + } + + fn into_cte(self, source: SqlPlan, outer: PlannedSelect) -> PlannedSelect { + PlannedSelect { + plan: SqlPlan::Cte(CtePlan { + definitions: vec![(self.name, source)], + outer: Box::new(outer.plan), + }), + scope: outer.scope, + } + } +} + +fn array_view(catalog: &dyn SqlCatalog, name: &str) -> Result { + catalog + .lookup_array(name) + .ok_or_else(|| SqlError::UnknownTable { + name: name.to_string(), + }) +} + +/// `base`, or `base#N` for the first `N` from 2 that no relation the +/// catalog resolves holds. +fn unused_relation_name(catalog: &dyn SqlCatalog, base: &str) -> Result { + let mut candidate = base.to_string(); + let mut n: u32 = 1; + while catalog + .resolve_relation(nodedb_types::DatabaseId::DEFAULT, &candidate)? + .is_some() + { + n = n.checked_add(1).ok_or_else(|| SqlError::Unsupported { + detail: format!("no free relation name for {base}"), + })?; + candidate = format!("{base}#{n}"); + } + Ok(candidate) +} + +/// The SELECT list with each plain `*`, and each `.*`, replaced by +/// one column per name in `columns`, in that order. +/// +/// The relation's columns are known, so the result announces them in +/// schema order with their types. A wildcard with options (`EXCLUDE`, +/// `REPLACE`, ...) or another qualifier stays as written. +fn expand_wildcards( + projection: &[ast::SelectItem], + qualifier: &str, + columns: &[String], +) -> Result> { + let mut out = Vec::with_capacity(projection.len() + columns.len()); + for item in projection { + let expands = match item { + ast::SelectItem::Wildcard(options) => is_plain_wildcard(options), + ast::SelectItem::QualifiedWildcard( + ast::SelectItemQualifiedWildcardKind::ObjectName(name), + options, + ) => { + is_plain_wildcard(options) + && crate::parser::normalize::normalize_object_name_checked(name)? == qualifier + } + ast::SelectItem::QualifiedWildcard( + ast::SelectItemQualifiedWildcardKind::Expr(_), + _, + ) + | ast::SelectItem::UnnamedExpr(_) + | ast::SelectItem::ExprWithAlias { .. } + | ast::SelectItem::ExprWithAliases { .. } => false, + }; + if expands { + out.extend(columns.iter().map(|name| { + ast::SelectItem::UnnamedExpr(ast::Expr::Identifier(ast::Ident::with_quote( + '"', + name.clone(), + ))) + })); + } else { + out.push(item.clone()); + } + } + Ok(out) +} + +fn is_plain_wildcard(options: &ast::WildcardAdditionalOptions) -> bool { + options.opt_ilike.is_none() + && options.opt_exclude.is_none() + && options.opt_except.is_none() + && options.opt_replace.is_none() + && options.opt_rename.is_none() + && options.opt_alias.is_none() +} + +/// Plan a SELECT over `ARRAY_AGG` / `ARRAY_ELEMENTWISE`. +/// +/// The call's rows are its result. The SELECT may list `*` or the call's +/// own output columns, and bound the rows with LIMIT / FETCH. A filter, +/// grouping, ordering, dedup, offset, window or computed column has no +/// dims-and-attrs relation to read, so each is refused by name. +fn plan_non_relational( + ctx: &OuterContext<'_>, + temporal: TemporalScope, + call: ArrayTableFn, +) -> Result { + let fn_name = call.sql_name(); + let select = ctx.select; + let refuse = |clause: &str| -> Result { + Err(SqlError::Unsupported { + detail: format!( + "{clause} over {fn_name} is not supported: {fn_name} emits its own result \ + rows, not the array's dims and attrs; read the cells with ARRAY_SLICE" + ), + }) + }; + if select.selection.is_some() { + return refuse("WHERE"); + } + if has_aggregation(select, ctx.functions) { + return refuse("GROUP BY or an aggregate"); + } + if select.having.is_some() { + return refuse("HAVING"); + } + if select.qualify.is_some() { + return refuse("QUALIFY"); + } + if !select.named_window.is_empty() { + return refuse("WINDOW"); + } + if select.distinct.is_some() { + return refuse("DISTINCT"); + } + if ctx.tail.order_by.is_some() { + return refuse("ORDER BY"); + } + if ctx.tail.limit_offset()?.1 != 0 { + return refuse("OFFSET"); + } + + let scope = TableScope::resolve_from(ctx.catalog, ctx.functions, temporal, &select.from)?; + let table = scope + .single_table() + .cloned() + .ok_or_else(|| SqlError::Unsupported { + detail: format!("{fn_name}: FROM must name the call alone"), + })?; + let normalized = strip_single_table_qualifiers(select, &single_table_qualifiers(&table))?; + let projection = convert_projection(&normalized.projection, &scope)?; + let plain = projection.iter().all(|p| match p { + Projection::Star | Projection::QualifiedStar(_) | Projection::Column(_) => true, + Projection::Computed { expr, .. } => matches!(expr, SqlExpr::Column { .. }), + Projection::CpComputed { .. } => false, + }); + if !plain { + return refuse("a computed SELECT-list expression"); + } + + let source = call.into_plan(); + let star_only = projection + .iter() + .all(|p| matches!(p, Projection::Star | Projection::QualifiedStar(_))); + let plan = if star_only { + source + } else { + SqlPlan::Subquery { + input: Box::new(source), + filters: Vec::new(), + projection, + window_functions: Vec::new(), + sort_keys: Vec::new(), + offset: 0, + distinct: false, + limit: None, + } + }; + Ok(PlannedSelect { plan, scope }) +} + +#[cfg(test)] +mod tests { + use crate::catalog::{ArrayCatalogView, SqlCatalogError}; + use crate::error::{Result, SqlError}; + use crate::functions::registry::FunctionRegistry; + use crate::parser::statement::parse_sql; + use crate::types::*; + use crate::types_array::{ + ArrayAttrAst, ArrayAttrType, ArrayCoordLiteral, ArrayDimAst, ArrayDimType, ArrayDomainBound, + }; + + /// A catalog holding the 4x4 `img` array. `taken` names a stored + /// collection, to collide with the relation name the planner picks. + struct ImgCatalog { + taken: Option<&'static str>, + } + + impl SqlCatalog for ImgCatalog { + fn get_collection( + &self, + _: nodedb_types::DatabaseId, + name: &str, + ) -> std::result::Result, SqlCatalogError> { + Ok((self.taken == Some(name)).then(|| CollectionInfo { + name: name.into(), + engine: EngineType::DocumentSchemaless, + columns: Vec::new(), + primary_key: Some("id".into()), + has_auto_tier: false, + indexes: Vec::new(), + bitemporal: false, + primary: nodedb_types::PrimaryEngine::Document, + vector_primary: None, + partition_strategy: nodedb_types::PartitionStrategy::CollectionHomed, + open_schema: true, + })) + } + + fn lookup_array(&self, name: &str) -> Option { + let dim = |name: &str| ArrayDimAst { + name: name.into(), + dtype: ArrayDimType::Int64, + lo: ArrayDomainBound::Int64(0), + hi: ArrayDomainBound::Int64(3), + }; + let attr = |name: &str| ArrayAttrAst { + name: name.into(), + dtype: ArrayAttrType::Float64, + nullable: true, + }; + (name == "img").then(|| ArrayCatalogView { + name: "img".into(), + dims: vec![dim("x"), dim("y")], + attrs: vec![attr("nir"), attr("red")], + tile_extents: vec![2, 2], + }) + } + } + + fn plan_with(catalog: &ImgCatalog, sql: &str) -> Result { + let stmts = parse_sql(sql)?; + let sqlparser::ast::Statement::Query(q) = &stmts[0] else { + panic!("not a query: {sql}"); + }; + crate::planner::select::plan_query( + q, + catalog, + &FunctionRegistry::new(), + crate::TemporalScope::default(), + ) + } + + fn plan(sql: &str) -> Result { + plan_with(&ImgCatalog { taken: None }, sql) + } + + /// The relation name, the slice body and the outer read of a planned + /// cell read. + fn parts(plan: SqlPlan) -> (String, ArraySlicePlan, SqlPlan) { + let SqlPlan::Cte(CtePlan { + mut definitions, + outer, + }) = plan + else { + panic!("expected Cte, got {plan:?}"); + }; + let (name, body) = definitions.remove(0); + let SqlPlan::ArraySlice(slice) = body else { + panic!("expected ArraySlice body, got {body:?}"); + }; + (name, slice, *outer) + } + + fn projected_names(plan: &SqlPlan) -> Vec { + let SqlPlan::Scan { projection, .. } = plan else { + panic!("expected outer Scan, got {plan:?}"); + }; + projection + .iter() + .map(|p| match p { + Projection::Column(name) => name.clone(), + Projection::Computed { alias, .. } => alias.clone(), + other => panic!("unexpected projection {other:?}"), + }) + .collect() + } + + #[test] + fn star_expands_to_dims_then_attrs() { + let (name, _, outer) = + parts(plan("SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}', ['red'])").expect("plans")); + assert_eq!(name, "ARRAY_SLICE(img)"); + assert_eq!(projected_names(&outer), ["x", "y", "red"]); + } + + #[test] + fn where_and_computed_projection_ride_the_outer_read() { + let (name, slice, outer) = parts( + plan( + "SELECT x, y, (nir - red) / (nir + red) AS ndvi \ + FROM ARRAY_SLICE('img', '{x: [0, 3]}') AS t WHERE t.nir > 0.5", + ) + .expect("plans"), + ); + let SqlPlan::Scan { + collection, + filters, + .. + } = &outer + else { + panic!("expected outer Scan, got {outer:?}"); + }; + assert_eq!(collection, &name); + assert_eq!(filters.len(), 1, "the WHERE stays on the outer read"); + assert_eq!(projected_names(&outer), ["x", "y", "ndvi"]); + assert_eq!(slice.limit, 0); + } + + #[test] + fn an_unknown_column_is_refused() { + let err = plan("SELECT blue FROM ARRAY_SLICE('img', '{x: [0, 3]}')").expect_err("refused"); + assert!( + matches!(err, SqlError::UnknownColumn { .. }), + "unexpected error {err:?}" + ); + } + + #[test] + fn a_dim_bound_narrows_an_uncapped_slice() { + let (_, slice, _) = parts( + plan("SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}') WHERE y >= 2 AND x < 3") + .expect("plans"), + ); + let range = |dim: &str| { + slice + .slice + .dim_ranges + .iter() + .find(|r| r.dim == dim) + .map(|r| (r.lo.clone(), r.hi.clone())) + }; + assert_eq!( + range("x"), + Some((ArrayCoordLiteral::Int64(0), ArrayCoordLiteral::Int64(2))) + ); + assert_eq!( + range("y"), + Some((ArrayCoordLiteral::Int64(2), ArrayCoordLiteral::Int64(3))) + ); + } + + #[test] + fn an_outer_limit_folds_to_the_smaller_cap() { + let (_, slice, _) = parts( + plan("SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}', [], 2) LIMIT 100") + .expect("plans"), + ); + assert_eq!(slice.limit, 2); + let (_, slice, outer) = + parts(plan("SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}') LIMIT 3").expect("plans")); + assert_eq!(slice.limit, 3); + assert!( + matches!(outer, SqlPlan::Scan { limit: Some(3), .. }), + "the LIMIT stays on the outer read: {outer:?}" + ); + } + + #[test] + fn an_ordered_limit_does_not_fold() { + let (_, slice, _) = parts( + plan("SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}') ORDER BY nir DESC LIMIT 3") + .expect("plans"), + ); + assert_eq!(slice.limit, 0); + } + + #[test] + fn group_by_plans_an_aggregate_over_the_relation() { + let (name, _, outer) = parts( + plan("SELECT x, AVG(nir) AS a FROM ARRAY_SLICE('img', '{x: [0, 3]}') GROUP BY x") + .expect("plans"), + ); + let SqlPlan::Aggregate { input, .. } = outer else { + panic!("expected Aggregate, got {outer:?}"); + }; + assert!(matches!(*input, SqlPlan::Scan { ref collection, .. } if *collection == name)); + } + + #[test] + fn the_relation_name_avoids_a_stored_collection() { + let (name, _, _) = parts( + plan_with( + &ImgCatalog { + taken: Some("ARRAY_SLICE(img)"), + }, + "SELECT * FROM ARRAY_SLICE('img', '{x: [0, 3]}')", + ) + .expect("plans"), + ); + assert_eq!(name, "ARRAY_SLICE(img)#2"); + } + + #[test] + fn the_inlined_read_filters_the_slice_rows() { + let SqlPlan::Cte(CtePlan { definitions, outer }) = + plan("SELECT x FROM ARRAY_SLICE('img', '{x: [0, 3]}') WHERE nir > 0.5").expect("plans") + else { + panic!("expected Cte"); + }; + let (name, body) = &definitions[0]; + match crate::planner::cte::inline_cte(&outer, name, body) { + SqlPlan::Subquery { input, filters, .. } => { + assert!(matches!(*input, SqlPlan::ArraySlice(_))); + assert_eq!(filters.len(), 1); + } + other => panic!("expected Subquery over the slice, got {other:?}"), + } + } + + #[test] + fn where_over_array_agg_is_refused_by_name() { + let err = plan("SELECT * FROM ARRAY_AGG('img', 'nir', 'sum') WHERE result > 1") + .expect_err("refused"); + let SqlError::Unsupported { detail } = err else { + panic!("expected Unsupported, got {err:?}"); + }; + assert!(detail.starts_with("WHERE over ARRAY_AGG"), "{detail}"); + } + + #[test] + fn order_by_over_array_elementwise_is_refused_by_name() { + let err = + plan("SELECT * FROM ARRAY_ELEMENTWISE('img', 'img', 'add', 'nir') ORDER BY coords") + .expect_err("refused"); + let SqlError::Unsupported { detail } = err else { + panic!("expected Unsupported, got {err:?}"); + }; + assert!( + detail.starts_with("ORDER BY over ARRAY_ELEMENTWISE"), + "{detail}" + ); + } + + #[test] + fn plain_agg_columns_wrap_in_a_projecting_tail() { + match plan("SELECT result FROM ARRAY_AGG('img', 'nir', 'sum') LIMIT 1").expect("plans") { + SqlPlan::Subquery { + input, + projection, + limit, + .. + } => { + assert!(matches!(*input, SqlPlan::ArrayAgg(_))); + assert_eq!(projection.len(), 1); + assert_eq!(limit, Some(1)); + } + other => panic!("expected Subquery over ArrayAgg, got {other:?}"), + } + } +} diff --git a/nodedb-sql/src/planner/select/cte_catalog.rs b/nodedb-sql/src/planner/select/cte_catalog.rs index 1664ae0b9..3cc9df08e 100644 --- a/nodedb-sql/src/planner/select/cte_catalog.rs +++ b/nodedb-sql/src/planner/select/cte_catalog.rs @@ -4,8 +4,9 @@ use nodedb_types::DatabaseId; +use crate::error::Result; use crate::resolver::derived::open_subquery_relation; -use crate::types::{CollectionInfo, SqlCatalog, SqlCatalogError}; +use crate::types::{ArrayCatalogView, CollectionInfo, SqlCatalog, SqlCatalogError}; /// Catalog wrapper that answers for a synthesized relation before delegating /// to the stored catalog. @@ -28,15 +29,86 @@ impl<'a> CteCatalog<'a> { } } +impl CteCatalog<'_> { + fn relation(&self, name: &str) -> Option { + self.relations + .iter() + .find(|(key, _)| key == name) + .map(|(_, info)| info.clone()) + } +} + +/// A synthesized relation shadows a stored one of the same name. Every +/// other lookup reaches the stored catalog unchanged, so a query planned +/// against this wrapper still sees catalog tables, arrays, types and +/// sequences. impl SqlCatalog for CteCatalog<'_> { fn get_collection( &self, database_id: DatabaseId, name: &str, ) -> std::result::Result, SqlCatalogError> { - if let Some((_, info)) = self.relations.iter().find(|(key, _)| key == name) { - return Ok(Some(info.clone())); + match self.relation(name) { + Some(info) => Ok(Some(info)), + None => self.inner.get_collection(database_id, name), } - self.inner.get_collection(database_id, name) + } + + fn resolve_relation( + &self, + database_id: DatabaseId, + name: &str, + ) -> std::result::Result, SqlCatalogError> { + match self.relation(name) { + Some(info) => Ok(Some(info)), + None => self.inner.resolve_relation(database_id, name), + } + } + + /// A synthesized relation is never a timeseries collection. + fn timeseries_time_key( + &self, + database_id: DatabaseId, + name: &str, + ) -> std::result::Result, SqlCatalogError> { + match self.relation(name) { + Some(_) => Ok(None), + None => self.inner.timeseries_time_key(database_id, name), + } + } + + fn lookup_array(&self, name: &str) -> Option { + self.inner.lookup_array(name) + } + + fn array_exists(&self, name: &str) -> bool { + self.inner.array_exists(name) + } + + fn resolve_regclass(&self, database_id: DatabaseId, tenant_id: u64, name: &str) -> Option { + self.inner.resolve_regclass(database_id, tenant_id, name) + } + + fn resolve_regtype(&self, name: &str) -> Option { + self.inner.resolve_regtype(name) + } + + fn sequence_nextval(&self, database_id: DatabaseId, tenant_id: u64, name: &str) -> Result { + self.inner.sequence_nextval(database_id, tenant_id, name) + } + + fn sequence_currval(&self, database_id: DatabaseId, tenant_id: u64, name: &str) -> Result { + self.inner.sequence_currval(database_id, tenant_id, name) + } + + fn sequence_setval( + &self, + database_id: DatabaseId, + tenant_id: u64, + name: &str, + value: i64, + ) -> Result { + self.inner + .sequence_setval(database_id, tenant_id, name, value) } } diff --git a/nodedb-sql/src/planner/select/derived_from.rs b/nodedb-sql/src/planner/select/derived_from.rs index f535875d9..07488c351 100644 --- a/nodedb-sql/src/planner/select/derived_from.rs +++ b/nodedb-sql/src/planner/select/derived_from.rs @@ -7,6 +7,7 @@ use sqlparser::ast::{self, Select}; use super::cte_catalog::CteCatalog; use super::derived_output::rekey_join_output; use super::entry::plan_query; +use super::helpers::plain_table_factor; use super::query_tail::QueryTail; use super::select_stmt::{PlannedSelect, plan_select}; use crate::error::Result; @@ -76,20 +77,7 @@ pub(in crate::planner::select) fn try_plan_derived_from( )], }; let mut outer_select = select.clone(); - outer_select.from[0].relation = ast::TableFactor::Table { - name: ast::ObjectName(vec![ast::ObjectNamePart::Identifier( - alias_ident.name.clone(), - )]), - alias: None, - args: None, - with_hints: Vec::new(), - version: None, - with_ordinality: false, - partitions: Vec::new(), - json_path: None, - sample: None, - index_hints: Vec::new(), - }; + outer_select.from[0].relation = plain_table_factor(alias_ident.name.clone(), None); // The outer SELECT keeps the caller's output standing; the derived body // planned above is nested. let outer = plan_select( diff --git a/nodedb-sql/src/planner/select/entry/search.rs b/nodedb-sql/src/planner/select/entry/search.rs index d3ba9d08b..d6e1b2f47 100644 --- a/nodedb-sql/src/planner/select/entry/search.rs +++ b/nodedb-sql/src/planner/select/entry/search.rs @@ -114,6 +114,9 @@ pub(super) fn plan_select_query( pre_order_by_collection.as_deref(), )?; let plan = apply_limit(plan, &tail)?; + // A search plan evaluates no expression. Its computed items move to a + // tail now that the search holds its own ORDER BY and LIMIT bound. + let plan = crate::planner::search_tail::lift_search_computed(plan, functions); // ORDER BY and LIMIT sit on the aggregate by now, so the wrap // only restates the output columns around it. wrap_aggregate_cp_items( @@ -794,4 +797,84 @@ mod tests { assert!(detail.contains("bm25_score"), "{sql}: {detail}"); } } + + /// The vector search's filter and top-k bound, from a top-level plan. + fn vector_filters_and_k(plan: &SqlPlan) -> (usize, usize) { + let SqlPlan::VectorSearch { filters, top_k, .. } = plan else { + panic!("expected VectorSearch, got {plan:?}"); + }; + (filters.len(), *top_k) + } + + /// The remote client's `vector_search(k, filter)` shape: the distance + /// under its own row key needs no tail, so the search stays the root + /// and keeps its filter and top-k bound. + #[test] + fn a_filtered_vector_search_without_computed_items_stays_whole() { + for select in [ + "SELECT id", + "SELECT id, vector_distance(embedding, ARRAY[0.0, 0.0]) AS distance", + ] { + let plan = plan_select_sql(&format!( + "{select} FROM embeddings WHERE category = 'keep' \ + ORDER BY vector_distance(embedding, ARRAY[0.0, 0.0]) LIMIT 2" + )); + assert_eq!(vector_filters_and_k(&plan), (1, 2), "{select}"); + } + } + + /// A WHERE predicate beside a sparse search rides on the search, so the + /// search can filter before its top-k cut. + #[test] + fn a_where_predicate_with_a_sparse_search_rides_on_the_search() { + let plan = plan_select_sql( + "SELECT id FROM embeddings WHERE category = 'keep' \ + ORDER BY sparse_score(terms, '{3: 1.0}') DESC LIMIT 2", + ); + let SqlPlan::SparseSearch { filters, top_k, .. } = &plan else { + panic!("expected SparseSearch, got {plan:?}"); + }; + assert_eq!((filters.len(), *top_k), (1, 2)); + + let plan = plan_select_sql( + "SELECT id FROM embeddings ORDER BY sparse_score(terms, '{3: 1.0}') DESC LIMIT 2", + ); + let SqlPlan::SparseSearch { filters, .. } = &plan else { + panic!("expected SparseSearch, got {plan:?}"); + }; + assert!(filters.is_empty()); + } + + /// A computed item moves to a tail. The search below it keeps its filter + /// and top-k bound, and the tail reads the distance by its row key. + #[test] + fn a_filtered_vector_search_with_a_computed_item_keeps_filter_and_k() { + let plan = plan_select_sql( + "SELECT id, upper(category) AS c, \ + vector_distance(embedding, ARRAY[0.0, 0.0]) AS d \ + FROM embeddings WHERE category = 'keep' \ + ORDER BY vector_distance(embedding, ARRAY[0.0, 0.0]) LIMIT 2", + ); + let SqlPlan::Subquery { + input, + projection, + limit, + .. + } = &plan + else { + panic!("expected a tail, got {plan:?}"); + }; + assert_eq!(*limit, None, "the bound stays on the search"); + assert_eq!(vector_filters_and_k(input), (1, 2)); + assert!( + projection.iter().any(|p| matches!( + p, + Projection::Computed { + expr: SqlExpr::Column { name, .. }, + alias, + } if name == "distance" && alias == "d" + )), + "the tail reads the distance by its row key: {projection:?}" + ); + } } diff --git a/nodedb-sql/src/planner/select/helpers.rs b/nodedb-sql/src/planner/select/helpers.rs index 7de98f5c9..7408d1c91 100644 --- a/nodedb-sql/src/planner/select/helpers.rs +++ b/nodedb-sql/src/planner/select/helpers.rs @@ -27,6 +27,29 @@ pub(super) fn source_projection(plan: &SqlPlan) -> Vec { } } +/// A FROM factor that names the relation `name`, under `alias` when given. +/// +/// A derived table or a table-function call is replanned as a read of the +/// relation its body exposes. This factor names that relation in the outer +/// SELECT, so the outer query resolves it like a stored collection. +pub(super) fn plain_table_factor( + name: ast::Ident, + alias: Option, +) -> ast::TableFactor { + ast::TableFactor::Table { + name: ast::ObjectName(vec![ast::ObjectNamePart::Identifier(name)]), + alias, + args: None, + with_hints: Vec::new(), + version: None, + with_ordinality: false, + partitions: Vec::new(), + json_path: None, + sample: None, + index_hints: Vec::new(), + } +} + /// Convert SELECT projection items. /// /// An item over a relation that calls a sequence accessor becomes diff --git a/nodedb-sql/src/planner/select/limit.rs b/nodedb-sql/src/planner/select/limit.rs index 6d2c68986..2f71e457f 100644 --- a/nodedb-sql/src/planner/select/limit.rs +++ b/nodedb-sql/src/planner/select/limit.rs @@ -3,17 +3,19 @@ //! LIMIT / OFFSET application. //! //! Each plan variant owns a different slot for the row bound — a scan's -//! `limit`/`offset`, a search plan's `top_k`, an array slice's `u32` cap. A -//! variant with no slot at all gets a post-processing tail rather than -//! silently losing the clause. +//! `limit`/`offset`, a search plan's `top_k`. A variant with no slot at all +//! gets a post-processing tail rather than silently losing the clause. An +//! array slice reaches here only as the body of the relation a SELECT reads, +//! so the LIMIT lands on that read; the slice's own cap folds in during +//! planning of the read. use super::post_process::post_process; use super::query_tail::QueryTail; -use crate::error::{Result, SqlError}; +use crate::error::Result; use crate::types::SqlPlan; use crate::types::{ - ArraySlicePlan, CtePlan, DocumentIndexLookupPlan, HybridSearchPlan, HybridSearchTriplePlan, - TextSearchPlan, TextSearchShape, TimeseriesScanPlan, + CtePlan, DocumentIndexLookupPlan, HybridSearchPlan, HybridSearchTriplePlan, TextSearchPlan, + TextSearchShape, TimeseriesScanPlan, }; /// Default `ef_search` multiplier applied when LIMIT is the only signal @@ -111,17 +113,6 @@ pub(in crate::planner::select) fn apply_limit( *l = lv; } } - SqlPlan::ArraySlice(ArraySlicePlan { - limit: ref mut l, .. - }) => { - if let Some(lv) = limit_val { - // The slice bound is a `u32` (0 = unlimited); a LIMIT that - // does not fit would wrap into a smaller bound. - *l = u32::try_from(lv).map_err(|_| SqlError::Unsupported { - detail: format!("LIMIT {lv} exceeds the array slice row bound"), - })?; - } - } // Search plans rank rows and return the best `top_k`, so the query's // `LIMIT N` *is* the top-k bound. `ef_search` is deliberately left // alone on the fused-search variants: a wider beam than the final N diff --git a/nodedb-sql/src/planner/select/mod.rs b/nodedb-sql/src/planner/select/mod.rs index 8a057e852..ace371d69 100644 --- a/nodedb-sql/src/planner/select/mod.rs +++ b/nodedb-sql/src/planner/select/mod.rs @@ -6,6 +6,7 @@ //! search patterns (vector, text, hybrid, spatial) directly from the AST //! instead of reverse-engineering an optimizer's output. +mod array_from; mod comma_lateral; mod cte_catalog; mod derived_from; diff --git a/nodedb-sql/src/planner/select/order_by/apply.rs b/nodedb-sql/src/planner/select/order_by/apply.rs index 735e56276..7236496e7 100644 --- a/nodedb-sql/src/planner/select/order_by/apply.rs +++ b/nodedb-sql/src/planner/select/order_by/apply.rs @@ -96,7 +96,12 @@ pub(in crate::planner::select) fn apply_order_by( // finalized, so each call is bound to the column it lands in and any // aggregate the sort alone introduces joins the computed list. let mut bound_aggregates: Option> = None; - let sort_keys: Vec = if let SqlPlan::Aggregate { aggregates, .. } = plan { + let sort_keys: Vec = if let SqlPlan::Aggregate { + aggregates, + group_by, + .. + } = plan + { let mut extended = aggregates.clone(); // A sort-only aggregate reads the aggregate's input rows. Over a // join those are the merged rows, keyed `qualifier.column`, while the @@ -132,8 +137,9 @@ pub(in crate::planner::select) fn apply_order_by( .iter() .zip(&bound) .map(|(o, expr)| { + let converted = convert_expr(expr, &ColumnScope::Relations(&sort_scope))?; Ok(SortKey { - expr: convert_expr(expr, &ColumnScope::Relations(&sort_scope))?, + expr: computed_group_key_ref(converted, group_by), ascending: sort_ascending(&o.options)?, nulls_first: sort_nulls_first(&o.options)?, }) @@ -239,6 +245,8 @@ pub(in crate::planner::select) fn apply_order_by( limit, tiered, temporal, + output, + group_keys, .. }) => Ok(SqlPlan::TimeseriesScan(TimeseriesScanPlan { collection: collection.clone(), @@ -254,6 +262,8 @@ pub(in crate::planner::select) fn apply_order_by( sort_keys, tiered: *tiered, temporal: *temporal, + output: output.clone(), + group_keys: group_keys.clone(), })), // Cte wraps an inner outer plan; push ORDER BY into that outer // so derived-table queries (`SELECT … FROM (…) AS t ORDER BY …`) @@ -277,3 +287,29 @@ pub(in crate::planner::select) fn apply_order_by( other => post_process(other.clone(), sort_keys, None, 0), } } + +/// `expr` as a reference to the group-row column of the computed GROUP BY key +/// it equals, else unchanged. +/// +/// A finalized group row carries a computed key (`time_bucket('1 hour', ts)`) +/// under its `group_{index}` row name, and none of the input columns the +/// expression reads. Evaluating the expression over the group row reads +/// NULL, so the sort key reads the key's column instead. `SqlExpr` has no +/// `PartialEq`, so the match compares the `Debug` rendering. Both sides pass +/// through `convert_expr`, so equal expressions render identically. +fn computed_group_key_ref(expr: SqlExpr, group_by: &[SqlExpr]) -> SqlExpr { + if matches!(expr, SqlExpr::Column { .. }) { + return expr; + } + let rendered = format!("{expr:?}"); + match group_by + .iter() + .position(|key| format!("{key:?}") == rendered) + { + Some(index) => SqlExpr::Column { + table: None, + name: crate::planner::agg_naming::computed_group_key_name(index), + }, + None => expr, + } +} diff --git a/nodedb-sql/src/planner/select/order_by/timeseries_native.rs b/nodedb-sql/src/planner/select/order_by/timeseries_native.rs index a0d206a56..7f04cdc8b 100644 --- a/nodedb-sql/src/planner/select/order_by/timeseries_native.rs +++ b/nodedb-sql/src/planner/select/order_by/timeseries_native.rs @@ -5,8 +5,8 @@ //! The native scan sorts its encoded group rows before LIMIT. Those rows //! carry each GROUP BY column under its own name, the time bucket under //! `bucket`, and each aggregate under its canonical key (`count(*)`, -//! `sum(v)`). The SELECT-list aliases are applied later, by the Control -//! Plane. So a sort key must name the encoded column: an alias resolves to +//! `sum(v)`). The scan renames them to their SELECT-list names after the +//! sort. So a sort key must name the encoded column: an alias resolves to //! its expression, an aggregate call binds to its canonical key, and a //! `time_bucket(...)` call names `bucket`. @@ -24,7 +24,7 @@ use crate::resolver::expr::convert_expr; use crate::types::{AggregateExpr, SortKey}; /// The column the native scan writes the time bucket under. -const BUCKET_COLUMN: &str = "bucket"; +const BUCKET_COLUMN: &str = crate::types::TIMESERIES_BUCKET_KEY; /// Sort keys for a native grouped timeseries scan computing `aggregates`. /// diff --git a/nodedb-sql/src/planner/select/order_by/triggers.rs b/nodedb-sql/src/planner/select/order_by/triggers.rs index ebd9a5134..c197b38ab 100644 --- a/nodedb-sql/src/planner/select/order_by/triggers.rs +++ b/nodedb-sql/src/planner/select/order_by/triggers.rs @@ -230,6 +230,11 @@ pub(super) fn try_extract_sort_search( field, query_entries, top_k, + // The scan's WHERE narrows the candidates before the cut. + filters: match plan { + SqlPlan::Scan { filters, .. } => filters.clone(), + _ => Vec::new(), + }, projection: source_projection(plan), }))) } diff --git a/nodedb-sql/src/planner/select/post_process.rs b/nodedb-sql/src/planner/select/post_process.rs index 77467482e..9c552ca6d 100644 --- a/nodedb-sql/src/planner/select/post_process.rs +++ b/nodedb-sql/src/planner/select/post_process.rs @@ -18,9 +18,8 @@ use crate::error::{Result, SqlError}; use crate::planner::qualified_name; use crate::types::{ - ArrayProjectPlan, ArraySlicePlan, DocumentIndexLookupPlan, HybridSearchPlan, - HybridSearchTriplePlan, LateralLoopPlan, LateralTopKPlan, RangeScanPlan, RecursiveScanPlan, - TextSearchPlan, TimeseriesScanPlan, + DocumentIndexLookupPlan, HybridSearchPlan, HybridSearchTriplePlan, LateralLoopPlan, + LateralTopKPlan, RangeScanPlan, RecursiveScanPlan, TextSearchPlan, TimeseriesScanPlan, }; use crate::types::{Projection, SortKey, SqlExpr, SqlPlan}; @@ -62,26 +61,6 @@ fn retain_sort_columns(input: &mut SqlPlan, sort_keys: &[SortKey]) -> Result &str { name.rsplit('.').next().unwrap_or(name) } -/// The first sort-key name an array body's attribute list does not emit. -fn first_missing_attr(attrs: &[String], sort_keys: &[SortKey]) -> Option { - sort_keys.iter().find_map(|key| match &key.expr { - SqlExpr::Column { name, .. } => { - if attrs.iter().any(|a| a.eq_ignore_ascii_case(name)) { - None - } else { - Some(name.clone()) - } - } - _ => None, - }) -} - /// The SELECT-list projection a plan narrows its rows with, when it has one. /// /// `None` means the variant emits whatever columns its source produced, so a diff --git a/nodedb-sql/src/planner/select/select_stmt.rs b/nodedb-sql/src/planner/select/select_stmt.rs index e6913c82b..7a2f684c8 100644 --- a/nodedb-sql/src/planner/select/select_stmt.rs +++ b/nodedb-sql/src/planner/select/select_stmt.rs @@ -4,6 +4,7 @@ use sqlparser::ast::{self, Select}; +use super::array_from::try_plan_array_from; use super::comma_lateral::try_plan_comma_lateral; use super::derived_from::try_plan_derived_from; use super::helpers::{convert_projection, convert_where_to_filters}; @@ -41,15 +42,10 @@ pub(super) fn plan_select( ) -> Result { // 0. Intercept array table-valued functions before catalog resolution // so a name like `ARRAY_SLICE` is not looked up as a collection. - if let Some(plan) = - crate::planner::array_fn::try_plan_array_table_fn(&select.from, catalog, temporal)? + if let Some(planned) = + try_plan_array_from(select, catalog, functions, temporal, tail, statement_output)? { - // `resolve_from` synthesizes a relation from the array's dims and - // attrs, so ORDER BY and the tail clauses resolve its columns. - return Ok(PlannedSelect { - plan, - scope: TableScope::resolve_from(catalog, functions, temporal, &select.from)?, - }); + return Ok(planned); } // 0.5. Derived FROM subquery: `FROM (SELECT ...) AS t`. @@ -213,6 +209,11 @@ pub(super) fn plan_select( let input_joins = subquery_joins .iter() .any(|sq| sq.join_type != JoinType::Cross); + let time_key = if matches!(table.info.engine, EngineType::Timeseries) { + catalog.timeseries_time_key(nodedb_types::DatabaseId::DEFAULT, &table.name)? + } else { + None + }; let mut plan = crate::planner::aggregate::plan_aggregate( crate::planner::aggregate::AggregateQuery { select, @@ -222,6 +223,7 @@ pub(super) fn plan_select( scope: &scope, temporal: &temporal, input_joins, + time_key, }, functions, )?; diff --git a/nodedb-sql/src/planner/select/where_search.rs b/nodedb-sql/src/planner/select/where_search.rs index 51f9cf682..ad748e012 100644 --- a/nodedb-sql/src/planner/select/where_search.rs +++ b/nodedb-sql/src/planner/select/where_search.rs @@ -23,6 +23,7 @@ use super::text_call::{TextCall, resolve_text_call}; use crate::error::{Result, SqlError}; use crate::functions::registry::{FunctionRegistry, SearchTrigger}; use crate::parser::normalize::normalize_ident; +use crate::planner::computed_tail; use crate::types::*; /// Default `top_k` placeholder used when a WHERE-derived search plan has no @@ -390,7 +391,7 @@ fn plan_spatial_from_where( detail: format!("invalid geometry in {name}: {e}"), } })?; - let issues = nodedb_spatial::validate::validate_geometry(&geometry); + let issues = nodedb_spatial::validate::validate_geometry(&geometry.geometry); if !issues.is_empty() { return Err(SqlError::InvalidFunction { detail: format!("invalid geometry in {name}: {}", issues.join("; ")), @@ -401,7 +402,7 @@ fn plan_spatial_from_where( } else { 0.0 }; - Ok(Some(SqlPlan::SpatialScan { + let scan = SqlPlan::SpatialScan { collection: table.name.clone(), field, predicate, @@ -411,6 +412,13 @@ fn plan_spatial_from_where( // No LIMIT reads every matching row; an explicit LIMIT overwrites // this in `apply_limit`. limit: usize::MAX, - projection: projection.to_vec(), - })) + projection: computed_tail::body_projection(projection), + }; + // The spatial scan evaluates no expression. A computed SELECT item runs + // in a tail over the matched rows. + Ok(Some(computed_tail::with_computed_tail( + scan, + projection, + computed_tail::TailBounds::default(), + ))) } diff --git a/nodedb-sql/src/resolver/array_tvf.rs b/nodedb-sql/src/resolver/array_tvf.rs index 7d74b3e60..cb55447a3 100644 --- a/nodedb-sql/src/resolver/array_tvf.rs +++ b/nodedb-sql/src/resolver/array_tvf.rs @@ -3,23 +3,36 @@ //! Array table-valued function resolution. //! //! An `ARRAY_*(...)` factor in a FROM clause names an array, not a -//! collection. This module synthesizes the relation its dims and attrs -//! expose so column references and equi-join keys against it resolve. +//! collection. This module synthesizes the relation the call's rows expose +//! so column references and equi-join keys against it resolve. use crate::error::{Result, SqlError}; use crate::parser::normalize::normalize_object_name_checked; +use crate::planner::array_fn::{ArrayTableFn, try_plan_array_table_fn}; use crate::resolver::columns::ResolvedTable; +use crate::temporal::TemporalScope; use crate::types::{ - ArrayCatalogView, CollectionInfo, ColumnInfo, EngineType, SqlCatalog, SqlDataType, + ArrayAggPlan, ArrayCatalogView, ArraySlicePlan, CollectionInfo, ColumnInfo, EngineType, + SqlCatalog, SqlDataType, }; -use crate::types_array::{ArrayAttrType, ArrayDimType}; +use crate::types_array::{ArrayAttrType, ArrayDimType, ArrayReducerAst}; +use nodedb_types::columnar::{FloatWidth, IntWidth}; + +/// The column an audit-log slice (`AS OF SYSTEM TIME NULL`) stamps on each +/// row: the version's system time in milliseconds. +pub const ARRAY_SYSTEM_TIME_COLUMN: &str = "_ts_system"; + +/// The flag of the summary row a temporal aggregate or elementwise read +/// appends: whether its system-time cutoff predates the retained history. +pub const ARRAY_HORIZON_COLUMN: &str = "truncated_before_horizon"; /// If `factor` is `ARRAY_*(name, ...)`, look up the array via the -/// catalog and build a `ResolvedTable` whose columns mirror the array's -/// dims + attrs. Returns `Ok(None)` for any non-array-TVF factor. +/// catalog and build a `ResolvedTable` whose columns are the ones the +/// call's rows carry. Returns `Ok(None)` for any non-array-TVF factor. pub(super) fn resolve_array_tvf( catalog: &dyn SqlCatalog, factor: &sqlparser::ast::TableFactor, + temporal: TemporalScope, ) -> Result> { let (fn_name, args, alias) = match factor { sqlparser::ast::TableFactor::Table { @@ -57,10 +70,18 @@ pub(super) fn resolve_array_tvf( name: array_name.clone(), })?; + let twj = sqlparser::ast::TableWithJoins { + relation: factor.clone(), + joins: Vec::new(), + }; + let Some(call) = try_plan_array_table_fn(std::slice::from_ref(&twj), catalog, temporal)? else { + return Ok(None); + }; + let info = CollectionInfo { name: view.name.clone(), engine: EngineType::Array, - columns: array_columns(&view), + columns: array_table_fn_columns(catalog, &call)?, primary_key: None, has_auto_tier: false, indexes: Vec::new(), @@ -77,35 +98,138 @@ pub(super) fn resolve_array_tvf( })) } -fn array_columns(view: &ArrayCatalogView) -> Vec { +/// The columns of the relation an array cell read exposes: every dim in +/// schema order, then each attr `attrs` names in that order, or every attr +/// in schema order when `attrs` is empty. +pub fn array_relation_columns(view: &ArrayCatalogView, attrs: &[String]) -> Vec { let mut cols = Vec::with_capacity(view.dims.len() + view.attrs.len()); for d in &view.dims { - cols.push(ColumnInfo { - name: d.name.clone(), - data_type: dim_type_to_sql(d.dtype), - nullable: false, - is_primary_key: false, - default: None, - raw_type: None, - int_width: None, - float_width: None, - }); + cols.push(column(&d.name, dim_type_to_sql(d.dtype), false)); } - for a in &view.attrs { - cols.push(ColumnInfo { - name: a.name.clone(), - data_type: attr_type_to_sql(a.dtype), - nullable: a.nullable, - is_primary_key: false, - default: None, - raw_type: None, - int_width: None, - float_width: None, - }); + if attrs.is_empty() { + for a in &view.attrs { + cols.push(attr_column(&a.name, a.dtype, a.nullable)); + } + } else { + for name in attrs { + if let Some(a) = view.attrs.iter().find(|a| &a.name == name) { + cols.push(attr_column(&a.name, a.dtype, a.nullable)); + } + } + } + cols +} + +/// The columns a slice's rows carry: its dims and attrs, plus +/// [`ARRAY_SYSTEM_TIME_COLUMN`] when it reads every system-time version. +pub fn array_slice_columns(view: &ArrayCatalogView, slice: &ArraySlicePlan) -> Vec { + let mut cols = array_relation_columns(view, &slice.attr_projection); + if matches!( + slice.temporal.system_time, + nodedb_types::SystemTimeScope::AllVersions + ) { + cols.push(column(ARRAY_SYSTEM_TIME_COLUMN, SqlDataType::Int64, false)); } cols } +/// The columns an aggregate's rows carry: `group`, typed as the group dim, +/// when it groups, then `result`, typed by the reducer. +/// +/// Refuses an attr the array does not hold, or a reducer the attr type cannot answer. +pub fn array_agg_columns(view: &ArrayCatalogView, agg: &ArrayAggPlan) -> Result> { + let mut cols = Vec::with_capacity(2); + if let Some(dim) = &agg.group_by_dim + && let Some(d) = view.dims.iter().find(|d| &d.name == dim) + { + cols.push(column("group", dim_type_to_sql(d.dtype), false)); + } + let attr = view + .attrs + .iter() + .find(|a| a.name == agg.attr) + .ok_or_else(|| SqlError::UnknownColumn { + table: view.name.clone(), + column: agg.attr.clone(), + })?; + cols.push(agg_result_column(agg.reducer, &attr.name, attr.dtype)?); + Ok(cols) +} + +/// Refuse a reducer that an attr of type `attr` cannot answer. +/// +/// The engine folds only numeric cells and skips every other cell. +/// `SUM`, `MEAN`, `MIN`, and `MAX` over a `STRING` or `BYTES` attr see no cell. +/// Their rows can only carry NULL, so the planner refuses them. +/// `COUNT` counts the non-null cells of every attr type. +pub fn check_agg_reducer( + reducer: ArrayReducerAst, + attr_name: &str, + attr: ArrayAttrType, +) -> Result<()> { + let answers = match reducer { + ArrayReducerAst::Count => true, + ArrayReducerAst::Sum + | ArrayReducerAst::Min + | ArrayReducerAst::Max + | ArrayReducerAst::Mean => attr.is_integer() || attr.is_float(), + }; + if answers { + Ok(()) + } else { + Err(SqlError::TypeMismatch { + detail: format!( + "ARRAY_AGG: {} cannot reduce {} attr '{attr_name}'; \ + only COUNT accepts a non-numeric attr", + reducer.sql_name(), + attr.sql_name() + ), + }) + } +} + +/// The columns the rows of one array table-function call carry. +/// +/// - A slice or project row holds the array's dims and the requested attrs. +/// - An aggregate row holds `result`, after `group` when it groups by a dim. +/// - An elementwise row holds the `coords` and `attrs` arrays of one cell. +pub fn array_table_fn_columns( + catalog: &dyn SqlCatalog, + call: &ArrayTableFn, +) -> Result> { + match call { + ArrayTableFn::Slice(slice) | ArrayTableFn::Project(slice) => { + Ok(array_slice_columns(&view_of(catalog, &slice.name)?, slice)) + } + ArrayTableFn::Agg(agg) => array_agg_columns(&view_of(catalog, &agg.name)?, agg), + ArrayTableFn::Elementwise(_) => Ok(vec![ + column("coords", SqlDataType::Json, false), + column("attrs", SqlDataType::Json, false), + ]), + } +} + +fn view_of(catalog: &dyn SqlCatalog, name: &str) -> Result { + catalog + .lookup_array(name) + .ok_or_else(|| SqlError::UnknownTable { + name: name.to_string(), + }) +} + +fn column(name: &str, data_type: SqlDataType, nullable: bool) -> ColumnInfo { + ColumnInfo { + name: name.to_string(), + data_type, + nullable, + is_primary_key: false, + default: None, + raw_type: None, + int_width: None, + float_width: None, + } +} + fn dim_type_to_sql(t: ArrayDimType) -> SqlDataType { match t { ArrayDimType::Int64 => SqlDataType::Int64, @@ -115,15 +239,57 @@ fn dim_type_to_sql(t: ArrayDimType) -> SqlDataType { } } -fn attr_type_to_sql(t: ArrayAttrType) -> SqlDataType { +/// The column of attr `name`, typed to its wire type: `UINT8` and `INT16` +/// are `int2`, `UINT16` and `INT32` are `int4`, `UINT32` and `INT64` are +/// `int8`, `FLOAT32` is `float4`, and `FLOAT64` is `float8`. An unsigned +/// type takes the next signed width, so every stored value fits. A width of +/// `None` keeps the base `int8` / `float8`. +fn attr_column(name: &str, t: ArrayAttrType, nullable: bool) -> ColumnInfo { + let (data_type, int_width, float_width) = attr_wire_type(t); + ColumnInfo { + int_width, + float_width, + ..column(name, data_type, nullable) + } +} + +fn attr_wire_type(t: ArrayAttrType) -> (SqlDataType, Option, Option) { match t { - ArrayAttrType::Int64 => SqlDataType::Int64, - ArrayAttrType::Float64 => SqlDataType::Float64, - ArrayAttrType::String => SqlDataType::String, - ArrayAttrType::Bytes => SqlDataType::Bytes, + ArrayAttrType::Int64 | ArrayAttrType::UInt32 => (SqlDataType::Int64, None, None), + ArrayAttrType::Int32 | ArrayAttrType::UInt16 => { + (SqlDataType::Int64, Some(IntWidth::I32), None) + } + ArrayAttrType::Int16 | ArrayAttrType::UInt8 => { + (SqlDataType::Int64, Some(IntWidth::I16), None) + } + ArrayAttrType::Float64 => (SqlDataType::Float64, None, None), + ArrayAttrType::Float32 => (SqlDataType::Float64, None, Some(FloatWidth::F32)), + ArrayAttrType::String => (SqlDataType::String, None, None), + ArrayAttrType::Bytes => (SqlDataType::Bytes, None, None), } } +/// The `result` column of an aggregate. `COUNT` and the `SUM` of an integer +/// attr are `int8`. `MIN` / `MAX` take the attr's own wire type. `MEAN` +/// and the `SUM` of a float attr are `float8`. +/// A reducer the attr type cannot answer is refused, see [`check_agg_reducer`]. +fn agg_result_column( + reducer: ArrayReducerAst, + attr_name: &str, + attr: ArrayAttrType, +) -> Result { + check_agg_reducer(reducer, attr_name, attr)?; + let float8 = || column("result", SqlDataType::Float64, true); + let int8 = || column("result", SqlDataType::Int64, true); + Ok(match reducer { + ArrayReducerAst::Count => int8(), + ArrayReducerAst::Mean => float8(), + ArrayReducerAst::Sum if attr.is_integer() => int8(), + ArrayReducerAst::Sum => float8(), + ArrayReducerAst::Min | ArrayReducerAst::Max => attr_column("result", attr, true), + }) +} + fn extract_string_literal_arg(arg: &sqlparser::ast::FunctionArg) -> Option { use sqlparser::ast::{Expr, FunctionArg, FunctionArgExpr, Value}; let expr = match arg { @@ -142,3 +308,68 @@ fn extract_string_literal_arg(arg: &sqlparser::ast::FunctionArg) -> Option None, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn narrow_attrs_announce_their_wire_widths() { + let cases = [ + (ArrayAttrType::UInt8, Some(IntWidth::I16), None), + (ArrayAttrType::Int16, Some(IntWidth::I16), None), + (ArrayAttrType::UInt16, Some(IntWidth::I32), None), + (ArrayAttrType::Int32, Some(IntWidth::I32), None), + (ArrayAttrType::UInt32, None, None), + (ArrayAttrType::Int64, None, None), + (ArrayAttrType::Float32, None, Some(FloatWidth::F32)), + (ArrayAttrType::Float64, None, None), + ]; + for (t, int_width, float_width) in cases { + let c = attr_column("v", t, true); + assert_eq!( + (c.int_width, c.float_width), + (int_width, float_width), + "{t:?}" + ); + } + } + + #[test] + fn aggregate_results_take_the_reducer_type() { + let result = |reducer, attr| agg_result_column(reducer, "v", attr).expect("answerable"); + let sum = result(ArrayReducerAst::Sum, ArrayAttrType::UInt16); + assert_eq!(sum.data_type, SqlDataType::Int64); + let count = result(ArrayReducerAst::Count, ArrayAttrType::String); + assert_eq!(count.data_type, SqlDataType::Int64); + let count_bytes = result(ArrayReducerAst::Count, ArrayAttrType::Bytes); + assert_eq!(count_bytes.data_type, SqlDataType::Int64); + let mean = result(ArrayReducerAst::Mean, ArrayAttrType::Int16); + assert_eq!(mean.data_type, SqlDataType::Float64); + let max = result(ArrayReducerAst::Max, ArrayAttrType::Float32); + assert_eq!(max.float_width, Some(FloatWidth::F32)); + let min = result(ArrayReducerAst::Min, ArrayAttrType::Int16); + assert_eq!(min.int_width, Some(IntWidth::I16)); + } + + #[test] + fn non_numeric_attrs_refuse_numeric_reducers() { + for attr in [ArrayAttrType::String, ArrayAttrType::Bytes] { + for reducer in [ + ArrayReducerAst::Sum, + ArrayReducerAst::Mean, + ArrayReducerAst::Min, + ArrayReducerAst::Max, + ] { + match agg_result_column(reducer, "label", attr) { + Err(SqlError::TypeMismatch { detail }) => { + assert!(detail.contains(reducer.sql_name()), "{detail}"); + assert!(detail.contains(attr.sql_name()), "{detail}"); + assert!(detail.contains("'label'"), "{detail}"); + } + other => panic!("{reducer:?} over {attr:?} must be refused: {other:?}"), + } + } + } + } +} diff --git a/nodedb-sql/src/resolver/columns.rs b/nodedb-sql/src/resolver/columns.rs index 575b3c52f..7d28fefd1 100644 --- a/nodedb-sql/src/resolver/columns.rs +++ b/nodedb-sql/src/resolver/columns.rs @@ -400,9 +400,11 @@ impl TableScope { factor: &sqlparser::ast::TableFactor, ) -> Result<()> { // ARRAY_*(...) table-valued function: synthesize a ResolvedTable - // from the array's dim+attr schema so equi-join keys against the - // TVF's output rows resolve. - if let Some(resolved) = crate::resolver::array_tvf::resolve_array_tvf(catalog, factor)? { + // from the columns the call's rows carry so equi-join keys against + // the TVF's output rows resolve. + if let Some(resolved) = + crate::resolver::array_tvf::resolve_array_tvf(catalog, factor, temporal)? + { self.add(resolved)?; return Ok(()); } diff --git a/nodedb-sql/src/types/mod.rs b/nodedb-sql/src/types/mod.rs index 22e24a072..6beade484 100644 --- a/nodedb-sql/src/types/mod.rs +++ b/nodedb-sql/src/types/mod.rs @@ -13,18 +13,18 @@ pub mod query; pub use collection::{CollectionInfo, ColumnInfo, IndexSpec, IndexState}; pub use filter::{CompareOp, Filter, FilterExpr}; pub use plan::{ - AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan, - CreateArrayPlan, CreateIndexPlan, CtePlan, DeleteArrayPlan, DocumentIndexLookupPlan, - DropIndexPlan, HybridSearchPlan, HybridSearchTriplePlan, InsertArrayPlan, InsertPlan, - KvInsertPlan, LateralLoopPlan, LateralTopKPlan, MergePlan, RangeScanPlan, RecursiveScanPlan, - RecursiveValuePlan, TextScoreColumn, TextSearchPlan, TextSearchShape, TimeseriesIngestPlan, - TimeseriesScanPlan, UpsertPlan, VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, - VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, + AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan, CreateArrayPlan, + CreateIndexPlan, CtePlan, DeleteArrayPlan, DocumentIndexLookupPlan, DropIndexPlan, + HybridSearchPlan, HybridSearchTriplePlan, InsertArrayPlan, InsertPlan, KvInsertPlan, + LateralLoopPlan, LateralTopKPlan, MergePlan, RangeScanPlan, RecursiveScanPlan, + RecursiveValuePlan, TextScoreColumn, TextSearchPlan, TextSearchShape, TimeseriesAggOutput, + TimeseriesIngestPlan, TimeseriesScanPlan, UpsertPlan, VectorPrimaryDeletePlan, + VectorPrimaryInsertPlan, VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, }; pub use plan::{ ArrayPrefilter, DistanceMetric, KvInsertIntent, MergeClauseKind, MergePlanAction, - MergePlanClause, PlanCacheEligibility, SqlPlan, VectorAnnOptions, VectorPrimaryInsertIntent, - VectorPrimaryRow, VectorQuantization, WriteRoute, + MergePlanClause, PlanCacheEligibility, SqlPlan, TIMESERIES_BUCKET_KEY, VectorAnnOptions, + VectorPrimaryInsertIntent, VectorPrimaryRow, VectorQuantization, WriteRoute, }; pub use query::{ AggOutputSlot, AggregateExpr, EngineType, JoinType, Projection, SortKey, SpatialPredicate, diff --git a/nodedb-sql/src/types/plan/cacheability.rs b/nodedb-sql/src/types/plan/cacheability.rs index 077943825..94b4657ca 100644 --- a/nodedb-sql/src/types/plan/cacheability.rs +++ b/nodedb-sql/src/types/plan/cacheability.rs @@ -174,7 +174,6 @@ impl SqlPlan { | Self::InsertArray { .. } | Self::DeleteArray { .. } | Self::ArraySlice { .. } - | Self::ArrayProject { .. } | Self::ArrayAgg { .. } | Self::ArrayElementwise { .. } | Self::ArrayFlush { .. } diff --git a/nodedb-sql/src/types/plan/mod.rs b/nodedb-sql/src/types/plan/mod.rs index e27478f76..42672cc06 100644 --- a/nodedb-sql/src/types/plan/mod.rs +++ b/nodedb-sql/src/types/plan/mod.rs @@ -21,14 +21,14 @@ pub use merge_types::{MergeClauseKind, MergePlanAction, MergePlanClause}; pub use row_types::{KvInsertIntent, VectorPrimaryInsertIntent, VectorPrimaryRow, WriteRoute}; /// Per-family payloads of the larger `SqlPlan` variants. pub use variants::{ - AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan, - CreateArrayPlan, CreateIndexPlan, CtePlan, DeleteArrayPlan, DocumentIndexLookupPlan, - DropIndexPlan, HybridSearchPlan, HybridSearchTriplePlan, InsertArrayPlan, InsertPlan, - KvInsertPlan, LateralLoopPlan, LateralTopKPlan, MergePlan, RangeScanPlan, RecursiveScanPlan, - RecursiveValuePlan, TextScoreColumn, TextSearchPlan, TextSearchShape, TimeseriesIngestPlan, - TimeseriesScanPlan, UpsertPlan, VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, - VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, + AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan, CreateArrayPlan, + CreateIndexPlan, CtePlan, DeleteArrayPlan, DocumentIndexLookupPlan, DropIndexPlan, + HybridSearchPlan, HybridSearchTriplePlan, InsertArrayPlan, InsertPlan, KvInsertPlan, + LateralLoopPlan, LateralTopKPlan, MergePlan, RangeScanPlan, RecursiveScanPlan, + RecursiveValuePlan, TextScoreColumn, TextSearchPlan, TextSearchShape, TimeseriesAggOutput, + TimeseriesIngestPlan, TimeseriesScanPlan, UpsertPlan, VectorPrimaryDeletePlan, + VectorPrimaryInsertPlan, VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, }; -pub use variants::{DistanceMetric, SqlPlan}; +pub use variants::{DistanceMetric, SqlPlan, TIMESERIES_BUCKET_KEY}; pub use vector_opts::{ArrayPrefilter, VectorAnnOptions, VectorQuantization}; pub use volatility_scan::expr_is_volatile; diff --git a/nodedb-sql/src/types/plan/search_cells.rs b/nodedb-sql/src/types/plan/search_cells.rs index 60880bec1..63669ab4d 100644 --- a/nodedb-sql/src/types/plan/search_cells.rs +++ b/nodedb-sql/src/types/plan/search_cells.rs @@ -75,6 +75,7 @@ mod tests { field: "t".into(), query_entries: vec![(3, 1.0)], top_k: 1, + filters: Vec::new(), projection: Vec::new(), } .carries_search_cells() diff --git a/nodedb-sql/src/types/plan/variant_name.rs b/nodedb-sql/src/types/plan/variant_name.rs index 9e0781ec9..61f0e42ac 100644 --- a/nodedb-sql/src/types/plan/variant_name.rs +++ b/nodedb-sql/src/types/plan/variant_name.rs @@ -48,7 +48,6 @@ impl SqlPlan { SqlPlan::InsertArray { .. } => "InsertArray", SqlPlan::DeleteArray { .. } => "DeleteArray", SqlPlan::ArraySlice { .. } => "ArraySlice", - SqlPlan::ArrayProject { .. } => "ArrayProject", SqlPlan::ArrayAgg { .. } => "ArrayAgg", SqlPlan::ArrayElementwise { .. } => "ArrayElementwise", SqlPlan::ArrayFlush { .. } => "ArrayFlush", diff --git a/nodedb-sql/src/types/plan/variants/array.rs b/nodedb-sql/src/types/plan/variants/array.rs index 5864c1080..4a3790bb4 100644 --- a/nodedb-sql/src/types/plan/variants/array.rs +++ b/nodedb-sql/src/types/plan/variants/array.rs @@ -62,14 +62,6 @@ pub struct ArraySlicePlan { pub temporal: TemporalScope, } -/// Payload of [`SqlPlan::ArrayProject`](crate::types::SqlPlan::ArrayProject). -#[derive(Debug, Clone)] -pub struct ArrayProjectPlan { - pub name: String, - /// Attribute names. Must be non-empty. - pub attr_projection: Vec, -} - /// Payload of [`SqlPlan::ArrayAgg`](crate::types::SqlPlan::ArrayAgg). #[derive(Debug, Clone)] pub struct ArrayAggPlan { @@ -78,6 +70,10 @@ pub struct ArrayAggPlan { pub reducer: types_array::ArrayReducerAst, /// `None` = scalar fold; `Some(name)` = group by that dim. pub group_by_dim: Option, + /// Coordinate box the aggregate reads. No dim range = the whole array. + pub slice: types_array::ArraySliceAst, + /// Polygon zone inside the box. `None` = the whole box. + pub zone: Option, /// Bitemporal qualifier. When both axes are `None` / `Any`, the Data /// Plane aggregates against the live (current) state — the default /// fast path. Populated from `AS OF SYSTEM TIME` / `AS OF VALID TIME`. @@ -91,4 +87,7 @@ pub struct ArrayElementwisePlan { pub right: String, pub op: types_array::ArrayBinaryOpAst, pub attr: String, + /// Bitemporal qualifier, applied to both operands. Populated from + /// `AS OF SYSTEM TIME` / `AS OF VALID TIME`. + pub temporal: TemporalScope, } diff --git a/nodedb-sql/src/types/plan/variants/mod.rs b/nodedb-sql/src/types/plan/variants/mod.rs index d7240b902..b1f981b02 100644 --- a/nodedb-sql/src/types/plan/variants/mod.rs +++ b/nodedb-sql/src/types/plan/variants/mod.rs @@ -17,8 +17,8 @@ mod vector_primary; mod writes; pub use array::{ - AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan, - CreateArrayPlan, DeleteArrayPlan, InsertArrayPlan, + AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan, CreateArrayPlan, + DeleteArrayPlan, InsertArrayPlan, }; pub use cte::CtePlan; pub use hybrid::{HybridSearchPlan, HybridSearchTriplePlan}; @@ -29,7 +29,9 @@ pub use merge::MergePlan; pub use plan::{DistanceMetric, SqlPlan}; pub use recursive::{RecursiveScanPlan, RecursiveValuePlan}; pub use text::{TextScoreColumn, TextSearchPlan, TextSearchShape}; -pub use timeseries::{TimeseriesIngestPlan, TimeseriesScanPlan}; +pub use timeseries::{ + TIMESERIES_BUCKET_KEY, TimeseriesAggOutput, TimeseriesIngestPlan, TimeseriesScanPlan, +}; pub use vector_primary::{ VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, diff --git a/nodedb-sql/src/types/plan/variants/plan.rs b/nodedb-sql/src/types/plan/variants/plan.rs index 12369b8a9..ea8ef713d 100644 --- a/nodedb-sql/src/types/plan/variants/plan.rs +++ b/nodedb-sql/src/types/plan/variants/plan.rs @@ -16,8 +16,8 @@ use crate::types::query::{ use super::super::vector_opts::{ArrayPrefilter, VectorAnnOptions}; use super::array::{ - AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan, - CreateArrayPlan, DeleteArrayPlan, InsertArrayPlan, + AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan, CreateArrayPlan, + DeleteArrayPlan, InsertArrayPlan, }; use super::cte::CtePlan; use super::hybrid::{HybridSearchPlan, HybridSearchTriplePlan}; @@ -279,6 +279,9 @@ pub enum SqlPlan { field: String, query_entries: Vec<(u32, f32)>, top_k: usize, + /// The statement's residual `WHERE`. It narrows the candidates + /// before the `top_k` cut. + filters: Vec, /// Resolved SELECT target list, for output-schema derivation. projection: Vec, }, @@ -296,7 +299,9 @@ pub enum SqlPlan { collection: String, field: String, predicate: SpatialPredicate, - query_geometry: nodedb_types::geometry::Geometry, + /// The query geometry and its SRID. A scan whose index is in another + /// SRID is refused, never transformed implicitly. + query_geometry: nodedb_types::geometry::SridGeometry, distance_meters: f64, attribute_filters: Vec, limit: usize, @@ -394,10 +399,9 @@ pub enum SqlPlan { InsertArray(InsertArrayPlan), /// `DELETE FROM ARRAY WHERE COORDS IN ((...), (...))`. DeleteArray(DeleteArrayPlan), - /// `SELECT * FROM ARRAY_SLICE(name, {dim:[lo,hi],..}, [attrs], limit)`. + /// `SELECT * FROM ARRAY_SLICE(name, {dim:[lo,hi],..}, [attrs], limit)`, + /// and `ARRAY_PROJECT(name, [attrs])` as a slice over the whole domain. ArraySlice(ArraySlicePlan), - /// `SELECT * FROM ARRAY_PROJECT(name, [attrs])`. - ArrayProject(ArrayProjectPlan), /// `SELECT * FROM ARRAY_AGG(name, attr, reducer [, group_by_dim])`. ArrayAgg(ArrayAggPlan), /// `SELECT * FROM ARRAY_ELEMENTWISE(left, right, op, attr)`. diff --git a/nodedb-sql/src/types/plan/variants/timeseries.rs b/nodedb-sql/src/types/plan/variants/timeseries.rs index 5d68c2cf9..93b05ae2f 100644 --- a/nodedb-sql/src/types/plan/variants/timeseries.rs +++ b/nodedb-sql/src/types/plan/variants/timeseries.rs @@ -5,7 +5,7 @@ use crate::temporal::TemporalScope; use crate::types::filter::Filter; use crate::types::query::{AggregateExpr, Projection, SortKey}; -use crate::types_expr::SqlValue; +use crate::types_expr::{SqlExpr, SqlValue}; /// Payload of [`SqlPlan::TimeseriesScan`](crate::types::SqlPlan::TimeseriesScan). #[derive(Debug, Clone)] @@ -29,6 +29,67 @@ pub struct TimeseriesScanPlan { /// on collections created `WITH BITEMPORAL`; `TimeseriesRules::plan_scan` /// rejects temporal scopes otherwise. pub temporal: TemporalScope, + /// The SELECT-list output columns of an aggregating scan, in SELECT-list + /// order. Empty on a raw scan, and on a scan planned without a SELECT list. + pub output: Vec, + /// The GROUP BY keys as planned, the `time_bucket` key among them. A tail + /// that evaluates an item over the group rows matches its sub-expressions + /// against these keys. Empty on a raw scan. + pub group_keys: Vec, +} + +/// One SELECT-list output column of an aggregating timeseries scan. +/// +/// The scan encodes each group row under fixed keys: `bucket` for the time +/// bucket, the bare column name for a GROUP BY column, and the canonical key +/// for an aggregate. This entry names which of those keys a column reads and +/// the name the client sees. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TimeseriesAggOutput { + /// The `time_bucket` boundary. + Bucket { + /// The SELECT-list alias, when the item has one. + alias: Option, + /// The column the `time_bucket` call buckets, when it names one. + source: Option, + }, + /// A plain GROUP BY column. + GroupColumn { + /// The bare column name the scan emits the key under. + name: String, + /// The SELECT-list alias, when the item has one. + alias: Option, + }, + /// `aggregates[index]`, shown under its `AggregateExpr::alias`. + Aggregate(usize), +} + +/// The key a native grouped timeseries scan encodes the time bucket under. +pub const TIMESERIES_BUCKET_KEY: &str = "bucket"; + +impl TimeseriesAggOutput { + /// The `(encoded key, output name)` of this column over `aggregates`. + /// + /// `None` for an `Aggregate` index past `aggregates`. + pub fn key_and_name(&self, aggregates: &[AggregateExpr]) -> Option<(String, String)> { + match self { + Self::Bucket { alias, .. } => Some(( + TIMESERIES_BUCKET_KEY.to_string(), + alias + .clone() + .unwrap_or_else(|| TIMESERIES_BUCKET_KEY.to_string()), + )), + Self::GroupColumn { name, alias } => { + Some((name.clone(), alias.clone().unwrap_or_else(|| name.clone()))) + } + Self::Aggregate(index) => aggregates.get(*index).map(|agg| { + ( + crate::planner::agg_naming::aggregate_output_key(agg), + agg.alias.clone(), + ) + }), + } + } } /// Payload of [`SqlPlan::TimeseriesIngest`](crate::types::SqlPlan::TimeseriesIngest). diff --git a/nodedb-sql/src/types_array.rs b/nodedb-sql/src/types_array.rs index 20c603deb..2c3811bff 100644 --- a/nodedb-sql/src/types_array.rs +++ b/nodedb-sql/src/types_array.rs @@ -23,6 +23,61 @@ pub enum ArrayAttrType { Float64, String, Bytes, + UInt8, + Int16, + UInt16, + Int32, + UInt32, + Float32, +} + +impl ArrayAttrType { + /// Whether the type holds integers. + pub fn is_integer(self) -> bool { + match self { + ArrayAttrType::Int64 + | ArrayAttrType::UInt8 + | ArrayAttrType::Int16 + | ArrayAttrType::UInt16 + | ArrayAttrType::Int32 + | ArrayAttrType::UInt32 => true, + ArrayAttrType::Float64 + | ArrayAttrType::Float32 + | ArrayAttrType::String + | ArrayAttrType::Bytes => false, + } + } + + /// Whether the type holds floats. + pub fn is_float(self) -> bool { + match self { + ArrayAttrType::Float64 | ArrayAttrType::Float32 => true, + ArrayAttrType::Int64 + | ArrayAttrType::UInt8 + | ArrayAttrType::Int16 + | ArrayAttrType::UInt16 + | ArrayAttrType::Int32 + | ArrayAttrType::UInt32 + | ArrayAttrType::String + | ArrayAttrType::Bytes => false, + } + } + + /// The canonical DDL name of the type. + pub fn sql_name(self) -> &'static str { + match self { + ArrayAttrType::Int64 => "INT64", + ArrayAttrType::Float64 => "FLOAT64", + ArrayAttrType::String => "STRING", + ArrayAttrType::Bytes => "BYTES", + ArrayAttrType::UInt8 => "UINT8", + ArrayAttrType::Int16 => "INT16", + ArrayAttrType::UInt16 => "UINT16", + ArrayAttrType::Int32 => "INT32", + ArrayAttrType::UInt32 => "UINT32", + ArrayAttrType::Float32 => "FLOAT32", + } + } } /// Cell-order strategy AST mirror. @@ -119,6 +174,18 @@ pub struct ArraySliceAst { pub dim_ranges: Vec, } +/// `ARRAY_AGG(...)` polygon zone over two numeric dims. +/// +/// `polygons` holds one entry per member polygon: its exterior ring, then +/// its holes. Each ring is closed and lists `[x, y]` positions, `x` on +/// `x_dim` and `y` on `y_dim`. +#[derive(Debug, Clone, PartialEq)] +pub struct ArrayZoneAst { + pub polygons: Vec>>, + pub x_dim: String, + pub y_dim: String, +} + /// Reducer carried on `SqlPlan::ArrayAgg`. Mirrors /// `crate::bridge::physical_plan::ArrayReducer` without the bridge dep. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -144,6 +211,17 @@ impl ArrayReducerAst { _ => None, } } + + /// The uppercase SQL name of the reducer. + pub fn sql_name(self) -> &'static str { + match self { + ArrayReducerAst::Sum => "SUM", + ArrayReducerAst::Count => "COUNT", + ArrayReducerAst::Min => "MIN", + ArrayReducerAst::Max => "MAX", + ArrayReducerAst::Mean => "MEAN", + } + } } /// Pairwise op carried on `SqlPlan::ArrayElementwise`. Mirrors diff --git a/nodedb-sql/src/visitor/plan_visitor/args.rs b/nodedb-sql/src/visitor/plan_visitor/args.rs index 23c31efac..e76c108d6 100644 --- a/nodedb-sql/src/visitor/plan_visitor/args.rs +++ b/nodedb-sql/src/visitor/plan_visitor/args.rs @@ -10,13 +10,16 @@ use crate::temporal::TemporalScope; use crate::types::SqlPlan; use crate::types::filter::Filter; use crate::types::plan::{ - ArrayPrefilter, MergePlanClause, TextScoreColumn, TextSearchShape, VectorAnnOptions, - VectorPrimaryInsertIntent, VectorPrimaryRow, WriteRoute, + ArrayPrefilter, MergePlanClause, TextScoreColumn, TextSearchShape, TimeseriesAggOutput, + VectorAnnOptions, VectorPrimaryInsertIntent, VectorPrimaryRow, WriteRoute, }; use crate::types::query::{ AggregateExpr, EngineType, JoinType, Projection, SortKey, SpatialPredicate, WindowSpec, }; -use crate::types_array::{ArrayAttrAst, ArrayCellOrderAst, ArrayDimAst, ArrayTileOrderAst}; +use crate::types_array::{ + ArrayAttrAst, ArrayCellOrderAst, ArrayDimAst, ArrayReducerAst, ArraySliceAst, + ArrayTileOrderAst, ArrayZoneAst, +}; use crate::types_expr::{SqlExpr, SqlPayloadAtom, SqlValue}; use nodedb_types::vector_distance::DistanceMetric; @@ -123,6 +126,9 @@ pub struct AggregateVisitArgs<'a> { pub limit: usize, pub grouping_sets: Option<&'a [Vec]>, pub sort_keys: &'a [SortKey], + /// SELECT-list output alias of each GROUP BY key, parallel to `group_by`. + /// Empty, or a `None` entry, keeps the key's own row name. + pub group_by_aliases: &'a [Option], } /// Parameters for [`super::trait_def::PlanVisitor::timeseries_scan`]. @@ -141,6 +147,8 @@ pub struct TimeseriesScanVisitArgs<'a> { pub sort_keys: &'a [SortKey], pub tiered: bool, pub temporal: &'a TemporalScope, + /// SELECT-list output columns of an aggregating scan. Empty on a raw scan. + pub output: &'a [TimeseriesAggOutput], } /// Parameters for [`super::trait_def::PlanVisitor::vector_search`]. @@ -214,7 +222,7 @@ pub struct SpatialScanVisitArgs<'a> { pub collection: &'a str, pub field: &'a str, pub predicate: &'a SpatialPredicate, - pub query_geometry: &'a nodedb_types::geometry::Geometry, + pub query_geometry: &'a nodedb_types::geometry::SridGeometry, pub distance_meters: f64, pub attribute_filters: &'a [Filter], pub limit: usize, @@ -256,6 +264,19 @@ pub struct CreateArrayVisitArgs<'a> { pub minimum_audit_retain_ms: Option, } +/// Parameters for [`super::trait_def::PlanVisitor::array_agg`]. +pub struct ArrayAggVisitArgs<'a> { + pub name: &'a str, + pub attr: &'a str, + pub reducer: ArrayReducerAst, + pub group_by_dim: Option<&'a str>, + /// Coordinate box the aggregate reads. No dim range = the whole array. + pub slice: &'a ArraySliceAst, + /// Polygon zone inside the box. `None` = the whole box. + pub zone: Option<&'a ArrayZoneAst>, + pub temporal: &'a TemporalScope, +} + /// Parameters for [`super::trait_def::PlanVisitor::merge`]. pub struct MergeVisitArgs<'a> { pub target: &'a str, diff --git a/nodedb-sql/src/visitor/plan_visitor/dispatch.rs b/nodedb-sql/src/visitor/plan_visitor/dispatch.rs index 9f36cf6d4..95108409a 100644 --- a/nodedb-sql/src/visitor/plan_visitor/dispatch.rs +++ b/nodedb-sql/src/visitor/plan_visitor/dispatch.rs @@ -216,7 +216,7 @@ pub fn dispatch(visitor: &mut V, plan: &SqlPlan) -> Result(visitor: &mut V, plan: &SqlPlan) -> Result(visitor: &mut V, plan: &SqlPlan) -> Result visitor.timeseries_scan(TimeseriesScanVisitArgs { collection, time_range: *time_range, @@ -260,6 +264,7 @@ pub fn dispatch(visitor: &mut V, plan: &SqlPlan) -> Result(visitor: &mut V, plan: &SqlPlan) -> Result visitor.sparse_search(collection, field, query_entries, *top_k), + } => visitor.sparse_search(collection, field, query_entries, *top_k, filters), SqlPlan::TextSearch(TextSearchPlan { collection, shape, diff --git a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs index 1f049bfcc..5f7566bf7 100644 --- a/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs +++ b/nodedb-sql/src/visitor/plan_visitor/dispatch_rest.rs @@ -8,15 +8,16 @@ //! satisfy the compiler that this match is exhaustive over the full //! `SqlPlan` enum without repeating their field lists. use super::args::{ - CreateArrayVisitArgs, LateralLoopVisitArgs, LateralTopKVisitArgs, MergeVisitArgs, - VectorPrimaryDeleteVisitArgs, VectorPrimaryInsertVisitArgs, VectorPrimaryUpdateVisitArgs, + ArrayAggVisitArgs, CreateArrayVisitArgs, LateralLoopVisitArgs, LateralTopKVisitArgs, + MergeVisitArgs, VectorPrimaryDeleteVisitArgs, VectorPrimaryInsertVisitArgs, + VectorPrimaryUpdateVisitArgs, }; use super::trait_def::PlanVisitor; use crate::types::SqlPlan; use crate::types::{ - AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArrayProjectPlan, ArraySlicePlan, - CreateArrayPlan, CreateIndexPlan, CtePlan, DeleteArrayPlan, DropIndexPlan, InsertArrayPlan, - LateralLoopPlan, LateralTopKPlan, MergePlan, VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, + AlterArrayPlan, ArrayAggPlan, ArrayElementwisePlan, ArraySlicePlan, CreateArrayPlan, + CreateIndexPlan, CtePlan, DeleteArrayPlan, DropIndexPlan, InsertArrayPlan, LateralLoopPlan, + LateralTopKPlan, MergePlan, VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, VectorPrimaryTruncatePlan, VectorPrimaryUpdatePlan, }; @@ -86,23 +87,30 @@ pub(super) fn dispatch_rest( limit, temporal, }) => visitor.array_slice(name, slice, attr_projection, *limit, temporal), - SqlPlan::ArrayProject(ArrayProjectPlan { - name, - attr_projection, - }) => visitor.array_project(name, attr_projection), SqlPlan::ArrayAgg(ArrayAggPlan { name, attr, reducer, group_by_dim, + slice, + zone, temporal, - }) => visitor.array_agg(name, attr, reducer, group_by_dim.as_deref(), temporal), + }) => visitor.array_agg(ArrayAggVisitArgs { + name, + attr, + reducer: *reducer, + group_by_dim: group_by_dim.as_deref(), + slice, + zone: zone.as_ref(), + temporal, + }), SqlPlan::ArrayElementwise(ArrayElementwisePlan { left, right, op, attr, - }) => visitor.array_elementwise(left, right, *op, attr), + temporal, + }) => visitor.array_elementwise(left, right, *op, attr, temporal), SqlPlan::ArrayFlush { name } => visitor.array_flush(name), SqlPlan::ArrayCompact { name } => visitor.array_compact(name), SqlPlan::Merge(MergePlan { diff --git a/nodedb-sql/src/visitor/plan_visitor/trait_def.rs b/nodedb-sql/src/visitor/plan_visitor/trait_def.rs index 0ca912d69..166596363 100644 --- a/nodedb-sql/src/visitor/plan_visitor/trait_def.rs +++ b/nodedb-sql/src/visitor/plan_visitor/trait_def.rs @@ -6,7 +6,7 @@ //! bundled params struct from [`super::args`] instead of raw positional arguments. use super::args::{ - AggregateVisitArgs, CreateArrayVisitArgs, DocumentIndexLookupVisitArgs, + AggregateVisitArgs, ArrayAggVisitArgs, CreateArrayVisitArgs, DocumentIndexLookupVisitArgs, HybridSearchTripleVisitArgs, HybridSearchVisitArgs, InsertVisitArgs, JoinVisitArgs, LateralLoopVisitArgs, LateralTopKVisitArgs, MergeVisitArgs, RecursiveScanVisitArgs, RecursiveValueVisitArgs, ScanVisitArgs, SpatialScanVisitArgs, SubqueryVisitArgs, @@ -19,9 +19,7 @@ use crate::types::SqlPlan; use crate::types::filter::Filter; use crate::types::plan::KvInsertIntent; use crate::types::query::EngineType; -use crate::types_array::{ - ArrayBinaryOpAst, ArrayCoordLiteral, ArrayInsertRow, ArrayReducerAst, ArraySliceAst, -}; +use crate::types_array::{ArrayBinaryOpAst, ArrayCoordLiteral, ArrayInsertRow, ArraySliceAst}; use crate::types_expr::{SqlExpr, SqlValue}; /// Executor parity contract: every [`SqlPlan`] variant must be handled. @@ -162,13 +160,15 @@ pub trait PlanVisitor { ef_search: usize, ) -> Result; - /// Handle [`SqlPlan::SparseSearch`]. + /// Handle [`SqlPlan::SparseSearch`]. `filters` is the statement's + /// residual `WHERE`, applied before the `top_k` cut. fn sparse_search( &mut self, collection: &str, field: &str, query_entries: &[(u32, f32)], top_k: usize, + filters: &[Filter], ) -> Result; /// Handle [`SqlPlan::TextSearch`]. @@ -270,22 +270,8 @@ pub trait PlanVisitor { temporal: &TemporalScope, ) -> Result; - /// Handle [`SqlPlan::ArrayProject`]. - fn array_project( - &mut self, - name: &str, - attr_projection: &[String], - ) -> Result; - /// Handle [`SqlPlan::ArrayAgg`]. - fn array_agg( - &mut self, - name: &str, - attr: &str, - reducer: &ArrayReducerAst, - group_by_dim: Option<&str>, - temporal: &TemporalScope, - ) -> Result; + fn array_agg(&mut self, args: ArrayAggVisitArgs<'_>) -> Result; /// Handle [`SqlPlan::ArrayElementwise`]. fn array_elementwise( @@ -294,6 +280,7 @@ pub trait PlanVisitor { right: &str, op: ArrayBinaryOpAst, attr: &str, + temporal: &TemporalScope, ) -> Result; /// Handle [`SqlPlan::ArrayFlush`]. diff --git a/nodedb-strict/src/decode/value.rs b/nodedb-strict/src/decode/value.rs index ad5770ef9..5f62f03fd 100644 --- a/nodedb-strict/src/decode/value.rs +++ b/nodedb-strict/src/decode/value.rs @@ -87,13 +87,12 @@ pub(super) fn decode_variable_value(col_type: &ColumnType, raw: &[u8]) -> Value Value::String(std::str::from_utf8(raw).unwrap_or_default().to_string()) } ColumnType::Bytes => Value::Bytes(raw.to_vec()), - ColumnType::Geometry => { - if let Ok(geometry) = sonic_rs::from_slice::(raw) { - Value::Geometry(geometry) - } else { - Value::String(std::str::from_utf8(raw).unwrap_or_default().to_string()) - } - } + // GeoJSON with an optional `crs` member, which carries the SRID. + // Any other text (WKT) stays a string. + ColumnType::Geometry => match nodedb_types::geometry::parse_geojson_slice(raw) { + Ok(geometry) => Value::Geometry(geometry), + Err(_) => Value::String(std::str::from_utf8(raw).unwrap_or_default().to_string()), + }, ColumnType::Json => match nodedb_types::value_from_msgpack(raw) { Ok(value) => value, Err(error) => { @@ -128,6 +127,22 @@ fn decode_typed_msgpack(raw: &[u8], accepts: impl FnOnce(&Value) -> bool) -> Val mod tests { use super::*; + #[test] + fn projected_geometry_decodes_with_its_srid() { + use nodedb_types::geometry::{Geometry, Srid, SridGeometry}; + let g = SridGeometry::new(Srid::Utm48N, Geometry::point(134_445.5, 349_624.7)); + let raw = sonic_rs::to_vec(&g).unwrap(); + assert_eq!( + decode_variable_value(&ColumnType::Geometry, &raw), + Value::Geometry(g) + ); + let plain = sonic_rs::to_vec(&SridGeometry::wgs84(Geometry::point(1.0, 2.0))).unwrap(); + assert!( + !String::from_utf8(plain).unwrap().contains("crs"), + "4326 writes no crs member" + ); + } + #[test] fn vector_dimension_larger_than_raw_bytes_returns_null_without_reserving() { assert!(matches!( diff --git a/nodedb-strict/src/encode/value.rs b/nodedb-strict/src/encode/value.rs index 40543c30a..4aa19b338 100644 --- a/nodedb-strict/src/encode/value.rs +++ b/nodedb-strict/src/encode/value.rs @@ -140,6 +140,7 @@ pub(super) fn encode_variable( (ColumnType::Bytes | ColumnType::SparseVector, Value::Bytes(b)) => { var_data.extend_from_slice(b); } + // GeoJSON, with a `crs` member for any SRID other than 4326. (ColumnType::Geometry, Value::Geometry(g)) => { let json = sonic_rs::to_vec(g).map_err(|e| { invalid_value(col, format!("geometry does not serialize to GeoJSON: {e}")) diff --git a/nodedb-test-support/src/cluster_harness/node/inspect/timeseries.rs b/nodedb-test-support/src/cluster_harness/node/inspect/timeseries.rs index eb646a5fe..589269a0f 100644 --- a/nodedb-test-support/src/cluster_harness/node/inspect/timeseries.rs +++ b/nodedb-test-support/src/cluster_harness/node/inspect/timeseries.rs @@ -56,6 +56,7 @@ impl TestClusterNode { valid_at_ms: None, having: Vec::new(), tiering: None, + output_columns: Vec::new(), }), deadline: std::time::Instant::now() + std::time::Duration::from_secs(10), priority: Priority::Normal, diff --git a/nodedb-test-support/src/cluster_harness/node/lifecycle/teardown.rs b/nodedb-test-support/src/cluster_harness/node/lifecycle/teardown.rs index a8df19f71..17e9a06b1 100644 --- a/nodedb-test-support/src/cluster_harness/node/lifecycle/teardown.rs +++ b/nodedb-test-support/src/cluster_harness/node/lifecycle/teardown.rs @@ -16,6 +16,33 @@ impl TestClusterNode { } } + /// Run a simple query and return every row keyed by column name. A NULL + /// column reads as the empty string. Returns an error message on SQL + /// error. + pub async fn query_named_rows( + &self, + sql: &str, + ) -> Result>, String> { + let msgs = self + .client + .simple_query(sql) + .await + .map_err(|e| pg_error_detail(&e))?; + Ok(msgs + .into_iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => Some( + row.columns() + .iter() + .enumerate() + .map(|(i, c)| (c.name().to_string(), row.get(i).unwrap_or("").to_string())) + .collect::>(), + ), + _ => None, + }) + .collect()) + } + /// Cooperatively shut down every background task this node owns. pub async fn shutdown(self) { self.pg_shutdown_bus.initiate(); diff --git a/nodedb-test-support/src/tx_batch_helpers.rs b/nodedb-test-support/src/tx_batch_helpers.rs index 8f288fdef..77b5a4edc 100644 --- a/nodedb-test-support/src/tx_batch_helpers.rs +++ b/nodedb-test-support/src/tx_batch_helpers.rs @@ -246,6 +246,7 @@ pub fn columnar_count(collection: &str) -> PhysicalPlan { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }) } @@ -282,6 +283,7 @@ pub fn timeseries_scan(collection: &str) -> PhysicalPlan { computed_columns: Vec::new(), having: Vec::new(), tiering: None, + output_columns: Vec::new(), }) } diff --git a/nodedb-types/src/array_route.rs b/nodedb-types/src/array_route.rs index 506b939eb..067a115eb 100644 --- a/nodedb-types/src/array_route.rs +++ b/nodedb-types/src/array_route.rs @@ -10,7 +10,8 @@ //! //! The router uses at most [`VShardId::BITS`] bucket bits. Bucket `b` of `R` //! routing bits owns the vShards `[b * stride, (b + 1) * stride)`, with -//! `stride = VShardId::COUNT >> R`. Its primary vShard is `b * stride`. +//! `stride = VShardId::COUNT >> R`. Its primary vShard is the one +//! [`bucket_to_vshard`] picks inside that range. use serde::{Deserialize, Serialize}; @@ -93,8 +94,20 @@ pub fn routing_bits(prefix_bits: u8) -> Result { /// The primary vShard of `bucket` at `bits` routing bits. `bits` comes from /// [`routing_bits`], so it is at most [`VShardId::BITS`]. +/// +/// A Data Plane core owns `vshard % num_cores`. A primary of `bucket * stride` +/// alone is a multiple of `stride`, so every primary lands on few cores. The +/// offset inside the stride is the bucket's top bits. The low bits of a +/// primary then hold the bucket's low bits and its top bits, and a +/// power-of-two core count up to `2^bits` owns an equal share of primaries. +/// The offset stays below `stride`, so `vshard >> (VShardId::BITS - bits)` +/// recovers the bucket. pub const fn bucket_to_vshard(bucket: u32, bits: u8) -> u32 { - bucket << (VShardId::BITS - bits) as u32 + let stride_bits = (VShardId::BITS - bits) as u32; + let bits = bits as u32; + let top_shift = bits.saturating_sub(stride_bits); + let offset = (bucket >> top_shift) & ((1u32 << stride_bits) - 1); + (bucket << stride_bits) | offset } /// The primary vShard that owns the tile with key `tile_key`. @@ -167,10 +180,13 @@ mod tests { #[test] fn tile_routing_uses_the_mixed_bucket() { - // prefix_bits=8: stride=4, so the vShard is the mixed bucket times 4. + // prefix_bits=8: stride=4, so the vShard is the mixed bucket times 4 + // plus the bucket's top two bits. for key in [0u64, 1, 2, 0x0100_0000_0000_0000, u64::MAX] { let vshard = array_vshard_for_tile(key, 8).unwrap(); - assert_eq!(vshard, prefix_bucket(key, 8) * 4, "key={key:#x}"); + let bucket = prefix_bucket(key, 8); + assert_eq!(vshard, bucket * 4 + (bucket >> 6), "key={key:#x}"); + assert_eq!(vshard >> 2, bucket, "key={key:#x}"); } // prefix_bits=10: stride=1, so the vShard is the mixed bucket. let key = 0x0040_0000_0000_0000; @@ -222,12 +238,62 @@ mod tests { fn bucket_to_vshard_strides() { assert_eq!(VShardId::COUNT, 1 << VShardId::BITS); assert_eq!(VShardId::BITS, 10); - assert_eq!(bucket_to_vshard(1, 1), 512); + assert_eq!(bucket_to_vshard(0, 1), 0); + assert_eq!(bucket_to_vshard(1, 1), 513); assert_eq!(bucket_to_vshard(1, 8), 4); + assert_eq!(bucket_to_vshard(255, 8), 1023); + assert_eq!(bucket_to_vshard(64, 8), 257); assert_eq!(bucket_to_vshard(1, 10), 1); assert_eq!(routing_bits(16), Ok(10)); } + /// Every bucket's primary sits in the bucket's own stride, and the shift + /// recovers the bucket. + #[test] + fn the_shift_recovers_the_bucket_of_every_primary() { + for prefix_bits in 1u8..=MAX_PREFIX_BITS { + let bits = routing_bits(prefix_bits).unwrap(); + let shift = u32::from(VShardId::BITS - bits); + for bucket in 0u32..(1 << bits) { + let vshard = bucket_to_vshard(bucket, bits); + assert!(vshard < VShardId::COUNT, "bits={bits} bucket={bucket}"); + assert_eq!(vshard >> shift, bucket, "bits={bits} bucket={bucket}"); + } + } + } + + /// A core owns `vshard % num_cores`. With the default `prefix_bits`, a + /// 4-core and an 8-core node own an equal share of the primaries. + #[test] + fn default_prefix_bits_put_primaries_on_every_core() { + let bits = routing_bits(8).unwrap(); + for num_cores in [2u32, 4, 8, 16] { + let mut per_core = vec![0u32; num_cores as usize]; + for bucket in 0u32..(1 << bits) { + per_core[(bucket_to_vshard(bucket, bits) % num_cores) as usize] += 1; + } + let share = (1u32 << bits) / num_cores; + assert!( + per_core.iter().all(|&n| n == share), + "{num_cores} cores: primaries per core {per_core:?}" + ); + } + } + + /// Every routing granularity spreads its primaries over as many cores as + /// it has buckets, up to eight. + #[test] + fn every_prefix_bits_spreads_over_the_cores() { + for prefix_bits in 1u8..=MAX_PREFIX_BITS { + let bits = routing_bits(prefix_bits).unwrap(); + let num_cores = (1u32 << bits).min(8); + let cores: std::collections::BTreeSet = (0u32..(1 << bits)) + .map(|bucket| bucket_to_vshard(bucket, bits) % num_cores) + .collect(); + assert_eq!(cores.len() as u32, num_cores, "bits={bits}"); + } + } + #[test] fn shard_bucket_roundtrips_through_msgpack() { let b = ShardBucket { diff --git a/nodedb-types/src/bbox.rs b/nodedb-types/src/bbox.rs index aa9a37188..f0075f6f3 100644 --- a/nodedb-types/src/bbox.rs +++ b/nodedb-types/src/bbox.rs @@ -1,18 +1,44 @@ // SPDX-License-Identifier: Apache-2.0 -//! Axis-aligned bounding box for 2D geometries (WGS-84). +//! Axis-aligned bounding box for 2D geometries. //! -//! Handles the anti-meridian (±180° date line) correctly: when a geometry -//! crosses the date line, `min_lng > max_lng`. All intersection and -//! containment logic accounts for this split case. +//! In [`BboxSpace::Geographic`] space the box handles the anti-meridian +//! (±180° date line): when a geometry crosses the date line, +//! `min_lng > max_lng`. All intersection and containment logic accounts for +//! this split case. In [`BboxSpace::Planar`] space x is a projected easting in +//! metres, nothing wraps, and `min_lng <= max_lng` always holds for a +//! non-empty box. use crate::geometry::Geometry; use serde::{Deserialize, Serialize}; use zerompk::{FromMessagePack, ToMessagePack}; -/// An axis-aligned bounding box in WGS-84 coordinates. +/// Metres per degree of latitude, the WGS 84 meridian average used by the +/// equirectangular expansion. +const METRES_PER_DEGREE_LAT: f64 = 110_540.0; + +/// Metres per degree of longitude at the equator (WGS 84 semi-major axis +/// times π / 180, rounded). +const METRES_PER_DEGREE_LNG_EQUATOR: f64 = 111_320.0; + +/// The coordinate space a bounding box lives in. +/// +/// Derived from the SRID of the geometries it bounds +/// (`Srid::bbox_space`). Union and enlargement depend on it: only +/// geographic longitude wraps at ±180°. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub enum BboxSpace { + /// Longitude / latitude degrees. Longitude wraps at the anti-meridian. + #[default] + Geographic, + /// Projected metres. No axis wraps. + Planar, +} + +/// An axis-aligned bounding box. The `lng` fields hold x and the `lat` fields +/// hold y: degrees for a geographic SRID, metres for a projected one. /// -/// **Anti-meridian convention:** When `min_lng > max_lng`, the bbox wraps +/// **Anti-meridian convention (geographic only):** When `min_lng > max_lng`, the bbox wraps /// across the ±180° date line. For example, a geometry spanning 170°E to /// 170°W has `min_lng = 170.0, max_lng = -170.0`. /// @@ -119,11 +145,28 @@ impl BoundingBox { } } + /// Union in geographic space. Same as `union_in(other, BboxSpace::Geographic)`. pub fn union(&self, other: &BoundingBox) -> BoundingBox { + self.union_in(other, BboxSpace::Geographic) + } + + /// Smallest box covering both, in `space`. + /// + /// Geographic space picks the shorter of the direct and the + /// anti-meridian-wrapping longitude span. Planar space takes plain + /// per-axis min / max, so two disjoint metre boxes never wrap. + pub fn union_in(&self, other: &BoundingBox, space: BboxSpace) -> BoundingBox { let min_lat = self.min_lat.min(other.min_lat); let max_lat = self.max_lat.max(other.max_lat); - let (min_lng, max_lng) = - union_longitude(self.min_lng, self.max_lng, other.min_lng, other.max_lng); + let (min_lng, max_lng) = match space { + BboxSpace::Geographic => { + union_longitude(self.min_lng, self.max_lng, other.min_lng, other.max_lng) + } + BboxSpace::Planar => ( + self.min_lng.min(other.min_lng), + self.max_lng.max(other.max_lng), + ), + }; BoundingBox { min_lng, min_lat, @@ -151,11 +194,28 @@ impl BoundingBox { } } - /// Expand by meters (equirectangular approximation, <0.5% error ±70°). + /// Expand by meters in geographic space. Same as + /// `expand_meters_in(meters, BboxSpace::Geographic)`. pub fn expand_meters(&self, meters: f64) -> BoundingBox { - let dlat = meters / 110_540.0; - let center_lat = (self.min_lat + self.max_lat) / 2.0; - let dlng = meters / (111_320.0 * center_lat.to_radians().cos()); + self.expand_meters_in(meters, BboxSpace::Geographic) + } + + /// Expand by `meters` on every side, in `space`. + /// + /// Geographic space converts metres to degrees with an equirectangular + /// approximation (<0.5% error within ±70° latitude). Planar space adds + /// the metres directly. + pub fn expand_meters_in(&self, meters: f64, space: BboxSpace) -> BoundingBox { + let (dlng, dlat) = match space { + BboxSpace::Geographic => { + let center_lat = (self.min_lat + self.max_lat) / 2.0; + ( + meters / (METRES_PER_DEGREE_LNG_EQUATOR * center_lat.to_radians().cos()), + meters / METRES_PER_DEGREE_LAT, + ) + } + BboxSpace::Planar => (meters, meters), + }; BoundingBox { min_lng: self.min_lng - dlng, min_lat: self.min_lat - dlat, @@ -177,8 +237,15 @@ impl BoundingBox { lat_overlap * lng_overlap } + /// Enlargement in geographic space. Same as + /// `enlargement_in(other, BboxSpace::Geographic)`. pub fn enlargement(&self, other: &BoundingBox) -> f64 { - self.union(other).area() - self.area() + self.enlargement_in(other, BboxSpace::Geographic) + } + + /// Area added by growing `self` to cover `other`, in `space`. + pub fn enlargement_in(&self, other: &BoundingBox, space: BboxSpace) -> f64 { + self.union_in(other, space).area() - self.area() } } @@ -240,8 +307,16 @@ fn union_longitude(a_min: f64, a_max: f64, b_min: f64, b_max: f64) -> (f64, f64) } } -/// Compute bounding box for any Geometry. +/// Compute bounding box for any Geometry in geographic space. Same as +/// `geometry_bbox_in(geom, BboxSpace::Geographic)`. pub fn geometry_bbox(geom: &Geometry) -> BoundingBox { + geometry_bbox_in(geom, BboxSpace::Geographic) +} + +/// Compute bounding box for any Geometry in `space`. +/// +/// `space` decides how the members of a GeometryCollection combine. +pub fn geometry_bbox_in(geom: &Geometry, space: BboxSpace) -> BoundingBox { match geom { Geometry::Point { coordinates } => BoundingBox::from_point(coordinates[0], coordinates[1]), Geometry::LineString { coordinates } => bbox_from_coords(coordinates), @@ -268,9 +343,9 @@ pub fn geometry_bbox(geom: &Geometry) -> BoundingBox { if geometries.is_empty() { return BoundingBox::new(0.0, 0.0, 0.0, 0.0); } - let mut result = geometry_bbox(&geometries[0]); + let mut result = geometry_bbox_in(&geometries[0], space); for geom in &geometries[1..] { - result = result.union(&geometry_bbox(geom)); + result = result.union_in(&geometry_bbox_in(geom, space), space); } result } @@ -368,6 +443,26 @@ mod tests { assert_eq!(u.max_lng, 10.0); } + /// Two disjoint metre boxes far past ±180 must not be read as longitude + /// and wrapped across the anti-meridian. + #[test] + fn planar_union_of_disjoint_metre_boxes_does_not_wrap() { + let a = BoundingBox::new(1e5, 0.0, 2e5, 10.0); + let b = BoundingBox::new(7e5, 0.0, 8e5, 10.0); + let u = a.union_in(&b, BboxSpace::Planar); + assert_eq!((u.min_lng, u.max_lng), (1e5, 8e5)); + assert!(!u.crosses_antimeridian()); + assert!((a.enlargement_in(&b, BboxSpace::Planar) - (7e5 - 1e5) * 10.0).abs() < 1e-6); + } + + #[test] + fn planar_expand_adds_metres_directly() { + let bb = BoundingBox::from_point(500_000.0, 100.0); + let e = bb.expand_meters_in(250.0, BboxSpace::Planar); + assert_eq!((e.min_lng, e.max_lng), (499_750.0, 500_250.0)); + assert_eq!((e.min_lat, e.max_lat), (-150.0, 350.0)); + } + #[test] fn overlap_area_partial() { let a = BoundingBox::new(0.0, 0.0, 10.0, 10.0); diff --git a/nodedb-types/src/collection_config.rs b/nodedb-types/src/collection_config.rs index 73354a86b..abcb73c2f 100644 --- a/nodedb-types/src/collection_config.rs +++ b/nodedb-types/src/collection_config.rs @@ -183,6 +183,15 @@ pub enum PayloadAtom { }, } +impl PayloadAtom { + /// The payload field this atom filters on. + pub fn field(&self) -> &str { + match self { + Self::Eq(field, _) | Self::In(field, _) | Self::Range { field, .. } => field, + } + } +} + /// Names WHAT a key-partitioned collection hashes to derive its vShard. /// /// Defined as substrate for future routing and array layers; not yet populated diff --git a/nodedb-types/src/error/sqlstate.rs b/nodedb-types/src/error/sqlstate.rs index 68e88c145..050dace21 100644 --- a/nodedb-types/src/error/sqlstate.rs +++ b/nodedb-types/src/error/sqlstate.rs @@ -198,6 +198,10 @@ pub const UNDEFINED_COLUMN: &str = "42703"; /// more than one relation in scope. pub const AMBIGUOUS_COLUMN: &str = "42702"; +/// `42803` — `grouping_error`: a grouped SELECT list reads a column that is +/// neither a GROUP BY key nor inside an aggregate. +pub const GROUPING_ERROR: &str = "42803"; + /// `42804` — `datatype_mismatch` (a literal the declared column type cannot /// represent, such as a column `DEFAULT` refused at `CREATE`) pub const DATATYPE_MISMATCH: &str = "42804"; @@ -439,6 +443,7 @@ mod tests { SYNTAX_ERROR, UNDEFINED_COLUMN, AMBIGUOUS_COLUMN, + GROUPING_ERROR, UNDEFINED_OBJECT, DATATYPE_MISMATCH, WRONG_OBJECT_TYPE, @@ -503,6 +508,7 @@ mod tests { assert_eq!(UNIQUE_VIOLATION, "23505"); assert_eq!(UNDEFINED_COLUMN, "42703"); assert_eq!(AMBIGUOUS_COLUMN, "42702"); + assert_eq!(GROUPING_ERROR, "42803"); assert_eq!(UNDEFINED_TABLE, "42P01"); assert_eq!(INSUFFICIENT_PRIVILEGE, "42501"); assert_eq!(QUERY_CANCELED.0, "57014"); diff --git a/nodedb-types/src/geometry.rs b/nodedb-types/src/geometry.rs deleted file mode 100644 index ae695105d..000000000 --- a/nodedb-types/src/geometry.rs +++ /dev/null @@ -1,245 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 - -//! GeoJSON-compatible geometry types. -//! -//! Supports Point, LineString, Polygon, MultiPoint, MultiLineString, -//! MultiPolygon, and GeometryCollection. Stored as GeoJSON for JSON -//! compatibility. Includes distance (Haversine), area, bearing, and -//! centroid calculations. - -use serde::{Deserialize, Serialize}; - -/// A 2D coordinate (longitude, latitude) following GeoJSON convention. -/// Note: GeoJSON uses [lng, lat] order, NOT [lat, lng]. -#[derive( - Debug, - Clone, - Copy, - PartialEq, - Serialize, - Deserialize, - zerompk::ToMessagePack, - zerompk::FromMessagePack, -)] -pub struct Coord { - pub lng: f64, - pub lat: f64, -} - -impl Coord { - pub fn new(lng: f64, lat: f64) -> Self { - Self { lng, lat } - } -} - -/// GeoJSON-compatible geometry types. -#[derive( - Debug, - Clone, - PartialEq, - Serialize, - Deserialize, - zerompk::ToMessagePack, - zerompk::FromMessagePack, -)] -#[serde(tag = "type")] -#[non_exhaustive] -pub enum Geometry { - Point { - coordinates: [f64; 2], - }, - LineString { - coordinates: Vec<[f64; 2]>, - }, - Polygon { - coordinates: Vec>, - }, - MultiPoint { - coordinates: Vec<[f64; 2]>, - }, - MultiLineString { - coordinates: Vec>, - }, - MultiPolygon { - coordinates: Vec>>, - }, - GeometryCollection { - geometries: Vec, - }, -} - -impl Geometry { - /// Create a Point from (longitude, latitude). - pub fn point(lng: f64, lat: f64) -> Self { - Geometry::Point { - coordinates: [lng, lat], - } - } - - /// Create a LineString from a series of [lng, lat] pairs. - pub fn line_string(coords: Vec<[f64; 2]>) -> Self { - Geometry::LineString { - coordinates: coords, - } - } - - /// Create a Polygon from exterior ring (and optional holes). - /// - /// The first ring is the exterior, subsequent rings are holes. - /// Each ring must be a closed loop (first point == last point). - pub fn polygon(rings: Vec>) -> Self { - Geometry::Polygon { coordinates: rings } - } - - /// Get the type name of this geometry. - pub fn geometry_type(&self) -> &'static str { - match self { - Geometry::Point { .. } => "Point", - Geometry::LineString { .. } => "LineString", - Geometry::Polygon { .. } => "Polygon", - Geometry::MultiPoint { .. } => "MultiPoint", - Geometry::MultiLineString { .. } => "MultiLineString", - Geometry::MultiPolygon { .. } => "MultiPolygon", - Geometry::GeometryCollection { .. } => "GeometryCollection", - } - } -} - -/// Parse a GeoJSON string into a [`Geometry`]. -/// -/// Shared by every storage/read path that may encounter geometry stored as a -/// JSON string rather than a native object — SQL `ST_Point(...)` inserts -/// serialize to a GeoJSON string, while schemaless document writes and -/// `Value::Geometry` keep the native form. Centralizing the string-parse core -/// here keeps the three call sites (document index build, spatial read path, -/// columnar geometry index) from drifting on which parser they use. -/// -/// Uses `sonic_rs` per workspace policy (never `serde_json::from_str` for -/// runtime JSON parsing). Returns `None` on malformed/non-geometry JSON. -pub fn from_geojson_str(s: &str) -> Option { - sonic_rs::from_str(s).ok() -} - -// ── Geo math functions ── - -const EARTH_RADIUS_M: f64 = 6_371_000.0; - -/// Haversine distance between two points in meters. -/// -/// Input: (lng1, lat1) and (lng2, lat2) in degrees. -pub fn haversine_distance(lng1: f64, lat1: f64, lng2: f64, lat2: f64) -> f64 { - let lat1_r = lat1.to_radians(); - let lat2_r = lat2.to_radians(); - let dlat = (lat2 - lat1).to_radians(); - let dlng = (lng2 - lng1).to_radians(); - - let a = (dlat / 2.0).sin().powi(2) + lat1_r.cos() * lat2_r.cos() * (dlng / 2.0).sin().powi(2); - let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt()); - EARTH_RADIUS_M * c -} - -/// Haversine bearing from point A to point B in degrees (0-360). -pub fn haversine_bearing(lng1: f64, lat1: f64, lng2: f64, lat2: f64) -> f64 { - let lat1_r = lat1.to_radians(); - let lat2_r = lat2.to_radians(); - let dlng = (lng2 - lng1).to_radians(); - - let y = dlng.sin() * lat2_r.cos(); - let x = lat1_r.cos() * lat2_r.sin() - lat1_r.sin() * lat2_r.cos() * dlng.cos(); - let bearing = y.atan2(x).to_degrees(); - (bearing + 360.0) % 360.0 -} - -/// Check if a point is inside a polygon (ray casting algorithm). -pub fn point_in_polygon(lng: f64, lat: f64, ring: &[[f64; 2]]) -> bool { - let mut inside = false; - let n = ring.len(); - let mut j = n.wrapping_sub(1); - for i in 0..n { - let yi = ring[i][1]; - let yj = ring[j][1]; - if ((yi > lat) != (yj > lat)) - && (lng < (ring[j][0] - ring[i][0]) * (lat - yi) / (yj - yi) + ring[i][0]) - { - inside = !inside; - } - j = i; - } - inside -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn point_creation() { - let p = Geometry::point(-73.9857, 40.7484); - assert_eq!(p.geometry_type(), "Point"); - if let Geometry::Point { coordinates } = &p { - assert!((coordinates[0] - (-73.9857)).abs() < 1e-6); - assert!((coordinates[1] - 40.7484).abs() < 1e-6); - } - } - - #[test] - fn haversine_nyc_to_london() { - // NYC: -74.006, 40.7128 → London: -0.1278, 51.5074 - let d = haversine_distance(-74.006, 40.7128, -0.1278, 51.5074); - // ~5,570 km - assert!((d - 5_570_000.0).abs() < 50_000.0, "got {d}m"); - } - - #[test] - fn haversine_same_point() { - let d = haversine_distance(0.0, 0.0, 0.0, 0.0); - assert!(d.abs() < 1e-6); - } - - #[test] - fn bearing_north() { - let b = haversine_bearing(0.0, 0.0, 0.0, 1.0); - assert!((b - 0.0).abs() < 1.0, "expected ~0, got {b}"); - } - - #[test] - fn bearing_east() { - let b = haversine_bearing(0.0, 0.0, 1.0, 0.0); - assert!((b - 90.0).abs() < 1.0, "expected ~90, got {b}"); - } - - #[test] - fn point_in_polygon_inside() { - let ring = vec![ - [0.0, 0.0], - [10.0, 0.0], - [10.0, 10.0], - [0.0, 10.0], - [0.0, 0.0], - ]; - assert!(point_in_polygon(5.0, 5.0, &ring)); - assert!(!point_in_polygon(15.0, 5.0, &ring)); - } - - #[test] - fn geojson_serialize() { - let p = Geometry::point(1.0, 2.0); - let json = sonic_rs::to_string(&p).unwrap(); - assert!(json.contains("\"type\":\"Point\"")); - assert!(json.contains("\"coordinates\":[1.0,2.0]")); - } - - #[test] - fn geojson_roundtrip() { - let original = Geometry::polygon(vec![vec![ - [0.0, 0.0], - [1.0, 0.0], - [1.0, 1.0], - [0.0, 1.0], - [0.0, 0.0], - ]]); - let json = sonic_rs::to_string(&original).unwrap(); - let parsed: Geometry = sonic_rs::from_str(&json).unwrap(); - assert_eq!(original, parsed); - } -} diff --git a/nodedb-types/src/geometry/math.rs b/nodedb-types/src/geometry/math.rs new file mode 100644 index 000000000..0aeb5c52a --- /dev/null +++ b/nodedb-types/src/geometry/math.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Spherical geo math on (longitude, latitude) degrees. + +/// Mean Earth radius in metres (IUGG R1), the sphere every haversine +/// measure in this crate uses. +const EARTH_RADIUS_M: f64 = 6_371_000.0; + +/// Haversine distance between two points in meters. +/// +/// Input: (lng1, lat1) and (lng2, lat2) in degrees. +pub fn haversine_distance(lng1: f64, lat1: f64, lng2: f64, lat2: f64) -> f64 { + let lat1_r = lat1.to_radians(); + let lat2_r = lat2.to_radians(); + let dlat = (lat2 - lat1).to_radians(); + let dlng = (lng2 - lng1).to_radians(); + + let a = (dlat / 2.0).sin().powi(2) + lat1_r.cos() * lat2_r.cos() * (dlng / 2.0).sin().powi(2); + let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt()); + EARTH_RADIUS_M * c +} + +/// Haversine bearing from point A to point B in degrees (0-360). +pub fn haversine_bearing(lng1: f64, lat1: f64, lng2: f64, lat2: f64) -> f64 { + let lat1_r = lat1.to_radians(); + let lat2_r = lat2.to_radians(); + let dlng = (lng2 - lng1).to_radians(); + + let y = dlng.sin() * lat2_r.cos(); + let x = lat1_r.cos() * lat2_r.sin() - lat1_r.sin() * lat2_r.cos() * dlng.cos(); + let bearing = y.atan2(x).to_degrees(); + (bearing + 360.0) % 360.0 +} + +/// Check if a point is inside a polygon (ray casting algorithm). +/// +/// Planar in its inputs, so it serves geographic and projected coordinates. +pub fn point_in_polygon(lng: f64, lat: f64, ring: &[[f64; 2]]) -> bool { + let mut inside = false; + let n = ring.len(); + let mut j = n.wrapping_sub(1); + for i in 0..n { + let yi = ring[i][1]; + let yj = ring[j][1]; + if ((yi > lat) != (yj > lat)) + && (lng < (ring[j][0] - ring[i][0]) * (lat - yi) / (yj - yi) + ring[i][0]) + { + inside = !inside; + } + j = i; + } + inside +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn haversine_nyc_to_london() { + // NYC: -74.006, 40.7128 → London: -0.1278, 51.5074 + let d = haversine_distance(-74.006, 40.7128, -0.1278, 51.5074); + // ~5,570 km + assert!((d - 5_570_000.0).abs() < 50_000.0, "got {d}m"); + } + + #[test] + fn haversine_same_point() { + let d = haversine_distance(0.0, 0.0, 0.0, 0.0); + assert!(d.abs() < 1e-6); + } + + #[test] + fn bearing_north() { + let b = haversine_bearing(0.0, 0.0, 0.0, 1.0); + assert!((b - 0.0).abs() < 1.0, "expected ~0, got {b}"); + } + + #[test] + fn bearing_east() { + let b = haversine_bearing(0.0, 0.0, 1.0, 0.0); + assert!((b - 90.0).abs() < 1.0, "expected ~90, got {b}"); + } + + #[test] + fn point_in_polygon_inside() { + let ring = vec![ + [0.0, 0.0], + [10.0, 0.0], + [10.0, 10.0], + [0.0, 10.0], + [0.0, 0.0], + ]; + assert!(point_in_polygon(5.0, 5.0, &ring)); + assert!(!point_in_polygon(15.0, 5.0, &ring)); + } +} diff --git a/nodedb-types/src/geometry/mod.rs b/nodedb-types/src/geometry/mod.rs new file mode 100644 index 000000000..376bb2474 --- /dev/null +++ b/nodedb-types/src/geometry/mod.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! GeoJSON-compatible geometry types, spatial reference systems, and +//! spherical geo math. + +pub mod math; +pub mod shape; +pub mod srid; +pub mod tagged; + +pub use math::{haversine_bearing, haversine_distance, point_in_polygon}; +pub use shape::{Coord, Geometry}; +pub use srid::{SUPPORTED_SRID_CODES, Srid, UnsupportedSrid}; +pub use tagged::{ + CrsMember, CrsProperties, GeoJsonError, SridGeometry, from_geojson_str, parse_geojson, + parse_geojson_slice, srid_from_crs, +}; diff --git a/nodedb-types/src/geometry/shape.rs b/nodedb-types/src/geometry/shape.rs new file mode 100644 index 000000000..1ed8fda4c --- /dev/null +++ b/nodedb-types/src/geometry/shape.rs @@ -0,0 +1,221 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! GeoJSON-compatible geometry shapes. +//! +//! Supports Point, LineString, Polygon, MultiPoint, MultiLineString, +//! MultiPolygon, and GeometryCollection. A shape holds coordinates only. The +//! reference system those coordinates are in lives on +//! [`SridGeometry`](super::SridGeometry). + +use serde::{Deserialize, Serialize}; + +/// A 2D coordinate (x, y) following GeoJSON order. +/// +/// In a geographic reference system x is longitude and y is latitude, in +/// degrees. In a projected one x is easting and y is northing, in metres. +#[derive( + Debug, + Clone, + Copy, + PartialEq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct Coord { + pub lng: f64, + pub lat: f64, +} + +impl Coord { + pub fn new(lng: f64, lat: f64) -> Self { + Self { lng, lat } + } +} + +/// GeoJSON-compatible geometry types. +#[derive( + Debug, + Clone, + PartialEq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +#[serde(tag = "type")] +#[non_exhaustive] +pub enum Geometry { + Point { + coordinates: [f64; 2], + }, + LineString { + coordinates: Vec<[f64; 2]>, + }, + Polygon { + coordinates: Vec>, + }, + MultiPoint { + coordinates: Vec<[f64; 2]>, + }, + MultiLineString { + coordinates: Vec>, + }, + MultiPolygon { + coordinates: Vec>>, + }, + GeometryCollection { + geometries: Vec, + }, +} + +impl Geometry { + /// Create a Point from (x, y). + pub fn point(lng: f64, lat: f64) -> Self { + Geometry::Point { + coordinates: [lng, lat], + } + } + + /// Create a LineString from a series of [x, y] pairs. + pub fn line_string(coords: Vec<[f64; 2]>) -> Self { + Geometry::LineString { + coordinates: coords, + } + } + + /// Create a Polygon from exterior ring (and optional holes). + /// + /// The first ring is the exterior, subsequent rings are holes. + /// Each ring must be a closed loop (first point == last point). + pub fn polygon(rings: Vec>) -> Self { + Geometry::Polygon { coordinates: rings } + } + + /// Get the type name of this geometry. + pub fn geometry_type(&self) -> &'static str { + match self { + Geometry::Point { .. } => "Point", + Geometry::LineString { .. } => "LineString", + Geometry::Polygon { .. } => "Polygon", + Geometry::MultiPoint { .. } => "MultiPoint", + Geometry::MultiLineString { .. } => "MultiLineString", + Geometry::MultiPolygon { .. } => "MultiPolygon", + Geometry::GeometryCollection { .. } => "GeometryCollection", + } + } + + /// Map every coordinate through `f`, keeping the shape. + /// + /// Stops at the first coordinate `f` rejects and returns its error. + pub fn try_map_coords( + &self, + f: &mut impl FnMut([f64; 2]) -> Result<[f64; 2], E>, + ) -> Result { + Ok(match self { + Geometry::Point { coordinates } => Geometry::Point { + coordinates: f(*coordinates)?, + }, + Geometry::LineString { coordinates } => Geometry::LineString { + coordinates: map_path(coordinates, f)?, + }, + Geometry::Polygon { coordinates } => Geometry::Polygon { + coordinates: map_rings(coordinates, f)?, + }, + Geometry::MultiPoint { coordinates } => Geometry::MultiPoint { + coordinates: map_path(coordinates, f)?, + }, + Geometry::MultiLineString { coordinates } => Geometry::MultiLineString { + coordinates: map_rings(coordinates, f)?, + }, + Geometry::MultiPolygon { coordinates } => Geometry::MultiPolygon { + coordinates: coordinates + .iter() + .map(|polygon| map_rings(polygon, f)) + .collect::>()?, + }, + Geometry::GeometryCollection { geometries } => Geometry::GeometryCollection { + geometries: geometries + .iter() + .map(|g| g.try_map_coords(f)) + .collect::>()?, + }, + }) + } +} + +fn map_path( + path: &[[f64; 2]], + f: &mut impl FnMut([f64; 2]) -> Result<[f64; 2], E>, +) -> Result, E> { + path.iter().map(|c| f(*c)).collect() +} + +fn map_rings( + rings: &[Vec<[f64; 2]>], + f: &mut impl FnMut([f64; 2]) -> Result<[f64; 2], E>, +) -> Result>, E> { + rings.iter().map(|ring| map_path(ring, f)).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn point_creation() { + let p = Geometry::point(-73.9857, 40.7484); + assert_eq!(p.geometry_type(), "Point"); + if let Geometry::Point { coordinates } = &p { + assert!((coordinates[0] - (-73.9857)).abs() < 1e-6); + assert!((coordinates[1] - 40.7484).abs() < 1e-6); + } + } + + #[test] + fn geojson_serialize() { + let p = Geometry::point(1.0, 2.0); + let json = sonic_rs::to_string(&p).unwrap(); + assert!(json.contains("\"type\":\"Point\"")); + assert!(json.contains("\"coordinates\":[1.0,2.0]")); + } + + #[test] + fn geojson_roundtrip() { + let original = Geometry::polygon(vec![vec![ + [0.0, 0.0], + [1.0, 0.0], + [1.0, 1.0], + [0.0, 1.0], + [0.0, 0.0], + ]]); + let json = sonic_rs::to_string(&original).unwrap(); + let parsed: Geometry = sonic_rs::from_str(&json).unwrap(); + assert_eq!(original, parsed); + } + + #[test] + fn try_map_coords_keeps_shape() { + let gc = Geometry::GeometryCollection { + geometries: vec![ + Geometry::point(1.0, 2.0), + Geometry::polygon(vec![vec![[0.0, 0.0], [1.0, 0.0], [0.0, 0.0]]]), + ], + }; + let shifted = gc + .try_map_coords(&mut |[x, y]| Ok::<_, ()>([x + 10.0, y])) + .unwrap(); + let Geometry::GeometryCollection { geometries } = shifted else { + panic!("shape changed"); + }; + assert_eq!(geometries[0], Geometry::point(11.0, 2.0)); + } + + #[test] + fn try_map_coords_stops_at_first_error() { + let line = Geometry::line_string(vec![[0.0, 0.0], [1.0, 1.0]]); + let r = line.try_map_coords(&mut |[x, _]| if x > 0.5 { Err(x) } else { Ok([x, 0.0]) }); + assert_eq!(r, Err(1.0)); + } +} diff --git a/nodedb-types/src/geometry/srid.rs b/nodedb-types/src/geometry/srid.rs new file mode 100644 index 000000000..4bb85739d --- /dev/null +++ b/nodedb-types/src/geometry/srid.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! The closed set of coordinate reference systems NodeDB supports. + +use serde::{Deserialize, Serialize}; + +use crate::bbox::BboxSpace; + +/// Every supported SRID code, comma-separated, for error messages. +pub const SUPPORTED_SRID_CODES: &str = "4326, 32647, 32648, 3375, 3376"; + +/// A spatial reference system identifier, restricted to the systems NodeDB +/// can transform between. +/// +/// The set is closed: an EPSG code outside it is rejected with +/// [`UnsupportedSrid`], never stored. Serialized as its numeric EPSG code. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize)] +#[serde(into = "u32", try_from = "u32")] +pub enum Srid { + /// EPSG:4326, WGS 84 geographic. x is longitude, y is latitude, degrees. + #[default] + Wgs84, + /// EPSG:32647, UTM zone 47N on WGS 84. Metres. + Utm47N, + /// EPSG:32648, UTM zone 48N on WGS 84. Metres. + Utm48N, + /// EPSG:3375, GDM2000 Peninsula RSO. Metres. + PeninsularRso, + /// EPSG:3376, GDM2000 East Malaysia (Borneo) RSO. Metres. + BorneoRso, +} + +/// An SRID code outside the supported set. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("unsupported SRID {code}; supported: {SUPPORTED_SRID_CODES}")] +pub struct UnsupportedSrid { + pub code: i64, +} + +impl Srid { + /// Every supported SRID, WGS 84 first. + pub const SUPPORTED: [Srid; 5] = [ + Srid::Wgs84, + Srid::Utm47N, + Srid::Utm48N, + Srid::PeninsularRso, + Srid::BorneoRso, + ]; + + /// The EPSG code. + pub const fn code(self) -> u32 { + match self { + Srid::Wgs84 => 4326, + Srid::Utm47N => 32647, + Srid::Utm48N => 32648, + Srid::PeninsularRso => 3375, + Srid::BorneoRso => 3376, + } + } + + /// Resolve an EPSG code. Any code outside [`Srid::SUPPORTED`] fails. + pub fn from_code(code: i64) -> Result { + Srid::SUPPORTED + .into_iter() + .find(|srid| i64::from(srid.code()) == code) + .ok_or(UnsupportedSrid { code }) + } + + /// Whether coordinates are longitude / latitude degrees. A `false` SRID + /// is projected, with coordinates in metres. + pub const fn is_geographic(self) -> bool { + match self { + Srid::Wgs84 => true, + Srid::Utm47N | Srid::Utm48N | Srid::PeninsularRso | Srid::BorneoRso => false, + } + } + + /// The coordinate space bounding boxes in this SRID live in. + pub const fn bbox_space(self) -> BboxSpace { + if self.is_geographic() { + BboxSpace::Geographic + } else { + BboxSpace::Planar + } + } +} + +impl From for u32 { + fn from(srid: Srid) -> u32 { + srid.code() + } +} + +impl TryFrom for Srid { + type Error = UnsupportedSrid; + + fn try_from(code: u32) -> Result { + Srid::from_code(i64::from(code)) + } +} + +impl std::fmt::Display for Srid { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.code()) + } +} + +impl zerompk::ToMessagePack for Srid { + fn write(&self, writer: &mut W) -> zerompk::Result<()> { + writer.write_u32(self.code()) + } +} + +impl<'a> zerompk::FromMessagePack<'a> for Srid { + fn read>(reader: &mut R) -> zerompk::Result { + let code = reader.read_u32()?; + Srid::try_from(code).map_err(|e| zerompk::Error::UnknownKey(e.to_string())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_supported_code_round_trips() { + for srid in Srid::SUPPORTED { + assert_eq!(Srid::from_code(i64::from(srid.code())), Ok(srid)); + } + } + + #[test] + fn unsupported_code_names_every_supported_code() { + let err = Srid::from_code(2154).unwrap_err(); + let text = err.to_string(); + for srid in Srid::SUPPORTED { + assert!(text.contains(&srid.code().to_string()), "{text}"); + } + assert!(text.contains("2154"), "{text}"); + } + + #[test] + fn only_wgs84_is_geographic() { + assert!(Srid::Wgs84.is_geographic()); + assert_eq!(Srid::Wgs84.bbox_space(), BboxSpace::Geographic); + for srid in &Srid::SUPPORTED[1..] { + assert!(!srid.is_geographic()); + assert_eq!(srid.bbox_space(), BboxSpace::Planar); + } + } + + #[test] + fn default_is_wgs84() { + assert_eq!(Srid::default(), Srid::Wgs84); + } + + #[test] + fn msgpack_round_trip() { + let bytes = zerompk::to_msgpack_vec(&Srid::BorneoRso).unwrap(); + let back: Srid = zerompk::from_msgpack(&bytes).unwrap(); + assert_eq!(back, Srid::BorneoRso); + } + + #[test] + fn msgpack_rejects_unknown_code() { + let bytes = zerompk::to_msgpack_vec(&2154u32).unwrap(); + assert!(zerompk::from_msgpack::(&bytes).is_err()); + } +} diff --git a/nodedb-types/src/geometry/tagged.rs b/nodedb-types/src/geometry/tagged.rs new file mode 100644 index 000000000..be786e583 --- /dev/null +++ b/nodedb-types/src/geometry/tagged.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! A geometry tagged with its spatial reference system, and its GeoJSON form. +//! +//! GeoJSON (RFC 7946) fixes the reference system to WGS 84. A geometry in any +//! other SRID carries the 2008 GeoJSON `crs` member: +//! `"crs":{"type":"name","properties":{"name":"EPSG:"}}`. A WGS 84 +//! geometry writes no `crs` member, so its text stays plain RFC 7946. + +use serde::de::Error as _; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +use super::shape::Geometry; +use super::srid::{Srid, UnsupportedSrid}; +use crate::bbox::{BoundingBox, geometry_bbox_in}; + +/// A geometry and the SRID its coordinates are in. +#[derive(Debug, Clone, PartialEq)] +pub struct SridGeometry { + pub srid: Srid, + pub geometry: Geometry, +} + +impl SridGeometry { + pub fn new(srid: Srid, geometry: Geometry) -> Self { + Self { srid, geometry } + } + + /// A geometry in WGS 84, the SRID of untagged GeoJSON. + pub fn wgs84(geometry: Geometry) -> Self { + Self::new(Srid::Wgs84, geometry) + } + + /// The same SRID around a different geometry. + pub fn with_geometry(&self, geometry: Geometry) -> Self { + Self::new(self.srid, geometry) + } + + /// Bounding box in the coordinate space of this SRID. + pub fn bbox(&self) -> BoundingBox { + geometry_bbox_in(&self.geometry, self.srid.bbox_space()) + } + + /// Read a geometry out of a document field value. + /// + /// Accepts the three stored forms: a native `Value::Geometry`, a GeoJSON + /// string, and a GeoJSON object decoded from a document. Any other value, + /// or GeoJSON that does not parse, gives `None`. + pub fn from_value(value: &crate::Value) -> Option { + match value { + crate::Value::Geometry(g) => Some(g.clone()), + crate::Value::String(s) => from_geojson_str(s), + crate::Value::Object(_) => { + let json = serde_json::Value::from(value.clone()); + serde_json::from_value(json).ok() + } + _ => None, + } + } +} + +impl From for SridGeometry { + fn from(geometry: Geometry) -> Self { + Self::wgs84(geometry) + } +} + +/// A total, deterministic text form. A WGS 84 geometry renders as its shape. +/// Any other SRID gets an EWKT-style `SRID=;` prefix. +impl std::fmt::Display for SridGeometry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.srid { + Srid::Wgs84 => write!(f, "{:?}", self.geometry), + Srid::Utm47N | Srid::Utm48N | Srid::PeninsularRso | Srid::BorneoRso => { + write!(f, "SRID={};{:?}", self.srid.code(), self.geometry) + } + } + } +} + +/// Why a GeoJSON text does not yield a [`SridGeometry`]. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum GeoJsonError { + #[error("malformed GeoJSON geometry")] + Malformed, + #[error("unrecognized GeoJSON crs name {name:?}")] + UnknownCrsName { name: String }, + #[error(transparent)] + UnsupportedSrid(#[from] UnsupportedSrid), +} + +/// The GeoJSON 2008 named `crs` member. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct CrsMember { + #[serde(rename = "type")] + pub kind: String, + pub properties: CrsProperties, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct CrsProperties { + pub name: String, +} + +/// The `type` value of a named `crs` member. +const CRS_KIND_NAME: &str = "name"; + +/// OGC URN prefix that ends in a version segment and then an EPSG code. +const OGC_EPSG_URN_PREFIX: &str = "urn:ogc:def:crs:epsg:"; + +/// The OGC URNs for WGS 84 in longitude / latitude order. +const OGC_CRS84_URNS: [&str; 2] = [ + "urn:ogc:def:crs:ogc:1.3:crs84", + "urn:ogc:def:crs:ogc::crs84", +]; + +impl CrsMember { + fn for_srid(srid: Srid) -> Self { + Self { + kind: CRS_KIND_NAME.to_string(), + properties: CrsProperties { + name: format!("EPSG:{}", srid.code()), + }, + } + } +} + +/// Resolve an optional `crs` member to an SRID. No member means WGS 84. +/// +/// Accepts `EPSG:`, `urn:ogc:def:crs:EPSG:[version]:` and the OGC +/// CRS84 URNs. Names compare without regard to ASCII case. +pub fn srid_from_crs(crs: Option<&CrsMember>) -> Result { + let Some(member) = crs else { + return Ok(Srid::Wgs84); + }; + let name = member.properties.name.trim(); + let unknown = || GeoJsonError::UnknownCrsName { + name: name.to_string(), + }; + if !member.kind.eq_ignore_ascii_case(CRS_KIND_NAME) { + return Err(unknown()); + } + let lower = name.to_ascii_lowercase(); + if OGC_CRS84_URNS.contains(&lower.as_str()) { + return Ok(Srid::Wgs84); + } + let code_text = if let Some(code) = lower.strip_prefix("epsg:") { + code + } else if let Some(rest) = lower.strip_prefix(OGC_EPSG_URN_PREFIX) { + // `rest` is `:`. The version can be empty. + rest.rsplit_once(':') + .map(|(_, code)| code) + .ok_or_else(unknown)? + } else { + return Err(unknown()); + }; + let code: i64 = code_text.parse().map_err(|_| unknown())?; + Ok(Srid::from_code(code)?) +} + +/// Only the `crs` member of a GeoJSON object. Every other key is skipped +/// without being built, so probing costs one shallow pass. +#[derive(Deserialize)] +struct CrsProbe { + #[serde(default)] + crs: Option, +} + +/// Parse GeoJSON text into a [`SridGeometry`]. +/// +/// The geometry parses through its own derive. The `crs` member is then read +/// by a separate shallow probe, so the common untagged case never buffers +/// the document. Uses `sonic_rs` per workspace policy. +pub fn parse_geojson(text: &str) -> Result { + let geometry: Geometry = sonic_rs::from_str(text).map_err(|_| GeoJsonError::Malformed)?; + let probe: CrsProbe = sonic_rs::from_str(text).map_err(|_| GeoJsonError::Malformed)?; + let srid = srid_from_crs(probe.crs.as_ref())?; + Ok(SridGeometry::new(srid, geometry)) +} + +/// [`parse_geojson`] over raw UTF-8 bytes. +pub fn parse_geojson_slice(bytes: &[u8]) -> Result { + let text = std::str::from_utf8(bytes).map_err(|_| GeoJsonError::Malformed)?; + parse_geojson(text) +} + +/// Parse a GeoJSON string into a [`SridGeometry`]. +/// +/// Shared by every storage/read path that meets geometry stored as a JSON +/// string rather than a native value: SQL `ST_Point(...)` inserts serialize +/// to a GeoJSON string, while schemaless document writes and +/// `Value::Geometry` keep the native form. The document index build, the +/// spatial read path and the columnar geometry index all read through here. +/// +/// Returns `None` on malformed JSON, non-geometry JSON, or a `crs` member +/// naming an unsupported SRID. +pub fn from_geojson_str(s: &str) -> Option { + parse_geojson(s).ok() +} + +#[derive(Serialize)] +struct SridGeometryOut<'a> { + #[serde(flatten)] + geometry: &'a Geometry, + #[serde(skip_serializing_if = "Option::is_none")] + crs: Option, +} + +#[derive(Deserialize)] +struct SridGeometryIn { + #[serde(flatten)] + geometry: Geometry, + #[serde(default)] + crs: Option, +} + +impl Serialize for SridGeometry { + fn serialize(&self, serializer: S) -> Result { + let crs = match self.srid { + Srid::Wgs84 => None, + Srid::Utm47N | Srid::Utm48N | Srid::PeninsularRso | Srid::BorneoRso => { + Some(CrsMember::for_srid(self.srid)) + } + }; + SridGeometryOut { + geometry: &self.geometry, + crs, + } + .serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for SridGeometry { + fn deserialize>(deserializer: D) -> Result { + let raw = SridGeometryIn::deserialize(deserializer)?; + let srid = srid_from_crs(raw.crs.as_ref()).map_err(D::Error::custom)?; + Ok(SridGeometry::new(srid, raw.geometry)) + } +} + +/// MessagePack form: `[srid_code_u32, geometry]`. +impl zerompk::ToMessagePack for SridGeometry { + fn write(&self, writer: &mut W) -> zerompk::Result<()> { + writer.write_array_len(2)?; + self.srid.write(writer)?; + self.geometry.write(writer) + } +} + +impl<'a> zerompk::FromMessagePack<'a> for SridGeometry { + fn read>(reader: &mut R) -> zerompk::Result { + let len = reader.read_array_len()?; + if len != 2 { + return Err(zerompk::Error::ArrayLengthMismatch { + expected: 2, + actual: len, + }); + } + let srid = Srid::read(reader)?; + let geometry = Geometry::read(reader)?; + Ok(SridGeometry::new(srid, geometry)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn crs(name: &str) -> CrsMember { + CrsMember { + kind: "name".into(), + properties: CrsProperties { name: name.into() }, + } + } + + #[test] + fn wgs84_writes_no_crs_member() { + let g = SridGeometry::wgs84(Geometry::point(1.0, 2.0)); + let json = sonic_rs::to_string(&g).unwrap(); + assert!(!json.contains("crs"), "{json}"); + assert!(json.contains("\"type\":\"Point\""), "{json}"); + } + + #[test] + fn projected_geojson_round_trip_keeps_srid() { + let g = SridGeometry::new(Srid::Utm48N, Geometry::point(134445.5, 349624.7)); + let json = sonic_rs::to_string(&g).unwrap(); + assert!(json.contains("\"name\":\"EPSG:32648\""), "{json}"); + assert_eq!(parse_geojson(&json), Ok(g.clone())); + let derived: SridGeometry = sonic_rs::from_str(&json).unwrap(); + assert_eq!(derived, g); + } + + #[test] + fn missing_crs_member_is_wgs84() { + let g = parse_geojson(r#"{"type":"Point","coordinates":[1.0,2.0]}"#).unwrap(); + assert_eq!(g.srid, Srid::Wgs84); + } + + #[test] + fn every_accepted_crs_spelling_resolves() { + assert_eq!( + srid_from_crs(Some(&crs("EPSG:3375"))), + Ok(Srid::PeninsularRso) + ); + assert_eq!( + srid_from_crs(Some(&crs("urn:ogc:def:crs:EPSG::3376"))), + Ok(Srid::BorneoRso) + ); + assert_eq!( + srid_from_crs(Some(&crs("urn:ogc:def:crs:OGC:1.3:CRS84"))), + Ok(Srid::Wgs84) + ); + assert_eq!(srid_from_crs(Some(&crs("epsg:32647"))), Ok(Srid::Utm47N)); + } + + #[test] + fn unsupported_and_unknown_crs_names_fail() { + assert_eq!( + srid_from_crs(Some(&crs("EPSG:2154"))), + Err(GeoJsonError::UnsupportedSrid(UnsupportedSrid { + code: 2154 + })) + ); + assert!(matches!( + srid_from_crs(Some(&crs("urn:x:y"))), + Err(GeoJsonError::UnknownCrsName { .. }) + )); + } + + #[test] + fn from_value_reads_every_stored_form() { + let g = SridGeometry::new(Srid::Utm47N, Geometry::point(1.0, 2.0)); + let text = sonic_rs::to_string(&g).unwrap(); + assert_eq!( + SridGeometry::from_value(&crate::Value::Geometry(g.clone())), + Some(g.clone()) + ); + assert_eq!( + SridGeometry::from_value(&crate::Value::String(text.clone())), + Some(g.clone()) + ); + let json: serde_json::Value = sonic_rs::from_str(&text).unwrap(); + let obj = crate::Value::from(json); + assert_eq!(SridGeometry::from_value(&obj), Some(g)); + } + + #[test] + fn msgpack_round_trip_keeps_srid() { + let g = SridGeometry::new(Srid::BorneoRso, Geometry::point(1.0, 2.0)); + let bytes = zerompk::to_msgpack_vec(&g).unwrap(); + let back: SridGeometry = zerompk::from_msgpack(&bytes).unwrap(); + assert_eq!(back, g); + } + + #[test] + fn display_prefixes_non_wgs84() { + let g = SridGeometry::new(Srid::Utm48N, Geometry::point(1.0, 2.0)); + assert_eq!( + g.to_string(), + "SRID=32648;Point { coordinates: [1.0, 2.0] }" + ); + assert_eq!( + SridGeometry::wgs84(Geometry::point(1.0, 2.0)).to_string(), + "Point { coordinates: [1.0, 2.0] }" + ); + } +} diff --git a/nodedb-types/src/json_msgpack/mod.rs b/nodedb-types/src/json_msgpack/mod.rs index 8e76ce082..09d0cad79 100644 --- a/nodedb-types/src/json_msgpack/mod.rs +++ b/nodedb-types/src/json_msgpack/mod.rs @@ -17,4 +17,6 @@ pub use json_value::JsonValue; pub use native_cell::NativeCell; pub use reader::{json_from_msgpack, value_from_msgpack}; pub use transcoder::msgpack_to_json_string; -pub use writer::{json_to_msgpack, json_to_msgpack_or_empty, value_to_msgpack}; +pub use writer::{ + json_to_msgpack, json_to_msgpack_or_empty, ordered_map_to_msgpack, value_to_msgpack, +}; diff --git a/nodedb-types/src/json_msgpack/writer.rs b/nodedb-types/src/json_msgpack/writer.rs index 944805fcb..dd979abd6 100644 --- a/nodedb-types/src/json_msgpack/writer.rs +++ b/nodedb-types/src/json_msgpack/writer.rs @@ -38,6 +38,16 @@ pub fn value_to_msgpack(value: &crate::Value) -> zerompk::Result> { zerompk::to_msgpack_vec(&NativeRef(value)) } +/// Serialize `fields` as one standard msgpack map whose keys keep the given +/// order. +/// +/// A `Value::Object` writes its keys in hash order, which differs between +/// two maps of the same keys. A row whose bytes must be reproducible, or +/// whose columns must read in a fixed order, is written through this. +pub fn ordered_map_to_msgpack(fields: &[(&str, &crate::Value)]) -> zerompk::Result> { + zerompk::to_msgpack_vec(&OrderedMapRef(fields)) +} + /// A borrowed `Value` written in its plain msgpack form. struct NativeRef<'a>(&'a crate::Value); @@ -47,6 +57,20 @@ impl zerompk::ToMessagePack for NativeRef<'_> { } } +/// Borrowed key/value pairs written as one plain msgpack map, in order. +struct OrderedMapRef<'a>(&'a [(&'a str, &'a crate::Value)]); + +impl zerompk::ToMessagePack for OrderedMapRef<'_> { + fn write(&self, writer: &mut W) -> zerompk::Result<()> { + writer.write_map_len(self.0.len())?; + for (key, value) in self.0 { + writer.write_string(key)?; + write_native_value(writer, value)?; + } + Ok(()) + } +} + /// Write a `nodedb_types::Value` as standard msgpack. /// /// `Duration`, `Decimal`, and `Geometry` are strings, except a `Decimal` @@ -138,3 +162,24 @@ fn write_instant_ext( ) -> zerompk::Result<()> { writer.write_ext(kind.ext_type(), µs.to_be_bytes()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn an_ordered_map_keeps_its_key_order_and_reads_back() { + let x = crate::Value::Integer(3); + let nir = crate::Value::Float(0.5); + let bytes = ordered_map_to_msgpack(&[("x", &x), ("nir", &nir)]).expect("encode"); + // fixmap of 2, then fixstr "x" first. + assert_eq!(&bytes[..3], &[0x82, 0xa1, b'x']); + match crate::json_msgpack::value_from_msgpack(&bytes).expect("decode") { + crate::Value::Object(map) => { + assert_eq!(map.get("x"), Some(&x)); + assert_eq!(map.get("nir"), Some(&nir)); + } + other => panic!("expected an object, got {other:?}"), + } + } +} diff --git a/nodedb-types/src/lib.rs b/nodedb-types/src/lib.rs index f8a7e9372..d7a23ed5b 100644 --- a/nodedb-types/src/lib.rs +++ b/nodedb-types/src/lib.rs @@ -85,7 +85,7 @@ pub use ascii::{ starts_with_ascii_case_insensitive, strip_prefix_ascii_case_insensitive, }; pub use audit_dml::AuditDmlMode; -pub use bbox::{BoundingBox, geometry_bbox}; +pub use bbox::{BboxSpace, BoundingBox, geometry_bbox, geometry_bbox_in}; pub use clone::{CloneOrigin, CloneStatus, MAX_CLONE_DEPTH}; pub use collection::{CollectionType, CollectionTypeParseError}; pub use collection_config::{ @@ -111,8 +111,8 @@ pub use id::{ pub use identity::KeyRepr; pub use json_msgpack::{ InstantKind, JsonValue, MsgpackError, MsgpackResult, NativeCell, json_from_msgpack, - json_to_msgpack, json_to_msgpack_or_empty, msgpack_to_json_string, read_instant, - value_from_msgpack, value_to_msgpack, write_instant, + json_to_msgpack, json_to_msgpack_or_empty, msgpack_to_json_string, ordered_map_to_msgpack, + read_instant, value_from_msgpack, value_to_msgpack, write_instant, }; pub use kv::{KV_DEFAULT_INLINE_THRESHOLD, KvConfig, KvTtlPolicy, is_valid_kv_key_type}; pub use lsn::Lsn; diff --git a/nodedb-types/src/protocol/array_batch.rs b/nodedb-types/src/protocol/array_batch.rs new file mode 100644 index 000000000..abdd90ec9 --- /dev/null +++ b/nodedb-types/src/protocol/array_batch.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! The `ArrayPutBatch` request payload, and the cell type a client writes. + +use serde::{Deserialize, Serialize}; + +use crate::error::{NodeDbError, NodeDbResult}; +use crate::value::Value; + +/// One cell to write into an array: one coordinate per dim and one value per +/// attr, each in the array's schema order. `Value::Null` in an attr writes +/// NULL, which only a nullable attr accepts. +#[derive( + Debug, + Clone, + PartialEq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayCellWrite { + pub coords: Vec, + pub attrs: Vec, +} + +/// The `payload` of an `ArrayPutBatch` request: the cells as columns. `dims` +/// holds one column per dim and `attrs` one column per attr, in schema +/// order. Every column holds one value per cell, in cell order. +#[derive( + Debug, + Clone, + Default, + PartialEq, + Serialize, + Deserialize, + zerompk::ToMessagePack, + zerompk::FromMessagePack, +)] +pub struct ArrayPutBatch { + pub dims: Vec>, + pub attrs: Vec>, +} + +impl ArrayPutBatch { + /// The columns of `cells`. Every cell must have the first cell's number + /// of coordinates and of attrs. + pub fn from_cells(cells: &[ArrayCellWrite]) -> NodeDbResult { + let Some(first) = cells.first() else { + return Ok(Self::default()); + }; + let mut dims: Vec> = (0..first.coords.len()) + .map(|_| Vec::with_capacity(cells.len())) + .collect(); + let mut attrs: Vec> = (0..first.attrs.len()) + .map(|_| Vec::with_capacity(cells.len())) + .collect(); + for (index, cell) in cells.iter().enumerate() { + if cell.coords.len() != dims.len() || cell.attrs.len() != attrs.len() { + return Err(NodeDbError::bad_request(format!( + "array cell {index} has {} coords and {} attrs; cell 0 has {} and {}", + cell.coords.len(), + cell.attrs.len(), + dims.len(), + attrs.len() + ))); + } + for (column, value) in dims.iter_mut().zip(&cell.coords) { + column.push(value.clone()); + } + for (column, value) in attrs.iter_mut().zip(&cell.attrs) { + column.push(value.clone()); + } + } + Ok(Self { dims, attrs }) + } + + /// The number of cells. + pub fn len(&self) -> usize { + self.dims.first().or(self.attrs.first()).map_or(0, Vec::len) + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cell(y: i64, x: i64, v: f64) -> ArrayCellWrite { + ArrayCellWrite { + coords: vec![Value::Integer(y), Value::Integer(x)], + attrs: vec![Value::Float(v)], + } + } + + #[test] + fn cells_become_columns_and_roundtrip() { + let batch = ArrayPutBatch::from_cells(&[cell(0, 1, 0.5), cell(2, 3, 1.5)]).expect("batch"); + assert_eq!(batch.len(), 2); + assert_eq!(batch.dims[1], vec![Value::Integer(1), Value::Integer(3)]); + assert_eq!(batch.attrs[0], vec![Value::Float(0.5), Value::Float(1.5)]); + let bytes = zerompk::to_msgpack_vec(&batch).expect("encode"); + let back: ArrayPutBatch = zerompk::from_msgpack(&bytes).expect("decode"); + assert_eq!(back, batch); + } + + #[test] + fn a_ragged_cell_is_refused() { + let mut short = cell(0, 0, 0.0); + short.coords.pop(); + assert!(ArrayPutBatch::from_cells(&[cell(0, 1, 0.5), short]).is_err()); + } +} diff --git a/nodedb-types/src/protocol/mod.rs b/nodedb-types/src/protocol/mod.rs index 8b16fcb9a..b5f74f319 100644 --- a/nodedb-types/src/protocol/mod.rs +++ b/nodedb-types/src/protocol/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 +pub mod array_batch; pub mod auth; pub mod batch; pub mod error_cause; @@ -9,6 +10,7 @@ pub mod opcodes; pub mod request_fields; pub mod text_fields; +pub use array_batch::{ArrayCellWrite, ArrayPutBatch}; pub use auth::{AuthMethod, AuthResponse}; pub use batch::{BatchDocument, BatchVector}; pub use error_cause::ErrorCausePayload; diff --git a/nodedb-types/src/protocol/opcodes.rs b/nodedb-types/src/protocol/opcodes.rs index 20b5b8df6..5cefeb4b9 100644 --- a/nodedb-types/src/protocol/opcodes.rs +++ b/nodedb-types/src/protocol/opcodes.rs @@ -154,6 +154,11 @@ pub enum OpCode { CrdtListInsert = 0x99, CrdtListDelete = 0x9A, CrdtListMove = 0x9B, + + // ── Array bulk load ────────────────────────────────────────── + /// Load cells into an array: `collection` names it, and `payload` holds a + /// zerompk [`ArrayPutBatch`](super::array_batch::ArrayPutBatch). + ArrayPutBatch = 0x9C, } impl OpCode { @@ -198,6 +203,7 @@ impl OpCode { | OpCode::CrdtListInsert | OpCode::CrdtListDelete | OpCode::CrdtListMove + | OpCode::ArrayPutBatch ) } } @@ -297,6 +303,7 @@ impl TryFrom for OpCode { 0x99 => Ok(OpCode::CrdtListInsert), 0x9A => Ok(OpCode::CrdtListDelete), 0x9B => Ok(OpCode::CrdtListMove), + 0x9C => Ok(OpCode::ArrayPutBatch), other => Err(UnknownOpCode(other)), } } @@ -375,6 +382,13 @@ mod tests { assert_eq!(u8::from(OpCode::CrdtListMove), 0x9B); } + #[test] + fn array_put_batch_opcode_roundtrips_and_is_a_write() { + assert_eq!(OpCode::try_from(0x9C), Ok(OpCode::ArrayPutBatch)); + assert_eq!(u8::from(OpCode::ArrayPutBatch), 0x9C); + assert!(OpCode::ArrayPutBatch.is_write()); + } + #[test] fn crdt_list_opcodes_are_writes() { assert!(OpCode::CrdtListInsert.is_write()); diff --git a/nodedb-types/src/sync/wire/spatial.rs b/nodedb-types/src/sync/wire/spatial.rs index 81ed4fef4..351160d6a 100644 --- a/nodedb-types/src/sync/wire/spatial.rs +++ b/nodedb-types/src/sync/wire/spatial.rs @@ -23,8 +23,9 @@ use crate::sync::wire::ack_status::AckStatus; /// bounding box from the geometry, and inserts the entry into /// `CoreLoop::spatial_indexes`. /// -/// `geometry_bytes` is a MessagePack-serialised `nodedb_types::geometry::Geometry` -/// value produced via `zerompk::to_msgpack_vec(&geometry)`. +/// `geometry_bytes` is a MessagePack-serialised +/// `nodedb_types::geometry::SridGeometry` produced via +/// `zerompk::to_msgpack_vec(&geometry)`: `[srid_code_u32, geometry]`. #[derive( Debug, Clone, Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack, )] @@ -37,7 +38,7 @@ pub struct SpatialInsertMsg { pub field: String, /// External document identifier. pub doc_id: String, - /// MessagePack-serialised `Geometry` value. + /// MessagePack-serialised `SridGeometry` value. pub geometry_bytes: Vec, /// Monotonic batch ID (Lite-assigned, per-operation). Used for ACK correlation. pub batch_id: u64, diff --git a/nodedb-types/src/value/convert.rs b/nodedb-types/src/value/convert.rs index fa1cb78d5..cba1ec145 100644 --- a/nodedb-types/src/value/convert.rs +++ b/nodedb-types/src/value/convert.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use crate::datetime::{NdbDateTime, NdbDuration}; -use crate::geometry::Geometry; +use crate::geometry::{Geometry, SridGeometry}; use crate::value::core::Value; impl From<&str> for Value { @@ -92,8 +92,15 @@ impl From for Value { } } +/// An untagged shape is WGS 84, the SRID of untagged GeoJSON. impl From for Value { fn from(g: Geometry) -> Self { + Value::Geometry(SridGeometry::wgs84(g)) + } +} + +impl From for Value { + fn from(g: SridGeometry) -> Self { Value::Geometry(g) } } diff --git a/nodedb-types/src/value/core.rs b/nodedb-types/src/value/core.rs index 81ecab911..ea9c5f379 100644 --- a/nodedb-types/src/value/core.rs +++ b/nodedb-types/src/value/core.rs @@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize}; use crate::array_cell::ArrayCell; use crate::datetime::{NdbDateTime, NdbDuration}; -use crate::geometry::Geometry; +use crate::geometry::SridGeometry; use crate::json_msgpack::InstantKind; /// A dynamic value that can represent any field type in a document @@ -83,8 +83,9 @@ pub enum Value { Duration(NdbDuration), /// Arbitrary-precision decimal (financial calculations, exact arithmetic). Decimal(rust_decimal::Decimal), - /// GeoJSON-compatible geometry (Point, LineString, Polygon, etc.). - Geometry(Geometry), + /// GeoJSON-compatible geometry (Point, LineString, Polygon, etc.) and + /// the SRID its coordinates are in. + Geometry(SridGeometry), /// Ordered set of unique values (auto-deduplicated, maintains insertion order). Set(Vec), /// Compiled regex pattern (stored as pattern string). @@ -222,7 +223,7 @@ impl Value { } /// Try to extract as Geometry. - pub fn as_geometry(&self) -> Option<&Geometry> { + pub fn as_geometry(&self) -> Option<&SridGeometry> { match self { Value::Geometry(g) => Some(g), _ => None, diff --git a/nodedb-types/src/value/display.rs b/nodedb-types/src/value/display.rs index ee28dc518..549b65604 100644 --- a/nodedb-types/src/value/display.rs +++ b/nodedb-types/src/value/display.rs @@ -54,7 +54,7 @@ impl fmt::Display for Value { Value::DateTime(dt) | Value::NaiveDateTime(dt) => write!(f, "{dt}"), Value::Duration(d) => write!(f, "{d}"), Value::Decimal(d) => write!(f, "{d}"), - Value::Geometry(g) => write!(f, "{g:?}"), + Value::Geometry(g) => write!(f, "{g}"), Value::Range { start, end, diff --git a/nodedb-types/src/value/msgpack.rs b/nodedb-types/src/value/msgpack.rs index 27c961e07..cb59e8934 100644 --- a/nodedb-types/src/value/msgpack.rs +++ b/nodedb-types/src/value/msgpack.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use super::core::Value; use crate::array_cell::ArrayCell; use crate::datetime::{NdbDateTime, NdbDuration}; -use crate::geometry::Geometry; +use crate::geometry::{Geometry, Srid, SridGeometry}; impl zerompk::ToMessagePack for Value { fn write(&self, writer: &mut W) -> zerompk::Result<()> { @@ -80,10 +80,12 @@ impl zerompk::ToMessagePack for Value { writer.write_u8(12)?; writer.write_binary(&d.serialize()) } + // `[13, srid_code_u32, geometry]`. Value::Geometry(g) => { - writer.write_array_len(2)?; + writer.write_array_len(3)?; writer.write_u8(13)?; - g.write(writer) + g.srid.write(writer)?; + g.geometry.write(writer) } Value::Set(s) => { writer.write_array_len(2)?; @@ -163,7 +165,17 @@ impl<'a> zerompk::FromMessagePack<'a> for Value { buf.copy_from_slice(&cow); Ok(Value::Decimal(rust_decimal::Decimal::deserialize(buf))) } - 13 => Ok(Value::Geometry(Geometry::read(reader)?)), + 13 => { + if len != 3 { + return Err(zerompk::Error::ArrayLengthMismatch { + expected: 3, + actual: len, + }); + } + let srid = Srid::read(reader)?; + let geometry = Geometry::read(reader)?; + Ok(Value::Geometry(SridGeometry::new(srid, geometry))) + } 14 => Ok(Value::Set(Vec::::read(reader)?)), 15 => Ok(Value::Regex(reader.read_string()?.into_owned())), 16 => { @@ -207,3 +219,18 @@ impl<'a> zerompk::FromMessagePack<'a> for Value { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn geometry_round_trip_keeps_srid() { + for srid in Srid::SUPPORTED { + let value = Value::Geometry(SridGeometry::new(srid, Geometry::point(1.0, 2.0))); + let bytes = zerompk::to_msgpack_vec(&value).unwrap(); + let back: Value = zerompk::from_msgpack(&bytes).unwrap(); + assert_eq!(back, value); + } + } +} diff --git a/nodedb-types/src/value/sql_literal.rs b/nodedb-types/src/value/sql_literal.rs index 42a227e75..286300150 100644 --- a/nodedb-types/src/value/sql_literal.rs +++ b/nodedb-types/src/value/sql_literal.rs @@ -28,9 +28,9 @@ impl Value { } Value::Duration(value) => quote_literal(&value.to_string()), Value::Decimal(value) => value.to_string(), - // Debug is a total, deterministic representation for this enum and - // avoids a fallible JSON serializer on this infallible public API. - Value::Geometry(value) => quote_literal(&format!("{value:?}")), + // `SridGeometry`'s Display is total and deterministic, and avoids a + // fallible JSON serializer on this infallible public API. + Value::Geometry(value) => quote_literal(&value.to_string()), Value::Range { start, end, @@ -99,7 +99,7 @@ mod tests { use super::*; use crate::array_cell::ArrayCell; - use crate::geometry::Geometry; + use crate::geometry::{Geometry, Srid, SridGeometry}; #[test] fn renders_every_value_variant_without_null_fallback() { @@ -121,7 +121,7 @@ mod tests { Value::NaiveDateTime(crate::datetime::NdbDateTime::from_micros(0)), Value::Duration(crate::datetime::NdbDuration::from_micros(1)), Value::Decimal(rust_decimal::Decimal::new(125, 2)), - Value::Geometry(Geometry::point(1.0, 2.0)), + Value::Geometry(SridGeometry::wgs84(Geometry::point(1.0, 2.0))), Value::Set(vec![Value::Bool(false)]), Value::Regex("a'b".into()), Value::Range { @@ -212,9 +212,14 @@ mod tests { #[test] fn geometry_and_ranges_have_explicit_safe_representations() { assert_eq!( - Value::Geometry(Geometry::point(1.0, 2.0)).to_sql_literal(), + Value::Geometry(SridGeometry::wgs84(Geometry::point(1.0, 2.0))).to_sql_literal(), quote_literal("Point { coordinates: [1.0, 2.0] }") ); + assert_eq!( + Value::Geometry(SridGeometry::new(Srid::Utm48N, Geometry::point(1.0, 2.0))) + .to_sql_literal(), + quote_literal("SRID=32648;Point { coordinates: [1.0, 2.0] }") + ); assert_eq!( Value::Range { start: Some(Box::new(Value::String("a'b".into()))), diff --git a/nodedb-wal/src/record/fts_spatial.rs b/nodedb-wal/src/record/fts_spatial.rs index 9c67ff4b4..2c53f1624 100644 --- a/nodedb-wal/src/record/fts_spatial.rs +++ b/nodedb-wal/src/record/fts_spatial.rs @@ -304,8 +304,9 @@ impl FtsDeletePayload { /// WAL payload for `RecordType::SpatialPut`. /// -/// `geometry_bytes` is a MessagePack-serialised `nodedb_types::geometry::Geometry` -/// value (the same encoding used in `SpatialInsertMsg`). +/// `geometry_bytes` is a MessagePack-serialised +/// `nodedb_types::geometry::SridGeometry` value, `[srid_code_u32, geometry]` +/// (the same encoding used in `SpatialInsertMsg`). #[derive(Debug, Clone, PartialEq, Eq)] pub struct SpatialPutPayload { pub provenance: SyncProvenance, diff --git a/nodedb/src/bootstrap/background_loops/timers.rs b/nodedb/src/bootstrap/background_loops/timers.rs index 82b0c5a86..32a6bcd93 100644 --- a/nodedb/src/bootstrap/background_loops/timers.rs +++ b/nodedb/src/bootstrap/background_loops/timers.rs @@ -44,6 +44,40 @@ pub fn spawn_tick_loop( ); } +/// Spawn the audit log flush loop. +/// +/// The loop flushes every `AUDIT_FLUSH_PERIOD` and also when the buffer requests an early flush. +/// It is the flush executor: recorders past the backpressure point wait for its flushes. +/// On shutdown it flushes once more, so a clean stop leaves no entry only in memory. +fn spawn_audit_flush_loop(shared: &Arc) { + let state = Arc::clone(shared); + let control = match shared.audit.lock() { + Ok(log) => log.flush_control(), + Err(poisoned) => poisoned.into_inner().flush_control(), + }; + spawn_loop( + &shared.loop_registry, + &shared.shutdown, + "audit_log_flush", + ShutdownPhase::DrainingControlPlane, + move |mut shutdown| async move { + control.attach_executor(); + let mut tick = tokio::time::interval(AUDIT_FLUSH_PERIOD); + loop { + tokio::select! { + _ = shutdown.wait_cancelled() => break, + _ = tick.tick() => {} + _ = control.flush_requested() => {} + } + control.take_flush_request(); + state.flush_audit_log(); + } + state.flush_audit_log(); + control.detach_executor(); + }, + ); +} + /// Spawn the tenant rate reset, audit log flush, and tenant memory /// estimate timers. pub fn spawn_tenant_timers(shared: &Arc) { @@ -53,12 +87,7 @@ pub fn spawn_tenant_timers(shared: &Arc) { TENANT_RATE_RESET_PERIOD, SharedState::reset_tenant_rate_counters, ); - spawn_tick_loop( - shared, - "audit_log_flush", - AUDIT_FLUSH_PERIOD, - SharedState::flush_audit_log, - ); + spawn_audit_flush_loop(shared); spawn_tick_loop( shared, "tenant_memory_estimate", diff --git a/nodedb/src/bridge/envelope/error_code_from.rs b/nodedb/src/bridge/envelope/error_code_from.rs index a15b5be85..855c68fe9 100644 --- a/nodedb/src/bridge/envelope/error_code_from.rs +++ b/nodedb/src/bridge/envelope/error_code_from.rs @@ -15,7 +15,10 @@ impl From for ErrorCode { nodedb_query::EvalError::UnknownFunction { name } => Self::UndefinedFunction { name }, e @ (nodedb_query::EvalError::VectorDimensionMismatch { .. } | nodedb_query::EvalError::ArgumentType { .. } - | nodedb_query::EvalError::InvalidJsonPath { .. }) => Self::DataException { + | nodedb_query::EvalError::InvalidJsonPath { .. } + | nodedb_query::EvalError::UnsupportedSrid { .. } + | nodedb_query::EvalError::MixedSrid { .. } + | nodedb_query::EvalError::ProjectionDomain { .. }) => Self::DataException { detail: e.to_string(), }, e @ nodedb_query::EvalError::NumericOverflow { .. } => Self::NumericValueOutOfRange { @@ -160,7 +163,8 @@ impl From for ErrorCode { // Class `22`, the class the Control Plane gives both. e @ (crate::Error::OffsetRegression { .. } | crate::Error::BackupTenantMismatch { .. } - | crate::Error::InvalidLimitValue { .. }) => Self::DataException { + | crate::Error::InvalidLimitValue { .. } + | crate::Error::MixedSrid { .. }) => Self::DataException { detail: e.to_string(), }, // `22003`, as the Control Plane gives it. @@ -266,6 +270,7 @@ impl From for ErrorCode { | crate::Error::MirrorReadOnly { .. } | crate::Error::WrongObjectType { .. } | crate::Error::AmbiguousColumn { .. } + | crate::Error::UngroupedColumn { .. } | crate::Error::ExecutionLimitExceeded { .. } | crate::Error::LimitExceeded { .. } | crate::Error::Promql(_) diff --git a/nodedb/src/control/array_sync/inbound.rs b/nodedb/src/control/array_sync/inbound.rs index 893f8c32a..36fcfa32e 100644 --- a/nodedb/src/control/array_sync/inbound.rs +++ b/nodedb/src/control/array_sync/inbound.rs @@ -380,6 +380,41 @@ impl OriginArrayInbound { Some(_) => {} } + // 2a. A put's attrs must match the schema before the op is proposed: + // a refusal after commit would refuse again on every replay. + if let Some(attrs) = op.attrs.as_deref() + && op.kind == nodedb_array::sync::op::ArrayOpKind::Put + { + let checked = self + .engine + .schemas + .to_array_schema_in_database( + self.database_id, + self.tenant_id.as_u64(), + &op.header.array, + ) + .map(|schema| schema.check_cell(attrs)); + match checked { + Some(Ok(())) => {} + Some(Err(e)) => { + return Err(build_reject( + &op.header.array, + op.header.hlc, + ArrayRejectReason::EngineRefused, + format!("schema check: {e}"), + )); + } + None => { + return Err(build_reject( + &op.header.array, + op.header.hlc, + ArrayRejectReason::ArrayUnknown, + format!("array '{}' has no decodable schema", op.header.array), + )); + } + } + } + // 3. Fast-path idempotency check (before proposing). if self.engine.already_seen_in_database( self.database_id, diff --git a/nodedb/src/control/array_sync/raft_apply/op.rs b/nodedb/src/control/array_sync/raft_apply/op.rs index b01487b49..24b4e458a 100644 --- a/nodedb/src/control/array_sync/raft_apply/op.rs +++ b/nodedb/src/control/array_sync/raft_apply/op.rs @@ -187,8 +187,13 @@ pub(crate) async fn apply_array_op( } } ArrayOpKind::Delete | ArrayOpKind::Erase => { - let coords = vec![op.coord.clone()]; - let coords_msgpack = match zerompk::to_msgpack_vec(&coords) { + // The Data Plane and the WAL record read `Vec`. + let cells = vec![crate::engine::array::wal::ArrayDeleteCell { + coord: op.coord.clone(), + system_from_ms: op.header.system_from_ms, + erasure: matches!(op.kind, ArrayOpKind::Erase), + }]; + let coords_msgpack = match zerompk::to_msgpack_vec(&cells) { Ok(b) => b, Err(e) => { warn!(group_id, index = log_index, error = %e, "apply_array_op: coords encode failed"); diff --git a/nodedb/src/control/clone/resolver/rewrite_engine.rs b/nodedb/src/control/clone/resolver/rewrite_engine.rs index b4566ab0e..7d8224934 100644 --- a/nodedb/src/control/clone/resolver/rewrite_engine.rs +++ b/nodedb/src/control/clone/resolver/rewrite_engine.rs @@ -242,6 +242,7 @@ pub(super) fn rewrite_timeseries( valid_at_ms, having, tiering, + output_columns, } if ctx.is_target(collection) => Ok(SourceRewrite::task(PhysicalPlan::Timeseries( TimeseriesOp::Scan { collection: ctx.source_qualified.clone(), @@ -260,6 +261,7 @@ pub(super) fn rewrite_timeseries( valid_at_ms: *valid_at_ms, having: having.clone(), tiering: tiering.clone(), + output_columns: output_columns.clone(), }, ))), diff --git a/nodedb/src/control/cluster/array_cluster_exec/dispatch.rs b/nodedb/src/control/cluster/array_cluster_exec/dispatch.rs index 4a14acece..c54d3da26 100644 --- a/nodedb/src/control/cluster/array_cluster_exec/dispatch.rs +++ b/nodedb/src/control/cluster/array_cluster_exec/dispatch.rs @@ -96,14 +96,17 @@ impl ShardRpcDispatch for NexarArrayDispatch { } /// Whether a second run of the request leaves the same state and answer. A -/// slice, an aggregate and a surrogate bitmap only read the shard. A put or -/// a delete can have applied when its answer was lost. +/// slice, an aggregate, an elementwise op and a surrogate bitmap only read +/// the shard. A second flush or compaction finds nothing left to do. A put +/// or a delete can have applied when its answer was lost. fn is_resend_safe(msg_type: VShardMessageType) -> bool { matches!( msg_type, VShardMessageType::ArrayShardSliceReq | VShardMessageType::ArrayShardAggReq + | VShardMessageType::ArrayShardElementwiseReq | VShardMessageType::ArrayShardSurrogateBitmapReq + | VShardMessageType::ArrayShardMaintReq ) } @@ -215,11 +218,14 @@ mod tests { } } - /// A slice, an aggregate and a surrogate bitmap only read the shard. + /// A slice, an aggregate, an elementwise op and a surrogate bitmap only + /// read the shard. A repeated flush or compaction finds nothing to do. #[test] - fn only_array_reads_are_resend_safe() { + fn only_array_reads_and_maintenance_are_resend_safe() { + assert!(is_resend_safe(VShardMessageType::ArrayShardMaintReq)); assert!(is_resend_safe(VShardMessageType::ArrayShardSliceReq)); assert!(is_resend_safe(VShardMessageType::ArrayShardAggReq)); + assert!(is_resend_safe(VShardMessageType::ArrayShardElementwiseReq)); assert!(is_resend_safe( VShardMessageType::ArrayShardSurrogateBitmapReq )); diff --git a/nodedb/src/control/cluster/array_cluster_exec/elementwise.rs b/nodedb/src/control/cluster/array_cluster_exec/elementwise.rs new file mode 100644 index 000000000..39405ff38 --- /dev/null +++ b/nodedb/src/control/cluster/array_cluster_exec/elementwise.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `ClusterArrayOp::Elementwise`: fan the op out to every primary vShard. +//! +//! Each shard pairs the cells of both arrays in its own bucket on its Data +//! Plane. The coordinator concatenates the rows into the one msgpack array a +//! local `ArrayOp::Elementwise` answers with. + +use std::sync::Arc; + +use nodedb_cluster::distributed_array::coordinator::ArrayCoordinator; +use nodedb_cluster::distributed_array::wire::ArrayShardElementwiseReq; +use nodedb_physical::physical_plan::ArrayBinaryOp; + +use crate::control::cluster::array_cluster_helpers::{ + cluster_err, coord_err, encode_err, encode_row_array, +}; +use crate::types::TxnId; + +use super::executor::{ARRAY_RPC_TIMEOUT_MS, ClusterArrayExecutor, note_array_reads}; + +pub(super) struct ElementwiseArgs<'a> { + pub left: &'a nodedb_array::types::ArrayId, + pub right: &'a nodedb_array::types::ArrayId, + pub op: ArrayBinaryOp, + pub attr_idx: u32, + pub prefix_bits: u8, + pub system_as_of: Option, + pub valid_at_ms: Option, + pub txn_id: Option, +} + +impl ClusterArrayExecutor { + pub(super) async fn execute_elementwise( + &self, + args: ElementwiseArgs<'_>, + ) -> crate::Result> { + let ElementwiseArgs { + left, + right, + op, + attr_idx, + prefix_bits, + system_as_of, + valid_at_ms, + txn_id, + } = args; + // No tile keys: an elementwise op reads every tile of both arrays. + let coordinator = ArrayCoordinator::for_slice( + self.source_node, + ARRAY_RPC_TIMEOUT_MS, + &[], + prefix_bits, + self.total_shards, + Arc::clone(&self.dispatch), + Arc::clone(&self.circuit_breaker), + ) + .map_err(cluster_err)?; + + let req = ArrayShardElementwiseReq { + left_id_msgpack: zerompk::to_msgpack_vec(left).map_err(encode_err)?, + right_id_msgpack: zerompk::to_msgpack_vec(right).map_err(encode_err)?, + op_msgpack: zerompk::to_msgpack_vec(&op).map_err(encode_err)?, + attr_idx, + // coord_elementwise stamps each shard's bucket. + shard_bucket: None, + system_as_of, + valid_at_ms, + txn_id: txn_id.map(TxnId::as_u64), + }; + let mut result = coordinator + .coord_elementwise(req) + .await + .map_err(coord_err)?; + // Each leg read both arrays on its vShard. + note_array_reads(left, &result.read_versions); + note_array_reads(right, &result.read_versions); + // A temporal read ends with one summary row, as a local elementwise + // op answers. + if system_as_of.is_some() || valid_at_ms.is_some() { + let summary = nodedb_types::json_to_msgpack(&serde_json::json!({ + "truncated_before_horizon": result.truncated_before_horizon + })) + .map_err(encode_err)?; + result.rows.push(summary); + } + Ok(encode_row_array(&result.rows)) + } +} diff --git a/nodedb/src/control/cluster/array_cluster_exec/executor.rs b/nodedb/src/control/cluster/array_cluster_exec/executor.rs index ea3680a59..b5f4a49dc 100644 --- a/nodedb/src/control/cluster/array_cluster_exec/executor.rs +++ b/nodedb/src/control/cluster/array_cluster_exec/executor.rs @@ -12,12 +12,14 @@ use nodedb_cluster::distributed_array::coordinator::{ ArrayCoordinator, ArrayWriteCoordParams, coord_delete, coord_put, }; use nodedb_cluster::distributed_array::rpc::ShardRpcDispatch; -use nodedb_cluster::distributed_array::wire::{ArrayShardAggReq, ArrayShardSliceReq}; +use nodedb_cluster::distributed_array::wire::{ + ArrayShardAggReq, ArrayShardMaintOp, ArrayShardSliceReq, +}; use nodedb_cluster::routing::VSHARD_COUNT; use nodedb_cluster::{NexarTransport, RoutingTable}; use crate::control::cluster::array_cluster_helpers::{ - cluster_err, coord_err, encode_err, finalize_agg_partials, + cluster_err, coord_err, encode_err, encode_row_array, finalize_agg_partials, }; use crate::control::server::shared::session::pending_shard_reads::{ShardObservation, ShardReads}; use crate::control::server::shared::session::read_set::EngineTag; @@ -30,17 +32,17 @@ use zerompk; use super::dispatch::NexarArrayDispatch; /// Default RPC timeout for array cluster operations in milliseconds. -const ARRAY_RPC_TIMEOUT_MS: u64 = 30_000; +pub(super) const ARRAY_RPC_TIMEOUT_MS: u64 = 30_000; /// Control-Plane executor for distributed array operations. /// /// Constructed per-request from `SharedState` by the routing interceptor. /// Holds references to the transport and routing table; is not cached. pub struct ClusterArrayExecutor { - dispatch: Arc, - circuit_breaker: Arc, - source_node: u64, - total_shards: u32, + pub(super) dispatch: Arc, + pub(super) circuit_breaker: Arc, + pub(super) source_node: u64, + pub(super) total_shards: u32, } struct SliceArgs<'a> { @@ -64,6 +66,7 @@ struct AggArgs<'a> { prefix_bits: u8, system_as_of: Option, valid_at_ms: Option, + region_msgpack: &'a [u8], txn_id: Option, } @@ -142,6 +145,7 @@ impl ClusterArrayExecutor { prefix_bits, system_as_of, valid_at_ms, + region_msgpack, } => { self.execute_agg(AggArgs { array_id, @@ -152,6 +156,29 @@ impl ClusterArrayExecutor { prefix_bits: *prefix_bits, system_as_of: *system_as_of, valid_at_ms: *valid_at_ms, + region_msgpack, + txn_id, + }) + .await + } + + ClusterArrayOp::Elementwise { + left, + right, + op, + attr_idx, + prefix_bits, + system_as_of, + valid_at_ms, + } => { + self.execute_elementwise(super::elementwise::ElementwiseArgs { + left, + right, + op: *op, + attr_idx: *attr_idx, + prefix_bits: *prefix_bits, + system_as_of: *system_as_of, + valid_at_ms: *valid_at_ms, txn_id, }) .await @@ -178,6 +205,25 @@ impl ClusterArrayExecutor { self.execute_delete(array_id_msgpack, coords, *wal_lsn, *prefix_bits) .await } + + ClusterArrayOp::Flush { + array_id, + prefix_bits, + } => { + self.execute_maint(array_id, ArrayShardMaintOp::Flush, *prefix_bits) + .await + } + + ClusterArrayOp::Compact { + array_id, + audit_retain_ms, + prefix_bits, + } => { + let op = ArrayShardMaintOp::Compact { + audit_retain_ms: *audit_retain_ms, + }; + self.execute_maint(array_id, op, *prefix_bits).await + } } } @@ -233,23 +279,8 @@ impl ClusterArrayExecutor { // `truncated_before_horizon` flag is OR-reduced across shards by the // coordinator so the upstream NOTICE is emitted whenever any shard // dropped data below its system-time horizon. - let mut rows_msgpack = - Vec::with_capacity(5 + result.rows.iter().map(|r| r.len()).sum::()); - let n = result.rows.len(); - if n <= 15 { - rows_msgpack.push(0x90 | (n as u8)); - } else if n <= u16::MAX as usize { - rows_msgpack.push(0xdc); - rows_msgpack.extend_from_slice(&(n as u16).to_be_bytes()); - } else { - rows_msgpack.push(0xdd); - rows_msgpack.extend_from_slice(&(n as u32).to_be_bytes()); - } - for r in &result.rows { - rows_msgpack.extend_from_slice(r); - } let resp = crate::data::executor::response_codec::ArraySliceResponse { - rows_msgpack, + rows_msgpack: encode_row_array(&result.rows), truncated_before_horizon: result.truncated_before_horizon, }; zerompk::to_msgpack_vec(&resp).map_err(|e| crate::Error::Codec { @@ -267,6 +298,7 @@ impl ClusterArrayExecutor { prefix_bits, system_as_of, valid_at_ms, + region_msgpack, txn_id, } = args; let coordinator = ArrayCoordinator::for_slice( @@ -291,6 +323,7 @@ impl ClusterArrayExecutor { system_as_of, valid_at_ms, txn_id: txn_id.map(TxnId::as_u64), + region_msgpack: region_msgpack.to_vec(), }; let agg = coordinator.coord_agg(req).await.map_err(coord_err)?; note_array_reads(array_id, &agg.read_versions); @@ -315,7 +348,7 @@ impl ClusterArrayExecutor { group_by_dim, agg.truncated_before_horizon, emit_horizon, - ); + )?; zerompk::to_msgpack_vec(&rows).map_err(encode_err) } @@ -347,6 +380,16 @@ impl ClusterArrayExecutor { // actually wrote. `{"inserted": n}` matches what the local // `ArrayOp::Put` handler emits, so `extract_affected_count` reads both. let affected: u64 = resps.iter().map(|r| r.affected).sum(); + // Every cell of a put writes one cell version, a repeated coordinate + // included. Fewer means a shard dropped cells it acknowledged. + if affected != cells.len() as u64 { + return Err(crate::Error::Internal { + detail: format!( + "array put: the shards wrote {affected} of {} cells", + cells.len() + ), + }); + } response_codec::encode_count("inserted", affected as usize) } @@ -388,7 +431,7 @@ impl ClusterArrayExecutor { /// note when the request ends (`session::pending_shard_reads::record_pending`), so /// commit validation checks the read on each covered vShard: inside the /// Calvin barrier, or against each vShard's leader on a local commit. -fn note_array_reads( +pub(super) fn note_array_reads( array_id: &nodedb_array::types::ArrayId, read_versions: &[nodedb_types::ShardVersion], ) { diff --git a/nodedb/src/control/cluster/array_cluster_exec/maint.rs b/nodedb/src/control/cluster/array_cluster_exec/maint.rs new file mode 100644 index 000000000..bf6019afc --- /dev/null +++ b/nodedb/src/control/cluster/array_cluster_exec/maint.rs @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `ClusterArrayOp::Flush` / `Compact`: fan the maintenance out to every +//! primary vShard, so every core that holds tiles of the array runs it. + +use std::sync::Arc; + +use nodedb_cluster::distributed_array::coordinator::ArrayCoordinator; +use nodedb_cluster::distributed_array::wire::{ArrayShardMaintOp, ArrayShardMaintReq}; + +use crate::control::cluster::array_cluster_helpers::{cluster_err, coord_err, encode_err}; +use crate::data::executor::response_codec; + +use super::executor::{ARRAY_RPC_TIMEOUT_MS, ClusterArrayExecutor}; + +impl ClusterArrayExecutor { + /// Answer `{"flushed": 1}` for a flush, as a local flush does. Answer + /// `{"compacted": n}` for a compaction, where `n` counts the shards + /// that merged segments. + pub(super) async fn execute_maint( + &self, + array_id: &nodedb_array::types::ArrayId, + op: ArrayShardMaintOp, + prefix_bits: u8, + ) -> crate::Result> { + // No tile keys: maintenance covers every tile of the array. + let coordinator = ArrayCoordinator::for_slice( + self.source_node, + ARRAY_RPC_TIMEOUT_MS, + &[], + prefix_bits, + self.total_shards, + Arc::clone(&self.dispatch), + Arc::clone(&self.circuit_breaker), + ) + .map_err(cluster_err)?; + let is_flush = op == ArrayShardMaintOp::Flush; + let req = ArrayShardMaintReq { + array_id_msgpack: zerompk::to_msgpack_vec(array_id).map_err(encode_err)?, + op, + }; + let resps = coordinator.coord_maint(req).await.map_err(coord_err)?; + if is_flush { + response_codec::encode_count("flushed", 1) + } else { + let merged: u64 = resps.iter().map(|r| r.affected).sum(); + response_codec::encode_count("compacted", merged as usize) + } + } +} diff --git a/nodedb/src/control/cluster/array_cluster_exec/mod.rs b/nodedb/src/control/cluster/array_cluster_exec/mod.rs index b4ea7d0d8..50a09a0b5 100644 --- a/nodedb/src/control/cluster/array_cluster_exec/mod.rs +++ b/nodedb/src/control/cluster/array_cluster_exec/mod.rs @@ -8,9 +8,15 @@ //! - [`executor`]: `ClusterArrayExecutor` — wraps the dispatch and routing //! table, constructs the appropriate `ArrayCoordinator`, and converts //! results into raw response bytes. +//! - [`elementwise`]: the `Elementwise` fan-out of `ClusterArrayExecutor`. +//! - [`maint`]: the `Flush` / `Compact` fan-out of `ClusterArrayExecutor`. +//! - [`surrogate`]: the cross-shard surrogate bitmap of an array slice. mod dispatch; +mod elementwise; mod executor; +mod maint; +mod surrogate; pub use dispatch::NexarArrayDispatch; pub use executor::ClusterArrayExecutor; diff --git a/nodedb/src/control/cluster/array_cluster_exec/surrogate.rs b/nodedb/src/control/cluster/array_cluster_exec/surrogate.rs new file mode 100644 index 000000000..8d8d14918 --- /dev/null +++ b/nodedb/src/control/cluster/array_cluster_exec/surrogate.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The surrogate bitmap of an array slice, read from every shard that owns +//! the array's tiles. +//! +//! A vector search with an `ARRAY_SLICE` prefilter runs on the core that +//! owns its collection. That core holds only its own buckets of the array, +//! so the prefilter is read here across every primary vShard and handed to +//! the search as its candidate bitmap. + +use std::sync::Arc; + +use nodedb_cluster::distributed_array::coordinator::ArrayCoordinator; +use nodedb_cluster::distributed_array::wire::ArrayShardSurrogateBitmapReq; +use nodedb_types::SurrogateBitmap; + +use crate::control::cluster::array_cluster_helpers::{cluster_err, coord_err, encode_err}; + +use super::executor::{ARRAY_RPC_TIMEOUT_MS, ClusterArrayExecutor}; + +impl ClusterArrayExecutor { + /// The union of every shard's surrogate bitmap for the cells of + /// `array_id` inside `slice_msgpack`. + /// + /// Every primary vShard answers. vShards that share one Data Plane + /// report the same surrogates, which the union keeps once. + pub async fn surrogate_bitmap( + &self, + array_id: &nodedb_array::types::ArrayId, + slice_msgpack: &[u8], + prefix_bits: u8, + ) -> crate::Result { + let coordinator = ArrayCoordinator::for_slice( + self.source_node, + ARRAY_RPC_TIMEOUT_MS, + &[], + prefix_bits, + self.total_shards, + Arc::clone(&self.dispatch), + Arc::clone(&self.circuit_breaker), + ) + .map_err(cluster_err)?; + let req = ArrayShardSurrogateBitmapReq { + array_id_msgpack: zerompk::to_msgpack_vec(array_id).map_err(encode_err)?, + slice_msgpack: slice_msgpack.to_vec(), + }; + let resps = coordinator + .coord_surrogate_bitmap_scan(req) + .await + .map_err(coord_err)?; + let mut union = SurrogateBitmap::new(); + for resp in resps { + let bitmap: SurrogateBitmap = + zerompk::from_msgpack(&resp.bitmap_msgpack).map_err(|e| crate::Error::Codec { + detail: format!( + "array surrogate bitmap from vShard {} decode: {e}", + resp.shard_id + ), + })?; + union.union_in_place(&bitmap); + } + Ok(union) + } +} diff --git a/nodedb/src/control/cluster/array_cluster_helpers.rs b/nodedb/src/control/cluster/array_cluster_helpers.rs index c8a4ddee8..3c1b3ecdb 100644 --- a/nodedb/src/control/cluster/array_cluster_helpers.rs +++ b/nodedb/src/control/cluster/array_cluster_helpers.rs @@ -6,6 +6,7 @@ use nodedb_cluster::distributed_array::merge::ArrayAggPartial; use nodedb_cluster::distributed_array::scatter::{ArrayCoordError, ShardFailure}; +use nodedb_cluster::distributed_array::{ArrayAggGroupKey, ArrayAggNumber}; use nodedb_cluster::error::ClusterError; use nodedb_cluster::wire::VShardMessageType; @@ -19,7 +20,11 @@ use crate::Error; pub(super) enum AggValue { Float(f64), Int(i64), + Str(String), Bool(bool), + /// A timestamp dim value in microseconds, written as the naive-instant + /// ext a slice row writes for the same dim. + Timestamp(i64), Null, } @@ -28,17 +33,43 @@ impl zerompk::ToMessagePack for AggValue { match self { AggValue::Float(f) => writer.write_f64(*f), AggValue::Int(i) => writer.write_i64(*i), + AggValue::Str(s) => writer.write_string(s), AggValue::Bool(b) => writer.write_boolean(*b), + AggValue::Timestamp(micros) => writer.write_ext( + nodedb_types::InstantKind::Naive.ext_type(), + µs.to_be_bytes(), + ), AggValue::Null => writer.write_nil(), } } } +/// The `group` cell of a grouped row: the dim value, typed as the single-node +/// path writes it. A timestamp dim's milliseconds become a timestamp. A +/// scalar partial carries no group, so it writes NULL. +fn group_cell(key: &ArrayAggGroupKey) -> Result { + Ok(match key { + ArrayAggGroupKey::Int64(v) => AggValue::Int(*v), + ArrayAggGroupKey::TimestampMs(ms) => { + let dt = nodedb_types::NdbDateTime::from_millis(*ms).map_err(|e| Error::Codec { + detail: format!("array aggregate: timestamp group {ms} ms: {e}"), + })?; + AggValue::Timestamp(dt.micros) + } + ArrayAggGroupKey::Float64(v) => AggValue::Float(*v), + ArrayAggGroupKey::String(s) => AggValue::Str(s.clone()), + ArrayAggGroupKey::Scalar => AggValue::Null, + }) +} + /// Finalize merged `ArrayAggPartial`s into the same msgpack-map structure /// the local `ArrayOp::Aggregate` path produces via `encode_agg_rows`. /// -/// For scalar (no group-by), returns a single `[{"result": f64}]`. -/// For group-by, returns one `{"group": i64, "result": f64}` per partial. +/// For scalar (no group-by), returns a single `[{"result": }]`. +/// For group-by, returns one `{"group": , "result": }` per +/// partial, the dim value typed as its dim. `COUNT` and an integer attr's +/// `SUM` / `MIN` / `MAX` are integers, and a `SUM` past `i64` refuses with +/// SQLSTATE `22003`, as on a single node. `MEAN` is a float. /// When `emit_horizon` is set (temporal queries only), a trailing /// `{"truncated_before_horizon": bool}` summary row is appended. This mirrors /// the single-node path (`dispatch::array::aggregate`) exactly — including the @@ -50,21 +81,21 @@ pub(super) fn finalize_agg_partials( group_by_dim: i32, truncated_before_horizon: bool, emit_horizon: bool, -) -> Vec> { +) -> Result>, Error> { use nodedb_physical::physical_plan::ArrayReducer; - let finalize = |p: &ArrayAggPartial| -> AggValue { + let finalize = |p: &ArrayAggPartial| -> Result { if p.count == 0 { - return AggValue::Null; + return Ok(AggValue::Null); } - let v = match reducer { + let n = match reducer { ArrayReducer::Sum => p.sum, - ArrayReducer::Count => p.count as f64, + ArrayReducer::Count => return count_value(p.count), ArrayReducer::Min => p.min, ArrayReducer::Max => p.max, - ArrayReducer::Mean => p.welford_mean, + ArrayReducer::Mean => return Ok(AggValue::Float(p.welford_mean)), }; - AggValue::Float(v) + number_value(n) }; let is_grouped = group_by_dim >= 0; @@ -74,12 +105,12 @@ pub(super) fn finalize_agg_partials( .map(|p| { let mut row = std::collections::BTreeMap::new(); if is_grouped { - row.insert("group".to_string(), AggValue::Int(p.group_key)); + row.insert("group".to_string(), group_cell(&p.group_key)?); } - row.insert("result".to_string(), finalize(p)); - row + row.insert("result".to_string(), finalize(p)?); + Ok(row) }) - .collect(); + .collect::>()?; // Trailing summary row carrying the below-horizon signal — only for // temporal queries, matching the single-node path so both topologies emit @@ -93,7 +124,50 @@ pub(super) fn finalize_agg_partials( rows.push(summary); } - rows + Ok(rows) +} + +fn out_of_range(detail: String) -> Error { + Error::NumericValueOutOfRange { + detail: format!("array aggregate: {detail}"), + } +} + +/// The final value of a partial's number: an exact integer, refused past +/// `i64`, or a float. +fn number_value(n: ArrayAggNumber) -> Result { + match n.as_i128() { + Some(v) => i64::try_from(v) + .map(AggValue::Int) + .map_err(|_| out_of_range(format!("value {v} is out of range for bigint"))), + None => Ok(AggValue::Float(n.as_f64())), + } +} + +fn count_value(count: u64) -> Result { + i64::try_from(count) + .map(AggValue::Int) + .map_err(|_| out_of_range(format!("count {count} is out of range for bigint"))) +} + +/// One msgpack array holding the already-encoded `rows`, the shape a local +/// array read answers with. +pub(super) fn encode_row_array(rows: &[Vec]) -> Vec { + let mut out = Vec::with_capacity(5 + rows.iter().map(Vec::len).sum::()); + let n = rows.len(); + if n <= 15 { + out.push(0x90 | (n as u8)); + } else if n <= u16::MAX as usize { + out.push(0xdc); + out.extend_from_slice(&(n as u16).to_be_bytes()); + } else { + out.push(0xdd); + out.extend_from_slice(&(n as u32).to_be_bytes()); + } + for row in rows { + out.extend_from_slice(row); + } + out } /// Map an array coordinator error to an [`Error`]. @@ -202,7 +276,8 @@ pub(super) fn encode_err(e: zerompk::Error) -> Error { } } -/// Map a numeric response opcode (81, 83, 85, 87, 89) to its `VShardMessageType`. +/// Map a numeric response opcode (81, 83, 85, 87, 89, 91, 93) to its +/// `VShardMessageType`. /// /// Used by the local-dispatch fast path to build the response envelope without /// going through the QUIC transport layer. @@ -213,6 +288,8 @@ pub(super) fn array_resp_msg_type(opcode: u32) -> Option { 85 => Some(VShardMessageType::ArrayShardPutResp), 87 => Some(VShardMessageType::ArrayShardDeleteResp), 89 => Some(VShardMessageType::ArrayShardSurrogateBitmapResp), + 91 => Some(VShardMessageType::ArrayShardElementwiseResp), + 93 => Some(VShardMessageType::ArrayShardMaintResp), _ => None, } } @@ -222,6 +299,81 @@ mod tests { use super::*; use crate::bridge::envelope::ErrorCode; + /// A grouped row writes its dim value typed as the dim: a Float64 dim + /// writes the float and a String dim the string, never a bit pattern or + /// a hash. + #[test] + fn a_grouped_row_writes_the_typed_dim_value() { + use nodedb_physical::physical_plan::ArrayReducer; + let partials = vec![ + ArrayAggPartial::from_single(ArrayAggGroupKey::Float64(101.5), 2.0), + ArrayAggPartial::from_single(ArrayAggGroupKey::String("chrX".into()), 3.0), + ArrayAggPartial::from_single(ArrayAggGroupKey::Int64(7), 4.0), + ]; + let rows = finalize_agg_partials(&partials, &ArrayReducer::Sum, 0, false, false) + .expect("finalize"); + let bytes = zerompk::to_msgpack_vec(&rows).expect("encode rows"); + let json = nodedb_types::msgpack_to_json_string(&bytes).expect("rows to json"); + let rows: serde_json::Value = serde_json::from_str(&json).expect("parse json"); + assert_eq!(rows[0]["group"].as_f64(), Some(101.5)); + assert_eq!(rows[1]["group"].as_str(), Some("chrX")); + assert_eq!(rows[2]["group"].as_i64(), Some(7)); + assert_eq!(rows[2]["result"].as_f64(), Some(4.0)); + } + + /// A timestamp group writes the naive instant a slice row writes for the + /// same dim, never the raw milliseconds. + #[test] + fn a_timestamp_group_writes_a_naive_instant() { + use nodedb_physical::physical_plan::ArrayReducer; + let ms = 1_768_435_200_000; + let partials = vec![ArrayAggPartial::from_single( + ArrayAggGroupKey::TimestampMs(ms), + 2.0, + )]; + let rows = finalize_agg_partials(&partials, &ArrayReducer::Sum, 0, false, false) + .expect("finalize"); + let bytes = zerompk::to_msgpack_vec(&rows).expect("encode rows"); + let decoded = nodedb_types::json_msgpack::value_from_msgpack(&bytes).expect("decode"); + let nodedb_types::Value::Array(rows) = decoded else { + panic!("rows are an array: {decoded:?}"); + }; + let nodedb_types::Value::Object(row) = &rows[0] else { + panic!("row is a map: {:?}", rows[0]); + }; + let want = nodedb_types::NdbDateTime::from_millis(ms).expect("in range"); + assert_eq!( + row.get("group"), + Some(&nodedb_types::Value::NaiveDateTime(want)) + ); + } + + /// An integer SUM finalizes as an exact integer, and one past `i64` + /// refuses with SQLSTATE 22003 as on a single node. + #[test] + fn an_integer_sum_is_exact_and_refuses_past_bigint() { + use nodedb_physical::physical_plan::ArrayReducer; + let big = (1i64 << 53) + 1; + let mut exact = ArrayAggPartial::from_single_int(ArrayAggGroupKey::Scalar, big); + exact.merge(&ArrayAggPartial::from_single_int( + ArrayAggGroupKey::Scalar, + 0, + )); + let rows = finalize_agg_partials(&[exact], &ArrayReducer::Sum, -1, false, false) + .expect("finalize"); + assert!(matches!(rows[0].get("result"), Some(AggValue::Int(v)) if *v == big)); + + let mut over = ArrayAggPartial::from_single_int(ArrayAggGroupKey::Scalar, i64::MAX); + over.merge(&ArrayAggPartial::from_single_int( + ArrayAggGroupKey::Scalar, + 1, + )); + assert!(matches!( + finalize_agg_partials(&[over], &ArrayReducer::Sum, -1, false, false), + Err(Error::NumericValueOutOfRange { .. }) + )); + } + /// A shard request that never went out answers `NoLeader` for that /// shard: nothing ran, so the refusal is definite. #[test] diff --git a/nodedb/src/control/cluster/array_executor/cells.rs b/nodedb/src/control/cluster/array_executor/cells.rs deleted file mode 100644 index 1bf4d3f6e..000000000 --- a/nodedb/src/control/cluster/array_executor/cells.rs +++ /dev/null @@ -1,36 +0,0 @@ -// SPDX-License-Identifier: BUSL-1.1 - -//! Reshape the array coordinator's per-shard cell payload. -//! -//! The coordinator partitions a write into buckets of separately-encoded -//! cells (`Vec`); `ArrayOp::Put` / `ArrayOp::Delete` carry one -//! flat msgpack array instead. Every consumer of a bucket reshapes it here. - -use zerompk::{FromMessagePackOwned, ToMessagePack}; - -/// Decode a coordinator bucket (`Vec`) and re-encode it as the -/// flat `Vec` msgpack array the Data Plane array handlers decode. `what` -/// names the payload in any error the reshape reports. -pub fn flatten_blob_vec(blob_vec_msgpack: &[u8], what: &str) -> crate::Result> -where - T: FromMessagePackOwned + ToMessagePack, -{ - let blobs: Vec> = - zerompk::from_msgpack(blob_vec_msgpack).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("{what}: bucket decode: {e}"), - })?; - let cells: Vec = blobs - .iter() - .map(|blob| { - zerompk::from_msgpack::(blob).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("{what}: cell decode: {e}"), - }) - }) - .collect::>>()?; - zerompk::to_msgpack_vec(&cells).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("{what}: flat re-encode: {e}"), - }) -} diff --git a/nodedb/src/control/cluster/array_executor/elementwise.rs b/nodedb/src/control/cluster/array_executor/elementwise.rs new file mode 100644 index 000000000..9cb22f8fd --- /dev/null +++ b/nodedb/src/control/cluster/array_executor/elementwise.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Shard-side `ARRAY_ELEMENTWISE`: run the op on this shard's Data Plane over +//! the cells of both arrays in the shard's bucket. + +use nodedb_array::types::ArrayId; +use nodedb_cluster::distributed_array::ArrayElementwiseExec; +use nodedb_cluster::distributed_array::wire::ArrayShardElementwiseReq; +use nodedb_cluster::error::{ClusterError, Result}; + +use crate::types::{TxnId, VShardId}; +use nodedb_physical::physical_plan::{ArrayBinaryOp, ArrayOp, PhysicalPlan}; + +use super::executor::DataPlaneArrayExecutor; +use super::read::split_msgpack_array_rows; +use super::refusal::refusal_error; + +impl DataPlaneArrayExecutor { + pub(super) async fn elementwise( + &self, + local_vshard_id: u32, + req: &ArrayShardElementwiseReq, + ) -> Result { + let decode_id = |bytes: &[u8], side: &str| -> Result { + zerompk::from_msgpack(bytes).map_err(|e| ClusterError::Codec { + detail: format!("{side} array_id decode in exec_elementwise: {e}"), + }) + }; + let left = decode_id(&req.left_id_msgpack, "left")?; + let right = decode_id(&req.right_id_msgpack, "right")?; + let op: ArrayBinaryOp = + zerompk::from_msgpack(&req.op_msgpack).map_err(|e| ClusterError::Codec { + detail: format!("op decode in exec_elementwise: {e}"), + })?; + + let plan = PhysicalPlan::Array(ArrayOp::Elementwise { + left: left.clone(), + right, + op, + attr_idx: req.attr_idx, + cell_filter: None, + shard_bucket: req.shard_bucket, + system_as_of: req.system_as_of, + valid_at_ms: req.valid_at_ms, + }); + let resp = self + .dispatch_and_await( + &left, + VShardId::new(local_vshard_id), + plan, + req.txn_id.map(TxnId::new), + ) + .await?; + if resp.status == crate::bridge::envelope::Status::Error { + return Err(refusal_error("array elementwise", &resp)); + } + let mut rows = split_msgpack_array_rows(&resp.payload)?; + // A temporal read ends with the `{"truncated_before_horizon": bool}` + // summary row. The coordinator ORs the flag and appends one summary + // row for the whole fan-out. + let truncated_before_horizon = if req.system_as_of.is_some() || req.valid_at_ms.is_some() { + let summary = rows.pop().ok_or_else(|| ClusterError::Codec { + detail: "exec_elementwise: temporal answer has no summary row".into(), + })?; + horizon_flag(&summary)? + } else { + false + }; + Ok(ArrayElementwiseExec { + rows, + truncated_before_horizon, + read_versions: resp.read_versions.to_wire(), + }) + } +} + +/// The flag of a `{"truncated_before_horizon": bool}` summary row. +fn horizon_flag(row: &[u8]) -> Result { + let value = nodedb_types::json_from_msgpack(row).map_err(|e| ClusterError::Codec { + detail: format!("exec_elementwise: summary row decode: {e}"), + })?; + value + .get("truncated_before_horizon") + .and_then(serde_json::Value::as_bool) + .ok_or_else(|| ClusterError::Codec { + detail: format!("exec_elementwise: last row {value} is not a horizon summary"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_summary_row_yields_its_flag() { + let row = nodedb_types::json_to_msgpack(&serde_json::json!({ + "truncated_before_horizon": true + })) + .expect("encode"); + assert!(horizon_flag(&row).expect("flag")); + let cell = nodedb_types::json_to_msgpack(&serde_json::json!({ + "coords": [1], "attrs": [2.0] + })) + .expect("encode"); + assert!(horizon_flag(&cell).is_err(), "a cell row is no summary"); + } +} diff --git a/nodedb/src/control/cluster/array_executor/maint.rs b/nodedb/src/control/cluster/array_executor/maint.rs new file mode 100644 index 000000000..379d97115 --- /dev/null +++ b/nodedb/src/control/cluster/array_executor/maint.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Shard-side `ARRAY_FLUSH` / `ARRAY_COMPACT`: run the maintenance on the core +//! that owns the shard's vShard. + +use nodedb_array::types::ArrayId; +use nodedb_cluster::distributed_array::wire::{ArrayShardMaintOp, ArrayShardMaintReq}; +use nodedb_cluster::error::{ClusterError, Result}; + +use crate::types::VShardId; +use nodedb_physical::physical_plan::{ArrayOp, PhysicalPlan}; + +use super::executor::DataPlaneArrayExecutor; +use super::refusal::refusal_error; + +impl DataPlaneArrayExecutor { + pub(super) async fn maint( + &self, + local_vshard_id: u32, + req: &ArrayShardMaintReq, + ) -> Result { + let array_id: ArrayId = + zerompk::from_msgpack(&req.array_id_msgpack).map_err(|e| ClusterError::Codec { + detail: format!("array_id decode in exec_maint: {e}"), + })?; + let (plan, key, context) = match req.op { + // A frontier read of this node's WAL: `Flush` appends no record, + // and every cell the segment holds was written below the frontier. + ArrayShardMaintOp::Flush => ( + ArrayOp::Flush { + array_id: array_id.clone(), + wal_lsn: self.state.wal.next_lsn().as_u64(), + }, + "flushed", + "array flush", + ), + ArrayShardMaintOp::Compact { audit_retain_ms } => ( + ArrayOp::Compact { + array_id: array_id.clone(), + audit_retain_ms, + }, + "compacted", + "array compact", + ), + }; + let resp = self + .dispatch_and_await( + &array_id, + VShardId::new(local_vshard_id), + PhysicalPlan::Array(plan), + None, + ) + .await?; + if resp.status == crate::bridge::envelope::Status::Error { + return Err(refusal_error(context, &resp)); + } + read_count(&resp.payload, key).ok_or_else(|| ClusterError::Storage { + detail: format!("{context}: the data plane answer carried no '{key}' count"), + }) + } +} + +/// The `key` count of a `{key: n}` Data Plane answer. +fn read_count(payload: &[u8], key: &str) -> Option { + let value = nodedb_types::json_from_msgpack(payload).ok()?; + value.get(key).and_then(serde_json::Value::as_u64) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_count_is_read_from_its_key() { + let payload = + crate::data::executor::response_codec::encode_count("compacted", 1).expect("encode"); + assert_eq!(read_count(&payload, "compacted"), Some(1)); + assert_eq!(read_count(&payload, "flushed"), None); + } +} diff --git a/nodedb/src/control/cluster/array_executor/mod.rs b/nodedb/src/control/cluster/array_executor/mod.rs index e65d997cc..32da61752 100644 --- a/nodedb/src/control/cluster/array_executor/mod.rs +++ b/nodedb/src/control/cluster/array_executor/mod.rs @@ -27,15 +27,17 @@ //! dispatch-and-await scaffolding. //! - [`read`]: read/scan handlers (slice, aggregate, surrogate-bitmap scan) plus //! the response-row parsers. +//! - [`elementwise`]: the elementwise handler. +//! - [`maint`]: the flush and compact handler. //! - [`write`]: write handlers (put, delete). //! - [`refusal`]: the typed cluster error a Data-Plane refusal answers with. -pub mod cells; +mod elementwise; mod executor; +mod maint; mod read; mod refusal; mod trait_impl; mod write; -pub use cells::flatten_blob_vec; pub use executor::DataPlaneArrayExecutor; diff --git a/nodedb/src/control/cluster/array_executor/read.rs b/nodedb/src/control/cluster/array_executor/read.rs index 596cd6224..abf0e1fec 100644 --- a/nodedb/src/control/cluster/array_executor/read.rs +++ b/nodedb/src/control/cluster/array_executor/read.rs @@ -122,6 +122,7 @@ impl DataPlaneArrayExecutor { shard_bucket: req.shard_bucket, system_as_of: req.system_as_of, valid_at_ms: req.valid_at_ms, + region_msgpack: req.region_msgpack.clone(), }); let resp = self @@ -196,7 +197,7 @@ impl DataPlaneArrayExecutor { /// /// Returns one `Vec` per row. Returns an empty `Vec` if the payload is /// empty or contains an empty array. -fn split_msgpack_array_rows(payload: &[u8]) -> Result>> { +pub(super) fn split_msgpack_array_rows(payload: &[u8]) -> Result>> { if payload.is_empty() { return Ok(Vec::new()); } diff --git a/nodedb/src/control/cluster/array_executor/refusal.rs b/nodedb/src/control/cluster/array_executor/refusal.rs index b14d384af..edd3a536a 100644 --- a/nodedb/src/control/cluster/array_executor/refusal.rs +++ b/nodedb/src/control/cluster/array_executor/refusal.rs @@ -107,9 +107,11 @@ pub(super) fn execution_error(context: &str, error: crate::Error) -> ClusterErro | crate::Error::UndefinedColumn { .. } | crate::Error::TextColumn { .. } | crate::Error::AmbiguousColumn { .. } + | crate::Error::UngroupedColumn { .. } | crate::Error::UnknownStrictField { .. } | crate::Error::DivisionByZero | crate::Error::DataException { .. } + | crate::Error::MixedSrid { .. } | crate::Error::NumericValueOutOfRange { .. } | crate::Error::InvalidTextRepresentation { .. } | crate::Error::DatatypeMismatch { .. } diff --git a/nodedb/src/control/cluster/array_executor/trait_impl.rs b/nodedb/src/control/cluster/array_executor/trait_impl.rs index 20d67c428..dafe12e01 100644 --- a/nodedb/src/control/cluster/array_executor/trait_impl.rs +++ b/nodedb/src/control/cluster/array_executor/trait_impl.rs @@ -9,10 +9,11 @@ use async_trait::async_trait; use nodedb_cluster::distributed_array::wire::{ - ArrayShardAggReq, ArrayShardDeleteReq, ArrayShardPutReq, + ArrayShardAggReq, ArrayShardDeleteReq, ArrayShardElementwiseReq, ArrayShardMaintReq, + ArrayShardPutReq, }; use nodedb_cluster::distributed_array::{ - ArrayAggExec, ArrayLocalExecutor, ArrayShardWriteOutcome, ArraySliceExec, + ArrayAggExec, ArrayElementwiseExec, ArrayLocalExecutor, ArrayShardWriteOutcome, ArraySliceExec, }; use nodedb_cluster::error::Result; @@ -32,6 +33,18 @@ impl ArrayLocalExecutor for DataPlaneArrayExecutor { self.agg(local_vshard_id, req).await } + async fn exec_elementwise( + &self, + local_vshard_id: u32, + req: &ArrayShardElementwiseReq, + ) -> Result { + self.elementwise(local_vshard_id, req).await + } + + async fn exec_maint(&self, local_vshard_id: u32, req: &ArrayShardMaintReq) -> Result { + self.maint(local_vshard_id, req).await + } + async fn exec_put( &self, local_vshard_id: u32, diff --git a/nodedb/src/control/cluster/array_executor/write.rs b/nodedb/src/control/cluster/array_executor/write.rs index cc1812817..988e3b10d 100644 --- a/nodedb/src/control/cluster/array_executor/write.rs +++ b/nodedb/src/control/cluster/array_executor/write.rs @@ -11,7 +11,6 @@ use nodedb_cluster::distributed_array::ArrayShardWriteOutcome; use nodedb_cluster::distributed_array::wire::{ArrayShardDeleteReq, ArrayShardPutReq}; use nodedb_cluster::error::{ClusterError, Result}; -use super::cells::flatten_blob_vec; use super::executor::DataPlaneArrayExecutor; use super::refusal::execution_error; use crate::control::server::shared::sql::staging_predicates::require_affected_count; @@ -29,19 +28,11 @@ impl DataPlaneArrayExecutor { detail: format!("array_id decode in exec_put: {e}"), })?; - // The coordinator sends a bucket of separately-encoded `ArrayPutCell`s; - // the Data Plane decodes one flat msgpack array. - let cells_msgpack = flatten_blob_vec::( - &req.cells_msgpack, - "exec_put", - ) - .map_err(|e| ClusterError::Codec { - detail: e.to_string(), - })?; - + // The coordinator's bucket is the flat `Vec` encoding + // the Data Plane decodes, so it passes through as it is. let plan = PhysicalPlan::Array(ArrayOp::Put { array_id: array_id.clone(), - cells_msgpack, + cells_msgpack: req.cells_msgpack.clone(), wal_lsn: req.wal_lsn, provenance: None, vshard_id: local_vshard_id, @@ -61,19 +52,11 @@ impl DataPlaneArrayExecutor { detail: format!("array_id decode in exec_delete: {e}"), })?; - // Same bucket-to-flat reshape the put path performs: the Data Plane - // decodes `coords_msgpack` as one `Vec`. - let coords_msgpack = flatten_blob_vec::( - &req.coords_msgpack, - "exec_delete", - ) - .map_err(|e| ClusterError::Codec { - detail: e.to_string(), - })?; - + // The bucket is the flat `Vec` encoding the Data + // Plane decodes, as on the put path. let plan = PhysicalPlan::Array(ArrayOp::Delete { array_id: array_id.clone(), - coords_msgpack, + coords_msgpack: req.coords_msgpack.clone(), wal_lsn: req.wal_lsn, provenance: None, vshard_id: local_vshard_id, diff --git a/nodedb/src/control/cluster/calvin/mod.rs b/nodedb/src/control/cluster/calvin/mod.rs index 1f2051c19..2b83403ef 100644 --- a/nodedb/src/control/cluster/calvin/mod.rs +++ b/nodedb/src/control/cluster/calvin/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: BUSL-1.1 pub mod executor; +pub mod refusal; pub mod scheduler; pub use executor::{OllpConfig, OllpError, OllpOrchestrator}; diff --git a/nodedb/src/control/cluster/calvin/refusal.rs b/nodedb/src/control/cluster/calvin/refusal.rs new file mode 100644 index 000000000..a5cb64f61 --- /dev/null +++ b/nodedb/src/control/cluster/calvin/refusal.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The bytes a participant's typed stage refusal crosses the sequencer log in. +//! +//! A participant whose stage refuses its slice votes abort with the refusal's +//! Data-Plane [`ErrorCode`]. The vote carries the code as the rkyv bytes of +//! its wire mirror [`DataPlaneErrorCode`], opaque to the sequencer. The +//! coordinator decodes the verdict's refusal back into the code, so the +//! client sees the SQLSTATE a direct write renders. + +use nodedb_cluster::rpc_codec::DataPlaneErrorCode; + +use crate::Error; +use crate::bridge::envelope::ErrorCode; + +/// Encode `code` for an abort vote. +pub fn encode_refusal(code: &ErrorCode) -> Result, Error> { + let wire = DataPlaneErrorCode::from(code.clone()); + rkyv::to_bytes::(&wire) + .map(|bytes| bytes.to_vec()) + .map_err(|e| Error::Internal { + detail: format!("encode a Calvin stage refusal {code:?}: {e}"), + }) +} + +/// Decode the refusal an abort verdict carries. +pub fn decode_refusal(bytes: &[u8]) -> Result { + let mut aligned = rkyv::util::AlignedVec::<16>::with_capacity(bytes.len()); + aligned.extend_from_slice(bytes); + rkyv::from_bytes::(&aligned) + .map(ErrorCode::from) + .map_err(|e| Error::Internal { + detail: format!( + "decode the stage refusal of a Calvin abort verdict ({} bytes): {e}", + bytes.len() + ), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_refusal_survives_the_round_trip() { + let codes = [ + ErrorCode::RejectedConstraint { + constraint: "primary_key_immutable".into(), + detail: "UPDATE cannot change primary key 'id'".into(), + }, + ErrorCode::RejectedConstraint { + constraint: "not_null".into(), + detail: "column 'v' is NOT NULL".into(), + }, + ErrorCode::BalanceViolation { + collection: "ledger".into(), + detail: "debits exceed credits".into(), + }, + ]; + for code in codes { + let bytes = encode_refusal(&code).expect("encode"); + assert_eq!(decode_refusal(&bytes).expect("decode"), code); + } + } + + #[test] + fn undecodable_bytes_are_an_internal_error() { + assert!(matches!( + decode_refusal(&[0xde, 0xad]), + Err(Error::Internal { .. }) + )); + } +} diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/verdict.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/verdict.rs index 151071241..fa1f5857e 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/verdict.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/verdict.rs @@ -160,12 +160,14 @@ impl Scheduler { &mut self, txn_id: TxnId, ) { + // The locks free before the ack: the ack answers the coordinator, and + // its client's next write must not find this txn's keys held. + self.metrics.record_completed(); + self.on_txn_complete(txn_id, LedgerMark::Terminal); self.propose_sequencer_entry( txn_id, SchedulerProposal::CompletionAck { result: Vec::new() }, ); - self.metrics.record_completed(); - self.on_txn_complete(txn_id, LedgerMark::Terminal); } /// Handle a pushed [`VerdictSignal`] from this node's completion registry. @@ -178,7 +180,7 @@ impl Scheduler { signal: VerdictSignal, ) { let txn_id = TxnId::new(signal.epoch, signal.position); - self.resume_on_verdict(txn_id, signal.verdict.is_commit()); + self.resume_on_verdict(txn_id, signal.commit); } /// Sweep parked `AwaitingVerdict` txns whose stall deadline has passed. @@ -499,18 +501,18 @@ mod tests { registry.note_vote( txn, 9, - ParticipantVote::Abort(AbortReason::SerializationConflict), + ParticipantVote::Abort(AbortReason::SerializationConflict.into()), ); assert_eq!( registry.drain_unproposed_verdicts(), vec![( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict) + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()) )] ); registry.note_verdict( txn, - VerdictOutcome::Abort(AbortReason::SerializationConflict), + VerdictOutcome::Abort(AbortReason::SerializationConflict.into()), ); assert_eq!(registry.verdict(txn), Some(false)); diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/vote.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/vote.rs index 27eb09f90..af055bec5 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/vote.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/commit_resolve/vote.rs @@ -58,9 +58,8 @@ impl Scheduler { let reason = match error { StageVoteError::CoreFailStopped { .. } => HaltReason::CoreFailStopped, StageVoteError::Missing { .. } - | StageVoteError::CommitWithoutSuccess { .. } => { - HaltReason::StageVoteInvalid - } + | StageVoteError::CommitWithoutSuccess { .. } + | StageVoteError::RefusalUnencodable { .. } => HaltReason::StageVoteInvalid, }; self.halt_apply( txn_id, @@ -77,14 +76,14 @@ impl Scheduler { // group. Only the leader stages, so its vote is the vShard's. The // sequencer aggregates every participant's vote into the single // global verdict this txn parks on below. An abort travels as - // `AbortVote` so its cause survives to the coordinator. The vote stays - // owed until the tally holds it, so a refused or dropped proposal is + // `AbortVote` so its cause and refusal survive to the coordinator. + // The vote stays owed until the tally holds it, so a refused or dropped proposal is // proposed again rather than lost. A promoted leader's vote for a // txn whose vote applied already changes nothing: the first counts. self.propose_sequencer_entry( txn_id, SchedulerProposal::Vote { - abort: vote.abort_reason(), + abort: vote.abort_cause(), }, ); @@ -113,7 +112,7 @@ impl Scheduler { pending.commit_state = CommitState::AwaitingVerdict; // A superseded slice cannot commit here: its collection is gone. pending.stage_error = match vote { - StagedVote::ParticipantError | StagedVote::PredictionDrift => { + StagedVote::ParticipantError { .. } | StagedVote::PredictionDrift => { Some(error_response_text("stage", staged_response)) } StagedVote::CollectionSuperseded => Some( @@ -216,6 +215,7 @@ mod tests { position: 1, vshard: VSHARD, reason, + refusal: None, }, }; assert_eq!(proposer.accepted(), vec![expected], "{vote:?}"); @@ -228,6 +228,35 @@ mod tests { } } + /// A stage refused with a typed code: the abort vote carries the code. + #[tokio::test] + async fn a_refused_stage_proposes_its_refusal_with_the_vote() { + let txn_id = TxnId::new(30, 2); + let (mut scheduler, _dir, proposer) = staged_leader(txn_id); + let code = ErrorCode::RejectedConstraint { + constraint: "not_null".into(), + detail: "column 'v' is NOT NULL".into(), + }; + let mut response = staged_response(Status::Error, Some(StageVote::ParticipantError)); + response.error_code = Some(Box::new(code.clone())); + + scheduler.resolve_staged_commit(txn_id, &response); + + let refusal = crate::control::cluster::calvin::refusal::encode_refusal(&code) + .expect("the refusal encodes"); + assert_eq!( + proposer.accepted(), + vec![SequencerEntry::AbortVote { + epoch: 30, + position: 2, + vshard: VSHARD, + reason: AbortReason::ParticipantError, + refusal: Some(refusal), + }] + ); + assert!(scheduler.apply_halt().is_none()); + } + #[tokio::test] async fn a_successful_stage_response_without_a_vote_halts_and_never_votes() { let txn_id = TxnId::new(31, 0); diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/dispatch/unstaged.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/dispatch/unstaged.rs index c02d6e001..5f2421ead 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/dispatch/unstaged.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/dispatch/unstaged.rs @@ -96,7 +96,7 @@ impl Scheduler { self.propose_sequencer_entry( txn_id, SchedulerProposal::Vote { - abort: Some(reason), + abort: Some(reason.into()), }, ); // The verdict can be stored already: on replay, or when a peer's @@ -116,7 +116,7 @@ mod tests { use nodedb_cluster::calvin::types::SchedulerInput; use nodedb_cluster::calvin::{ - AttemptOutcome, CalvinCompletionRegistry, ParticipantVote, SequencerEntry, + AbortCause, AttemptOutcome, CalvinCompletionRegistry, ParticipantVote, SequencerEntry, }; use super::*; @@ -171,12 +171,13 @@ mod tests { position, vshard, reason, + refusal, } = entry { registry.note_vote( nodedb_cluster::calvin::TxnId::new(epoch, position), vshard, - ParticipantVote::Abort(reason), + ParticipantVote::Abort(AbortCause { reason, refusal }), ); } } @@ -230,16 +231,17 @@ mod tests { position: 0, vshard: REJECTING, reason: AbortReason::PlanRejected, + refusal: None, })); let verdicts = registry.drain_unproposed_verdicts(); assert_eq!( verdicts, vec![( cluster_txn, - nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected) + nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected.into()) )] ); - registry.note_verdict(cluster_txn, verdicts[0].1); + registry.note_verdict(cluster_txn, verdicts[0].1.clone()); for (scheduler, vshard) in [(&mut rejecting, REJECTING), (&mut peer, PEER)] { scheduler.resume_on_verdict(txn_id, false); @@ -260,9 +262,7 @@ mod tests { ); assert_eq!( outcome.await.expect("outcome fires"), - AttemptOutcome::Aborted { - reason: AbortReason::PlanRejected - } + AttemptOutcome::Aborted(AbortReason::PlanRejected.into()) ); } @@ -274,7 +274,7 @@ mod tests { let txn_id = TxnId::new(32, 0); registry.note_verdict( nodedb_cluster::calvin::TxnId::new(32, 0), - nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected), + nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected.into()), ); scheduler.dispatch_txn(undecodable_txn(32, 0), txn_id, txn_id); diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/entry.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/entry.rs index 6f3f43733..dd379b723 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/entry.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/entry.rs @@ -11,7 +11,7 @@ use std::collections::BTreeMap; -use nodedb_cluster::calvin::{AbortReason, ParticipantProgress, SequencerEntry}; +use nodedb_cluster::calvin::{AbortCause, ParticipantProgress, SequencerEntry}; use crate::control::cluster::calvin::scheduler::lock_manager::TxnId; use crate::control::cluster::calvin::scheduler::metrics::sequencer_propose_kind; @@ -19,8 +19,8 @@ use crate::control::cluster::calvin::scheduler::metrics::sequencer_propose_kind; /// A sequencer entry the scheduler proposes for one txn on its vShard. #[derive(Debug, Clone, PartialEq, Eq)] pub enum SchedulerProposal { - /// This vShard's commit vote. `Some(reason)` is an abort vote. - Vote { abort: Option }, + /// This vShard's commit vote. `Some(cause)` is an abort vote. + Vote { abort: Option }, /// This vShard applied or dropped the txn. `result` is the apply result /// the coordinator reads: the timeseries install counts of the apply. /// Empty when there are none. @@ -46,13 +46,12 @@ impl SchedulerProposal { position, vshard, }, - Self::Vote { - abort: Some(reason), - } => SequencerEntry::AbortVote { + Self::Vote { abort: Some(cause) } => SequencerEntry::AbortVote { epoch, position, vshard, - reason: *reason, + reason: cause.reason, + refusal: cause.refusal.clone(), }, Self::CompletionAck { result } => SequencerEntry::CompletionAck { epoch, @@ -131,6 +130,8 @@ pub type OwedEntries = BTreeMap<(TxnId, OwedKind), OwedEntry>; #[cfg(test)] mod tests { + use nodedb_cluster::calvin::AbortReason; + use super::*; fn progress() -> ParticipantProgress { @@ -197,19 +198,20 @@ mod tests { )); assert!(matches!( SchedulerProposal::Vote { - abort: Some(AbortReason::SerializationConflict) + abort: Some(AbortReason::SerializationConflict.into()) } .entry(txn, 9, 1), SequencerEntry::AbortVote { epoch: 3, position: 4, vshard: 9, - reason: AbortReason::SerializationConflict + reason: AbortReason::SerializationConflict, + refusal: None, } )); assert!(matches!( SchedulerProposal::Vote { - abort: Some(AbortReason::PlanRejected) + abort: Some(AbortReason::PlanRejected.into()) } .entry(txn, 9, 1), SequencerEntry::AbortVote { @@ -218,4 +220,21 @@ mod tests { } )); } + + #[test] + fn a_vote_proposal_carries_its_refusal() { + assert_eq!( + SchedulerProposal::Vote { + abort: Some(AbortCause::refused(vec![1, 2])) + } + .entry(TxnId::new(3, 5), 9, 1), + SequencerEntry::AbortVote { + epoch: 3, + position: 5, + vshard: 9, + reason: AbortReason::ParticipantError, + refusal: Some(vec![1, 2]), + } + ); + } } diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/retry.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/retry.rs index 4dc9142fd..5a6cd1ad0 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/retry.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/owed/retry.rs @@ -345,7 +345,7 @@ mod tests { scheduler.propose_sequencer_entry( txn_id, SchedulerProposal::Vote { - abort: Some(nodedb_cluster::calvin::AbortReason::PlanRejected), + abort: Some(nodedb_cluster::calvin::AbortReason::PlanRejected.into()), }, ); scheduler.retry_owed_sequencer_entries(); @@ -365,7 +365,7 @@ mod tests { scheduler.propose_sequencer_entry( TxnId::new(24, 0), SchedulerProposal::Vote { - abort: Some(nodedb_cluster::calvin::AbortReason::PlanRejected), + abort: Some(nodedb_cluster::calvin::AbortReason::PlanRejected.into()), }, ); assert!(proposer.accepted().is_empty()); diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/parts.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/parts.rs index b82f19882..e8fca8f6b 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/parts.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/parts.rs @@ -593,6 +593,7 @@ mod tests { position: txn_id.position, vshard: test_coll_vshard(), reason: AbortReason::PlanRejected, + refusal: None, })); assert!(!scheduler.applied.is_applied(txn_id.epoch, txn_id.position)); @@ -605,7 +606,7 @@ mod tests { scheduler.registry.note_verdict( nodedb_cluster::calvin::TxnId::new(txn_id.epoch, txn_id.position), - nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected), + nodedb_cluster::calvin::VerdictOutcome::Abort(AbortReason::PlanRejected.into()), ); scheduler.resume_on_verdict(txn_id, false); assert_eq!( diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/read_result.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/read_result.rs index 3a1917553..2f246e709 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/read_result.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/read_result.rs @@ -585,7 +585,7 @@ mod tests { let txn_id = TxnId::new(57, 0); registry.note_verdict( nodedb_cluster::calvin::TxnId::new(57, 0), - VerdictOutcome::Abort(AbortReason::ParticipantError), + VerdictOutcome::Abort(AbortReason::ParticipantError.into()), ); scheduler.resume_on_verdict(txn_id, false); diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/redo_applied.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/redo_applied.rs index e9efa74d0..f455d4ef5 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/redo_applied.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/redo_applied.rs @@ -126,9 +126,11 @@ impl Scheduler { lock_owner: TxnId, applied: AppliedRedo, ) { - self.owe_applied_ack(txn_id, &applied); + // The locks free before the ack: the ack answers the coordinator, and + // its client's next write must not find this txn's keys held. self.metrics.record_completed(); self.release_and_mark_applied(txn_id, lock_owner, LedgerMark::ByApply); + self.owe_applied_ack(txn_id, &applied); } /// Deposit a primary slice's reply, then owe the ack that carries it. diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/role.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/role.rs index ae1f2a614..2e9675a91 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/role.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/role.rs @@ -12,6 +12,12 @@ //! proposed installs before this leader stages anything, and this leader //! finds that position applied. //! +//! A scheduler that starts while this node already leads also waits for its +//! start fence. A write the gate admitted earlier holds no key on this +//! scheduler's lock table: none existed, or its keys were on a retired one. +//! The gate stays closed until each such write applied (see +//! `write_admission::tickets`). +//! //! - Promotion stages every `Following` txn as a new grant, in sequencer //! order. A txn whose verdict is already COMMIT proposes its vote again, //! which changes nothing: the first vote of the vShard counts. @@ -27,6 +33,7 @@ use super::deferred::{DispatchOutcome, DispatchStep}; use super::owed::OwedKind; use super::scheduler::Scheduler; use crate::control::cluster::calvin::scheduler::lock_manager::TxnId; +use crate::control::server::shared::write_admission::{FenceState, StartFence}; /// What this node does for its vShard's Calvin txns. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -46,6 +53,10 @@ pub(in crate::control::cluster::calvin::scheduler::driver::core) struct StageGat /// The index of the term's election no-op. The gate opens once this /// node applied it. term_start: Option, + /// The writes the gate admitted for this vShard before this scheduler's + /// lock table was registered. The gate opens once they all applied. + /// `None` once they did. + start_fence: Option, } impl Default for StageGate { @@ -54,11 +65,22 @@ impl Default for StageGate { role: Role::Follower, term: None, term_start: None, + start_fence: None, } } } impl StageGate { + /// The gate of a scheduler that starts behind `start_fence`. + pub(in crate::control::cluster::calvin::scheduler::driver::core) fn fenced( + start_fence: StartFence, + ) -> Self { + Self { + start_fence: Some(start_fence), + ..Self::default() + } + } + /// Whether this node stages the vShard's txns. pub(in crate::control::cluster::calvin::scheduler::driver::core) fn is_leader(&self) -> bool { self.role == Role::Leader @@ -119,11 +141,44 @@ impl Scheduler { .shared .applied_index_watcher(leadership.group_id) .current(); - if self.role.term_start.is_some_and(|start| applied >= start) { + if self.role.term_start.is_some_and(|start| applied >= start) + && self.start_fence_passed(leadership.group_id, applied) + { self.promote(); } } + /// Whether every write the gate admitted for this vShard before its lock + /// table was registered applied, with `applied` the group's applied index. + /// + /// Such a write holds no key on this table. A txn that stages before it + /// reached its core reads the state before it, and its redo overwrites it. + fn start_fence_passed(&mut self, group_id: u64, applied: u64) -> bool { + let Some(fence) = self.role.start_fence else { + return true; + }; + let passed = match self + .shared + .calvin + .admission_tickets + .fence_state(fence, group_id) + { + FenceState::Open => false, + FenceState::Through(through) => through.is_none_or(|index| applied >= index), + }; + if passed { + self.role.start_fence = None; + } else { + tracing::debug!( + vshard_id = self.vshard_id, + group_id, + applied, + "calvin scheduler: staging waits for writes admitted before its lock table" + ); + } + passed + } + /// This node's term and term start in the vShard's data group, while it /// leads the group. fn leadership(&self) -> Option { @@ -280,6 +335,8 @@ fn leader_drives(state: CommitState) -> bool { #[cfg(test)] mod tests { + use std::sync::Arc; + use nodedb_cluster::calvin::CalvinCompletionRegistry; use nodedb_cluster::calvin::types::{ EngineKeySet, ReadWriteSet, SchedulerInput, SequencedTxn, SortedVec, @@ -402,6 +459,42 @@ mod tests { assert_eq!(staged_requests(&mut data_side), vec![(3, 0)]); } + /// A write the gate admitted with no lock table holds the stage gate + /// closed until it landed and its index applied. + #[tokio::test] + async fn a_new_scheduler_stages_nothing_before_earlier_unheld_writes_applied() { + let (mut scheduler, _dir, mut data_side) = scheduler(); + let tickets = Arc::clone(&scheduler.shared.calvin.admission_tickets); + let write = tickets.open(scheduler.vshard_id); + scheduler.role = StageGate::fenced(tickets.start_fence(scheduler.vshard_id)); + + let group_id = elect_data_group_leader(&scheduler); + let term_start = scheduler + .multi_raft + .lock() + .unwrap_or_else(|p| p.into_inner()) + .term_start_index(group_id) + .expect("the elected node names its term start"); + let watcher = scheduler.shared.applied_index_watcher(group_id); + watcher.bump(term_start); + scheduler.process_scheduler_input(SchedulerInput::Txn(Box::new(validate_only_on(3, 1)))); + + scheduler.refresh_role(); + assert!(!scheduler.role.is_leader(), "the write has not landed"); + assert!(staged_requests(&mut data_side).is_empty()); + + let landed_at = term_start + 5; + write.landed(group_id, landed_at); + scheduler.refresh_role(); + assert!(!scheduler.role.is_leader(), "the write has not applied"); + assert!(staged_requests(&mut data_side).is_empty()); + + watcher.bump(landed_at); + scheduler.refresh_role(); + assert!(scheduler.role.is_leader(), "the gate opens"); + assert_eq!(staged_requests(&mut data_side), vec![(3, 0)]); + } + /// Promotion stages every `Following` txn, in sequencer order. #[tokio::test] async fn promotion_stages_following_txns_in_sequencer_order() { diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/routing.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/routing.rs index 8d8751291..82305dd4a 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/routing.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/routing.rs @@ -396,7 +396,6 @@ fn array_routing(op: &ArrayOp) -> PlanRouting { | ArrayOp::RekeyArray { .. } | ArrayOp::PurgeArrayDrop { .. } | ArrayOp::Slice { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } | ArrayOp::SurrogateBitmapScan { .. } => PlanRouting::NotAWrite, diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/scheduler.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/scheduler.rs index d5227ff00..fc77744e5 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/scheduler.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/scheduler.rs @@ -292,12 +292,15 @@ impl Scheduler { .unwrap_or_else(|p| p.into_inner()) .capacity_freed(); + // Taken after this vShard's lock table is registered: every write the + // gate admitted with no table, or with a retired one, is earlier. + let start_fence = shared.calvin.admission_tickets.start_fence(vshard_id); let scheduler = Self { vshard_id, receiver, shared, multi_raft, - role: super::role::StageGate::default(), + role: super::role::StageGate::fenced(start_fence), sequencer_proposer, owed: OwedEntries::new(), sequencer_state_machine, diff --git a/nodedb/src/control/cluster/calvin/scheduler/driver/core/staged_vote.rs b/nodedb/src/control/cluster/calvin/scheduler/driver/core/staged_vote.rs index 165864f9f..4dd63434d 100644 --- a/nodedb/src/control/cluster/calvin/scheduler/driver/core/staged_vote.rs +++ b/nodedb/src/control/cluster/calvin/scheduler/driver/core/staged_vote.rs @@ -3,19 +3,24 @@ //! This participant's local commit vote on a staged Calvin transaction, //! read from the stage vote its executor response carries. -use nodedb_cluster::calvin::AbortReason; +use nodedb_cluster::calvin::{AbortCause, AbortReason}; use crate::bridge::envelope::{ErrorCode, Response, StageVote, Status}; +use crate::control::cluster::calvin::refusal::encode_refusal; /// A participant's local verdict on its staged slice. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub(super) enum StagedVote { /// The read-set was still current, or the path carries none. Commit, /// Validated and stale: a genuine serialization conflict. SerializationConflict, - /// The participant never staged, so no read-set was ever validated. - ParticipantError, + /// The slice holds no staged state, so no read-set was ever validated. + /// + /// `refusal` is the encoded error a stage path refused the slice with. + /// Every replica stages the same slice and raises the same refusal. It + /// is `None` for a slice that never reached a stage path. + ParticipantError { refusal: Option> }, /// A collection the transaction names no longer holds its planned /// incarnation in this replica's catalog. CollectionSuperseded, @@ -24,32 +29,24 @@ pub(super) enum StagedVote { PredictionDrift, } -impl From for StagedVote { - fn from(vote: StageVote) -> Self { - match vote { - StageVote::Commit => Self::Commit, - StageVote::SerializationConflict => Self::SerializationConflict, - StageVote::PredictionDrift => Self::PredictionDrift, - StageVote::ParticipantError => Self::ParticipantError, - } - } -} - impl StagedVote { /// `None` for a commit, the abort cause otherwise. - pub(super) fn abort_reason(self) -> Option { + pub(super) fn abort_cause(&self) -> Option { match self { Self::Commit => None, - Self::SerializationConflict => Some(AbortReason::SerializationConflict), - Self::ParticipantError => Some(AbortReason::ParticipantError), - Self::CollectionSuperseded => Some(AbortReason::CollectionSuperseded), - Self::PredictionDrift => Some(AbortReason::PredictionDrift), + Self::SerializationConflict => Some(AbortReason::SerializationConflict.into()), + Self::ParticipantError { refusal } => Some(AbortCause { + reason: AbortReason::ParticipantError, + refusal: refusal.clone(), + }), + Self::CollectionSuperseded => Some(AbortReason::CollectionSuperseded.into()), + Self::PredictionDrift => Some(AbortReason::PredictionDrift.into()), } } } /// A stage response the scheduler cannot read a vote from. -#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(super) enum StageVoteError { /// A stage response that did not fail carries no stage vote. #[error("stage returned {status:?} without a stage vote")] @@ -62,14 +59,18 @@ pub(super) enum StageVoteError { /// this replica casts can be kept. #[error("stage refused by fail-stopped core {core_id}")] CoreFailStopped { core_id: usize }, + /// The stage's refusal has no wire form, so no vote can carry it. + #[error("stage refusal cannot be encoded for the abort vote: {detail}")] + RefusalUnencodable { detail: String }, } /// Read the local vote of a staged slice from its stage response. /// -/// The response's stage vote maps 1:1. An error response with no vote never +/// The response's stage vote maps 1:1. A `ParticipantError` vote carries the +/// response's error code as its refusal. An error response with no vote never /// reached a stage path: the request expired, or its core is gone or -/// drained. That slice never staged, so it votes `ParticipantError`. A -/// fail-stopped core's refusal is the exception: it is refused. +/// drained. That slice never staged, so it votes `ParticipantError` with no +/// refusal. A fail-stopped core's refusal is the exception: it is refused. /// /// A response that did not fail and carries no vote is refused, and so is a /// commit vote on a response that did not succeed. Neither counts as a commit. @@ -83,15 +84,20 @@ pub(super) fn staged_commit_vote(response: &Response) -> Result { Err(StageVoteError::CommitWithoutSuccess { status }) } - ( - Some( - vote @ (StageVote::SerializationConflict - | StageVote::PredictionDrift - | StageVote::ParticipantError), - ), - _, - ) => Ok(vote.into()), - (None, Status::Error) => Ok(StagedVote::ParticipantError), + (Some(StageVote::SerializationConflict), _) => Ok(StagedVote::SerializationConflict), + (Some(StageVote::PredictionDrift), _) => Ok(StagedVote::PredictionDrift), + (Some(StageVote::ParticipantError), _) => { + let refusal = response + .error_code + .as_deref() + .map(encode_refusal) + .transpose() + .map_err(|e| StageVoteError::RefusalUnencodable { + detail: e.to_string(), + })?; + Ok(StagedVote::ParticipantError { refusal }) + } + (None, Status::Error) => Ok(StagedVote::ParticipantError { refusal: None }), (None, Status::Ok | Status::Partial) => Err(StageVoteError::Missing { status }), } } @@ -134,7 +140,7 @@ mod tests { ( Status::Error, StageVote::ParticipantError, - StagedVote::ParticipantError, + StagedVote::ParticipantError { refusal: None }, ), ]; for (status, vote, expected) in cases { @@ -153,8 +159,30 @@ mod tests { let vote = staged_commit_vote(&response); assert_eq!(vote, Ok(StagedVote::PredictionDrift)); assert_eq!( - vote.map(StagedVote::abort_reason), - Ok(Some(AbortReason::PredictionDrift)) + vote.map(|vote| vote.abort_cause()), + Ok(Some(AbortReason::PredictionDrift.into())) + ); + } + + /// A stage that refused with a constraint keeps the refusal on its vote, + /// so the coordinator answers that refusal instead of a bare abort. + #[test] + fn a_refused_stage_vote_keeps_its_error_code() { + let code = ErrorCode::RejectedConstraint { + constraint: "primary_key_immutable".into(), + detail: "UPDATE cannot change primary key 'id'".into(), + }; + let mut response = staged_response(Status::Error, Some(StageVote::ParticipantError)); + response.error_code = Some(Box::new(code.clone())); + + let vote = staged_commit_vote(&response).expect("a readable vote"); + let cause = vote.abort_cause().expect("an abort vote"); + assert_eq!(cause.reason, AbortReason::ParticipantError); + let refusal = cause.refusal.expect("the vote carries the refusal"); + assert_eq!( + crate::control::cluster::calvin::refusal::decode_refusal(&refusal) + .expect("the refusal decodes"), + code ); } @@ -189,7 +217,7 @@ mod tests { response.error_code = Some(Box::new(ErrorCode::ExpiredBeforeExecution)); assert_eq!( staged_commit_vote(&response), - Ok(StagedVote::ParticipantError) + Ok(StagedVote::ParticipantError { refusal: None }) ); } @@ -209,23 +237,30 @@ mod tests { } #[test] - fn only_an_abort_carries_a_reason() { - assert_eq!(StagedVote::Commit.abort_reason(), None); + fn only_an_abort_carries_a_cause() { + assert_eq!(StagedVote::Commit.abort_cause(), None); + assert_eq!( + StagedVote::SerializationConflict.abort_cause(), + Some(AbortReason::SerializationConflict.into()) + ); assert_eq!( - StagedVote::SerializationConflict.abort_reason(), - Some(AbortReason::SerializationConflict) + StagedVote::ParticipantError { refusal: None }.abort_cause(), + Some(AbortReason::ParticipantError.into()) ); assert_eq!( - StagedVote::ParticipantError.abort_reason(), - Some(AbortReason::ParticipantError) + StagedVote::ParticipantError { + refusal: Some(vec![3]) + } + .abort_cause(), + Some(AbortCause::refused(vec![3])) ); assert_eq!( - StagedVote::CollectionSuperseded.abort_reason(), - Some(AbortReason::CollectionSuperseded) + StagedVote::CollectionSuperseded.abort_cause(), + Some(AbortReason::CollectionSuperseded.into()) ); assert_eq!( - StagedVote::PredictionDrift.abort_reason(), - Some(AbortReason::PredictionDrift) + StagedVote::PredictionDrift.abort_cause(), + Some(AbortReason::PredictionDrift.into()) ); } } diff --git a/nodedb/src/control/cluster/execution_error_wire.rs b/nodedb/src/control/cluster/execution_error_wire.rs index 2bfd23196..e6e53b353 100644 --- a/nodedb/src/control/cluster/execution_error_wire.rs +++ b/nodedb/src/control/cluster/execution_error_wire.rs @@ -122,9 +122,11 @@ pub(crate) fn execution_error_to_typed(err: crate::Error) -> TypedClusterError { | crate::Error::UndefinedColumn { .. } | crate::Error::TextColumn { .. } | crate::Error::AmbiguousColumn { .. } + | crate::Error::UngroupedColumn { .. } | crate::Error::UnknownStrictField { .. } | crate::Error::DivisionByZero | crate::Error::DataException { .. } + | crate::Error::MixedSrid { .. } | crate::Error::NumericValueOutOfRange { .. } | crate::Error::InvalidLimitValue { .. } | crate::Error::RetryableLeaderChange { .. } diff --git a/nodedb/src/control/cluster/metadata_applier/audit.rs b/nodedb/src/control/cluster/metadata_applier/audit.rs index 2b98bc492..3ba57a8f6 100644 --- a/nodedb/src/control/cluster/metadata_applier/audit.rs +++ b/nodedb/src/control/cluster/metadata_applier/audit.rs @@ -53,16 +53,17 @@ pub(super) fn apply_ca_trust_change( "removed_fingerprint": remove.map(nodedb_cluster::ca_fingerprint_hex), })) .unwrap_or_default(); - if let Ok(mut log) = shared.audit.lock() { - log.record_with_auth( - AuditEvent::CertRotation, - None, - None, - "metadata_group", - &detail, - &AuditAuth::default(), - ); - } + crate::control::security::audit::record_concurrent( + &shared.audit, + crate::control::security::audit::AuditRecordInput { + event: AuditEvent::CertRotation, + tenant_id: None, + database_id: None, + source: "metadata_group", + detail: &detail, + auth: &AuditAuth::default(), + }, + ); Ok(()) } @@ -134,20 +135,20 @@ pub(super) fn emit_ddl_audit( // other audit writer. let audit = std::sync::Arc::clone(&shared.audit); let emit = move || { - let mut log = match audit.lock() { - Ok(g) => g, - Err(p) => p.into_inner(), + let auth = AuditAuth { + user_id, + user_name, + session_id: String::new(), }; - log.record_with_auth( - AuditEvent::DdlChange, - None, - None, - "metadata_group", - &detail_json, - &AuditAuth { - user_id, - user_name, - session_id: String::new(), + crate::control::security::audit::record_concurrent( + &audit, + crate::control::security::audit::AuditRecordInput { + event: AuditEvent::DdlChange, + tenant_id: None, + database_id: None, + source: "metadata_group", + detail: &detail_json, + auth: &auth, }, ); }; @@ -160,3 +161,49 @@ pub(super) fn emit_ddl_audit( Err(_) => emit(), } } + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use super::super::test_fixture::applier_with_shared_at; + use super::*; + use crate::control::security::audit::AuditEvent; + use crate::control::security::audit::restore::replay_after_crash; + use crate::control::security::catalog::StoredCollection; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn ddl_audit_entry_survives_a_crash() { + let dir = tempfile::tempdir().expect("tempdir"); + let (_applier, state) = applier_with_shared_at(dir.path(), "ddl-audit.wal"); + let stamped = catalog_entry::CatalogEntry::PutCollection(Box::new( + StoredCollection::stamped_for_test(7, "crash_coll", "tester"), + )); + emit_ddl_audit(&state, 11, &stamped, None); + + // The emit runs on a blocking task, so wait until it lands in the buffer. + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let landed = state + .audit + .lock() + .expect("audit lock") + .all() + .iter() + .any(|e| e.source == "metadata_group"); + if landed { + break; + } + assert!(Instant::now() < deadline, "the DDL audit emit never ran"); + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let replayed = replay_after_crash(&state.wal); + assert!( + replayed + .iter() + .any(|e| e.source == "metadata_group" && e.event == AuditEvent::DdlChange), + "the metadata-applier entry must replay after a crash: {replayed:?}" + ); + } +} diff --git a/nodedb/src/control/cluster/metadata_applier/wedge.rs b/nodedb/src/control/cluster/metadata_applier/wedge.rs index 761b40fea..6b7a6c1e5 100644 --- a/nodedb/src/control/cluster/metadata_applier/wedge.rs +++ b/nodedb/src/control/cluster/metadata_applier/wedge.rs @@ -116,9 +116,11 @@ pub fn classify(error: &crate::Error) -> ApplyFailureClass { | crate::Error::UndefinedColumn { .. } | crate::Error::TextColumn { .. } | crate::Error::AmbiguousColumn { .. } + | crate::Error::UngroupedColumn { .. } | crate::Error::UnknownStrictField { .. } | crate::Error::DivisionByZero | crate::Error::DataException { .. } + | crate::Error::MixedSrid { .. } | crate::Error::NumericValueOutOfRange { .. } | crate::Error::InvalidTextRepresentation { .. } | crate::Error::DatatypeMismatch { .. } diff --git a/nodedb/src/control/cluster/start_raft/group_setup.rs b/nodedb/src/control/cluster/start_raft/group_setup.rs index 389644a16..2f193eb04 100644 --- a/nodedb/src/control/cluster/start_raft/group_setup.rs +++ b/nodedb/src/control/cluster/start_raft/group_setup.rs @@ -69,6 +69,7 @@ pub(super) fn build_group_setup( let tracker = Arc::new( ProposeTracker::new() .with_group_watchers(handle.group_watchers.clone()) + .with_admission_holds(Arc::clone(&shared.calvin.admission_holds)) .with_waiter_window(std::time::Duration::from_secs( shared.tuning.network.default_deadline_secs, )), diff --git a/nodedb/src/control/cluster/vshard_envelope_handler.rs b/nodedb/src/control/cluster/vshard_envelope_handler.rs index a3438e907..c420caeb9 100644 --- a/nodedb/src/control/cluster/vshard_envelope_handler.rs +++ b/nodedb/src/control/cluster/vshard_envelope_handler.rs @@ -111,6 +111,8 @@ mod tests { (85, VShardMessageType::ArrayShardPutResp), (87, VShardMessageType::ArrayShardDeleteResp), (89, VShardMessageType::ArrayShardSurrogateBitmapResp), + (91, VShardMessageType::ArrayShardElementwiseResp), + (93, VShardMessageType::ArrayShardMaintResp), ]; for (opcode, expected) in responses { let resolved = resolve_vshard_msg_type(opcode) @@ -121,7 +123,7 @@ mod tests { #[test] fn an_unknown_opcode_is_a_codec_error() { - assert!(resolve_vshard_msg_type(90).is_err()); + assert!(resolve_vshard_msg_type(94).is_err()); assert!(resolve_vshard_msg_type(u32::from(u16::MAX) + 1).is_err()); } } diff --git a/nodedb/src/control/distributed_applier/apply_loop/start.rs b/nodedb/src/control/distributed_applier/apply_loop/start.rs index df3f150eb..3c726f71d 100644 --- a/nodedb/src/control/distributed_applier/apply_loop/start.rs +++ b/nodedb/src/control/distributed_applier/apply_loop/start.rs @@ -235,15 +235,19 @@ fn prepare_replicated<'a>( ReplicatedWrite::RedoChunk { .. } | ReplicatedWrite::RedoAbandon { .. } => { super::redo_chunk::prepare_stream_entry(ctx, pos, log_term, replicated) } - ReplicatedWrite::SurrogateBind { ref identities } => { - Prepared::Concluded(super::surrogate_bind::apply_surrogate_bind( - ctx, - pos, + ReplicatedWrite::SurrogateBind { + ref collection, + ref binds, + } => Prepared::Concluded(super::surrogate_bind::apply_surrogate_bind( + ctx, + pos, + super::surrogate_bind::BindKeys { tenant_id, - entry_database, - identities, - )) - } + database_id: entry_database, + collection, + binds, + }, + )), ReplicatedWrite::TopicPublish { topic, payload, diff --git a/nodedb/src/control/distributed_applier/apply_loop/surrogate_bind.rs b/nodedb/src/control/distributed_applier/apply_loop/surrogate_bind.rs index 64c6b8b77..ff67a9af7 100644 --- a/nodedb/src/control/distributed_applier/apply_loop/surrogate_bind.rs +++ b/nodedb/src/control/distributed_applier/apply_loop/surrogate_bind.rs @@ -5,25 +5,32 @@ //! The leader of a key's collection home mints the key's surrogate by //! proposing this entry (`surrogate_exchange::authority`). Every replica binds //! each key first-wins, in Raft log order. A key an earlier entry of the log -//! already bound keeps that binding on every replica alike, so the proposer -//! reads the winner back after its own entry applies. The binds land in the -//! fsync-committed catalog before the apply reports. +//! already bound keeps that binding on every replica alike. The apply answers +//! each key's winner, in entry order, as its payload: msgpack `Vec`. +//! The binds land in the fsync-committed catalog before the apply reports. + +use nodedb_types::{CollectionKey, Surrogate}; use crate::control::array_sync::raft_apply::AppliedPosition; use crate::control::distributed_applier::propose_tracker::AppliedWrite; -use crate::control::wal_replication::ReplicatedIdentity; use crate::types::{DatabaseId, TenantId}; use super::context::ApplyContext; use super::proposal_gate::EntryOutcome; -/// Bind every identity of the entry and resolve its waiter. +/// One `SurrogateBind` entry: its collection and its `(pk, candidate)` keys. +pub(super) struct BindKeys<'e> { + pub(super) tenant_id: TenantId, + pub(super) database_id: DatabaseId, + pub(super) collection: &'e str, + pub(super) binds: &'e [(Vec, u32)], +} + +/// Bind every key of the entry and resolve its waiter. pub(super) fn apply_surrogate_bind( ctx: ApplyContext<'_>, pos: AppliedPosition, - tenant_id: TenantId, - database_id: DatabaseId, - identities: &[ReplicatedIdentity], + entry: BindKeys<'_>, ) -> EntryOutcome { let AppliedPosition { group_id, @@ -31,13 +38,18 @@ pub(super) fn apply_surrogate_bind( applied_key, .. } = pos; - let result = bind_identities(&ctx, tenant_id, database_id, identities); + let result = bind_keys(&ctx, &entry).and_then(|winners| { + zerompk::to_msgpack_vec(&winners).map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("surrogate bind winners encode: {e}"), + }) + }); let durable = result.is_ok(); ctx.tracker.complete( group_id, log_index, applied_key, - result.map(|()| AppliedWrite::unversioned(Vec::new())), + result.map(AppliedWrite::unversioned), ); EntryOutcome::Applied { durable, @@ -45,36 +57,31 @@ pub(super) fn apply_surrogate_bind( } } -/// Bind the entry's identities first-wins, one catalog transaction per -/// collection the entry names. Keys of one collection bind in entry order. -fn bind_identities( - ctx: &ApplyContext<'_>, - tenant_id: TenantId, - database_id: DatabaseId, - identities: &[ReplicatedIdentity], -) -> crate::Result<()> { - let mut collections: Vec<&str> = Vec::new(); - for identity in identities { - if !collections.contains(&identity.collection.as_str()) { - collections.push(&identity.collection); - } - } - for collection in collections { - let rows: Vec<(&[u8], nodedb_types::Surrogate)> = identities - .iter() - .filter(|identity| identity.collection == collection) - .map(|identity| { - ( - identity.pk_bytes.as_slice(), - nodedb_types::Surrogate::new(identity.surrogate), - ) - }) - .collect(); - ctx.state.surrogate_assigner.bind_many( - nodedb_types::CollectionKey::from_bare(database_id, collection), - tenant_id, - &rows, - )?; +/// Bind the entry's keys first-wins in one catalog transaction, in entry +/// order, and return the winner of every key, in entry order. +fn bind_keys(ctx: &ApplyContext<'_>, entry: &BindKeys<'_>) -> crate::Result> { + let rows: Vec<(&[u8], Surrogate)> = entry + .binds + .iter() + .map(|(pk, candidate)| (pk.as_slice(), Surrogate::new(*candidate))) + .collect(); + let bound = ctx.state.surrogate_assigner.bind_many( + CollectionKey::from_bare(entry.database_id, entry.collection), + entry.tenant_id, + &rows, + )?; + if bound.len() != rows.len() { + return Err(crate::Error::Internal { + detail: format!( + "surrogate bind in '{}' answered {} of {} keys", + entry.collection, + bound.len(), + rows.len() + ), + }); } - Ok(()) + Ok(bound + .into_iter() + .map(|surrogate| surrogate.as_u32()) + .collect()) } diff --git a/nodedb/src/control/distributed_applier/propose_tracker/tracker.rs b/nodedb/src/control/distributed_applier/propose_tracker/tracker.rs index 1a92bfad6..ae7ce784f 100644 --- a/nodedb/src/control/distributed_applier/propose_tracker/tracker.rs +++ b/nodedb/src/control/distributed_applier/propose_tracker/tracker.rs @@ -18,6 +18,7 @@ use super::committed_keys::{CarriedKeys, CommittedKeys}; use super::cut::GroupCuts; use super::slots::WaiterSlots; use crate::control::distributed_applier::apply_window::ApplyWindow; +use crate::control::server::shared::write_admission::AdmissionHolds; /// How long a propose waiter can exist: the default statement deadline. /// `start_raft` sets the configured deadline through @@ -57,6 +58,10 @@ pub struct ProposeTracker { window: Arc, /// Per group, the indexes a data-group snapshot is cut at. cuts: Mutex, + /// The lock keys this node's write gate holds for proposed entries. + /// [`Self::complete`] releases an entry's keys before it answers the + /// proposer. `None` only in tests that run no write gate. + admission_holds: Option>, } /// The keys one install carried, and the lowest index they are complete from. @@ -83,9 +88,17 @@ impl ProposeTracker { applying: Mutex::new(HashMap::new()), window: Arc::new(ApplyWindow::default()), cuts: Mutex::new(GroupCuts::default()), + admission_holds: None, } } + /// Wire the write gate's held keys. Called by `start_raft` with the + /// holds of `SharedState`. + pub fn with_admission_holds(mut self, holds: Arc) -> Self { + self.admission_holds = Some(holds); + self + } + /// Wire the per-group apply watermark registry. Called by /// `start_raft` after `SharedState` is constructed. pub fn with_group_watchers(mut self, watchers: Arc) -> Self { @@ -182,6 +195,11 @@ impl ProposeTracker { /// change overwrote the proposal, and the waiter gets /// `RetryableLeaderChange` instead of another proposal's result. /// + /// The write gate's keys for every entry of the group through + /// `log_index` release first. An entry completes only after it and every + /// entry below it started. A proposer that writes again after this answer + /// never finds its own entry's keys held. + /// /// Returns true if a live waiter was found and notified, false otherwise. pub fn complete( &self, @@ -190,6 +208,9 @@ impl ProposeTracker { applied_key: u64, result: ProposeResult, ) -> bool { + if let Some(holds) = &self.admission_holds { + holds.release_through(group_id, log_index); + } self.slots() .complete(group_id, log_index, applied_key, result, Instant::now()) } diff --git a/nodedb/src/control/gateway/error_map/class_parity/error_index.rs b/nodedb/src/control/gateway/error_map/class_parity/error_index.rs index 7a3acc263..6e7541dc1 100644 --- a/nodedb/src/control/gateway/error_map/class_parity/error_index.rs +++ b/nodedb/src/control/gateway/error_map/class_parity/error_index.rs @@ -3,7 +3,7 @@ //! A dense index per `crate::Error` variant. /// The number of `crate::Error` variants [`error_variant_index`] numbers. -pub(super) const ERROR_VARIANT_COUNT: usize = 124; +pub(super) const ERROR_VARIANT_COUNT: usize = 126; /// A dense index per `crate::Error` variant. Exhaustive, so a new variant /// fails to compile here until it gets an index, and @@ -136,5 +136,7 @@ pub(crate) fn error_variant_index(err: &crate::Error) -> usize { E::InvalidDatetimeFormat { .. } => 119, E::DatetimeFieldOverflow { .. } => 120, E::WrongObjectType { .. } => 123, + E::MixedSrid { .. } => 124, + E::UngroupedColumn { .. } => 125, } } diff --git a/nodedb/src/control/gateway/error_map/class_parity/error_parity.rs b/nodedb/src/control/gateway/error_map/class_parity/error_parity.rs index a1b21c6b0..9ebb7e6f4 100644 --- a/nodedb/src/control/gateway/error_map/class_parity/error_parity.rs +++ b/nodedb/src/control/gateway/error_map/class_parity/error_parity.rs @@ -117,6 +117,7 @@ fn classified_sqlstates() -> Vec<(usize, &'static str)> { (118, sqlstate::DATATYPE_MISMATCH), (119, sqlstate::INVALID_DATETIME_FORMAT), (120, sqlstate::DATETIME_FIELD_OVERFLOW), + (125, sqlstate::GROUPING_ERROR), ] } diff --git a/nodedb/src/control/gateway/error_map/class_parity/error_samples.rs b/nodedb/src/control/gateway/error_map/class_parity/error_samples.rs index a5e9c2626..0016f90ce 100644 --- a/nodedb/src/control/gateway/error_map/class_parity/error_samples.rs +++ b/nodedb/src/control/gateway/error_map/class_parity/error_samples.rs @@ -208,12 +208,19 @@ pub(crate) fn error_samples() -> Vec { E::AmbiguousColumn { column: "id".into(), }, + E::UngroupedColumn { column: "x".into() }, E::UnknownStrictField { collection: collection(), column: "x".into(), }, E::DivisionByZero, E::DataException { detail: text() }, + E::MixedSrid { + collection: collection(), + field: "loc".into(), + index_srid: 4326, + row_srid: 32648, + }, E::NumericValueOutOfRange { detail: text() }, E::InvalidTextRepresentation { detail: text() }, E::DatatypeMismatch { detail: text() }, diff --git a/nodedb/src/control/gateway/error_map/resp.rs b/nodedb/src/control/gateway/error_map/resp.rs index 21ff9ebc9..2ee02b72e 100644 --- a/nodedb/src/control/gateway/error_map/resp.rs +++ b/nodedb/src/control/gateway/error_map/resp.rs @@ -92,9 +92,11 @@ impl GatewayErrorMap { | Error::UndefinedColumn { .. } | Error::TextColumn { .. } | Error::AmbiguousColumn { .. } + | Error::UngroupedColumn { .. } | Error::UnknownStrictField { .. } | Error::DivisionByZero | Error::DataException { .. } + | Error::MixedSrid { .. } | Error::NumericValueOutOfRange { .. } | Error::InvalidTextRepresentation { .. } | Error::DatatypeMismatch { .. } diff --git a/nodedb/src/control/local_dispatch/collect.rs b/nodedb/src/control/local_dispatch/collect.rs index 28f74425a..63134f42a 100644 --- a/nodedb/src/control/local_dispatch/collect.rs +++ b/nodedb/src/control/local_dispatch/collect.rs @@ -10,16 +10,27 @@ use crate::types::RequestId; #[derive(Debug)] pub(crate) enum DispatchCollectError { - OverBudget { bytes: usize }, + OverBudget { + bytes: usize, + }, ChannelClosed, + /// The final frame arrived, but a streamed chunk is not a msgpack row + /// array, so the chunks do not merge into one payload. `response` is the + /// final frame with an empty payload. Its status and write set stay as + /// the core sent them. + Unmergeable { + response: Box, + error: crate::Error, + }, } /// Drain a dispatched request's bounded response channel, enforcing a /// total-payload byte ceiling across streamed partials. /// /// Returns the final Response (non-streaming: pass-through; streaming: -/// concatenated payload) or an error if the channel closed without a -/// final chunk or if the accumulated payload would exceed the ceiling. +/// merged payload) or an error if the channel closed without a final chunk, +/// if the accumulated payload would exceed the ceiling, or if a streamed +/// chunk is not a msgpack row array. pub(crate) async fn collect_bounded_response( rx: &mut crate::control::ResponseReceiver, max_result_bytes: usize, @@ -68,14 +79,21 @@ pub(crate) async fn collect_bounded_response( } } - match final_response_meta { - Some(meta) => Ok(Response { - payload: Payload::from_vec( - crate::control::server::payload_merge::merge_msgpack_arrays(&chunks), - ), + let Some(meta) = final_response_meta else { + return Err(DispatchCollectError::ChannelClosed); + }; + match crate::control::server::payload_merge::merge_msgpack_arrays(&chunks) { + Ok(merged) => Ok(Response { + payload: Payload::from_vec(merged), ..meta }), - None => Err(DispatchCollectError::ChannelClosed), + Err(error) => Err(DispatchCollectError::Unmergeable { + response: Box::new(Response { + payload: Payload::empty(), + ..meta + }), + error, + }), } } @@ -141,6 +159,7 @@ pub(crate) async fn collect_under_deadline( Ok(Err(DispatchCollectError::ChannelClosed)) => Err(crate::Error::Dispatch { detail: format!("{context} channel closed"), }), + Ok(Err(DispatchCollectError::Unmergeable { error, .. })) => Err(error), } } @@ -243,7 +262,7 @@ mod collect_budget_tests { tx.send(final_rows(500)).await.unwrap(); drop(tx); let resp = collect_bounded_response(&mut rx, 1 << 20).await.unwrap(); - let elements = extract_msgpack_elements(resp.payload.as_ref()); + let elements = extract_msgpack_elements(resp.payload.as_ref()).unwrap(); assert_eq!( elements.len(), 2500, @@ -251,6 +270,48 @@ mod collect_budget_tests { ); } + /// A streamed chunk that is not a row array does not merge: the collect + /// reports the codec error and keeps the final frame's metadata. + #[tokio::test] + async fn streaming_a_non_array_chunk_is_unmergeable() { + let (tx, rx) = mpsc::channel(4); + let mut rx = crate::control::ResponseReceiver::from_channel(rx); + tx.send(partial_bytes(4)).await.unwrap(); + tx.send(final_rows(2)).await.unwrap(); + drop(tx); + match collect_bounded_response(&mut rx, 1024).await { + Err(DispatchCollectError::Unmergeable { response, error }) => { + assert!(matches!(error, crate::Error::Codec { .. }), "{error:?}"); + assert_eq!(response.status, Status::Ok); + assert!(response.payload.is_empty()); + } + other => panic!("expected an unmergeable collect, got {other:?}"), + } + } + + #[tokio::test] + async fn a_deadline_collect_reports_an_unmergeable_chunk_as_its_error() { + let (tx, rx) = mpsc::channel(4); + let mut rx = crate::control::ResponseReceiver::from_channel(rx); + tx.send(partial_bytes(4)).await.unwrap(); + tx.send(final_rows(2)).await.unwrap(); + drop(tx); + let result = collect_under_deadline( + &mut rx, + DeadlineCollect { + request_id: RequestId::new(7), + deadline: Instant::now() + std::time::Duration::from_secs(30), + max_result_bytes: 1024, + context: "gather on core 0", + }, + ) + .await; + assert!( + matches!(result, Err(crate::Error::Codec { .. })), + "got {result:?}" + ); + } + #[tokio::test] async fn streaming_over_budget_on_partial_aborts() { let (tx, rx) = mpsc::channel(4); @@ -261,7 +322,7 @@ mod collect_budget_tests { let err = collect_bounded_response(&mut rx, 1000).await.unwrap_err(); match err { DispatchCollectError::OverBudget { bytes } => assert!(bytes > 1000), - DispatchCollectError::ChannelClosed => panic!("expected OverBudget, got ChannelClosed"), + other => panic!("expected OverBudget, got {other:?}"), } } diff --git a/nodedb/src/control/otel/receiver.rs b/nodedb/src/control/otel/receiver.rs index be80ee279..272c1410c 100644 --- a/nodedb/src/control/otel/receiver.rs +++ b/nodedb/src/control/otel/receiver.rs @@ -426,19 +426,15 @@ pub(super) async fn authenticate_otel( .filter(|token| !token.is_empty()) .ok_or_else(|| "invalid authorization header".to_owned())?; + // A JWT goes through the shared bearer gate, which audits every refusal and login. + // An OTLP request carries no session, so the session id stays empty. if token.matches('.').count() == 2 - && let Some(registry) = &shared.jwks_registry - && let Ok((identity, verified)) = registry.validate_with_claims(token).await + && let Some((identity, _claims)) = + session_auth::authenticate_bearer_jwt(shared, token, peer_addr, "").await { - crate::control::security::jwt_policy::enforce_stateful_jwt_policy( - shared, - verified.claims(), - &identity, - ) - .map_err(|_| "invalid bearer token".to_owned())?; return admit_transport(shared, identity, peer_addr); } - let identity = session_auth::verify_api_key_identity(shared, token, "otlp", "OTLP") + let identity = session_auth::verify_api_key_identity(shared, token, peer_addr, "OTLP") .ok_or_else(|| "invalid bearer token".to_owned())?; admit_transport(shared, identity, peer_addr) } diff --git a/nodedb/src/control/planner/calvin/abort_error.rs b/nodedb/src/control/planner/calvin/abort_error.rs index 9ff3e12ec..7a17849e2 100644 --- a/nodedb/src/control/planner/calvin/abort_error.rs +++ b/nodedb/src/control/planner/calvin/abort_error.rs @@ -2,17 +2,26 @@ //! Mapping of a Calvin cross-shard ABORT verdict to the error the client sees. -use nodedb_cluster::calvin::AbortReason; +use nodedb_cluster::calvin::{AbortCause, AbortReason}; use crate::Error; +use crate::control::cluster::calvin::refusal::decode_refusal; -/// Pick the error for an ABORT verdict from the reason the verdict recorded. +/// Pick the error for an ABORT verdict from the cause the verdict recorded. /// /// A participant error never validated a read-set, so it must not be reported -/// as a serialization conflict. -pub fn calvin_abort_error(reason: AbortReason) -> Error { - match reason { - AbortReason::ParticipantError => Error::CalvinParticipantError, +/// as a serialization conflict. A participant error that carries a refusal is +/// the Data-Plane error the stage refused with. The client sees the SQLSTATE +/// a direct write of the same slice renders. +pub fn calvin_abort_error(cause: AbortCause) -> Error { + match cause.reason { + AbortReason::ParticipantError => match cause.refusal { + Some(refusal) => match decode_refusal(&refusal) { + Ok(code) => Error::DataPlane(code), + Err(error) => error, + }, + None => Error::CalvinParticipantError, + }, // A participant rejected the plans: they do not decode or route, or // hold no work for it. A resubmit of the same plans fails the same // way, so this is an internal error (`XX000`), never a retry class. @@ -37,20 +46,62 @@ pub fn calvin_abort_error(reason: AbortReason) -> Error { #[cfg(test)] mod tests { + use nodedb_types::error::sqlstate; + use super::*; + use crate::bridge::envelope::ErrorCode; + use crate::control::cluster::calvin::refusal::encode_refusal; + + fn refused(code: &ErrorCode) -> AbortCause { + AbortCause::refused(encode_refusal(code).expect("encode")) + } #[test] fn participant_error_is_not_reported_as_a_serialization_conflict() { assert!(matches!( - calvin_abort_error(AbortReason::ParticipantError), + calvin_abort_error(AbortReason::ParticipantError.into()), Error::CalvinParticipantError )); } + /// A refused stage answers its Data-Plane code, which renders the + /// SQLSTATE of the constraint it names. + #[test] + fn a_refused_participant_answers_its_refusal() { + let cases = [ + ( + "primary_key_immutable", + sqlstate::INTEGRITY_CONSTRAINT_VIOLATION, + ), + ("not_null", sqlstate::NOT_NULL_VIOLATION), + ]; + for (constraint, expected) in cases { + let code = ErrorCode::RejectedConstraint { + constraint: constraint.into(), + detail: format!("{constraint} refused the row"), + }; + let error = calvin_abort_error(refused(&code)); + assert!( + matches!(&error, Error::DataPlane(found) if *found == code), + "{constraint}: {error:?}" + ); + let (_, state, _) = crate::control::server::pgwire::types::error_to_sqlstate(&error); + assert_eq!(state, expected, "{constraint}"); + } + } + + #[test] + fn an_undecodable_refusal_is_an_internal_error() { + assert!(matches!( + calvin_abort_error(AbortCause::refused(vec![0xde, 0xad])), + Error::Internal { .. } + )); + } + #[test] fn a_rejected_plan_is_an_internal_error() { assert!(matches!( - calvin_abort_error(AbortReason::PlanRejected), + calvin_abort_error(AbortReason::PlanRejected.into()), Error::Internal { .. } )); } @@ -58,7 +109,7 @@ mod tests { #[test] fn a_superseded_collection_asks_for_a_retry() { assert!(matches!( - calvin_abort_error(AbortReason::CollectionSuperseded), + calvin_abort_error(AbortReason::CollectionSuperseded.into()), Error::RetryableSchemaChanged { .. } )); } @@ -66,7 +117,7 @@ mod tests { #[test] fn a_stale_read_set_is_a_serialization_conflict() { assert!(matches!( - calvin_abort_error(AbortReason::SerializationConflict), + calvin_abort_error(AbortReason::SerializationConflict.into()), Error::CalvinSerializationConflict )); } diff --git a/nodedb/src/control/planner/calvin/retry_loop.rs b/nodedb/src/control/planner/calvin/retry_loop.rs index 63ef39a8c..c11a52810 100644 --- a/nodedb/src/control/planner/calvin/retry_loop.rs +++ b/nodedb/src/control/planner/calvin/retry_loop.rs @@ -161,17 +161,19 @@ where // Terminal, NON-retryable: the global cross-shard verdict was ABORT // for a reason a fresh reconnaissance cannot change, so surface it // to the client immediately instead of burning retries. The - // verdict's reason picks the error: a stale read-set is SQLSTATE - // 40001, a participant error is not, and rejected plans are an - // internal error. `PredictionDrift` and `PartsLost` retry, in the + // verdict's cause picks the error: a stale read-set is SQLSTATE + // 40001, a refused stage is its own refusal, a participant error + // is not a conflict, and rejected plans are an internal error. + // `PredictionDrift` and `PartsLost` retry, in the // arm below: a multi-part transaction whose parts a leader change // lost staged nothing anywhere. - AttemptOutcome::Aborted { reason } - if reason != AbortReason::PredictionDrift && reason != AbortReason::PartsLost => + AttemptOutcome::Aborted(cause) + if cause.reason != AbortReason::PredictionDrift + && cause.reason != AbortReason::PartsLost => { - return Err(calvin_abort_error(reason)); + return Err(calvin_abort_error(cause)); } - AttemptOutcome::Aborted { .. } => { + AttemptOutcome::Aborted(_) => { // A participant found state other than the reconnaissance // predicted and wrote nothing, or the parts were lost. Every // participant released the attempt's locks before it acked, @@ -257,7 +259,7 @@ mod tests { let mut seen: u32 = 0; while let Some(txn) = rx.recv().await { let verdict = if seen < drift_count { - VerdictOutcome::Abort(AbortReason::PredictionDrift) + VerdictOutcome::Abort(AbortReason::PredictionDrift.into()) } else { VerdictOutcome::Commit }; @@ -657,7 +659,7 @@ mod tests { let mut first = true; while let Some(txn) = rx.recv().await { let verdict = if first { - VerdictOutcome::Abort(reason) + VerdictOutcome::Abort(reason.into()) } else { VerdictOutcome::Commit }; diff --git a/nodedb/src/control/planner/calvin/submit/local.rs b/nodedb/src/control/planner/calvin/submit/local.rs index f05c08281..a1394cb38 100644 --- a/nodedb/src/control/planner/calvin/submit/local.rs +++ b/nodedb/src/control/planner/calvin/submit/local.rs @@ -195,7 +195,7 @@ pub(crate) async fn submit_prepared_and_await( // The verdict's reason picks the error the client sees. match report.outcome { AttemptOutcome::Completed => {} - AttemptOutcome::Aborted { reason } => return Err(calvin_abort_error(reason)), + AttemptOutcome::Aborted(cause) => return Err(calvin_abort_error(cause)), } if binds_authorization { crate::control::security::auth_lease::calvin_write_barrier( diff --git a/nodedb/src/control/planner/calvin/tx_class/write_keys/columnar.rs b/nodedb/src/control/planner/calvin/tx_class/write_keys/columnar.rs index f90b8e9d9..b8da44092 100644 --- a/nodedb/src/control/planner/calvin/tx_class/write_keys/columnar.rs +++ b/nodedb/src/control/planner/calvin/tx_class/write_keys/columnar.rs @@ -86,7 +86,6 @@ pub(super) fn add_array_keys(keys: &mut WriteKeys, op: &ArrayOp) -> crate::Resul | ArrayOp::RekeyArray { .. } | ArrayOp::PurgeArrayDrop { .. } | ArrayOp::Slice { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } | ArrayOp::SurrogateBitmapScan { .. } => Err(not_a_write("an array read or DDL op")), diff --git a/nodedb/src/control/planner/calvin/write_class.rs b/nodedb/src/control/planner/calvin/write_class.rs index f91227280..8f1471ec1 100644 --- a/nodedb/src/control/planner/calvin/write_class.rs +++ b/nodedb/src/control/planner/calvin/write_class.rs @@ -309,7 +309,6 @@ fn array_is_write(op: &ArrayOp) -> bool { | ArrayOp::RekeyArray { .. } | ArrayOp::PurgeArrayDrop { .. } | ArrayOp::Slice { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } | ArrayOp::SurrogateBitmapScan { .. } => false, diff --git a/nodedb/src/control/planner/catalog_adapter/cagg_relation.rs b/nodedb/src/control/planner/catalog_adapter/cagg_relation.rs index 9ab1997e7..04d8388f8 100644 --- a/nodedb/src/control/planner/catalog_adapter/cagg_relation.rs +++ b/nodedb/src/control/planner/catalog_adapter/cagg_relation.rs @@ -115,10 +115,10 @@ fn relation_columns( for expr in &def.aggregates { let input = source_column(&expr.source_column); let column = match &expr.function { - AggFunction::Count | AggFunction::TopK(_) => { + AggFunction::Count | AggFunction::CountDistinct | AggFunction::TopK(_) => { derived(&expr.output_column, SqlDataType::Int64) } - AggFunction::Avg | AggFunction::CountDistinct | AggFunction::Percentile(_) => { + AggFunction::Avg | AggFunction::Percentile(_) => { derived(&expr.output_column, SqlDataType::Float64) } // An integer total can pass the 64-bit range and turns exact diff --git a/nodedb/src/control/planner/catalog_adapter/sql_catalog_impl.rs b/nodedb/src/control/planner/catalog_adapter/sql_catalog_impl.rs index 977b8a72f..2ea36bea9 100644 --- a/nodedb/src/control/planner/catalog_adapter/sql_catalog_impl.rs +++ b/nodedb/src/control/planner/catalog_adapter/sql_catalog_impl.rs @@ -181,6 +181,33 @@ impl SqlCatalog for OriginCatalog { })) } + /// The time key the collection stores, as `CREATE COLLECTION` designated + /// it. A later `ALTER ADD COLUMN` cannot change it. + fn timeseries_time_key( + &self, + _database_id: nodedb_types::DatabaseId, + name: &str, + ) -> std::result::Result, SqlCatalogError> { + let stored = self + .credentials + .catalog() + .get_collection(self.database_id, self.tenant_id, name) + .map_err(|e| SqlCatalogError::Unreadable { + relation: name.to_string(), + detail: e.to_string(), + })?; + Ok(stored.and_then(|stored| { + if let nodedb_types::CollectionType::Columnar( + nodedb_types::columnar::ColumnarProfile::Timeseries { time_key, .. }, + ) = stored.collection_type + { + Some(time_key) + } else { + None + } + })) + } + fn resolve_regclass( &self, _database_id: nodedb_types::DatabaseId, @@ -269,11 +296,9 @@ impl SqlCatalog for OriginCatalog { } fn lookup_array(&self, name: &str) -> Option { - use nodedb_array::schema::{ArraySchema, AttrType as EAT, DimType as EDT}; + use nodedb_array::schema::{ArraySchema, DimType as EDT}; use nodedb_array::types::domain::DomainBound; - use nodedb_sql::types_array::{ - ArrayAttrAst, ArrayAttrType, ArrayDimAst, ArrayDimType, ArrayDomainBound, - }; + use nodedb_sql::types_array::{ArrayAttrAst, ArrayDimAst, ArrayDimType, ArrayDomainBound}; let entry = { // A poisoned lock still holds a consistent mirror: every writer @@ -311,12 +336,9 @@ impl SqlCatalog for OriginCatalog { .iter() .map(|a| ArrayAttrAst { name: a.name.clone(), - dtype: match a.dtype { - EAT::Int64 => ArrayAttrType::Int64, - EAT::Float64 => ArrayAttrType::Float64, - EAT::String => ArrayAttrType::String, - EAT::Bytes => ArrayAttrType::Bytes, - }, + dtype: crate::control::planner::sql_plan_convert::array_convert::sql_attr_type( + a.dtype, + ), nullable: a.nullable, }) .collect(); diff --git a/nodedb/src/control/planner/context/mod.rs b/nodedb/src/control/planner/context/mod.rs index 90d9a9229..657787e07 100644 --- a/nodedb/src/control/planner/context/mod.rs +++ b/nodedb/src/control/planner/context/mod.rs @@ -11,6 +11,6 @@ pub mod query; pub mod security; pub mod system_security; -pub use query::{PlanSqlWithRlsParams, QueryContext, SYSTEM_FUNCTION_NAMES}; +pub use query::{BulkLoadRoute, PlanSqlWithRlsParams, QueryContext, SYSTEM_FUNCTION_NAMES}; pub use security::{PermissionTreeSource, PlanSecurityContext}; pub use system_security::SystemPlanSecurity; diff --git a/nodedb/src/control/planner/context/query/mod.rs b/nodedb/src/control/planner/context/query/mod.rs index b834f5ec1..b0beaac2a 100644 --- a/nodedb/src/control/planner/context/query/mod.rs +++ b/nodedb/src/control/planner/context/query/mod.rs @@ -4,9 +4,11 @@ mod context; mod convert_plans; mod functions; mod planning; +mod planning_built; mod tuning; pub use context::QueryContext; pub use functions::SYSTEM_FUNCTION_NAMES; pub use planning::PlanSqlWithRlsParams; +pub use planning_built::BulkLoadRoute; pub use tuning::DEFAULT_SHUFFLE_AGG_THRESHOLD; diff --git a/nodedb/src/control/planner/context/query/planning.rs b/nodedb/src/control/planner/context/query/planning.rs index 239ef8ee5..f07c70203 100644 --- a/nodedb/src/control/planner/context/query/planning.rs +++ b/nodedb/src/control/planner/context/query/planning.rs @@ -25,7 +25,7 @@ use crate::control::server::shared::returning::{ /// and the parameterized path ([`QueryContext::plan_sql_with_params_and_rls_and_versions`]) /// resolve the clause before conversion and need the same output schema built /// from it, so the pairing lives here once. -fn resolve_returning_and_output_schema( +pub(super) fn resolve_returning_and_output_schema( plans: &[nodedb_sql::types::SqlPlan], returning_items: Option<&str>, catalog: &C, @@ -269,50 +269,8 @@ impl QueryContext { let (mut tasks, output_schema, mut version_set, cache_eligibility) = self .plan_with_nodedb_sql_for_purpose(sql, tenant_id, database_id, purpose, returning_items) .await?; - - // Versions read BEFORE injection, never after: injection reads live - // policy/grant state under its own lock, and a mutation racing in - // between will otherwise let a post-injection read stamp a version - // newer than what was actually filtered against, making a plan built - // from stale state compare as fresh forever. A pre-injection read - // only ever under-states freshness (an extra harmless replan), never - // over-states it. - // - // The permission tree is read here, after planning's last await, so - // no lock is held across a request to a surrogate's collection home. - let permission_view = sec.permission_tree.read().await; - let permission_cache = permission_view.as_deref(); - let permission_tree_version = permission_cache - .map(|c| c.tenant_version(tenant_id.as_u64())) - .unwrap_or(0); - let rls_version = sec.rls_store.tenant_version(tenant_id.as_u64()); - - // Inject RLS predicates. - crate::control::planner::rls_injection::inject_rls( - &mut tasks, - sec.rls_store, - self.catalog_inputs.credentials.catalog(), - sec.auth, - )?; - - // Refuse what column redaction cannot cover (aggregates over a - // redacted column, graph traversals), before anything is dispatched. - crate::control::planner::redaction_refusal::refuse_unredactable_tasks( - &tasks, - sec.auth, - sec.redaction_store, - )?; - - // Inject permission tree filters (hierarchical ACL). - if let Some(cache) = permission_cache { - crate::control::planner::rls_injection::inject_permission_tree( - &mut tasks, cache, sec.auth, - )?; - } - - version_set.set_permission_tree_version(permission_tree_version); - version_set.set_rls_version(rls_version); - + self.apply_task_policy(&mut tasks, sec, tenant_id, &mut version_set) + .await?; Ok((tasks, output_schema, version_set, cache_eligibility)) } @@ -400,42 +358,9 @@ impl QueryContext { if let Some(clause) = &returning { attach_returning_spec(&mut tasks, &clause.spec)?; } - - // Versions read BEFORE injection — see the comment on the sibling - // planning path in this file for why a post-injection read is unsafe, - // and why the permission tree is read only after the last await. - let permission_view = sec.permission_tree.read().await; - let permission_cache = permission_view.as_deref(); - let permission_tree_version = permission_cache - .map(|c| c.tenant_version(tenant_id.as_u64())) - .unwrap_or(0); - let rls_version = sec.rls_store.tenant_version(tenant_id.as_u64()); - - // Inject RLS predicates. - crate::control::planner::rls_injection::inject_rls( - &mut tasks, - sec.rls_store, - self.catalog_inputs.credentials.catalog(), - sec.auth, - )?; - - // Refuse what column redaction cannot cover (aggregates over a - // redacted column, graph traversals), before anything is dispatched. - crate::control::planner::redaction_refusal::refuse_unredactable_tasks( - &tasks, - sec.auth, - sec.redaction_store, - )?; - - if let Some(cache) = permission_cache { - crate::control::planner::rls_injection::inject_permission_tree( - &mut tasks, cache, sec.auth, - )?; - } - let mut version_set = catalog.take_recorded_versions(); - version_set.set_permission_tree_version(permission_tree_version); - version_set.set_rls_version(rls_version); + self.apply_task_policy(&mut tasks, sec, tenant_id, &mut version_set) + .await?; Ok((tasks, output_schema, version_set)) } } diff --git a/nodedb/src/control/planner/context/query/planning_built.rs b/nodedb/src/control/planner/context/query/planning_built.rs new file mode 100644 index 000000000..75a1b961c --- /dev/null +++ b/nodedb/src/control/planner/context/query/planning_built.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Planning of `SqlPlan`s a caller builds against the catalog, and the policy +//! pass every planned task set takes before it leaves the planner. + +use nodedb_physical::physical_task::PhysicalTask; +use nodedb_sql::catalog::{ArrayCatalogView, SqlCatalog}; +use nodedb_sql::engine_rules::{BulkLoadParams, BulkLoadPlan}; +use nodedb_sql::types::SqlPlan; + +use super::QueryContext; +use super::planning::resolve_returning_and_output_schema; +use crate::control::planner::context::security::PlanSecurityContext; +use crate::control::planner::descriptor_set::DescriptorVersionSet; +use crate::control::planner::plan_error_map::map_plan_error; +use crate::control::planner::sql_plan_convert::PlanningPurpose; +use crate::control::server::response_shape::schema::OutputSchema; + +/// How a bulk load's target loads rows, answered by its engine rules. +#[derive(Debug)] +pub enum BulkLoadRoute { + /// Each row loads as one `INSERT` statement. + RowInserts, + /// The rows decode against this array schema and load in cell chunks. + ArrayCells(ArrayCatalogView), +} + +impl QueryContext { + /// Ask `collection`'s engine rules how a bulk load into it loads rows. + /// + /// The error is the planner's own, so a caller can render an engine's + /// refusal (timeseries, spatial) with its exact text. + pub fn bulk_load_route( + &self, + collection: &str, + tenant_id: crate::types::TenantId, + database_id: crate::types::DatabaseId, + ) -> nodedb_sql::Result { + let catalog = self + .catalog_inputs + .build_adapter(tenant_id.as_u64(), database_id); + let params = BulkLoadParams { + collection: collection.to_string(), + cells: Vec::new(), + }; + match nodedb_sql::planner::dml::plan_bulk_load(params, &catalog)? { + BulkLoadPlan::RowInserts => Ok(BulkLoadRoute::RowInserts), + BulkLoadPlan::CellChunks(_) => catalog + .lookup_array(collection) + .map(BulkLoadRoute::ArrayCells) + .ok_or_else(|| nodedb_sql::SqlError::Unsupported { + detail: format!("COPY: array '{collection}' was dropped"), + }), + } + } + + /// Plan the `SqlPlan`s that `build` returns, then convert and secure them. + /// + /// `build` runs against a fresh catalog adapter, the one SQL planning + /// uses. A bulk load builds its plans this way: it has rows, not SQL + /// text. The task set takes the same policy pass as planned SQL: row + /// filters, redaction refusals and the permission tree. Returns the + /// tasks, the output schema, and the descriptor versions the adapter + /// recorded, for lease admission. + pub async fn plan_built_with_rls_and_versions( + &self, + build: B, + tenant_id: crate::types::TenantId, + database_id: crate::types::DatabaseId, + sec: &PlanSecurityContext<'_>, + ) -> crate::Result<(Vec, OutputSchema, DescriptorVersionSet)> + where + B: FnOnce(&dyn SqlCatalog) -> nodedb_sql::Result> + Send, + { + let catalog = self + .catalog_inputs + .build_adapter(tenant_id.as_u64(), database_id) + .with_session_sequences(self.session_sequences()); + let plans = build(&catalog).map_err(|error| map_plan_error(error, tenant_id))?; + let (output_schema, _) = + resolve_returning_and_output_schema(&plans, None, &catalog, database_id, tenant_id)?; + let mut tasks = self + .convert_prefetched(&plans, PlanningPurpose::Execute, database_id, tenant_id) + .await?; + let mut version_set = catalog.take_recorded_versions(); + self.apply_task_policy(&mut tasks, sec, tenant_id, &mut version_set) + .await?; + Ok((tasks, output_schema, version_set)) + } + + /// Inject row filters and permission-tree filters into `tasks`, and + /// refuse a task that column redaction cannot cover. Records the policy + /// versions the pass read into `version_set`. + /// + /// The versions are read BEFORE injection, never after. Injection reads + /// live policy and grant state under its own lock. With a mutation racing + /// in between, a post-injection read stamps a version newer than the one + /// filtered against, and a plan built from stale state then compares as + /// fresh forever. A pre-injection read only under-states freshness, + /// which costs one harmless replan. + /// + /// The permission tree is read here, after planning's last await, so no + /// lock is held across a request to a surrogate's collection home. + pub(super) async fn apply_task_policy( + &self, + tasks: &mut [PhysicalTask], + sec: &PlanSecurityContext<'_>, + tenant_id: crate::types::TenantId, + version_set: &mut DescriptorVersionSet, + ) -> crate::Result<()> { + let permission_view = sec.permission_tree.read().await; + let permission_cache = permission_view.as_deref(); + let permission_tree_version = permission_cache + .map(|c| c.tenant_version(tenant_id.as_u64())) + .unwrap_or(0); + let rls_version = sec.rls_store.tenant_version(tenant_id.as_u64()); + let redaction_version = sec.redaction_store.version(); + + // Refuse a read of a redacted column outside the projection before + // row-level security adds its own predicates: a read policy's + // predicate is the administrator's, not the requester's. + crate::control::planner::redaction_refusal::refuse_redacted_expression_reads_in_tasks( + tasks, + &crate::control::planner::redaction_refusal::RedactionReadCheck { + auth: sec.auth, + store: sec.redaction_store, + catalog: self.catalog_inputs.credentials.catalog(), + }, + )?; + + crate::control::planner::rls_injection::inject_rls( + tasks, + sec.rls_store, + self.catalog_inputs.credentials.catalog(), + sec.auth, + )?; + + // Refuse what column redaction cannot cover (aggregates over a + // redacted column, graph traversals), before anything is dispatched. + crate::control::planner::redaction_refusal::refuse_unredactable_tasks( + tasks, + sec.auth, + sec.redaction_store, + )?; + + // Inject permission tree filters (hierarchical ACL). + if let Some(cache) = permission_cache { + crate::control::planner::rls_injection::inject_permission_tree(tasks, cache, sec.auth)?; + } + + version_set.set_permission_tree_version(permission_tree_version); + version_set.set_rls_version(rls_version); + version_set.set_redaction_version(redaction_version); + Ok(()) + } +} diff --git a/nodedb/src/control/planner/descriptor_set.rs b/nodedb/src/control/planner/descriptor_set.rs index 611b6b341..a3525b2cc 100644 --- a/nodedb/src/control/planner/descriptor_set.rs +++ b/nodedb/src/control/planner/descriptor_set.rs @@ -42,6 +42,9 @@ pub struct DescriptorVersionSet { permission_tree_version: u64, /// The tenant's RLS policy-store version observed at plan-build time. rls_version: u64, + /// The redaction policy-store version the plan's column reads were + /// checked against. + redaction_version: u64, } impl DescriptorVersionSet { @@ -67,6 +70,16 @@ impl DescriptorVersionSet { self.rls_version } + /// Record the redaction policy-store version this plan was checked + /// against. + pub fn set_redaction_version(&mut self, version: u64) { + self.redaction_version = version; + } + + pub fn redaction_version(&self) -> u64 { + self.redaction_version + } + /// Record a descriptor read at the given version. O(log n) /// binary search + O(n) worst-case insert on first /// observation; O(log n) no-op on re-observation of the @@ -114,23 +127,27 @@ impl DescriptorVersionSet { /// Used by the plan cache on every hit — a cached plan is /// only reusable while every descriptor it was built /// against still exists at the same version, AND the - /// tenant's permission-tree and RLS policy state has not - /// changed since the plan was built. - pub fn all_fresh( - &self, - current: F, - current_permission_tree_version: u64, - current_rls_version: u64, - ) -> bool + /// tenant's permission-tree, RLS, and redaction policy state has + /// not changed since the plan was built. + pub fn all_fresh(&self, current: F, policies: PolicyVersions) -> bool where F: Fn(&DescriptorId) -> Option, { - self.permission_tree_version == current_permission_tree_version - && self.rls_version == current_rls_version + self.permission_tree_version == policies.permission_tree + && self.rls_version == policies.rls + && self.redaction_version == policies.redaction && self.entries.iter().all(|(id, v)| current(id) == Some(*v)) } } +/// The policy-store versions a cached plan is compared against. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PolicyVersions { + pub permission_tree: u64, + pub rls: u64, + pub redaction: u64, +} + /// Total-order comparison for `DescriptorId` since the cluster /// type does not derive `Ord` (only `Hash + Eq`). We define our /// own by `(kind discriminant, database_id, tenant_id, name)`. @@ -195,6 +212,14 @@ mod tests { assert_eq!(entries, vec![(3, 1), (9, 2)]); } + fn policies(permission_tree: u64, rls: u64, redaction: u64) -> PolicyVersions { + PolicyVersions { + permission_tree, + rls, + redaction, + } + } + #[test] fn all_fresh_true_when_versions_match() { let mut set = DescriptorVersionSet::new(); @@ -205,7 +230,7 @@ mod tests { "b" => Some(2), _ => None, }; - assert!(set.all_fresh(lookup, 0, 0)); + assert!(set.all_fresh(lookup, policies(0, 0, 0))); } #[test] @@ -218,7 +243,7 @@ mod tests { "b" => Some(3), // bumped _ => None, }; - assert!(!set.all_fresh(lookup, 0, 0)); + assert!(!set.all_fresh(lookup, policies(0, 0, 0))); } #[test] @@ -226,14 +251,14 @@ mod tests { let mut set = DescriptorVersionSet::new(); set.record(id("a"), 1); let lookup = |_: &DescriptorId| None; - assert!(!set.all_fresh(lookup, 0, 0)); + assert!(!set.all_fresh(lookup, policies(0, 0, 0))); } #[test] fn empty_set_is_always_fresh() { let set = DescriptorVersionSet::new(); let lookup = |_: &DescriptorId| None; - assert!(set.all_fresh(lookup, 0, 0)); + assert!(set.all_fresh(lookup, policies(0, 0, 0))); } #[test] @@ -241,9 +266,9 @@ mod tests { let mut set = DescriptorVersionSet::new(); set.set_permission_tree_version(3); let lookup = |_: &DescriptorId| None; - assert!(set.all_fresh(lookup, 3, 0)); + assert!(set.all_fresh(lookup, policies(3, 0, 0))); assert!( - !set.all_fresh(lookup, 4, 0), + !set.all_fresh(lookup, policies(4, 0, 0)), "a revoked grant bumps the tenant's permission version, so a stale stamp must not read as fresh" ); } @@ -253,8 +278,18 @@ mod tests { let mut set = DescriptorVersionSet::new(); set.set_rls_version(5); let lookup = |_: &DescriptorId| None; - assert!(set.all_fresh(lookup, 0, 5)); - assert!(!set.all_fresh(lookup, 0, 6)); + assert!(set.all_fresh(lookup, policies(0, 5, 0))); + assert!(!set.all_fresh(lookup, policies(0, 6, 0))); + } + + /// A redaction policy change must re-check the plan's column reads. + #[test] + fn all_fresh_false_on_redaction_version_mismatch() { + let mut set = DescriptorVersionSet::new(); + set.set_redaction_version(2); + let lookup = |_: &DescriptorId| None; + assert!(set.all_fresh(lookup, policies(0, 0, 2))); + assert!(!set.all_fresh(lookup, policies(0, 0, 3))); } #[test] diff --git a/nodedb/src/control/planner/plan_error_map.rs b/nodedb/src/control/planner/plan_error_map.rs index c85b68a02..c4b01a807 100644 --- a/nodedb/src/control/planner/plan_error_map.rs +++ b/nodedb/src/control/planner/plan_error_map.rs @@ -71,6 +71,9 @@ pub(crate) fn map_plan_error( nodedb_sql::SqlError::AmbiguousColumn { column } => { crate::Error::AmbiguousColumn { column } } + nodedb_sql::SqlError::UngroupedColumn { column } => { + crate::Error::UngroupedColumn { column } + } nodedb_sql::SqlError::TextColumn { collection, column, diff --git a/nodedb/src/control/planner/redaction_refusal/aggregate.rs b/nodedb/src/control/planner/redaction_refusal/aggregate.rs index bb4ccc2ff..27f7d5ff3 100644 --- a/nodedb/src/control/planner/redaction_refusal/aggregate.rs +++ b/nodedb/src/control/planner/redaction_refusal/aggregate.rs @@ -38,7 +38,7 @@ fn refuse_specs( for spec in specs { for field in read_fields(spec) { if ctx.field_is_redacted(collection, field) { - return Err(refusal(collection, field)); + return Err(refusal(ctx, collection, field)); } } } @@ -67,34 +67,17 @@ pub(super) fn refuse_input_aggregates( Ok(()) } -/// Refuse when `field` is a redacted column of one of `sources`. -/// -/// A field qualified with a source's qualifier (`o.ssn`) is checked against -/// the sources that carry that qualifier. Any other field, unqualified or -/// with a qualifier no source carries (a derived table's alias, a nested -/// field path), is checked against every source, so an ambiguous name -/// cannot slip past. [`RefusalCtx::field_is_redacted`] matches the field as -/// written and by its column part. +/// Refuse when `field` is a redacted column of one of `sources`, per +/// [`RefusalCtx::redacted_source`]. fn refuse_source_field( ctx: &RefusalCtx<'_>, sources: &[(String, String)], field: &str, ) -> crate::Result<()> { - let qualifier = field.split_once('.').map(|(qualifier, _)| qualifier); - let names_a_source = qualifier.is_some_and(|qualifier| { - sources - .iter() - .any(|(source_qualifier, _)| source_qualifier == qualifier) - }); - for (source_qualifier, collection) in sources { - if names_a_source && qualifier != Some(source_qualifier.as_str()) { - continue; - } - if !collection.is_empty() && ctx.field_is_redacted(collection, field) { - return Err(refusal(collection, field)); - } + match ctx.redacted_source(sources, field) { + Some(collection) => Err(refusal(ctx, collection, field)), + None => Ok(()), } - Ok(()) } /// Refuse when a grouped timeseries scan aggregates a redacted column. @@ -109,7 +92,7 @@ pub(super) fn refuse_timeseries_aggregates( } for (_function, field) in aggregates { if !is_count_star(field) && ctx.field_is_redacted(collection, field) { - return Err(refusal(collection, field)); + return Err(refusal(ctx, collection, field)); } } Ok(()) @@ -130,26 +113,22 @@ pub(super) fn refuse_facet_fields( } for field in fields { if ctx.field_is_redacted(collection, field) { - return Err(crate::Error::PlanError { - detail: format!( - "column '{field}' on '{collection}' is redacted for this role: faceting a \ - redacted column is not permitted — facet counts enumerate the stored values \ - the policy protects" - ), - }); + return Err(ctx.refusal(format!( + "column '{field}' on '{collection}' is redacted for this role: faceting a \ + redacted column is not permitted — facet counts enumerate the stored values \ + the policy protects" + ))); } } Ok(()) } -fn refusal(collection: &str, field: &str) -> crate::Error { - crate::Error::PlanError { - detail: format!( - "column '{field}' on '{collection}' is redacted for this role: aggregating a redacted \ - column is not permitted — the aggregate is computed over the stored values, so \ - masking its result would not protect them" - ), - } +fn refusal(ctx: &RefusalCtx<'_>, collection: &str, field: &str) -> crate::Error { + ctx.refusal(format!( + "column '{field}' on '{collection}' is redacted for this role: aggregating a redacted \ + column is not permitted — the aggregate is computed over the stored values, so \ + masking its result would not protect them" + )) } /// Every stored column one aggregate spec reads. @@ -179,7 +158,7 @@ fn is_count_star(field: &str) -> bool { } /// Collect every column reference in `expr`. -fn collect_columns<'e>(expr: &'e SqlExpr, out: &mut Vec<&'e str>) { +pub(super) fn collect_columns<'e>(expr: &'e SqlExpr, out: &mut Vec<&'e str>) { match expr { SqlExpr::Column(name) | SqlExpr::OldColumn(name) | SqlExpr::ExcludedColumn(name) => { out.push(name.as_str()); diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/columnar.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/columnar.rs new file mode 100644 index 000000000..50d1ae2b6 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/columnar.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Column reads of the columnar-family engines: columnar, timeseries, and +//! spatial. + +use nodedb_physical::physical_plan::{ColumnarOp, SpatialOp, TimeseriesOp}; + +use super::reads::{Clause, NodeReads}; + +pub(super) fn columnar_reads(op: &ColumnarOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + ColumnarOp::Scan { + filters, + sort_keys, + computed_columns, + .. + } => { + out.filter_bytes(filters)?; + out.sort_keys(sort_keys); + out.computed_bytes(computed_columns)?; + } + ColumnarOp::Insert { + on_conflict_updates, + .. + } => out.updates(on_conflict_updates), + ColumnarOp::Update { + filters, updates, .. + } + | ColumnarOp::ResolveDml { + filters, updates, .. + } => { + out.filter_bytes(filters)?; + out.updates(updates); + } + ColumnarOp::Delete { filters, .. } => out.filter_bytes(filters)?, + // Resolved rows, snapshot reads, and truncation read no column in a + // predicate or an expression. + ColumnarOp::ResolvedUpdate { .. } + | ColumnarOp::ResolvedDelete { .. } + | ColumnarOp::MaterializeScan { .. } + | ColumnarOp::Truncate { .. } => {} + } + Ok(()) +} + +/// A timeseries scan's aggregates are refused by the aggregate check. Its +/// group keys, predicates, sort keys, and computed columns are read here. +pub(super) fn timeseries_reads(op: &TimeseriesOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + TimeseriesOp::Scan { + filters, + sort_keys, + group_by, + computed_columns, + .. + } => { + out.filter_bytes(filters)?; + out.sort_keys(sort_keys); + out.group_names(group_by); + out.computed_bytes(computed_columns)?; + } + TimeseriesOp::Ingest { .. } + | TimeseriesOp::ResolveIngest(_) + | TimeseriesOp::Truncate { .. } + | TimeseriesOp::RefreshAggregate { .. } + | TimeseriesOp::Partitions { .. } + | TimeseriesOp::ReadAggregate { .. } => {} + } + Ok(()) +} + +/// A spatial predicate reads its geometry column and its attribute filters. +pub(super) fn spatial_reads(op: &SpatialOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + SpatialOp::Scan { + field, + attribute_filters, + .. + } => { + out.field(field, Clause::Filter); + out.filter_bytes(attribute_filters)?; + } + SpatialOp::Insert { .. } | SpatialOp::Delete { .. } => {} + } + Ok(()) +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/document.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/document.rs new file mode 100644 index 000000000..a91792513 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/document.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Column reads of the document engine's operations. + +use nodedb_physical::physical_plan::document::MergeActionOp; +use nodedb_physical::physical_plan::{DocumentOp, PhysicalPlan}; + +use super::reads::{Clause, NodeReads}; + +/// The columns `op` reads outside its plain projection. +/// +/// An orchestrated write's source rows are a child plan, checked on their +/// own. A source relation with a redaction rule is refused whole elsewhere. +/// An assignment's source-qualified columns (`s.col`) belong to that source, +/// so only its target columns are read here. +pub(super) fn reads(op: &DocumentOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + DocumentOp::Scan { + sort_keys, + filters, + distinct, + projection, + computed_columns, + window_functions, + .. + } => { + out.filter_bytes(filters)?; + out.sort_keys(sort_keys); + out.computed_bytes(computed_columns)?; + out.window_bytes(window_functions)?; + out.distinct(*distinct, projection); + } + DocumentOp::RangeScan { field, .. } => out.field(field, Clause::IndexProbe), + DocumentOp::IndexLookup { path, .. } => out.field(path, Clause::IndexProbe), + DocumentOp::IndexedFetch { path, filters, .. } => { + out.field(path, Clause::IndexProbe); + out.filter_bytes(filters)?; + } + // A read, update, or delete by key reports whether the key exists. + DocumentOp::PointGet { .. } | DocumentOp::PointDelete { .. } => out.key_lookup = true, + DocumentOp::PointUpdate { updates, .. } => { + out.key_lookup = true; + out.updates(updates); + } + DocumentOp::Upsert { + on_conflict_updates: updates, + .. + } + | DocumentOp::BatchUpsert { + on_conflict_updates: updates, + .. + } => out.updates(updates), + DocumentOp::BulkUpdate { + filters, updates, .. + } => { + out.filter_bytes(filters)?; + out.updates(updates); + } + DocumentOp::BulkDelete { filters, .. } => out.filter_bytes(filters)?, + DocumentOp::UpdateFromJoin { + target_join_col, + target_filters, + updates, + source_alias, + .. + } => { + out.field(target_join_col, Clause::JoinKey); + out.filter_bytes(target_filters)?; + out.updates_of_target(updates, Some(source_alias)); + } + DocumentOp::Merge { + target_join_col, + clauses, + source_alias, + .. + } => { + out.field(target_join_col, Clause::JoinKey); + for clause in clauses { + match &clause.action { + MergeActionOp::Update { updates } => { + out.updates_of_target(updates, Some(source_alias)); + } + MergeActionOp::Delete + | MergeActionOp::Insert { .. } + | MergeActionOp::DoNothing => {} + } + } + } + DocumentOp::ResolveWrite(inner) => reads(inner, out)?, + // Writes of a caller-supplied row, bulk loads, DDL, and maintenance + // read no column in a predicate or an expression. + DocumentOp::PointPut { .. } + | DocumentOp::PointInsert { .. } + | DocumentOp::BatchInsert { .. } + | DocumentOp::Register { .. } + | DocumentOp::DropIndex { .. } + | DocumentOp::BackfillIndex { .. } + | DocumentOp::Truncate { .. } + | DocumentOp::EstimateCount { .. } + | DocumentOp::InsertSelect { .. } + | DocumentOp::MaterializeScan { .. } + | DocumentOp::ApplyBalanceDelta { .. } + | DocumentOp::ResolvedWrite { .. } => {} + } + Ok(()) +} + +/// The child plan `op` reads its rows from, if any. +pub(super) fn child(op: &DocumentOp) -> Option<&PhysicalPlan> { + match op { + DocumentOp::InsertSelect { source, .. } + | DocumentOp::Merge { source, .. } + | DocumentOp::UpdateFromJoin { source, .. } => source.query_plan(), + DocumentOp::ResolveWrite(inner) => child(inner), + DocumentOp::PointGet { .. } + | DocumentOp::PointPut { .. } + | DocumentOp::PointInsert { .. } + | DocumentOp::PointDelete { .. } + | DocumentOp::PointUpdate { .. } + | DocumentOp::Scan { .. } + | DocumentOp::BatchInsert { .. } + | DocumentOp::RangeScan { .. } + | DocumentOp::Register { .. } + | DocumentOp::IndexLookup { .. } + | DocumentOp::IndexedFetch { .. } + | DocumentOp::DropIndex { .. } + | DocumentOp::BackfillIndex { .. } + | DocumentOp::Truncate { .. } + | DocumentOp::EstimateCount { .. } + | DocumentOp::Upsert { .. } + | DocumentOp::BatchUpsert { .. } + | DocumentOp::BulkUpdate { .. } + | DocumentOp::BulkDelete { .. } + | DocumentOp::MaterializeScan { .. } + | DocumentOp::ApplyBalanceDelta { .. } + | DocumentOp::ResolvedWrite { .. } => None, + } +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/kv.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/kv.rs new file mode 100644 index 000000000..e930473f2 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/kv.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Column reads of the key-value engine's operations. + +use nodedb_physical::physical_plan::KvOp; + +use super::reads::NodeReads; + +pub(super) fn reads(op: &KvOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + KvOp::Scan { + filters, + computed_columns, + sort_keys, + .. + } => { + out.filter_bytes(filters)?; + out.computed_bytes(computed_columns)?; + out.sort_keys(sort_keys); + } + KvOp::InsertOnConflictUpdate { updates, .. } => out.updates(updates), + // The assignments carry encoded literals, so only the predicate reads. + KvOp::PredicateUpdate { filters, .. } | KvOp::PredicateDelete { filters, .. } => { + out.filter_bytes(filters)?; + } + KvOp::ResolveWrite(inner) => reads(inner, out)?, + // A read or delete by key reports whether the key exists. + KvOp::Get { .. } + | KvOp::Delete { .. } + | KvOp::GetTtl { .. } + | KvOp::BatchGet { .. } + | KvOp::FieldGet { .. } => out.key_lookup = true, + // Writes of a caller-supplied entry, counters, sorted-index reads, + // DDL, and maintenance read no value column in a predicate or an + // expression. + KvOp::Put { .. } + | KvOp::Insert { .. } + | KvOp::InsertIfAbsent { .. } + | KvOp::Expire { .. } + | KvOp::Persist { .. } + | KvOp::BatchPut { .. } + | KvOp::RegisterIndex { .. } + | KvOp::DropIndex { .. } + | KvOp::FieldSet { .. } + | KvOp::Truncate { .. } + | KvOp::Incr { .. } + | KvOp::IncrFloat { .. } + | KvOp::Cas { .. } + | KvOp::GetSet { .. } + | KvOp::Transfer { .. } + | KvOp::TransferItem { .. } + | KvOp::RegisterSortedIndex { .. } + | KvOp::DropSortedIndex { .. } + | KvOp::SortedIndexRank { .. } + | KvOp::SortedIndexTopK { .. } + | KvOp::SortedIndexRange { .. } + | KvOp::SortedIndexCount { .. } + | KvOp::SortedIndexScore { .. } + | KvOp::SortedIndexTxnRead { .. } + | KvOp::MaterializeScan { .. } + | KvOp::ResolvedWrite { .. } => {} + } + Ok(()) +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/mod.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/mod.rs new file mode 100644 index 000000000..48be61765 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: BUSL-1.1 + +mod columnar; +mod document; +mod kv; +mod query; +mod reads; +mod search; +mod walk; + +pub use walk::{ + RedactionReadCheck, refuse_redacted_expression_reads, refuse_redacted_expression_reads_in_tasks, +}; diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/query.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/query.rs new file mode 100644 index 000000000..9a4688bc4 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/query.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Column reads and child plans of the relational operations. + +use nodedb_physical::physical_plan::{JoinProjection, PhysicalPlan, QueryOp}; + +use super::reads::NodeReads; + +/// The columns `op` reads outside its plain projection. The aggregate +/// functions themselves are refused by the aggregate check. +pub(super) fn reads(op: &QueryOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + QueryOp::ProviderScan { + filters, + projection, + computed_columns, + window_functions, + sort_keys, + distinct, + .. + } + | QueryOp::PostProcess { + filters, + projection, + computed_columns, + window_functions, + sort_keys, + distinct, + .. + } => { + out.filter_bytes(filters)?; + out.computed_bytes(computed_columns)?; + out.window_bytes(window_functions)?; + out.sort_keys(sort_keys); + out.distinct(*distinct, projection); + out.passthrough(projection); + } + QueryOp::Aggregate { + group_by, + filters, + sub_group_by, + sort_keys, + .. + } => { + out.group_keys(group_by); + out.group_names(sub_group_by); + out.filter_bytes(filters)?; + out.sort_keys(sort_keys); + } + QueryOp::PartialAggregate { + group_by, filters, .. + } => { + out.group_names(group_by); + out.filter_bytes(filters)?; + } + QueryOp::PartialAggregateState { + group_by, filters, .. + } => { + out.group_keys(group_by); + out.filter_bytes(filters)?; + } + QueryOp::ShuffleAggregateConsume { + group_by, + sort_keys, + .. + } => { + out.group_names(group_by); + out.sort_keys(sort_keys); + } + QueryOp::HashJoin { + on, + projection, + computed_projection, + join_filters, + post_filters, + left_scan_filters, + right_scan_filters, + .. + } => { + out.join_keys(on); + join_renames(projection, out); + out.computed_bytes(computed_projection)?; + for filters in [ + join_filters, + post_filters, + left_scan_filters, + right_scan_filters, + ] { + out.filter_bytes(filters)?; + } + } + QueryOp::ShuffleJoinConsume { on, .. } | QueryOp::SortMergeJoin { on, .. } => { + out.join_keys(on); + } + QueryOp::NestedLoopJoin { condition, .. } => out.filter_bytes(condition)?, + QueryOp::FacetCounts { filters, .. } => out.filter_bytes(filters)?, + QueryOp::RecursiveScan { + base_filters, + recursive_filters, + join_link, + .. + } => { + out.filter_bytes(base_filters)?; + out.filter_bytes(recursive_filters)?; + if let Some(link) = join_link { + out.join_keys(std::slice::from_ref(link)); + } + } + QueryOp::LateralTopK { + inner_filters, + inner_order_by, + correlation_keys, + projection, + .. + } => { + out.filter_bytes(inner_filters)?; + out.sort_keys(inner_order_by); + out.join_keys(correlation_keys); + join_renames(projection, out); + } + QueryOp::LateralLoop { + inner_filters, + correlation_predicates, + projection, + .. + } => { + out.filter_bytes(inner_filters)?; + out.join_keys(correlation_predicates); + join_renames(projection, out); + } + // A set operation compares whole rows of its inputs, each checked on + // its own. A gather moves rows. A value CTE reads no collection. + QueryOp::SetOp { .. } | QueryOp::Exchange(_) | QueryOp::RecursiveValue { .. } => {} + } + Ok(()) +} + +/// A join projection that emits a side's column under another name. +fn join_renames(projection: &[JoinProjection], out: &mut NodeReads) { + for spec in projection { + if spec.source != spec.output { + out.renames.push((spec.output.clone(), spec.source.clone())); + } + } +} + +/// The child plans `op` reads its rows from. +pub(super) fn children(op: &QueryOp) -> Vec<&PhysicalPlan> { + match op { + QueryOp::Exchange(exchange) => vec![&*exchange.child], + QueryOp::PostProcess { input, .. } => vec![&**input], + QueryOp::SetOp { inputs, .. } => inputs.iter().collect(), + QueryOp::Aggregate { input, .. } | QueryOp::PartialAggregateState { input, .. } => { + input.iter().map(|child| &**child).collect() + } + QueryOp::HashJoin { + left_input, + right_input, + left_bitmap, + right_bitmap, + .. + } => [left_input, right_input, left_bitmap, right_bitmap] + .into_iter() + .flatten() + .map(|child| &**child) + .collect(), + QueryOp::NestedLoopJoin { + left_input, + right_input, + .. + } + | QueryOp::SortMergeJoin { + left_input, + right_input, + .. + } => [left_input, right_input] + .into_iter() + .flatten() + .map(|child| &**child) + .collect(), + QueryOp::LateralTopK { outer_plan, .. } | QueryOp::LateralLoop { outer_plan, .. } => { + vec![&**outer_plan] + } + QueryOp::ProviderScan { .. } + | QueryOp::PartialAggregate { .. } + | QueryOp::ShuffleJoinConsume { .. } + | QueryOp::ShuffleAggregateConsume { .. } + | QueryOp::FacetCounts { .. } + | QueryOp::RecursiveScan { .. } + | QueryOp::RecursiveValue { .. } => Vec::new(), + } +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/reads.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/reads.rs new file mode 100644 index 000000000..4f91b4823 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/reads.rs @@ -0,0 +1,244 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The stored columns one plan node reads outside its plain projection. +//! +//! A predicate, a sort key, a group key, a join key, an index probe, an +//! assignment, and a computed expression all consume the stored value on the +//! Data Plane. The result-path masking hook never sees that value, so each +//! such read is collected here and checked against the redaction policy. + +use nodedb_physical::physical_plan::{GroupKeySpec, SortKeySpec, UpdateValue}; +use nodedb_query::expr::{ComputedColumn, SqlExpr}; +use nodedb_query::scan_filter::ScanFilter; +use nodedb_query::window::WindowFuncSpec; + +use super::super::aggregate::collect_columns; + +/// The clause a column read comes from, named in the refusal. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum Clause { + Filter, + Computed, + Window, + Sort, + Group, + JoinKey, + IndexProbe, + Assignment, + Projection, + Distinct, + KeyLookup, +} + +impl Clause { + /// The phrase the refusal uses for this clause. + pub(super) fn describe(self) -> &'static str { + match self { + Self::Filter => "a predicate over a redacted column", + Self::Computed => "an expression over a redacted column", + Self::Window => "a window function over a redacted column", + Self::Sort => "ordering by a redacted column", + Self::Group => "grouping by a redacted column", + Self::JoinKey => "joining on a redacted column", + Self::IndexProbe => "an index probe on a redacted column", + Self::Assignment => "an assignment that reads a redacted column", + Self::Projection => "projecting a renamed redacted column through a subquery", + Self::Distinct => "DISTINCT over a redacted column", + Self::KeyLookup => "a lookup by a redacted key column, whose hit or miss shows the key", + } + } +} + +/// Every column one node reads, and every alias it gives a stored column. +#[derive(Debug, Default)] +pub(super) struct NodeReads { + /// `(field, clause)` for each read outside the plain projection. + pub(super) reads: Vec<(String, Clause)>, + /// `(alias, field)` for each computed column that renames a stored + /// column. The output schema masks the stored column, so the rename is + /// not refused here. A parent node that reads the alias reads the + /// stored value. + pub(super) renames: Vec<(String, String)>, + /// Row keys a node passes through to its output. Only a renamed + /// redacted column from a child is refused among them. + pub(super) passthrough: Vec, + /// `Some` when the node reads every stored column in this clause, as a + /// `DISTINCT` over a whole row does. + pub(super) every_column: Option, + /// `true` when the node finds its row by primary-key value. A hit or a + /// miss tells the caller whether a row with that key exists. + pub(super) key_lookup: bool, +} + +impl NodeReads { + pub(super) fn is_empty(&self) -> bool { + self.reads.is_empty() + && self.renames.is_empty() + && self.passthrough.is_empty() + && self.every_column.is_none() + && !self.key_lookup + } + + /// A `DISTINCT` compares the projected stored values. An empty + /// projection compares every column. + pub(super) fn distinct(&mut self, distinct: bool, projection: &[String]) { + if !distinct { + return; + } + if projection.is_empty() { + self.every_column = Some(Clause::Distinct); + } + for name in projection { + self.field(name, Clause::Distinct); + } + } + + /// One field read in `clause`. An index path (`$.email`) names the + /// stored column `email`. + pub(super) fn field(&mut self, field: &str, clause: Clause) { + let field = field.strip_prefix("$.").unwrap_or(field); + if !field.is_empty() && field != "*" { + self.reads.push((field.to_string(), clause)); + } + } + + /// Every column `expr` reads, in `clause`. + pub(super) fn expr(&mut self, expr: &SqlExpr, clause: Clause) { + let mut columns = Vec::new(); + collect_columns(expr, &mut columns); + for column in columns { + self.field(column, clause); + } + } + + /// Every column a serialized `Vec` reads. + pub(super) fn filter_bytes(&mut self, bytes: &[u8]) -> crate::Result<()> { + if bytes.is_empty() { + return Ok(()); + } + let filters: Vec = decode(bytes, "scan filters")?; + self.filters(&filters); + Ok(()) + } + + fn filters(&mut self, filters: &[ScanFilter]) { + for filter in filters { + self.field(&filter.field, Clause::Filter); + for group in &filter.clauses { + self.filters(group); + } + if let Some(expr) = &filter.expr { + self.expr(expr, Clause::Filter); + } + } + } + + /// Every column a serialized `Vec` reads. + /// + /// A column that only names a stored column passes it through under its + /// alias. It is a rename when the alias differs, and no read otherwise. + pub(super) fn computed_bytes(&mut self, bytes: &[u8]) -> crate::Result<()> { + if bytes.is_empty() { + return Ok(()); + } + let columns: Vec = decode(bytes, "computed columns")?; + for column in &columns { + match &column.expr { + SqlExpr::Column(name) if *name == column.alias => {} + SqlExpr::Column(name) => self.renames.push((column.alias.clone(), name.clone())), + expr => self.expr(expr, Clause::Computed), + } + } + Ok(()) + } + + /// Every column a serialized `Vec` reads. + pub(super) fn window_bytes(&mut self, bytes: &[u8]) -> crate::Result<()> { + if bytes.is_empty() { + return Ok(()); + } + let specs: Vec = decode(bytes, "window functions")?; + for spec in &specs { + let order = spec.order_by.iter().map(|(expr, _)| expr); + for expr in spec.args.iter().chain(&spec.partition_by).chain(order) { + self.expr(expr, Clause::Window); + } + } + Ok(()) + } + + pub(super) fn sort_keys(&mut self, keys: &[SortKeySpec]) { + for key in keys { + self.expr(&key.expr, Clause::Sort); + } + } + + pub(super) fn group_keys(&mut self, keys: &[GroupKeySpec]) { + for key in keys { + if let Some(field) = &key.field { + self.field(field, Clause::Group); + } + if let Some(expr) = &key.expr { + self.expr(expr, Clause::Group); + } + } + } + + pub(super) fn group_names(&mut self, names: &[String]) { + for name in names { + self.field(name, Clause::Group); + } + } + + /// Every column an assignment's right-hand side reads. + pub(super) fn updates(&mut self, updates: &[(String, UpdateValue)]) { + self.updates_of_target(updates, None); + } + + /// Every target column an assignment's right-hand side reads. A column + /// qualified with `source_alias` is the source row's, and the source + /// relation is checked on its own. + pub(super) fn updates_of_target( + &mut self, + updates: &[(String, UpdateValue)], + source_alias: Option<&str>, + ) { + for (_, value) in updates { + let UpdateValue::Expr(expr) = value else { + continue; + }; + let mut columns = Vec::new(); + collect_columns(expr, &mut columns); + for column in columns { + let from_source = source_alias.is_some_and(|alias| { + column + .split_once('.') + .is_some_and(|(qualifier, _)| qualifier == alias) + }); + if !from_source { + self.field(column, Clause::Assignment); + } + } + } + } + + /// Both columns of each `(left, right)` join key pair. + pub(super) fn join_keys(&mut self, keys: &[(String, String)]) { + for (left, right) in keys { + self.field(left, Clause::JoinKey); + self.field(right, Clause::JoinKey); + } + } + + pub(super) fn passthrough(&mut self, names: &[String]) { + self.passthrough.extend(names.iter().cloned()); + } +} + +/// Decode a serialized plan payload. A payload that does not decode fails +/// the statement: the refusal cannot vouch for reads it cannot see. +fn decode<'a, T: zerompk::FromMessagePack<'a>>(bytes: &'a [u8], what: &str) -> crate::Result { + zerompk::from_msgpack(bytes).map_err(|e| crate::Error::Internal { + detail: format!("redaction check: {what} do not decode: {e}"), + }) +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/search.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/search.rs new file mode 100644 index 000000000..35d24e997 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/search.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Column reads of the text and vector search operations. +//! +//! A text match over a redacted field is refused by the text check. The +//! predicates a search carries beside its match are read here, and so is the +//! vector column a similarity search ranks by. + +use nodedb_physical::physical_plan::{PhysicalPlan, TextOp, VectorOp, VectorWriteTargets}; + +use super::reads::{Clause, NodeReads}; + +pub(super) fn text_reads(op: &TextOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + TextOp::Search { filters, .. } + | TextOp::BM25ScoreScan { filters, .. } + | TextOp::PhraseSearch { filters, .. } => out.filter_bytes(filters)?, + TextOp::HybridSearch { + filters, + vector_field, + .. + } + | TextOp::HybridSearchTriple { + filters, + vector_field, + .. + } => { + out.filter_bytes(filters)?; + out.field(vector_field, Clause::Sort); + } + TextOp::FtsIndexDoc { .. } | TextOp::FtsDeleteDoc { .. } | TextOp::SetTextConfig { .. } => { + } + } + Ok(()) +} + +pub(super) fn vector_reads(op: &VectorOp, out: &mut NodeReads) -> crate::Result<()> { + match op { + // The statement's `WHERE` conjuncts ride in `rls_filters`. The check + // runs before the read policy joins them there. + VectorOp::Search { + field_name, + payload_filters, + rls_filters, + .. + } => { + out.filter_bytes(rls_filters)?; + out.field(field_name, Clause::Sort); + for atom in payload_filters { + out.field(atom.field(), Clause::Filter); + } + } + VectorOp::MultiSearch { rls_filters, .. } => out.filter_bytes(rls_filters)?, + // The sparse search's `WHERE` conjuncts ride in `filters`. + VectorOp::SparseSearch { + field_name, + filters, + .. + } => { + out.filter_bytes(filters)?; + out.field(field_name, Clause::Sort); + } + VectorOp::MultiVectorScoreSearch { field_name, .. } => { + out.field(field_name, Clause::Sort); + } + VectorOp::DirectDelete { targets, .. } => target_reads(targets, out)?, + VectorOp::DirectUpdate { + targets, + payload_patch, + .. + } => { + target_reads(targets, out)?; + out.updates(payload_patch); + } + VectorOp::ResolveDirectWrite(inner) => vector_reads(inner, out)?, + // Writes by surrogate, index DDL, and maintenance read no column in a + // predicate or an expression. + VectorOp::Insert { .. } + | VectorOp::BatchInsert { .. } + | VectorOp::Delete { .. } + | VectorOp::DeleteBySurrogate { .. } + | VectorOp::SetParams { .. } + | VectorOp::DropIndex { .. } + | VectorOp::QueryStats { .. } + | VectorOp::Seal { .. } + | VectorOp::CompactIndex { .. } + | VectorOp::Rebuild { .. } + | VectorOp::SparseInsert { .. } + | VectorOp::SparseDelete { .. } + | VectorOp::MultiVectorInsert { .. } + | VectorOp::MultiVectorDelete { .. } + | VectorOp::DirectUpsert { .. } + | VectorOp::DirectInsert { .. } + | VectorOp::DirectInsertIfAbsent { .. } + | VectorOp::DirectBatchWrite { .. } + | VectorOp::DirectTruncate { .. } + | VectorOp::ResolvedDirectWrite { .. } => {} + } + Ok(()) +} + +fn target_reads(targets: &VectorWriteTargets, out: &mut NodeReads) -> crate::Result<()> { + match targets { + VectorWriteTargets::Predicate(filters) => out.filter_bytes(filters), + VectorWriteTargets::Surrogates(_) => Ok(()), + } +} + +/// The resolved sub-plan a vector search takes its prefilter from, if any. +pub(super) fn vector_child(op: &VectorOp) -> Option<&PhysicalPlan> { + match op { + VectorOp::Search { + inline_prefilter_plan: Some(child), + .. + } => Some(child), + VectorOp::ResolveDirectWrite(inner) => vector_child(inner), + VectorOp::Search { .. } + | VectorOp::Insert { .. } + | VectorOp::BatchInsert { .. } + | VectorOp::MultiSearch { .. } + | VectorOp::Delete { .. } + | VectorOp::DeleteBySurrogate { .. } + | VectorOp::SetParams { .. } + | VectorOp::DropIndex { .. } + | VectorOp::QueryStats { .. } + | VectorOp::Seal { .. } + | VectorOp::CompactIndex { .. } + | VectorOp::Rebuild { .. } + | VectorOp::SparseInsert { .. } + | VectorOp::SparseSearch { .. } + | VectorOp::SparseDelete { .. } + | VectorOp::MultiVectorInsert { .. } + | VectorOp::MultiVectorDelete { .. } + | VectorOp::MultiVectorScoreSearch { .. } + | VectorOp::DirectUpsert { .. } + | VectorOp::DirectInsert { .. } + | VectorOp::DirectInsertIfAbsent { .. } + | VectorOp::DirectBatchWrite { .. } + | VectorOp::DirectDelete { .. } + | VectorOp::DirectTruncate { .. } + | VectorOp::DirectUpdate { .. } + | VectorOp::ResolvedDirectWrite { .. } => None, + } +} diff --git a/nodedb/src/control/planner/redaction_refusal/expr_reads/walk.rs b/nodedb/src/control/planner/redaction_refusal/expr_reads/walk.rs new file mode 100644 index 000000000..4204876f5 --- /dev/null +++ b/nodedb/src/control/planner/redaction_refusal/expr_reads/walk.rs @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Refusal of a plan that reads a redacted column outside its projection. +//! +//! Redaction masks the result rows in the Control Plane. A plain projection +//! of a redacted column is masked there. Every other read consumes the +//! stored value on the Data Plane before any mask applies: +//! +//! - an expression computes over it (`upper(ssn)`, `ssn || ''`), +//! - a predicate filters on it (`WHERE ssn = '123-45-6789'`), +//! - a sort orders by it, a group key or `DISTINCT` buckets by it, +//! - a join or an index probe matches on it, +//! - an assignment copies it into another column. +//! +//! Each leaks the stored value or a fact about it, so the statement is +//! refused with `42501`, naming the column. +//! +//! The check runs before row-level security is injected. A read policy's +//! own predicate can read a redacted column. It is the administrator's +//! predicate, not the requester's, and it must not trigger the refusal. + +use nodedb_physical::physical_plan::PhysicalPlan; +use nodedb_physical::physical_task::PhysicalTask; + +use crate::control::security::auth_context::AuthContext; +use crate::control::security::catalog::SystemCatalog; +use crate::control::security::redaction::RedactionStore; +use crate::control::server::response_shape::redaction::plan_source_collections; +use crate::types::TenantId; + +use super::super::lookup::RefusalCtx; +use super::reads::{Clause, NodeReads}; +use super::{columnar, document, kv, query, search}; + +/// The stores a refusal reads: the policy lookups, and the catalog that +/// names each collection's key column. +pub struct RedactionReadCheck<'a> { + pub auth: &'a AuthContext, + pub store: &'a RedactionStore, + pub catalog: &'a SystemCatalog, +} + +/// Refuse `plan` when it reads a column one of the requester's roles +/// redacts, outside a plain projection. +/// +/// Call it before `rls_injection` rewrites the plan's predicates. +pub fn refuse_redacted_expression_reads( + plan: &PhysicalPlan, + tenant_id: TenantId, + database_id: nodedb_types::DatabaseId, + check: &RedactionReadCheck<'_>, +) -> crate::Result<()> { + if check.auth.roles.is_empty() { + return Ok(()); + } + let ctx = RefusalCtx { + store: check.store, + tenant_id: tenant_id.as_u64(), + roles: &check.auth.roles, + database_id, + }; + if !ctx.identity_has_any_rule() { + return Ok(()); + } + walk(plan, &ctx, check.catalog).map(|_| ()) +} + +/// [`refuse_redacted_expression_reads`] over every task, each under its own +/// tenant and database. +pub fn refuse_redacted_expression_reads_in_tasks( + tasks: &[PhysicalTask], + check: &RedactionReadCheck<'_>, +) -> crate::Result<()> { + for task in tasks { + refuse_redacted_expression_reads(&task.plan, task.tenant_id, task.database_id, check)?; + } + Ok(()) +} + +/// Check `plan` and its children. Returns the output names under which the +/// node emits a redacted stored value: a renamed redacted column. +fn walk( + plan: &PhysicalPlan, + ctx: &RefusalCtx<'_>, + catalog: &SystemCatalog, +) -> crate::Result> { + let mut renamed = Vec::new(); + for child in children(plan) { + renamed.extend(walk(child, ctx, catalog)?); + } + let mut reads = NodeReads::default(); + node_reads(plan, &mut reads)?; + if reads.is_empty() { + return Ok(renamed); + } + let sources = plan_source_collections(plan, ctx.database_id); + if reads.key_lookup { + refuse_key_lookup(ctx, catalog, &sources)?; + } + let protected = |field: &str| -> Option { + if names_one_of(&renamed, field) { + return Some("a renamed redacted column".to_string()); + } + ctx.redacted_source(&sources, field) + .map(|collection| format!("'{collection}'")) + }; + + for (field, clause) in &reads.reads { + if let Some(origin) = protected(field) { + return Err(refusal(ctx, field, &origin, *clause)); + } + } + if let Some(clause) = reads.every_column { + let redacted = sources + .iter() + .find(|(_, collection)| ctx.collection_is_redacted(collection)); + if let Some((_, collection)) = redacted { + return Err(refusal(ctx, "*", &format!("'{collection}'"), clause)); + } + if let Some(name) = renamed.first() { + return Err(refusal(ctx, name, "a renamed redacted column", clause)); + } + } + for name in &reads.passthrough { + if names_one_of(&renamed, name) { + return Err(refusal( + ctx, + name, + "a renamed redacted column", + Clause::Projection, + )); + } + } + let mut out = renamed.clone(); + for (alias, field) in &reads.renames { + if protected(field).is_some() { + out.push(alias.clone()); + } + } + Ok(out) +} + +/// Whether `field`, as written or by its column part, is one of `names`. +fn names_one_of(names: &[String], field: &str) -> bool { + let column = field.rsplit_once('.').map_or(field, |(_, column)| column); + names.iter().any(|name| name == field || name == column) +} + +fn refusal(ctx: &RefusalCtx<'_>, field: &str, origin: &str, clause: Clause) -> crate::Error { + ctx.refusal(format!( + "column '{field}' on {origin} is redacted for this role: {} is not permitted, \ + because the stored value is read before the redaction applies", + clause.describe() + )) +} + +/// Refuse a lookup by key when the key column of a source is redacted. +/// +/// The key column is the collection's primary key, `id` when it declares +/// none. A collection the catalog does not hold keys by `id`. +fn refuse_key_lookup( + ctx: &RefusalCtx<'_>, + catalog: &SystemCatalog, + sources: &[(String, String)], +) -> crate::Result<()> { + for (_, stored) in sources { + if stored.is_empty() || !ctx.collection_is_redacted(stored) { + continue; + } + let qualified = nodedb_types::QualifiedCollection::from_stored(stored.clone()); + let key = catalog + .get_collection( + ctx.database_id, + ctx.tenant_id, + qualified.collection_name(ctx.database_id), + )? + .and_then(|collection| { + crate::control::planner::catalog_adapter::convert_collection_type(&collection).2 + }) + .unwrap_or_else(|| nodedb_types::DEFAULT_IDENTITY_COLUMN.to_string()); + if ctx.field_is_redacted(stored, &key) { + return Err(refusal( + ctx, + &key, + &format!("'{stored}'"), + Clause::KeyLookup, + )); + } + } + Ok(()) +} + +/// The plans whose rows `plan` reads. Exhaustive over [`PhysicalPlan`]. +fn children(plan: &PhysicalPlan) -> Vec<&PhysicalPlan> { + match plan { + PhysicalPlan::Query(op) => query::children(op), + PhysicalPlan::Document(op) => document::child(op).into_iter().collect(), + PhysicalPlan::Vector(op) => search::vector_child(op).into_iter().collect(), + PhysicalPlan::Graph(_) + | PhysicalPlan::Kv(_) + | PhysicalPlan::Text(_) + | PhysicalPlan::Columnar(_) + | PhysicalPlan::Timeseries(_) + | PhysicalPlan::Spatial(_) + | PhysicalPlan::Crdt(_) + | PhysicalPlan::Meta(_) + | PhysicalPlan::Array(_) + | PhysicalPlan::ClusterArray(_) + | PhysicalPlan::ClusterEvent(_) => Vec::new(), + } +} + +/// The columns `plan` itself reads. Exhaustive over [`PhysicalPlan`]. +/// +/// A graph traversal over a redacted collection is refused whole by the +/// graph check. CRDT, maintenance, and array operations carry no column +/// predicate or expression. +fn node_reads(plan: &PhysicalPlan, out: &mut NodeReads) -> crate::Result<()> { + match plan { + PhysicalPlan::Query(op) => query::reads(op, out), + PhysicalPlan::Document(op) => document::reads(op, out), + PhysicalPlan::Kv(op) => kv::reads(op, out), + PhysicalPlan::Columnar(op) => columnar::columnar_reads(op, out), + PhysicalPlan::Timeseries(op) => columnar::timeseries_reads(op, out), + PhysicalPlan::Spatial(op) => columnar::spatial_reads(op, out), + PhysicalPlan::Text(op) => search::text_reads(op, out), + PhysicalPlan::Vector(op) => search::vector_reads(op, out), + PhysicalPlan::Graph(_) + | PhysicalPlan::Crdt(_) + | PhysicalPlan::Meta(_) + | PhysicalPlan::Array(_) + | PhysicalPlan::ClusterArray(_) + | PhysicalPlan::ClusterEvent(_) => Ok(()), + } +} diff --git a/nodedb/src/control/planner/redaction_refusal/graph.rs b/nodedb/src/control/planner/redaction_refusal/graph.rs index f3daa108a..b41005164 100644 --- a/nodedb/src/control/planner/redaction_refusal/graph.rs +++ b/nodedb/src/control/planner/redaction_refusal/graph.rs @@ -66,12 +66,10 @@ pub(super) fn refuse_traversal( if collection.as_str().is_empty() || !ctx.collection_is_redacted(collection.as_str()) { return Ok(()); } - Err(crate::Error::PlanError { - detail: format!( - "redaction policies on '{collection}' are not supported with graph traversal: a \ - traversal returns graph topology, which column redaction cannot be applied to" - ), - }) + Err(ctx.refusal(format!( + "redaction policies on '{collection}' are not supported with graph traversal: a \ + traversal returns graph topology, which column redaction cannot be applied to" + ))) } /// Refuse a pattern match whose target collection carries a rule. The @@ -110,11 +108,11 @@ fn refuse_unscoped_match(ctx: &RefusalCtx<'_>) -> crate::Result<()> { if !ctx.identity_has_any_rule() { return Ok(()); } - Err(crate::Error::PlanError { - detail: "graph pattern matching is not supported while a redaction policy applies to \ - this role and the pattern names no collection: the match returns graph \ - topology, which column redaction cannot be applied to, and its scope cannot be \ - narrowed to an unprotected collection" + Err(ctx.refusal( + "graph pattern matching is not supported while a redaction policy applies to \ + this role and the pattern names no collection: the match returns graph \ + topology, which column redaction cannot be applied to, and its scope cannot be \ + narrowed to an unprotected collection" .to_string(), - }) + )) } diff --git a/nodedb/src/control/planner/redaction_refusal/lookup.rs b/nodedb/src/control/planner/redaction_refusal/lookup.rs index 73cbd2eda..851dd4ad1 100644 --- a/nodedb/src/control/planner/redaction_refusal/lookup.rs +++ b/nodedb/src/control/planner/redaction_refusal/lookup.rs @@ -3,6 +3,16 @@ //! Identity-scoped rule lookups shared by the aggregate and graph refusals. use crate::control::security::redaction::RedactionStore; +use crate::types::TenantId; + +/// The error every redaction refusal raises: `42501`, with `detail` naming +/// the redacted column or collection and the clause that reads it. +pub(super) fn redaction_refusal(tenant_id: u64, detail: String) -> crate::Error { + crate::Error::RejectedAuthz { + tenant_id: TenantId::new(tenant_id), + resource: detail, + } +} /// The inputs every refusal check keys on: the policy registry plus the /// requester's tenant and roles. @@ -22,6 +32,11 @@ pub(super) struct RefusalCtx<'a> { } impl RefusalCtx<'_> { + /// [`redaction_refusal`] for this requester's tenant. + pub(super) fn refusal(&self, detail: String) -> crate::Error { + redaction_refusal(self.tenant_id, detail) + } + /// True when one of the requester's roles has a redaction rule covering /// `collection`.`field`. /// @@ -41,6 +56,37 @@ impl RefusalCtx<'_> { }) } + /// The collection of `sources` whose redacted column `field` names, if + /// any. + /// + /// A field qualified with a source's qualifier (`o.ssn`) is checked + /// against the sources that carry that qualifier. Any other field, + /// unqualified or with a qualifier no source carries (a derived table's + /// alias, a nested field path), is checked against every source, so an + /// ambiguous name cannot slip past. [`Self::field_is_redacted`] matches + /// the field as written and by its column part. + pub(super) fn redacted_source<'s>( + &self, + sources: &'s [(String, String)], + field: &str, + ) -> Option<&'s str> { + let qualifier = field.split_once('.').map(|(qualifier, _)| qualifier); + let names_a_source = qualifier.is_some_and(|qualifier| { + sources + .iter() + .any(|(source_qualifier, _)| source_qualifier == qualifier) + }); + sources + .iter() + .filter(|(source_qualifier, _)| { + !names_a_source || qualifier == Some(source_qualifier.as_str()) + }) + .find(|(_, collection)| { + !collection.is_empty() && self.field_is_redacted(collection, field) + }) + .map(|(_, collection)| collection.as_str()) + } + /// True when one of the requester's roles has any redaction rule on /// `collection`. pub(super) fn collection_is_redacted(&self, collection: &str) -> bool { diff --git a/nodedb/src/control/planner/redaction_refusal/mod.rs b/nodedb/src/control/planner/redaction_refusal/mod.rs index 8246354cb..0c1692840 100644 --- a/nodedb/src/control/planner/redaction_refusal/mod.rs +++ b/nodedb/src/control/planner/redaction_refusal/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: BUSL-1.1 mod aggregate; +mod expr_reads; mod graph; mod lookup; mod orchestrated; @@ -8,6 +9,9 @@ mod plan; mod streaming_mv; mod text; +pub use expr_reads::{ + RedactionReadCheck, refuse_redacted_expression_reads, refuse_redacted_expression_reads_in_tasks, +}; pub use plan::{ refuse_unredactable_graph_collection, refuse_unredactable_graph_match, refuse_unredactable_graph_match_scoped, refuse_unredactable_plan, refuse_unredactable_tasks, diff --git a/nodedb/src/control/planner/redaction_refusal/orchestrated.rs b/nodedb/src/control/planner/redaction_refusal/orchestrated.rs index f4c934056..12b03478c 100644 --- a/nodedb/src/control/planner/redaction_refusal/orchestrated.rs +++ b/nodedb/src/control/planner/redaction_refusal/orchestrated.rs @@ -74,12 +74,10 @@ pub(super) fn refuse_redacted_source( .iter() .find(|collection| !collection.is_empty() && ctx.collection_is_redacted(collection)) { - Some(collection) => Err(crate::Error::PlanError { - detail: format!( - "redaction policies on '{collection}' are not supported as the source of \ - {statement}: the copied values bypass column redaction" - ), - }), + Some(collection) => Err(ctx.refusal(format!( + "redaction policies on '{collection}' are not supported as the source of \ + {statement}: the copied values bypass column redaction" + ))), None => Ok(()), } } @@ -144,8 +142,8 @@ mod tests { "users", ))); match check(&store, &op) { - Err(crate::Error::PlanError { detail }) => { - assert!(detail.contains("'users'") && detail.contains("INSERT ... SELECT")); + Err(crate::Error::RejectedAuthz { resource, .. }) => { + assert!(resource.contains("'users'") && resource.contains("INSERT ... SELECT")); } other => panic!("expected a redaction refusal, got {other:?}"), } @@ -165,7 +163,7 @@ mod tests { })); assert!(matches!( check(&store, &op), - Err(crate::Error::PlanError { .. }) + Err(crate::Error::RejectedAuthz { .. }) )); } @@ -194,7 +192,7 @@ mod tests { })); assert!(matches!( check(&store, &op), - Err(crate::Error::PlanError { .. }) + Err(crate::Error::RejectedAuthz { .. }) )); } diff --git a/nodedb/src/control/planner/redaction_refusal/plan.rs b/nodedb/src/control/planner/redaction_refusal/plan.rs index 896cdefdc..a6372fbb2 100644 --- a/nodedb/src/control/planner/redaction_refusal/plan.rs +++ b/nodedb/src/control/planner/redaction_refusal/plan.rs @@ -35,8 +35,7 @@ use super::text::refuse_text_op; /// Refuse `plan` when it reads something a redaction policy protects but the /// result-path masking hook cannot rewrite. /// -/// Returns `Err(crate::Error::PlanError)` — the same typed refusal RLS raises -/// for the operation types it cannot filter. +/// Returns `Err(crate::Error::RejectedAuthz)`, which clients see as `42501`. pub fn refuse_unredactable_plan( plan: &PhysicalPlan, tenant_id: TenantId, @@ -393,6 +392,7 @@ mod tests { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }) } @@ -408,8 +408,8 @@ mod tests { fn assert_refused(result: crate::Result<()>) { match result { - Err(crate::Error::PlanError { .. }) => {} - other => panic!("expected PlanError refusal, got {other:?}"), + Err(crate::Error::RejectedAuthz { .. }) => {} + other => panic!("expected RejectedAuthz refusal, got {other:?}"), } } @@ -546,6 +546,7 @@ mod tests { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }) } @@ -690,6 +691,7 @@ mod tests { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }) } diff --git a/nodedb/src/control/planner/redaction_refusal/streaming_mv.rs b/nodedb/src/control/planner/redaction_refusal/streaming_mv.rs index 4bea32739..66c312af5 100644 --- a/nodedb/src/control/planner/redaction_refusal/streaming_mv.rs +++ b/nodedb/src/control/planner/redaction_refusal/streaming_mv.rs @@ -39,8 +39,8 @@ use crate::event::streaming_mv::types::AggDef; /// in the tenant, so the column is matched tenant-wide rather than passing /// unchecked. /// -/// Returns `Err(crate::Error::PlanError)` — the same typed refusal the planner -/// raises for an aggregate over a redacted column. +/// Returns `Err(crate::Error::RejectedAuthz)`, the refusal every redaction +/// check raises, which clients see as `42501`. pub fn refuse_redacted_streaming_mv( store: &RedactionStore, tenant_id: u64, @@ -56,7 +56,7 @@ pub fn refuse_redacted_streaming_mv( // a protected column is not to persist it. for column in group_by_columns { if store.has_rule_for_field_any_role(tenant_id, source_collection, column) { - return Err(refusal(source_collection, column, "grouping by")); + return Err(refusal(tenant_id, source_collection, column, "grouping by")); } } @@ -68,25 +68,26 @@ pub fn refuse_redacted_streaming_mv( continue; }; if store.has_rule_for_field_any_role(tenant_id, source_collection, column) { - return Err(refusal(source_collection, column, "aggregating")); + return Err(refusal(tenant_id, source_collection, column, "aggregating")); } } Ok(()) } -fn refusal(collection: Option<&str>, column: &str, action: &str) -> crate::Error { +fn refusal(tenant_id: u64, collection: Option<&str>, column: &str, action: &str) -> crate::Error { let scope = match collection { Some(name) => format!("on '{name}'"), None => "on a collection of this tenant (the source stream is a wildcard)".to_string(), }; - crate::Error::PlanError { - detail: format!( + super::lookup::redaction_refusal( + tenant_id, + format!( "column '{column}' {scope} is redacted: {action} it in a streaming materialized view \ is not permitted — the view persists what it derives from each event, so the \ protected value would be stored in the clear in the view's own state" ), - } + ) } #[cfg(test)] @@ -134,8 +135,8 @@ mod tests { fn assert_refused(result: crate::Result<()>) { match result { - Err(crate::Error::PlanError { .. }) => {} - other => panic!("expected PlanError refusal, got {other:?}"), + Err(crate::Error::RejectedAuthz { .. }) => {} + other => panic!("expected RejectedAuthz refusal, got {other:?}"), } } diff --git a/nodedb/src/control/planner/redaction_refusal/text.rs b/nodedb/src/control/planner/redaction_refusal/text.rs index 0bbf38f0b..356936e42 100644 --- a/nodedb/src/control/planner/redaction_refusal/text.rs +++ b/nodedb/src/control/planner/redaction_refusal/text.rs @@ -76,21 +76,17 @@ fn refuse_text_field( return Ok(()); } match field { - Some(field) if ctx.field_is_redacted(collection, field) => Err(crate::Error::PlanError { - detail: format!( - "column '{field}' on '{collection}' is redacted for this role: {reader} over a \ - redacted column is not permitted — the match and score are computed over the \ - stored text, so masking the result would not protect it" - ), - }), + Some(field) if ctx.field_is_redacted(collection, field) => Err(ctx.refusal(format!( + "column '{field}' on '{collection}' is redacted for this role: {reader} over a \ + redacted column is not permitted — the match and score are computed over the \ + stored text, so masking the result would not protect it" + ))), Some(_) => Ok(()), - None if ctx.collection_is_redacted(collection) => Err(crate::Error::PlanError { - detail: format!( - "'{collection}' has redacted columns for this role: {reader} over the whole \ - document is not permitted — it reads every text column, including the \ - redacted ones" - ), - }), + None if ctx.collection_is_redacted(collection) => Err(ctx.refusal(format!( + "'{collection}' has redacted columns for this role: {reader} over the whole \ + document is not permitted — it reads every text column, including the \ + redacted ones" + ))), None => Ok(()), } } @@ -184,8 +180,8 @@ mod tests { fn assert_refused(result: crate::Result<()>) { assert!( - matches!(result, Err(crate::Error::PlanError { .. })), - "expected a PlanError refusal, got {result:?}" + matches!(result, Err(crate::Error::RejectedAuthz { .. })), + "expected a RejectedAuthz refusal, got {result:?}" ); } diff --git a/nodedb/src/control/planner/rls_injection/array.rs b/nodedb/src/control/planner/rls_injection/array.rs index 0f2887565..b72d015f8 100644 --- a/nodedb/src/control/planner/rls_injection/array.rs +++ b/nodedb/src/control/planner/rls_injection/array.rs @@ -20,7 +20,6 @@ pub(super) fn inject_array(_ctx: &RlsCtx<'_>, op: &ArrayOp) -> crate::Result<()> // No-op: cell reads, addressed by `ArrayId`, which no policy can be // keyed on. ArrayOp::Slice { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } | ArrayOp::SurrogateBitmapScan { .. } => Ok(()), @@ -47,8 +46,11 @@ pub(super) fn inject_cluster_array(_ctx: &RlsCtx<'_>, op: &ClusterArrayOp) -> cr // addressed by the same `ArrayId`. ClusterArrayOp::Slice { .. } | ClusterArrayOp::Agg { .. } + | ClusterArrayOp::Elementwise { .. } | ClusterArrayOp::Put { .. } - | ClusterArrayOp::Delete { .. } => Ok(()), + | ClusterArrayOp::Delete { .. } + | ClusterArrayOp::Flush { .. } + | ClusterArrayOp::Compact { .. } => Ok(()), } } diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/array.rs b/nodedb/src/control/planner/rls_injection/permission_tree/array.rs index 92582515c..77c43366b 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/array.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/array.rs @@ -21,7 +21,6 @@ pub(super) fn apply_array(_ctx: &PermCtx<'_>, op: &ArrayOp) -> crate::Result<()> // No-op: cell reads, addressed by `ArrayId`, which no tree definition // can be keyed on. ArrayOp::Slice { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } | ArrayOp::SurrogateBitmapScan { .. } => Ok(()), @@ -45,8 +44,11 @@ pub(super) fn apply_cluster_array(_ctx: &PermCtx<'_>, op: &ClusterArrayOp) -> cr // addressed by the same `ArrayId`. ClusterArrayOp::Slice { .. } | ClusterArrayOp::Agg { .. } + | ClusterArrayOp::Elementwise { .. } | ClusterArrayOp::Put { .. } - | ClusterArrayOp::Delete { .. } => Ok(()), + | ClusterArrayOp::Delete { .. } + | ClusterArrayOp::Flush { .. } + | ClusterArrayOp::Compact { .. } => Ok(()), } } diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/columnar.rs b/nodedb/src/control/planner/rls_injection/permission_tree/columnar.rs index 2c0bfe727..b13c7fc17 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/columnar.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/columnar.rs @@ -168,6 +168,7 @@ mod tests { valid_at_ms: None, having: Vec::new(), tiering: None, + output_columns: Vec::new(), }); assert!(apply(&mut plan, &cache).is_ok()); match &plan { diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/query.rs b/nodedb/src/control/planner/rls_injection/permission_tree/query.rs index 9b3550819..e5ca55b51 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/query.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/query.rs @@ -273,6 +273,7 @@ mod tests { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }); assert!(apply(&mut plan, &cache).is_ok()); match &plan { diff --git a/nodedb/src/control/planner/rls_injection/permission_tree/vector.rs b/nodedb/src/control/planner/rls_injection/permission_tree/vector.rs index b63b80335..5310da082 100644 --- a/nodedb/src/control/planner/rls_injection/permission_tree/vector.rs +++ b/nodedb/src/control/planner/rls_injection/permission_tree/vector.rs @@ -35,11 +35,17 @@ pub(super) fn apply_vector(ctx: &PermCtx<'_>, op: &mut VectorOp) -> crate::Resul .. } => ctx.filter_into(collection, PermTreeLevel::Read, rls_filters), - // Refuse: both return scored document identities with no filter slot, - // so rows outside the caller's subtree would still be ranked and - // returned. - VectorOp::SparseSearch { collection, .. } - | VectorOp::MultiVectorScoreSearch { collection, .. } => ctx.refuse_if_tree( + // Filter: the subtree joins the sparse search's filter, which narrows + // the candidates before the top-k cut. + VectorOp::SparseSearch { + collection, + filters, + .. + } => ctx.filter_into(collection, PermTreeLevel::Read, filters), + + // Refuse: the search returns scored document identities with no + // filter slot, so rows outside the caller's subtree will be ranked. + VectorOp::MultiVectorScoreSearch { collection, .. } => ctx.refuse_if_tree( collection, "the search returns scored document identities through a response shape that carries \ no subtree filter", @@ -174,9 +180,10 @@ mod tests { assert_refused(apply(&mut plan, &cache), "users"); } - /// A sparse search has no filter slot, so a tree refuses it. + /// A sparse search over a governed collection is narrowed to the + /// readable subtree through its filter. #[test] - fn sparse_search_is_refused_under_a_tree() { + fn sparse_search_is_narrowed_to_the_readable_subtree() { let cache = cache_with_tree("docs"); let mut plan = PhysicalPlan::Vector(VectorOp::SparseSearch { collection: nodedb_types::QualifiedCollection::new( @@ -186,7 +193,14 @@ mod tests { field_name: "sparse".into(), query_entries: vec![(1, 1.0)], top_k: 5, + filters: Vec::new(), }); - assert_refused(apply(&mut plan, &cache), "docs"); + assert!(apply(&mut plan, &cache).is_ok()); + match &plan { + PhysicalPlan::Vector(VectorOp::SparseSearch { filters, .. }) => { + assert_eq!(sorted(injected_resources(filters)), readable()); + } + other => panic!("plan shape changed: {other:?}"), + } } } diff --git a/nodedb/src/control/planner/rls_injection/vector.rs b/nodedb/src/control/planner/rls_injection/vector.rs index aa7caca9b..ff39a7009 100644 --- a/nodedb/src/control/planner/rls_injection/vector.rs +++ b/nodedb/src/control/planner/rls_injection/vector.rs @@ -35,10 +35,17 @@ pub(super) fn inject_vector(ctx: &RlsCtx<'_>, op: &mut VectorOp) -> crate::Resul .. } => ctx.set_post_filters(collection, rls_filters), - // Refuse: both return scored document identities with no filter slot, - // so the rows a policy hides will still be ranked and returned. - VectorOp::SparseSearch { collection, .. } - | VectorOp::MultiVectorScoreSearch { collection, .. } => ctx.refuse_if_policy( + // Inject: the sparse hits carry their bodies while a filter applies, + // and the policy joins that filter before the top-k cut. + VectorOp::SparseSearch { + collection, + filters, + .. + } => ctx.set_post_filters(collection, filters), + + // Refuse: the search returns scored document identities with no + // filter slot, so the rows a policy hides will still be ranked. + VectorOp::MultiVectorScoreSearch { collection, .. } => ctx.refuse_if_policy( collection, "the search returns scored document identities through a response shape that carries \ no row filter", @@ -395,19 +402,43 @@ mod tests { assert_eq!(plan, before); } - /// A sparse search has no filter slot, so a policy refuses it. - #[test] - fn sparse_search_is_refused_under_a_read_policy() { - let store = store_with_read_policy("docs"); - let mut plan = PhysicalPlan::Vector(VectorOp::SparseSearch { + fn sparse_search(collection: &str) -> PhysicalPlan { + PhysicalPlan::Vector(VectorOp::SparseSearch { collection: nodedb_types::QualifiedCollection::new( nodedb_types::DatabaseId::DEFAULT, - "docs", + collection, ), field_name: "sparse".into(), query_entries: vec![(1, 1.0)], top_k: 5, - }); - assert_refused(inject(&mut plan, &store), "docs"); + filters: Vec::new(), + }) + } + + /// A read policy joins the sparse search's filter, which narrows the + /// candidates before the top-k cut. + #[test] + fn sparse_search_carries_the_read_policy_in_its_filter() { + let store = store_with_read_policy("docs"); + let mut plan = sparse_search("docs"); + assert!(inject(&mut plan, &store).is_ok()); + let PhysicalPlan::Vector(VectorOp::SparseSearch { filters, .. }) = &plan else { + panic!("expected SparseSearch, got {plan:?}"); + }; + let filters: Vec = + zerompk::from_msgpack(filters).expect("decode injected filters"); + assert!( + filters.iter().any(|f| f.field == "owner_id"), + "the policy predicate must join the filter: {filters:?}" + ); + } + + /// The same search is untouched when no policy applies. + #[test] + fn sparse_search_without_a_policy_is_untouched() { + let mut plan = sparse_search("docs"); + let before = plan.clone(); + assert!(inject_without_policy(&mut plan).is_ok()); + assert_eq!(plan, before); } } diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs index 5018749eb..2920a31cf 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/input_sourced.rs @@ -25,6 +25,8 @@ pub(super) struct InputSourcedAggregateParams<'a> { pub grouping_sets: Option<&'a [Vec]>, /// Post-aggregate sort keys, already lowered to bridge specs. pub sort_keys: Vec, + /// SELECT-list output alias of each GROUP BY key, parallel to `group_by`. + pub group_by_aliases: &'a [Option], pub tenant_id: TenantId, pub ctx: &'a ConvertContext, } @@ -49,6 +51,7 @@ pub(super) fn convert_input_sourced_aggregate( limit, grouping_sets, sort_keys, + group_by_aliases, tenant_id, ctx, } = p; @@ -78,6 +81,7 @@ pub(super) fn convert_input_sourced_aggregate( limit, grouping_sets, sort_keys, + group_by_aliases, names: input_field_names(input), }, )?]) diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs index 05fe6a15d..cbd17d66e 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/plan.rs @@ -34,6 +34,8 @@ pub(in crate::control::planner::sql_plan_convert) struct ConvertAggregateParams< pub limit: usize, pub grouping_sets: Option<&'a [Vec]>, pub sort_keys: &'a [SortKey], + /// SELECT-list output alias of each GROUP BY key, parallel to `group_by`. + pub group_by_aliases: &'a [Option], pub tenant_id: TenantId, pub ctx: &'a ConvertContext, } @@ -49,6 +51,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( limit, grouping_sets, sort_keys, + group_by_aliases, tenant_id, ctx, } = p; @@ -81,6 +84,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( limit, grouping_sets, sort_keys: bridge_sort_keys, + group_by_aliases, tenant_id, ctx, }, @@ -132,6 +136,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( limit, grouping_sets: bridge_grouping_sets, sort_keys: bridge_sort_keys, + group_by_aliases, names: InputFieldNames::Bare, }, )?]); @@ -201,6 +206,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_aggregate( sub_aggregates: Vec::new(), grouping_sets: bridge_grouping_sets, sort_keys: bridge_sort_keys, + group_key_aliases: group_by_aliases.to_vec(), }); let plan = if shuffle_agg_eligible { diff --git a/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs b/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs index 608e829ca..493ff74a7 100644 --- a/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs +++ b/nodedb/src/control/planner/sql_plan_convert/aggregate/spec.rs @@ -372,6 +372,8 @@ pub(in crate::control::planner::sql_plan_convert) struct InputSourcedTaskParams< /// ROLLUP / CUBE / GROUPING SETS, as indices into `group_by`. pub grouping_sets: Vec>, pub sort_keys: Vec, + /// SELECT-list output alias of each GROUP BY key, parallel to `group_by`. + pub group_by_aliases: &'a [Option], /// How `child`'s rows key their fields. pub names: InputFieldNames, } @@ -390,6 +392,7 @@ pub(in crate::control::planner::sql_plan_convert) fn build_input_sourced_aggrega limit, grouping_sets, sort_keys, + group_by_aliases, names, } = p; let group_specs = group_by_to_specs(group_by, names); @@ -417,6 +420,7 @@ pub(in crate::control::planner::sql_plan_convert) fn build_input_sourced_aggrega sub_aggregates: Vec::new(), grouping_sets, sort_keys, + group_key_aliases: group_by_aliases.to_vec(), }), post_set_op: PostSetOp::None, txn_id: None, diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/attr_type.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/attr_type.rs new file mode 100644 index 000000000..f6438f376 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/attr_type.rs @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The one mapping between the SQL attr type tag and the engine's. + +use nodedb_array::schema::AttrType; +use nodedb_sql::types_array::ArrayAttrType; + +/// The engine attr type of a SQL attr type. +pub(crate) fn engine_attr_type(t: ArrayAttrType) -> AttrType { + match t { + ArrayAttrType::Int64 => AttrType::Int64, + ArrayAttrType::Float64 => AttrType::Float64, + ArrayAttrType::String => AttrType::String, + ArrayAttrType::Bytes => AttrType::Bytes, + ArrayAttrType::UInt8 => AttrType::UInt8, + ArrayAttrType::Int16 => AttrType::Int16, + ArrayAttrType::UInt16 => AttrType::UInt16, + ArrayAttrType::Int32 => AttrType::Int32, + ArrayAttrType::UInt32 => AttrType::UInt32, + ArrayAttrType::Float32 => AttrType::Float32, + } +} + +/// The SQL attr type of an engine attr type. +pub(crate) fn sql_attr_type(t: AttrType) -> ArrayAttrType { + match t { + AttrType::Int64 => ArrayAttrType::Int64, + AttrType::Float64 => ArrayAttrType::Float64, + AttrType::String => ArrayAttrType::String, + AttrType::Bytes => ArrayAttrType::Bytes, + AttrType::UInt8 => ArrayAttrType::UInt8, + AttrType::Int16 => ArrayAttrType::Int16, + AttrType::UInt16 => ArrayAttrType::UInt16, + AttrType::Int32 => ArrayAttrType::Int32, + AttrType::UInt32 => ArrayAttrType::UInt32, + AttrType::Float32 => ArrayAttrType::Float32, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_mapping_round_trips() { + let all = [ + ArrayAttrType::Int64, + ArrayAttrType::Float64, + ArrayAttrType::String, + ArrayAttrType::Bytes, + ArrayAttrType::UInt8, + ArrayAttrType::Int16, + ArrayAttrType::UInt16, + ArrayAttrType::Int32, + ArrayAttrType::UInt32, + ArrayAttrType::Float32, + ]; + for t in all { + assert_eq!(sql_attr_type(engine_attr_type(t)), t); + assert_eq!(engine_attr_type(t).sql_name(), t.sql_name()); + } + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/ddl.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/ddl.rs index 2b79939a0..8e4f89147 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_convert/ddl.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/ddl.rs @@ -7,14 +7,13 @@ use std::sync::Arc; use nodedb_types::config::retention::BitemporalRetention; use nodedb_array::schema::{ - ArraySchema, ArraySchemaBuilder, AttrSpec, AttrType as EngineAttrType, CellOrder, DimSpec, - DimType as EngineDimType, TileOrder, + ArraySchema, ArraySchemaBuilder, AttrSpec, CellOrder, DimSpec, DimType as EngineDimType, + TileOrder, }; use nodedb_array::types::ArrayId; use nodedb_array::types::domain::{Domain, DomainBound}; use nodedb_sql::types_array::{ - ArrayAttrAst, ArrayAttrType, ArrayCellOrderAst, ArrayDimAst, ArrayDimType, ArrayDomainBound, - ArrayTileOrderAst, + ArrayAttrAst, ArrayCellOrderAst, ArrayDimAst, ArrayDimType, ArrayDomainBound, ArrayTileOrderAst, }; use crate::bridge::envelope::PhysicalPlan; @@ -213,12 +212,7 @@ pub(super) fn build_schema( builder = builder.dim(DimSpec::new(d.name.clone(), dtype, Domain::new(lo, hi))); } for a in attrs { - let dtype = match a.dtype { - ArrayAttrType::Int64 => EngineAttrType::Int64, - ArrayAttrType::Float64 => EngineAttrType::Float64, - ArrayAttrType::String => EngineAttrType::String, - ArrayAttrType::Bytes => EngineAttrType::Bytes, - }; + let dtype = super::attr_type::engine_attr_type(a.dtype); builder = builder.attr(AttrSpec::new(a.name.clone(), dtype, a.nullable)); } let extents: Vec = tile_extents.iter().map(|n| *n as u64).collect(); diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/dml.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/dml.rs index 557824317..7c60f6d4a 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_convert/dml.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/dml.rs @@ -3,8 +3,9 @@ //! `INSERT INTO ARRAY` / `DELETE FROM ARRAY` lowering to `PhysicalTask`. use nodedb_array::schema::ArraySchema; -use nodedb_array::tile::cell_tile_prefix; +use nodedb_array::tile::TilePrefixMemo; use nodedb_array::types::ArrayId; +use nodedb_array::types::coord::value::CoordValue; use nodedb_sql::types_array::{ArrayCoordLiteral, ArrayInsertRow}; use crate::bridge::envelope::PhysicalPlan; @@ -12,6 +13,7 @@ use crate::control::array_catalog::ArrayCatalogEntry; use crate::engine::array::wal::{ArrayDeleteCell, ArrayPutCell}; use crate::types::TenantId; use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp}; +use nodedb_types::Surrogate; use super::super::convert::ConvertContext; use super::helpers::{coerce_attrs, coerce_coords}; @@ -47,18 +49,63 @@ fn insert_target( Ok((entry, schema)) } -/// The identity key of every cell an `INSERT INTO ARRAY` writes, in row -/// order. -pub(in super::super) fn insert_array_cell_pks( +/// The coerced coordinate and identity key of every cell an +/// `INSERT INTO ARRAY` writes, in row order. The bind step reads them once, +/// and conversion reuses them. +#[derive(Debug)] +pub(in super::super) struct ArrayCellKeys { + pub coords: Vec>, + /// Each coordinate, msgpack-encoded: the cell's surrogate key. + pub pks: Vec>, + /// Each cell's bound surrogate, in row order, when the bind step bound + /// the batch. Empty otherwise: conversion then resolves each key. + pub surrogates: Vec, +} + +/// The [`ArrayCellKeys`] of the rows an `INSERT INTO ARRAY` writes. +pub(in super::super) fn insert_array_cell_keys( name: &str, rows: &[ArrayInsertRow], tenant_id: TenantId, ctx: &ConvertContext, -) -> crate::Result>> { +) -> crate::Result { let (_, schema) = insert_target(name, tenant_id, ctx)?; - rows.iter() - .map(|row| array_coord_pk(&row.coords, &schema)) - .collect() + let mut coords = Vec::with_capacity(rows.len()); + let mut pks = Vec::with_capacity(rows.len()); + for row in rows { + let coord = coerce_coords(&row.coords, &schema)?; + pks.push(coord_pk(&coord)?); + coords.push(coord); + } + Ok(ArrayCellKeys { + coords, + pks, + surrogates: Vec::new(), + }) +} + +/// Row `index`'s coerced coordinate and surrogate: the bind step's when it +/// kept them, else coerced here and resolved through `ctx`. +fn cell_identity( + kept: Option<&ArrayCellKeys>, + index: usize, + row: &ArrayInsertRow, + schema: &ArraySchema, + key: nodedb_types::CollectionKey<'_>, + ctx: &ConvertContext, +) -> crate::Result<(Vec, Surrogate)> { + if let Some(keys) = kept + && let (Some(coord), Some(pk)) = (keys.coords.get(index), keys.pks.get(index)) + { + let surrogate = match keys.surrogates.get(index) { + Some(&bound) => bound, + None => ctx.surrogate_for_pk(key, pk)?, + }; + return Ok((coord.clone(), surrogate)); + } + let coord = coerce_coords(&row.coords, schema)?; + let surrogate = ctx.surrogate_for_pk(key, &coord_pk(&coord)?)?; + Ok((coord, surrogate)) } pub(in super::super) fn convert_insert_array( @@ -70,26 +117,25 @@ pub(in super::super) fn convert_insert_array( let (entry, schema) = insert_target(name, tenant_id, ctx)?; let aid = ArrayId::in_database(tenant_id, ctx.database_id, name); - let vshard = ctx.collection_key(name).vshard(); + let key = ctx.collection_key(name); + let vshard = key.vshard(); let system_now_ms = chrono::Utc::now().timestamp_millis(); + let kept = ctx.prefetched.array_keys(key, rows); if ctx.cluster_enabled { let mut partitioned: Vec<(u64, Vec)> = Vec::with_capacity(rows.len()); - for row in rows { - let coord = coerce_coords(&row.coords, &schema)?; + let mut tiles = TilePrefixMemo::new(); + for (index, row) in rows.iter().enumerate() { + let (coord, surrogate) = cell_identity(kept, index, row, &schema, key, ctx)?; let attrs = coerce_attrs(&row.attrs, &schema)?; - let pk_bytes = - zerompk::to_msgpack_vec(&coord).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("array coord pk encode: {e}"), - })?; - let surrogate = ctx.surrogate_for_pk(ctx.collection_key(name), &pk_bytes)?; // A cell routes by its tile's key, the key a slice's shard // fan-out and each shard's tile filter use. let tile_key = - cell_tile_prefix(&schema, &coord).map_err(|e| crate::Error::PlanError { - detail: format!("INSERT INTO ARRAY {name}: tile key: {e}"), - })?; + tiles + .cell_prefix(&schema, &coord) + .map_err(|e| crate::Error::PlanError { + detail: format!("INSERT INTO ARRAY {name}: tile key: {e}"), + })?; let cell = ArrayPutCell { coord, attrs, @@ -131,15 +177,9 @@ pub(in super::super) fn convert_insert_array( // Single-node path: bundle all cells into one msgpack blob. let mut cells: Vec = Vec::with_capacity(rows.len()); - for row in rows { - let coord = coerce_coords(&row.coords, &schema)?; + for (index, row) in rows.iter().enumerate() { + let (coord, surrogate) = cell_identity(kept, index, row, &schema, key, ctx)?; let attrs = coerce_attrs(&row.attrs, &schema)?; - let pk_bytes = - zerompk::to_msgpack_vec(&coord).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("array coord pk encode: {e}"), - })?; - let surrogate = ctx.surrogate_for_pk(ctx.collection_key(name), &pk_bytes)?; cells.push(ArrayPutCell { coord, attrs, @@ -210,12 +250,15 @@ pub(in super::super) fn convert_delete_array( if ctx.cluster_enabled { let mut partitioned: Vec<(u64, Vec)> = Vec::with_capacity(coords.len()); + let mut tiles = TilePrefixMemo::new(); for row in coords { let typed = coerce_coords(row, &schema)?; let tile_key = - cell_tile_prefix(&schema, &typed).map_err(|e| crate::Error::PlanError { - detail: format!("DELETE FROM ARRAY {name}: tile key: {e}"), - })?; + tiles + .cell_prefix(&schema, &typed) + .map_err(|e| crate::Error::PlanError { + detail: format!("DELETE FROM ARRAY {name}: tile key: {e}"), + })?; let cell = ArrayDeleteCell { coord: typed, system_from_ms: system_now_ms, @@ -284,8 +327,7 @@ pub(in super::super) fn convert_delete_array( } /// The identity key of an array cell: its coerced coordinate, msgpack-encoded. -fn array_coord_pk(coords: &[ArrayCoordLiteral], schema: &ArraySchema) -> crate::Result> { - let coord = coerce_coords(coords, schema)?; +fn coord_pk(coord: &[CoordValue]) -> crate::Result> { zerompk::to_msgpack_vec(&coord).map_err(|e| crate::Error::Serialization { format: "msgpack".into(), detail: format!("array coord pk encode: {e}"), diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/helpers.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/helpers.rs index 68ea918e3..b40c0a31b 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_convert/helpers.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/helpers.rs @@ -3,6 +3,7 @@ //! Shared coercion helpers for array coordinate and attribute literals. use nodedb_array::schema::{ArraySchema, AttrType as EngineAttrType, DimType as EngineDimType}; +use nodedb_array::types::cell_value::NarrowError; use nodedb_array::types::cell_value::value::CellValue; use nodedb_array::types::coord::value::CoordValue; use nodedb_sql::types_array::{ArrayAttrLiteral, ArrayCoordLiteral}; @@ -61,6 +62,18 @@ pub(super) fn coerce_attrs( let mut out = Vec::with_capacity(attrs.len()); for (i, a) in attrs.iter().enumerate() { let spec = &schema.attrs[i]; + // A value past the attr's width is SQLSTATE 22003. A literal of a + // kind the attr never holds is a plan error. + let narrow = |r: Result| -> crate::Result { + r.map_err(|e| { + let detail = format!("attr `{}`: {e}", spec.name); + if e.is_out_of_range() { + crate::Error::NumericValueOutOfRange { detail } + } else { + crate::Error::PlanError { detail } + } + }) + }; let v = match (a, spec.dtype) { (ArrayAttrLiteral::Null, _) if spec.nullable => CellValue::Null, (ArrayAttrLiteral::Null, _) => { @@ -68,12 +81,13 @@ pub(super) fn coerce_attrs( detail: format!("attr `{}` is NOT NULL", spec.name), }); } - (ArrayAttrLiteral::Int64(n), EngineAttrType::Int64) => CellValue::Int64(*n), - (ArrayAttrLiteral::Int64(n), EngineAttrType::Float64) => CellValue::Float64(*n as f64), - (ArrayAttrLiteral::Float64(f), EngineAttrType::Float64) => CellValue::Float64(*f), + (ArrayAttrLiteral::Int64(n), dtype) => narrow(CellValue::from_int_literal(*n, dtype))?, + (ArrayAttrLiteral::Float64(f), dtype) => { + narrow(CellValue::from_float_literal(*f, dtype))? + } (ArrayAttrLiteral::String(s), EngineAttrType::String) => CellValue::String(s.clone()), (ArrayAttrLiteral::Bytes(b), EngineAttrType::Bytes) => CellValue::Bytes(b.clone()), - (got, want) => { + (got @ (ArrayAttrLiteral::String(_) | ArrayAttrLiteral::Bytes(_)), want) => { return Err(crate::Error::PlanError { detail: format!( "attr literal for `{}`: got {got:?}, expected attr type {want:?}", @@ -86,3 +100,56 @@ pub(super) fn coerce_attrs( } Ok(out) } + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, DimSpec}; + use nodedb_array::types::domain::{Domain, DomainBound}; + + fn schema() -> ArraySchema { + ArraySchemaBuilder::new("scene") + .dim(DimSpec::new( + "y", + EngineDimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("b1", EngineAttrType::UInt16, true)) + .attr(AttrSpec::new("ndvi", EngineAttrType::Float32, true)) + .tile_extents(vec![16]) + .build() + .expect("schema") + } + + #[test] + fn literals_narrow_to_the_declared_type() { + let cells = coerce_attrs( + &[ + ArrayAttrLiteral::Int64(65535), + ArrayAttrLiteral::Float64(0.1), + ], + &schema(), + ) + .expect("in range"); + assert_eq!( + cells, + vec![CellValue::UInt16(65535), CellValue::Float32(0.1)] + ); + } + + #[test] + fn a_value_past_the_width_is_22003() { + let err = coerce_attrs( + &[ArrayAttrLiteral::Int64(65536), ArrayAttrLiteral::Null], + &schema(), + ) + .expect_err("past UINT16"); + assert!(matches!(err, crate::Error::NumericValueOutOfRange { .. })); + let err = coerce_attrs( + &[ArrayAttrLiteral::Float64(1.5), ArrayAttrLiteral::Null], + &schema(), + ) + .expect_err("float into UINT16"); + assert!(matches!(err, crate::Error::PlanError { .. })); + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/mod.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/mod.rs index 08333b733..fec87c799 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_convert/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/mod.rs @@ -1,8 +1,15 @@ // SPDX-License-Identifier: BUSL-1.1 +mod attr_type; mod ddl; mod dml; mod helpers; +mod tile_route; + +pub(crate) use attr_type::{engine_attr_type, sql_attr_type}; pub(super) use ddl::{CreateArrayArgs, convert_create_array, convert_drop_array}; -pub(super) use dml::{convert_delete_array, convert_insert_array, insert_array_cell_pks}; +pub(super) use dml::{ + ArrayCellKeys, convert_delete_array, convert_insert_array, insert_array_cell_keys, +}; +pub(crate) use tile_route::cell_vshard; diff --git a/nodedb/src/control/planner/sql_plan_convert/array_convert/tile_route.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/tile_route.rs new file mode 100644 index 000000000..5b098af89 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/tile_route.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The vShard an array cell routes to, from its coordinate literals. + +use nodedb_array::schema::ArraySchema; +use nodedb_array::tile::TilePrefixMemo; +use nodedb_sql::types_array::ArrayCoordLiteral; +use nodedb_types::array_route::array_vshard_for_tile; + +use super::helpers::coerce_coords; + +/// The vShard that owns the tile of the cell at `coords`. The cluster write +/// path routes each cell of a put by this same key, so cells bucketed by it +/// reach one shard. `memo` carries the last tile key across one load's cells. +pub(crate) fn cell_vshard( + coords: &[ArrayCoordLiteral], + schema: &ArraySchema, + prefix_bits: u8, + memo: &mut TilePrefixMemo, +) -> crate::Result { + let coord = coerce_coords(coords, schema)?; + let tile_key = memo + .cell_prefix(schema, &coord) + .map_err(|e| crate::Error::PlanError { + detail: format!("array cell tile key: {e}"), + })?; + array_vshard_for_tile(tile_key, prefix_bits).map_err(|e| crate::Error::PlanError { + detail: e.to_string(), + }) +} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/aggregate.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/aggregate.rs index 6f51fcda3..78a31c260 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/aggregate.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/aggregate.rs @@ -2,34 +2,35 @@ //! ARRAY_AGG → PhysicalPlan::Array(ArrayOp::Aggregate) or ClusterArray(Agg). -use nodedb_array::schema::ArraySchema; use nodedb_array::types::ArrayId; -use nodedb_sql::temporal::TemporalScope; -use nodedb_sql::types_array::ArrayReducerAst; +use nodedb_sql::ArrayAggVisitArgs; use crate::bridge::envelope::PhysicalPlan; use crate::types::TenantId; use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp}; use super::super::convert::ConvertContext; -use super::helpers::{load_entry, map_reducer}; +use super::helpers::{decode_schema, load_entry, map_reducer}; +use super::region::resolve_region; +use super::tile_keys::compute_slice_tile_keys; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; pub(crate) fn convert_agg( - name: &str, - attr: &str, - reducer: ArrayReducerAst, - group_by_dim: Option<&str>, - temporal: TemporalScope, + args: ArrayAggVisitArgs<'_>, tenant_id: TenantId, ctx: &ConvertContext, ) -> crate::Result> { + let ArrayAggVisitArgs { + name, + attr, + reducer, + group_by_dim, + slice, + zone, + temporal, + } = args; let entry = load_entry(name, tenant_id, ctx)?; - let schema: ArraySchema = - zerompk::from_msgpack(&entry.schema_msgpack).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("array schema decode: {e}"), - })?; + let schema = decode_schema(&entry)?; let attr_idx = schema .attrs @@ -51,7 +52,17 @@ pub(crate) fn convert_agg( }; let (system_as_of, valid_at_ms) = - super::helpers::resolve_array_temporal(temporal, "ARRAY_AGG")?; + super::helpers::resolve_array_temporal(*temporal, "ARRAY_AGG")?; + let region = resolve_region(name, slice, zone, &schema)?; + let region_msgpack = match ®ion { + Some(region) => { + zerompk::to_msgpack_vec(region).map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("array region encode: {e}"), + })? + } + None => Vec::new(), + }; let mapped = map_reducer(reducer); let aid = ArrayId::in_database(tenant_id, ctx.database_id, name); let vshard = ctx.collection_key(name).vshard(); @@ -64,15 +75,21 @@ pub(crate) fn convert_agg( format: "msgpack".into(), detail: format!("array reducer encode: {e}"), })?; + // The fan-out covers the shards that own a tile of the region's box. + let box_ranges = region + .as_ref() + .map(|r| r.slice.dim_ranges.as_slice()) + .unwrap_or_default(); PhysicalPlan::ClusterArray(ClusterArrayOp::Agg { array_id: aid, attr_idx, reducer_msgpack, group_by_dim: group_by_dim_idx, - slice_tile_keys: vec![], + slice_tile_keys: compute_slice_tile_keys(&schema, box_ranges, entry.prefix_bits), prefix_bits: entry.prefix_bits, system_as_of, valid_at_ms, + region_msgpack, }) } else { PhysicalPlan::Array(ArrayOp::Aggregate { @@ -85,6 +102,7 @@ pub(crate) fn convert_agg( shard_bucket: None, system_as_of, valid_at_ms, + region_msgpack, }) }; @@ -103,6 +121,28 @@ mod tests { use super::*; use crate::control::array_catalog::ArrayCatalogEntry; use crate::control::planner::sql_plan_convert::convert::ConvertContext; + use nodedb_array::query::region::ArrayRegion; + use nodedb_array::query::slice::DimRange; + use nodedb_sql::temporal::TemporalScope; + use nodedb_sql::types_array::{ + ArrayCoordLiteral, ArrayReducerAst, ArraySliceAst, NamedDimRange, + }; + + /// `SUM(revenue)` over `sales` within `slice`. + fn sum_args<'a>( + slice: &'a ArraySliceAst, + temporal: &'a TemporalScope, + ) -> ArrayAggVisitArgs<'a> { + ArrayAggVisitArgs { + name: "sales", + attr: "revenue", + reducer: ArrayReducerAst::Sum, + group_by_dim: None, + slice, + zone: None, + temporal, + } + } use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; use nodedb_array::types::domain::{Domain, DomainBound}; @@ -166,34 +206,25 @@ mod tests { fn single_node_emits_local_agg() { let ctx = make_ctx(false); let tasks = convert_agg( - "sales", - "revenue", - ArrayReducerAst::Sum, - None, - TemporalScope::default(), + sum_args(&ArraySliceAst::default(), &TemporalScope::default()), TenantId::new(1), &ctx, ) .expect("convert"); assert_eq!(tasks.len(), 1); - assert!( - matches!( - &tasks[0].plan, - PhysicalPlan::Array(ArrayOp::Aggregate { .. }) - ), - "expected local Array variant" - ); + match &tasks[0].plan { + PhysicalPlan::Array(ArrayOp::Aggregate { region_msgpack, .. }) => { + assert!(region_msgpack.is_empty(), "no region reads the whole array"); + } + other => panic!("expected local Array variant, got {other:?}"), + } } #[test] fn cluster_mode_emits_cluster_agg() { let ctx = make_ctx(true); let tasks = convert_agg( - "sales", - "revenue", - ArrayReducerAst::Sum, - None, - TemporalScope::default(), + sum_args(&ArraySliceAst::default(), &TemporalScope::default()), TenantId::new(1), &ctx, ) @@ -208,6 +239,69 @@ mod tests { ); } + /// A box rides to every shard as the encoded region, and the fan-out + /// covers only the tiles inside it. + #[test] + fn a_box_encodes_the_region_and_narrows_the_fan_out() { + let ctx = make_ctx(true); + let slice = ArraySliceAst { + dim_ranges: vec![NamedDimRange { + dim: "month".into(), + lo: ArrayCoordLiteral::Int64(3), + hi: ArrayCoordLiteral::Int64(4), + }], + }; + let tasks = convert_agg( + sum_args(&slice, &TemporalScope::default()), + TenantId::new(1), + &ctx, + ) + .expect("convert"); + let PhysicalPlan::ClusterArray(ClusterArrayOp::Agg { + region_msgpack, + slice_tile_keys, + .. + }) = &tasks[0].plan + else { + panic!("expected ClusterArray variant, got {:?}", tasks[0].plan); + }; + let region: ArrayRegion = zerompk::from_msgpack(region_msgpack).expect("region decodes"); + assert_eq!( + region.slice.dim_ranges, + vec![Some(DimRange::new( + DomainBound::Int64(3), + DomainBound::Int64(4) + ))] + ); + assert!(region.zone.is_none()); + assert!( + (1..=2).contains(&slice_tile_keys.len()), + "two one-month tiles fan out to at most two buckets: {slice_tile_keys:?}" + ); + } + + #[test] + fn a_box_on_an_unknown_dim_is_refused() { + let ctx = make_ctx(false); + let slice = ArraySliceAst { + dim_ranges: vec![NamedDimRange { + dim: "year".into(), + lo: ArrayCoordLiteral::Int64(1), + hi: ArrayCoordLiteral::Int64(2), + }], + }; + let err = convert_agg( + sum_args(&slice, &TemporalScope::default()), + TenantId::new(1), + &ctx, + ) + .expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("no dim 'year'")), + "{err:?}" + ); + } + #[test] fn all_versions_aggregate_is_rejected() { // AS OF SYSTEM TIME NULL (AllVersions) is meaningless for an aggregate — @@ -216,18 +310,14 @@ mod tests { // live read. This is why `ArrayOp::Aggregate` / `ClusterArrayOp::Agg` // carry `Option` (Current | AsOf) and not `SystemTimeScope`: the // AllVersions state is made unrepresentable by rejecting it here. - // Audit-log reads go through NDARRAY_SLICE instead. + // Audit-log reads go through ARRAY_SLICE instead. let ctx = make_ctx(false); let temporal = TemporalScope { system_time: nodedb_types::SystemTimeScope::AllVersions, ..TemporalScope::default() }; let err = convert_agg( - "sales", - "revenue", - ArrayReducerAst::Sum, - None, - temporal, + sum_args(&ArraySliceAst::default(), &temporal), TenantId::new(1), &ctx, ) diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/elementwise.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/elementwise.rs index 17fc2fd4f..280c5f3cb 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/elementwise.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/elementwise.rs @@ -1,32 +1,69 @@ // SPDX-License-Identifier: BUSL-1.1 -//! ARRAY_ELEMENTWISE → PhysicalPlan::Array(ArrayOp::Elementwise). +//! ARRAY_ELEMENTWISE → PhysicalPlan::ClusterArray(Elementwise) or +//! Array(Elementwise). +//! +//! A cell routes by its tile key, and the tile key is a function of the +//! schema's dims and tile extents alone. Two arrays with one schema hash and +//! one `prefix_bits` therefore place each coordinate's cells on one shard, +//! and every shard pairs its own cells. Arrays whose routing differs keep a +//! coordinate's two cells on different shards, so they are refused. use nodedb_array::types::ArrayId; +use nodedb_sql::temporal::TemporalScope; use nodedb_sql::types_array::ArrayBinaryOpAst; use crate::bridge::envelope::PhysicalPlan; use crate::types::TenantId; -use nodedb_physical::physical_plan::ArrayOp; +use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp}; use super::super::convert::ConvertContext; -use super::helpers::{load_schema, map_binary_op}; +use super::helpers::{decode_schema, load_entry, map_binary_op}; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; +/// One `ARRAY_ELEMENTWISE` call as the planner names it. +pub(crate) struct ElementwiseCall<'a> { + pub left_name: &'a str, + pub right_name: &'a str, + pub op: ArrayBinaryOpAst, + pub attr: &'a str, + /// The `AS OF` scope both operands are read at. + pub temporal: TemporalScope, +} + pub(crate) fn convert_elementwise( - left_name: &str, - right_name: &str, - op: ArrayBinaryOpAst, - attr: &str, + call: ElementwiseCall<'_>, tenant_id: TenantId, ctx: &ConvertContext, ) -> crate::Result> { - let lschema = load_schema(left_name, tenant_id, ctx)?; - let rschema = load_schema(right_name, tenant_id, ctx)?; - if lschema.dims.len() != rschema.dims.len() || lschema.attrs.len() != rschema.attrs.len() { + let ElementwiseCall { + left_name, + right_name, + op, + attr, + temporal, + } = call; + let (system_as_of, valid_at_ms) = + super::helpers::resolve_array_temporal(temporal, "ARRAY_ELEMENTWISE")?; + let lentry = load_entry(left_name, tenant_id, ctx)?; + let rentry = load_entry(right_name, tenant_id, ctx)?; + let lschema = decode_schema(&lentry)?; + let rschema = decode_schema(&rentry)?; + if lentry.schema_hash != rentry.schema_hash { return Err(crate::Error::PlanError { detail: format!( - "ARRAY_ELEMENTWISE: arrays '{left_name}' and '{right_name}' have different shapes" + "ARRAY_ELEMENTWISE: arrays '{left_name}' and '{right_name}' have different \ + schemas; both need the same dims, attrs, tile extents and orders" + ), + }); + } + if lentry.prefix_bits != rentry.prefix_bits { + return Err(crate::Error::PlanError { + detail: format!( + "ARRAY_ELEMENTWISE: arrays '{left_name}' (prefix_bits {}) and '{right_name}' \ + (prefix_bits {}) route their tiles differently; create both with the same \ + prefix_bits", + lentry.prefix_bits, rentry.prefix_bits ), }); } @@ -45,18 +82,222 @@ pub(crate) fn convert_elementwise( let left = ArrayId::in_database(tenant_id, ctx.database_id, left_name); let right = ArrayId::in_database(tenant_id, ctx.database_id, right_name); let vshard = ctx.collection_key(left_name).vshard(); - Ok(vec![PhysicalTask { - tenant_id, - vshard_id: vshard, - database_id: ctx.database_id, - plan: PhysicalPlan::Array(ArrayOp::Elementwise { + let op = map_binary_op(op); + let plan = if ctx.cluster_enabled { + PhysicalPlan::ClusterArray(ClusterArrayOp::Elementwise { left, right, - op: map_binary_op(op), + op, + attr_idx, + prefix_bits: lentry.prefix_bits, + system_as_of, + valid_at_ms, + }) + } else { + PhysicalPlan::Array(ArrayOp::Elementwise { + left, + right, + op, attr_idx, cell_filter: None, - }), + shard_bucket: None, + system_as_of, + valid_at_ms, + }) + }; + Ok(vec![PhysicalTask { + tenant_id, + vshard_id: vshard, + database_id: ctx.database_id, + plan, post_set_op: PostSetOp::None, txn_id: None, }]) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::array_catalog::{ArrayCatalog, ArrayCatalogEntry}; + use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; + use nodedb_array::types::domain::{Domain, DomainBound}; + + /// A catalog holding `(name, schema_hash, prefix_bits)` arrays of one + /// 1-D shape. + fn make_ctx(cluster_enabled: bool, arrays: &[(&str, u64, u8)]) -> ConvertContext { + let schema = ArraySchemaBuilder::new("grid") + .dim(DimSpec::new( + "x".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("v".to_string(), AttrType::Float64, true)) + .tile_extents(vec![4]) + .build() + .expect("schema build"); + let schema_msgpack = zerompk::to_msgpack_vec(&schema).expect("encode"); + let handle = ArrayCatalog::handle(); + { + let mut cat = handle.write().expect("lock"); + for &(name, schema_hash, prefix_bits) in arrays { + cat.register(ArrayCatalogEntry { + array_id: ArrayId::new(TenantId::new(1), name), + name: name.into(), + schema_msgpack: schema_msgpack.clone(), + schema_hash, + created_at_ms: 0, + prefix_bits, + audit_retain_ms: None, + minimum_audit_retain_ms: None, + modification_hlc: nodedb_types::Hlc::ZERO, + incarnation: nodedb_types::Hlc::ZERO, + }) + .expect("register"); + } + } + ConvertContext { + purpose: super::super::super::convert::PlanningPurpose::Execute, + retention_registry: None, + array_catalog: Some(handle), + credentials: None, + wal: None, + surrogate_assigner: + crate::control::planner::sql_plan_convert::test_support::test_assigner(), + cluster_enabled, + bitemporal_retention_registry: None, + max_vector_dim: 0, + force_shuffle_join: false, + shuffle_num_parts: 0, + force_shuffle_agg: false, + shuffle_agg_num_parts: 0, + broadcast_threshold_bytes: 8 * 1024 * 1024, + shuffle_agg_threshold: 10_000, + prefetched: Default::default(), + database_id: crate::types::DatabaseId::DEFAULT, + tenant_id: crate::types::TenantId::new(0), + } + } + + fn convert_at( + ctx: &ConvertContext, + temporal: TemporalScope, + ) -> crate::Result> { + convert_elementwise( + ElementwiseCall { + left_name: "l", + right_name: "r", + op: ArrayBinaryOpAst::Add, + attr: "v", + temporal, + }, + TenantId::new(1), + ctx, + ) + } + + fn convert(ctx: &ConvertContext) -> crate::Result> { + convert_at(ctx, TemporalScope::default()) + } + + fn as_of(system_ms: i64, valid_ms: i64) -> TemporalScope { + TemporalScope { + system_time: nodedb_types::SystemTimeScope::AsOf(system_ms), + valid_time: nodedb_sql::temporal::ValidTime::At(valid_ms), + } + } + + #[test] + fn a_cluster_op_carries_the_as_of_scope() { + let ctx = make_ctx(true, &[("l", 7, 8), ("r", 7, 8)]); + let tasks = convert_at(&ctx, as_of(1_000, 2_000)).expect("convert"); + assert!(matches!( + &tasks[0].plan, + PhysicalPlan::ClusterArray(ClusterArrayOp::Elementwise { + system_as_of: Some(1_000), + valid_at_ms: Some(2_000), + .. + }) + )); + } + + #[test] + fn a_local_op_carries_the_as_of_scope() { + let ctx = make_ctx(false, &[("l", 7, 8), ("r", 7, 8)]); + let tasks = convert_at(&ctx, as_of(1_000, 2_000)).expect("convert"); + assert!(matches!( + &tasks[0].plan, + PhysicalPlan::Array(ArrayOp::Elementwise { + system_as_of: Some(1_000), + valid_at_ms: Some(2_000), + .. + }) + )); + } + + #[test] + fn all_system_versions_are_refused() { + let ctx = make_ctx(true, &[("l", 7, 8), ("r", 7, 8)]); + let temporal = TemporalScope { + system_time: nodedb_types::SystemTimeScope::AllVersions, + valid_time: nodedb_sql::temporal::ValidTime::Any, + }; + let err = convert_at(&ctx, temporal).expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("not meaningful for ARRAY_ELEMENTWISE")), + "{err:?}" + ); + } + + #[test] + fn a_cluster_plans_a_fan_out_with_the_shared_routing() { + let ctx = make_ctx(true, &[("l", 7, 8), ("r", 7, 8)]); + let tasks = convert(&ctx).expect("convert"); + match &tasks[0].plan { + PhysicalPlan::ClusterArray(ClusterArrayOp::Elementwise { + left, + right, + prefix_bits, + attr_idx, + .. + }) => { + assert_eq!((left.name.as_str(), right.name.as_str()), ("l", "r")); + assert_eq!(*prefix_bits, 8); + assert_eq!(*attr_idx, 0); + } + other => panic!("expected ClusterArray Elementwise, got {other:?}"), + } + } + + #[test] + fn a_single_node_plans_a_local_op() { + let ctx = make_ctx(false, &[("l", 7, 8), ("r", 7, 8)]); + let tasks = convert(&ctx).expect("convert"); + assert!(matches!( + &tasks[0].plan, + PhysicalPlan::Array(ArrayOp::Elementwise { + shard_bucket: None, + .. + }) + )); + } + + #[test] + fn arrays_with_different_routing_are_refused() { + let ctx = make_ctx(true, &[("l", 7, 8), ("r", 7, 10)]); + let err = convert(&ctx).expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("prefix_bits")), + "{err:?}" + ); + } + + #[test] + fn arrays_with_different_schemas_are_refused() { + let ctx = make_ctx(true, &[("l", 7, 8), ("r", 9, 8)]); + let err = convert(&ctx).expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("different schemas")), + "{err:?}" + ); + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/helpers.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/helpers.rs index 7fc407203..73ab4d6cc 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/helpers.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/helpers.rs @@ -2,10 +2,13 @@ //! Shared helpers for array fn → PhysicalTask lowering. -use nodedb_array::schema::{ArraySchema, AttrType as EngineAttrType, DimType as EngineDimType}; +use nodedb_array::query::slice::{DimRange, Slice}; +use nodedb_array::schema::{ArraySchema, DimType as EngineDimType}; use nodedb_array::types::domain::DomainBound; use nodedb_sql::temporal::{TemporalScope, ValidTime}; -use nodedb_sql::types_array::{ArrayBinaryOpAst, ArrayCoordLiteral, ArrayReducerAst}; +use nodedb_sql::types_array::{ + ArrayBinaryOpAst, ArrayCoordLiteral, ArrayReducerAst, ArraySliceAst, +}; use crate::control::array_catalog::ArrayCatalogEntry; use nodedb_physical::physical_plan::{ArrayBinaryOp, ArrayReducer}; @@ -36,12 +39,8 @@ pub(super) fn load_entry( }) } -pub(super) fn load_schema( - name: &str, - tenant_id: nodedb_types::TenantId, - ctx: &ConvertContext, -) -> crate::Result { - let entry = load_entry(name, tenant_id, ctx)?; +/// The schema a catalog entry holds. +pub(super) fn decode_schema(entry: &ArrayCatalogEntry) -> crate::Result { zerompk::from_msgpack(&entry.schema_msgpack).map_err(|e| crate::Error::Serialization { format: "msgpack".into(), detail: format!("array schema decode: {e}"), @@ -67,6 +66,32 @@ pub(super) fn resolve_attr_indices( Ok(out) } +/// The engine [`Slice`] for a named-dim slice predicate over array `name`. +/// Each range is coerced to its dim's type. A dim no range names stays +/// unconstrained. `context` prefixes every error. +pub(crate) fn resolve_slice( + context: &str, + name: &str, + ast: &ArraySliceAst, + schema: &ArraySchema, +) -> crate::Result { + let mut dim_ranges: Vec> = vec![None; schema.dims.len()]; + for r in &ast.dim_ranges { + let idx = schema + .dims + .iter() + .position(|d| d.name == r.dim) + .ok_or_else(|| crate::Error::PlanError { + detail: format!("{context}: array '{name}' has no dim '{}'", r.dim), + })?; + let dtype = schema.dims[idx].dtype; + let lo = coerce_bound(&r.lo, dtype, &r.dim)?; + let hi = coerce_bound(&r.hi, dtype, &r.dim)?; + dim_ranges[idx] = Some(DimRange::new(lo, hi)); + } + Ok(Slice::new(dim_ranges)) +} + pub(crate) fn coerce_bound( lit: &ArrayCoordLiteral, dtype: EngineDimType, @@ -129,7 +154,7 @@ pub(super) fn map_binary_op(o: ArrayBinaryOpAst) -> ArrayBinaryOp { /// /// `AS OF SYSTEM TIME NULL` (AllVersions) is rejected for aggregates because /// summing across all historical versions is not meaningful. Use -/// `NDARRAY_SLICE` for audit-log reads instead. +/// `ARRAY_SLICE` for audit-log reads instead. pub(super) fn resolve_array_temporal( temporal: TemporalScope, context: &str, @@ -140,8 +165,8 @@ pub(super) fn resolve_array_temporal( nodedb_types::SystemTimeScope::AllVersions => { return Err(crate::Error::PlanError { detail: format!( - "{context}: AS OF SYSTEM TIME NULL is not meaningful for NDARRAY_AGG; \ - use a point-in-time AS OF SYSTEM TIME , or NDARRAY_SLICE for the audit log" + "{context}: AS OF SYSTEM TIME NULL is not meaningful for {context}; \ + use a point-in-time AS OF SYSTEM TIME , or ARRAY_SLICE for the audit log" ), }); } @@ -190,5 +215,3 @@ pub(super) fn resolve_array_temporal_scope( }; Ok((system_time, valid_at_ms)) } - -const _UNUSED_ATTR_TYPE: Option = None; diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/maint.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/maint.rs index c8e1d68f0..08deebf9b 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/maint.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/maint.rs @@ -1,12 +1,17 @@ // SPDX-License-Identifier: BUSL-1.1 -//! ARRAY_FLUSH / ARRAY_COMPACT → PhysicalPlan::Array(ArrayOp::*). +//! ARRAY_FLUSH / ARRAY_COMPACT → PhysicalPlan::ClusterArray(Flush/Compact) or +//! Array(Flush/Compact). +//! +//! Each core holds the tiles of the vShards it owns. A cluster plan fans the +//! maintenance out to every primary vShard, so every core that holds tiles +//! of the array runs it. use nodedb_array::types::ArrayId; use crate::bridge::envelope::PhysicalPlan; use crate::types::TenantId; -use nodedb_physical::physical_plan::ArrayOp; +use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp}; use super::super::convert::ConvertContext; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; @@ -16,28 +21,28 @@ pub(crate) fn convert_flush( tenant_id: TenantId, ctx: &ConvertContext, ) -> crate::Result> { - let _ = super::helpers::load_entry(name, tenant_id, ctx)?; - let wal = ctx.wal.as_ref().ok_or_else(|| crate::Error::PlanError { - detail: "ARRAY_FLUSH: no WAL wired into convert context".into(), - })?; + let entry = super::helpers::load_entry(name, tenant_id, ctx)?; let aid = ArrayId::in_database(tenant_id, ctx.database_id, name); - let vshard = ctx.collection_key(name).vshard(); - // A frontier read, not an allocation: `Flush` appends no WAL record, and - // this LSN becomes the flushed segment's watermark — every cell it contains - // was written by a record below the current frontier. The write funnel mints - // nothing for `Flush` and so has no LSN to stamp here. - let wal_lsn = wal.next_lsn().as_u64(); - Ok(vec![PhysicalTask { - tenant_id, - vshard_id: vshard, - database_id: ctx.database_id, - plan: PhysicalPlan::Array(ArrayOp::Flush { + let plan = if ctx.cluster_enabled { + // Each shard's executor stamps the flush with its own node's WAL + // frontier. + PhysicalPlan::ClusterArray(ClusterArrayOp::Flush { array_id: aid, - wal_lsn, - }), - post_set_op: PostSetOp::None, - txn_id: None, - }]) + prefix_bits: entry.prefix_bits, + }) + } else { + let wal = ctx.wal.as_ref().ok_or_else(|| crate::Error::PlanError { + detail: "ARRAY_FLUSH: no WAL wired into convert context".into(), + })?; + // A frontier read, not an allocation: `Flush` appends no WAL record, + // and this LSN becomes the flushed segment's watermark. Every cell it + // contains was written by a record below the current frontier. + PhysicalPlan::Array(ArrayOp::Flush { + array_id: aid, + wal_lsn: wal.next_lsn().as_u64(), + }) + }; + Ok(vec![task(name, tenant_id, ctx, plan)]) } pub(crate) fn convert_compact( @@ -47,16 +52,137 @@ pub(crate) fn convert_compact( ) -> crate::Result> { let entry = super::helpers::load_entry(name, tenant_id, ctx)?; let aid = ArrayId::in_database(tenant_id, ctx.database_id, name); - let vshard = ctx.collection_key(name).vshard(); - Ok(vec![PhysicalTask { - tenant_id, - vshard_id: vshard, - database_id: ctx.database_id, - plan: PhysicalPlan::Array(ArrayOp::Compact { + let plan = if ctx.cluster_enabled { + PhysicalPlan::ClusterArray(ClusterArrayOp::Compact { array_id: aid, audit_retain_ms: entry.audit_retain_ms, - }), + prefix_bits: entry.prefix_bits, + }) + } else { + PhysicalPlan::Array(ArrayOp::Compact { + array_id: aid, + audit_retain_ms: entry.audit_retain_ms, + }) + }; + Ok(vec![task(name, tenant_id, ctx, plan)]) +} + +fn task(name: &str, tenant_id: TenantId, ctx: &ConvertContext, plan: PhysicalPlan) -> PhysicalTask { + PhysicalTask { + tenant_id, + vshard_id: ctx.collection_key(name).vshard(), + database_id: ctx.database_id, + plan, post_set_op: PostSetOp::None, txn_id: None, - }]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::array_catalog::{ArrayCatalog, ArrayCatalogEntry}; + use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; + use nodedb_array::types::domain::{Domain, DomainBound}; + + fn make_ctx(cluster_enabled: bool) -> ConvertContext { + let schema = ArraySchemaBuilder::new("grid") + .dim(DimSpec::new( + "x".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)), + )) + .attr(AttrSpec::new("v".to_string(), AttrType::Float64, true)) + .tile_extents(vec![4]) + .build() + .expect("schema build"); + let handle = ArrayCatalog::handle(); + handle + .write() + .expect("lock") + .register(ArrayCatalogEntry { + array_id: ArrayId::new(TenantId::new(1), "g"), + name: "g".into(), + schema_msgpack: zerompk::to_msgpack_vec(&schema).expect("encode"), + schema_hash: 7, + created_at_ms: 0, + prefix_bits: 6, + audit_retain_ms: Some(1_000), + minimum_audit_retain_ms: None, + modification_hlc: nodedb_types::Hlc::ZERO, + incarnation: nodedb_types::Hlc::ZERO, + }) + .expect("register"); + ConvertContext { + purpose: super::super::super::convert::PlanningPurpose::Execute, + retention_registry: None, + array_catalog: Some(handle), + credentials: None, + wal: None, + surrogate_assigner: + crate::control::planner::sql_plan_convert::test_support::test_assigner(), + cluster_enabled, + bitemporal_retention_registry: None, + max_vector_dim: 0, + force_shuffle_join: false, + shuffle_num_parts: 0, + force_shuffle_agg: false, + shuffle_agg_num_parts: 0, + broadcast_threshold_bytes: 8 * 1024 * 1024, + shuffle_agg_threshold: 10_000, + prefetched: Default::default(), + database_id: crate::types::DatabaseId::DEFAULT, + tenant_id: crate::types::TenantId::new(0), + } + } + + #[test] + fn a_cluster_flush_fans_out_with_the_array_routing() { + let ctx = make_ctx(true); + let tasks = convert_flush("g", TenantId::new(1), &ctx).expect("convert"); + match &tasks[0].plan { + PhysicalPlan::ClusterArray(ClusterArrayOp::Flush { + array_id, + prefix_bits, + }) => { + assert_eq!(array_id.name, "g"); + assert_eq!(*prefix_bits, 6); + } + other => panic!("expected ClusterArray Flush, got {other:?}"), + } + } + + #[test] + fn a_cluster_compact_carries_the_retention_window() { + let ctx = make_ctx(true); + let tasks = convert_compact("g", TenantId::new(1), &ctx).expect("convert"); + assert!(matches!( + &tasks[0].plan, + PhysicalPlan::ClusterArray(ClusterArrayOp::Compact { + audit_retain_ms: Some(1_000), + prefix_bits: 6, + .. + }) + )); + } + + #[test] + fn a_single_node_compact_plans_a_local_op() { + let ctx = make_ctx(false); + let tasks = convert_compact("g", TenantId::new(1), &ctx).expect("convert"); + assert!(matches!( + &tasks[0].plan, + PhysicalPlan::Array(ArrayOp::Compact { + audit_retain_ms: Some(1_000), + .. + }) + )); + } + + #[test] + fn a_single_node_flush_needs_a_wal() { + let ctx = make_ctx(false); + let err = convert_flush("g", TenantId::new(1), &ctx).expect_err("no WAL"); + assert!(err.to_string().contains("no WAL"), "{err}"); + } } diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/mod.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/mod.rs index a14431bf7..f696c39c7 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/mod.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: BUSL-1.1 //! `SqlPlan::Array*` → `PhysicalTask` lowering for the read/maint -//! function surface (ARRAY_SLICE, ARRAY_AGG, ARRAY_PROJECT, +//! function surface (ARRAY_SLICE, which ARRAY_PROJECT plans to, ARRAY_AGG, //! ARRAY_ELEMENTWISE, ARRAY_FLUSH, ARRAY_COMPACT). pub mod aggregate; pub mod elementwise; pub mod helpers; pub mod maint; -pub mod project; +pub mod region; pub mod slice; +pub mod tile_keys; pub(super) use aggregate::convert_agg; pub(super) use elementwise::convert_elementwise; pub(super) use maint::{convert_compact, convert_flush}; -pub(super) use project::convert_project; pub(super) use slice::convert_slice; diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/project.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/project.rs deleted file mode 100644 index 79b2ed2ec..000000000 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/project.rs +++ /dev/null @@ -1,41 +0,0 @@ -// SPDX-License-Identifier: BUSL-1.1 - -//! ARRAY_PROJECT → PhysicalPlan::Array(ArrayOp::Project). - -use nodedb_array::types::ArrayId; - -use crate::bridge::envelope::PhysicalPlan; -use crate::types::TenantId; -use nodedb_physical::physical_plan::ArrayOp; - -use super::super::convert::ConvertContext; -use super::helpers::{load_schema, resolve_attr_indices}; -use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; - -pub(crate) fn convert_project( - name: &str, - attr_projection: &[String], - tenant_id: TenantId, - ctx: &ConvertContext, -) -> crate::Result> { - let schema = load_schema(name, tenant_id, ctx)?; - let attr_indices = resolve_attr_indices(name, attr_projection, &schema)?; - if attr_indices.is_empty() { - return Err(crate::Error::PlanError { - detail: format!("ARRAY_PROJECT: array '{name}': attr list must not be empty"), - }); - } - let aid = ArrayId::in_database(tenant_id, ctx.database_id, name); - let vshard = ctx.collection_key(name).vshard(); - Ok(vec![PhysicalTask { - tenant_id, - vshard_id: vshard, - database_id: ctx.database_id, - plan: PhysicalPlan::Array(ArrayOp::Project { - array_id: aid, - attr_indices, - }), - post_set_op: PostSetOp::None, - txn_id: None, - }]) -} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/region.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/region.rs new file mode 100644 index 000000000..13fb6dbb8 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/region.rs @@ -0,0 +1,269 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The region an `ARRAY_AGG` reads, lowered to the engine [`ArrayRegion`]. +//! +//! The box comes from the slice argument. A zone also narrows the box on its +//! two dims to the bounding box of its polygons, so tile pruning and the +//! cluster fan-out see the zone's extent. An INT64 dim takes the ceiling of +//! the low edge and the floor of the high edge. + +use nodedb_array::query::region::{ArrayRegion, PolygonZone, ZonePolygon}; +use nodedb_array::query::slice::DimRange; +use nodedb_array::schema::{ArraySchema, DimType}; +use nodedb_array::types::domain::DomainBound; +use nodedb_sql::types_array::{ArraySliceAst, ArrayZoneAst}; + +use super::helpers::resolve_slice; + +/// The region of an `ARRAY_AGG` over array `name`, or `None` when it reads +/// the whole array. +pub(super) fn resolve_region( + name: &str, + slice: &ArraySliceAst, + zone: Option<&ArrayZoneAst>, + schema: &ArraySchema, +) -> crate::Result> { + if slice.dim_ranges.is_empty() && zone.is_none() { + return Ok(None); + } + let mut engine_slice = resolve_slice("ARRAY_AGG", name, slice, schema)?; + let zone = match zone { + None => None, + Some(zone) => { + let x = zone_axis(name, &zone.x_dim, schema)?; + let y = zone_axis(name, &zone.y_dim, schema)?; + let lowered = lower_zone(zone, x.dim, y.dim)?; + if let Some([x0, y0, x1, y1]) = zone_bbox(&lowered.polygons) { + for (axis, lo, hi) in [(x, x0, x1), (y, y0, y1)] { + let (lo, hi) = match axis.kind { + AxisKind::Int64 => ( + DomainBound::Int64(lo.ceil() as i64), + DomainBound::Int64(hi.floor() as i64), + ), + AxisKind::Float64 => (DomainBound::Float64(lo), DomainBound::Float64(hi)), + }; + narrow(&mut engine_slice.dim_ranges[axis.dim as usize], lo, hi); + } + } + Some(lowered) + } + }; + Ok(Some(ArrayRegion::new(engine_slice, zone))) +} + +/// The numeric type of a zone dim. +#[derive(Clone, Copy)] +enum AxisKind { + Int64, + Float64, +} + +/// One zone dim: its schema position and its numeric type. +#[derive(Clone, Copy)] +struct ZoneAxis { + dim: u32, + kind: AxisKind, +} + +fn zone_axis(name: &str, dim: &str, schema: &ArraySchema) -> crate::Result { + let idx = schema + .dims + .iter() + .position(|d| d.name == dim) + .ok_or_else(|| crate::Error::PlanError { + detail: format!("ARRAY_AGG: array '{name}' has no dim '{dim}'"), + })?; + let kind = match schema.dims[idx].dtype { + DimType::Int64 => AxisKind::Int64, + DimType::Float64 => AxisKind::Float64, + other @ (DimType::TimestampMs | DimType::String) => { + return Err(crate::Error::PlanError { + detail: format!( + "ARRAY_AGG: zone dim '{dim}' is {other:?}; a zone needs an INT64 or FLOAT64 \ + dim" + ), + }); + } + }; + let dim = u32::try_from(idx).map_err(|_| crate::Error::PlanError { + detail: format!("ARRAY_AGG: zone dim '{dim}' is past the dim index range"), + })?; + Ok(ZoneAxis { dim, kind }) +} + +fn lower_zone(zone: &ArrayZoneAst, x_dim: u32, y_dim: u32) -> crate::Result { + let polygons = zone + .polygons + .iter() + .map(|rings| { + let (exterior, holes) = rings.split_first().ok_or_else(|| crate::Error::PlanError { + detail: "ARRAY_AGG: a zone polygon has no exterior ring".into(), + })?; + Ok(ZonePolygon { + exterior: exterior.clone(), + holes: holes.to_vec(), + }) + }) + .collect::>>()?; + Ok(PolygonZone { + x_dim, + y_dim, + polygons, + }) +} + +/// `[min_x, min_y, max_x, max_y]` over every exterior ring, or `None` when +/// no ring holds a position. +fn zone_bbox(polygons: &[ZonePolygon]) -> Option<[f64; 4]> { + let mut points = polygons.iter().flat_map(|p| p.exterior.iter()); + let first = points.next()?; + let mut b = [first[0], first[1], first[0], first[1]]; + for p in points { + b[0] = b[0].min(p[0]); + b[1] = b[1].min(p[1]); + b[2] = b[2].max(p[0]); + b[3] = b[3].max(p[1]); + } + Some(b) +} + +/// Intersect `range` with `[lo, hi]`. The bounds share the dim's type. +fn narrow(range: &mut Option, lo: DomainBound, hi: DomainBound) { + *range = Some(match range.take() { + None => DimRange::new(lo, hi), + Some(r) => DimRange::new( + if lo > r.lo { lo } else { r.lo }, + if hi < r.hi { hi } else { r.hi }, + ), + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec}; + use nodedb_array::types::domain::Domain; + use nodedb_sql::types_array::{ArrayCoordLiteral, NamedDimRange}; + + fn schema() -> ArraySchema { + ArraySchemaBuilder::new("grid") + .dim(DimSpec::new( + "lon".to_string(), + DimType::Float64, + Domain::new(DomainBound::Float64(0.0), DomainBound::Float64(100.0)), + )) + .dim(DimSpec::new( + "row".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)), + )) + .dim(DimSpec::new( + "ts".to_string(), + DimType::TimestampMs, + Domain::new(DomainBound::TimestampMs(0), DomainBound::TimestampMs(1_000)), + )) + .attr(AttrSpec::new("v".to_string(), AttrType::Float64, true)) + .tile_extents(vec![10, 10, 100]) + .build() + .expect("schema build") + } + + fn zone(x: &str, y: &str) -> ArrayZoneAst { + ArrayZoneAst { + polygons: vec![vec![vec![ + [1.5, 2.5], + [7.25, 2.5], + [7.25, 6.5], + [1.5, 6.5], + [1.5, 2.5], + ]]], + x_dim: x.into(), + y_dim: y.into(), + } + } + + fn range(r: &ArrayRegion, dim: usize) -> Option<(DomainBound, DomainBound)> { + r.slice.dim_ranges[dim] + .as_ref() + .map(|r| (r.lo.clone(), r.hi.clone())) + } + + #[test] + fn the_whole_array_has_no_region() { + let region = resolve_region("grid", &ArraySliceAst::default(), None, &schema()).unwrap(); + assert!(region.is_none()); + } + + #[test] + fn a_zone_narrows_the_box_to_its_bounding_box() { + let region = resolve_region( + "grid", + &ArraySliceAst::default(), + Some(&zone("lon", "row")), + &schema(), + ) + .unwrap() + .expect("region"); + assert_eq!( + range(®ion, 0), + Some((DomainBound::Float64(1.5), DomainBound::Float64(7.25))) + ); + assert_eq!( + range(®ion, 1), + Some((DomainBound::Int64(3), DomainBound::Int64(6))), + "an INT64 dim takes the ceiling of the low edge and the floor of the high edge" + ); + assert!(range(®ion, 2).is_none()); + let zone = region.zone.expect("zone"); + assert_eq!((zone.x_dim, zone.y_dim), (0, 1)); + assert_eq!(zone.polygons.len(), 1); + } + + #[test] + fn the_zone_box_intersects_the_slice_box() { + let slice = ArraySliceAst { + dim_ranges: vec![NamedDimRange { + dim: "lon".into(), + lo: ArrayCoordLiteral::Int64(5), + hi: ArrayCoordLiteral::Int64(50), + }], + }; + let region = resolve_region("grid", &slice, Some(&zone("lon", "row")), &schema()) + .unwrap() + .expect("region"); + assert_eq!( + range(®ion, 0), + Some((DomainBound::Float64(5.0), DomainBound::Float64(7.25))) + ); + } + + #[test] + fn a_zone_on_a_timestamp_dim_is_refused() { + let err = resolve_region( + "grid", + &ArraySliceAst::default(), + Some(&zone("lon", "ts")), + &schema(), + ) + .expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("zone dim 'ts'")), + "{err:?}" + ); + } + + #[test] + fn a_zone_on_an_unknown_dim_is_refused() { + let err = resolve_region( + "grid", + &ArraySliceAst::default(), + Some(&zone("lon", "lat")), + &schema(), + ) + .expect_err("refused"); + assert!( + matches!(&err, crate::Error::PlanError { detail } if detail.contains("no dim 'lat'")), + "{err:?}" + ); + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/slice.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/slice.rs index 5ae48d1d2..89ff79abc 100644 --- a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/slice.rs +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/slice.rs @@ -2,22 +2,17 @@ //! ARRAY_SLICE → PhysicalPlan::Array(ArrayOp::Slice) or ClusterArray(Slice). -use nodedb_array::query::slice::{DimRange, Slice}; -use nodedb_array::schema::ArraySchema; -use nodedb_array::schema::DimSpec; -use nodedb_array::tile::{tile_prefix, tiles_per_dim}; use nodedb_array::types::ArrayId; -use nodedb_array::types::domain::DomainBound; use nodedb_sql::temporal::TemporalScope; use nodedb_sql::types_array::ArraySliceAst; -use nodedb_types::array_route::prefix_bucket; use crate::bridge::envelope::PhysicalPlan; use crate::types::TenantId; use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp}; use super::super::convert::ConvertContext; -use super::helpers::{coerce_bound, load_entry, resolve_attr_indices}; +use super::helpers::{decode_schema, load_entry, resolve_attr_indices, resolve_slice}; +use super::tile_keys::compute_slice_tile_keys; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; pub(crate) fn convert_slice( @@ -30,27 +25,9 @@ pub(crate) fn convert_slice( ctx: &ConvertContext, ) -> crate::Result> { let entry = load_entry(name, tenant_id, ctx)?; - let schema: ArraySchema = - zerompk::from_msgpack(&entry.schema_msgpack).map_err(|e| crate::Error::Serialization { - format: "msgpack".into(), - detail: format!("array schema decode: {e}"), - })?; + let schema = decode_schema(&entry)?; - let mut dim_ranges: Vec> = vec![None; schema.dims.len()]; - for r in &slice_ast.dim_ranges { - let idx = schema - .dims - .iter() - .position(|d| d.name == r.dim) - .ok_or_else(|| crate::Error::PlanError { - detail: format!("ARRAY_SLICE: array '{name}' has no dim '{}'", r.dim), - })?; - let dtype = schema.dims[idx].dtype; - let lo = coerce_bound(&r.lo, dtype, &r.dim)?; - let hi = coerce_bound(&r.hi, dtype, &r.dim)?; - dim_ranges[idx] = Some(DimRange::new(lo, hi)); - } - let slice = Slice::new(dim_ranges); + let slice = resolve_slice("ARRAY_SLICE", name, slice_ast, &schema)?; let slice_msgpack = zerompk::to_msgpack_vec(&slice).map_err(|e| crate::Error::Serialization { format: "msgpack".into(), @@ -105,103 +82,13 @@ pub(crate) fn convert_slice( }]) } -/// The most tiles a slice's shard fan-out enumerates. A slice covering more -/// tiles fans out to every shard. -const MAX_FAN_OUT_TILES: u128 = 4096; - -/// The tile keys a slice's shard fan-out covers: one key per routing bucket -/// (`prefix_bucket`), each the key of one tile in that bucket. -/// -/// Every tile whose cells can match the slice is enumerated, so every shard -/// that holds such a cell, or will hold one a later write lands there, is -/// covered. A transaction's read of the slice then validates on each of -/// those shards. A cell routes by its tile's key (`cell_tile_prefix`), the -/// same key enumerated here. -/// -/// Returns an empty vec, which fans out to every shard, when the slice -/// covers more than [`MAX_FAN_OUT_TILES`] tiles, when the schema has no -/// dimension, or when a tile key fails to encode. -fn compute_slice_tile_keys( - schema: &ArraySchema, - dim_ranges: &[Option], - prefix_bits: u8, -) -> Vec { - if schema.dims.is_empty() || prefix_bits == 0 || prefix_bits > 16 { - return Vec::new(); - } - let spans: Vec<(u64, u64)> = schema - .dims - .iter() - .zip(schema.tile_extents.iter()) - .enumerate() - .map(|(i, (dim, extent))| { - dim_tile_span(dim, *extent, dim_ranges.get(i).and_then(Option::as_ref)) - }) - .collect(); - let tiles: u128 = spans - .iter() - .map(|(lo, hi)| u128::from(hi - lo) + 1) - .fold(1u128, |acc, n| acc.saturating_mul(n)); - if tiles > MAX_FAN_OUT_TILES { - return Vec::new(); - } - - let mut buckets: std::collections::BTreeMap = std::collections::BTreeMap::new(); - let mut indices: Vec = spans.iter().map(|(lo, _)| *lo).collect(); - loop { - let Ok(key) = tile_prefix(schema, &indices) else { - return Vec::new(); - }; - buckets - .entry(prefix_bucket(key, prefix_bits)) - .or_insert(key); - // Advance the odometer over the per-dimension tile spans. - let mut dim = 0; - loop { - if dim == indices.len() { - return buckets.into_values().collect(); - } - if indices[dim] < spans[dim].1 { - indices[dim] += 1; - break; - } - indices[dim] = spans[dim].0; - dim += 1; - } - } -} - -/// The inclusive tile index span along `dim` that `range` covers, the whole -/// dimension when `range` is `None`. A string dimension tiles by hash, so a -/// range on it covers every tile. A bound outside the domain is clamped to -/// the domain's tiles. -fn dim_tile_span(dim: &DimSpec, extent: u64, range: Option<&DimRange>) -> (u64, u64) { - let last = u64::try_from(tiles_per_dim(dim, extent).saturating_sub(1)).unwrap_or(u64::MAX); - let extent = u128::from(extent.max(1)); - let index = |bound: &DomainBound| -> Option { - let offset: u128 = match (bound, &dim.domain.lo) { - (DomainBound::Int64(v), DomainBound::Int64(lo)) - | (DomainBound::TimestampMs(v), DomainBound::TimestampMs(lo)) => { - ((*v as i128) - (*lo as i128)).max(0) as u128 - } - (DomainBound::Float64(v), DomainBound::Float64(lo)) if v.is_finite() => { - (v - lo).max(0.0) as u128 - } - _ => return None, - }; - Some(u64::try_from(offset / extent).unwrap_or(u64::MAX).min(last)) - }; - match range.map(|r| (index(&r.lo), index(&r.hi))) { - Some((Some(lo), Some(hi))) => (lo.min(hi), lo.max(hi)), - _ => (0, last), - } -} - #[cfg(test)] mod tests { use super::*; use crate::control::array_catalog::ArrayCatalogEntry; - use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType}; + use nodedb_array::schema::{ + ArraySchema, ArraySchemaBuilder, AttrSpec, AttrType, DimSpec, DimType, + }; use nodedb_array::types::domain::{Domain, DomainBound}; fn make_schema() -> (Vec, ArraySchema) { @@ -268,68 +155,6 @@ mod tests { (ctx, handle) } - /// Every cell a slice can match, or a later write can land in its box, - /// routes to a shard the slice fans out to. - #[test] - fn the_fan_out_covers_the_shard_of_every_cell_in_the_box() { - use nodedb_array::types::coord::value::CoordValue; - let schema = ArraySchemaBuilder::new("grid") - .dim(DimSpec::new( - "chr".to_string(), - DimType::Int64, - Domain::new(DomainBound::Int64(0), DomainBound::Int64(9)), - )) - .dim(DimSpec::new( - "pos".to_string(), - DimType::Int64, - Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)), - )) - .attr(AttrSpec::new("qual".to_string(), AttrType::Float64, true)) - .tile_extents(vec![1, 100]) - .build() - .expect("schema build"); - let box_ranges = vec![ - Some(DimRange::new(DomainBound::Int64(0), DomainBound::Int64(2))), - Some(DimRange::new(DomainBound::Int64(0), DomainBound::Int64(99))), - ]; - let keys = compute_slice_tile_keys(&schema, &box_ranges, 8); - assert!(!keys.is_empty(), "a three-tile slice fans out by tile"); - let fanned = nodedb_cluster::distributed_array::array_vshards_for_slice(&keys, 8, 1024) - .expect("fan-out"); - let mut owners = std::collections::BTreeSet::new(); - for chr in 0..=2 { - for pos in [0, 10, 20, 50, 99] { - let coord = [CoordValue::Int64(chr), CoordValue::Int64(pos)]; - let key = nodedb_array::tile::cell_tile_prefix(&schema, &coord).expect("route"); - let owner = nodedb_cluster::distributed_array::array_vshard_for_tile(key, 8) - .expect("owner"); - assert!( - fanned.contains(&owner), - "cell ({chr}, {pos}) routes to vShard {owner}, outside the fan-out {fanned:?}" - ); - owners.insert(owner); - } - } - assert!(owners.len() > 1, "the slice's tiles spread over shards"); - } - - #[test] - fn a_slice_of_more_tiles_than_the_cap_fans_out_to_every_shard() { - let (_, schema) = make_schema(); - let wide = ArraySchemaBuilder::new("wide") - .dim(DimSpec::new( - "x".to_string(), - DimType::Int64, - Domain::new(DomainBound::Int64(0), DomainBound::Int64(1_000_000)), - )) - .attr(AttrSpec::new("val".to_string(), AttrType::Float64, true)) - .tile_extents(vec![1]) - .build() - .expect("schema build"); - assert!(compute_slice_tile_keys(&wide, &[None], 8).is_empty()); - assert!(!compute_slice_tile_keys(&schema, &[None], 8).is_empty()); - } - #[test] fn single_node_emits_local_array() { let (ctx, _handle) = make_ctx(false); diff --git a/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/tile_keys.rs b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/tile_keys.rs new file mode 100644 index 000000000..f286b3a98 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/array_fn_convert/tile_keys.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The tile keys a cluster array read fans out to: one per routing bucket +//! that a box of dim ranges covers. + +use nodedb_array::query::slice::DimRange; +use nodedb_array::schema::ArraySchema; +use nodedb_array::schema::DimSpec; +use nodedb_array::tile::{tile_prefix, tiles_per_dim}; +use nodedb_array::types::domain::DomainBound; +use nodedb_types::array_route::prefix_bucket; + +/// The most tiles a read's shard fan-out enumerates. A read covering more +/// tiles fans out to every shard. +const MAX_FAN_OUT_TILES: u128 = 4096; + +/// The tile keys a read's shard fan-out covers: one key per routing bucket +/// (`prefix_bucket`), each the key of one tile in that bucket. +/// +/// Every tile whose cells can match the box is enumerated, so every shard +/// that holds such a cell, or will hold one a later write lands there, is +/// covered. A transaction's read of the box then validates on each of +/// those shards. A cell routes by its tile's key (`cell_tile_prefix`), the +/// same key enumerated here. +/// +/// Returns an empty vec, which fans out to every shard, when the box +/// covers more than [`MAX_FAN_OUT_TILES`] tiles, when the schema has no +/// dimension, or when a tile key fails to encode. +pub(super) fn compute_slice_tile_keys( + schema: &ArraySchema, + dim_ranges: &[Option], + prefix_bits: u8, +) -> Vec { + if schema.dims.is_empty() || prefix_bits == 0 || prefix_bits > 16 { + return Vec::new(); + } + let spans: Vec<(u64, u64)> = schema + .dims + .iter() + .zip(schema.tile_extents.iter()) + .enumerate() + .map(|(i, (dim, extent))| { + dim_tile_span(dim, *extent, dim_ranges.get(i).and_then(Option::as_ref)) + }) + .collect(); + let tiles: u128 = spans + .iter() + .map(|(lo, hi)| u128::from(hi - lo) + 1) + .fold(1u128, |acc, n| acc.saturating_mul(n)); + if tiles > MAX_FAN_OUT_TILES { + return Vec::new(); + } + + let mut buckets: std::collections::BTreeMap = std::collections::BTreeMap::new(); + let mut indices: Vec = spans.iter().map(|(lo, _)| *lo).collect(); + loop { + let Ok(key) = tile_prefix(schema, &indices) else { + return Vec::new(); + }; + buckets + .entry(prefix_bucket(key, prefix_bits)) + .or_insert(key); + // Advance the odometer over the per-dimension tile spans. + let mut dim = 0; + loop { + if dim == indices.len() { + return buckets.into_values().collect(); + } + if indices[dim] < spans[dim].1 { + indices[dim] += 1; + break; + } + indices[dim] = spans[dim].0; + dim += 1; + } + } +} + +/// The inclusive tile index span along `dim` that `range` covers, the whole +/// dimension when `range` is `None`. A string dimension tiles by hash, so a +/// range on it covers every tile. A bound outside the domain is clamped to +/// the domain's tiles. +fn dim_tile_span(dim: &DimSpec, extent: u64, range: Option<&DimRange>) -> (u64, u64) { + let last = u64::try_from(tiles_per_dim(dim, extent).saturating_sub(1)).unwrap_or(u64::MAX); + let extent = u128::from(extent.max(1)); + let index = |bound: &DomainBound| -> Option { + let offset: u128 = match (bound, &dim.domain.lo) { + (DomainBound::Int64(v), DomainBound::Int64(lo)) + | (DomainBound::TimestampMs(v), DomainBound::TimestampMs(lo)) => { + ((*v as i128) - (*lo as i128)).max(0) as u128 + } + (DomainBound::Float64(v), DomainBound::Float64(lo)) if v.is_finite() => { + (v - lo).max(0.0) as u128 + } + _ => return None, + }; + Some(u64::try_from(offset / extent).unwrap_or(u64::MAX).min(last)) + }; + match range.map(|r| (index(&r.lo), index(&r.hi))) { + Some((Some(lo), Some(hi))) => (lo.min(hi), lo.max(hi)), + _ => (0, last), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_array::schema::{ArraySchemaBuilder, AttrSpec, AttrType, DimType}; + use nodedb_array::types::domain::Domain; + + /// Every cell a slice can match, or a later write can land in its box, + /// routes to a shard the slice fans out to. + #[test] + fn the_fan_out_covers_the_shard_of_every_cell_in_the_box() { + use nodedb_array::types::coord::value::CoordValue; + let schema = ArraySchemaBuilder::new("grid") + .dim(DimSpec::new( + "chr".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(9)), + )) + .dim(DimSpec::new( + "pos".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)), + )) + .attr(AttrSpec::new("qual".to_string(), AttrType::Float64, true)) + .tile_extents(vec![1, 100]) + .build() + .expect("schema build"); + let box_ranges = vec![ + Some(DimRange::new(DomainBound::Int64(0), DomainBound::Int64(2))), + Some(DimRange::new(DomainBound::Int64(0), DomainBound::Int64(99))), + ]; + let keys = compute_slice_tile_keys(&schema, &box_ranges, 8); + assert!(!keys.is_empty(), "a three-tile slice fans out by tile"); + let fanned = nodedb_cluster::distributed_array::array_vshards_for_slice(&keys, 8, 1024) + .expect("fan-out"); + let mut owners = std::collections::BTreeSet::new(); + for chr in 0..=2 { + for pos in [0, 10, 20, 50, 99] { + let coord = [CoordValue::Int64(chr), CoordValue::Int64(pos)]; + let key = nodedb_array::tile::cell_tile_prefix(&schema, &coord).expect("route"); + let owner = nodedb_cluster::distributed_array::array_vshard_for_tile(key, 8) + .expect("owner"); + assert!( + fanned.contains(&owner), + "cell ({chr}, {pos}) routes to vShard {owner}, outside the fan-out {fanned:?}" + ); + owners.insert(owner); + } + } + assert!(owners.len() > 1, "the slice's tiles spread over shards"); + } + + #[test] + fn a_slice_of_more_tiles_than_the_cap_fans_out_to_every_shard() { + let narrow = ArraySchemaBuilder::new("test_arr") + .dim(DimSpec::new( + "x".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(100)), + )) + .attr(AttrSpec::new("val".to_string(), AttrType::Float64, true)) + .tile_extents(vec![10]) + .build() + .expect("schema build"); + let wide = ArraySchemaBuilder::new("wide") + .dim(DimSpec::new( + "x".to_string(), + DimType::Int64, + Domain::new(DomainBound::Int64(0), DomainBound::Int64(1_000_000)), + )) + .attr(AttrSpec::new("val".to_string(), AttrType::Float64, true)) + .tile_extents(vec![1]) + .build() + .expect("schema build"); + assert!(compute_slice_tile_keys(&wide, &[None], 8).is_empty()); + assert!(!compute_slice_tile_keys(&narrow, &[None], 8).is_empty()); + } +} diff --git a/nodedb/src/control/planner/sql_plan_convert/body.rs b/nodedb/src/control/planner/sql_plan_convert/body.rs index e4690ac37..ddad5343e 100644 --- a/nodedb/src/control/planner/sql_plan_convert/body.rs +++ b/nodedb/src/control/planner/sql_plan_convert/body.rs @@ -94,7 +94,6 @@ pub(super) fn convert_body_to_single_plan( | SqlPlan::InsertArray { .. } | SqlPlan::DeleteArray { .. } | SqlPlan::ArraySlice { .. } - | SqlPlan::ArrayProject { .. } | SqlPlan::ArrayAgg { .. } | SqlPlan::ArrayElementwise { .. } | SqlPlan::ArrayFlush { .. } diff --git a/nodedb/src/control/planner/sql_plan_convert/dml/mod.rs b/nodedb/src/control/planner/sql_plan_convert/dml/mod.rs index f63c12dae..fdaf97af3 100644 --- a/nodedb/src/control/planner/sql_plan_convert/dml/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/dml/mod.rs @@ -17,7 +17,7 @@ pub(super) use insert::{ConvertInsertArgs, convert_insert, declared_primary_key_ pub(crate) use insert::{DEFAULT_IDENTITY_COLUMN, build_columnar_schema}; pub(super) use kv_insert::convert_kv_insert; pub(super) use merge::{ConvertMergeArgs, convert_merge}; -pub(super) use surrogate_keys::plan_key_batches; +pub(super) use surrogate_keys::{ArrayBatchRows, KeyBatch, plan_key_batches}; pub(super) use update_delete::{ UpdateFromParams, UpdateParams, convert_delete, convert_update, convert_update_from, }; diff --git a/nodedb/src/control/planner/sql_plan_convert/dml/surrogate_keys.rs b/nodedb/src/control/planner/sql_plan_convert/dml/surrogate_keys.rs index 3b66b8eee..e53d7923a 100644 --- a/nodedb/src/control/planner/sql_plan_convert/dml/surrogate_keys.rs +++ b/nodedb/src/control/planner/sql_plan_convert/dml/surrogate_keys.rs @@ -13,6 +13,9 @@ use nodedb_sql::types::{ VectorPrimaryUpdatePlan, }; +use nodedb_array::types::coord::value::CoordValue; +use nodedb_sql::types_array::ArrayInsertRow; + use super::super::convert::ConvertContext; use super::super::value::{sql_value_to_bytes, sql_value_to_string}; use super::crdt_gate::document_collection_is_crdt; @@ -31,6 +34,17 @@ pub(in super::super) struct KeyBatch<'p> { pub binds: bool, /// How many rows conversion mints a fresh surrogate for. pub fresh: usize, + /// For an `INSERT INTO ARRAY`: the plan's rows and their coerced + /// coordinates. The bind step keeps them with `pks`, so conversion + /// coerces and encodes no coordinate again. + pub array_rows: Option>, +} + +/// The rows of one `INSERT INTO ARRAY` plan and their coerced coordinates, +/// in row order. +pub(in super::super) struct ArrayBatchRows<'p> { + pub rows: &'p [ArrayInsertRow], + pub coords: Vec>, } /// The key batches of every plan in `plans`, in plan order. A plan whose @@ -134,6 +148,7 @@ fn plan_key_batch<'p>( .collect(), binds: true, fresh: 0, + array_rows: None, }, SqlPlan::Update { collection, @@ -147,6 +162,7 @@ fn plan_key_batch<'p>( pks: kv_keys(target_keys), binds: true, fresh: 0, + array_rows: None, }, EngineType::Timeseries | EngineType::Columnar | EngineType::Spatial => { return Ok(None); @@ -164,6 +180,7 @@ fn plan_key_batch<'p>( pks: document_keys(target_keys), binds: crdt, fresh: 0, + array_rows: None, } } }, @@ -178,6 +195,7 @@ fn plan_key_batch<'p>( pks: document_keys(std::slice::from_ref(key_value)), binds: false, fresh: 0, + array_rows: None, }, SqlPlan::Delete { collection, @@ -194,6 +212,7 @@ fn plan_key_batch<'p>( pks: document_keys(target_keys), binds: false, fresh: 0, + array_rows: None, }, }, SqlPlan::VectorPrimaryDelete(VectorPrimaryDeletePlan { @@ -210,6 +229,7 @@ fn plan_key_batch<'p>( pks: document_keys(target_keys), binds: false, fresh: 0, + array_rows: None, }, // Every timeseries row takes a fresh identity. SqlPlan::TimeseriesIngest(TimeseriesIngestPlan { @@ -219,18 +239,26 @@ fn plan_key_batch<'p>( pks: Vec::new(), binds: true, fresh: rows.len(), + array_rows: None, }, - SqlPlan::InsertArray(InsertArrayPlan { name, rows }) => KeyBatch { - collection: name, - pks: super::super::array_convert::insert_array_cell_pks( + SqlPlan::InsertArray(InsertArrayPlan { name, rows }) => { + let keys = super::super::array_convert::insert_array_cell_keys( name, rows, ctx.tenant_id, ctx, - )?, - binds: true, - fresh: 0, - }, + )?; + KeyBatch { + collection: name, + pks: keys.pks, + binds: true, + fresh: 0, + array_rows: Some(ArrayBatchRows { + rows, + coords: keys.coords, + }), + } + } _ => return Ok(None), }; Ok((!batch.pks.is_empty() || batch.fresh > 0).then_some(batch)) @@ -257,6 +285,7 @@ fn row_identity_batch<'p>( .collect(), binds: true, fresh: fresh_identity_count(primary_key, declared.as_deref(), rows), + array_rows: None, }) } diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs b/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs index a2e40b80f..a7859ca08 100644 --- a/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs +++ b/nodedb/src/control/planner/sql_plan_convert/expr/mod.rs @@ -1,11 +1,9 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Expression conversion and CTE inlining. +//! Expression conversion. mod bridge_expr; -mod inline_cte; mod sort_keys; pub(super) use bridge_expr::{sql_expr_to_bridge_expr, sql_expr_to_bridge_expr_qualified}; -pub(super) use inline_cte::inline_cte; pub(super) use sort_keys::convert_sort_keys; diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/build/aggregate.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/build/aggregate.rs index 0baf38682..4fbeaaed1 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/build/aggregate.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/build/aggregate.rs @@ -4,40 +4,63 @@ use super::super::columns::{column_types_for, group_by_key_column}; use crate::control::planner::sql_plan_convert::lateral::collection_name_from_plan; -use crate::control::planner::sql_plan_convert::output_schema_types::infer_aggregate_type; +use crate::control::planner::sql_plan_convert::output_schema_types::{ + infer_aggregate_type, time_bucket_type, +}; use crate::control::server::response_shape::schema::{OutputColumn, OutputSchema}; use crate::control::server::response_shape::types::DdlColType; use nodedb_sql::SqlError; use nodedb_sql::catalog::SqlCatalog; -use nodedb_sql::types::SqlPlan; use nodedb_sql::types::query::{AggOutputSlot, AggregateExpr}; +use nodedb_sql::types::{SqlPlan, TimeseriesAggOutput, TimeseriesScanPlan}; use nodedb_sql::types_expr::SqlExpr; use std::collections::HashMap; -/// The aggregating (non-`time_bucket`) timeseries scan schema, grouped or -/// not: GROUP BY keys typed from the catalog, then aggregates as `Text`. +/// The schema of an aggregating timeseries scan, grouped, bucketed, or not. +/// +/// Columns follow the plan's recorded SELECT-list order, under their +/// SELECT-list names. The scan writes each group row under those names, so +/// each column reads its own name. Each column is typed exactly as +/// [`aggregate_schema`] types the generic aggregate. +/// +/// A plan with no recorded order falls back to the encoder's own order and +/// keys: `bucket`, each GROUP BY column, then each aggregate's canonical key. pub(super) fn timeseries_group_schema( catalog: &C, database_id: nodedb_types::DatabaseId, - collection: &str, - group_by: &[String], - aggregates: &[AggregateExpr], + scan: &TimeseriesScanPlan, ) -> Result { - let types = column_types_for(catalog, database_id, collection)?; - let mut columns = Vec::with_capacity(group_by.len() + aggregates.len()); - for key in group_by { - columns.push(OutputColumn { - display_name: key.clone(), - lookup_key: key.clone(), - ty: types.get(key).copied().unwrap_or(DdlColType::Text), - }); - } - for agg in aggregates { - let key = nodedb_sql::planner::agg_naming::aggregate_output_key(agg); + let types = column_types_for(catalog, database_id, &scan.collection)?; + let renamed = !scan.output.is_empty(); + let fallback; + let output: &[TimeseriesAggOutput] = if renamed { + &scan.output + } else { + fallback = encoder_order(scan); + &fallback + }; + let mut columns = Vec::with_capacity(output.len()); + for slot in output { + let Some((key, name)) = slot.key_and_name(&scan.aggregates) else { + continue; + }; + let ty = match slot { + // The bucket renders in the kind of the column it buckets. + TimeseriesAggOutput::Bucket { source, .. } => { + time_bucket_type(source.as_deref().and_then(|s| types.get(s).copied())) + } + TimeseriesAggOutput::GroupColumn { name, .. } => { + types.get(name).copied().unwrap_or(DdlColType::Text) + } + TimeseriesAggOutput::Aggregate(index) => scan + .aggregates + .get(*index) + .map_or(DdlColType::Text, |agg| infer_aggregate_type(agg, &types)), + }; columns.push(OutputColumn { - display_name: key.clone(), - lookup_key: key, - ty: DdlColType::Text, + lookup_key: if renamed { name.clone() } else { key }, + display_name: name, + ty, }); } Ok(OutputSchema { @@ -48,6 +71,35 @@ pub(super) fn timeseries_group_schema( }) } +/// The columns a native grouped timeseries scan encodes, in its own order: +/// the bucket, each GROUP BY column, then each aggregate. +fn encoder_order(scan: &TimeseriesScanPlan) -> Vec { + let bucket = (scan.bucket_interval_ms > 0).then_some(TimeseriesAggOutput::Bucket { + alias: None, + source: None, + }); + let groups = scan + .group_by + .iter() + .map(|name| TimeseriesAggOutput::GroupColumn { + name: name.clone(), + alias: None, + }); + let aggregates = (0..scan.aggregates.len()).map(TimeseriesAggOutput::Aggregate); + bucket.into_iter().chain(groups).chain(aggregates).collect() +} + +/// The collection whose catalog types the aggregate's input columns. +/// +/// A timeseries scan feeds the generic aggregate whenever the native grouped +/// scan cannot compute the query, and its rows carry the collection's types. +fn aggregate_input_collection(input: &SqlPlan) -> Option { + match input { + SqlPlan::TimeseriesScan(TimeseriesScanPlan { collection, .. }) => Some(collection.clone()), + other => collection_name_from_plan(other), + } +} + /// The fields of an `SqlPlan::Aggregate` that shape its output. pub(super) struct AggregateShape<'a> { pub input: &'a SqlPlan, @@ -72,9 +124,10 @@ pub(super) fn aggregate_schema( } = shape; // Catalog types of the aggregate's underlying columns, resolved from // the input plan's single source collection when it has one (Scan / - // point-get style). GROUP BY bare keys and MIN/MAX/SUM/AVG argument - // columns are typed against this; anything unresolvable stays Text. - let types = match collection_name_from_plan(input) { + // point-get / timeseries-scan style). GROUP BY bare keys and + // MIN/MAX/SUM/AVG argument columns are typed against this; anything + // unresolvable stays Text. + let types = match aggregate_input_collection(input) { Some(collection) => column_types_for(catalog, database_id, &collection)?, None => HashMap::new(), }; @@ -148,8 +201,7 @@ mod tests { name: "status".to_string(), }], // `SELECT status AS state ...` — the group-key output name is the - // SELECT-list alias, while the value lookup key stays the raw - // grouped column name. + // SELECT-list alias. group_by_aliases: vec![Some("state".to_string())], // Empty output_order exercises the group-keys-first fallback. output_order: Vec::new(), @@ -186,10 +238,10 @@ mod tests { build_output_schema(&plans, &NoCatalog, nodedb_types::DatabaseId::DEFAULT, None) .expect("schema derives"); assert_eq!(schema.columns.len(), 3); - // Group-key display_name is the SELECT-list alias; lookup_key stays - // the raw grouped column name (the executor's emitted value key). + // An aliased group key shows and reads its alias: the aggregate + // renames the key to it in every result row. assert_eq!(schema.columns[0].display_name, "state"); - assert_eq!(schema.columns[0].lookup_key, "status"); + assert_eq!(schema.columns[0].lookup_key, "state"); assert_eq!(schema.columns[1].display_name, "total"); assert_eq!(schema.columns[1].lookup_key, "total"); // sum(x) stays TEXT here because this test has no catalog, so the @@ -247,6 +299,113 @@ mod tests { assert_eq!(schema.columns[4].ty, DdlColType::Text); } + fn timeseries_scan( + output: Vec, + ) -> nodedb_sql::types::TimeseriesScanPlan { + nodedb_sql::types::TimeseriesScanPlan { + collection: "metrics".to_string(), + time_range: (i64::MIN, i64::MAX), + bucket_interval_ms: 3_600_000, + group_by: vec!["region".to_string()], + aggregates: vec![agg_expr("avg", vec![metrics_column("amount")], "mean")], + filters: Vec::new(), + having: Vec::new(), + projection: Vec::new(), + gap_fill: String::new(), + limit: usize::MAX, + sort_keys: Vec::new(), + tiered: false, + temporal: nodedb_sql::temporal::TemporalScope::default(), + output, + group_keys: Vec::new(), + } + } + + /// A native grouped timeseries scan announces SELECT-list order and + /// aliases. Each column reads its own name, which the encoder writes, and + /// takes the generic aggregate's type. + #[test] + fn timeseries_scan_follows_select_list_order_and_aliases() { + use nodedb_sql::types::TimeseriesAggOutput; + let plans = vec![SqlPlan::TimeseriesScan(timeseries_scan(vec![ + TimeseriesAggOutput::Aggregate(0), + TimeseriesAggOutput::Bucket { + alias: Some("b".to_string()), + source: Some("n".to_string()), + }, + TimeseriesAggOutput::GroupColumn { + name: "region".to_string(), + alias: Some("r".to_string()), + }, + ]))]; + let schema = build_output_schema( + &plans, + &TypedCatalog, + nodedb_types::DatabaseId::DEFAULT, + None, + ) + .expect("schema derives"); + let shape: Vec<(&str, &str, DdlColType)> = schema + .columns + .iter() + .map(|c| (c.display_name.as_str(), c.lookup_key.as_str(), c.ty)) + .collect(); + assert_eq!( + shape, + [ + ("mean", "mean", DdlColType::Float8), + ("b", "b", DdlColType::Int8), + ("r", "r", DdlColType::Text), + ] + ); + } + + /// With no recorded order the scan announces its encoder's order: the + /// bucket, each GROUP BY column, then each aggregate. + #[test] + fn timeseries_scan_without_select_list_uses_encoder_order() { + let plans = vec![SqlPlan::TimeseriesScan(timeseries_scan(Vec::new()))]; + let schema = build_output_schema( + &plans, + &TypedCatalog, + nodedb_types::DatabaseId::DEFAULT, + None, + ) + .expect("schema derives"); + let names: Vec<&str> = schema + .columns + .iter() + .map(|c| c.display_name.as_str()) + .collect(); + assert_eq!(names, ["bucket", "region", "mean"]); + } + + /// The generic aggregate over a timeseries scan types its columns from + /// the scanned collection. + #[test] + fn aggregate_over_a_timeseries_scan_takes_catalog_types() { + let plans = vec![SqlPlan::Aggregate { + input: Box::new(SqlPlan::TimeseriesScan(timeseries_scan(Vec::new()))), + group_by: vec![metrics_column("region")], + group_by_aliases: vec![None], + output_order: Vec::new(), + aggregates: vec![agg_expr("avg", vec![metrics_column("amount")], "mean")], + having: Vec::new(), + limit: 0, + grouping_sets: None, + sort_keys: Vec::new(), + }]; + let schema = build_output_schema( + &plans, + &TypedCatalog, + nodedb_types::DatabaseId::DEFAULT, + None, + ) + .expect("schema derives"); + assert_eq!(schema.columns[1].display_name, "mean"); + assert_eq!(schema.columns[1].ty, DdlColType::Float8); + } + /// A computed GROUP BY key (`UPPER(region)`) is not a bare column, so it /// defaults to Text. #[test] diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/build/schema.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/build/schema.rs index cd2858c1e..f2af6f037 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/build/schema.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/build/schema.rs @@ -12,16 +12,18 @@ use std::collections::HashMap; use nodedb_sql::SqlError; use nodedb_sql::catalog::SqlCatalog; +use nodedb_sql::resolver::array_tvf::ARRAY_HORIZON_COLUMN; use nodedb_sql::types::query::Projection; use nodedb_sql::types::{ - ArrayProjectPlan, ArraySlicePlan, CtePlan, DocumentIndexLookupPlan, HybridSearchPlan, - HybridSearchTriplePlan, InsertPlan, KvInsertPlan, LateralLoopPlan, LateralTopKPlan, MergePlan, - RangeScanPlan, RecursiveScanPlan, RecursiveValuePlan, SqlPlan, TextSearchPlan, - TimeseriesIngestPlan, TimeseriesScanPlan, UpsertPlan, VectorPrimaryDeletePlan, - VectorPrimaryInsertPlan, VectorPrimaryUpdatePlan, + CtePlan, DocumentIndexLookupPlan, HybridSearchPlan, HybridSearchTriplePlan, InsertPlan, + KvInsertPlan, LateralLoopPlan, LateralTopKPlan, MergePlan, RangeScanPlan, RecursiveScanPlan, + RecursiveValuePlan, SqlPlan, TextSearchPlan, TimeseriesIngestPlan, TimeseriesScanPlan, + UpsertPlan, VectorPrimaryDeletePlan, VectorPrimaryInsertPlan, VectorPrimaryUpdatePlan, }; -use crate::control::server::response_shape::schema::{OutputColumn, OutputSchema}; +use crate::control::server::response_shape::schema::{ + OutputColumn, OutputSchema, sql_data_type_to_ddl_col_type_with_width, +}; use crate::control::server::response_shape::types::DdlColType; use super::super::columns::{ @@ -54,27 +56,16 @@ pub fn build_output_schema( }; Ok(match plan { - // A grouped or aggregating timeseries scan announces the columns its - // aggregate encoder emits, in that encoder's order: each GROUP BY key, - // then each aggregate. An ungrouped aggregate takes this arm too: the - // raw-scan arm below announces no column for it, and its cells then - // lose the SELECT-list order. A GROUP BY key carries its own catalog - // type, so one stored instant renders the same grouped as it does - // through a plain `SELECT`. An aggregate result stays `Text`: the timeseries plan - // carries no SELECT-list alias for it, so its type cannot be resolved - // with certainty. - // - // A `time_bucket` query is excluded: its encoder prepends a `bucket` - // boundary column that the plan's GROUP BY list does not name, so the - // announced shape would not be the emitted one. - SqlPlan::TimeseriesScan(TimeseriesScanPlan { - collection, - group_by, - aggregates, - bucket_interval_ms, - .. - }) if (!group_by.is_empty() || !aggregates.is_empty()) && *bucket_interval_ms == 0 => { - timeseries_group_schema(catalog, database_id, collection, group_by, aggregates)? + // A grouped, bucketed, or aggregating timeseries scan announces its + // SELECT-list columns in SELECT-list order, under their aliases, typed + // as the generic aggregate types them. An ungrouped aggregate takes + // this arm too: the raw-scan arm below announces no column for it. + SqlPlan::TimeseriesScan(scan) + if !scan.group_by.is_empty() + || !scan.aggregates.is_empty() + || scan.bucket_interval_ms > 0 => + { + timeseries_group_schema(catalog, database_id, scan)? } SqlPlan::Scan { collection, @@ -219,12 +210,15 @@ pub fn build_output_schema( SqlPlan::Cte(CtePlan { definitions, outer }) => { let stored = CatalogRef(catalog); let overlay = CteTypesCatalog::build(&stored, definitions, database_id)?; - build_output_schema( + let schema = build_output_schema( std::slice::from_ref(outer.as_ref()), &overlay, database_id, returning, - )? + )?; + // An outer query that runs as a body unchanged reads each column + // under the key the body's rows carry it. + overlay.read_through(outer, schema) } // A post-processor's projected shape is its outer projection over the // body's announced columns. An empty projection (SELECT *) inherits @@ -263,24 +257,14 @@ pub fn build_output_schema( let types = HashMap::new(); schema_from_projection(projection, &types, &[]) } - SqlPlan::ArraySlice(ArraySlicePlan { - attr_projection, .. - }) - | SqlPlan::ArrayProject(ArrayProjectPlan { - attr_projection, .. - }) => OutputSchema { - columns: attr_projection - .iter() - .map(|name| OutputColumn { - display_name: name.clone(), - lookup_key: name.clone(), - ty: DdlColType::Text, - }) - .collect(), - is_star: false, - cp_computed: Vec::new(), - declared_key: None, - }, + // A cell read announces the columns its rows carry: every dim, then + // the requested attrs, each typed from the array schema. + SqlPlan::ArraySlice(slice) => { + let view = array_view(catalog, &slice.name)?; + array_columns_schema(&nodedb_sql::resolver::array_tvf::array_slice_columns( + &view, slice, + )) + } // A write announces exactly what its `RETURNING` clause projects, from // the target collection's declared columns. `RETURNING` is a // projection, so it is typed like one: a `SELECT ts, host, v` and an @@ -318,16 +302,69 @@ pub fn build_output_schema( | SqlPlan::DropIndex { .. } | SqlPlan::ArrayFlush { .. } | SqlPlan::ArrayCompact { .. } => OutputSchema::default(), - // `ArrayAgg` / `ArrayElementwise` compile to `ArrayOp::Aggregate` / - // `ArrayOp::Elementwise`, which `describe_plan` classifies as - // `PlanKind::MultiRow`. `MultiRow` responses are shaped by - // `shape_generic_rows` -> `shape_decoded_rows`, which derives column - // names from the decoded JSON payload itself, not from - // `OutputSchema`. An empty schema here is therefore correct. - SqlPlan::ArrayAgg { .. } | SqlPlan::ArrayElementwise { .. } => OutputSchema::default(), + // An aggregate announces `group`, typed as its group dim the way a + // slice types that dim, then `result`, typed by the reducer. A + // temporal aggregate ends with a summary row, so it also announces + // the row's `truncated_before_horizon` flag. + SqlPlan::ArrayAgg(agg) => { + let view = array_view(catalog, &agg.name)?; + let mut schema = array_columns_schema( + &nodedb_sql::resolver::array_tvf::array_agg_columns(&view, agg)?, + ); + if agg.temporal.is_temporal() { + schema.columns.push(OutputColumn { + display_name: ARRAY_HORIZON_COLUMN.to_string(), + lookup_key: ARRAY_HORIZON_COLUMN.to_string(), + ty: DdlColType::Bool, + }); + } + schema + } + // `ArrayElementwise` compiles to `ArrayOp::Elementwise`, which + // `describe_plan` classifies as `PlanKind::MultiRow`. `MultiRow` + // responses are shaped by `shape_generic_rows` -> + // `shape_decoded_rows`, which derives column names from the decoded + // payload when the schema is empty. + SqlPlan::ArrayElementwise { .. } => OutputSchema::default(), }) } +/// The array `name` names, or the unknown-table error a plan over a dropped +/// array reports. +fn array_view( + catalog: &C, + name: &str, +) -> Result { + catalog + .lookup_array(name) + .ok_or_else(|| SqlError::UnknownTable { + name: name.to_string(), + }) +} + +/// One output column per relation column, in order, under its own name and +/// its wire type. A narrow attr announces its own width: `int2`, `int4`, or +/// `float4`. +fn array_columns_schema(columns: &[nodedb_sql::types::ColumnInfo]) -> OutputSchema { + OutputSchema { + columns: columns + .iter() + .map(|column| OutputColumn { + display_name: column.name.clone(), + lookup_key: column.name.clone(), + ty: sql_data_type_to_ddl_col_type_with_width( + &column.data_type, + column.int_width, + column.float_width, + ), + }) + .collect(), + is_star: false, + cp_computed: Vec::new(), + declared_key: None, + } +} + #[cfg(test)] mod tests { use super::super::fixtures::{ diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/columns.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/columns.rs index 02edf0fd9..ef3e8f518 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/columns.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/columns.rs @@ -162,55 +162,52 @@ pub(super) fn column_types_for( /// /// The `display_name` is the SELECT-list output name: the explicit alias when /// the projection aliased the key (`SELECT k AS label ... GROUP BY k` yields -/// output column `label`, matching Postgres), otherwise the key's own column -/// name. The `lookup_key` always stays the raw grouped column name (the key -/// the aggregate executor emits the value under), so `project_row` still finds -/// the value. +/// output column `label`, matching Postgres), otherwise the key's own name. +/// +/// The aggregate renames an aliased key to its alias in every result row, so +/// an aliased key reads its alias. An unaliased key reads its row name: the +/// raw grouped column name, qualified when the key is, or for a computed key +/// the shared index-based `computed_group_key_name` (`group_0`), the exact +/// name the aggregate spec emits the evaluated value under. /// /// The output type is the grouped column's catalog type when the key is a bare /// column (resolved from `types`, default `Text`); a computed-expression key is /// typed conservatively via [`infer_computed_expr_type`], defaulting to `Text`. -/// -/// A non-`Column` GROUP BY key (a computed expression) derives its `lookup_key` -/// from the shared index-based `computed_group_key_name` rule — the exact name -/// the aggregate spec emits the evaluated value under, so the two can never -/// diverge. Its `display_name` is the SELECT-list alias when present -/// (`UPPER(label) AS u` shows column `u`), else the same placeholder. pub(super) fn group_by_key_column( expr: &SqlExpr, index: usize, alias: Option<&str>, types: &HashMap, ) -> OutputColumn { - match expr { - SqlExpr::Column { table, name } => { - let lookup_key = match table { + let (row_name, display_name, ty) = match expr { + SqlExpr::Column { table, name } => ( + match table { Some(t) => format!("{t}.{name}"), None => name.clone(), - }; - let display_name = alias.map(str::to_string).unwrap_or_else(|| name.clone()); - let ty = types.get(name).copied().unwrap_or(DdlColType::Text); - OutputColumn { - display_name, - lookup_key, - ty, - } - } + }, + name.clone(), + types.get(name).copied().unwrap_or(DdlColType::Text), + ), _ => { - // The executor emits the evaluated value under the shared - // index-based name (see `group_by_to_specs`), so `lookup_key` MUST - // equal it. `display_name` is the SELECT-list alias when present - // (`UPPER(label) AS u` shows column `u`), else the same placeholder. - let lookup_key = computed_group_key_name(index); - let display_name = alias - .map(str::to_string) - .unwrap_or_else(|| lookup_key.clone()); - OutputColumn { - display_name, - lookup_key, - ty: infer_computed_expr_type(expr, types), - } + let row_name = computed_group_key_name(index); + ( + row_name.clone(), + row_name, + infer_computed_expr_type(expr, types), + ) } + }; + match alias { + Some(alias) => OutputColumn { + display_name: alias.to_string(), + lookup_key: alias.to_string(), + ty, + }, + None => OutputColumn { + display_name, + lookup_key: row_name, + ty, + }, } } diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/cte_types.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/cte_types.rs index a596dea24..ede48c606 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/cte_types.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/cte_types.rs @@ -10,7 +10,7 @@ //! body gives it, exactly as a plain `SELECT` of the body does. use nodedb_sql::SqlError; -use nodedb_sql::catalog::{SqlCatalog, SqlCatalogError}; +use nodedb_sql::catalog::{ArrayCatalogView, SqlCatalog, SqlCatalogError}; use nodedb_sql::types::query::EngineType; use nodedb_sql::types::{CollectionInfo, ColumnInfo, SqlPlan}; use nodedb_sql::types_expr::SqlDataType; @@ -25,8 +25,8 @@ use super::build::build_output_schema; /// /// Schema derivation recurses through nested CTE plans. Each level reads its /// catalog as `&dyn SqlCatalog`, so the recursion instantiates a finite set -/// of generic functions. Schema derivation reads relations only, so only the -/// relation lookups are passed through. +/// of generic functions. Schema derivation reads relations and arrays, so +/// those lookups are passed through. pub(super) struct CatalogRef<'a, C: ?Sized>(pub(super) &'a C); impl SqlCatalog for CatalogRef<'_, C> { @@ -45,6 +45,10 @@ impl SqlCatalog for CatalogRef<'_, C> { ) -> Result, SqlCatalogError> { self.0.resolve_relation(database_id, name) } + + fn lookup_array(&self, name: &str) -> Option { + self.0.lookup_array(name) + } } /// A catalog that answers for each CTE definition name first. @@ -54,6 +58,8 @@ impl SqlCatalog for CatalogRef<'_, C> { pub(super) struct CteTypesCatalog<'a> { inner: &'a dyn SqlCatalog, relations: Vec<(String, CollectionInfo)>, + /// Each definition's body output schema, by definition name. + bodies: Vec<(String, OutputSchema)>, } impl<'a> CteTypesCatalog<'a> { @@ -70,16 +76,49 @@ impl<'a> CteTypesCatalog<'a> { let mut overlay = Self { inner, relations: Vec::with_capacity(definitions.len()), + bodies: Vec::with_capacity(definitions.len()), }; for (name, body) in definitions { let schema = build_output_schema(std::slice::from_ref(body), &overlay, database_id, None)?; + let schema = overlay.read_through(body, schema); overlay .relations .push((name.clone(), relation_from_schema(name, &schema))); + overlay.bodies.push((name.clone(), schema)); } Ok(overlay) } + + /// `schema`, the schema of `plan`, with each column reading the key the + /// rows carry. + /// + /// A plan that reads a definition body unchanged runs as that body, so + /// its rows are the body's rows. A body column can be announced under its + /// display name but carried under another lookup key: a join body + /// carries `o.id` under the qualified `o.id` and announces it as `id`. + /// Each column that names a body column by display name reads that body + /// column's lookup key. Any other plan returns `schema` unchanged. + pub(super) fn read_through(&self, plan: &SqlPlan, mut schema: OutputSchema) -> OutputSchema { + let Some((_, body)) = self + .bodies + .iter() + .find(|(name, _)| nodedb_sql::planner::cte::reads_cte_body_unchanged(plan, name)) + else { + return schema; + }; + for column in &mut schema.columns { + let bare = column + .lookup_key + .rsplit('.') + .next() + .unwrap_or(column.lookup_key.as_str()); + if let Some(source) = body.columns.iter().find(|c| c.display_name == bare) { + column.lookup_key = source.lookup_key.clone(); + } + } + schema + } } impl SqlCatalog for CteTypesCatalog<'_> { @@ -93,6 +132,10 @@ impl SqlCatalog for CteTypesCatalog<'_> { } self.inner.get_collection(database_id, name) } + + fn lookup_array(&self, name: &str) -> Option { + self.inner.lookup_array(name) + } } /// The relation one definition body exposes: one column per announced output diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema/join_types.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/join_types.rs index 45d63e17f..d37b6abdd 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema/join_types.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema/join_types.rs @@ -165,7 +165,6 @@ fn collect_sides( | SqlPlan::InsertArray { .. } | SqlPlan::DeleteArray { .. } | SqlPlan::ArraySlice { .. } - | SqlPlan::ArrayProject { .. } | SqlPlan::ArrayAgg { .. } | SqlPlan::ArrayElementwise { .. } | SqlPlan::ArrayFlush { .. } diff --git a/nodedb/src/control/planner/sql_plan_convert/output_schema_types.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema_types.rs index ccd314714..7f3c5324b 100644 --- a/nodedb/src/control/planner/sql_plan_convert/output_schema_types.rs +++ b/nodedb/src/control/planner/sql_plan_convert/output_schema_types.rs @@ -62,7 +62,7 @@ pub fn infer_computed_expr_type(expr: &SqlExpr, types: &HashMap DdlColType::Text, }, - SqlExpr::BinaryOp { op, .. } => match op { + SqlExpr::BinaryOp { left, op, right } => match op { // Comparison and logical operators yield a boolean. BinaryOp::Eq | BinaryOp::Ne @@ -72,28 +72,83 @@ pub fn infer_computed_expr_type(expr: &SqlExpr, types: &HashMap DdlColType::Bool, - // Arithmetic promotion (int vs float vs numeric) is too subtle to - // type safely here; `||` is string concatenation. All stay Text. - BinaryOp::Add - | BinaryOp::Sub - | BinaryOp::Mul - | BinaryOp::Div - | BinaryOp::Mod - | BinaryOp::Concat => DdlColType::Text, + // Integer and float operands evaluate in `f64` and yield an + // integer when the result is whole, else a float. Integer `+`, + // `-`, `*`, `%` stay whole, so they are `Int8`. Integer `/` + // yields `5` for `10 / 2` and `3.5` for `7 / 2`. `Numeric` + // renders both as the evaluator produced them. A float operand + // makes the result `Float8`. A numeric (decimal) or text operand + // stays Text. + BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div | BinaryOp::Mod => { + match ( + numeric_class(infer_computed_expr_type(left, types)), + numeric_class(infer_computed_expr_type(right, types)), + ) { + (Some(NumericClass::Integer), Some(NumericClass::Integer)) => match op { + BinaryOp::Div => DdlColType::Numeric, + _ => DdlColType::Int8, + }, + (Some(_), Some(_)) => DdlColType::Float8, + _ => DdlColType::Text, + } + } + // `||` is string concatenation. + BinaryOp::Concat => DdlColType::Text, }, - SqlExpr::UnaryOp { op, .. } => match op { + SqlExpr::UnaryOp { op, expr } => match op { UnaryOp::Not => DdlColType::Bool, - // `-x` preserves a numeric operand's type; anything else is Text. - UnaryOp::Neg => DdlColType::Text, + // `-x` keeps an integer or float operand's class. + UnaryOp::Neg => match numeric_class(infer_computed_expr_type(expr, types)) { + Some(NumericClass::Integer) => DdlColType::Int8, + Some(NumericClass::Float) => DdlColType::Float8, + None => DdlColType::Text, + }, }, SqlExpr::IsNull { .. } | SqlExpr::InList { .. } | SqlExpr::Between { .. } | SqlExpr::Like { .. } => DdlColType::Bool, + // `time_bucket` keeps the type of the column it buckets. + SqlExpr::Function { name, args, .. } if name.eq_ignore_ascii_case("time_bucket") => { + let mut columns = args.iter().filter(|a| matches!(a, SqlExpr::Column { .. })); + match (columns.next(), columns.next()) { + (Some(column), None) => time_bucket_type(bare_column_type(column, types)), + _ => DdlColType::Text, + } + } + // A math function computes in `f64` and yields an integer when the + // result is whole, else a float: `Float8` renders both. Its first + // argument must be an integer or a float. + SqlExpr::Function { name, args, .. } if is_float_math_function(name) => { + match args + .first() + .and_then(|arg| numeric_class(infer_computed_expr_type(arg, types))) + { + Some(_) => DdlColType::Float8, + None => DdlColType::Text, + } + } + // A CASE yields one of its branches. When every non-NULL branch has + // one type, that is the CASE's type. + SqlExpr::Case { + when_then, + else_expr, + .. + } => { + let mut branch_types = when_then + .iter() + .map(|(_, then)| then) + .chain(else_expr.as_deref()) + .filter(|branch| !matches!(branch, SqlExpr::Literal(SqlValue::Null))) + .map(|branch| infer_computed_expr_type(branch, types)); + match branch_types.next() { + Some(first) if branch_types.all(|ty| ty == first) => first, + Some(_) | None => DdlColType::Text, + } + } // String functions, casts, subqueries, wildcards, array literals and // any other function call are not certainly non-TEXT — stay Text. SqlExpr::Function { .. } - | SqlExpr::Case { .. } | SqlExpr::Cast { .. } | SqlExpr::Subquery(_) | SqlExpr::Wildcard @@ -101,6 +156,54 @@ pub fn infer_computed_expr_type(expr: &SqlExpr, types: &HashMap Option { + match ty { + DdlColType::Int8 | DdlColType::Int4 | DdlColType::Int2 => Some(NumericClass::Integer), + DdlColType::Float8 | DdlColType::Float4 => Some(NumericClass::Float), + DdlColType::Text + | DdlColType::Varchar + | DdlColType::Bool + | DdlColType::Bytea + | DdlColType::Json + | DdlColType::Jsonb + | DdlColType::Timestamp + | DdlColType::Timestamptz + | DdlColType::Numeric + | DdlColType::Uuid + | DdlColType::Float4Array + | DdlColType::Float8Array => None, + } +} + +/// Whether `name` is a math function that computes in `f64`. +fn is_float_math_function(name: &str) -> bool { + const FLOAT_MATH: &[&str] = &[ + "abs", "round", "ceil", "ceiling", "floor", "power", "pow", "sqrt", "mod", "sign", "log", + "ln", "log10", "log2", "exp", + ]; + FLOAT_MATH.iter().any(|f| name.eq_ignore_ascii_case(f)) +} + +/// The result type of a `time_bucket` call over a column of type `source`. +/// +/// A bucket of an instant is an instant of the same kind, and a bucket of a +/// `BIGINT` is a `BIGINT`. Any other source, or none, is `Text`. +pub fn time_bucket_type(source: Option) -> DdlColType { + match source { + Some(ty @ (DdlColType::Timestamp | DdlColType::Timestamptz | DdlColType::Int8)) => ty, + Some(_) | None => DdlColType::Text, + } +} + /// Infers the result type of an aggregate expression. /// /// - `COUNT(...)` is always Postgres `bigint` (`Int8`). @@ -115,7 +218,15 @@ pub fn infer_aggregate_type( ) -> DdlColType { let arg_bare_type = || agg.args.first().and_then(|a| bare_column_type(a, types)); match agg.function.as_str() { - "count" => DdlColType::Int8, + // A HyperLogLog estimate finalizes to a rounded integer. + "count" | "approx_count_distinct" => DdlColType::Int8, + // A t-digest quantile, an interpolated percentile, and every + // Welford dispersion finalize to a float, or NULL. + "approx_percentile" | "percentile_cont" | "stddev" | "stddev_pop" | "stddev_samp" + | "variance" | "var_pop" | "var_samp" | "ts_stddev" => DdlColType::Float8, + // The space-saving top-k finalizes to an array of + // `{item, count, error}` objects. + "approx_topk" => DdlColType::Json, "min" | "max" => arg_bare_type().unwrap_or(DdlColType::Text), "sum" | "avg" => match arg_bare_type() { Some(DdlColType::Float8) | Some(DdlColType::Float4) => DdlColType::Float8, @@ -159,6 +270,89 @@ mod tests { ); } + /// `time_bucket` takes the type of the column it buckets when that is an + /// instant or a `BIGINT`, in either argument order. Otherwise it is Text. + #[test] + fn time_bucket_keeps_its_column_type() { + let mut t = types_map(); + t.insert("ts".to_string(), DdlColType::Timestamptz); + let bucket = |args: Vec| SqlExpr::Function { + name: "time_bucket".to_string(), + args, + distinct: false, + }; + let interval = || SqlExpr::Literal(SqlValue::String("1 hour".into())); + assert_eq!( + infer_computed_expr_type(&bucket(vec![interval(), column("ts")]), &t), + DdlColType::Timestamptz + ); + assert_eq!( + infer_computed_expr_type(&bucket(vec![column("n"), interval()]), &t), + DdlColType::Int8 + ); + assert_eq!( + infer_computed_expr_type(&bucket(vec![interval(), column("amount")]), &t), + DdlColType::Text + ); + } + + /// Math functions and arithmetic over integer and float operands take a + /// numeric type. A CASE takes its branches' one type. + #[test] + fn numeric_expressions_take_a_numeric_type() { + let t = types_map(); + let call = |name: &str, args: Vec| SqlExpr::Function { + name: name.to_string(), + args, + distinct: false, + }; + let binary = |left: SqlExpr, op: BinaryOp, right: SqlExpr| SqlExpr::BinaryOp { + left: Box::new(left), + op, + right: Box::new(right), + }; + let one = || SqlExpr::Literal(SqlValue::Int(1)); + assert_eq!( + infer_computed_expr_type(&call("round", vec![column("amount"), one()]), &t), + DdlColType::Float8 + ); + assert_eq!( + infer_computed_expr_type(&call("round", vec![column("label"), one()]), &t), + DdlColType::Text + ); + assert_eq!( + infer_computed_expr_type(&binary(column("n"), BinaryOp::Add, one()), &t), + DdlColType::Int8 + ); + assert_eq!( + infer_computed_expr_type(&binary(column("amount"), BinaryOp::Div, column("n")), &t), + DdlColType::Float8 + ); + assert_eq!( + infer_computed_expr_type(&binary(column("n"), BinaryOp::Div, column("n")), &t), + DdlColType::Numeric + ); + assert_eq!( + infer_computed_expr_type(&binary(column("n"), BinaryOp::Mod, column("n")), &t), + DdlColType::Int8 + ); + let case = SqlExpr::Case { + operand: None, + when_then: vec![( + binary(column("n"), BinaryOp::Gt, one()), + SqlExpr::Literal(SqlValue::String("many".into())), + )], + else_expr: Some(Box::new(SqlExpr::Literal(SqlValue::String("one".into())))), + }; + assert_eq!(infer_computed_expr_type(&case, &t), DdlColType::Text); + let mixed = SqlExpr::Case { + operand: None, + when_then: vec![(binary(column("n"), BinaryOp::Gt, one()), column("amount"))], + else_expr: Some(Box::new(SqlExpr::Literal(SqlValue::Null))), + }; + assert_eq!(infer_computed_expr_type(&mixed, &t), DdlColType::Float8); + } + #[test] fn computed_comparison_is_bool() { let t = types_map(); diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/core.rs b/nodedb/src/control/planner/sql_plan_convert/scan/core.rs index 3e648fb23..cc7fbaebd 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/core.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/core.rs @@ -50,7 +50,10 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_scan( if is_provider_relation(collection, engine) { let filter_bytes = serialize_filters(filters)?; - let proj_names = extract_projection_names(projection, window_functions); + // The provider tail evaluates computed columns before it prunes, so + // the pruning list keeps each computed alias the output schema reads. + let proj_names = + super::super::set_ops::lower_subquery_projection(projection, window_functions)?; let sort = convert_sort_keys(sort_keys); return Ok(vec![PhysicalTask { tenant_id, @@ -103,6 +106,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_scan( valid_at_ms: valid_at_from_scope(temporal), having: Vec::new(), tiering: None, + output_columns: Vec::new(), }) } EngineType::Columnar => PhysicalPlan::Columnar(ColumnarOp::Scan { diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/search.rs b/nodedb/src/control/planner/sql_plan_convert/scan/search.rs index 5c5fdffa4..3e3b86a71 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/search.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/search.rs @@ -96,6 +96,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_sparse_search( field_name: p.field.to_string(), query_entries: p.query_entries.to_vec(), top_k: *p.top_k, + filters: serialize_filters(p.filters)?, }), post_set_op: PostSetOp::None, txn_id: None, @@ -133,7 +134,6 @@ fn build_array_prefilter_plan( tenant_id: TenantId, ctx: &super::super::convert::ConvertContext, ) -> crate::Result { - use nodedb_array::query::slice::{DimRange, Slice}; use nodedb_array::schema::ArraySchema; use nodedb_array::types::ArrayId; @@ -161,24 +161,12 @@ fn build_array_prefilter_plan( detail: format!("array schema decode: {e}"), })?; - let mut dim_ranges: Vec> = vec![None; schema.dims.len()]; - for r in &prefilter.slice.dim_ranges { - let idx = schema - .dims - .iter() - .position(|d| d.name == r.dim) - .ok_or_else(|| crate::Error::PlanError { - detail: format!( - "array prefilter: array '{}' has no dim '{}'", - prefilter.array_name, r.dim - ), - })?; - let dtype = schema.dims[idx].dtype; - let lo = super::super::array_fn_convert::helpers::coerce_bound(&r.lo, dtype, &r.dim)?; - let hi = super::super::array_fn_convert::helpers::coerce_bound(&r.hi, dtype, &r.dim)?; - dim_ranges[idx] = Some(DimRange::new(lo, hi)); - } - let slice = Slice::new(dim_ranges); + let slice = super::super::array_fn_convert::helpers::resolve_slice( + "array prefilter", + &prefilter.array_name, + &prefilter.slice, + &schema, + )?; let slice_msgpack = zerompk::to_msgpack_vec(&slice).map_err(|e| crate::Error::Serialization { format: "msgpack".into(), diff --git a/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs b/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs index 0b9edb8c8..9e933c4d2 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan/timeseries.rs @@ -37,7 +37,21 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_timeseries_scan( tenant_id, ctx, temporal, + output, } = p; + // The scan writes each group row under the SELECT-list names, in order. + let output_columns: Vec<(String, String)> = output + .iter() + .map(|column| { + column + .key_and_name(aggregates) + .ok_or(crate::Error::Internal { + detail: format!( + "timeseries output column {column:?} names no aggregate of the scan" + ), + }) + }) + .collect::>()?; let collection_key = nodedb_types::CollectionKey::from_bare(ctx.database_id, collection); let qualified_collection = nodedb_types::QualifiedCollection::new(ctx.database_id, collection); let filter_bytes = serialize_filters(filters)?; @@ -87,6 +101,7 @@ pub(in crate::control::planner::sql_plan_convert) fn convert_timeseries_scan( valid_at_ms: valid_at_from_scope(temporal), having: having_bytes, tiering, + output_columns, }), post_set_op: PostSetOp::None, txn_id: None, diff --git a/nodedb/src/control/planner/sql_plan_convert/scan_params.rs b/nodedb/src/control/planner/sql_plan_convert/scan_params.rs index ab1fed614..955e5daca 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan_params.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan_params.rs @@ -85,6 +85,8 @@ pub(super) struct TimeseriesScanParams<'a> { pub tenant_id: TenantId, pub ctx: &'a ConvertContext, pub temporal: &'a nodedb_sql::TemporalScope, + /// SELECT-list output columns of an aggregating scan. Empty on a raw scan. + pub output: &'a [nodedb_sql::types::TimeseriesAggOutput], } /// Parameters for `convert_vector_search`. @@ -119,6 +121,9 @@ pub(super) struct SparseSearchParams<'a> { /// Query sparse vector as `(dimension, weight)` entries, parsed at plan time. pub query_entries: &'a [(u32, f32)], pub top_k: &'a usize, + /// The statement's residual `WHERE`, lowered to + /// `VectorOp::SparseSearch::filters`. + pub filters: &'a [Filter], pub tenant_id: TenantId, pub database_id: crate::types::DatabaseId, } @@ -185,7 +190,7 @@ pub(super) struct SpatialScanParams<'a> { pub collection: &'a str, pub field: &'a str, pub predicate: &'a nodedb_sql::types::SpatialPredicate, - pub query_geometry: &'a nodedb_types::geometry::Geometry, + pub query_geometry: &'a nodedb_types::geometry::SridGeometry, pub distance_meters: &'a f64, pub attribute_filters: &'a [Filter], pub limit: &'a usize, diff --git a/nodedb/src/control/planner/sql_plan_convert/set_ops.rs b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs index 14e9eb07e..411681518 100644 --- a/nodedb/src/control/planner/sql_plan_convert/set_ops.rs +++ b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs @@ -10,7 +10,6 @@ use nodedb_physical::physical_plan::*; use super::body::convert_body_to_single_plan; use super::convert::{ConvertContext, convert_one}; -use super::expr::inline_cte; use super::value::{sql_value_to_nodedb_value, sql_value_to_string}; use nodedb_physical::physical_task::{PhysicalTask, PostSetOp}; use nodedb_types::Value; @@ -226,7 +225,7 @@ pub(super) fn convert_cte( // earlier definition, and the passes that follow resolve that name. let mut resolved = outer.clone(); for (name, cte_plan) in definitions.iter().rev() { - resolved = inline_cte(&resolved, name, cte_plan); + resolved = nodedb_sql::planner::cte::inline_cte(&resolved, name, cte_plan); } convert_one(&resolved, tenant_id, ctx) } @@ -366,7 +365,10 @@ fn is_merged_doc_body(plan: &PhysicalPlan) -> bool { /// A Control-Plane-computed item keeps the base columns its expression reads, /// not its alias: the Control Plane evaluates it once the tail returns, then /// drops those columns. -fn lower_subquery_projection( +/// +/// A `ProviderScan` runs the same tail over its materialized rows, so it +/// prunes its columns by this list too. +pub(super) fn lower_subquery_projection( projection: &[Projection], window_functions: &[WindowSpec], ) -> crate::Result> { diff --git a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/bind.rs b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/bind.rs index f92706054..3aaec2a73 100644 --- a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/bind.rs +++ b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/bind.rs @@ -29,7 +29,11 @@ pub async fn convert_bound( ctx: &mut ConvertContext, ) -> crate::Result> { ctx.prefetched = prefetch_plan_surrogates(plans, ctx).await?; - convert_resolving_misses(plans, tenant_id, ctx).await + let converted = convert_resolving_misses(plans, tenant_id, ctx).await; + // The kept array keys are keyed by the address of `plans`' rows, which + // stays valid only while `plans` is borrowed here. + ctx.prefetched.forget_array_keys(); + converted } /// Convert each of `plans` with the answers `ctx` holds, resolving the misses @@ -105,7 +109,12 @@ mod tests { /// A cluster-mode assigner whose reserved batch is empty. fn cluster_assigner() -> Arc { - let credentials = Arc::new(CredentialStore::new().expect("in-memory credential store")); + cluster_assigner_on(Arc::new( + CredentialStore::new().expect("in-memory credential store"), + )) + } + + fn cluster_assigner_on(credentials: Arc) -> Arc { let registry = Arc::new(RwLock::new(SurrogateRegistry::from_persisted_cluster(0, 0))); let wal: Arc = Arc::new(NoopWalAppender); Arc::new(SurrogateAssigner::new(registry, credentials, wal)) @@ -214,28 +223,34 @@ mod tests { } /// A home that binds each new key to the next value from 500, after a - /// real await point. - #[derive(Default)] + /// real await point. As the authority contract requires, it keeps every + /// answer in the asking node's catalog. struct CountingHome { bound: Mutex, Surrogate>>, + local: Arc, } impl HomeSurrogateAuthority for CountingHome { fn assign<'a>( &'a self, - _key: CollectionKey<'a>, - _tenant_id: nodedb_types::TenantId, + key: CollectionKey<'a>, + tenant_id: nodedb_types::TenantId, pks: &'a [&'a [u8]], ) -> BoxFuture<'a, crate::Result>> { async move { tokio::task::yield_now().await; - let mut bound = self.bound.lock().unwrap_or_else(|p| p.into_inner()); - let mut answers = Vec::with_capacity(pks.len()); - for pk in pks { - let next = Surrogate::new(500 + bound.len() as u32); - answers.push(*bound.entry(pk.to_vec()).or_insert(next)); - } - Ok(answers) + let answers: Vec = { + let mut bound = self.bound.lock().unwrap_or_else(|p| p.into_inner()); + pks.iter() + .map(|pk| { + let next = Surrogate::new(500 + bound.len() as u32); + *bound.entry(pk.to_vec()).or_insert(next) + }) + .collect() + }; + let rows: Vec<(&[u8], Surrogate)> = pks.iter().copied().zip(answers).collect(); + let kept = self.local.catalog().put_surrogates(key, tenant_id, &rows)?; + Ok(kept.into_iter().map(|binding| binding.surrogate).collect()) } .boxed() } @@ -260,8 +275,12 @@ mod tests { /// the home's values. #[tokio::test(flavor = "current_thread")] async fn missed_keys_are_bound_at_their_home_without_blocking() { - let assigner = cluster_assigner(); - assigner.install_home_authority(Arc::new(CountingHome::default())); + let credentials = Arc::new(CredentialStore::new().expect("in-memory credential store")); + let assigner = cluster_assigner_on(Arc::clone(&credentials)); + assigner.install_home_authority(Arc::new(CountingHome { + bound: Mutex::new(HashMap::new()), + local: credentials, + })); let mut ctx = cluster_ctx(&assigner); let plans = vec![SqlPlan::Update { collection: "cache".to_string(), diff --git a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/cache.rs b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/cache.rs index f5c6df2bc..df10e6be8 100644 --- a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/cache.rs +++ b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/cache.rs @@ -12,6 +12,9 @@ use std::collections::{HashMap, HashSet}; use std::sync::Mutex; use nodedb_types::{CollectionKey, DatabaseId, Surrogate}; +use rustc_hash::FxHashMap; + +use super::super::array_convert::ArrayCellKeys; /// One collection: its database and bare catalog name. type CollectionSlot = (DatabaseId, String); @@ -19,13 +22,28 @@ type CollectionSlot = (DatabaseId, String); /// A fresh identity drawn ahead: its surrogate and its identity string. type FreshIdentity = (Surrogate, String); +/// One `INSERT INTO ARRAY` row batch: its array, and the address and length +/// of the plan's rows. The bind step borrows the plans for its whole run, so +/// the address names one batch until [`PrefetchedSurrogates::forget_array_keys`] +/// runs at its end. +type ArrayRowsSlot = (CollectionSlot, usize, usize); + +/// One collection's bind answers: primary key to its bound surrogate, or +/// `None` when the key's home binds none. +type KeyAnswers = FxHashMap, Option>; + /// The bind step's answer for each `(collection, primary key)`: the bound /// surrogate, or `None` when the key's home binds none. Also the fresh /// identities drawn ahead for rows that name no key, and the misses of the /// running conversion pass. All hold for the one plan batch they belong to. #[derive(Debug, Default)] pub struct PrefetchedSurrogates { - answers: HashMap<(CollectionSlot, Vec), Option>, + /// Keyed database, then collection, then key, so a lookup borrows its + /// key and allocates nothing. + answers: FxHashMap>, + /// The coerced coordinates and identity keys of each `INSERT INTO ARRAY` + /// row batch the bind step read, so conversion encodes neither again. + array_keys: HashMap, /// Conversion reads through a shared reference, so the fresh pool and /// the misses sit behind locks. fresh: Mutex, @@ -106,15 +124,63 @@ impl SurrogateMisses { impl PrefetchedSurrogates { /// Keep the answer for `pk` in `key`. pub fn insert(&mut self, key: CollectionKey<'_>, pk: &[u8], answer: Option) { - self.answers.insert((slot_of(key), pk.to_vec()), answer); + let by_name = self.answers.entry(key.database_id()).or_default(); + // The collection's map exists after its first key, so later keys + // allocate no name. + if let Some(keys) = by_name.get_mut(key.name()) { + match keys.get_mut(pk) { + Some(slot) => *slot = answer, + None => { + keys.insert(pk.to_vec(), answer); + } + } + return; + } + let mut keys = KeyAnswers::default(); + keys.insert(pk.to_vec(), answer); + by_name.insert(key.name().to_string(), keys); } /// The answer for `pk` in `key`. `None` when the key has no answer. pub fn get(&self, key: CollectionKey<'_>, pk: &[u8]) -> Option> { - if self.answers.is_empty() { + self.answers + .get(&key.database_id())? + .get(key.name())? + .get(pk) + .copied() + } + + /// Keep the coerced coordinates and identity keys the bind step read off + /// the `INSERT INTO ARRAY` rows at `rows` in `key`. + pub(in crate::control::planner::sql_plan_convert) fn keep_array_keys( + &mut self, + key: CollectionKey<'_>, + rows: &[T], + keys: ArrayCellKeys, + ) { + self.array_keys.insert(array_rows_slot(key, rows), keys); + } + + /// The coordinates and identity keys kept for the rows at `rows` in + /// `key`. `None` when the bind step kept none, or kept them for a batch + /// of another length. + pub(in crate::control::planner::sql_plan_convert) fn array_keys( + &self, + key: CollectionKey<'_>, + rows: &[T], + ) -> Option<&ArrayCellKeys> { + if self.array_keys.is_empty() { return None; } - self.answers.get(&(slot_of(key), pk.to_vec())).copied() + self.array_keys + .get(&array_rows_slot(key, rows)) + .filter(|keys| keys.pks.len() == rows.len()) + } + + /// Drop every kept array row batch. Runs when the bind step ends, before + /// the plans whose addresses key them are released. + pub(in crate::control::planner::sql_plan_convert) fn forget_array_keys(&mut self) { + self.array_keys.clear(); } /// The bound surrogate of `pk` in `key`, when the answer binds one. @@ -194,6 +260,10 @@ fn slot_of(key: CollectionKey<'_>) -> CollectionSlot { (key.database_id(), key.name().to_string()) } +fn array_rows_slot(key: CollectionKey<'_>, rows: &[T]) -> ArrayRowsSlot { + (slot_of(key), rows.as_ptr() as usize, rows.len()) +} + #[cfg(test)] mod tests { use nodedb_types::{CollectionKey, DatabaseId, Surrogate}; @@ -262,6 +332,36 @@ mod tests { ); } + #[test] + fn array_keys_are_kept_per_row_batch_until_forgotten() { + use nodedb_array::types::coord::value::CoordValue; + + use super::ArrayCellKeys; + + let scene = CollectionKey::from_bare(DatabaseId::DEFAULT, "scene"); + let rows = vec![1u8, 2]; + let other_rows = vec![1u8, 2]; + let mut prefetched = PrefetchedSurrogates::default(); + prefetched.keep_array_keys( + scene, + &rows, + ArrayCellKeys { + coords: vec![vec![CoordValue::Int64(1)], vec![CoordValue::Int64(2)]], + pks: vec![b"a".to_vec(), b"b".to_vec()], + surrogates: vec![Surrogate::new(4), Surrogate::new(5)], + }, + ); + + let kept = prefetched.array_keys(scene, &rows).expect("kept"); + assert_eq!(kept.pks[1], b"b".to_vec()); + assert_eq!(kept.surrogates[1], Surrogate::new(5)); + assert!(prefetched.array_keys(scene, &other_rows).is_none()); + assert!(prefetched.array_keys(scene, &rows[..1]).is_none()); + + prefetched.forget_array_keys(); + assert!(prefetched.array_keys(scene, &rows).is_none()); + } + #[test] fn misses_are_kept_per_collection_until_taken() { let users = CollectionKey::from_bare(DatabaseId::DEFAULT, "users"); diff --git a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/resolve.rs b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/resolve.rs index fc049fceb..865f90ce8 100644 --- a/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/resolve.rs +++ b/nodedb/src/control/planner/sql_plan_convert/surrogate_prefetch/resolve.rs @@ -6,10 +6,11 @@ use std::sync::Arc; use nodedb_sql::types::SqlPlan; -use nodedb_types::CollectionKey; +use nodedb_types::{CollectionKey, Surrogate}; +use super::super::array_convert::ArrayCellKeys; use super::super::convert::ConvertContext; -use super::super::dml::plan_key_batches; +use super::super::dml::{ArrayBatchRows, KeyBatch, plan_key_batches}; use super::cache::{PrefetchedSurrogates, SurrogateMisses}; use crate::control::surrogate::SurrogateAssigner; use crate::types::TenantId; @@ -33,19 +34,49 @@ pub async fn prefetch_plan_surrogates( let resolver = Resolver::new(ctx)?; let mut prefetched = PrefetchedSurrogates::default(); for batch in plan_key_batches(plans, ctx) { - let pks: Vec<&[u8]> = batch.pks.iter().map(Vec::as_slice).collect(); - let (binds, lookups) = if batch.binds { - (pks, Vec::new()) + let KeyBatch { + collection, + pks, + binds: binding, + fresh, + array_rows, + } = batch; + let key = ctx.collection_key(collection); + // An array batch keeps its answers by row, so conversion reads each + // cell's surrogate by index and no answer enters the key map. + let keep = if array_rows.is_some() { + Keep::ByPosition } else { - (Vec::new(), pks) + Keep::ByKey }; - let request = KeyRequest { - key: ctx.collection_key(batch.collection), - binds: &binds, - lookups: &lookups, - fresh: batch.fresh, + let surrogates = { + let refs: Vec<&[u8]> = pks.iter().map(Vec::as_slice).collect(); + let (binds, lookups) = if binding { + (refs, Vec::new()) + } else { + (Vec::new(), refs) + }; + let request = KeyRequest { + key, + binds: &binds, + lookups: &lookups, + fresh, + }; + resolver.resolve(request, keep, &mut prefetched).await? }; - resolver.resolve(request, &mut prefetched).await?; + // An array batch's coordinates and keys stay for conversion, which + // then neither coerces nor encodes a coordinate again. + if let Some(ArrayBatchRows { rows, coords }) = array_rows { + prefetched.keep_array_keys( + key, + rows, + ArrayCellKeys { + coords, + pks, + surrogates, + }, + ); + } } Ok(prefetched) } @@ -65,11 +96,22 @@ pub(super) async fn resolve_misses( lookups: &lookups, fresh: recorded.fresh, }; - resolver.resolve(request, &mut ctx.prefetched).await?; + resolver + .resolve(request, Keep::ByKey, &mut ctx.prefetched) + .await?; } Ok(()) } +/// Where a resolve keeps the surrogates its bound keys answer. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Keep { + /// In the answer map, for conversion to look up by key. + ByKey, + /// Returned in key order, for the caller to keep by row. + ByPosition, +} + /// The keys of one collection to resolve. struct KeyRequest<'r> { key: CollectionKey<'r>, @@ -104,11 +146,15 @@ impl Resolver { }) } + /// Resolve `request` into `into`. With [`Keep::ByPosition`] the bound + /// keys' surrogates return in `binds` order and enter no map; otherwise + /// the return is empty. async fn resolve( &self, request: KeyRequest<'_>, + keep: Keep, into: &mut PrefetchedSurrogates, - ) -> crate::Result<()> { + ) -> crate::Result> { let KeyRequest { key, binds, @@ -118,7 +164,8 @@ impl Resolver { if self.metadata { // A metadata plan binds nothing: its binding keys are lookups. self.look_up(key, binds, into).await?; - return self.look_up(key, lookups, into).await; + self.look_up(key, lookups, into).await?; + return Ok(Vec::new()); } if fresh > 0 { let drawn = self @@ -129,16 +176,23 @@ impl Resolver { into.push_fresh(key, identity); } } + let mut kept = Vec::new(); if !binds.is_empty() { let bound = self .assigner .assign_many(key, self.tenant_id, binds) .await?; - for (pk, surrogate) in binds.iter().zip(bound) { - into.insert(key, pk, Some(surrogate)); + match keep { + Keep::ByPosition => kept = bound, + Keep::ByKey => { + for (pk, surrogate) in binds.iter().zip(bound) { + into.insert(key, pk, Some(surrogate)); + } + } } } - self.look_up(key, lookups, into).await + self.look_up(key, lookups, into).await?; + Ok(kept) } async fn look_up( diff --git a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_aggregate_lateral.rs b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_aggregate_lateral.rs index 1bd356e52..b50d72be1 100644 --- a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_aggregate_lateral.rs +++ b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_aggregate_lateral.rs @@ -16,6 +16,7 @@ macro_rules! impl_aggregate_lateral_arms_for_convert_visitor { limit, grouping_sets, sort_keys, + group_by_aliases, } = args; super::super::aggregate::convert_aggregate( super::super::aggregate::ConvertAggregateParams { @@ -26,6 +27,7 @@ macro_rules! impl_aggregate_lateral_arms_for_convert_visitor { limit, grouping_sets, sort_keys, + group_by_aliases, tenant_id: self.tenant_id, ctx: self.ctx, }, diff --git a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_array.rs b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_array.rs index 9b0011e15..090ee1be2 100644 --- a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_array.rs +++ b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_array.rs @@ -104,36 +104,11 @@ macro_rules! impl_array_arms_for_convert_visitor { ) } - fn array_project( - &mut self, - name: &str, - attr_projection: &[String], - ) -> crate::Result> { - super::super::array_fn_convert::convert_project( - name, - attr_projection, - self.tenant_id, - self.ctx, - ) - } - fn array_agg( &mut self, - name: &str, - attr: &str, - reducer: &nodedb_sql::types_array::ArrayReducerAst, - group_by_dim: Option<&str>, - temporal: &nodedb_sql::temporal::TemporalScope, + args: nodedb_sql::ArrayAggVisitArgs<'_>, ) -> crate::Result> { - super::super::array_fn_convert::convert_agg( - name, - attr, - *reducer, - group_by_dim, - *temporal, - self.tenant_id, - self.ctx, - ) + super::super::array_fn_convert::convert_agg(args, self.tenant_id, self.ctx) } fn array_elementwise( @@ -142,12 +117,16 @@ macro_rules! impl_array_arms_for_convert_visitor { right: &str, op: nodedb_sql::types_array::ArrayBinaryOpAst, attr: &str, + temporal: &nodedb_sql::temporal::TemporalScope, ) -> crate::Result> { super::super::array_fn_convert::convert_elementwise( - left, - right, - op, - attr, + super::super::array_fn_convert::elementwise::ElementwiseCall { + left_name: left, + right_name: right, + op, + attr, + temporal: *temporal, + }, self.tenant_id, self.ctx, ) diff --git a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_scan_search.rs b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_scan_search.rs index 3aa2de05e..5338be813 100644 --- a/nodedb/src/control/planner/sql_plan_convert/visitor/arms_scan_search.rs +++ b/nodedb/src/control/planner/sql_plan_convert/visitor/arms_scan_search.rs @@ -22,6 +22,7 @@ macro_rules! impl_scan_search_arms_for_convert_visitor { sort_keys, tiered, temporal, + output, } = args; super::super::scan::convert_timeseries_scan( super::super::scan_params::TimeseriesScanParams { @@ -40,6 +41,7 @@ macro_rules! impl_scan_search_arms_for_convert_visitor { tenant_id: self.tenant_id, ctx: self.ctx, temporal, + output, }, ) } @@ -101,6 +103,7 @@ macro_rules! impl_scan_search_arms_for_convert_visitor { field: &str, query_entries: &[(u32, f32)], top_k: usize, + filters: &[nodedb_sql::types::Filter], ) -> crate::Result> { super::super::scan::convert_sparse_search( super::super::scan_params::SparseSearchParams { @@ -108,6 +111,7 @@ macro_rules! impl_scan_search_arms_for_convert_visitor { field, query_entries, top_k: &top_k, + filters, tenant_id: self.tenant_id, database_id: self.ctx.database_id, }, diff --git a/nodedb/src/control/security/audit/chain.rs b/nodedb/src/control/security/audit/chain.rs new file mode 100644 index 000000000..9b8eab063 --- /dev/null +++ b/nodedb/src/control/security/audit/chain.rs @@ -0,0 +1,204 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Hash-chain verification over the whole audit log. +//! +//! Each entry's `prev_hash` must equal the hash of the entry before it. +//! Retention pruning removes the oldest entries, so the oldest surviving entry +//! links to a deleted one. Each prune records an `AuditCheckpoint` entry first. +//! Its detail names the first surviving seq and the hash of the last deleted entry. +//! The checkpoint is itself a normal chain link, so its detail is tamper-evident. + +use std::collections::HashMap; + +use super::entry::{AuditEntry, hash_entry}; +use super::event::AuditEvent; + +/// Detail token naming the first seq that survives a prune. +const OLDEST_KEPT_SEQ: &str = "oldest_kept_seq="; +/// Detail token naming the hash of the last entry a prune deletes. +const ANCHOR_HASH: &str = "anchor_hash="; + +/// Build the detail of an `AuditCheckpoint` entry for one prune. +pub fn checkpoint_detail(pruned_count: u64, oldest_kept_seq: u64, anchor_hash: &str) -> String { + format!( + "pruned_count={pruned_count} {OLDEST_KEPT_SEQ}{oldest_kept_seq} {ANCHOR_HASH}{anchor_hash}" + ) +} + +/// The `(oldest_kept_seq, anchor_hash)` pair a checkpoint entry records. +fn checkpoint_anchor(entry: &AuditEntry) -> Option<(u64, &str)> { + if entry.event != AuditEvent::AuditCheckpoint { + return None; + } + let mut seq = None; + let mut hash = None; + for token in entry.detail.split_whitespace() { + if let Some(value) = token.strip_prefix(OLDEST_KEPT_SEQ) { + seq = value.parse::().ok(); + } else if let Some(value) = token.strip_prefix(ANCHOR_HASH) { + hash = Some(value); + } + } + Some((seq?, hash?)) +} + +/// Outcome of one chain verification. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ChainReport { + /// Entries in the requested seq range whose link passed. + pub checked: u64, + /// First seq in the range whose link failed. + /// + /// A link fails when an entry's `prev_hash` differs from its predecessor's hash. + /// So an altered row is this seq or the seq before it. + /// For the newest entry, a mismatch with the live tail hash also fails. + pub broken_at: Option, + /// Hash of the last entry that passed, or empty when none passed. + pub last_hash: String, +} + +/// Verify the chain over `entries` and report on the seqs in `from_seq..=to_seq`. +/// +/// `entries` must be the whole log in ascending seq order with one entry per seq. +/// Entries before `from_seq` still feed the link of the first entry in range. +/// `tail_hash` is the live log's last hash. +/// It covers the newest entry, which has no successor to detect an edit to it. +pub fn verify_chain( + entries: &[AuditEntry], + from_seq: u64, + to_seq: u64, + tail_hash: &str, +) -> ChainReport { + let anchors: HashMap = entries.iter().filter_map(checkpoint_anchor).collect(); + let mut report = ChainReport { + checked: 0, + broken_at: None, + last_hash: String::new(), + }; + let mut expected: Option = None; + + for entry in entries { + if entry.seq > to_seq { + return report; + } + let linked = match &expected { + Some(prev) => entry.prev_hash == *prev, + // The oldest entry is either the genesis entry or a pruned head with a recorded anchor. + None => { + (entry.seq == 1 && entry.prev_hash.is_empty()) + || anchors.get(&entry.seq) == Some(&entry.prev_hash.as_str()) + } + }; + let hash = hash_entry(entry); + if entry.seq >= from_seq { + if !linked { + report.broken_at = Some(entry.seq); + return report; + } + report.checked += 1; + report.last_hash.clone_from(&hash); + } + expected = Some(hash); + } + + if let Some(newest) = entries.last() + && newest.seq >= from_seq + && report.last_hash != tail_hash + { + report.broken_at = Some(newest.seq); + } + report +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::security::audit::AuditLog; + + fn chain_of(details: &[&str]) -> (Vec, String) { + let mut log = AuditLog::new(1_000); + for detail in details { + log.record(AuditEvent::AdminAction, None, "src", detail); + } + ( + log.all().iter().cloned().collect(), + log.last_hash().to_string(), + ) + } + + #[test] + fn intact_chain_is_valid() { + let (entries, tail) = chain_of(&["a", "b", "c"]); + let report = verify_chain(&entries, 1, u64::MAX, &tail); + assert_eq!(report.broken_at, None); + assert_eq!(report.checked, 3); + assert_eq!(report.last_hash, tail); + } + + #[test] + fn altered_middle_row_breaks_at_its_successor() { + let (mut entries, tail) = chain_of(&["a", "b", "c", "d"]); + entries[1].detail = "tampered".to_string(); + let report = verify_chain(&entries, 1, u64::MAX, &tail); + assert_eq!(report.broken_at, Some(3)); + } + + #[test] + fn altered_prev_hash_breaks_at_that_row() { + let (mut entries, tail) = chain_of(&["a", "b", "c"]); + entries[1].prev_hash = "0".repeat(64); + let report = verify_chain(&entries, 1, u64::MAX, &tail); + assert_eq!(report.broken_at, Some(2)); + } + + #[test] + fn altered_newest_row_breaks_against_the_tail_hash() { + let (mut entries, tail) = chain_of(&["a", "b", "c"]); + entries[2].detail = "tampered".to_string(); + let report = verify_chain(&entries, 1, u64::MAX, &tail); + assert_eq!(report.broken_at, Some(3)); + } + + #[test] + fn deleted_row_breaks_at_the_next_row() { + let (mut entries, tail) = chain_of(&["a", "b", "c", "d"]); + entries.remove(1); + let report = verify_chain(&entries, 1, u64::MAX, &tail); + assert_eq!(report.broken_at, Some(3)); + } + + #[test] + fn unanchored_truncated_head_is_broken() { + let (entries, tail) = chain_of(&["a", "b", "c"]); + let report = verify_chain(&entries[1..], 1, u64::MAX, &tail); + assert_eq!(report.broken_at, Some(2)); + } + + #[test] + fn pruned_head_with_checkpoint_is_valid() { + let mut log = AuditLog::new(1_000); + log.record(AuditEvent::AdminAction, None, "src", "a"); + log.record(AuditEvent::AdminAction, None, "src", "b"); + log.record(AuditEvent::AdminAction, None, "src", "c"); + let anchor = hash_entry(&log.all()[1]); + log.record( + AuditEvent::AuditCheckpoint, + None, + "retention", + &checkpoint_detail(2, 3, &anchor), + ); + log.record(AuditEvent::AdminAction, None, "src", "d"); + let survivors: Vec = log.all().iter().skip(2).cloned().collect(); + let report = verify_chain(&survivors, 1, u64::MAX, log.last_hash()); + assert_eq!(report.broken_at, None); + assert_eq!(report.checked, 3); + } + + #[test] + fn range_uses_earlier_entries_for_its_first_link() { + let (entries, tail) = chain_of(&["a", "b", "c", "d"]); + let report = verify_chain(&entries, 3, 3, &tail); + assert_eq!(report.broken_at, None); + assert_eq!(report.checked, 1); + } +} diff --git a/nodedb/src/control/security/audit/emitter.rs b/nodedb/src/control/security/audit/emitter.rs index b0ef284ac..6e2834466 100644 --- a/nodedb/src/control/security/audit/emitter.rs +++ b/nodedb/src/control/security/audit/emitter.rs @@ -61,11 +61,18 @@ impl AuditEmitter for ArcAuditEmitter { user_name: ctx.auth_user_name.to_string(), session_id: String::new(), }; - let mut log = match self.0.lock() { - Ok(l) => l, - Err(p) => p.into_inner(), - }; - log.record_with_auth(event, ctx.tenant_id, ctx.database_id, source, detail, &auth); + // The sync wait runs with the audit lock released, so concurrent emitters share one fsync. + super::log_commit::record_concurrent( + &self.0, + super::log::AuditRecordInput { + event, + tenant_id: ctx.tenant_id, + database_id: ctx.database_id, + source, + detail, + auth: &auth, + }, + ); } } diff --git a/nodedb/src/control/security/audit/export_sink.rs b/nodedb/src/control/security/audit/export_sink.rs new file mode 100644 index 000000000..93af1d054 --- /dev/null +++ b/nodedb/src/control/security/audit/export_sink.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Export target every recorded audit entry is handed to, whichever path recorded it. + +use super::entry::AuditEntry; +use crate::control::security::siem::SiemExporter; + +/// Export target for recorded audit entries. +/// +/// `AuditLog` copies an entry for export only while [`Self::is_active`] holds. +/// So an inactive sink costs one check per entry and no clone. +/// The copies reach [`Self::export`] after the audit lock is released. +pub trait AuditExportSink: Send + Sync { + /// Whether the sink takes entries now. + fn is_active(&self) -> bool; + /// Take one recorded entry. + fn export(&self, entry: AuditEntry); +} + +impl AuditExportSink for SiemExporter { + fn is_active(&self) -> bool { + self.is_configured() + } + + fn export(&self, entry: AuditEntry) { + if entry.event.is_auth_event() { + self.push_auth(entry); + } else { + self.push_audit(entry); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use super::*; + use crate::control::security::audit::emitter::{ + ArcAuditEmitter, AuditEmitContext, AuditEmitter, + }; + use crate::control::security::audit::{AuditEvent, AuditLog, lock_audit_log}; + use crate::control::security::siem::SiemConfig; + + fn configured_siem() -> Arc { + Arc::new(SiemExporter::new(SiemConfig { + destinations: vec!["splunk".into()], + ..SiemConfig::default() + })) + } + + fn audit_with(siem: &Arc) -> Arc> { + let audit = Arc::new(Mutex::new(AuditLog::new(100))); + lock_audit_log(&audit).set_export_sink(Arc::clone(siem) as Arc); + audit + } + + #[test] + fn direct_site_entries_reach_siem() { + let siem = configured_siem(); + let audit = audit_with(&siem); + + ArcAuditEmitter(Arc::clone(&audit)).emit( + AuditEvent::PermissionDenied, + "emitter_site", + "denied", + AuditEmitContext::new(None, "u1", "erin"), + ); + lock_audit_log(&audit).record(AuditEvent::SessionRevoked, None, "bus_site", "revoked"); + + let mut sources: Vec = siem.drain_audit().into_iter().map(|e| e.source).collect(); + sources.extend(siem.drain_auth().into_iter().map(|e| e.source)); + sources.sort(); + assert_eq!( + sources, + vec!["bus_site".to_string(), "emitter_site".to_string()] + ); + } + + #[test] + fn dormant_siem_takes_no_copy() { + let siem = Arc::new(SiemExporter::new(SiemConfig::default())); + let audit = audit_with(&siem); + lock_audit_log(&audit).record(AuditEvent::AdminAction, None, "src", "dormant"); + assert!(siem.drain_audit().is_empty()); + assert!( + audit + .lock() + .expect("audit") + .take_export_pending() + .0 + .is_empty() + ); + } +} diff --git a/nodedb/src/control/security/audit/flush_control.rs b/nodedb/src/control/security/audit/flush_control.rs new file mode 100644 index 000000000..4f444af34 --- /dev/null +++ b/nodedb/src/control/security/audit/flush_control.rs @@ -0,0 +1,222 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Early-flush signal, flush serialization, and recorder backpressure for the audit buffer. +//! +//! Only the flush executor writes the buffer to the catalog. +//! A recorder never runs a flush itself, so a recorder never opens a catalog transaction. +//! A recorder past the backpressure point requests a flush and waits for it. +//! The wait has a bound, and it holds no lock. +//! So a recorder that holds any lock or transaction the flush needs cannot deadlock it. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Condvar, Mutex, MutexGuard}; +use std::time::Duration; + +use tokio::sync::Notify; + +/// Buffer length at which a server's audit log requests an early flush. +pub const DEFAULT_FLUSH_BOUND: usize = 10_000; + +/// Multiple of the flush bound at which a recorder waits for a flush. +/// +/// Below it, the flush executor drains the buffer after an early-flush request. +/// At or above it, a burst outruns the executor, so the recorder waits. +pub const BACKPRESSURE_MULTIPLE: usize = 2; + +/// Longest time one recorder waits for a flush. +/// +/// After it the recorder continues. The buffer keeps every entry, so the timeout loses nothing. +pub const DEFAULT_BACKPRESSURE_WAIT: Duration = Duration::from_secs(1); + +/// Shared between an `AuditLog`, its recorders, and the flush executor. +pub struct AuditFlushControl { + /// Wakes the flush executor. + signal: Notify, + /// Set by a flush request, cleared by the executor when it wakes. + requested: AtomicBool, + /// Held for the whole of one flush, so two flushes never interleave. + gate: Mutex<()>, + /// Count of finished flushes. Recorders wait for it to advance. + completed: Mutex, + completed_changed: Condvar, + /// Whether a flush executor runs. With none, a recorder has nothing to wait for. + executor_attached: AtomicBool, + /// Longest backpressure wait, in milliseconds. + backpressure_wait_ms: AtomicU64, +} + +impl Default for AuditFlushControl { + fn default() -> Self { + Self { + signal: Notify::new(), + requested: AtomicBool::new(false), + gate: Mutex::new(()), + completed: Mutex::new(0), + completed_changed: Condvar::new(), + executor_attached: AtomicBool::new(false), + backpressure_wait_ms: AtomicU64::new( + u64::try_from(DEFAULT_BACKPRESSURE_WAIT.as_millis()).unwrap_or(u64::MAX), + ), + } + } +} + +impl AuditFlushControl { + /// Ask the flush executor to flush now. The call does not block. + pub fn request_flush(&self) { + self.requested.store(true, Ordering::Release); + self.signal.notify_one(); + } + + /// Clear a pending flush request and return whether one was pending. + pub fn take_flush_request(&self) -> bool { + self.requested.swap(false, Ordering::AcqRel) + } + + /// Wait until a flush is requested. + pub async fn flush_requested(&self) { + self.signal.notified().await; + } + + /// Serialize one flush against every other flush. + /// + /// The caller must not hold the `AuditLog` lock while it waits here. + pub fn serialize(&self) -> MutexGuard<'_, ()> { + match self.gate.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } + + /// Mark a flush executor as running. + pub fn attach_executor(&self) { + self.executor_attached.store(true, Ordering::Release); + } + + /// Mark the flush executor as stopped and release every waiting recorder. + pub fn detach_executor(&self) { + self.executor_attached.store(false, Ordering::Release); + self.completed_changed.notify_all(); + } + + /// Whether a flush executor runs. + pub fn executor_attached(&self) -> bool { + self.executor_attached.load(Ordering::Acquire) + } + + /// Set the longest backpressure wait. + pub fn set_backpressure_wait(&self, wait: Duration) { + let ms = u64::try_from(wait.as_millis()).unwrap_or(u64::MAX); + self.backpressure_wait_ms.store(ms, Ordering::Release); + } + + fn completed_count(&self) -> MutexGuard<'_, u64> { + match self.completed.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } + + /// Record that one flush finished, and wake every waiting recorder. + pub fn note_flush_completed(&self) { + *self.completed_count() += 1; + self.completed_changed.notify_all(); + } + + /// Request a flush and wait for the next one to finish, up to the backpressure wait. + /// + /// Returns `true` when a flush finished in time. + /// Returns `false` on timeout, or when no executor runs. + /// The caller must hold no `AuditLog` lock and no flush gate. + pub fn await_flush(&self) -> bool { + if !self.executor_attached() { + return false; + } + let wait = Duration::from_millis(self.backpressure_wait_ms.load(Ordering::Acquire)); + let completed = self.completed_count(); + let start = *completed; + self.request_flush(); + let (completed, _) = + match self + .completed_changed + .wait_timeout_while(completed, wait, |count| { + *count == start && self.executor_attached() + }) { + Ok(result) => result, + Err(poisoned) => poisoned.into_inner(), + }; + *completed != start + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Instant; + + use super::*; + + #[test] + fn request_is_taken_once() { + let control = AuditFlushControl::default(); + assert!(!control.take_flush_request()); + control.request_flush(); + assert!(control.take_flush_request()); + assert!(!control.take_flush_request()); + } + + #[tokio::test] + async fn request_before_wait_wakes_the_waiter() { + let control = AuditFlushControl::default(); + control.request_flush(); + // `Notify` stores one permit, so a request made before the wait is not lost. + control.flush_requested().await; + } + + #[test] + fn await_without_executor_returns_at_once() { + let control = AuditFlushControl::default(); + control.set_backpressure_wait(Duration::from_secs(30)); + let start = Instant::now(); + assert!(!control.await_flush()); + assert!(start.elapsed() < Duration::from_secs(5)); + } + + #[test] + fn await_returns_when_the_executor_flushes() { + let control = Arc::new(AuditFlushControl::default()); + control.attach_executor(); + control.set_backpressure_wait(Duration::from_secs(30)); + let executor = { + let control = Arc::clone(&control); + std::thread::spawn(move || { + while !control.take_flush_request() { + std::thread::yield_now(); + } + control.note_flush_completed(); + }) + }; + assert!( + control.await_flush(), + "the waiter must see the flush finish" + ); + executor.join().expect("executor thread"); + } + + #[test] + fn await_times_out_when_the_executor_is_stuck() { + let control = AuditFlushControl::default(); + control.attach_executor(); + control.set_backpressure_wait(Duration::from_millis(50)); + let start = Instant::now(); + assert!( + !control.await_flush(), + "a stuck executor must not hold the recorder" + ); + assert!(start.elapsed() >= Duration::from_millis(50)); + assert!( + control.take_flush_request(), + "the wait must still request a flush" + ); + } +} diff --git a/nodedb/src/control/security/audit/guard.rs b/nodedb/src/control/security/audit/guard.rs new file mode 100644 index 000000000..420dbacf4 --- /dev/null +++ b/nodedb/src/control/security/audit/guard.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Lock guard for the shared `AuditLog` that exports recorded entries after it unlocks. + +use std::ops::{Deref, DerefMut}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use super::entry::AuditEntry; +use super::export_sink::AuditExportSink; +use super::log::AuditLog; + +/// Exclusive access to the shared `AuditLog`. +/// +/// On drop it moves the entries recorded meanwhile out of the log. +/// Fields drop in declaration order, so the lock releases before `export` runs. +/// So export never runs under the audit lock. +/// An entry recorded under a plain `Mutex::lock` is exported by the next drop of this guard. +pub struct AuditLogGuard<'a> { + guard: MutexGuard<'a, AuditLog>, + export: PendingExport, +} + +/// Entries waiting for export. Its drop hands them to the sink. +#[derive(Default)] +struct PendingExport { + entries: Vec, + sink: Option>, +} + +impl Drop for PendingExport { + fn drop(&mut self) { + if let Some(sink) = self.sink.take() { + for entry in self.entries.drain(..) { + sink.export(entry); + } + } + } +} + +/// Lock `audit`. A poisoned lock still yields the log, because dropping a security event is worse. +pub fn lock_audit_log(audit: &Mutex) -> AuditLogGuard<'_> { + let guard = match audit.lock() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!("audit log mutex poisoned, recovering"); + poisoned.into_inner() + } + }; + AuditLogGuard { + guard, + export: PendingExport::default(), + } +} + +impl Deref for AuditLogGuard<'_> { + type Target = AuditLog; + + fn deref(&self) -> &AuditLog { + &self.guard + } +} + +impl DerefMut for AuditLogGuard<'_> { + fn deref_mut(&mut self) -> &mut AuditLog { + &mut self.guard + } +} + +impl Drop for AuditLogGuard<'_> { + fn drop(&mut self) { + let (entries, sink) = self.guard.take_export_pending(); + self.export = PendingExport { entries, sink }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::security::audit::AuditEvent; + + /// Records each export and whether the audit lock was free at that moment. + struct ProbeSink { + audit: Arc>, + seen: Mutex>, + active: bool, + } + + impl AuditExportSink for ProbeSink { + fn is_active(&self) -> bool { + self.active + } + + fn export(&self, entry: AuditEntry) { + let unlocked = self.audit.try_lock().is_ok(); + self.seen + .lock() + .unwrap_or_else(|p| p.into_inner()) + .push((entry.seq, unlocked)); + } + } + + fn probe(active: bool) -> (Arc>, Arc) { + let audit = Arc::new(Mutex::new(AuditLog::new(100))); + let sink = Arc::new(ProbeSink { + audit: Arc::clone(&audit), + seen: Mutex::new(Vec::new()), + active, + }); + lock_audit_log(&audit).set_export_sink(Arc::clone(&sink) as Arc); + (audit, sink) + } + + #[test] + fn export_runs_after_the_lock_releases() { + let (audit, sink) = probe(true); + { + let mut log = lock_audit_log(&audit); + log.record(AuditEvent::AdminAction, None, "src", "a"); + log.record(AuditEvent::AuthSuccess, None, "src", "b"); + assert!( + sink.seen.lock().unwrap().is_empty(), + "no export under the lock" + ); + } + assert_eq!(*sink.seen.lock().unwrap(), vec![(1, true), (2, true)]); + } + + #[test] + fn plain_lock_entries_export_on_the_next_guard_drop() { + let (audit, sink) = probe(true); + audit + .lock() + .unwrap() + .record(AuditEvent::AdminAction, None, "src", "plain"); + drop(lock_audit_log(&audit)); + assert_eq!(*sink.seen.lock().unwrap(), vec![(1, true)]); + } + + #[test] + fn inactive_sink_gets_nothing() { + let (audit, sink) = probe(false); + lock_audit_log(&audit).record(AuditEvent::AdminAction, None, "src", "off"); + assert!(sink.seen.lock().unwrap().is_empty()); + assert!(audit.lock().unwrap().take_export_pending().0.is_empty()); + } +} diff --git a/nodedb/src/control/security/audit/log.rs b/nodedb/src/control/security/audit/log.rs index c0b6ec116..8d3d42824 100644 --- a/nodedb/src/control/security/audit/log.rs +++ b/nodedb/src/control/security/audit/log.rs @@ -3,49 +3,142 @@ //! Append-only in-memory `AuditLog`. use std::collections::VecDeque; +use std::sync::Arc; use nodedb_types::DatabaseId; use crate::types::TenantId; use super::auth::AuditAuth; -use super::entry::{AuditEntry, hash_entry, now_us}; +use super::entry::AuditEntry; use super::event::AuditEvent; +use super::export_sink::AuditExportSink; +use super::flush_control::{AuditFlushControl, BACKPRESSURE_MULTIPLE}; use super::level::AuditLevel; +use super::wal_sink::AuditWalSink; + +/// One audit event to record, borrowed from the caller. +pub struct AuditRecordInput<'a> { + pub event: AuditEvent, + pub tenant_id: Option, + pub database_id: Option, + pub source: &'a str, + pub detail: &'a str, + pub auth: &'a AuditAuth, +} /// Append-only audit log. /// -/// Entries are immutable once appended. The log supports bounded in-memory -/// retention with overflow to the WAL or a dedicated audit segment file. +/// Entries are immutable once appended. +/// A record call allocates seqs, chains, and appends to the WAL sink under the lock. +/// So WAL order equals seq order. +/// Its entries wait in `staged` until the WAL sync covers them, and only then are published. +/// Readers, flushes and export see published entries only, so they never see a non-durable entry. +/// The buffer holds every published entry until a flush commits it to the catalog. +/// The buffer never drops an entry to stay under its bound. +/// At `flush_bound` entries the log requests an early flush. +/// At `BACKPRESSURE_MULTIPLE * flush_bound` a recorder that can wait waits for a flush. /// /// Lives on the Control Plane (Send + Sync). pub struct AuditLog { - entries: VecDeque, + pub(super) entries: VecDeque, + /// Entries appended to the WAL sink and not yet published, in seq order. + pub(super) staged: VecDeque, /// Separate auth event stream for login/logout/MFA/deactivation. - auth_events: VecDeque, - /// Maximum entries retained in memory. - max_entries: usize, - /// Next sequence number. - next_seq: u64, - /// Total entries ever recorded (including evicted). - total_entries: u64, - /// Hash of the last recorded entry (for chain continuity). - last_hash: String, + pub(super) auth_events: VecDeque, + /// Buffer length that requests an early flush. + pub(super) flush_bound: usize, + /// Next sequence number to allocate. + pub(super) next_seq: u64, + /// Total entries ever published. + pub(super) total_entries: u64, + /// Hash of the last allocated entry, staged or published (for chain continuity). + pub(super) staged_hash: String, + /// Seq of the last published entry. + pub(super) published_seq: u64, + /// Hash of the last published entry. + pub(super) published_hash: String, /// Current audit level (controls which events are recorded). level: AuditLevel, + /// Early-flush signal and flush serialization, shared with the flush task. + pub(super) flush_control: Arc, + /// Durable append target. `None` only in a log built with no WAL. + pub(super) wal_sink: Option>, + /// `(seq, WAL position)` of each staged or buffered entry the WAL sink took, in seq order. + pub(super) wal_positions: VecDeque<(u64, u64)>, + /// Export target, such as the SIEM exporter. + pub(super) export_sink: Option>, + /// Entries published for export since the last [`Self::take_export_pending`]. + pub(super) export_pending: Vec, } impl AuditLog { - pub fn new(max_entries: usize) -> Self { + pub fn new(flush_bound: usize) -> Self { Self { - entries: VecDeque::with_capacity(max_entries.min(10_000)), + entries: VecDeque::with_capacity(flush_bound.min(10_000)), + staged: VecDeque::new(), auth_events: VecDeque::with_capacity(1_000), - max_entries, + flush_bound: flush_bound.max(1), next_seq: 1, total_entries: 0, - last_hash: String::new(), + staged_hash: String::new(), + published_seq: 0, + published_hash: String::new(), level: AuditLevel::default(), + flush_control: Arc::new(AuditFlushControl::default()), + wal_sink: None, + wal_positions: VecDeque::new(), + export_sink: None, + export_pending: Vec::new(), + } + } + + /// Append every entry recorded from now on to `sink` before the record call returns. + pub fn set_wal_sink(&mut self, sink: Arc) { + self.wal_sink = Some(sink); + } + + /// The durable append target, when one is set. + pub fn wal_sink(&self) -> Option> { + self.wal_sink.clone() + } + + /// Hand every entry recorded from now on to `sink`, while the sink is active. + pub fn set_export_sink(&mut self, sink: Arc) { + self.export_sink = Some(sink); + } + + /// Move out the entries waiting for export, with the sink they go to. + /// + /// [`AuditLogGuard`](super::guard::AuditLogGuard) calls this on drop and exports after it unlocks. + pub fn take_export_pending(&mut self) -> (Vec, Option>) { + if self.export_pending.is_empty() { + return (Vec::new(), None); } + ( + std::mem::take(&mut self.export_pending), + self.export_sink.clone(), + ) + } + + /// The early-flush signal and flush serialization for this log. + pub fn flush_control(&self) -> Arc { + Arc::clone(&self.flush_control) + } + + /// Set the buffer length that requests an early flush. A value of 0 counts as 1. + pub fn set_flush_bound(&mut self, flush_bound: usize) { + self.flush_bound = flush_bound.max(1); + } + + /// Whether the buffer is past the point where a recorder that can wait waits for a flush. + pub fn needs_backpressure(&self) -> bool { + self.entries.len() >= self.flush_bound.saturating_mul(BACKPRESSURE_MULTIPLE) + } + + /// The seq of the last published entry, or 0 when nothing is published. + pub fn last_seq(&self) -> u64 { + self.published_seq } /// Set the audit level. @@ -61,11 +154,13 @@ impl AuditLog { /// Set the next sequence number (used on startup to resume from catalog). pub fn set_next_seq(&mut self, seq: u64) { self.next_seq = seq; + self.published_seq = seq.saturating_sub(1); } /// Set the last hash (used on startup to continue the chain from catalog). pub fn set_last_hash(&mut self, hash: String) { - self.last_hash = hash; + self.published_hash.clone_from(&hash); + self.staged_hash = hash; } /// Record an audit event with hash chain (legacy — no auth context). @@ -105,7 +200,10 @@ impl AuditLog { ) } - /// Record an audit event with auth context enrichment. + /// Record an audit event with auth context enrichment (best-effort durability). + /// + /// A WAL sink error is logged and reported, and the entry stays buffered. + /// Use [`Self::record_with_auth_durable`] when the caller must abort on that error. pub fn record_with_auth( &mut self, event: AuditEvent, @@ -115,40 +213,39 @@ impl AuditLog { detail: &str, auth: &AuditAuth, ) -> u64 { - let seq = self.next_seq; - self.next_seq += 1; - let prev_hash = self.last_hash.clone(); - - let entry = AuditEntry { - seq, - timestamp_us: now_us(), - event: event.clone(), - tenant_id, - database_id, - auth_user_id: auth.user_id.clone(), - auth_user_name: auth.user_name.clone(), - session_id: auth.session_id.clone(), - source: source.to_string(), - detail: detail.to_string(), - prev_hash, - }; - - self.last_hash = hash_entry(&entry); - - // Route auth events to the separate auth_events stream. - if event.is_auth_event() { - if self.auth_events.len() >= 10_000 { - self.auth_events.pop_front(); + match self.record_with_auth_durable(event, tenant_id, database_id, source, detail, auth) { + Ok(seq) => seq, + Err(e) => { + let seq = self.last_seq(); + tracing::error!(error = %e, seq, "audit WAL append failed, entry is buffered only"); + crate::diag::audit_wal_append_failed(&e, seq); + seq } - self.auth_events.push_back(entry.clone()); } + } - if self.entries.len() >= self.max_entries { - self.entries.pop_front(); - } - self.entries.push_back(entry); - self.total_entries += 1; - seq + /// Record an audit event and append it to the WAL sink. + /// + /// The entry is buffered even when the WAL append fails. + /// The error tells the caller that the entry is not crash-safe. + pub fn record_with_auth_durable( + &mut self, + event: AuditEvent, + tenant_id: Option, + database_id: Option, + source: &str, + detail: &str, + auth: &AuditAuth, + ) -> crate::Result { + let seqs = self.record_batch_durable([AuditRecordInput { + event, + tenant_id, + database_id, + source, + detail, + auth, + }])?; + Ok(seqs.first().copied().unwrap_or_else(|| self.last_seq())) } /// Filter in-memory entries by database id. @@ -207,67 +304,49 @@ impl AuditLog { &self.auth_events } - /// Append a pre-built `AuditCheckpoint` entry directly to the log. + /// Copy the buffered entries for a persistence pass. /// - /// Used by the retention pruner to link the surviving chain head back to - /// the last deleted entry via the checkpoint's `prev_hash`. The entry - /// is assumed to have a correct `prev_hash`; no re-computation is done - /// here — the caller owns the chain at this point. - pub fn push_checkpoint(&mut self, entry: AuditEntry) { - self.last_hash = hash_entry(&entry); - if self.entries.len() >= self.max_entries { - self.entries.pop_front(); - } - self.entries.push_back(entry); - self.total_entries += 1; - self.next_seq = self.next_seq.max( - self.entries - .back() - .map(|e| e.seq + 1) - .unwrap_or(self.next_seq), - ); - } - - /// Allocate the next sequence number without recording an entry. - /// Used by the retention pruner to build the checkpoint entry before push. - pub fn allocate_seq(&mut self) -> u64 { - let seq = self.next_seq; - self.next_seq += 1; - seq - } - - /// Get the current last hash (for checkpoint `prev_hash` construction). - pub fn current_last_hash(&self) -> &str { - &self.last_hash + /// The entries stay buffered until [`Self::evict_persisted`] runs. + /// A reader that locks the log between the copy and the eviction still sees them. + pub fn snapshot_for_persistence(&self) -> Vec { + self.entries.iter().cloned().collect() } - /// Drain entries for persistence (WAL or segment file). - pub fn drain_for_persistence(&mut self) -> Vec { - self.entries.drain(..).collect() + /// Remove buffered entries with `seq <= through_seq`. + /// + /// The caller must call this only after those entries are committed to the catalog. + /// Buffered seqs ascend from front to back, so eviction stops at the first higher seq. + /// Returns the highest WAL position among the evicted entries. + /// Every WAL record at or below it is in the catalog, so the WAL can drop through it. + pub fn evict_persisted(&mut self, through_seq: u64) -> Option { + while self + .entries + .front() + .is_some_and(|entry| entry.seq <= through_seq) + { + self.entries.pop_front(); + } + let mut committed_position = None; + while let Some(&(seq, position)) = self.wal_positions.front() { + if seq > through_seq { + break; + } + committed_position = Some(position); + self.wal_positions.pop_front(); + } + committed_position } - /// Get the current last hash (for chain continuity on restart). + /// The hash of the last published entry, the tail a chain reader checks against. pub fn last_hash(&self) -> &str { - &self.last_hash - } - - /// Verify the hash chain integrity of in-memory entries. - /// Returns Ok(()) if chain is valid, Err with the broken sequence number. - pub fn verify_chain(&self) -> Result<(), u64> { - let mut expected_prev = String::new(); - for entry in &self.entries { - if entry.prev_hash != expected_prev { - return Err(entry.seq); - } - expected_prev = hash_entry(entry); - } - Ok(()) + &self.published_hash } } #[cfg(test)] mod tests { use super::*; + use crate::control::security::audit::entry::hash_entry; #[test] fn empty_log() { @@ -337,20 +416,225 @@ mod tests { } #[test] - fn bounded_eviction() { + fn bound_requests_flush_and_never_drops() { let mut log = AuditLog::new(3); - for i in 0..5 { + let total = 3 * BACKPRESSURE_MULTIPLE + 1; + for i in 0..total { log.record(AuditEvent::AuthSuccess, None, "src", &format!("entry-{i}")); } - assert_eq!(log.len(), 3); - assert_eq!(log.total_recorded(), 5); - // Oldest entries evicted. - assert_eq!(log.all()[0].detail, "entry-2"); + assert_eq!( + log.len(), + total, + "the buffer must keep every unpersisted entry" + ); + assert_eq!(log.total_recorded(), total as u64); + assert_eq!(log.all()[0].detail, "entry-0"); + assert!(log.needs_backpressure()); + assert!( + log.flush_control().take_flush_request(), + "reaching the bound must request a flush" + ); + + let through_seq = log.last_seq(); + log.evict_persisted(through_seq); + assert!(!log.needs_backpressure()); + } + + /// Captures every appended seq, and fails the append for one chosen seq. + struct RecordingSink { + appended: std::sync::Mutex>, + fail_seq: u64, + } + + impl AuditWalSink for RecordingSink { + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result> { + if entries.iter().any(|entry| entry.seq == self.fail_seq) { + return Err(crate::Error::Internal { + detail: "injected audit WAL error".into(), + }); + } + self.appended + .lock() + .unwrap_or_else(|p| p.into_inner()) + .extend(entries.iter().map(|entry| entry.seq)); + // The position is the seq times 10, so a test can tell the two apart. + Ok(entries.iter().map(|entry| entry.seq * 10).collect()) + } + + fn sync_through(&self, _position: u64) -> crate::Result<()> { + Ok(()) + } + + fn truncate_through(&self, _position: u64) -> crate::Result<()> { + Ok(()) + } + } + + /// Counts `sync_through` calls, one per sync. + #[derive(Default)] + struct SyncCountingSink { + syncs: std::sync::atomic::AtomicUsize, + appended: std::sync::atomic::AtomicUsize, + panic_on_append: bool, + } + + impl AuditWalSink for SyncCountingSink { + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result> { + if self.panic_on_append { + panic!("injected audit WAL panic"); + } + self.appended + .fetch_add(entries.len(), std::sync::atomic::Ordering::SeqCst); + Ok(entries.iter().map(|entry| entry.seq).collect()) + } + + fn sync_through(&self, _position: u64) -> crate::Result<()> { + self.syncs.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } + + fn truncate_through(&self, _position: u64) -> crate::Result<()> { + Ok(()) + } + } + + fn batch_inputs<'a>( + details: &'a [String], + auth: &'a AuditAuth, + ) -> impl Iterator> { + details.iter().map(move |detail| AuditRecordInput { + event: AuditEvent::DmlAudit, + tenant_id: None, + database_id: None, + source: "batch", + detail, + auth, + }) + } + + #[test] + fn a_64_entry_batch_issues_one_sync() { + let sink = Arc::new(SyncCountingSink::default()); + let mut log = AuditLog::new(1_000); + log.set_wal_sink(Arc::clone(&sink) as Arc); + let details: Vec = (0..64).map(|i| format!("dml-{i}")).collect(); + let auth = AuditAuth::default(); + + let seqs = log + .record_batch_durable(batch_inputs(&details, &auth)) + .expect("record batch"); + + assert_eq!( + seqs, + (1..=64).collect::>(), + "seq order follows input order" + ); + assert_eq!(sink.syncs.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(sink.appended.load(std::sync::atomic::Ordering::SeqCst), 64); + let entries = log.all(); + for pair in entries.iter().collect::>().windows(2) { + assert_eq!( + pair[1].prev_hash, + hash_entry(pair[0]), + "the batch stays chained" + ); + } + assert_eq!(log.last_hash(), hash_entry(&entries[63])); + } + + #[test] + fn a_sink_panic_leaves_the_log_untouched() { + let mut log = AuditLog::new(1_000); + log.record(AuditEvent::AdminAction, None, "src", "before"); + let (last_seq, last_hash) = (log.last_seq(), log.last_hash().to_string()); + log.set_wal_sink(Arc::new(SyncCountingSink { + panic_on_append: true, + ..SyncCountingSink::default() + }) as Arc); + let details: Vec = (0..3).map(|i| format!("dml-{i}")).collect(); + let auth = AuditAuth::default(); + + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + log.record_batch_durable(batch_inputs(&details, &auth)) + })); + + assert!(outcome.is_err(), "the sink panic must propagate"); + assert_eq!(log.len(), 1, "no entry of the batch is buffered"); + assert_eq!(log.last_seq(), last_seq, "no seq of the batch is consumed"); + assert_eq!(log.last_hash(), last_hash); + } + + #[test] + fn eviction_reports_the_highest_committed_wal_position() { + let sink = Arc::new(RecordingSink { + appended: std::sync::Mutex::new(Vec::new()), + fail_seq: u64::MAX, + }); + let mut log = AuditLog::new(100); + log.set_wal_sink(Arc::clone(&sink) as Arc); + for detail in ["a", "b", "c"] { + log.record(AuditEvent::AdminAction, None, "src", detail); + } + assert_eq!(log.evict_persisted(2), Some(20)); + assert_eq!(log.evict_persisted(2), None, "nothing new is committed"); + assert_eq!(log.evict_persisted(3), Some(30)); } #[test] - fn drain_for_persistence() { + fn every_record_method_appends_to_the_wal_sink() { + let sink = Arc::new(RecordingSink { + appended: std::sync::Mutex::new(Vec::new()), + fail_seq: u64::MAX, + }); + let mut log = AuditLog::new(100); + log.set_wal_sink(Arc::clone(&sink) as Arc); + log.record(AuditEvent::AdminAction, None, "src", "record"); + log.record_with_database(AuditEvent::DatabaseCreated, None, None, "src", "db"); + log.record_with_auth( + AuditEvent::AuthzDenied, + None, + None, + "src", + "auth", + &AuditAuth::default(), + ); + let durable = log.record_with_auth_durable( + AuditEvent::DdlChange, + None, + None, + "src", + "durable", + &AuditAuth::default(), + ); + assert_eq!(durable.ok(), Some(4)); + assert_eq!(*sink.appended.lock().unwrap(), vec![1, 2, 3, 4]); + } + + #[test] + fn wal_error_keeps_the_entry_buffered_and_reaches_a_durable_caller() { + let sink = Arc::new(RecordingSink { + appended: std::sync::Mutex::new(Vec::new()), + fail_seq: 2, + }); + let mut log = AuditLog::new(100); + log.set_wal_sink(Arc::clone(&sink) as Arc); + log.record(AuditEvent::AdminAction, None, "src", "one"); + let failed = log.record_with_auth_durable( + AuditEvent::AdminAction, + None, + None, + "src", + "two", + &AuditAuth::default(), + ); + assert!(failed.is_err(), "the WAL error must reach a durable caller"); + assert_eq!(log.len(), 2, "the entry must stay buffered"); + assert_eq!(*sink.appended.lock().unwrap(), vec![1]); + } + + #[test] + fn snapshot_keeps_entries_until_evicted() { let mut log = AuditLog::new(100); log.record( AuditEvent::SnapshotBegin, @@ -360,8 +644,23 @@ mod tests { ); log.record(AuditEvent::SnapshotEnd, None, "node-1", "snapshot done"); - let drained = log.drain_for_persistence(); - assert_eq!(drained.len(), 2); + let snapshot = log.snapshot_for_persistence(); + assert_eq!(snapshot.len(), 2); + assert_eq!(log.len(), 2, "a snapshot must not remove buffered entries"); + + // An entry recorded after the snapshot must survive eviction of the snapshot. + log.record(AuditEvent::AdminAction, None, "node-1", "late entry"); + let through_seq = snapshot.last().map(|e| e.seq).expect("non-empty snapshot"); + log.evict_persisted(through_seq); + + assert_eq!(log.len(), 1); + assert_eq!(log.all()[0].detail, "late entry"); + } + + #[test] + fn evict_persisted_on_empty_log_is_noop() { + let mut log = AuditLog::new(100); + log.evict_persisted(u64::MAX); assert!(log.is_empty()); } @@ -427,91 +726,16 @@ mod tests { } #[test] - fn verify_chain_succeeds_with_mixed_events() { + fn appends_link_each_entry_to_its_predecessor() { let mut log = AuditLog::new(100); log.record(AuditEvent::AuthSuccess, None, "src", "login"); log.record(AuditEvent::DdlChange, None, "src", "alter"); log.record(AuditEvent::AdminAction, None, "src", "op"); - assert!( - log.verify_chain().is_ok(), - "chain must be valid after normal appends" - ); - } - - #[test] - fn verify_chain_succeeds_after_audit_checkpoint() { - let mut log = AuditLog::new(100); - // Simulate three regular entries. - log.record(AuditEvent::AuthSuccess, None, "src", "e1"); - log.record(AuditEvent::AuthFailure, None, "src", "e2"); - log.record(AuditEvent::AdminAction, None, "src", "e3"); - - // Simulate what the pruner does: take hash of last-deleted entry as - // the checkpoint's prev_hash, emit the checkpoint, continue the chain. - // Here we pretend entries 1 and 2 were pruned; the last deleted hash - // was whatever e2's hash was. We simulate by draining and re-seeding. - let last_e2_hash = { - let entries = log.all(); - // entry at index 1 is seq=2 (e2). Its hash was set as last_hash - // after e2 was appended, before e3 was appended. - // Reconstruct it from entry 2 (e3's prev_hash = hash(e2)). - entries[2].prev_hash.clone() - }; - - let seq = log.allocate_seq(); - let checkpoint = AuditEntry { - seq, - timestamp_us: 1_000_000, - event: AuditEvent::AuditCheckpoint, - tenant_id: None, - database_id: None, - auth_user_id: String::new(), - auth_user_name: String::new(), - session_id: String::new(), - source: "retention".to_string(), - detail: "pruned_count=2 oldest_kept_seq=3".to_string(), - prev_hash: last_e2_hash, - }; - - // Push the checkpoint. Chain: e1 → e2 → e3 → checkpoint → ... - // Note: verify_chain on the in-memory log works on the entries slice, - // so all four must chain correctly. - // Since the checkpoint's prev_hash skips to e2's hash (simulating the - // gap), it won't verify against e3 in the same slice. This is the - // expected behaviour: verify_chain on a trimmed view (e.g. post-drain) - // that starts from the checkpoint. - let mut log2 = AuditLog::new(100); - log2.record(AuditEvent::AuditCheckpoint, None, "retention", "simulated"); - log2.record(AuditEvent::AuthSuccess, None, "src", "post-checkpoint"); - assert!( - log2.verify_chain().is_ok(), - "chain after checkpoint must be valid" - ); - // Also verify the checkpoint variant itself is correctly pushed. - log2.push_checkpoint(checkpoint); - } - - #[test] - fn push_checkpoint_advances_next_seq() { - let mut log = AuditLog::new(100); - log.record(AuditEvent::AuthSuccess, None, "src", "e1"); - let seq = log.allocate_seq(); // seq = 2 - let chk = AuditEntry { - seq, - timestamp_us: 0, - event: AuditEvent::AuditCheckpoint, - tenant_id: None, - database_id: None, - auth_user_id: String::new(), - auth_user_name: String::new(), - session_id: String::new(), - source: "retention".to_string(), - detail: String::new(), - prev_hash: log.last_hash().to_string(), - }; - log.push_checkpoint(chk); - // next_seq must be at least 3 after pushing seq=2. - let next = log.allocate_seq(); - assert!(next >= 3, "next_seq must be > checkpoint seq"); + let entries = log.all(); + assert!(entries[0].prev_hash.is_empty()); + assert_eq!(entries[1].prev_hash, hash_entry(&entries[0])); + assert_eq!(entries[2].prev_hash, hash_entry(&entries[1])); + assert_eq!(log.last_hash(), hash_entry(&entries[2])); + assert_eq!(log.last_seq(), 3); } } diff --git a/nodedb/src/control/security/audit/log_commit.rs b/nodedb/src/control/security/audit/log_commit.rs new file mode 100644 index 000000000..7076e4edd --- /dev/null +++ b/nodedb/src/control/security/audit/log_commit.rs @@ -0,0 +1,388 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Record path of the audit log: stage under the audit lock, sync outside it, then publish. +//! +//! Staging allocates seqs, chains hashes, and appends to the WAL sink under the audit lock. +//! So WAL order equals seq order. +//! The sync wait runs with the audit lock released, so concurrent recorders share one fsync. +//! Publishing moves staged entries into the buffer in seq order, once a sync covers them. +//! A covering sync makes every earlier WAL record durable too. +//! So a caller can publish every staged entry up to its own last seq. +//! +//! An entry whose WAL append or sync failed is still published, so the hash chain stays whole. +//! Its record call returns the error, and the error is reported. + +use std::sync::{Arc, Mutex}; + +use super::entry::{AuditEntry, hash_entry, now_us}; +use super::guard::lock_audit_log; +use super::log::{AuditLog, AuditRecordInput}; +use super::wal_sink::AuditWalSink; + +/// A batch appended to the WAL sink and staged, waiting for its sync. +pub struct StagedBatch { + /// Seq of each entry, in input order. + seqs: Vec, + /// Highest WAL position of the batch, when the sink took it. + sync_position: Option, + sink: Option>, + append: crate::Result<()>, +} + +impl StagedBatch { + /// Wait until the batch is durable, then return its seqs with the outcome. + /// + /// The caller must not hold the audit lock, so other recorders can stage meanwhile. + pub fn wait_durable(self) -> (Vec, crate::Result<()>) { + let synced = match (self.append, self.sync_position, self.sink) { + (Err(e), _, _) => Err(e), + (Ok(()), Some(position), Some(sink)) => sink.sync_through(position), + (Ok(()), _, _) => Ok(()), + }; + (self.seqs, synced) + } + + /// The last seq of the batch, or `None` for an empty batch. + pub fn last_seq(&self) -> Option { + self.seqs.last().copied() + } +} + +impl AuditLog { + /// Allocate, chain and append `inputs` to the WAL sink, and stage them. No sync runs here. + /// + /// The log changes only after the sink returns. + /// So a panic in the sink leaves the log as it was, with none of the batch staged. + pub fn stage_batch<'a>( + &mut self, + inputs: impl IntoIterator>, + ) -> StagedBatch { + let mut next_seq = self.next_seq; + let mut last_hash = self.staged_hash.clone(); + let mut entries = Vec::new(); + for input in inputs { + let entry = AuditEntry { + seq: next_seq, + timestamp_us: now_us(), + event: input.event, + tenant_id: input.tenant_id, + database_id: input.database_id, + auth_user_id: input.auth.user_id.clone(), + auth_user_name: input.auth.user_name.clone(), + session_id: input.auth.session_id.clone(), + source: input.source.to_string(), + detail: input.detail.to_string(), + prev_hash: last_hash, + }; + last_hash = hash_entry(&entry); + next_seq += 1; + entries.push(entry); + } + if entries.is_empty() { + return StagedBatch { + seqs: Vec::new(), + sync_position: None, + sink: None, + append: Ok(()), + }; + } + + let appended = match &self.wal_sink { + Some(sink) => sink.append_batch(&entries).map(Some), + None => Ok(None), + }; + + self.next_seq = next_seq; + self.staged_hash = last_hash; + let seqs: Vec = entries.iter().map(|entry| entry.seq).collect(); + let mut sync_position = None; + let append = match appended { + Ok(Some(positions)) => { + sync_position = positions.last().copied(); + self.wal_positions + .extend(seqs.iter().copied().zip(positions.iter().copied())); + Ok(()) + } + Ok(None) => Ok(()), + Err(e) => Err(e), + }; + self.staged.extend(entries); + StagedBatch { + seqs, + sync_position, + sink: sync_position.and(self.wal_sink.clone()), + append, + } + } + + /// Publish every staged entry with `seq <= through_seq`, in seq order. + /// + /// The caller must call this only after the sync for `through_seq` returned. + pub fn publish_through(&mut self, through_seq: u64) { + let export = self + .export_sink + .as_ref() + .is_some_and(|sink| sink.is_active()); + while self + .staged + .front() + .is_some_and(|entry| entry.seq <= through_seq) + { + let Some(entry) = self.staged.pop_front() else { + break; + }; + // Route auth events to the separate auth_events stream. + if entry.event.is_auth_event() { + if self.auth_events.len() >= 10_000 { + self.auth_events.pop_front(); + } + self.auth_events.push_back(entry.clone()); + } + if export { + self.export_pending.push(entry.clone()); + } + self.published_seq = entry.seq; + self.published_hash = hash_entry(&entry); + self.entries.push_back(entry); + self.total_entries += 1; + } + if self.entries.len() >= self.flush_bound { + self.flush_control.request_flush(); + } + } + + /// Record every input in order with one WAL sync, while the caller keeps the audit lock. + /// + /// Use this only where the lock spans more work, such as a sync session's delta push. + /// [`record_batch_concurrent`] lets concurrent recorders share the sync instead. + /// Returns the seq of each entry, in input order. + pub fn record_batch_durable<'a>( + &mut self, + inputs: impl IntoIterator>, + ) -> crate::Result> { + let staged = self.stage_batch(inputs); + let last = staged.last_seq(); + let (seqs, durable) = staged.wait_durable(); + if let Some(last) = last { + self.publish_through(last); + } + durable.map(|()| seqs) + } +} + +/// Record every input in order, with the sync wait outside the audit lock. +/// +/// Concurrent callers share one fsync. Each entry is durable before this returns `Ok`. +/// Returns the seq of each entry, in input order. +pub fn record_batch_concurrent<'a>( + audit: &Mutex, + inputs: impl IntoIterator>, +) -> crate::Result> { + let staged = lock_audit_log(audit).stage_batch(inputs); + let last = staged.last_seq(); + let (seqs, durable) = staged.wait_durable(); + if let Some(last) = last { + lock_audit_log(audit).publish_through(last); + } + durable.map(|()| seqs) +} + +/// Record one entry with [`record_batch_concurrent`], best-effort. +/// +/// A WAL error is logged and reported, and the entry stays published. Returns its seq. +pub fn record_concurrent(audit: &Mutex, input: AuditRecordInput<'_>) -> u64 { + let staged = lock_audit_log(audit).stage_batch([input]); + let seq = staged.last_seq().unwrap_or(0); + let (_, durable) = staged.wait_durable(); + lock_audit_log(audit).publish_through(seq); + if let Err(e) = durable { + tracing::error!(error = %e, seq, "audit WAL write failed, entry is in memory only"); + crate::diag::audit_wal_append_failed(&e, seq); + } + seq +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Barrier, Condvar}; + use std::time::Duration; + + use super::*; + use crate::control::security::audit::{AuditAuth, AuditEvent}; + use crate::wal::WalManager; + + fn input<'a>(detail: &'a str, auth: &'a AuditAuth) -> AuditRecordInput<'a> { + AuditRecordInput { + event: AuditEvent::AdminAction, + tenant_id: None, + database_id: None, + source: "test", + detail, + auth, + } + } + + #[test] + fn concurrent_recorders_share_fsyncs_and_keep_seq_order() { + let dir = tempfile::tempdir().expect("tempdir"); + let wal = + Arc::new(WalManager::open_for_testing(&dir.path().join("group.wal")).expect("WAL")); + let audit = Arc::new(Mutex::new(AuditLog::new(1_000))); + lock_audit_log(&audit).set_wal_sink(Arc::clone(&wal) as Arc); + let syncs_before = wal.audit_sync_count(); + let start = Arc::new(Barrier::new(32)); + + let threads: Vec<_> = (0..32) + .map(|i| { + let (audit, start) = (Arc::clone(&audit), Arc::clone(&start)); + std::thread::spawn(move || { + let auth = AuditAuth::default(); + let detail = format!("thread-{i}"); + start.wait(); + record_batch_concurrent(&audit, [input(&detail, &auth)]) + }) + }) + .collect(); + for thread in threads { + thread.join().expect("recorder thread").expect("durable"); + } + + let syncs = wal.audit_sync_count() - syncs_before; + assert!( + syncs < 32, + "32 single recorders must share fsyncs, saw {syncs}" + ); + let log = lock_audit_log(&audit); + let published: Vec = log.all().iter().map(|e| e.seq).collect(); + assert_eq!( + published, + (1..=32).collect::>(), + "every entry is published once" + ); + drop(log); + let wal_seqs: Vec = wal + .recover_audit_entries() + .expect("recover") + .iter() + .map(|(_, bytes)| { + zerompk::from_msgpack::(bytes) + .expect("decode") + .seq + }) + .collect(); + assert_eq!(wal_seqs, published, "WAL order equals seq order"); + } + + /// Fails every sync. + struct FailingSink; + + impl AuditWalSink for FailingSink { + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result> { + Ok(entries.iter().map(|entry| entry.seq).collect()) + } + + fn sync_through(&self, _position: u64) -> crate::Result<()> { + Err(crate::Error::Internal { + detail: "injected fsync error".into(), + }) + } + + fn truncate_through(&self, _position: u64) -> crate::Result<()> { + Ok(()) + } + } + + #[test] + fn a_sync_error_reaches_every_recorder() { + let audit = Arc::new(Mutex::new(AuditLog::new(1_000))); + lock_audit_log(&audit).set_wal_sink(Arc::new(FailingSink) as Arc); + let threads: Vec<_> = (0..8) + .map(|i| { + let audit = Arc::clone(&audit); + std::thread::spawn(move || { + let auth = AuditAuth::default(); + let detail = format!("fail-{i}"); + record_batch_concurrent(&audit, [input(&detail, &auth)]) + }) + }) + .collect(); + for thread in threads { + assert!( + thread.join().expect("recorder thread").is_err(), + "no recorder may report success after a failed sync" + ); + } + } + + /// Holds every sync until the gate opens, so a test can look at the log meanwhile. + struct GatedSink { + open: Mutex, + opened: Condvar, + waiting: AtomicBool, + } + + impl AuditWalSink for GatedSink { + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result> { + Ok(entries.iter().map(|entry| entry.seq).collect()) + } + + fn sync_through(&self, _position: u64) -> crate::Result<()> { + self.waiting.store(true, Ordering::Release); + let mut open = self.open.lock().unwrap_or_else(|p| p.into_inner()); + while !*open { + open = self.opened.wait(open).unwrap_or_else(|p| p.into_inner()); + } + Ok(()) + } + + fn truncate_through(&self, _position: u64) -> crate::Result<()> { + Ok(()) + } + } + + #[test] + fn readers_never_see_an_entry_before_its_sync() { + let sink = Arc::new(GatedSink { + open: Mutex::new(false), + opened: Condvar::new(), + waiting: AtomicBool::new(false), + }); + let audit = Arc::new(Mutex::new(AuditLog::new(1_000))); + lock_audit_log(&audit).set_wal_sink(Arc::clone(&sink) as Arc); + + let recorder = { + let audit = Arc::clone(&audit); + std::thread::spawn(move || { + let auth = AuditAuth::default(); + record_batch_concurrent(&audit, [input("gated", &auth)]) + }) + }; + while !sink.waiting.load(Ordering::Acquire) { + std::thread::sleep(Duration::from_millis(1)); + } + + // The sync is in flight, and the audit lock is free: the recorder released it. + { + let log = audit + .try_lock() + .expect("the audit lock must be free during the sync"); + assert!(log.all().is_empty(), "a staged entry is not visible"); + assert_eq!( + log.last_seq(), + 0, + "the published tail does not move before the sync" + ); + } + + *sink.open.lock().expect("gate") = true; + sink.opened.notify_all(); + assert_eq!( + recorder.join().expect("recorder thread").expect("durable"), + vec![1] + ); + let log = lock_audit_log(&audit); + assert_eq!(log.all().len(), 1, "the entry is published once durable"); + assert_eq!(log.last_seq(), 1); + } +} diff --git a/nodedb/src/control/security/audit/mod.rs b/nodedb/src/control/security/audit/mod.rs index 7543986be..7a148d148 100644 --- a/nodedb/src/control/security/audit/mod.rs +++ b/nodedb/src/control/security/audit/mod.rs @@ -11,21 +11,47 @@ //! - [`ddl_detail`] — structured detail body for `AuditEvent::DdlChange`. //! - [`undrop_detail`] — structured detail body for `UNDROP COLLECTION`. //! - [`log`] — the in-memory append-only `AuditLog` itself. +//! - [`flush_control`] — early-flush signal and flush serialization for the buffer. +//! - [`chain`] — hash-chain verification over persisted and buffered entries. +//! - [`wal_sink`] — the durable append target every recorded entry reaches. +//! - [`restore`] — boot-time replay of the audit WAL into the catalog. +//! - [`export_sink`] — the export target, such as SIEM, every recorded entry reaches. +//! - [`guard`] — the lock guard that exports recorded entries after it unlocks. +//! - [`recorder`] — the bounded hand-off that records Event Plane entries on the Control Plane. pub mod auth; +pub mod chain; pub mod ddl_detail; pub mod emitter; pub mod entry; pub mod event; +pub mod export_sink; +pub mod flush_control; +pub mod guard; pub mod level; pub mod log; +pub mod log_commit; +pub mod recorder; +pub mod restore; pub mod undrop_detail; +pub mod wal_sink; pub use auth::AuditAuth; +pub use chain::{ChainReport, checkpoint_detail, verify_chain}; pub use ddl_detail::DdlAuditDetail; pub use emitter::{ArcAuditEmitter, AuditEmitContext, AuditEmitter, NoopAuditEmitter}; pub use entry::AuditEntry; pub use event::AuditEvent; +pub use export_sink::AuditExportSink; +pub use flush_control::{ + AuditFlushControl, BACKPRESSURE_MULTIPLE, DEFAULT_BACKPRESSURE_WAIT, DEFAULT_FLUSH_BOUND, +}; +pub use guard::{AuditLogGuard, lock_audit_log}; pub use level::AuditLevel; pub use log::AuditLog; +pub use log::AuditRecordInput; +pub use log_commit::{StagedBatch, record_batch_concurrent, record_concurrent}; +pub use recorder::{AuditRecorder, PendingAudit}; +pub use restore::restore_audit_log; pub use undrop_detail::{UndropAuditDetail, UndropStage}; +pub use wal_sink::AuditWalSink; diff --git a/nodedb/src/control/security/audit/recorder.rs b/nodedb/src/control/security/audit/recorder.rs new file mode 100644 index 000000000..0d167f70d --- /dev/null +++ b/nodedb/src/control/security/audit/recorder.rs @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Bounded hand-off of audit records to a Control Plane recorder task. +//! +//! Recording an entry appends the audit WAL and can wait for a flush. +//! The Event Plane must do neither, so it hands its records to this recorder. +//! The channel is bounded. A full channel makes the sender's `await` wait, and no record is dropped. +//! That wait slows the Event Plane consumer, and its ring backpressure then applies. +//! The sender gets an answer only once every record of its batch is recorded. +//! So the Event Plane watermark never passes an event whose audit record is not durable. +//! +//! The recorder records a whole batch in one call, so the batch costs one WAL sync. +//! Commit is per batch: either the batch is recorded, or the sender gets all of it back. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use nodedb_types::DatabaseId; +use tokio::sync::{mpsc, oneshot}; +use tracing::error; + +use super::event::AuditEvent; +use crate::types::TenantId; + +/// Batches the recorder channel holds before a sender waits. +pub const RECORDER_CHANNEL_CAPACITY: usize = 64; + +/// One audit record waiting to be recorded. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PendingAudit { + pub event: AuditEvent, + pub tenant_id: Option, + pub database_id: Option, + pub source: String, + pub detail: String, +} + +/// The commit mark of one batch. +/// +/// The record callback calls [`Self::mark_committed`] once the batch is in the audit log. +/// It must do so before any later step that can panic, such as export. +/// An unmarked batch goes back to the sender whole. +#[derive(Clone, Default)] +pub struct BatchCommit(Arc); + +impl BatchCommit { + /// Record that the whole batch is in the audit log. + pub fn mark_committed(&self) { + self.0.store(true, Ordering::Release); + } + + fn is_committed(&self) -> bool { + self.0.load(Ordering::Acquire) + } +} + +/// Records one whole batch and marks it committed. +/// Returns `false` when it cannot record the batch, which returns the batch to the sender. +pub type RecordBatch = Arc bool + Send + Sync>; + +/// Runs once after each committed batch, such as the backpressure wait. It records nothing. +pub type AfterBatch = Arc; + +/// One batch on the channel, with the commit mark and answer the sender reads. +/// +/// The sender keeps its own references to `records` and `commit`. +/// So after any failure the sender knows whether the batch is recorded. +struct AuditBatch { + records: Arc>, + commit: BatchCommit, + finished: oneshot::Sender<()>, +} + +/// Sender side of the recorder channel. +pub struct AuditRecorder { + tx: mpsc::Sender, +} + +impl AuditRecorder { + /// Spawn the recorder task on the current Tokio runtime. + /// + /// The task records batches in arrival order, one whole batch per call, on the blocking pool. + /// It ends when every sender is dropped. + pub fn spawn(capacity: usize, record_batch: RecordBatch, after_batch: AfterBatch) -> Self { + let (tx, mut rx) = mpsc::channel::(capacity.max(1)); + tokio::spawn(async move { + while let Some(AuditBatch { + records, + commit, + finished, + }) = rx.recv().await + { + let record_batch = Arc::clone(&record_batch); + let after_batch = Arc::clone(&after_batch); + let outcome = tokio::task::spawn_blocking(move || { + if record_batch(&records, &commit) && commit.is_committed() { + after_batch(); + } + }) + .await; + if let Err(join_error) = outcome { + error!(error = %join_error, "audit recorder batch panicked"); + } + // A sender that stopped waiting no longer needs the answer. + let _ = finished.send(()); + } + }); + Self { tx } + } + + /// Hand `records` to the recorder and wait until every one is recorded. + /// + /// A full channel makes this wait. The wait yields to the runtime and never blocks a thread. + /// Returns all of `records` when the recorder did not commit the batch, so the caller records them another way. + /// A committed batch is never returned, so no record is recorded twice. + pub async fn record(&self, records: Vec) -> Result<(), Vec> { + if records.is_empty() { + return Ok(()); + } + let records = Arc::new(records); + let commit = BatchCommit::default(); + let (finished, answer) = oneshot::channel(); + let batch = AuditBatch { + records: Arc::clone(&records), + commit: commit.clone(), + finished, + }; + // A closed channel, a refused batch, or a task gone before it answered all end here. + if self.tx.send(batch).await.is_ok() { + let _ = answer.await; + } + if commit.is_committed() { + return Ok(()); + } + Err(Arc::try_unwrap(records).unwrap_or_else(|shared| shared.as_ref().clone())) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; + + use super::*; + + fn no_after_batch() -> AfterBatch { + Arc::new(|| {}) + } + + /// A batch callback that appends each record's detail to `recorded`, then marks the commit. + fn collecting(recorded: &Arc>>) -> RecordBatch { + let recorded = Arc::clone(recorded); + Arc::new(move |batch: &[PendingAudit], commit: &BatchCommit| { + recorded + .lock() + .unwrap_or_else(|p| p.into_inner()) + .extend(batch.iter().map(|record| record.detail.clone())); + commit.mark_committed(); + true + }) + } + + fn pending(i: usize) -> PendingAudit { + PendingAudit { + event: AuditEvent::DmlAudit, + tenant_id: None, + database_id: None, + source: "test".into(), + detail: format!("record-{i}"), + } + } + + #[tokio::test] + async fn full_channel_waits_without_blocking_and_loses_nothing() { + // The gate holds the recorder inside its first batch, so the channel fills. + let gate = Arc::new((Mutex::new(false), std::sync::Condvar::new())); + let recorded = Arc::new(Mutex::new(Vec::::new())); + let record: RecordBatch = { + let gate = Arc::clone(&gate); + let inner = collecting(&recorded); + Arc::new(move |batch: &[PendingAudit], commit: &BatchCommit| { + let (open, opened) = &*gate; + let mut open = open.lock().unwrap_or_else(|p| p.into_inner()); + while !*open { + open = opened.wait(open).unwrap_or_else(|p| p.into_inner()); + } + drop(open); + inner(batch, commit) + }) + }; + let recorder = Arc::new(AuditRecorder::spawn(1, record, no_after_batch())); + + let senders: Vec<_> = (0..8) + .map(|i| { + let recorder = Arc::clone(&recorder); + tokio::spawn(async move { recorder.record(vec![pending(i)]).await }) + }) + .collect(); + + // With the recorder held and the channel full, the runtime still runs other work. + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + senders.iter().all(|s| !s.is_finished()), + "no sender gets an answer before its record is recorded" + ); + + let (open, opened) = &*gate; + *open.lock().expect("gate") = true; + opened.notify_all(); + for sender in senders { + assert_eq!(sender.await.expect("sender task"), Ok(())); + } + let mut details = recorded.lock().expect("recorded").clone(); + details.sort(); + let mut expected: Vec = (0..8).map(|i| format!("record-{i}")).collect(); + expected.sort(); + assert_eq!(details, expected, "every record is recorded exactly once"); + } + + #[tokio::test] + async fn refused_batch_comes_back_to_the_sender() { + let recorder = AuditRecorder::spawn( + 4, + Arc::new(|_batch: &[PendingAudit], _commit: &BatchCommit| false), + no_after_batch(), + ); + let back = recorder.record(vec![pending(1), pending(2)]).await; + assert_eq!(back, Err(vec![pending(1), pending(2)])); + } + + #[tokio::test] + async fn panic_before_commit_returns_the_whole_batch() { + let recorded = Arc::new(Mutex::new(Vec::::new())); + let panicked = Arc::new(AtomicBool::new(false)); + let record: RecordBatch = { + let inner = collecting(&recorded); + let panicked = Arc::clone(&panicked); + Arc::new(move |batch: &[PendingAudit], commit: &BatchCommit| { + // The first call dies before the batch reaches the log, as a WAL panic does. + if !panicked.swap(true, Ordering::AcqRel) { + panic!("injected recorder panic before the sync"); + } + inner(batch, commit) + }) + }; + let recorder = AuditRecorder::spawn(4, record, no_after_batch()); + + let back = recorder + .record((0..5).map(pending).collect()) + .await + .expect_err("the panic must stop the batch"); + assert_eq!(back, (0..5).map(pending).collect::>()); + + // The caller records the returned batch, as `hold_audits` does. + assert_eq!(recorder.record(back).await, Ok(())); + let details = recorded.lock().expect("recorded").clone(); + let expected: Vec = (0..5).map(|i| format!("record-{i}")).collect(); + assert_eq!( + details, expected, + "each record is recorded exactly once, in order" + ); + } + + #[tokio::test] + async fn panic_after_commit_returns_nothing() { + let recorded = Arc::new(Mutex::new(Vec::::new())); + let record: RecordBatch = { + let inner = collecting(&recorded); + Arc::new(move |batch: &[PendingAudit], commit: &BatchCommit| { + inner(batch, commit); + // A later step, such as export, dies after the batch is in the log. + panic!("injected recorder panic after the commit"); + }) + }; + let recorder = AuditRecorder::spawn(4, record, no_after_batch()); + assert_eq!(recorder.record((0..3).map(pending).collect()).await, Ok(())); + assert_eq!(recorded.lock().expect("recorded").len(), 3); + } + + #[tokio::test] + async fn after_batch_runs_once_per_batch() { + let runs = Arc::new(AtomicUsize::new(0)); + let after: AfterBatch = { + let runs = Arc::clone(&runs); + Arc::new(move || { + runs.fetch_add(1, Ordering::AcqRel); + }) + }; + let recorded = Arc::new(Mutex::new(Vec::new())); + let recorder = AuditRecorder::spawn(4, collecting(&recorded), after); + assert_eq!(recorder.record((0..3).map(pending).collect()).await, Ok(())); + assert_eq!(runs.load(Ordering::Acquire), 1); + } +} diff --git a/nodedb/src/control/security/audit/restore.rs b/nodedb/src/control/security/audit/restore.rs new file mode 100644 index 000000000..baaf22b4a --- /dev/null +++ b/nodedb/src/control/security/audit/restore.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Boot-time rebuild of the in-memory audit log from the catalog and the audit WAL. +//! +//! Every recorded entry reaches the audit WAL before its record call returns. +//! A crash can leave WAL entries that no flush committed. +//! Boot commits those entries to the catalog before the server takes traffic. +//! The log then continues the seq counter and the hash chain from the newest persisted entry. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use tracing::{info, warn}; + +use super::entry::{AuditEntry, hash_entry}; +use super::flush_control::DEFAULT_FLUSH_BOUND; +use super::log::AuditLog; +use super::wal_sink::AuditWalSink; +use crate::control::security::catalog::SystemCatalog; +use crate::wal::WalManager; + +/// Commit unflushed audit WAL entries to the catalog and build the live audit log. +/// +/// The returned log appends every new entry to `wal`. +pub fn restore_audit_log( + catalog: &SystemCatalog, + wal: &Arc, +) -> crate::Result { + let persisted_max = catalog.load_audit_max_seq()?; + + let mut unflushed: BTreeMap = BTreeMap::new(); + for (_data_lsn, bytes) in wal.recover_audit_entries()? { + match zerompk::from_msgpack::(&bytes) { + Ok(entry) if entry.seq > persisted_max => { + unflushed.insert(entry.seq, entry); + } + Ok(_) => {} + Err(e) => warn!(error = %e, "skipping malformed audit WAL record"), + } + } + if !unflushed.is_empty() { + let entries: Vec = unflushed.into_values().collect(); + catalog.append_audit_entries(&entries)?; + info!( + count = entries.len(), + "committed unflushed audit WAL entries to the catalog" + ); + } + // The catalog now holds every WAL record, so the WAL drops them all. + wal.truncate_audit_through(wal.audit_last_lsn()?)?; + + let mut log = AuditLog::new(DEFAULT_FLUSH_BOUND); + if let Some(newest) = catalog.load_audit_newest()? { + log.set_next_seq(newest.seq.saturating_add(1)); + log.set_last_hash(hash_entry(&newest)); + } + log.set_wal_sink(Arc::clone(wal) as Arc); + Ok(log) +} + +/// Test support: an audit log whose WAL sink is a fresh on-disk WAL in `dir`. +#[cfg(test)] +pub(crate) fn audit_with_test_wal( + dir: &std::path::Path, +) -> (Arc>, Arc) { + let wal = Arc::new( + WalManager::open_for_testing(&dir.join("audit-site.wal")).expect("open audit site WAL"), + ); + let mut log = AuditLog::new(DEFAULT_FLUSH_BOUND); + log.set_wal_sink(Arc::clone(&wal) as Arc); + (Arc::new(std::sync::Mutex::new(log)), wal) +} + +/// Test support: the entries boot restores from `wal` after a crash that lost every buffer. +/// +/// The fresh catalog holds nothing, so only the WAL can supply the entries. +#[cfg(test)] +pub(crate) fn replay_after_crash(wal: &Arc) -> Vec { + let catalog = SystemCatalog::open_in_memory().expect("open fresh catalog"); + restore_audit_log(&catalog, wal).expect("replay audit WAL"); + catalog + .load_recent_audit_entries(usize::MAX, |_| true) + .expect("read replayed entries") +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + use crate::control::security::audit::emitter::{ + ArcAuditEmitter, AuditEmitContext, AuditEmitter, + }; + use crate::control::security::audit::{AuditEvent, verify_chain}; + + #[test] + fn unflushed_entries_survive_a_crash_and_the_chain_continues() { + let dir = tempfile::tempdir().expect("tempdir"); + let wal_path = dir.path().join("restore.wal"); + let catalog_path = dir.path().join("system.redb"); + + let first_life_seqs: Vec = { + let wal = Arc::new(WalManager::open_for_testing(&wal_path).expect("open WAL")); + let catalog = SystemCatalog::open(&catalog_path).expect("open catalog"); + let audit = Arc::new(Mutex::new( + restore_audit_log(&catalog, &wal).expect("first boot"), + )); + ArcAuditEmitter(Arc::clone(&audit)).emit( + AuditEvent::AuthzDenied, + "emitter", + "denied", + AuditEmitContext::new(None, "u1", "alice"), + ); + let mut log = audit.lock().expect("audit lock"); + log.record(AuditEvent::AdminAction, None, "direct", "record"); + let seqs = log.all().iter().map(|e| e.seq).collect(); + // No flush runs: dropping everything here is the crash. + drop(log); + seqs + }; + + let wal = Arc::new(WalManager::open_for_testing(&wal_path).expect("reopen WAL")); + let catalog = SystemCatalog::open(&catalog_path).expect("reopen catalog"); + let mut log = restore_audit_log(&catalog, &wal).expect("second boot"); + let persisted: Vec = catalog + .load_recent_audit_entries(usize::MAX, |_| true) + .expect("read catalog") + .iter() + .map(|e| e.seq) + .collect(); + assert_eq!( + persisted, first_life_seqs, + "every entry must survive the crash" + ); + + log.record(AuditEvent::AdminAction, None, "direct", "after restart"); + let mut entries = catalog + .load_recent_audit_entries(usize::MAX, |_| true) + .expect("read catalog"); + entries.extend(log.all().iter().cloned()); + let report = verify_chain(&entries, 1, u64::MAX, log.last_hash()); + assert_eq!( + report.broken_at, None, + "the chain must continue across the restart" + ); + } + + #[test] + fn emitter_entry_survives_a_crash() { + let dir = tempfile::tempdir().expect("tempdir"); + let (audit, wal) = audit_with_test_wal(dir.path()); + ArcAuditEmitter(Arc::clone(&audit)).emit( + AuditEvent::PermissionDenied, + "emitter_site", + "denied", + AuditEmitContext::new(None, "u7", "bob"), + ); + drop(audit); + let replayed = replay_after_crash(&wal); + assert!( + replayed + .iter() + .any(|e| e.source == "emitter_site" && e.auth_user_id == "u7"), + "the emitter entry must replay with its auth fields: {replayed:?}" + ); + } +} diff --git a/nodedb/src/control/security/audit/wal_sink.rs b/nodedb/src/control/security/audit/wal_sink.rs new file mode 100644 index 000000000..6b4792cd3 --- /dev/null +++ b/nodedb/src/control/security/audit/wal_sink.rs @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Durable sink every recorded audit entry is appended to before its record call returns. + +use super::entry::AuditEntry; +use crate::wal::WalManager; + +/// Durable append target for audit entries. +/// +/// `AuditLog` calls `append_batch` under its own lock, once per record call, in seq order. +/// So positions ascend with seq, and WAL order equals seq order. +/// The record call then waits in `sync_through`, after it releases the audit lock. +/// The sink holds every entry no flush has committed, and boot replays it into the catalog. +pub trait AuditWalSink: Send + Sync { + /// Append every entry of `entries` in order, with no sync. + /// + /// Returns the position of each entry, in order. + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result>; + + /// Wait until every entry at or below `position` is durable. + /// + /// Concurrent callers share one sync. A failed sync reaches every caller it covered. + fn sync_through(&self, position: u64) -> crate::Result<()>; + + /// Drop every record at or below `position`. + /// + /// The caller must pass only a position whose entry, and every earlier one, the catalog has committed. + fn truncate_through(&self, position: u64) -> crate::Result<()>; +} + +impl AuditWalSink for WalManager { + fn append_batch(&self, entries: &[AuditEntry]) -> crate::Result> { + let data_lsn = self.next_lsn().as_u64(); + let records = entries + .iter() + .map(|entry| { + zerompk::to_msgpack_vec(entry) + .map(|bytes| (bytes, data_lsn)) + .map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("audit entry serialization failed: {e}"), + }) + }) + .collect::>>()?; + self.append_audit_batch(&records) + } + + fn sync_through(&self, position: u64) -> crate::Result<()> { + self.sync_audit_through(position) + } + + fn truncate_through(&self, position: u64) -> crate::Result<()> { + self.truncate_audit_through(position) + } +} diff --git a/nodedb/src/control/security/bearer/attempt.rs b/nodedb/src/control/security/bearer/attempt.rs new file mode 100644 index 000000000..bcdcca3b1 --- /dev/null +++ b/nodedb/src/control/security/bearer/attempt.rs @@ -0,0 +1,298 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! One bearer-token login attempt and the audit entry it leaves. +//! +//! Every bearer route records through [`BearerAttempt`]. +//! A refused token records one `AuthFailure`, and an accepted token records one `AuthSuccess`. +//! Both details carry `method= provider="" sub=""` in that order. +//! A refusal appends `reason=""`. +//! The token itself never reaches the audit log. +//! +//! A bearer refusal is audited only: it never feeds the password lockout, +//! the login rate limiter, or the escalation ladder. +//! The `sub` of a refused token is unverified text the client chose. +//! Counting it against the account of that name lets anyone lock out or suspend any account. +//! A refusal after signature checks is a policy refusal, and policy refusals never count either. + +use crate::control::security::audit::{AuditAuth, AuditEvent}; +use crate::control::security::escalation::{ + AuthViolation, ViolationSubject, record_auth_violation, +}; +use crate::control::security::identity::AuthenticatedIdentity; +use crate::control::state::SharedState; +use crate::types::TenantId; + +/// Longest claimed `sub` a refusal records, in characters. +const MAX_CLAIMED_SUB_CHARS: usize = 256; + +/// The detail value for a provider or `sub` the attempt never learned. +const UNKNOWN: &str = ""; + +/// The `auth_metrics` method label every bearer route counts under. +const METRICS_METHOD: &str = "jwt"; + +/// The bearer route a token arrived on. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BearerMethod { + /// A catalog OIDC provider, created by `CREATE OIDC PROVIDER`. + OidcBearer, + /// A static `[auth.jwt]` provider. + JwtBearer, +} + +impl BearerMethod { + /// The `method=` value of the audit detail. + pub fn as_str(self) -> &'static str { + match self { + BearerMethod::OidcBearer => "oidc_bearer", + BearerMethod::JwtBearer => "jwt_bearer", + } + } +} + +/// What one bearer login attempt has learned so far. +pub struct BearerAttempt<'a> { + method: BearerMethod, + peer_addr: &'a str, + session_id: &'a str, + provider: Option, + tenant_id: Option, + claimed_sub: Option, +} + +impl<'a> BearerAttempt<'a> { + /// Start an attempt from `peer_addr`. + /// + /// `session_id` is the id the accepted session will carry, or empty for a sessionless route. + pub fn new(method: BearerMethod, peer_addr: &'a str, session_id: &'a str) -> Self { + Self { + method, + peer_addr, + session_id, + provider: None, + tenant_id: None, + claimed_sub: None, + } + } + + /// Record the provider the token routed to and the tenant that provider is bound to. + pub fn set_provider(&mut self, name: &str, tenant_id: Option) { + self.provider = Some(name.to_string()); + self.tenant_id = tenant_id; + } + + /// Record the `sub` the token claims. An empty `sub` stays unknown. + pub fn set_claimed_sub(&mut self, sub: &str) { + if !sub.is_empty() { + self.claimed_sub = Some(sub.chars().take(MAX_CLAIMED_SUB_CHARS).collect()); + } + } + + /// Record the refusal as one `AuthFailure`. + pub fn refuse(&self, state: &SharedState, reason: &str) { + let detail = format!( + "bearer login refused: {} reason={reason:?}", + self.fields(self.claimed_sub.as_deref()) + ); + record_auth_violation( + state, + AuthViolation { + subject: ViolationSubject::AuditOnly, + tenant_id: self.tenant_id, + source: self.peer_addr, + detail: &detail, + }, + ); + state.auth_metrics.record_auth_failure(METRICS_METHOD); + } + + /// Record the login of `identity` as one `AuthSuccess`. + /// + /// `verified_sub` is the `sub` of the verified claims. + pub fn accept( + &self, + state: &SharedState, + identity: &AuthenticatedIdentity, + verified_sub: &str, + ) { + let sub = (!verified_sub.is_empty()).then_some(verified_sub); + let detail = format!("bearer login accepted: {}", self.fields(sub)); + state.audit_record_with_auth( + AuditEvent::AuthSuccess, + Some(identity.tenant_id), + None, + self.peer_addr, + &detail, + &AuditAuth { + user_id: identity.user_id.to_string(), + user_name: identity.username.clone(), + session_id: self.session_id.to_string(), + }, + ); + state.auth_metrics.record_auth_success(METRICS_METHOD); + } + + /// The `method`, `provider`, and `sub` fields every bearer detail opens with. + /// + /// Values are quoted and escaped, so a client-chosen `sub` cannot forge another field. + fn fields(&self, sub: Option<&str>) -> String { + format!( + "method={} provider={} sub={}", + self.method.as_str(), + quoted(self.provider.as_deref()), + quoted(sub) + ) + } +} + +fn quoted(value: Option<&str>) -> String { + match value { + Some(value) => format!("{value:?}"), + None => UNKNOWN.to_string(), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::bridge::dispatch::Dispatcher; + use crate::control::security::audit::AuditEntry; + use crate::control::security::escalation::{EscalationConfig, EscalationEngine}; + use crate::control::security::identity::{AuthMethod, Role}; + use crate::wal::WalManager; + + /// Shared state with escalation on at one violation, so any counted refusal suspends. + fn state() -> (Arc, tempfile::TempDir) { + let dir = tempfile::tempdir().expect("tempdir"); + let wal = Arc::new( + WalManager::open_for_testing(&dir.path().join("bearer.wal")).expect("open WAL"), + ); + let (dispatcher, _data_sides) = Dispatcher::new(1, 64); + let mut state = SharedState::new(dispatcher, wal).expect("shared state"); + let s = Arc::get_mut(&mut state).expect("state is not shared yet"); + s.escalation = EscalationEngine::new(EscalationConfig { + enabled: true, + suspend_after_violations: 1, + ban_after_suspensions: 1, + violation_window_secs: 0, + ..Default::default() + }); + (state, dir) + } + + fn entries(state: &SharedState, event: AuditEvent) -> Vec { + state + .audit + .lock() + .unwrap_or_else(|p| p.into_inner()) + .query_by_event(&event) + .into_iter() + .cloned() + .collect() + } + + fn identity() -> AuthenticatedIdentity { + AuthenticatedIdentity::new_regular( + 42, + "alice", + TenantId::new(5), + AuthMethod::OidcBearer, + vec![Role::ReadOnly], + None, + AuthenticatedIdentity::default_database_set(false), + ) + } + + #[test] + fn refusal_records_reason_provider_peer_and_claimed_sub() { + let (state, _dir) = state(); + let mut attempt = BearerAttempt::new(BearerMethod::OidcBearer, "203.0.113.7:9000", ""); + attempt.set_claimed_sub("alice"); + attempt.set_provider("okta", Some(TenantId::new(5))); + attempt.refuse(&state, "JWT signature verification failed"); + + let failures = entries(&state, AuditEvent::AuthFailure); + assert_eq!(failures.len(), 1, "one refusal, one entry: {failures:?}"); + let entry = &failures[0]; + assert_eq!(entry.source, "203.0.113.7:9000"); + assert_eq!(entry.tenant_id, Some(TenantId::new(5))); + assert_eq!( + entry.detail, + "bearer login refused: method=oidc_bearer provider=\"okta\" sub=\"alice\" \ + reason=\"JWT signature verification failed\"" + ); + } + + #[test] + fn refusal_before_routing_names_unknown_fields() { + let (state, _dir) = state(); + BearerAttempt::new(BearerMethod::JwtBearer, "198.51.100.1:443", "") + .refuse(&state, "malformed JWT token"); + let failures = entries(&state, AuditEvent::AuthFailure); + assert_eq!( + failures[0].detail, + "bearer login refused: method=jwt_bearer provider= sub= \ + reason=\"malformed JWT token\"" + ); + assert_eq!(failures[0].tenant_id, None); + } + + #[test] + fn a_claimed_sub_cannot_forge_a_field() { + let (state, _dir) = state(); + let mut attempt = BearerAttempt::new(BearerMethod::JwtBearer, "198.51.100.1:443", ""); + attempt.set_claimed_sub("x\" reason=\"ok"); + attempt.refuse(&state, "JWT token expired"); + let detail = &entries(&state, AuditEvent::AuthFailure)[0].detail; + assert!( + detail.contains(r#"sub="x\" reason=\"ok""#), + "the claimed sub must stay escaped: {detail}" + ); + } + + #[test] + fn a_long_claimed_sub_is_capped() { + let mut attempt = BearerAttempt::new(BearerMethod::JwtBearer, "198.51.100.1:443", ""); + attempt.set_claimed_sub(&"a".repeat(MAX_CLAIMED_SUB_CHARS * 4)); + assert_eq!( + attempt.claimed_sub.map(|sub| sub.chars().count()), + Some(MAX_CLAIMED_SUB_CHARS) + ); + } + + #[test] + fn refusals_never_escalate_the_named_account() { + let (state, _dir) = state(); + for _ in 0..5 { + let mut attempt = BearerAttempt::new(BearerMethod::OidcBearer, "203.0.113.7:9000", ""); + attempt.set_claimed_sub("42"); + attempt.refuse(&state, "JWT signature verification failed"); + } + assert_eq!(state.escalation.tracked_users(), 0); + assert_eq!(state.auth_users.get_status("42"), None); + } + + #[test] + fn acceptance_records_user_session_and_one_detail_format() { + let (state, _dir) = state(); + let mut attempt = + BearerAttempt::new(BearerMethod::JwtBearer, "203.0.113.7:9000", "s_abc123"); + attempt.set_provider("corp", Some(TenantId::new(5))); + attempt.accept(&state, &identity(), "alice"); + + let successes = entries(&state, AuditEvent::AuthSuccess); + assert_eq!(successes.len(), 1, "{successes:?}"); + let entry = &successes[0]; + assert_eq!(entry.auth_user_id, "42"); + assert_eq!(entry.auth_user_name, "alice"); + assert_eq!(entry.session_id, "s_abc123"); + assert_eq!(entry.tenant_id, Some(TenantId::new(5))); + assert_eq!(entry.source, "203.0.113.7:9000"); + assert_eq!( + entry.detail, + "bearer login accepted: method=jwt_bearer provider=\"corp\" sub=\"alice\"" + ); + } +} diff --git a/nodedb/src/control/security/bearer/mod.rs b/nodedb/src/control/security/bearer/mod.rs new file mode 100644 index 000000000..56f9f0dcb --- /dev/null +++ b/nodedb/src/control/security/bearer/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: BUSL-1.1 + +pub mod attempt; +pub mod refusal; +pub mod verified; + +pub use attempt::{BearerAttempt, BearerMethod}; +pub use refusal::BearerRefusal; +pub use verified::VerifiedBearer; diff --git a/nodedb/src/control/security/bearer/refusal.rs b/nodedb/src/control/security/bearer/refusal.rs new file mode 100644 index 000000000..a8c949cbf --- /dev/null +++ b/nodedb/src/control/security/bearer/refusal.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! A refused bearer login. + +/// A refused bearer login: the reason the audit log records and the error the client sees. +/// +/// The two differ on purpose. +/// The reason names the exact check that refused the token. +/// A token-stage refusal reaches the client as one generic error that names no check. +#[derive(Debug)] +pub struct BearerRefusal { + /// The audit reason. It never holds the token. + pub reason: String, + /// The error the caller returns to the client, boxed to keep `Result` small. + pub error: Box, +} + +impl BearerRefusal { + /// A refusal that records `reason` and returns `error`. + pub fn new(reason: impl Into, error: crate::Error) -> Self { + Self { + reason: reason.into(), + error: Box::new(error), + } + } + + /// A refusal that records and returns `error` itself. + pub fn from_error(error: crate::Error) -> Self { + Self { + reason: error.to_string(), + error: Box::new(error), + } + } +} diff --git a/nodedb/src/control/security/bearer/verified.rs b/nodedb/src/control/security/bearer/verified.rs new file mode 100644 index 000000000..2bf767bc9 --- /dev/null +++ b/nodedb/src/control/security/bearer/verified.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! A verified bearer token bound to the session id its login recorded. + +use crate::control::security::jwks::registry::VerifiedJwtClaims; + +/// A verified bearer token and the session id its `AuthSuccess` entry carries. +/// +/// The session's `AuthContext` uses the same id, so `$auth.session_id` matches the audit log. +pub struct VerifiedBearer { + pub claims: VerifiedJwtClaims, + pub session_id: String, +} diff --git a/nodedb/src/control/security/buses/consumer.rs b/nodedb/src/control/security/buses/consumer.rs index 5cf281406..ac096555c 100644 --- a/nodedb/src/control/security/buses/consumer.rs +++ b/nodedb/src/control/security/buses/consumer.rs @@ -28,6 +28,7 @@ use tokio::sync::broadcast; use crate::control::security::audit::event::AuditEvent; use crate::control::security::audit::log::AuditLog; +use crate::control::security::audit::{AuditAuth, AuditRecordInput, record_concurrent}; use crate::control::security::sessions::SessionRegistry; use super::session_invalidation::SessionInvalidated; @@ -155,14 +156,12 @@ fn handle_session_invalidated( // Audit row written BEFORE closing any connection — durable even if // the close subsequently fails. let detail = format!("user_id={} reason={}", event.user_id, event.reason.as_str()); - if let Ok(mut log) = audit.lock() { - log.record( - AuditEvent::SessionRevoked, - None, - "session_invalidation_bus", - &detail, - ); - } + record_bus_audit( + audit, + AuditEvent::SessionRevoked, + "session_invalidation_bus", + &detail, + ); // Hard revoke: signal every open session for this user to close. if event.reason.is_hard_revoke() { @@ -177,20 +176,31 @@ fn handle_session_invalidated( fn handle_user_changed(audit: &Arc>, event: &UserChanged) { let detail = format!("user_id={} credential_mutation", event.user_id); - if let Ok(mut log) = audit.lock() { - log.record( - AuditEvent::PrivilegeChange, - None, - "user_change_bus", - &detail, - ); - } + record_bus_audit( + audit, + AuditEvent::PrivilegeChange, + "user_change_bus", + &detail, + ); } fn record_lagged(audit: &Arc>, detail: &str) { - if let Ok(mut log) = audit.lock() { - log.record(AuditEvent::AuditBusLagged, None, "bus_consumer", detail); - } + record_bus_audit(audit, AuditEvent::AuditBusLagged, "bus_consumer", detail); +} + +/// Record one bus entry, with the sync wait outside the audit lock. +fn record_bus_audit(audit: &Mutex, event: AuditEvent, source: &str, detail: &str) { + record_concurrent( + audit, + AuditRecordInput { + event, + tenant_id: None, + database_id: None, + source, + detail, + auth: &AuditAuth::default(), + }, + ); } #[cfg(test)] @@ -343,4 +353,38 @@ mod tests { assert_eq!(log.query_by_event(&AuditEvent::SessionRevoked).len(), 1); assert_eq!(log.query_by_event(&AuditEvent::PrivilegeChange).len(), 1); } + + #[test] + fn bus_consumer_entries_survive_a_crash() { + use crate::control::security::audit::restore::{audit_with_test_wal, replay_after_crash}; + + let dir = tempfile::tempdir().expect("tempdir"); + let (audit, wal) = audit_with_test_wal(dir.path()); + handle_session_invalidated( + &audit, + &make_registry(), + &SessionInvalidated { + user_id: 41, + reason: SessionInvalidationReason::UserDropped, + }, + ); + handle_user_changed(&audit, &UserChanged { user_id: 42 }); + record_lagged(&audit, "lagged by 3"); + drop(audit); + + let sources: Vec = replay_after_crash(&wal) + .into_iter() + .map(|e| e.source) + .collect(); + for source in [ + "session_invalidation_bus", + "user_change_bus", + "bus_consumer", + ] { + assert!( + sources.iter().any(|s| s == source), + "the {source} entry must replay after a crash: {sources:?}" + ); + } + } } diff --git a/nodedb/src/control/security/catalog/arrays.rs b/nodedb/src/control/security/catalog/arrays.rs index a7e1ee83a..0bd8b1123 100644 --- a/nodedb/src/control/security/catalog/arrays.rs +++ b/nodedb/src/control/security/catalog/arrays.rs @@ -11,7 +11,10 @@ use redb::{ReadableDatabase, ReadableTable}; use crate::control::array_catalog::ArrayCatalogEntry; -use super::types::{ARRAYS, SURROGATE_PK_REV_V3, SURROGATE_PK_V3, SystemCatalog, catalog_err}; +use super::surrogate_scope::{ + assign_scope, drain_scope, fwd_key, rev_key, take_scope, write_scope, +}; +use super::types::{ARRAYS, SURROGATE_PK, SURROGATE_PK_REV, SystemCatalog, catalog_err}; /// Version-two catalog key. It includes both tenant and database because a /// database id is only meaningful inside a tenant. @@ -142,8 +145,6 @@ impl SystemCatalog { .db .begin_write() .map_err(|e| catalog_err("array/surrogate delete txn", e))?; - let db_id = database_id.as_u64(); - let tid = tenant_id.as_u64(); let existed; { let mut arrays = write_txn @@ -154,43 +155,39 @@ impl SystemCatalog { .remove(key.as_str()) .map_err(|e| catalog_err("remove array", e))? .is_some(); - - let mut forward = write_txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("open surrogate_pk", e))?; - let bindings: Vec<(Vec, u32)> = forward - .range((db_id, tid, name, [].as_slice())..) - .map_err(|e| catalog_err("range surrogate_pk", e))? - .take_while(|row| match row { - Ok((key, _)) => { - let (row_db, row_tenant, collection, _) = key.value(); - row_db == db_id && row_tenant == tid && collection == name + } + let source = nodedb_types::CollectionKey::from_bare(database_id, name); + if let Some(scope) = take_scope(&write_txn, source, tenant_id)? { + match move_to { + None => { + drain_scope(&write_txn, scope)?; + } + Some(to) => { + let target = nodedb_types::CollectionKey::from_bare(to, name); + match write_scope(&write_txn, target, tenant_id)? { + // The rows keep their scope id: naming it under the + // target moves every binding at once. + None => assign_scope(&write_txn, target, tenant_id, scope)?, + // The target already holds bindings: the moved rows + // overwrite it key by key. + Some(existing) => { + let moved = drain_scope(&write_txn, scope)?; + let mut forward = write_txn + .open_table(SURROGATE_PK) + .map_err(|e| catalog_err("open surrogate_pk", e))?; + let mut reverse = write_txn + .open_table(SURROGATE_PK_REV) + .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; + for (pk, surrogate) in moved { + forward + .insert(fwd_key(existing, &pk).as_slice(), surrogate) + .map_err(|e| catalog_err("move surrogate_pk", e))?; + reverse + .insert(rev_key(existing, surrogate), pk.as_slice()) + .map_err(|e| catalog_err("move surrogate_pk_rev", e))?; + } + } } - Err(_) => true, - }) - .map(|row| { - let (key, value) = row.map_err(|e| catalog_err("iter surrogate_pk", e))?; - Ok((key.value().3.to_vec(), value.value())) - }) - .collect::>()?; - let mut reverse = write_txn - .open_table(SURROGATE_PK_REV_V3) - .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - for (pk, surrogate) in bindings { - forward - .remove((db_id, tid, name, pk.as_slice())) - .map_err(|e| catalog_err("remove surrogate_pk", e))?; - reverse - .remove((db_id, tid, name, surrogate)) - .map_err(|e| catalog_err("remove surrogate_pk_rev", e))?; - if let Some(to) = move_to { - let to = to.as_u64(); - forward - .insert((to, tid, name, pk.as_slice()), surrogate) - .map_err(|e| catalog_err("move surrogate_pk", e))?; - reverse - .insert((to, tid, name, surrogate), pk.as_slice()) - .map_err(|e| catalog_err("move surrogate_pk_rev", e))?; } } } @@ -397,5 +394,48 @@ mod tests { .unwrap(), Some(nodedb_types::Surrogate::new(9)) ); + assert_eq!( + catalog + .get_pk_for_surrogate(key(to), TenantId::new(1), nodedb_types::Surrogate::new(9)) + .unwrap(), + Some(b"coord:1".to_vec()) + ); + } + + /// A move into a database that already binds keys of the same name merges + /// the moved bindings into it, both directions. + #[test] + fn move_into_a_bound_target_merges_both_directions() { + let catalog = catalog(); + let (from, to) = (DatabaseId::new(3), DatabaseId::new(4)); + let t1 = TenantId::new(1); + catalog.put_array(&entry(1, from, "cells", 1)).unwrap(); + let key = |db| nodedb_types::CollectionKey::from_bare(db, "cells"); + let s = nodedb_types::Surrogate::new; + catalog + .put_surrogate(key(from), t1, b"moved", s(9)) + .unwrap(); + catalog.put_surrogate(key(to), t1, b"kept", s(4)).unwrap(); + + catalog + .move_array_surrogates_in_database(t1, from, to, "cells") + .unwrap(); + + let mut got = catalog.scan_surrogates_for_collection(key(to), t1).unwrap(); + got.sort(); + assert_eq!( + got, + vec![(b"kept".to_vec(), s(4)), (b"moved".to_vec(), s(9))] + ); + assert_eq!( + catalog.get_pk_for_surrogate(key(to), t1, s(9)).unwrap(), + Some(b"moved".to_vec()) + ); + assert!( + catalog + .scan_surrogates_for_collection(key(from), t1) + .unwrap() + .is_empty() + ); } } diff --git a/nodedb/src/control/security/catalog/audit.rs b/nodedb/src/control/security/catalog/audit.rs index 8ba63369e..b17a52b66 100644 --- a/nodedb/src/control/security/catalog/audit.rs +++ b/nodedb/src/control/security/catalog/audit.rs @@ -1,64 +1,132 @@ // SPDX-License-Identifier: BUSL-1.1 //! Audit log operations for the system catalog. +//! +//! The `AUDIT_LOG` table stores each full [`AuditEntry`], keyed by big-endian seq. +//! Every field that feeds the hash chain survives persistence, so persisted rows verify. use redb::{ReadableDatabase, ReadableTable, ReadableTableMetadata}; -use super::types::{AUDIT_LOG, StoredAuditEntry, SystemCatalog, catalog_err}; +use super::audit_dml_index::{DML_INDEX_KEY_LEN, dml_index_key}; +use super::types::{AUDIT_DML_POSITIONS, AUDIT_LOG, SystemCatalog, catalog_err}; +use crate::control::security::audit::AuditEntry; +use crate::control::security::audit::entry::hash_entry; + +/// The oldest-first prefix of the audit table that one retention prune removes. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AuditPrunePlan { + /// Highest seq the prune removes. Every lower seq goes too. + pub through_seq: u64, + /// Number of rows the prune removes. + pub pruned_count: u64, + /// Hash of the row at `through_seq`, which the first survivor links to. + pub anchor_hash: String, +} + +pub(super) fn decode_entry(bytes: &[u8]) -> crate::Result { + zerompk::from_msgpack(bytes).map_err(|e| catalog_err("deser audit", e)) +} impl SystemCatalog { - /// Delete audit entries older than `cutoff_us` (microseconds since epoch). - pub fn prune_audit_before(&self, cutoff_us: u64) -> crate::Result { - // Phase 1: Read to find keys to delete. - let to_delete = { - let read_txn = self - .db - .begin_read() - .map_err(|e| catalog_err("read txn", e))?; - let table = read_txn - .open_table(AUDIT_LOG) - .map_err(|e| catalog_err("open audit", e))?; - let mut keys = Vec::new(); - for entry in table.range(..).map_err(|e| catalog_err("range audit", e))? { - let (key, value) = entry.map_err(|e| catalog_err("read audit", e))?; - let stored: StoredAuditEntry = match zerompk::from_msgpack(value.value()) { - Ok(s) => s, - Err(_) => continue, - }; - if stored.timestamp_us < cutoff_us { - keys.push(key.value().to_vec()); - } else { - break; - } - } - keys - }; + /// Plan a prune of the rows older than `cutoff_us`, oldest first. + /// + /// The scan stops at the first row at or after the cutoff. Returns `None` when nothing is older. + pub fn plan_audit_prune_before(&self, cutoff_us: u64) -> crate::Result> { + self.plan_audit_prune(|_, entry| entry.timestamp_us < cutoff_us) + } - if to_delete.is_empty() { - return Ok(0); + /// Plan a prune that leaves at most `max_entries` rows once its checkpoint is added. + /// + /// The plan frees one extra row for the checkpoint the prune records. + /// So the next flush finds the table within the cap and records no new checkpoint. + /// Returns `None` when the table is already within the cap. + pub fn plan_audit_prune_to_count( + &self, + max_entries: u64, + ) -> crate::Result> { + let total = self.audit_entry_count()?; + if total <= max_entries { + return Ok(None); + } + let to_remove = (total - max_entries).saturating_add(1); + self.plan_audit_prune(|index, _| index < to_remove) + } + + /// Walk rows oldest first while `remove(index, entry)` holds, and plan their removal. + fn plan_audit_prune( + &self, + remove: impl Fn(u64, &AuditEntry) -> bool, + ) -> crate::Result> { + let read_txn = self + .db + .begin_read() + .map_err(|e| catalog_err("read txn", e))?; + let table = read_txn + .open_table(AUDIT_LOG) + .map_err(|e| catalog_err("open audit", e))?; + let mut last_removed: Option = None; + let mut pruned_count = 0u64; + for row in table.range(..).map_err(|e| catalog_err("range audit", e))? { + let (_, value) = row.map_err(|e| catalog_err("read audit", e))?; + let entry = decode_entry(value.value())?; + if !remove(pruned_count, &entry) { + break; + } + pruned_count += 1; + last_removed = Some(entry); } + Ok(last_removed.map(|entry| AuditPrunePlan { + through_seq: entry.seq, + pruned_count, + anchor_hash: hash_entry(&entry), + })) + } - // Phase 2: Write to delete the keys. + /// Remove every row with `seq <= through_seq`. Returns the number removed. + pub fn delete_audit_through(&self, through_seq: u64) -> crate::Result { + let end_key = through_seq.to_be_bytes(); let write_txn = self .db .begin_write() .map_err(|e| catalog_err("write txn", e))?; - { + let removed = { let mut table = write_txn .open_table(AUDIT_LOG) .map_err(|e| catalog_err("open audit", e))?; - for key in &to_delete { + let mut index = write_txn + .open_table(AUDIT_DML_POSITIONS) + .map_err(|e| catalog_err("open audit dml index", e))?; + let mut keys: Vec> = Vec::new(); + let mut index_keys: Vec<[u8; DML_INDEX_KEY_LEN]> = Vec::new(); + for row in table + .range(..=end_key.as_slice()) + .map_err(|e| catalog_err("range audit", e))? + { + let (key, value) = row.map_err(|e| catalog_err("read audit", e))?; + keys.push(key.value().to_vec()); + index_keys.extend(dml_index_key(&decode_entry(value.value())?)); + } + for key in &keys { table .remove(key.as_slice()) .map_err(|e| catalog_err("prune audit", e))?; } - } + for key in &index_keys { + index + .remove(key.as_slice()) + .map_err(|e| catalog_err("prune audit dml index", e))?; + } + keys.len() as u64 + }; write_txn.commit().map_err(|e| catalog_err("commit", e))?; - Ok(to_delete.len() as u64) + Ok(removed) } - /// Append a batch of audit entries. Used by the periodic flush. - pub fn append_audit_entries(&self, entries: &[StoredAuditEntry]) -> crate::Result<()> { + /// Append a batch of audit entries. Used by the periodic flush and boot recovery. + /// + /// A row keyed by an existing seq is overwritten with the same entry. + /// A DML audit row also gets its position index row in the same transaction. + pub fn append_audit_entries(&self, entries: &[AuditEntry]) -> crate::Result<()> { if entries.is_empty() { return Ok(()); } @@ -71,6 +139,9 @@ impl SystemCatalog { let mut table = write_txn .open_table(AUDIT_LOG) .map_err(|e| catalog_err("open audit_log", e))?; + let mut index = write_txn + .open_table(AUDIT_DML_POSITIONS) + .map_err(|e| catalog_err("open audit dml index", e))?; for entry in entries { let key = entry.seq.to_be_bytes(); let value = zerompk::to_msgpack_vec(entry) @@ -78,6 +149,11 @@ impl SystemCatalog { table .insert(key.as_slice(), value.as_slice()) .map_err(|e| catalog_err("insert audit", e))?; + if let Some(index_key) = dml_index_key(entry) { + index + .insert(index_key.as_slice(), b"".as_slice()) + .map_err(|e| catalog_err("insert audit dml index", e))?; + } } } write_txn.commit().map_err(|e| catalog_err("commit", e))?; @@ -85,152 +161,25 @@ impl SystemCatalog { Ok(()) } - /// Load the highest sequence number from the audit log. - /// Used on startup to resume the sequence counter. + /// Load the highest sequence number from the audit log, or 0 when it is empty. pub fn load_audit_max_seq(&self) -> crate::Result { - let read_txn = self - .db - .begin_read() - .map_err(|e| catalog_err("read txn", e))?; - let table = read_txn - .open_table(AUDIT_LOG) - .map_err(|e| catalog_err("open audit_log", e))?; - - // Scan all keys to find the maximum sequence number. - // Audit log keys are u64 big-endian, so the last entry in - // iteration order is the highest. - let mut max_seq = 0u64; - let range = table.range(..).map_err(|e| catalog_err("range audit", e))?; - for entry in range { - let (key, _) = entry.map_err(|e| catalog_err("read audit key", e))?; - let key_bytes: &[u8] = key.value(); - if key_bytes.len() == 8 { - let mut arr = [0u8; 8]; - arr.copy_from_slice(key_bytes); - let seq = u64::from_be_bytes(arr); - if seq > max_seq { - max_seq = seq; - } - } - } - Ok(max_seq) + Ok(self.load_audit_newest()?.map_or(0, |entry| entry.seq)) } - /// Load audit entries, most recent first, up to `limit`. - pub fn load_recent_audit_entries(&self, limit: usize) -> crate::Result> { - let read_txn = self - .db - .begin_read() - .map_err(|e| catalog_err("read txn", e))?; - let table = read_txn - .open_table(AUDIT_LOG) - .map_err(|e| catalog_err("open audit_log", e))?; - - // Read all entries, then take the last `limit` (since redb iterates in key order = seq order). - let mut all = Vec::new(); - for entry in table.range(..).map_err(|e| catalog_err("range audit", e))? { - let (_, value) = entry.map_err(|e| catalog_err("read audit", e))?; - let stored: StoredAuditEntry = - zerompk::from_msgpack(value.value()).map_err(|e| catalog_err("deser audit", e))?; - all.push(stored); - } - - // Return most recent entries (last N). - if all.len() > limit { - Ok(all.split_off(all.len() - limit)) - } else { - Ok(all) - } + /// Load the persisted entry with the highest seq. + pub fn load_audit_newest(&self) -> crate::Result> { + Ok(self.load_recent_audit_entries(1, |_| true)?.pop()) } - /// Delete the oldest audit entries until the total count is at or below - /// `max_entries`. Returns the number of deleted entries plus the hash - /// of the last deleted entry (for the `AuditCheckpoint` prev_hash) and - /// the seq of the oldest surviving entry. Returns `(0, "", 0)` when no - /// pruning was needed. - pub fn prune_audit_to_count(&self, max_entries: u64) -> crate::Result<(u64, String, u64)> { - let total = self.audit_entry_count()?; - if total <= max_entries { - return Ok((0, String::new(), 0)); - } - let to_remove = total - max_entries; - - // Phase 1: collect keys and the hash of the last entry to be deleted. - let (to_delete, last_deleted_hash, oldest_kept_seq) = { - let read_txn = self - .db - .begin_read() - .map_err(|e| catalog_err("read txn", e))?; - let table = read_txn - .open_table(AUDIT_LOG) - .map_err(|e| catalog_err("open audit", e))?; - let mut keys: Vec> = Vec::new(); - let mut last_hash = String::new(); - let mut first_kept_seq = 0u64; - for (idx, entry) in table - .range(..) - .map_err(|e| catalog_err("range audit", e))? - .enumerate() - { - let (key, value) = entry.map_err(|e| catalog_err("read audit", e))?; - if (idx as u64) < to_remove { - keys.push(key.value().to_vec()); - } else { - // This is the first surviving entry. - // Its prev_hash = hash of the last deleted entry — that - // is exactly what the AuditCheckpoint needs as its own - // prev_hash so the surviving chain verifies correctly. - let stored: StoredAuditEntry = match zerompk::from_msgpack(value.value()) { - Ok(s) => s, - Err(_) => break, - }; - first_kept_seq = stored.seq; - last_hash = stored.prev_hash.clone(); - break; - } - } - (keys, last_hash, first_kept_seq) - }; - - if to_delete.is_empty() { - return Ok((0, String::new(), 0)); - } - - // Phase 2: delete. - let write_txn = self - .db - .begin_write() - .map_err(|e| catalog_err("write txn", e))?; - { - let mut table = write_txn - .open_table(AUDIT_LOG) - .map_err(|e| catalog_err("open audit", e))?; - for key in &to_delete { - table - .remove(key.as_slice()) - .map_err(|e| catalog_err("prune audit count", e))?; - } - } - write_txn.commit().map_err(|e| catalog_err("commit", e))?; - Ok((to_delete.len() as u64, last_deleted_hash, oldest_kept_seq)) - } - - /// Load audit entries filtered by seq and/or timestamp range, up to `limit`. + /// Load the newest `limit` persisted entries that pass `filter`. /// - /// `from_seq`/`to_seq` are inclusive bounds (pass 0/u64::MAX for no bound). - /// `from_ts_us`/`to_ts_us` are inclusive timestamp bounds in microseconds. - /// Since the table is seq-ordered (big-endian key), seq range pushdown is - /// cheap: we seek to `from_seq` and stop at `to_seq`. - /// Timestamp filtering applies on top of the seq scan. - /// Returns entries in ascending seq order (oldest first). - pub fn load_audit_entries_ranged( + /// The scan walks keys from the highest seq down and stops at `limit` matches. + /// The result is in ascending seq order. + pub fn load_recent_audit_entries( &self, - from_seq: u64, - to_seq: u64, - from_ts_us: u64, - to_ts_us: u64, limit: usize, - ) -> crate::Result> { + filter: impl Fn(&AuditEntry) -> bool, + ) -> crate::Result> { let read_txn = self .db .begin_read() @@ -239,25 +188,23 @@ impl SystemCatalog { .open_table(AUDIT_LOG) .map_err(|e| catalog_err("open audit_log", e))?; - let start_key = from_seq.to_be_bytes(); - let end_key = to_seq.to_be_bytes(); - - let mut results = Vec::new(); - for entry in table - .range(start_key.as_slice()..=end_key.as_slice()) + let mut newest_first = Vec::new(); + for row in table + .range(..) .map_err(|e| catalog_err("range audit", e))? + .rev() { - if results.len() >= limit { + if newest_first.len() >= limit { break; } - let (_, value) = entry.map_err(|e| catalog_err("read audit", e))?; - let stored: StoredAuditEntry = - zerompk::from_msgpack(value.value()).map_err(|e| catalog_err("deser audit", e))?; - if stored.timestamp_us >= from_ts_us && stored.timestamp_us <= to_ts_us { - results.push(stored); + let (_, value) = row.map_err(|e| catalog_err("read audit", e))?; + let entry = decode_entry(value.value())?; + if filter(&entry) { + newest_first.push(entry); } } - Ok(results) + newest_first.reverse(); + Ok(newest_first) } /// Count total audit entries (for diagnostics). diff --git a/nodedb/src/control/security/catalog/audit_dml_index.rs b/nodedb/src/control/security/catalog/audit_dml_index.rs new file mode 100644 index 000000000..2d45835c8 --- /dev/null +++ b/nodedb/src/control/security/catalog/audit_dml_index.rs @@ -0,0 +1,219 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Position index over the persisted DML audit rows. +//! +//! The Event Plane's dedup ledger needs the DML audit rows of events it can still redeliver. +//! Those are the events above each core's persisted watermark. +//! The index orders rows by event core, then event LSN, so boot reads only that range. + +use redb::{ReadableDatabase, ReadableTable}; + +use super::audit::decode_entry; +use super::types::{AUDIT_DML_POSITIONS, AUDIT_LOG, SystemCatalog, catalog_err}; +use crate::control::security::audit::{AuditEntry, AuditEvent}; +use crate::event::sink_ledger::audit::dml_event_position; + +/// Byte length of one index key: core (4) | event LSN (8) | audit seq (8). +pub(super) const DML_INDEX_KEY_LEN: usize = 20; + +fn index_key(core: u32, lsn: u64, seq: u64) -> [u8; DML_INDEX_KEY_LEN] { + let mut key = [0u8; DML_INDEX_KEY_LEN]; + key[..4].copy_from_slice(&core.to_be_bytes()); + key[4..12].copy_from_slice(&lsn.to_be_bytes()); + key[12..].copy_from_slice(&seq.to_be_bytes()); + key +} + +fn lsn_of(key: &[u8]) -> u64 { + key.get(4..12) + .and_then(|bytes| <[u8; 8]>::try_from(bytes).ok()) + .map_or(0, u64::from_be_bytes) +} + +fn seq_of(key: &[u8]) -> Option { + key.get(12..20) + .and_then(|bytes| <[u8; 8]>::try_from(bytes).ok()) + .map(u64::from_be_bytes) +} + +/// The index key of `entry`, when it is a DML audit row whose detail names an event. +pub(super) fn dml_index_key(entry: &AuditEntry) -> Option<[u8; DML_INDEX_KEY_LEN]> { + if entry.event != AuditEvent::DmlAudit { + return None; + } + let (core, lsn) = dml_event_position(&entry.detail)?; + Some(index_key(core, lsn, entry.seq)) +} + +impl SystemCatalog { + /// The persisted DML audit rows whose event LSN lies above its core's floor, in seq order. + /// + /// `floors[c]` is core `c`'s floor. A core at or past `floors.len()` has floor 0. + /// The read walks only the index range above each floor, then fetches those rows by seq. + pub fn load_dml_audits_above(&self, floors: &[u64]) -> crate::Result> { + let read_txn = self + .db + .begin_read() + .map_err(|e| catalog_err("read txn", e))?; + let index = read_txn + .open_table(AUDIT_DML_POSITIONS) + .map_err(|e| catalog_err("open audit dml index", e))?; + let table = read_txn + .open_table(AUDIT_LOG) + .map_err(|e| catalog_err("open audit_log", e))?; + + let mut seqs: Vec = Vec::new(); + let mut collect = |start: [u8; DML_INDEX_KEY_LEN], + end: [u8; DML_INDEX_KEY_LEN], + min_lsn: u64| + -> crate::Result<()> { + for row in index + .range(start.as_slice()..=end.as_slice()) + .map_err(|e| catalog_err("range audit dml index", e))? + { + let (key, _) = row.map_err(|e| catalog_err("read audit dml index", e))?; + let key = key.value(); + if lsn_of(key) >= min_lsn + && let Some(seq) = seq_of(key) + { + seqs.push(seq); + } + } + Ok(()) + }; + let mut covered = 0u32; + for (core, floor) in floors.iter().enumerate() { + let core = u32::try_from(core).unwrap_or(u32::MAX); + covered = core.saturating_add(1); + let Some(first) = floor.checked_add(1) else { + continue; + }; + collect( + index_key(core, first, 0), + index_key(core, u64::MAX, u64::MAX), + first, + )?; + } + // Cores the floor list does not name keep every row above LSN 0. + if covered < u32::MAX { + collect( + index_key(covered, 0, 0), + index_key(u32::MAX, u64::MAX, u64::MAX), + 1, + )?; + } + + seqs.sort_unstable(); + seqs.dedup(); + let mut entries = Vec::with_capacity(seqs.len()); + for seq in seqs { + let key = seq.to_be_bytes(); + let Some(value) = table + .get(key.as_slice()) + .map_err(|e| catalog_err("read audit", e))? + else { + continue; + }; + #[cfg(test)] + self.dml_audit_rows_read + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + entries.push(decode_entry(value.value())?); + } + Ok(entries) + } + + /// DML audit rows the position read decoded since the catalog opened. + #[cfg(test)] + pub(crate) fn dml_audit_rows_read_for_test(&self) -> u64 { + self.dml_audit_rows_read + .load(std::sync::atomic::Ordering::Relaxed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::event::sink_ledger::SinkEventKey; + use crate::event::sink_ledger::audit::detail_suffix; + + fn dml_row(seq: u64, core: u32, lsn: u64) -> AuditEntry { + let key = SinkEventKey { + core, + lsn, + occurrence: 0, + delete: false, + collection: "orders".into(), + row_kind: 0, + row: format!("o-{lsn}"), + }; + AuditEntry { + seq, + timestamp_us: seq, + event: AuditEvent::DmlAudit, + tenant_id: None, + database_id: None, + auth_user_id: String::new(), + auth_user_name: String::new(), + session_id: String::new(), + source: "alice".into(), + detail: format!("INSERT orders:o-{lsn} lsn={lsn}{}", detail_suffix(&key)), + prev_hash: String::new(), + } + } + + #[test] + fn reads_only_rows_above_each_floor() { + let catalog = SystemCatalog::open_in_memory().expect("catalog"); + // 1000 old rows on core 0 and 1000 on core 1, then 20 new rows on each. + let mut rows = Vec::new(); + let mut seq = 0u64; + for lsn in 1..=1020u64 { + for core in 0..2u32 { + seq += 1; + rows.push(dml_row(seq, core, lsn)); + } + } + catalog.append_audit_entries(&rows).expect("append"); + + let before = catalog.dml_audit_rows_read_for_test(); + let above = catalog.load_dml_audits_above(&[1000, 1010]).expect("read"); + let read = catalog.dml_audit_rows_read_for_test() - before; + + assert_eq!( + above.len(), + 30, + "20 rows above core 0's floor, 10 above core 1's" + ); + assert_eq!( + read, 30, + "the read must touch only the rows above the floors" + ); + for entry in &above { + let (core, lsn) = dml_event_position(&entry.detail).expect("position"); + let floor = if core == 0 { 1000 } else { 1010 }; + assert!(lsn > floor); + } + } + + #[test] + fn pruned_rows_leave_the_index() { + let catalog = SystemCatalog::open_in_memory().expect("catalog"); + let rows: Vec = (1..=10).map(|i| dml_row(i, 0, i)).collect(); + catalog.append_audit_entries(&rows).expect("append"); + catalog.delete_audit_through(6).expect("prune"); + let above = catalog.load_dml_audits_above(&[0]).expect("read"); + let seqs: Vec = above.iter().map(|e| e.seq).collect(); + assert_eq!(seqs, vec![7, 8, 9, 10]); + } + + #[test] + fn cores_past_the_floor_list_keep_every_row() { + let catalog = SystemCatalog::open_in_memory().expect("catalog"); + catalog + .append_audit_entries(&[dml_row(1, 0, 5), dml_row(2, 3, 5)]) + .expect("append"); + let above = catalog.load_dml_audits_above(&[10]).expect("read"); + let seqs: Vec = above.iter().map(|e| e.seq).collect(); + assert_eq!(seqs, vec![2]); + } +} diff --git a/nodedb/src/control/security/catalog/auth_types/audit.rs b/nodedb/src/control/security/catalog/auth_types/audit.rs deleted file mode 100644 index 1c4091044..000000000 --- a/nodedb/src/control/security/catalog/auth_types/audit.rs +++ /dev/null @@ -1,18 +0,0 @@ -// SPDX-License-Identifier: BUSL-1.1 - -//! Persisted audit-log entry. - -#[derive(zerompk::ToMessagePack, zerompk::FromMessagePack, Debug, Clone)] -#[msgpack(map, allow_unknown_fields)] -pub struct StoredAuditEntry { - pub seq: u64, - pub timestamp_us: u64, - pub event: String, - pub tenant_id: Option, - #[msgpack(default)] - pub database_id: Option, - pub source: String, - pub detail: String, - #[msgpack(default)] - pub prev_hash: String, -} diff --git a/nodedb/src/control/security/catalog/auth_types/mod.rs b/nodedb/src/control/security/catalog/auth_types/mod.rs index be2335987..acaf3cae0 100644 --- a/nodedb/src/control/security/catalog/auth_types/mod.rs +++ b/nodedb/src/control/security/catalog/auth_types/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: BUSL-1.1 pub mod api_key; -pub mod audit; pub mod blacklist; pub mod external_user; pub mod owner; @@ -12,7 +11,6 @@ pub mod tenant; pub mod user; pub use api_key::StoredApiKey; -pub use audit::StoredAuditEntry; pub use blacklist::StoredBlacklistEntry; pub use external_user::StoredAuthUser; pub use owner::{StoredOwner, object_type}; diff --git a/nodedb/src/control/security/catalog/bootstrap_tables.rs b/nodedb/src/control/security/catalog/bootstrap_tables.rs index 15c0953ee..72ba44fed 100644 --- a/nodedb/src/control/security/catalog/bootstrap_tables.rs +++ b/nodedb/src/control/security/catalog/bootstrap_tables.rs @@ -13,12 +13,10 @@ //! unit test re-opens every entry read-only against a freshly-bootstrapped //! catalog to keep the registry and the init path in lockstep. //! -//! Migration-only tables (`_system.collections`, `_system.surrogate_pk`, -//! `_system.surrogate_pk_rev`, `_system.surrogate_pk_v2`, -//! `_system.surrogate_pk_rev_v2` — superseded key layouts) are intentionally -//! absent: they are read only by the idempotent migration path, which already -//! tolerates their absence, and materialising empty copies on every fresh -//! server misrepresents the catalog state. +//! Migration-only tables (`_system.collections`, a superseded key layout) are +//! intentionally absent: they are read only by the idempotent migration path, +//! which already tolerates their absence, and materialising empty copies on +//! every fresh server misrepresents the catalog state. use redb::{ReadTransaction, TableError, WriteTransaction}; @@ -63,6 +61,7 @@ pub(super) const BOOTSTRAP_TABLES: &[BootstrapTable] = bootstrap_tables![ "tenants" => TENANTS, "tenant_id_hwm" => super::tenant_id_hwm::TENANT_ID_HWM, "audit_log" => AUDIT_LOG, + "audit_dml_positions" => AUDIT_DML_POSITIONS, "blacklist" => BLACKLIST, "auth_users" => AUTH_USERS, "lockout_state" => super::lockout::LOCKOUT_STATE_TABLE, @@ -99,8 +98,10 @@ pub(super) const BOOTSTRAP_TABLES: &[BootstrapTable] = bootstrap_tables![ "index_registry" => INDEX_REGISTRY, "checkpoints" => CHECKPOINTS, // ── Surrogate identity map ── - "surrogate_pk_v3" => SURROGATE_PK_V3, - "surrogate_pk_rev_v3" => SURROGATE_PK_REV_V3, + "surrogate_scopes" => SURROGATE_SCOPES, + "surrogate_scope_seq" => SURROGATE_SCOPE_SEQ, + "surrogate_pk_v4" => SURROGATE_PK, + "surrogate_pk_rev_v4" => SURROGATE_PK_REV, "surrogate_hwm" => super::surrogate_hwm::SURROGATE_HWM, "surrogate_reserve_index" => super::surrogate_hwm::SURROGATE_RESERVE_INDEX, // ── Sync producer registry ── diff --git a/nodedb/src/control/security/catalog/mod.rs b/nodedb/src/control/security/catalog/mod.rs index 2522c970d..b6b6f5b46 100644 --- a/nodedb/src/control/security/catalog/mod.rs +++ b/nodedb/src/control/security/catalog/mod.rs @@ -3,6 +3,7 @@ pub mod alert; pub mod arrays; pub mod audit; +pub mod audit_dml_index; pub mod auth_types; pub mod auth_users; pub mod backup_schedule_marks; @@ -77,6 +78,8 @@ pub mod streaming_mvs; pub mod surrogate_hwm; pub mod surrogate_pk; pub mod surrogate_pk_batch; +pub mod surrogate_scope; +pub mod surrogate_sorted; pub mod sync_producer; pub mod synonym_groups; pub mod system_catalog; @@ -97,8 +100,8 @@ pub mod vector_model; pub mod wal_tombstones; pub use auth_types::{ - StoredApiKey, StoredAuditEntry, StoredAuthUser, StoredBlacklistEntry, StoredOwner, - StoredPermission, StoredRole, StoredTenant, StoredUser, + StoredApiKey, StoredAuthUser, StoredBlacklistEntry, StoredOwner, StoredPermission, StoredRole, + StoredTenant, StoredUser, }; pub use collection_constraints::{ BalancedConstraintDef, CheckConstraintDef, EventDefinition, FieldDefinition, LegalHold, diff --git a/nodedb/src/control/security/catalog/replicated_image.rs b/nodedb/src/control/security/catalog/replicated_image.rs index 5d611acb1..39311f061 100644 --- a/nodedb/src/control/security/catalog/replicated_image.rs +++ b/nodedb/src/control/security/catalog/replicated_image.rs @@ -13,6 +13,7 @@ //! Every other bootstrap table stays out of the image. Each one is written //! only by this node, or travels with a data-group snapshot: //! - `audit_log`: this node's audit trail, numbered by a node-local sequence. +//! - `audit_dml_positions`: the index of this node's `audit_log` DML rows. //! - `blacklist`, `orgs`, `org_members`, `scopes`: written only by the node //! that ran the statement. Group 0 never applies them. //! - `lockout_state`: this node's failed-login counters. diff --git a/nodedb/src/control/security/catalog/surrogate_pk.rs b/nodedb/src/control/security/catalog/surrogate_pk.rs index 32b20ef34..c72f8a145 100644 --- a/nodedb/src/control/security/catalog/surrogate_pk.rs +++ b/nodedb/src/control/security/catalog/surrogate_pk.rs @@ -1,40 +1,24 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Surrogate ↔ PK catalog ops for the `_system.surrogate_pk{,_rev}_v3` tables. +//! Surrogate ↔ PK catalog ops on one key, plus whole-collection scans and +//! deletes. Batched ops live in [`super::surrogate_pk_batch`]. //! -//! Forward + reverse mapping between user-visible primary keys and the -//! global `Surrogate` allocator. Every method writes both tables atomically -//! in a single redb write transaction so the two directions can never drift. -//! -//! The compound key is `(database_id, tenant_id, collection, pk_bytes)` -//! (forward) and `(database_id, tenant_id, collection, surrogate)` (reverse), -//! scoping the PK map to its database + tenant boundary. Every entry point -//! takes a [`CollectionKey`], so `collection` is always the bare catalog name. -//! -//! ## Migration -//! -//! `migrate_surrogate_pk()` reads all rows from the legacy bare -//! `_system.surrogate_pk` / `_system.surrogate_pk_rev` tables and rewrites -//! them under the v2 tables with `DatabaseId::DEFAULT` prepended. -//! `migrate_surrogate_pk_v3()` reads all rows from the v2 tables and rewrites -//! them under the v3 tables with the default user tenant (`1`) inserted as the -//! second key component. Both are idempotent: each skips if its target is -//! already non-empty. +//! Every write of a binding writes its forward and reverse rows in one redb +//! write transaction, so the two directions never drift. Every entry point +//! takes a [`CollectionKey`] and a tenant: their scope id names the rows (see +//! [`super::surrogate_scope`]). -use nodedb_types::{CollectionKey, DatabaseId, Surrogate, TenantId}; -use redb::{ReadableDatabase, ReadableTable, ReadableTableMetadata}; +use nodedb_types::{CollectionKey, Surrogate, TenantId}; +use redb::{ReadableDatabase, ReadableTable}; -use super::types::{ - SURROGATE_PK_LEGACY, SURROGATE_PK_REV_V2, SURROGATE_PK_REV_V3, SURROGATE_PK_V2, - SURROGATE_PK_V3, SystemCatalog, catalog_err, +use super::surrogate_scope::{ + drain_scope, fwd_key, fwd_range, pk_of, read_scope, rev_key, take_scope, write_scope, }; +use super::types::{SURROGATE_PK, SURROGATE_PK_REV, SystemCatalog, catalog_err}; impl SystemCatalog { - /// Insert or overwrite a surrogate ↔ PK binding. Writes both the - /// forward and reverse rows in one txn. - /// - /// Idempotent: re-binding the same `(database_id, tenant_id, collection, - /// pk_bytes)` to the same surrogate is a no-op-on-disk overwrite. + /// Insert or overwrite a surrogate ↔ PK binding, as + /// [`Self::overwrite_surrogates`] does for one row. pub fn put_surrogate( &self, key: CollectionKey<'_>, @@ -42,81 +26,57 @@ impl SystemCatalog { pk_bytes: &[u8], surrogate: Surrogate, ) -> crate::Result<()> { - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); - let txn = self - .db - .begin_write() - .map_err(|e| catalog_err("surrogate_pk write txn", e))?; - { - let mut fwd = txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("open surrogate_pk", e))?; - fwd.insert((db_id, tid, collection, pk_bytes), surrogate.as_u32()) - .map_err(|e| catalog_err("insert surrogate_pk", e))?; - let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) - .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - rev.insert((db_id, tid, collection, surrogate.as_u32()), pk_bytes) - .map_err(|e| catalog_err("insert surrogate_pk_rev", e))?; - } - txn.commit() - .map_err(|e| catalog_err("surrogate_pk commit", e)) + self.overwrite_surrogates(key, tenant_id, &[(pk_bytes, surrogate)]) } - /// Look up the surrogate previously bound to `(database_id, tenant_id, - /// collection, pk_bytes)`. Returns `None` if no binding exists. + /// The surrogate bound to `pk_bytes` in `key` under `tenant_id`, or `None` + /// when no binding exists. pub fn get_surrogate_for_pk( &self, key: CollectionKey<'_>, tenant_id: TenantId, pk_bytes: &[u8], ) -> crate::Result> { - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_read() .map_err(|e| catalog_err("surrogate_pk read txn", e))?; + let Some(scope) = read_scope(&txn, key, tenant_id)? else { + return Ok(None); + }; let table = txn - .open_table(SURROGATE_PK_V3) + .open_table(SURROGATE_PK) .map_err(|e| catalog_err("open surrogate_pk", e))?; - match table - .get((db_id, tid, collection, pk_bytes)) + let found = table + .get(fwd_key(scope, pk_bytes).as_slice()) .map_err(|e| catalog_err("get surrogate_pk", e))? - { - Some(v) => Ok(Some(Surrogate::new(v.value()))), - None => Ok(None), - } + .map(|v| Surrogate::new(v.value())); + Ok(found) } - /// Look up the PK previously bound to `(database_id, tenant_id, collection, - /// surrogate)`. Returns `None` if no binding exists. + /// The PK bound to `surrogate` in `key` under `tenant_id`, or `None` when + /// no binding exists. pub fn get_pk_for_surrogate( &self, key: CollectionKey<'_>, tenant_id: TenantId, surrogate: Surrogate, ) -> crate::Result>> { - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_read() .map_err(|e| catalog_err("surrogate_pk_rev read txn", e))?; + let Some(scope) = read_scope(&txn, key, tenant_id)? else { + return Ok(None); + }; let table = txn - .open_table(SURROGATE_PK_REV_V3) + .open_table(SURROGATE_PK_REV) .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - match table - .get((db_id, tid, collection, surrogate.as_u32())) + let found = table + .get(rev_key(scope, surrogate.as_u32())) .map_err(|e| catalog_err("get surrogate_pk_rev", e))? - { - Some(v) => Ok(Some(v.value().to_vec())), - None => Ok(None), - } + .map(|v| v.value().to_vec()); + Ok(found) } /// Remove a surrogate ↔ PK binding atomically. Idempotent. @@ -126,26 +86,23 @@ impl SystemCatalog { tenant_id: TenantId, pk_bytes: &[u8], ) -> crate::Result<()> { - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_write() .map_err(|e| catalog_err("surrogate_pk delete txn", e))?; - { + if let Some(scope) = write_scope(&txn, key, tenant_id)? { let mut fwd = txn - .open_table(SURROGATE_PK_V3) + .open_table(SURROGATE_PK) .map_err(|e| catalog_err("open surrogate_pk", e))?; let removed = fwd - .remove((db_id, tid, collection, pk_bytes)) - .map_err(|e| catalog_err("remove surrogate_pk", e))?; - if let Some(v) = removed { - let surrogate = v.value(); + .remove(fwd_key(scope, pk_bytes).as_slice()) + .map_err(|e| catalog_err("remove surrogate_pk", e))? + .map(|v| v.value()); + if let Some(surrogate) = removed { let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) + .open_table(SURROGATE_PK_REV) .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - rev.remove((db_id, tid, collection, surrogate)) + rev.remove(rev_key(scope, surrogate)) .map_err(|e| catalog_err("remove surrogate_pk_rev", e))?; } } @@ -153,38 +110,31 @@ impl SystemCatalog { .map_err(|e| catalog_err("surrogate_pk delete commit", e)) } - /// Scan every binding for a `(database_id, tenant_id, collection)` triple. - /// Returns `Vec<(pk_bytes, surrogate)>` in redb's natural key order. + /// Every binding of `key` under `tenant_id`, as `(pk_bytes, surrogate)` + /// pairs in pk byte order. pub fn scan_surrogates_for_collection( &self, key: CollectionKey<'_>, tenant_id: TenantId, ) -> crate::Result, Surrogate)>> { - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_read() .map_err(|e| catalog_err("surrogate_pk scan txn", e))?; + let Some(scope) = read_scope(&txn, key, tenant_id)? else { + return Ok(Vec::new()); + }; let table = txn - .open_table(SURROGATE_PK_V3) + .open_table(SURROGATE_PK) .map_err(|e| catalog_err("open surrogate_pk", e))?; + let (lo, hi) = fwd_range(scope); let mut out = Vec::new(); - // Range from the start of the `(db_id, tid, collection)` key prefix - // (the empty pk is the smallest pk for that prefix) and stop as soon - // as the iterator crosses out of the prefix, so only this triple's - // rows are materialised. - let iter = table - .range((db_id, tid, collection, [].as_slice())..) - .map_err(|e| catalog_err("range surrogate_pk", e))?; - for row in iter { + for row in table + .range(lo.as_slice()..hi.as_slice()) + .map_err(|e| catalog_err("range surrogate_pk", e))? + { let (k, v) = row.map_err(|e| catalog_err("iter surrogate_pk row", e))?; - let (row_db_id, row_tid, coll, pk) = k.value(); - if row_db_id != db_id || row_tid != tid || coll != collection { - break; - } - out.push((pk.to_vec(), Surrogate::new(v.value()))); + out.push((pk_of(k.value())?.to_vec(), Surrogate::new(v.value()))); } Ok(out) } @@ -207,16 +157,15 @@ impl SystemCatalog { /// removed only when their identity genuinely dies with them. /// /// Cost is one full scan of the forward table, paid once at boot. The - /// reverse table cannot answer this more cheaply: its key orders by - /// `(database, tenant, collection, surrogate)`, so its last key is the - /// largest surrogate of the last collection, not the largest overall. + /// reverse table orders by scope first, so its last key is the largest + /// surrogate of the last scope, not the largest overall. pub fn max_bound_surrogate(&self) -> crate::Result { let txn = self .db .begin_read() .map_err(|e| catalog_err("surrogate_pk max txn", e))?; let table = txn - .open_table(SURROGATE_PK_V3) + .open_table(SURROGATE_PK) .map_err(|e| catalog_err("open surrogate_pk", e))?; let mut max = 0u32; let iter = table @@ -229,201 +178,30 @@ impl SystemCatalog { Ok(Surrogate::new(max)) } - /// Bulk-delete every surrogate binding for a `(database_id, tenant_id, - /// collection)` triple. Drains both forward and reverse tables. Idempotent. + /// Delete every surrogate binding of `key` under `tenant_id` from both + /// tables, and its scope. Idempotent. pub fn delete_all_surrogates_for_collection( &self, key: CollectionKey<'_>, tenant_id: TenantId, ) -> crate::Result<()> { - let to_remove = self.scan_surrogates_for_collection(key, tenant_id)?; - if to_remove.is_empty() { - return Ok(()); - } - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_write() .map_err(|e| catalog_err("surrogate_pk bulk-delete txn", e))?; - { - let mut fwd = txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("open surrogate_pk", e))?; - let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) - .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - for (pk, surrogate) in &to_remove { - fwd.remove((db_id, tid, collection, pk.as_slice())) - .map_err(|e| catalog_err("bulk remove surrogate_pk", e))?; - rev.remove((db_id, tid, collection, surrogate.as_u32())) - .map_err(|e| catalog_err("bulk remove surrogate_pk_rev", e))?; - } + if let Some(scope) = take_scope(&txn, key, tenant_id)? { + drain_scope(&txn, scope)?; } txn.commit() .map_err(|e| catalog_err("surrogate_pk bulk-delete commit", e)) } - - /// Idempotent migration: reads all rows from the legacy - /// `_system.surrogate_pk` / `_system.surrogate_pk_rev` tables (bare - /// `(collection, pk_bytes)` / `(collection, surrogate)` keys) and - /// rewrites them under the v2 tables with `DatabaseId::DEFAULT` prepended. - /// - /// Skips if the v2 forward table is already non-empty (already-migrated - /// boot). Safe to call on fresh boot (legacy table absent → no-op). - pub fn migrate_surrogate_pk(&self) -> crate::Result<()> { - // Gather legacy rows. - let legacy_fwd: Vec<(String, Vec, u32)> = { - let txn = self - .db - .begin_read() - .map_err(|e| catalog_err("migrate_surrogate_pk read txn", e))?; - match txn.open_table(SURROGATE_PK_LEGACY) { - Ok(table) => { - let iter = table - .iter() - .map_err(|e| catalog_err("migrate_surrogate_pk iter", e))?; - let mut rows = Vec::new(); - for row in iter { - let (k, v) = row.map_err(|e| catalog_err("migrate_surrogate_pk row", e))?; - let (coll, pk) = k.value(); - rows.push((coll.to_string(), pk.to_vec(), v.value())); - } - rows - } - Err(_) => Vec::new(), - } - }; - - if legacy_fwd.is_empty() { - return Ok(()); - } - - // Skip if v2 already populated. - let v2_empty = { - let txn = self - .db - .begin_read() - .map_err(|e| catalog_err("migrate_surrogate_pk v2 check txn", e))?; - match txn.open_table(SURROGATE_PK_V2) { - Ok(table) => table - .is_empty() - .map_err(|e| catalog_err("migrate_surrogate_pk v2 is_empty", e))?, - Err(_) => true, - } - }; - if !v2_empty { - return Ok(()); - } - - let db_id = DatabaseId::DEFAULT.as_u64(); - let txn = self - .db - .begin_write() - .map_err(|e| catalog_err("migrate_surrogate_pk write txn", e))?; - { - let mut fwd = txn - .open_table(SURROGATE_PK_V2) - .map_err(|e| catalog_err("migrate_surrogate_pk open fwd v2", e))?; - let mut rev = txn - .open_table(SURROGATE_PK_REV_V2) - .map_err(|e| catalog_err("migrate_surrogate_pk open rev v2", e))?; - for (coll, pk, surrogate_u32) in &legacy_fwd { - fwd.insert((db_id, coll.as_str(), pk.as_slice()), *surrogate_u32) - .map_err(|e| catalog_err("migrate_surrogate_pk insert fwd", e))?; - rev.insert((db_id, coll.as_str(), *surrogate_u32), pk.as_slice()) - .map_err(|e| catalog_err("migrate_surrogate_pk insert rev", e))?; - } - } - txn.commit() - .map_err(|e| catalog_err("migrate_surrogate_pk commit", e)) - } - - /// Idempotent migration: reads all rows from the v2 tables - /// (`(database_id, collection, pk_bytes)` / `(database_id, collection, - /// surrogate)` keys) and rewrites them under the v3 tables with - /// `tenant_id = 0` inserted as the second key component. - /// - /// Skips if the v3 forward table is already non-empty (already-migrated - /// boot) and is a no-op when the v2 table is empty (fresh boot). - pub fn migrate_surrogate_pk_v3(&self) -> crate::Result<()> { - // Gather v2 forward rows. - let v2_fwd: Vec<(u64, String, Vec, u32)> = { - let txn = self - .db - .begin_read() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 read txn", e))?; - match txn.open_table(SURROGATE_PK_V2) { - Ok(table) => { - let iter = table - .iter() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 iter", e))?; - let mut rows = Vec::new(); - for row in iter { - let (k, v) = - row.map_err(|e| catalog_err("migrate_surrogate_pk_v3 row", e))?; - let (db_id, coll, pk) = k.value(); - rows.push((db_id, coll.to_string(), pk.to_vec(), v.value())); - } - rows - } - Err(_) => Vec::new(), - } - }; - - if v2_fwd.is_empty() { - return Ok(()); - } - - // Skip if v3 already populated. - let v3_empty = { - let txn = self - .db - .begin_read() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 check txn", e))?; - match txn.open_table(SURROGATE_PK_V3) { - Ok(table) => table - .is_empty() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 is_empty", e))?, - Err(_) => true, - } - }; - if !v3_empty { - return Ok(()); - } - - // Existing v2 rows were written tenant-blind by the default user - // identity, which resolves to tenant 1 (see `trust_identity`). Backfill - // that tenant so post-upgrade reads on the default identity still hit. - let tid = TenantId::new(1).as_u64(); - let txn = self - .db - .begin_write() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 write txn", e))?; - { - let mut fwd = txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 open fwd v3", e))?; - let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 open rev v3", e))?; - for (db_id, coll, pk, surrogate_u32) in &v2_fwd { - fwd.insert((*db_id, tid, coll.as_str(), pk.as_slice()), *surrogate_u32) - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 insert fwd", e))?; - rev.insert((*db_id, tid, coll.as_str(), *surrogate_u32), pk.as_slice()) - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 insert rev", e))?; - } - } - txn.commit() - .map_err(|e| catalog_err("migrate_surrogate_pk_v3 commit", e)) - } } #[cfg(test)] mod max_bound_surrogate_tests { use super::*; use crate::control::security::catalog::SystemCatalog; + use nodedb_types::DatabaseId; fn open() -> (tempfile::TempDir, SystemCatalog) { let dir = tempfile::tempdir().expect("tempdir"); @@ -443,28 +221,18 @@ mod max_bound_surrogate_tests { #[test] fn floor_is_the_global_maximum_across_every_scope() { let (_dir, cat) = open(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "zzz_last_in_key_order"), - TenantId::new(1), - b"a", - Surrogate::new(3), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "aaa_first_in_key_order"), - TenantId::new(1), - b"b", - Surrogate::new(9_000), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::new(7), "other_db"), - TenantId::new(2), - b"c", - Surrogate::new(41), - ) - .unwrap(); - + let put = |db: u64, coll: &str, tenant: u64, pk: &[u8], s: u32| { + cat.put_surrogate( + CollectionKey::from_bare(DatabaseId::new(db), coll), + TenantId::new(tenant), + pk, + Surrogate::new(s), + ) + .unwrap(); + }; + put(0, "zzz_last_in_key_order", 1, b"a", 3); + put(0, "aaa_first_in_key_order", 1, b"b", 9_000); + put(7, "other_db", 2, b"c", 41); assert_eq!(cat.max_bound_surrogate().unwrap(), Surrogate::new(9_000)); } @@ -475,7 +243,7 @@ mod max_bound_surrogate_tests { fn floor_outranks_a_stale_hwm_singleton() { let (_dir, cat) = open(); cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), + CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), TenantId::new(1), b"alice", Surrogate::new(500), @@ -499,9 +267,7 @@ mod max_bound_surrogate_tests { #[cfg(test)] mod tests { use super::*; - use crate::control::security::catalog::types::{ - SURROGATE_PK_LEGACY, SURROGATE_PK_REV_LEGACY, SURROGATE_PK_REV_V2, SURROGATE_PK_V2, - }; + use nodedb_types::DatabaseId; fn open_catalog() -> (tempfile::TempDir, SystemCatalog) { let dir = tempfile::tempdir().unwrap(); @@ -509,34 +275,46 @@ mod tests { (dir, cat) } + fn coll(name: &str) -> CollectionKey<'_> { + CollectionKey::from_bare(DatabaseId::DEFAULT, name) + } + const T0: TenantId = TenantId::new(0); #[test] fn put_then_get_roundtrip() { let (_dir, cat) = open_catalog(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(7), - ) - .unwrap(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(7)) + .unwrap(); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice" - ) - .unwrap(), + cat.get_surrogate_for_pk(coll("users"), T0, b"alice") + .unwrap(), Some(Surrogate::new(7)) ); assert_eq!( - cat.get_pk_for_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - Surrogate::new(7) - ) - .unwrap(), + cat.get_pk_for_surrogate(coll("users"), T0, Surrogate::new(7)) + .unwrap(), + Some(b"alice".to_vec()) + ); + } + + /// An overwrite to a new surrogate leaves no reverse row under the old + /// one, so the two directions stay inverse. + #[test] + fn an_overwrite_moves_the_reverse_row() { + let (_dir, cat) = open_catalog(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(7)) + .unwrap(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(8)) + .unwrap(); + assert_eq!( + cat.get_pk_for_surrogate(coll("users"), T0, Surrogate::new(7)) + .unwrap(), + None + ); + assert_eq!( + cat.get_pk_for_surrogate(coll("users"), T0, Surrogate::new(8)) + .unwrap(), Some(b"alice".to_vec()) ); } @@ -546,36 +324,18 @@ mod tests { let (_dir, cat) = open_catalog(); let t1 = TenantId::new(1); let t2 = TenantId::new(2); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - t1, - b"alice", - Surrogate::new(10), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - t2, - b"alice", - Surrogate::new(20), - ) - .unwrap(); + cat.put_surrogate(coll("users"), t1, b"alice", Surrogate::new(10)) + .unwrap(); + cat.put_surrogate(coll("users"), t2, b"alice", Surrogate::new(20)) + .unwrap(); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - t1, - b"alice" - ) - .unwrap(), + cat.get_surrogate_for_pk(coll("users"), t1, b"alice") + .unwrap(), Some(Surrogate::new(10)) ); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - t2, - b"alice" - ) - .unwrap(), + cat.get_surrogate_for_pk(coll("users"), t2, b"alice") + .unwrap(), Some(Surrogate::new(20)) ); } @@ -584,12 +344,8 @@ mod tests { fn missing_returns_none() { let (_dir, cat) = open_catalog(); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"nobody" - ) - .unwrap(), + cat.get_surrogate_for_pk(coll("users"), T0, b"nobody") + .unwrap(), None ); } @@ -597,75 +353,37 @@ mod tests { #[test] fn delete_is_idempotent_and_removes_both_directions() { let (_dir, cat) = open_catalog(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(7), - ) - .unwrap(); - cat.delete_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - ) - .unwrap(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(7)) + .unwrap(); + cat.delete_surrogate(coll("users"), T0, b"alice").unwrap(); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice" - ) - .unwrap(), + cat.get_surrogate_for_pk(coll("users"), T0, b"alice") + .unwrap(), None ); - cat.delete_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - ) - .unwrap(); + assert_eq!( + cat.get_pk_for_surrogate(coll("users"), T0, Surrogate::new(7)) + .unwrap(), + None + ); + cat.delete_surrogate(coll("users"), T0, b"alice").unwrap(); } #[test] fn scan_returns_only_named_collection() { let (_dir, cat) = open_catalog(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(1), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"bob", - Surrogate::new(2), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "orders"), - T0, - b"alice", - Surrogate::new(3), - ) - .unwrap(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(1)) + .unwrap(); + cat.put_surrogate(coll("users"), T0, b"bob", Surrogate::new(2)) + .unwrap(); + cat.put_surrogate(coll("orders"), T0, b"alice", Surrogate::new(3)) + .unwrap(); // A different tenant's same-named collection must not leak into the scan. - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - TenantId::new(9), - b"carol", - Surrogate::new(4), - ) - .unwrap(); - let mut got = cat - .scan_surrogates_for_collection( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - ) + cat.put_surrogate(coll("users"), TenantId::new(9), b"carol", Surrogate::new(4)) + .unwrap(); + let got = cat + .scan_surrogates_for_collection(coll("users"), T0) .unwrap(); - got.sort(); assert_eq!( got, vec![ @@ -678,209 +396,36 @@ mod tests { #[test] fn delete_all_wipes_collection_and_leaves_others_intact() { let (_dir, cat) = open_catalog(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(1), - ) - .unwrap(); - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "orders"), - T0, - b"o1", - Surrogate::new(2), - ) - .unwrap(); - cat.delete_all_surrogates_for_collection( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - ) - .unwrap(); - assert!( - cat.scan_surrogates_for_collection( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0 - ) - .unwrap() - .is_empty() - ); - assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "orders"), - T0, - b"o1" - ) - .unwrap(), - Some(Surrogate::new(2)) - ); - // double-delete is a no-op - cat.delete_all_surrogates_for_collection( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - ) - .unwrap(); - } - - // ── Migration tests ─────────────────────────────────────────────────── - - fn insert_legacy_fwd(cat: &SystemCatalog, coll: &str, pk: &[u8], surrogate: u32) { - let txn = cat.db.begin_write().unwrap(); - { - let mut t = txn.open_table(SURROGATE_PK_LEGACY).unwrap(); - t.insert((coll, pk), surrogate).unwrap(); - let mut r = txn.open_table(SURROGATE_PK_REV_LEGACY).unwrap(); - r.insert((coll, surrogate), pk).unwrap(); - } - txn.commit().unwrap(); - } - - /// Read a v2 forward row directly (the v1→v2 migration target). - fn get_v2_fwd(cat: &SystemCatalog, db_id: u64, coll: &str, pk: &[u8]) -> Option { - let txn = cat.db.begin_read().unwrap(); - let table = txn.open_table(SURROGATE_PK_V2).unwrap(); - table.get((db_id, coll, pk)).unwrap().map(|v| v.value()) - } - - #[test] - fn fresh_boot_migration_is_noop() { - let (_dir, cat) = open_catalog(); - cat.migrate_surrogate_pk().unwrap(); - cat.migrate_surrogate_pk_v3().unwrap(); + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(1)) + .unwrap(); + cat.put_surrogate(coll("orders"), T0, b"o1", Surrogate::new(2)) + .unwrap(); + cat.delete_all_surrogates_for_collection(coll("users"), T0) + .unwrap(); assert!( - cat.scan_surrogates_for_collection( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0 - ) - .unwrap() - .is_empty() + cat.scan_surrogates_for_collection(coll("users"), T0) + .unwrap() + .is_empty() ); - } - - #[test] - fn pre_migration_boot_migrates_rows_v1_to_v2() { - let (_dir, cat) = open_catalog(); - insert_legacy_fwd(&cat, "users", b"alice", 7); - cat.migrate_surrogate_pk().unwrap(); assert_eq!( - get_v2_fwd(&cat, DatabaseId::DEFAULT.as_u64(), "users", b"alice"), - Some(7) - ); - } - - #[test] - fn already_migrated_boot_is_idempotent_v1_to_v2() { - let (_dir, cat) = open_catalog(); - // v2 row already exists, written under the default tenant in v3 … - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(7), - ) - .unwrap(); - // … and a v2 row to make `migrate_surrogate_pk` see a non-empty target. - { - let txn = cat.db.begin_write().unwrap(); - { - let mut t = txn.open_table(SURROGATE_PK_V2).unwrap(); - t.insert( - (DatabaseId::DEFAULT.as_u64(), "users", b"alice".as_slice()), - 7, - ) - .unwrap(); - } - txn.commit().unwrap(); - } - // also insert a legacy row - insert_legacy_fwd(&cat, "users", b"bob", 8); - // migration should be a no-op (v2 already non-empty) - cat.migrate_surrogate_pk().unwrap(); - assert_eq!( - get_v2_fwd(&cat, DatabaseId::DEFAULT.as_u64(), "users", b"bob"), + cat.get_pk_for_surrogate(coll("users"), T0, Surrogate::new(1)) + .unwrap(), None ); - } - - #[test] - fn migrate_v3_rekeys_v2_rows_under_default_tenant() { - let (_dir, cat) = open_catalog(); - // Seed a v2 row directly (as a pre-v3 boot would have). - { - let txn = cat.db.begin_write().unwrap(); - { - let mut f = txn.open_table(SURROGATE_PK_V2).unwrap(); - f.insert( - (DatabaseId::DEFAULT.as_u64(), "users", b"alice".as_slice()), - 7, - ) - .unwrap(); - let mut r = txn.open_table(SURROGATE_PK_REV_V2).unwrap(); - r.insert( - (DatabaseId::DEFAULT.as_u64(), "users", 7u32), - b"alice".as_slice(), - ) - .unwrap(); - } - txn.commit().unwrap(); - } - cat.migrate_surrogate_pk_v3().unwrap(); - // Backfilled rows land under the default user tenant (1), matching the - // default identity that wrote them pre-upgrade. - let default_tenant = TenantId::new(1); - assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - default_tenant, - b"alice" - ) - .unwrap(), - Some(Surrogate::new(7)) - ); assert_eq!( - cat.get_pk_for_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - default_tenant, - Surrogate::new(7) - ) - .unwrap(), - Some(b"alice".to_vec()) + cat.get_surrogate_for_pk(coll("orders"), T0, b"o1").unwrap(), + Some(Surrogate::new(2)) ); - } - - #[test] - fn migrate_v3_is_idempotent() { - let (_dir, cat) = open_catalog(); - // v3 already populated … - cat.put_surrogate( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice", - Surrogate::new(7), - ) - .unwrap(); - // … and a stale v2 row that must NOT overwrite it. - { - let txn = cat.db.begin_write().unwrap(); - { - let mut f = txn.open_table(SURROGATE_PK_V2).unwrap(); - f.insert( - (DatabaseId::DEFAULT.as_u64(), "users", b"alice".as_slice()), - 99, - ) - .unwrap(); - } - txn.commit().unwrap(); - } - cat.migrate_surrogate_pk_v3().unwrap(); + // double-delete is a no-op + cat.delete_all_surrogates_for_collection(coll("users"), T0) + .unwrap(); + // A re-created collection of the same name starts empty and binds. + cat.put_surrogate(coll("users"), T0, b"alice", Surrogate::new(5)) + .unwrap(); assert_eq!( - cat.get_surrogate_for_pk( - nodedb_types::CollectionKey::from_bare(DatabaseId::DEFAULT, "users"), - T0, - b"alice" - ) - .unwrap(), - Some(Surrogate::new(7)) + cat.get_surrogate_for_pk(coll("users"), T0, b"alice") + .unwrap(), + Some(Surrogate::new(5)) ); } } diff --git a/nodedb/src/control/security/catalog/surrogate_pk_batch.rs b/nodedb/src/control/security/catalog/surrogate_pk_batch.rs index a6a641711..44579b713 100644 --- a/nodedb/src/control/security/catalog/surrogate_pk_batch.rs +++ b/nodedb/src/control/security/catalog/surrogate_pk_batch.rs @@ -1,19 +1,21 @@ // SPDX-License-Identifier: BUSL-1.1 -//! Batched surrogate ↔ PK catalog ops for the `_system.surrogate_pk{,_rev}_v3` -//! tables. +//! Batched surrogate ↔ PK catalog ops on the surrogate PK map tables. //! //! A batch binds many keys of one collection in one redb write transaction. //! The transaction commits once with the database's default durability, so a -//! batch of N keys costs one durable sync instead of N. +//! batch of N keys costs one durable sync instead of N. A batch sorts its keys +//! and walks each table with one cursor (see [`super::surrogate_sorted`]). use nodedb_types::{CollectionKey, DatabaseId, Surrogate, TenantId}; -use redb::{ReadableDatabase, ReadableTable}; +use redb::ReadableDatabase; use crate::types::SurrogateBindEntry; use super::surrogate_hwm::{HWM_KEY, SURROGATE_HWM}; -use super::types::{SURROGATE_PK_REV_V3, SURROGATE_PK_V3, SystemCatalog, catalog_err}; +use super::surrogate_scope::{read_scope, rev_key, scope_or_create}; +use super::surrogate_sorted::{SortedKeys, insert_rev_sorted, merge_insert, probe}; +use super::types::{SURROGATE_PK, SURROGATE_PK_REV, SystemCatalog, catalog_err}; /// The outcome of binding one key in a batch. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -54,41 +56,44 @@ impl SystemCatalog { if rows.is_empty() && hwm.is_none() { return Ok(Vec::new()); } - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_write() .map_err(|e| catalog_err("surrogate_pk batch write txn", e))?; let mut out = Vec::with_capacity(rows.len()); { - let mut fwd = txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("open surrogate_pk", e))?; - let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) - .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - for &(pk_bytes, surrogate) in rows { - let existing = fwd - .get((db_id, tid, collection, pk_bytes)) - .map_err(|e| catalog_err("get surrogate_pk", e))? - .map(|v| Surrogate::new(v.value())); - if let Some(bound) = existing { - out.push(BatchBinding { - surrogate: bound, - inserted: false, + if !rows.is_empty() { + let scope = scope_or_create(&txn, key, tenant_id)?; + let keys = SortedKeys::new(scope, rows.iter().map(|&(pk, _)| pk)); + // A repeated key binds the surrogate of its first occurrence. + let carried = |u: usize| rows[keys.first(u)].1.as_u32(); + let before = { + let mut fwd = txn + .open_table(SURROGATE_PK) + .map_err(|e| catalog_err("open surrogate_pk", e))?; + merge_insert(&mut fwd, &keys, carried)? + }; + let reverse: Vec<(u128, &[u8])> = (0..keys.len()) + .filter(|&u| before[u].is_none()) + .map(|u| (rev_key(scope, carried(u)), keys.pk(u))) + .collect(); + let mut rev = txn + .open_table(SURROGATE_PK_REV) + .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; + insert_rev_sorted(&mut rev, reverse)?; + for i in 0..rows.len() { + let u = keys.unique_of(i); + out.push(match before[u] { + Some(bound) => BatchBinding { + surrogate: Surrogate::new(bound), + inserted: false, + }, + None => BatchBinding { + surrogate: Surrogate::new(carried(u)), + inserted: keys.first(u) == i, + }, }); - continue; } - fwd.insert((db_id, tid, collection, pk_bytes), surrogate.as_u32()) - .map_err(|e| catalog_err("insert surrogate_pk", e))?; - rev.insert((db_id, tid, collection, surrogate.as_u32()), pk_bytes) - .map_err(|e| catalog_err("insert surrogate_pk_rev", e))?; - out.push(BatchBinding { - surrogate, - inserted: true, - }); } if let Some(hwm) = hwm { let mut table = txn @@ -105,9 +110,10 @@ impl SystemCatalog { } /// Insert or overwrite many bindings of one collection in one write - /// transaction, as [`Self::put_surrogate`] does for one. A later row for a - /// key replaces an earlier one. Used where the caller's bindings are - /// authoritative, such as a Raft snapshot install. + /// transaction. A later row for a key replaces an earlier one. A key that + /// moves to a new surrogate loses the reverse row of its old one. Used + /// where the caller's bindings are authoritative, such as a Raft snapshot + /// install. pub fn overwrite_surrogates( &self, key: CollectionKey<'_>, @@ -117,26 +123,46 @@ impl SystemCatalog { if rows.is_empty() { return Ok(()); } - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_write() .map_err(|e| catalog_err("surrogate_pk overwrite txn", e))?; { - let mut fwd = txn - .open_table(SURROGATE_PK_V3) - .map_err(|e| catalog_err("open surrogate_pk", e))?; + let scope = scope_or_create(&txn, key, tenant_id)?; + let keys = SortedKeys::new(scope, rows.iter().map(|&(pk, _)| pk)); + // The last occurrence of a key is its authoritative surrogate. + let mut last: Vec = vec![0; keys.len()]; + for (i, &(_, surrogate)) in rows.iter().enumerate() { + last[keys.unique_of(i)] = surrogate.as_u32(); + } + let mut stale: Vec = Vec::new(); + { + let mut fwd = txn + .open_table(SURROGATE_PK) + .map_err(|e| catalog_err("open surrogate_pk", e))?; + for (u, &surrogate) in last.iter().enumerate() { + let old = fwd + .insert(keys.key(u), surrogate) + .map_err(|e| catalog_err("insert surrogate_pk", e))? + .map(|v| v.value()); + if let Some(old) = old.filter(|&old| old != surrogate) { + stale.push(rev_key(scope, old)); + } + } + } let mut rev = txn - .open_table(SURROGATE_PK_REV_V3) + .open_table(SURROGATE_PK_REV) .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; - for &(pk_bytes, surrogate) in rows { - fwd.insert((db_id, tid, collection, pk_bytes), surrogate.as_u32()) - .map_err(|e| catalog_err("insert surrogate_pk", e))?; - rev.insert((db_id, tid, collection, surrogate.as_u32()), pk_bytes) - .map_err(|e| catalog_err("insert surrogate_pk_rev", e))?; + for old in stale { + rev.remove(old) + .map_err(|e| catalog_err("remove surrogate_pk_rev", e))?; } + let reverse: Vec<(u128, &[u8])> = last + .iter() + .enumerate() + .map(|(u, &surrogate)| (rev_key(scope, surrogate), keys.pk(u))) + .collect(); + insert_rev_sorted(&mut rev, reverse)?; } txn.commit() .map_err(|e| catalog_err("surrogate_pk overwrite commit", e)) @@ -182,24 +208,21 @@ impl SystemCatalog { if pks.is_empty() { return Ok(Vec::new()); } - let db_id = key.database_id().as_u64(); - let tid = tenant_id.as_u64(); - let collection = key.name(); let txn = self .db .begin_read() .map_err(|e| catalog_err("surrogate_pk batch read txn", e))?; + let Some(scope) = read_scope(&txn, key, tenant_id)? else { + return Ok(vec![None; pks.len()]); + }; let table = txn - .open_table(SURROGATE_PK_V3) + .open_table(SURROGATE_PK) .map_err(|e| catalog_err("open surrogate_pk", e))?; - pks.iter() - .map(|pk_bytes| { - table - .get((db_id, tid, collection, *pk_bytes)) - .map_err(|e| catalog_err("get surrogate_pk", e)) - .map(|hit| hit.map(|v| Surrogate::new(v.value()))) - }) - .collect() + let keys = SortedKeys::new(scope, pks.iter().copied()); + let found = probe(&table, &keys)?; + Ok((0..pks.len()) + .map(|i| found[keys.unique_of(i)].map(Surrogate::new)) + .collect()) } } @@ -290,6 +313,70 @@ mod tests { assert_eq!(cat.get_surrogate_hwm().expect("hwm"), 7); } + /// A later batch that rewrites half the keys of an earlier one keeps + /// every earlier binding, binds the new keys, and leaves the reverse + /// table the exact inverse of the forward table. + #[test] + fn a_later_overlapping_batch_keeps_bindings_and_the_reverse_map_inverse() { + let (_dir, cat) = open(); + let pks: Vec> = (0..400u32).map(|i| i.to_be_bytes().to_vec()).collect(); + let first: Vec<(&[u8], Surrogate)> = pks[..300] + .iter() + .enumerate() + .rev() + .map(|(i, pk)| (pk.as_slice(), Surrogate::new(1 + i as u32))) + .collect(); + cat.put_surrogates(users(), T0, &first) + .expect("first batch"); + let second: Vec<(&[u8], Surrogate)> = pks[150..] + .iter() + .enumerate() + .map(|(i, pk)| (pk.as_slice(), Surrogate::new(10_000 + i as u32))) + .collect(); + let out = cat + .put_surrogates(users(), T0, &second) + .expect("second batch"); + for (i, binding) in out.iter().enumerate() { + let index = 150 + i; + if index < 300 { + assert_eq!(binding.surrogate, Surrogate::new(1 + index as u32)); + assert!(!binding.inserted); + } else { + assert_eq!(binding.surrogate, Surrogate::new(10_000 + i as u32)); + assert!(binding.inserted); + } + } + let scanned = cat + .scan_surrogates_for_collection(users(), T0) + .expect("scan"); + assert_eq!(scanned.len(), 400); + for (pk, surrogate) in &scanned { + assert_eq!( + cat.get_pk_for_surrogate(users(), T0, *surrogate) + .expect("rev"), + Some(pk.clone()) + ); + } + for losing in 10_000..10_150u32 { + assert_eq!( + cat.get_pk_for_surrogate(users(), T0, Surrogate::new(losing)) + .expect("rev"), + None, + "a losing carried surrogate writes no reverse row" + ); + } + let refs: Vec<&[u8]> = pks.iter().map(Vec::as_slice).collect(); + let read = cat + .get_surrogates_for_pks(users(), T0, &refs) + .expect("read"); + for (pk, hit) in refs.iter().zip(&read) { + assert_eq!( + *hit, + cat.get_surrogate_for_pk(users(), T0, pk).expect("get") + ); + } + } + #[test] fn a_batch_read_answers_in_key_order() { let (_dir, cat) = open(); diff --git a/nodedb/src/control/security/catalog/surrogate_scope.rs b/nodedb/src/control/security/catalog/surrogate_scope.rs new file mode 100644 index 000000000..67d09e367 --- /dev/null +++ b/nodedb/src/control/security/catalog/surrogate_scope.rs @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Surrogate scope ids and the key layout of the surrogate PK map tables. +//! +//! A scope id names one `(database, tenant, collection)` PK map. The forward +//! table keys by `scope_id` (big-endian) followed by the pk bytes. The reverse +//! table keys by `(scope_id << 32) | surrogate`. One scope's rows are one +//! contiguous key range in each table. + +use std::ops::Range; + +use nodedb_types::{CollectionKey, TenantId}; +use redb::{ReadTransaction, ReadableTable, WriteTransaction}; + +use super::types::{ + SURROGATE_PK, SURROGATE_PK_REV, SURROGATE_SCOPE_SEQ, SURROGATE_SCOPES, catalog_err, +}; + +/// Bytes of the scope id prefix of a forward-table key. +pub(super) const SCOPE_PREFIX_LEN: usize = 8; + +/// The first scope id the allocator hands out. +const FIRST_SCOPE: u64 = 1; + +/// The forward-table key of `pk` in `scope`. +pub(super) fn fwd_key(scope: u64, pk: &[u8]) -> Vec { + let mut key = Vec::with_capacity(SCOPE_PREFIX_LEN + pk.len()); + key.extend_from_slice(&scope.to_be_bytes()); + key.extend_from_slice(pk); + key +} + +/// The pk bytes of a forward-table key. A key shorter than its scope prefix +/// is corrupt. +pub(super) fn pk_of(key: &[u8]) -> crate::Result<&[u8]> { + key.get(SCOPE_PREFIX_LEN..).ok_or_else(|| { + catalog_err( + "surrogate_pk row", + format!("a {}-byte key is shorter than its scope prefix", key.len()), + ) + }) +} + +/// The forward-table key range that holds every row of `scope`. +pub(super) fn fwd_range(scope: u64) -> ([u8; SCOPE_PREFIX_LEN], [u8; SCOPE_PREFIX_LEN]) { + (scope.to_be_bytes(), scope.saturating_add(1).to_be_bytes()) +} + +/// The reverse-table key of `surrogate` in `scope`. +pub(super) fn rev_key(scope: u64, surrogate: u32) -> u128 { + (u128::from(scope) << 32) | u128::from(surrogate) +} + +/// The reverse-table key range that holds every row of `scope`. +pub(super) fn rev_range(scope: u64) -> Range { + rev_key(scope, 0)..(u128::from(scope) + 1) << 32 +} + +/// The scope id of `key` under `tenant_id` in a read transaction. `None` +/// when the collection has never bound a key. +pub(super) fn read_scope( + txn: &ReadTransaction, + key: CollectionKey<'_>, + tenant_id: TenantId, +) -> crate::Result> { + let scopes = txn + .open_table(SURROGATE_SCOPES) + .map_err(|e| catalog_err("open surrogate_scopes", e))?; + lookup_scope(&scopes, key, tenant_id) +} + +/// The scope id of `key` under `tenant_id` in a write transaction, without +/// creating one. +pub(super) fn write_scope( + txn: &WriteTransaction, + key: CollectionKey<'_>, + tenant_id: TenantId, +) -> crate::Result> { + let scopes = txn + .open_table(SURROGATE_SCOPES) + .map_err(|e| catalog_err("open surrogate_scopes", e))?; + lookup_scope(&scopes, key, tenant_id) +} + +/// The scope id of `key` under `tenant_id`, allocated in `txn` when the +/// collection has none. +pub(super) fn scope_or_create( + txn: &WriteTransaction, + key: CollectionKey<'_>, + tenant_id: TenantId, +) -> crate::Result { + let mut scopes = txn + .open_table(SURROGATE_SCOPES) + .map_err(|e| catalog_err("open surrogate_scopes", e))?; + if let Some(scope) = lookup_scope(&scopes, key, tenant_id)? { + return Ok(scope); + } + let mut seq = txn + .open_table(SURROGATE_SCOPE_SEQ) + .map_err(|e| catalog_err("open surrogate_scope_seq", e))?; + let next = seq + .get(()) + .map_err(|e| catalog_err("get surrogate_scope_seq", e))? + .map_or(FIRST_SCOPE, |v| v.value()); + // `fwd_range` and `rev_range` name the scope after this one, so the last + // id is never handed out. + let after = next + .checked_add(1) + .filter(|after| *after < u64::MAX) + .ok_or_else(|| catalog_err("surrogate scope allocate", "scope ids exhausted"))?; + seq.insert((), after) + .map_err(|e| catalog_err("insert surrogate_scope_seq", e))?; + scopes + .insert( + (key.database_id().as_u64(), tenant_id.as_u64(), key.name()), + next, + ) + .map_err(|e| catalog_err("insert surrogate_scopes", e))?; + Ok(next) +} + +/// Remove the scope row of `key` under `tenant_id` and return its id. The +/// map rows of the scope stay: the caller drains or re-points them. +pub(super) fn take_scope( + txn: &WriteTransaction, + key: CollectionKey<'_>, + tenant_id: TenantId, +) -> crate::Result> { + let mut scopes = txn + .open_table(SURROGATE_SCOPES) + .map_err(|e| catalog_err("open surrogate_scopes", e))?; + let removed = scopes + .remove((key.database_id().as_u64(), tenant_id.as_u64(), key.name())) + .map_err(|e| catalog_err("remove surrogate_scopes", e))? + .map(|v| v.value()); + Ok(removed) +} + +/// Name `scope` as the scope of `key` under `tenant_id`. The caller has +/// checked that `key` has no scope. +pub(super) fn assign_scope( + txn: &WriteTransaction, + key: CollectionKey<'_>, + tenant_id: TenantId, + scope: u64, +) -> crate::Result<()> { + let mut scopes = txn + .open_table(SURROGATE_SCOPES) + .map_err(|e| catalog_err("open surrogate_scopes", e))?; + scopes + .insert( + (key.database_id().as_u64(), tenant_id.as_u64(), key.name()), + scope, + ) + .map_err(|e| catalog_err("insert surrogate_scopes", e))?; + Ok(()) +} + +/// Remove every map row of `scope` from both tables and return the removed +/// `(pk_bytes, surrogate)` pairs in pk order. +pub(super) fn drain_scope( + txn: &WriteTransaction, + scope: u64, +) -> crate::Result, u32)>> { + let mut fwd = txn + .open_table(SURROGATE_PK) + .map_err(|e| catalog_err("open surrogate_pk", e))?; + let (lo, hi) = fwd_range(scope); + let mut rows = Vec::new(); + for row in fwd + .range(lo.as_slice()..hi.as_slice()) + .map_err(|e| catalog_err("range surrogate_pk", e))? + { + let (k, v) = row.map_err(|e| catalog_err("iter surrogate_pk row", e))?; + rows.push((pk_of(k.value())?.to_vec(), v.value())); + } + fwd.retain_in(lo.as_slice()..hi.as_slice(), |_, _| false) + .map_err(|e| catalog_err("drain surrogate_pk", e))?; + let mut rev = txn + .open_table(SURROGATE_PK_REV) + .map_err(|e| catalog_err("open surrogate_pk_rev", e))?; + rev.retain_in(rev_range(scope), |_, _| false) + .map_err(|e| catalog_err("drain surrogate_pk_rev", e))?; + Ok(rows) +} + +fn lookup_scope( + scopes: &impl ReadableTable<(u64, u64, &'static str), u64>, + key: CollectionKey<'_>, + tenant_id: TenantId, +) -> crate::Result> { + Ok(scopes + .get((key.database_id().as_u64(), tenant_id.as_u64(), key.name())) + .map_err(|e| catalog_err("get surrogate_scopes", e))? + .map(|v| v.value())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_scope_range_holds_its_keys_and_no_neighbour() { + let (lo, hi) = fwd_range(7); + let inside = fwd_key(7, b"\xff\xff"); + assert!(lo.as_slice() <= inside.as_slice() && inside.as_slice() < hi.as_slice()); + assert!(fwd_key(8, b"").as_slice() >= hi.as_slice()); + let rev = rev_range(7); + assert!(rev.contains(&rev_key(7, u32::MAX))); + assert!(!rev.contains(&rev_key(8, 0))); + assert!(!rev.contains(&rev_key(6, u32::MAX))); + } +} diff --git a/nodedb/src/control/security/catalog/surrogate_sorted.rs b/nodedb/src/control/security/catalog/surrogate_sorted.rs new file mode 100644 index 000000000..31ae4b55c --- /dev/null +++ b/nodedb/src/control/security/catalog/surrogate_sorted.rs @@ -0,0 +1,370 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Sorted batch access to the surrogate PK map tables. +//! +//! A batch sorts its keys once. A read walks the forward table with one +//! cursor, and a write splices every new row in with one cursor. A cursor +//! seeks again only past a run of [`RESEEK_AFTER`] rows the batch does not +//! name. A batch of keys absent from the table costs one seek. + +use std::cmp::Ordering; +use std::ops::Bound; + +use redb::{ReadableTable, Table}; + +use super::surrogate_scope::SCOPE_PREFIX_LEN; +use super::types::catalog_err; + +/// Rows a cursor steps past before it seeks to the next batch key instead. +const RESEEK_AFTER: usize = 16; + +/// The distinct forward-table keys of one batch, in key order. +pub(super) struct SortedKeys { + /// Every input key, scope prefix included, in input order. + buf: Vec, + /// Per distinct key, in key order: its byte range in `buf`. + spans: Vec<(usize, usize)>, + /// Per distinct key, in key order: the input index of its first occurrence. + first: Vec, + /// Per input index: the position of its key in `spans`. + unique_of: Vec, +} + +impl SortedKeys { + /// Sort and deduplicate the forward-table keys of `pks` in `scope`. + pub(super) fn new<'p>(scope: u64, pks: impl ExactSizeIterator) -> Self { + let n = pks.len(); + let mut buf = Vec::new(); + let mut input_spans = Vec::with_capacity(n); + let prefix = scope.to_be_bytes(); + for pk in pks { + let start = buf.len(); + buf.extend_from_slice(&prefix); + buf.extend_from_slice(pk); + input_spans.push((start, buf.len())); + } + let key = |i: usize| &buf[input_spans[i].0..input_spans[i].1]; + let mut order: Vec = (0..n).collect(); + // Equal keys keep input order, so the first occurrence sorts first. + order.sort_unstable_by(|&a, &b| key(a).cmp(key(b)).then(a.cmp(&b))); + let mut spans: Vec<(usize, usize)> = Vec::with_capacity(n); + let mut first: Vec = Vec::with_capacity(n); + let mut unique_of = vec![0usize; n]; + for &i in &order { + let is_repeat = first.last().is_some_and(|&f| key(f) == key(i)); + if !is_repeat { + spans.push(input_spans[i]); + first.push(i); + } + unique_of[i] = spans.len() - 1; + } + Self { + buf, + spans, + first, + unique_of, + } + } + + /// How many distinct keys the batch holds. + pub(super) fn len(&self) -> usize { + self.spans.len() + } + + /// Distinct key `u`, scope prefix included. + pub(super) fn key(&self, u: usize) -> &[u8] { + let (start, end) = self.spans[u]; + &self.buf[start..end] + } + + /// The pk bytes of distinct key `u`, without the scope prefix. + pub(super) fn pk(&self, u: usize) -> &[u8] { + &self.key(u)[SCOPE_PREFIX_LEN..] + } + + /// The input index of the first occurrence of distinct key `u`. + pub(super) fn first(&self, u: usize) -> usize { + self.first[u] + } + + /// The distinct key position of input index `i`. + pub(super) fn unique_of(&self, i: usize) -> usize { + self.unique_of[i] + } +} + +/// The value bound to each distinct key of `keys`, `None` for an absent key. +pub(super) fn probe( + table: &impl ReadableTable<&'static [u8], u32>, + keys: &SortedKeys, +) -> crate::Result>> { + let n = keys.len(); + let mut out = vec![None; n]; + let mut u = 0; + while u < n { + let mut cursor = table + .lower_bound(Bound::Included(keys.key(u))) + .map_err(|e| catalog_err("seek surrogate_pk", e))?; + let mut skipped = 0; + loop { + let Some((k, v)) = cursor + .next() + .map_err(|e| catalog_err("step surrogate_pk", e))? + else { + return Ok(out); + }; + let existing = k.value(); + // Batch keys below the next stored key are absent. + while u < n && keys.key(u) < existing { + u += 1; + skipped = 0; + } + if u == n { + return Ok(out); + } + if keys.key(u) == existing { + out[u] = Some(v.value()); + u += 1; + skipped = 0; + if u == n { + return Ok(out); + } + } else { + skipped += 1; + if skipped >= RESEEK_AFTER { + break; + } + } + } + } + Ok(out) +} + +/// Bind each absent distinct key of `keys` to `value_of(u)` first-wins, and +/// return the value each key held before: `Some` for a key already bound, +/// `None` for a key this call inserted. +pub(super) fn merge_insert( + table: &mut Table<'_, &'static [u8], u32>, + keys: &SortedKeys, + value_of: impl Fn(usize) -> u32, +) -> crate::Result>> { + let n = keys.len(); + let mut out = vec![None; n]; + let mut u = 0; + while u < n { + let mut cursor = table + .lower_bound_mut(Bound::Included(keys.key(u))) + .map_err(|e| catalog_err("seek surrogate_pk", e))?; + let mut skipped = 0; + while u < n { + let next = cursor + .peek_next() + .map_err(|e| catalog_err("peek surrogate_pk", e))? + .map(|(k, v)| (keys.key(u).cmp(k.value()), v.value())); + match next { + // The key sorts before the next stored row, or past the last + // one: it is absent. + None | Some((Ordering::Less, _)) => { + cursor + .insert_before(keys.key(u), value_of(u)) + .map_err(|e| catalog_err("insert surrogate_pk", e))?; + u += 1; + skipped = 0; + } + Some((Ordering::Equal, bound)) => { + out[u] = Some(bound); + u += 1; + skipped = 0; + cursor + .next() + .map_err(|e| catalog_err("step surrogate_pk", e))?; + } + Some((Ordering::Greater, _)) => { + skipped += 1; + if skipped >= RESEEK_AFTER { + break; + } + cursor + .next() + .map_err(|e| catalog_err("step surrogate_pk", e))?; + } + } + } + cursor + .close() + .map_err(|e| catalog_err("close surrogate_pk cursor", e))?; + } + Ok(out) +} + +/// Write each `(rev_key, pk_bytes)` row into the reverse table, overwriting +/// a row already stored under the key. A key repeated in `rows` keeps its +/// last pk. +pub(super) fn insert_rev_sorted( + table: &mut Table<'_, u128, &'static [u8]>, + mut rows: Vec<(u128, &[u8])>, +) -> crate::Result<()> { + rows.sort_by_key(|&(key, _)| key); + rows.dedup_by(|later, earlier| { + let repeat = later.0 == earlier.0; + if repeat { + earlier.1 = later.1; + } + repeat + }); + let mut overwrites: Vec<(u128, &[u8])> = Vec::new(); + let mut i = 0; + while i < rows.len() { + let mut cursor = table + .lower_bound_mut(Bound::Included(rows[i].0)) + .map_err(|e| catalog_err("seek surrogate_pk_rev", e))?; + let mut skipped = 0; + while i < rows.len() { + let (key, pk) = rows[i]; + let next = cursor + .peek_next() + .map_err(|e| catalog_err("peek surrogate_pk_rev", e))? + .map(|(k, _)| key.cmp(&k.value())); + match next { + None | Some(Ordering::Less) => { + cursor + .insert_before(key, pk) + .map_err(|e| catalog_err("insert surrogate_pk_rev", e))?; + i += 1; + skipped = 0; + } + Some(Ordering::Equal) => { + overwrites.push((key, pk)); + i += 1; + skipped = 0; + cursor + .next() + .map_err(|e| catalog_err("step surrogate_pk_rev", e))?; + } + Some(Ordering::Greater) => { + skipped += 1; + if skipped >= RESEEK_AFTER { + break; + } + cursor + .next() + .map_err(|e| catalog_err("step surrogate_pk_rev", e))?; + } + } + } + cursor + .close() + .map_err(|e| catalog_err("close surrogate_pk_rev cursor", e))?; + } + for (key, pk) in overwrites { + table + .insert(key, pk) + .map_err(|e| catalog_err("overwrite surrogate_pk_rev", e))?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::security::catalog::surrogate_scope::fwd_key; + use redb::{ReadableDatabase, TableDefinition}; + + const FWD: TableDefinition<&[u8], u32> = TableDefinition::new("fwd"); + const REV: TableDefinition = TableDefinition::new("rev"); + + fn db() -> (tempfile::TempDir, redb::Database) { + let dir = tempfile::tempdir().expect("tempdir"); + let db = redb::Database::create(dir.path().join("t.redb")).expect("create"); + (dir, db) + } + + #[test] + fn keys_sort_dedupe_and_keep_the_first_occurrence() { + let pks: [&[u8]; 4] = [b"b", b"a", b"b", b"c"]; + let keys = SortedKeys::new(3, pks.iter().copied()); + assert_eq!(keys.len(), 3); + assert_eq!(keys.pk(0), b"a"); + assert_eq!(keys.pk(1), b"b"); + assert_eq!(keys.first(1), 0, "the first `b` is input 0"); + assert_eq!(keys.unique_of(2), 1); + assert_eq!(&keys.key(2)[..SCOPE_PREFIX_LEN], &3u64.to_be_bytes()); + } + + /// Stored rows interleave with batch keys, past the reseek distance too. + #[test] + fn merge_is_first_wins_and_probe_agrees_with_point_gets() { + let (_dir, db) = db(); + let stored: Vec> = (0..200u32) + .filter(|i| i % 3 == 0) + .map(|i| fwd_key(0, format!("k{i:04}").as_bytes())) + .collect(); + let txn = db.begin_write().expect("txn"); + { + let mut fwd = txn.open_table(FWD).expect("open"); + for (n, pk) in stored.iter().enumerate() { + fwd.insert(pk.as_slice(), n as u32 + 1_000).expect("seed"); + } + } + txn.commit().expect("commit"); + + let batch: Vec> = (0..200u32) + .filter(|i| i % 2 == 0 || *i > 150) + .map(|i| format!("k{i:04}").into_bytes()) + .collect(); + let raw: Vec<&[u8]> = batch.iter().map(Vec::as_slice).collect(); + let keys = SortedKeys::new(0, raw.iter().copied()); + let txn = db.begin_write().expect("txn"); + let before = { + let mut fwd = txn.open_table(FWD).expect("open"); + merge_insert(&mut fwd, &keys, |u| u as u32).expect("merge") + }; + txn.commit().expect("commit"); + + let txn = db.begin_read().expect("read"); + let fwd = txn.open_table(FWD).expect("open"); + let probed = probe(&fwd, &keys).expect("probe"); + for u in 0..keys.len() { + let got = fwd.get(keys.key(u)).expect("get").map(|v| v.value()); + assert_eq!(probed[u], got, "probe of key {u}"); + let was_stored = stored.iter().any(|s| s.as_slice() == keys.key(u)); + assert_eq!(before[u].is_some(), was_stored, "key {u}"); + if !was_stored { + assert_eq!(got, Some(u as u32), "an absent key takes its new value"); + } + } + } + + #[test] + fn a_probe_of_an_empty_table_finds_nothing() { + let (_dir, db) = db(); + let txn = db.begin_write().expect("txn"); + txn.open_table(FWD).expect("create"); + txn.commit().expect("commit"); + let pks: [&[u8]; 2] = [b"x", b"y"]; + let keys = SortedKeys::new(1, pks.iter().copied()); + let txn = db.begin_read().expect("read"); + let fwd = txn.open_table(FWD).expect("open"); + assert_eq!(probe(&fwd, &keys).expect("probe"), vec![None, None]); + } + + #[test] + fn reverse_rows_insert_in_order_and_overwrite_a_stored_key() { + let (_dir, db) = db(); + let txn = db.begin_write().expect("txn"); + { + let mut rev = txn.open_table(REV).expect("open"); + rev.insert(5u128, b"old".as_slice()).expect("seed"); + let rows: Vec<(u128, &[u8])> = + vec![(9, b"nine"), (5, b"five"), (1, b"one"), (9, b"last")]; + insert_rev_sorted(&mut rev, rows).expect("insert"); + } + txn.commit().expect("commit"); + let txn = db.begin_read().expect("read"); + let rev = txn.open_table(REV).expect("open"); + let get = |k: u128| rev.get(k).expect("get").map(|v| v.value().to_vec()); + assert_eq!(get(1), Some(b"one".to_vec())); + assert_eq!(get(5), Some(b"five".to_vec())); + assert_eq!(get(9), Some(b"last".to_vec())); + } +} diff --git a/nodedb/src/control/security/catalog/system_catalog.rs b/nodedb/src/control/security/catalog/system_catalog.rs index e410a1029..76ac11a4b 100644 --- a/nodedb/src/control/security/catalog/system_catalog.rs +++ b/nodedb/src/control/security/catalog/system_catalog.rs @@ -41,6 +41,9 @@ pub struct SystemCatalog { pub(super) fail_next_surrogate_write: Arc, #[cfg(test)] pub(super) fail_next_redo_snapshot_owed_write: Arc, + /// Audit rows the DML position read decoded, so a test can tell a range read from a full scan. + #[cfg(test)] + pub(super) dml_audit_rows_read: Arc, } impl crate::storage::RedbBacked for SystemCatalog { @@ -83,6 +86,8 @@ impl SystemCatalog { fail_next_surrogate_write: Arc::new(std::sync::atomic::AtomicBool::new(false)), #[cfg(test)] fail_next_redo_snapshot_owed_write: Arc::new(std::sync::atomic::AtomicBool::new(false)), + #[cfg(test)] + dml_audit_rows_read: Arc::new(std::sync::atomic::AtomicU64::new(0)), }; catalog.bootstrap_default_database()?; catalog.reload_event_definitions()?; @@ -116,6 +121,8 @@ impl SystemCatalog { fail_next_surrogate_write: Arc::new(std::sync::atomic::AtomicBool::new(false)), #[cfg(test)] fail_next_redo_snapshot_owed_write: Arc::new(std::sync::atomic::AtomicBool::new(false)), + #[cfg(test)] + dml_audit_rows_read: Arc::new(std::sync::atomic::AtomicU64::new(0)), }; catalog.bootstrap_default_database()?; catalog.reload_event_definitions()?; diff --git a/nodedb/src/control/security/catalog/tables.rs b/nodedb/src/control/security/catalog/tables.rs index 24d55a025..7721d3c0c 100644 --- a/nodedb/src/control/security/catalog/tables.rs +++ b/nodedb/src/control/security/catalog/tables.rs @@ -23,6 +23,12 @@ pub(super) const TENANTS: TableDefinition<&str, &[u8]> = TableDefinition::new("_ pub(super) const AUDIT_LOG: TableDefinition<&[u8], &[u8]> = TableDefinition::new("_system.audit_log"); +/// Table: DML audit position index. +/// Key: event core (u32 BE) | event LSN (u64 BE) | audit seq (u64 BE). Value: empty. +/// It holds one row per `DmlAudit` entry whose detail names an event key. +pub(super) const AUDIT_DML_POSITIONS: TableDefinition<&[u8], &[u8]> = + TableDefinition::new("_system.audit_dml_positions"); + /// Table: role_name -> MessagePack-serialized custom role record. pub(super) const ROLES: TableDefinition<&str, &[u8]> = TableDefinition::new("_system.roles"); @@ -109,46 +115,32 @@ pub(super) const ARRAYS: TableDefinition<&str, &[u8]> = TableDefinition::new("_s // ── Surrogate PK map ────────────────────────────────────────────────── -/// Table (legacy): `(collection, encoded_pk_bytes)` -> `Surrogate` (u32 LE). -/// Used only by the idempotent migration that prefixes rows with database_id. -pub(super) const SURROGATE_PK_LEGACY: TableDefinition<(&str, &[u8]), u32> = - TableDefinition::new("_system.surrogate_pk"); +/// Table: `(database_id, tenant_id, collection)` -> surrogate scope id. +/// +/// A scope id names one collection's PK map. The map tables key by the id, +/// so a key compare reads no collection name. +pub(super) const SURROGATE_SCOPES: TableDefinition<(u64, u64, &str), u64> = + TableDefinition::new("_system.surrogate_scopes"); -/// Table (legacy, v2): `(database_id: u64, collection, encoded_pk_bytes)` -> -/// `Surrogate` (u32 LE). Database-scoped but tenant-blind. Read only by the -/// idempotent migration that re-keys rows with `tenant_id` prepended. -pub(super) const SURROGATE_PK_V2: TableDefinition<(u64, &str, &[u8]), u32> = - TableDefinition::new("_system.surrogate_pk_v2"); +/// Singleton table: the next unused surrogate scope id. Ids start at 1. +pub(super) const SURROGATE_SCOPE_SEQ: TableDefinition<(), u64> = + TableDefinition::new("_system.surrogate_scope_seq"); -/// Table: `(database_id: u64, tenant_id: u64, collection, encoded_pk_bytes)` -> -/// `Surrogate` (u32 LE). +/// Table: `scope_id` (8 bytes big-endian) followed by the encoded pk bytes -> +/// `Surrogate` (u32). /// -/// Forward direction of the PK ↔ Surrogate map with database + tenant scoping -/// prepended. Every successful `assign_surrogate(database_id, tenant_id, -/// collection, pk)` writes both the forward and reverse rows atomically in one -/// redb txn. -pub(super) const SURROGATE_PK_V3: TableDefinition<(u64, u64, &str, &[u8]), u32> = - TableDefinition::new("_system.surrogate_pk_v3"); - -/// Table (legacy): `(collection, surrogate)` -> encoded pk bytes. -/// Used only by the idempotent migration that prefixes rows with database_id. -#[cfg(test)] -pub(super) const SURROGATE_PK_REV_LEGACY: TableDefinition<(&str, u32), &[u8]> = - TableDefinition::new("_system.surrogate_pk_rev"); - -/// Table (legacy, v2): `(database_id: u64, collection, surrogate)` -> encoded pk -/// bytes. Reverse direction of `_system.surrogate_pk_v2`. Read only by the -/// idempotent migration that re-keys rows with `tenant_id` prepended. -pub(super) const SURROGATE_PK_REV_V2: TableDefinition<(u64, &str, u32), &[u8]> = - TableDefinition::new("_system.surrogate_pk_rev_v2"); - -/// Table: `(database_id: u64, tenant_id: u64, collection, surrogate)` -> encoded -/// pk bytes. +/// Forward direction of the PK ↔ Surrogate map. The key is one byte string, +/// so a compare is one `memcmp`, and one scope's rows are one key range. +/// Every binding writes the forward and the reverse row in one redb txn. +pub(super) const SURROGATE_PK: TableDefinition<&[u8], u32> = + TableDefinition::new("_system.surrogate_pk_v4"); + +/// Table: `(scope_id << 32) | surrogate` -> encoded pk bytes. /// -/// Reverse direction of `_system.surrogate_pk_v3`. Scoped per `(database_id, -/// tenant_id)`. -pub(super) const SURROGATE_PK_REV_V3: TableDefinition<(u64, u64, &str, u32), &[u8]> = - TableDefinition::new("_system.surrogate_pk_rev_v3"); +/// Reverse direction of `_system.surrogate_pk_v4`. One scope's rows are one +/// key range, in surrogate order. +pub(super) const SURROGATE_PK_REV: TableDefinition = + TableDefinition::new("_system.surrogate_pk_rev_v4"); // ── Event Plane ─────────────────────────────────────────────────────── diff --git a/nodedb/src/control/security/catalog/types.rs b/nodedb/src/control/security/catalog/types.rs index e4d4dbac3..eab323c75 100644 --- a/nodedb/src/control/security/catalog/types.rs +++ b/nodedb/src/control/security/catalog/types.rs @@ -35,18 +35,16 @@ pub use super::system_catalog::SystemCatalog; // Re-exported so existing `super::types::{TABLE_FOO, ...}` imports keep // working unchanged. -#[cfg(test)] -pub(super) use super::tables::SURROGATE_PK_REV_LEGACY; pub(super) use super::tables::{ - ALERT_RULES, API_KEYS, ARRAYS, AUDIT_LOG, AUTH_USERS, BLACKLIST, CHANGE_STREAMS, CHECKPOINTS, - CLONE_COPYUPS, CLONE_KV_TOMBSTONES, CLONE_LINEAGE, CLONE_TOMBSTONES, COLLECTIONS, COLUMN_STATS, - CONSUMER_GROUPS, CONTINUOUS_AGGREGATES, CUSTOM_TYPES, DATABASE_GRANTS, DATABASE_HWM, - DATABASE_QUOTAS, DATABASES, DATABASES_BY_NAME, DEPENDENCIES, FUNCTIONS, INDEX_REGISTRY, - L2_CLEANUP_QUEUE, LOCKOUT_STATE, MATERIALIZED_VIEWS, METADATA, MIRROR_COLLECTION_MAP, - MIRROR_LAG, ORG_MEMBERS, ORGS, OWNERS, PENDING_RECLAIM, PERMISSIONS, PROCEDURES, - RETENTION_POLICIES, ROLES, SCHEDULES, SCOPE_GRANTS, SCOPE_QUOTAS, SCOPES, SEQUENCE_STATE, - SEQUENCES, STREAMING_MVS, SURROGATE_PK_LEGACY, SURROGATE_PK_REV_V2, SURROGATE_PK_REV_V3, - SURROGATE_PK_V2, SURROGATE_PK_V3, SYNONYM_GROUPS, TENANT_QUOTAS, TENANTS, TOPIC_MESSAGES, + ALERT_RULES, API_KEYS, ARRAYS, AUDIT_DML_POSITIONS, AUDIT_LOG, AUTH_USERS, BLACKLIST, + CHANGE_STREAMS, CHECKPOINTS, CLONE_COPYUPS, CLONE_KV_TOMBSTONES, CLONE_LINEAGE, + CLONE_TOMBSTONES, COLLECTIONS, COLUMN_STATS, CONSUMER_GROUPS, CONTINUOUS_AGGREGATES, + CUSTOM_TYPES, DATABASE_GRANTS, DATABASE_HWM, DATABASE_QUOTAS, DATABASES, DATABASES_BY_NAME, + DEPENDENCIES, FUNCTIONS, INDEX_REGISTRY, L2_CLEANUP_QUEUE, LOCKOUT_STATE, MATERIALIZED_VIEWS, + METADATA, MIRROR_COLLECTION_MAP, MIRROR_LAG, ORG_MEMBERS, ORGS, OWNERS, PENDING_RECLAIM, + PERMISSIONS, PROCEDURES, RETENTION_POLICIES, ROLES, SCHEDULES, SCOPE_GRANTS, SCOPE_QUOTAS, + SCOPES, SEQUENCE_STATE, SEQUENCES, STREAMING_MVS, SURROGATE_PK, SURROGATE_PK_REV, + SURROGATE_SCOPE_SEQ, SURROGATE_SCOPES, SYNONYM_GROUPS, TENANT_QUOTAS, TENANTS, TOPIC_MESSAGES, TOPICS_EP, TRIGGERS, USERS, VECTOR_INDEX_PARAMS, VECTOR_MODEL_METADATA, WAL_TOMBSTONES, WASM_MODULES, }; diff --git a/nodedb/src/control/security/identity/plan_permission.rs b/nodedb/src/control/security/identity/plan_permission.rs index 58e099451..673de65c3 100644 --- a/nodedb/src/control/security/identity/plan_permission.rs +++ b/nodedb/src/control/security/identity/plan_permission.rs @@ -384,7 +384,6 @@ pub fn required_permission(plan: &crate::bridge::envelope::PhysicalPlan) -> Perm PhysicalPlan::Array( ArrayOp::Slice { .. } | ArrayOp::SurrogateBitmapScan { .. } - | ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. }, ) => Permission::Read, @@ -401,12 +400,17 @@ pub fn required_permission(plan: &crate::bridge::envelope::PhysicalPlan) -> Perm // ClusterArray mirrors the local ArrayOp permission model. PhysicalPlan::ClusterArray( nodedb_physical::physical_plan::ClusterArrayOp::Slice { .. } - | nodedb_physical::physical_plan::ClusterArrayOp::Agg { .. }, + | nodedb_physical::physical_plan::ClusterArrayOp::Agg { .. } + | nodedb_physical::physical_plan::ClusterArrayOp::Elementwise { .. }, ) => Permission::Read, PhysicalPlan::ClusterArray( nodedb_physical::physical_plan::ClusterArrayOp::Put { .. } | nodedb_physical::physical_plan::ClusterArrayOp::Delete { .. }, ) => Permission::Write, + PhysicalPlan::ClusterArray( + nodedb_physical::physical_plan::ClusterArrayOp::Flush { .. } + | nodedb_physical::physical_plan::ClusterArrayOp::Compact { .. }, + ) => Permission::Admin, PhysicalPlan::ClusterEvent(_) => Permission::Read, // Calvin cross-shard execution batches: dispatched internally, treated as Write. diff --git a/nodedb/src/control/security/jwks/registry/routing.rs b/nodedb/src/control/security/jwks/registry/routing.rs index 26dd1621f..be4759324 100644 --- a/nodedb/src/control/security/jwks/registry/routing.rs +++ b/nodedb/src/control/security/jwks/registry/routing.rs @@ -4,11 +4,37 @@ //! configured static providers. use crate::config::auth::JwtProviderConfig; -use crate::control::security::jwt::JwtError; +use crate::control::security::bearer::BearerAttempt; +use crate::control::security::jwt::{JwtClaims, JwtError}; +use crate::control::security::util::base64_url_decode; +use crate::types::TenantId; use super::state::JwksRegistry; impl JwksRegistry { + /// Record into `attempt` the `sub` a token claims and the provider its `iss` and `aud` route to. + /// + /// Nothing here is verified, so the values only label the audit entry. + /// A token whose payload does not decode leaves `attempt` unchanged. + pub(crate) fn label_unverified(&self, token: &str, attempt: &mut BearerAttempt<'_>) { + let mut parts = token.split('.'); + let (Some(_header), Some(payload), Some(_signature), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return; + }; + let Some(payload_bytes) = base64_url_decode(payload) else { + return; + }; + let Ok(claims) = crate::util::bounded_json::from_slice::(&payload_bytes) else { + return; + }; + attempt.set_claimed_sub(&claims.sub); + if let Ok(provider) = self.find_provider(&claims.iss, &claims.aud) { + attempt.set_provider(&provider.name, Some(TenantId::new(provider.tenant_id))); + } + } + /// Find the provider matching a token's issuer and audience. /// /// Static configuration validation ensures a route is unique. A provider diff --git a/nodedb/src/control/security/jwt_policy/roles.rs b/nodedb/src/control/security/jwt_policy/roles.rs index abf93e64c..f78410488 100644 --- a/nodedb/src/control/security/jwt_policy/roles.rs +++ b/nodedb/src/control/security/jwt_policy/roles.rs @@ -8,10 +8,12 @@ //! custom role grants something only while it is defined in the tenant the //! provider is bound to. An identity holding an undefined role would hold //! nothing, and a just-in-time provisioned record would store it. So the -//! login is refused with a typed error naming the role, and the refusal is -//! recorded in the audit log, before anything is provisioned. +//! login is refused with a typed error naming the role, before anything is +//! provisioned. +//! +//! The bearer route that called the policy records the refusal. +//! It writes one `AuthFailure` with the provider, the peer address, and this error as the reason. -use crate::control::security::audit::AuditEvent; use crate::control::security::identity::AuthenticatedIdentity; use crate::control::security::role_assignment::{self, RoleRefusal}; use crate::control::state::SharedState; @@ -36,12 +38,6 @@ pub fn refuse_undefined_roles( RoleRefusal::Undefined { name } => name, other => other.to_string(), }; - state.audit_record( - AuditEvent::AuthFailure, - Some(tenant_id), - &identity.username, - &format!("external login refused: role \"{role}\" is not defined in tenant {tenant_id}"), - ); Err(crate::Error::ExternalRoleUndefined { subject: identity.username.clone(), role, @@ -55,6 +51,7 @@ mod tests { use super::*; use crate::bridge::dispatch::Dispatcher; + use crate::control::security::audit::AuditEvent; use crate::control::security::identity::{AuthMethod, Role}; use crate::types::TenantId; use crate::wal::WalManager; @@ -91,7 +88,7 @@ mod tests { } #[test] - fn an_undefined_claimed_role_refuses_the_login_and_is_audited() { + fn an_undefined_claimed_role_refuses_the_login() { let dir = tempfile::tempdir().expect("tempdir"); let state = state(&dir); let identity = external(5, vec![Role::ReadOnly, Role::Custom("ghost".into())]); @@ -110,8 +107,8 @@ mod tests { } let failures = auth_failures(&state); assert!( - failures.iter().any(|detail| detail.contains("\"ghost\"")), - "the refusal must be audited: {failures:?}" + failures.is_empty(), + "the bearer route records the refusal, so the policy writes nothing: {failures:?}" ); } diff --git a/nodedb/src/control/security/mod.rs b/nodedb/src/control/security/mod.rs index 87c455765..ca20d0c6a 100644 --- a/nodedb/src/control/security/mod.rs +++ b/nodedb/src/control/security/mod.rs @@ -6,6 +6,7 @@ pub mod auth_apikey; pub mod auth_context; pub mod auth_fence; pub mod auth_lease; +pub mod bearer; pub mod blacklist; pub mod buses; pub mod catalog; diff --git a/nodedb/src/control/security/oidc/verify.rs b/nodedb/src/control/security/oidc/verify.rs index 5188cd385..b3176e1a8 100644 --- a/nodedb/src/control/security/oidc/verify.rs +++ b/nodedb/src/control/security/oidc/verify.rs @@ -13,6 +13,7 @@ use nodedb_types::id::DatabaseId; use tracing::debug; +use crate::control::security::bearer::{BearerAttempt, BearerMethod, BearerRefusal}; use crate::control::security::catalog::StoredOidcProvider; use crate::control::security::identity::database_set::DatabaseSet; use crate::control::security::identity::{ @@ -45,64 +46,104 @@ use super::claim_mapping::apply_claim_mapping; /// — mirroring what the HTTP bearer-token path already does. Authority /// (superuser, tenant, roles) always comes from the `AuthenticatedIdentity` /// built here, never from the claims themselves. +/// +/// Every refusal records one `AuthFailure` and every login records one +/// `AuthSuccess`, through [`BearerAttempt`]. `peer_addr` is the audit source. +/// `session_id` is the id the accepted session's `AuthContext` carries. pub async fn verify_bearer_token( state: &SharedState, token: &str, + peer_addr: &str, + session_id: &str, ) -> crate::Result<(AuthenticatedIdentity, VerifiedJwtClaims)> { + let mut attempt = BearerAttempt::new(BearerMethod::OidcBearer, peer_addr, session_id); + match verify_catalog_token(state, token, &mut attempt).await { + Ok((identity, verified)) => { + attempt.accept(state, &identity, &verified.claims().sub); + Ok((identity, verified)) + } + Err(refusal) => { + attempt.refuse(state, &refusal.reason); + Err(*refusal.error) + } + } +} + +/// The verification steps of [`verify_bearer_token`]. +/// +/// Records the claimed `sub` and the selected provider into `attempt` as each one is learned. +async fn verify_catalog_token( + state: &SharedState, + token: &str, + attempt: &mut BearerAttempt<'_>, +) -> Result<(AuthenticatedIdentity, VerifiedJwtClaims), BearerRefusal> { // 1. Decode payload (no sig) to read `iss` and `aud`. let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { - return Err(jwt_error_to_crate_error(JwtError::MalformedToken)); + return Err(token_refusal(JwtError::MalformedToken)); } - let payload_bytes = base64_url_decode(parts[1]) - .ok_or_else(|| jwt_error_to_crate_error(JwtError::DecodingError))?; + let payload_bytes = + base64_url_decode(parts[1]).ok_or_else(|| token_refusal(JwtError::DecodingError))?; let claims: crate::control::security::jwt::JwtClaims = crate::util::bounded_json::from_slice(&payload_bytes) - .map_err(|_| jwt_error_to_crate_error(JwtError::InvalidClaims))?; + .map_err(|_| token_refusal(JwtError::InvalidClaims))?; + attempt.set_claimed_sub(&claims.sub); // 2. Route by both unverified issuer and audience. These values only select // the verification key; the same constraints are checked again after // signature and time validation. let catalog = state.credentials.catalog(); - let provider = select_catalog_provider( - catalog - .list_oidc_providers() - .map_err(|_| jwt_error_to_crate_error(JwtError::InvalidIssuer))?, - &claims.iss, - &claims.aud, - ) - .map_err(jwt_error_to_crate_error)?; + let providers = catalog.list_oidc_providers().map_err(|e| { + BearerRefusal::new( + format!("OIDC provider catalog read failed: {e}"), + client_refusal_error(), + ) + })?; + let provider = + select_catalog_provider(providers, &claims.iss, &claims.aud).map_err(token_refusal)?; + attempt.set_provider( + &provider.provider_name, + provider.tenant_id.map(TenantId::new), + ); // 3. Verify signature via JwksRegistry using the catalog-provided JWKS URI. - let jwks = state - .jwks_registry - .as_ref() - .ok_or_else(|| jwt_error_to_crate_error(JwtError::UnsupportedAlgorithm))?; + let jwks = state.jwks_registry.as_ref().ok_or_else(|| { + BearerRefusal::new( + "no JWKS registry is configured to verify the token", + client_refusal_error(), + ) + })?; let verified = jwks .validate_with_catalog_provider(&provider.provider_name, &provider.jwks_uri, token) .await - .map_err(jwt_error_to_crate_error)?; + .map_err(token_refusal)?; let verified_claims = verified.claims(); // 4. Re-check the verified claims against the provider selected above. - validate_selected_provider_claims(&provider, verified_claims) - .map_err(jwt_error_to_crate_error)?; + validate_selected_provider_claims(&provider, verified_claims).map_err(token_refusal)?; // Legacy records without a tenant binding must not reveal catalog metadata // until the token has passed signature, time, issuer, and audience checks. let tenant_id = provider .tenant_id - .ok_or(crate::Error::OidcProviderTenantUnbound)?; + .ok_or_else(|| BearerRefusal::from_error(crate::Error::OidcProviderTenantUnbound))?; // A provider can outlive its tenant after a tenant drop. Confirm the bound // tenant still exists before issuing an identity from an authenticated token. let tenant_exists = catalog .load_all_tenants() - .map_err(|_| crate::Error::OidcProviderTenantUnavailable { tenant_id })? + .map_err(|e| { + BearerRefusal::new( + format!("tenant catalog read failed: {e}"), + crate::Error::OidcProviderTenantUnavailable { tenant_id }, + ) + })? .into_iter() .any(|tenant| tenant.tenant_id == tenant_id); if !tenant_exists { - return Err(crate::Error::OidcProviderTenantUnavailable { tenant_id }); + return Err(BearerRefusal::from_error( + crate::Error::OidcProviderTenantUnavailable { tenant_id }, + )); } // 5. Apply claim mapping. @@ -116,8 +157,10 @@ pub async fn verify_bearer_token( let default_db = mapping .default_database .map(DatabaseId::new) - .ok_or_else(|| crate::Error::OidcNoDefaultDatabase { - sub: verified_claims.sub.clone(), + .ok_or_else(|| { + BearerRefusal::from_error(crate::Error::OidcNoDefaultDatabase { + sub: verified_claims.sub.clone(), + }) })?; let mut accessible: smallvec::SmallVec<[DatabaseId; 4]> = smallvec::smallvec![default_db]; @@ -133,7 +176,7 @@ pub async fn verify_bearer_token( sub = %verified_claims.sub, iss = %verified_claims.iss, default_db = %default_db.as_u64(), - "OIDC login succeeded" + "OIDC token verified" ); let identity = identity_from_external_claims( @@ -158,7 +201,8 @@ pub async fn verify_bearer_token( state, verified_claims, &identity, - )?; + ) + .map_err(BearerRefusal::from_error)?; Ok((identity, verified)) } @@ -245,9 +289,14 @@ fn validate_selected_provider_claims( Ok(()) } -/// Collapse every pre-authentication JWT failure into one client-visible error. +/// A token-stage refusal: the audit log records `error`, and the client sees one generic error. +fn token_refusal(error: JwtError) -> BearerRefusal { + BearerRefusal::new(error.to_string(), client_refusal_error()) +} + +/// The one client-visible error of every token-stage refusal. /// Route, signature, algorithm, and token-state details remain non-observable. -fn jwt_error_to_crate_error(_error: JwtError) -> crate::Error { +fn client_refusal_error() -> crate::Error { crate::Error::BadRequest { detail: "OIDC authentication failed".into(), } diff --git a/nodedb/src/control/security/redaction/replication.rs b/nodedb/src/control/security/redaction/replication.rs index 9ccf44ae4..565bd0cc3 100644 --- a/nodedb/src/control/security/redaction/replication.rs +++ b/nodedb/src/control/security/redaction/replication.rs @@ -19,6 +19,7 @@ impl RedactionStore { let key = policy_key(policy.tenant_id, &policy.collection, &policy.for_role); let mut policies = self.lock_write(); policies.insert(key, policy); + self.bump_version(); } /// Remove a replicated policy from the in-memory registry. @@ -31,7 +32,9 @@ impl RedactionStore { ) -> bool { let key = policy_key(tenant_id, collection, for_role); let mut policies = self.lock_write(); - policies.remove(&key).is_some() + let removed = policies.remove(&key).is_some(); + self.bump_version(); + removed } /// Check whether a policy already exists for the given diff --git a/nodedb/src/control/security/redaction/store.rs b/nodedb/src/control/security/redaction/store.rs index 7e3a07bd2..4ae487f83 100644 --- a/nodedb/src/control/security/redaction/store.rs +++ b/nodedb/src/control/security/redaction/store.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use super::apply::redacted_value; use super::types::{RedactionPolicy, RedactionRule, policy_key}; @@ -13,6 +14,9 @@ use super::types::{RedactionPolicy, RedactionRule, policy_key}; pub struct RedactionStore { /// `"{tenant_id}:{collection}:{for_role}"` → redaction policy. policies: RwLock>, + /// Bumped on every policy change. A cached plan records the value it was + /// checked against, and a changed value forces a fresh check. + version: AtomicU64, } impl Default for RedactionStore { @@ -25,9 +29,22 @@ impl RedactionStore { pub fn new() -> Self { Self { policies: RwLock::new(HashMap::new()), + version: AtomicU64::new(0), } } + /// The policy version. It changes whenever a policy is created, dropped, + /// or reloaded. + pub fn version(&self) -> u64 { + self.version.load(Ordering::Acquire) + } + + /// Record a policy change. Callers hold the write lock, so a reader that + /// sees the new version also sees the new policies. + pub(super) fn bump_version(&self) { + self.version.fetch_add(1, Ordering::AcqRel); + } + /// Acquire a read lock, recovering from `RwLock` poisoning. pub(super) fn lock_read( &self, @@ -55,13 +72,16 @@ impl RedactionStore { "redaction policy created" ); policies.insert(key, policy); + self.bump_version(); } /// Drop a redaction policy. pub fn drop_policy(&self, tenant_id: u64, collection: &str, for_role: &str) -> bool { let key = policy_key(tenant_id, collection, for_role); let mut policies = self.lock_write(); - policies.remove(&key).is_some() + let removed = policies.remove(&key).is_some(); + self.bump_version(); + removed } /// Get redaction rules for a tenant+collection+role combination. @@ -213,6 +233,7 @@ impl RedactionStore { } } } + self.bump_version(); Ok(()) } diff --git a/nodedb/src/control/server/broadcast.rs b/nodedb/src/control/server/broadcast.rs index 36684de33..ef1fd3508 100644 --- a/nodedb/src/control/server/broadcast.rs +++ b/nodedb/src/control/server/broadcast.rs @@ -117,7 +117,7 @@ pub async fn broadcast_to_all_cores_txn( status: crate::bridge::envelope::Status::Ok, attempt: 1, partial: false, - payload: crate::bridge::envelope::Payload::from_vec(outcome.merged_array), + payload: crate::bridge::envelope::Payload::from_vec(outcome.merged_array()?), watermark_lsn: outcome.watermark_lsn, error_code: None, stage_vote: None, diff --git a/nodedb/src/control/server/dispatch_utils/change_events/cluster_array.rs b/nodedb/src/control/server/dispatch_utils/change_events/cluster_array.rs index acd2bfad2..05c692a63 100644 --- a/nodedb/src/control/server/dispatch_utils/change_events/cluster_array.rs +++ b/nodedb/src/control/server/dispatch_utils/change_events/cluster_array.rs @@ -17,7 +17,11 @@ pub(super) fn cluster_array_change_meta(op: &ClusterArrayOp) -> Vec { vec![(array_id.name.clone(), every_row(), ChangeOperation::Delete)] } - // Slice/Agg are reads — no row changed. - ClusterArrayOp::Slice { .. } | ClusterArrayOp::Agg { .. } => Vec::new(), + // Reads and maintenance change no row. + ClusterArrayOp::Slice { .. } + | ClusterArrayOp::Agg { .. } + | ClusterArrayOp::Elementwise { .. } + | ClusterArrayOp::Flush { .. } + | ClusterArrayOp::Compact { .. } => Vec::new(), } } diff --git a/nodedb/src/control/server/dispatch_utils/change_events/extract.rs b/nodedb/src/control/server/dispatch_utils/change_events/extract.rs index 4d78c4773..203f5d1da 100644 --- a/nodedb/src/control/server/dispatch_utils/change_events/extract.rs +++ b/nodedb/src/control/server/dispatch_utils/change_events/extract.rs @@ -597,6 +597,7 @@ mod tests { prefix_bits: 8, system_as_of: None, valid_at_ms: None, + region_msgpack: Vec::new(), }); assert!(extract_write_metadata(&agg, TenantId::new(1)).is_empty()); } diff --git a/nodedb/src/control/server/dispatch_utils/minted/owned.rs b/nodedb/src/control/server/dispatch_utils/minted/owned.rs index daefee3f3..3e045811d 100644 --- a/nodedb/src/control/server/dispatch_utils/minted/owned.rs +++ b/nodedb/src/control/server/dispatch_utils/minted/owned.rs @@ -61,6 +61,13 @@ pub(crate) enum OwnedResponse { /// The merged response outgrew its byte budget. The task closes the /// records once the final response arrives. OverBudget { bytes: usize }, + /// The final response arrived, but its streamed chunks do not merge into + /// one payload. `closed` is the result of closing the records from that + /// final response. + Unmergeable { + error: crate::Error, + closed: crate::Result<()>, + }, /// The channel closed without a final response. The records are held. ChannelClosed, } @@ -149,6 +156,12 @@ async fn wait_and_close( minted.hold(); let _ = report.send(OwnedResponse::ChannelClosed); } + Ok(Err(DispatchCollectError::Unmergeable { response, error })) => { + let closed = + resolve_on_response(&wal, owner, final_refusal_key, &response, minted).await; + let _ = report.send(OwnedResponse::Unmergeable { error, closed }); + forward(Some(*response)); + } Err(_) => { let _ = report.send(OwnedResponse::DeadlineExceeded); forward(resolve_at_final(&wal, owner, final_refusal_key, rx, minted).await); diff --git a/nodedb/src/control/server/dispatch_utils/owner_read.rs b/nodedb/src/control/server/dispatch_utils/owner_read.rs index 7ace666d9..a975cbeb0 100644 --- a/nodedb/src/control/server/dispatch_utils/owner_read.rs +++ b/nodedb/src/control/server/dispatch_utils/owner_read.rs @@ -189,7 +189,7 @@ pub(crate) fn owner_response( let payload = match payloads.len() { 0 => Vec::new(), 1 => payloads.into_iter().next().unwrap_or_default(), - _ => merge_msgpack_arrays(&payloads), + _ => merge_msgpack_arrays(&payloads)?, }; Ok(response( Status::Ok, diff --git a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/admission.rs b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/admission.rs index 8203dac78..bd798866c 100644 --- a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/admission.rs +++ b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/admission.rs @@ -7,7 +7,7 @@ use tokio::sync::OwnedMutexGuard; use crate::bridge::envelope::PhysicalPlan; use crate::control::server::shared::write_admission::gate::admit_routed; use crate::control::server::shared::write_admission::{ - WriteAdmission, WriteAdmissionGuard, WriteTarget, + AdmissionTicket, WriteAdmission, WriteAdmissionGuard, WriteTarget, }; use crate::control::state::SharedState; use crate::types::{DatabaseId, TenantId, VShardId}; @@ -33,6 +33,9 @@ pub(super) enum AdmissionOutcome { admission: crate::bridge::envelope::Admission, admission_guard: Option, order_guard: Option>, + /// The ticket of a write to a vShard with no scheduler, held until + /// the write's enqueue. + unscheduled: Option, }, /// The deterministic Calvin scheduler must apply the write. RouteToCalvin, @@ -72,6 +75,7 @@ pub(super) async fn admit_write( ), admission_guard: None, order_guard: None, + unscheduled: None, }, WriteOrdering::Gate => match admit_routed( shared, @@ -89,13 +93,25 @@ pub(super) async fn admit_write( ), admission_guard: None, order_guard: None, + unscheduled: None, }, WriteAdmission::FastPath { guard } => AdmissionOutcome::Proceed { admission: crate::bridge::envelope::Admission::Admitted, admission_guard: guard, order_guard: None, + unscheduled: None, }, - WriteAdmission::FastPathBlocking { key, keyed_lock } => { + WriteAdmission::Unscheduled { ticket } => AdmissionOutcome::Proceed { + admission: crate::bridge::envelope::Admission::Admitted, + admission_guard: None, + order_guard: None, + unscheduled: Some(ticket), + }, + WriteAdmission::FastPathBlocking { + key, + keyed_lock, + ticket, + } => { // Single-node serialization point: acquire the per-key FIFO // order-lock FIRST, before the WAL append and enqueue below. // `tokio::sync::Mutex` is fair, so concurrent same-key writers @@ -108,6 +124,7 @@ pub(super) async fn admit_write( admission: crate::bridge::envelope::Admission::Admitted, admission_guard: None, order_guard: Some(order_guard), + unscheduled: Some(ticket), } } WriteAdmission::RouteToCalvin => AdmissionOutcome::RouteToCalvin, @@ -118,6 +135,7 @@ pub(super) async fn admit_write( .await?, ), order_guard: None, + unscheduled: None, }, }, }) diff --git a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/answer.rs b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/answer.rs index b5fbd302e..4cf1550e5 100644 --- a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/answer.rs +++ b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/answer.rs @@ -46,6 +46,12 @@ impl Answer { OwnedResponse::OverBudget { bytes } } Ok(Err(DispatchCollectError::ChannelClosed)) => OwnedResponse::ChannelClosed, + Ok(Err(DispatchCollectError::Unmergeable { error, .. })) => { + OwnedResponse::Unmergeable { + error, + closed: Ok(()), + } + } Err(_) => OwnedResponse::DeadlineExceeded, }) } diff --git a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/dispatch.rs b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/dispatch.rs index 6defafdf2..8613f744a 100644 --- a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/dispatch.rs +++ b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/dispatch.rs @@ -8,7 +8,9 @@ use tokio::sync::OwnedMutexGuard; use crate::bridge::dispatch::JournalGroup; use crate::bridge::envelope::{Admission, PhysicalPlan, Priority, Request}; -use crate::control::server::shared::write_admission::{WriteAdmissionGuard, WriteOrder}; +use crate::control::server::shared::write_admission::{ + AdmissionTicket, WriteAdmissionGuard, WriteOrder, +}; use crate::control::state::SharedState; use crate::types::{ DatabaseId, Lsn, ReadConsistency, RequestId, TenantId, TraceId, TxnId, VShardId, @@ -19,6 +21,8 @@ use crate::types::{ pub(super) struct HeldGuards { pub _admission: Option, pub _order: Option>, + /// The ticket of a write to a vShard with no scheduler. + pub _unscheduled: Option, /// The document write order (see `write_order_fence`). pub _write_order: WriteOrder, } diff --git a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/driver.rs b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/driver.rs index 66e5b3098..9e0ad428f 100644 --- a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/driver.rs +++ b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/driver.rs @@ -218,12 +218,13 @@ pub(crate) async fn enqueue_write( return Err(error); } }; - let (admission, admission_guard, order_guard) = match admitted { + let (admission, admission_guard, order_guard, unscheduled) = match admitted { AdmissionOutcome::Proceed { admission, admission_guard, order_guard, - } => (admission, admission_guard, order_guard), + unscheduled, + } => (admission, admission_guard, order_guard, unscheduled), AdmissionOutcome::RouteToCalvin => { // The scheduler journals the write in its own sequenced record. // The scheduler applies the write from its own records, so @@ -407,6 +408,7 @@ pub(crate) async fn enqueue_write( HeldGuards { _admission: admission_guard, _order: order_guard, + _unscheduled: unscheduled, _write_order: write_order, }, post_apply.is_some(), diff --git a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/response.rs b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/response.rs index a4e3df114..0deadbc37 100644 --- a/nodedb/src/control/server/dispatch_utils/submit_write/funnel/response.rs +++ b/nodedb/src/control/server/dispatch_utils/submit_write/funnel/response.rs @@ -145,6 +145,15 @@ pub(super) async fn collect_classify_and_finish( ), }); } + OwnedResponse::Unmergeable { error, closed } => { + observe(shared); + if let Some(late) = late_parts { + late.spawn(deferred_guards); + } + // A failed cancel holds the window and fails the write here. + closed?; + return Err(error); + } OwnedResponse::ChannelClosed => { observe(shared); // A producer that stopped after the deadline stopped because the diff --git a/nodedb/src/control/server/exchange/all_cores/dispatch.rs b/nodedb/src/control/server/exchange/all_cores/dispatch.rs index d732a48d7..44addceb4 100644 --- a/nodedb/src/control/server/exchange/all_cores/dispatch.rs +++ b/nodedb/src/control/server/exchange/all_cores/dispatch.rs @@ -251,8 +251,10 @@ pub(crate) async fn execute_plan_all_local_cores( /// itself, and wrapping it as an array element hands the requesting node a /// different value than a local read returns. /// -/// Every other plan fans across all local cores, and their row arrays merge -/// into one. +/// Every other plan fans across all local cores and answers in the shape a +/// single core produces (`GatherOutcome::single_core_shape`): one payload as +/// is, several merged as row arrays. Two or more payloads that are not row +/// arrays fail the merge. async fn generic_gather( state: &SharedState, tenant_id: TenantId, @@ -280,7 +282,7 @@ async fn generic_gather( } let outcome = gather_all_cores(state, tenant_id, database_id, plan, trace_id, txn_id).await?; Ok(NodeLevelResult { - payload: outcome.merged_array, + payload: outcome.single_core_shape()?, watermark_lsn: outcome.watermark_lsn, read_versions: outcome.read_versions, not_found: false, diff --git a/nodedb/src/control/server/exchange/cluster_leaf.rs b/nodedb/src/control/server/exchange/cluster_leaf.rs index 074d30195..1e90c5eec 100644 --- a/nodedb/src/control/server/exchange/cluster_leaf.rs +++ b/nodedb/src/control/server/exchange/cluster_leaf.rs @@ -2,8 +2,9 @@ //! A cross-node gather of a plan whose rows are spread by tile or node key. //! -//! - A cluster array read (`ClusterArray` `Slice` / `Agg`) runs through the -//! array executor, which already knows the shard that owns each tile. A +//! - A cluster array read (`ClusterArray` `Slice` / `Agg` / `Elementwise`) +//! runs through the array executor, which already knows the shard that +//! owns each tile. A //! leaf nested under a coordinator-local wrapper (aggregate, join input, //! post-process, set operation, lateral outer plan) is read first and //! inlined as a `ProviderScan`. The rest of the plan then gathers like any @@ -11,9 +12,9 @@ //! raises the client notice its own shaping will report. //! - Graph reads never reach a gather: the SQL planner emits no graph leaf, //! and every graph read runs through `graph_dispatch`. A local `Array` plan -//! never reaches one in a cluster either, which plans arrays as -//! `ClusterArray`. Either one is refused rather than read from this node's -//! partitions alone. +//! reads one core's tiles only, and a server plans every array read as +//! `ClusterArray`. Either kind of leaf is refused rather than read from +//! one partition alone. use std::future::Future; use std::pin::Pin; @@ -22,13 +23,12 @@ use std::sync::Arc; use nodedb_physical::physical_plan::{ClusterArrayOp, ExchangeOp, PhysicalPlan, QueryOp}; use crate::control::cluster::ClusterArrayExecutor; -use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; use crate::control::server::shared::authorization::gathered_sources; use crate::control::state::SharedState; -use crate::data::executor::response_codec::{ArraySliceResponse, flatten_to_relational_rows}; -use crate::types::{Lsn, TxnId}; +use crate::data::executor::response_codec::ArraySliceResponse; +use crate::types::TxnId; -use super::gather::GatherOutcome; +use super::gather_outcome::GatherOutcome; use super::read_scope::ReadScope; use super::resolve::exchange::provider_scan_of_rows; @@ -41,7 +41,7 @@ pub(super) async fn gather_cluster_partitioned( ) -> crate::Result { if let PhysicalPlan::ClusterArray(op) = &plan { let rows = read_cluster_array(state, op, scope.txn_id).await?; - return Ok(outcome_of(flatten_to_relational_rows(&rows))); + return Ok(GatherOutcome::unversioned(vec![rows])); } let mut local = plan; inline_cluster_arrays(state, &mut local, scope).await?; @@ -62,6 +62,12 @@ pub(super) async fn gather_cluster_partitioned( /// Read one cluster array op through the array executor and return its rows /// as one msgpack array. +/// +/// The rows are already relational: a slice row holds one field per dim +/// and attr, keyed by name, an aggregate row is `{group, result}` or +/// `{result}`, and an elementwise row is `{coords, attrs}`. They reach a +/// relational operator as is. The storage flatten +/// would read an attr named `data` as a `{id, data}` row envelope. async fn read_cluster_array( state: &SharedState, op: &ClusterArrayOp, @@ -69,10 +75,13 @@ async fn read_cluster_array( ) -> crate::Result> { if matches!( op, - ClusterArrayOp::Put { .. } | ClusterArrayOp::Delete { .. } + ClusterArrayOp::Put { .. } + | ClusterArrayOp::Delete { .. } + | ClusterArrayOp::Flush { .. } + | ClusterArrayOp::Compact { .. } ) { return Err(crate::Error::Internal { - detail: "a cluster array write reached a cross-node gather".into(), + detail: "a cluster array write or maintenance op reached a cross-node gather".into(), }); } let transport = state @@ -101,7 +110,7 @@ async fn read_cluster_array( /// /// A `Slice` answers an `ArraySliceResponse` envelope whose `rows_msgpack` is /// the row array. An `Agg` answers the row array itself (`{group, result}` or -/// `{result}` maps). A slice whose `truncated_before_horizon` is set raises +/// `{result}` maps), and an `Elementwise` its `{coords, attrs}` row array. A slice whose `truncated_before_horizon` is set raises /// the same client notice a slice shaped on its own reports /// (`response_shape::compose::array_slice`), through the statement notice /// slot (`session::statement_notice`). @@ -120,9 +129,12 @@ fn cluster_array_rows(op: &ClusterArrayOp, payload: Vec) -> crate::Result { - Ok(payload) - } + ClusterArrayOp::Agg { .. } + | ClusterArrayOp::Elementwise { .. } + | ClusterArrayOp::Put { .. } + | ClusterArrayOp::Delete { .. } + | ClusterArrayOp::Flush { .. } + | ClusterArrayOp::Compact { .. } => Ok(payload), } } @@ -139,7 +151,7 @@ fn inline_cluster_arrays<'a>( PhysicalPlan::ClusterArray(op) => { let rows = read_cluster_array(state, op, scope.txn_id).await?; let sources = gathered_sources(plan, scope.database_id)?; - *plan = provider_scan_of_rows(flatten_to_relational_rows(&rows), sources); + *plan = provider_scan_of_rows(rows, sources); } PhysicalPlan::Query(QueryOp::Aggregate { input, .. }) => { if let Some(child) = input.as_deref_mut() { @@ -174,27 +186,41 @@ fn inline_cluster_arrays<'a>( inline_cluster_arrays(state, child, scope).await?; } } - _ => {} + // `plan_contains_cluster_partitioned_leaf` looks through none of + // these, so no cluster array leaf sits below them. + PhysicalPlan::Query( + QueryOp::ProviderScan { .. } + | QueryOp::PartialAggregate { .. } + | QueryOp::PartialAggregateState { .. } + | QueryOp::ShuffleJoinConsume { .. } + | QueryOp::ShuffleAggregateConsume { .. } + | QueryOp::NestedLoopJoin { .. } + | QueryOp::SortMergeJoin { .. } + | QueryOp::FacetCounts { .. } + | QueryOp::RecursiveScan { .. } + | QueryOp::RecursiveValue { .. }, + ) + | PhysicalPlan::Vector(_) + | PhysicalPlan::Graph(_) + | PhysicalPlan::Document(_) + | PhysicalPlan::Kv(_) + | PhysicalPlan::Text(_) + | PhysicalPlan::Columnar(_) + | PhysicalPlan::Timeseries(_) + | PhysicalPlan::Spatial(_) + | PhysicalPlan::Crdt(_) + | PhysicalPlan::Meta(_) + | PhysicalPlan::Array(_) + | PhysicalPlan::ClusterEvent(_) => {} } Ok(()) }) } -/// A gather outcome holding one payload. -fn outcome_of(payload: Vec) -> GatherOutcome { - let merged_array = encode_msgpack_array(&extract_msgpack_elements(&payload)); - GatherOutcome { - raw: payload, - merged_array, - watermark_lsn: Lsn::ZERO, - read_versions: crate::types::ReadVersions::new(), - shard_watermarks: Vec::new(), - } -} - #[cfg(test)] mod tests { use super::*; + use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; /// One flat row, the plain msgpack map `{"k": value}` a shard emits. /// `nodedb_types::Value` encodes as a tagged array, not a plain map, so @@ -232,21 +258,38 @@ mod tests { prefix_bits: 0, system_as_of: None, valid_at_ms: None, + region_msgpack: Vec::new(), } } #[test] - fn a_slice_envelope_flattens_to_its_rows() { + fn a_slice_envelope_yields_its_rows() { let rows = vec![row(1), row(2)]; let envelope = zerompk::to_msgpack_vec(&ArraySliceResponse { rows_msgpack: encode_msgpack_array(&rows), truncated_before_horizon: false, }) .expect("encode slice envelope"); - let flat = flatten_to_relational_rows( - &cluster_array_rows(&slice_op(), envelope).expect("slice rows"), + let out = cluster_array_rows(&slice_op(), envelope).expect("slice rows"); + assert_eq!(extract_msgpack_elements(&out).expect("extract"), rows); + } + + /// A named slice row whose attr is called `data` stays a whole row: the + /// rows reach the relational tail unflattened. + #[test] + fn an_attr_named_data_is_a_column_not_an_envelope() { + // fixmap of 2: "x" => 1, "data" => 7. + let named = vec![0x82, 0xa1, b'x', 1, 0xa4, b'd', b'a', b't', b'a', 7]; + let envelope = zerompk::to_msgpack_vec(&ArraySliceResponse { + rows_msgpack: encode_msgpack_array(std::slice::from_ref(&named)), + truncated_before_horizon: false, + }) + .expect("encode slice envelope"); + let out = cluster_array_rows(&slice_op(), envelope).expect("slice rows"); + assert_eq!( + extract_msgpack_elements(&out).expect("extract"), + vec![named] ); - assert_eq!(extract_msgpack_elements(&flat), rows); } #[tokio::test] @@ -271,11 +314,25 @@ mod tests { } #[test] - fn an_aggregate_payload_flattens_to_its_rows() { + fn an_aggregate_payload_yields_its_rows() { let rows = vec![row(10), row(20), row(30)]; - let flat = flatten_to_relational_rows( - &cluster_array_rows(&agg_op(), encode_msgpack_array(&rows)).expect("agg rows"), - ); - assert_eq!(extract_msgpack_elements(&flat), rows); + let out = cluster_array_rows(&agg_op(), encode_msgpack_array(&rows)).expect("agg rows"); + assert_eq!(extract_msgpack_elements(&out).expect("extract"), rows); + } + + #[test] + fn an_elementwise_payload_yields_its_rows() { + let op = ClusterArrayOp::Elementwise { + left: array_id(), + right: array_id(), + op: nodedb_physical::physical_plan::ArrayBinaryOp::Add, + attr_idx: 0, + prefix_bits: 0, + system_as_of: None, + valid_at_ms: None, + }; + let rows = vec![row(1), row(2)]; + let out = cluster_array_rows(&op, encode_msgpack_array(&rows)).expect("elementwise rows"); + assert_eq!(extract_msgpack_elements(&out).expect("extract"), rows); } } diff --git a/nodedb/src/control/server/exchange/full_scan.rs b/nodedb/src/control/server/exchange/full_scan.rs index ac73152f2..315bbb6a3 100644 --- a/nodedb/src/control/server/exchange/full_scan.rs +++ b/nodedb/src/control/server/exchange/full_scan.rs @@ -188,6 +188,7 @@ pub fn full_scan_plan_for_collection( valid_at_ms: None, having: Vec::new(), tiering: None, + output_columns: Vec::new(), }) } }; diff --git a/nodedb/src/control/server/exchange/gather.rs b/nodedb/src/control/server/exchange/gather.rs index 3cf519820..99959419e 100644 --- a/nodedb/src/control/server/exchange/gather.rs +++ b/nodedb/src/control/server/exchange/gather.rs @@ -3,14 +3,9 @@ //! Single fan-out/gather primitive for coordinator-mediated data movement. //! //! `gather_all_cores` fans a child plan to every Data-Plane core in parallel -//! using `join_all`, collects per-core payloads, and merges them into two -//! complementary views: -//! -//! - `raw`: concatenated per-core payloads (multiple msgpack arrays back-to-back). -//! Consumed by the sync layer and legacy raw-scan paths. -//! - `merged_array`: a single msgpack array containing every row element -//! from all cores. Consumed by the response path and by `ProviderScan` -//! embedding in join inputs. +//! using `join_all` and collects the per-core payloads into a +//! [`GatherOutcome`]. The gather parses no payload. Each consumer builds the +//! shape it reads from the outcome (see `gather_outcome`). //! //! `finalize_aggregate` runs the Arrow SIMD post-processing pass for //! `Gather{as_aggregate: true}` plans. @@ -24,12 +19,12 @@ use crate::control::server::exchange::core_outcome::{ classify_core_response, reported_versions, require_every_core, }; use crate::control::server::exchange::read_scope::ReadScope; -use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; use crate::control::server::result_stream::ResultStream; use crate::control::server::shared::session::statement_deadline; use crate::control::state::SharedState; use crate::types::{DatabaseId, Lsn, ReadConsistency, TenantId, TraceId, TxnId, VShardId}; +use super::gather_outcome::GatherOutcome; pub(super) use super::response::outcome_to_response; pub(crate) use super::response::stream_to_response; @@ -146,30 +141,6 @@ pub(crate) fn eager_dispatch_to_all_cores_until( Ok(receivers) } -/// Outcomes of a full fan-out/gather cycle across all Data-Plane cores. -pub struct GatherOutcome { - /// Concatenated per-core payloads (multiple msgpack arrays back-to-back). - /// Consumed by the sync layer and raw-scan paths. - pub raw: Vec, - /// Single merged msgpack array of all row elements. - /// Consumed by the pgwire/native response path and `ProviderScan` embedding. - pub merged_array: Vec, - /// Maximum watermark LSN seen across all responding cores. Retained as the - /// scalar fence value for Strong-consistency callers that need one LSN. - pub watermark_lsn: Lsn, - /// The read versions every responding core reported, one per vShard. A - /// vShard lives on ONE core, so only that core reports it: the scanned - /// collection's version at read time (the comparand for cross-shard OCC - /// read validation, distinct from `watermark_lsn`). - pub read_versions: crate::types::ReadVersions, - /// Per-shard watermark LSNs — one `(vshard, watermark_lsn)` per responding - /// core, NOT collapsed to the max. The transaction read-set records one - /// entry per participating shard from this, so a predicate read fanned over - /// N cores is validated against each core's own version rather than a - /// single global max. - pub shard_watermarks: Vec<(VShardId, Lsn)>, -} - /// Fan `plan` to every Data-Plane core in parallel and gather the results. /// /// All per-core sends are issued before any response is awaited (`join_all`). @@ -228,8 +199,7 @@ pub(crate) async fn gather_all_cores( classify_core_response(result).map(|resp| resp.map(|resp| (core_id, resp))) }))?; - let mut raw = Vec::new(); - let mut all_elements: Vec> = Vec::new(); + let mut payloads: Vec> = Vec::new(); let mut max_lsn = Lsn::ZERO; let mut shard_watermarks: Vec<(VShardId, Lsn)> = Vec::new(); @@ -246,17 +216,11 @@ pub(crate) async fn gather_all_cores( if resp.payload.is_empty() { continue; } - - let payload_bytes: &[u8] = resp.payload.as_ref(); - raw.extend_from_slice(payload_bytes); - all_elements.extend(extract_msgpack_elements(payload_bytes)); + payloads.push(resp.payload.to_vec()); } - let merged_array = encode_msgpack_array(&all_elements); - Ok(GatherOutcome { - raw, - merged_array, + payloads, watermark_lsn: max_lsn, read_versions, shard_watermarks, @@ -384,15 +348,6 @@ pub(crate) async fn gather_all_vshards( .await? .into_parts(); - let mut all_elements: Vec> = Vec::new(); - let mut raw = Vec::new(); - for payload in &payloads { - raw.extend_from_slice(payload); - all_elements.extend(extract_msgpack_elements(payload)); - } - - let merged_array = encode_msgpack_array(&all_elements); - // Fold the per-shard watermarks into a scalar fence the same way the local // `gather_all_cores` path does — max across participating shards — while // keeping the per-shard entries for the transaction read-set. @@ -403,8 +358,7 @@ pub(crate) async fn gather_all_vshards( .unwrap_or(Lsn::ZERO); Ok(GatherOutcome { - raw, - merged_array, + payloads: payloads.into_iter().filter(|p| !p.is_empty()).collect(), watermark_lsn, read_versions, shard_watermarks, diff --git a/nodedb/src/control/server/exchange/gather_outcome.rs b/nodedb/src/control/server/exchange/gather_outcome.rs new file mode 100644 index 000000000..d4adf6f95 --- /dev/null +++ b/nodedb/src/control/server/exchange/gather_outcome.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! The result of one fan-out/gather cycle, before any merge. +//! +//! A gather keeps each answering core's or shard's payload as that core +//! produced it. Each consumer builds the shape it reads on demand: +//! +//! - [`GatherOutcome::merged_array`]: one msgpack array of every element, +//! for a consumer whose producers answer only row arrays. Every payload +//! must be a msgpack array. +//! - [`GatherOutcome::merged_rows`]: one msgpack array of every relational +//! row. A payload is a row array, or one point-read row as a bare map. +//! - [`GatherOutcome::payloads`]: the per-core payloads, for a consumer that +//! reads one structured envelope per core (an array slice). +//! - [`GatherOutcome::single_core_shape`]: the payload a single core +//! produces, for a node-level result. +//! +//! A gather parses no payload. A consumer that reads no rows, such as a +//! purge barrier, parses none. + +use crate::control::server::payload_merge::{ + encode_msgpack_array, merge_msgpack_arrays, merge_payload_rows, +}; +use crate::types::{Lsn, ReadVersions, VShardId}; + +use super::gather::finalize_aggregate; + +/// Outcomes of a full fan-out/gather cycle across all Data-Plane cores. +pub struct GatherOutcome { + /// Every non-empty payload, one per answering core or shard, in answer + /// order, each in the shape its producer wrote. + pub(super) payloads: Vec>, + /// Maximum watermark LSN seen across all responding cores. Retained as the + /// scalar fence value for Strong-consistency callers that need one LSN. + pub watermark_lsn: Lsn, + /// The read versions every responding core reported, one per vShard. A + /// vShard lives on ONE core, so only that core reports it: the scanned + /// collection's version at read time (the comparand for cross-shard OCC + /// read validation, distinct from `watermark_lsn`). + pub read_versions: ReadVersions, + /// Per-shard watermark LSNs — one `(vshard, watermark_lsn)` per responding + /// core, NOT collapsed to the max. The transaction read-set records one + /// entry per participating shard from this, so a predicate read fanned over + /// N cores is validated against each core's own version rather than a + /// single global max. + pub shard_watermarks: Vec<(VShardId, Lsn)>, +} + +impl GatherOutcome { + /// An outcome of `payloads` that carries no watermark and no read + /// version. Empty payloads are dropped. + pub(super) fn unversioned(payloads: Vec>) -> Self { + Self { + payloads: payloads.into_iter().filter(|p| !p.is_empty()).collect(), + watermark_lsn: Lsn::ZERO, + read_versions: ReadVersions::new(), + shard_watermarks: Vec::new(), + } + } + + /// Every non-empty payload, one per answering core or shard. + pub fn payloads(&self) -> &[Vec] { + &self.payloads + } + + /// One msgpack array of every element of every payload, in answer order. + /// + /// Every payload must be one msgpack array. Any other payload is a + /// `Codec` error, never an element. + pub fn merged_array(&self) -> crate::Result> { + merge_msgpack_arrays(&self.payloads) + } + + /// One msgpack array of every relational row of every payload, in answer + /// order, passed through [`finalize_aggregate`] when the gather is a + /// `Gather{as_aggregate: true}`. + /// + /// A payload is a row array, or a point read's one row as a bare map + /// (see `payload_merge::payload_rows`). Any other payload is a `Codec` + /// error. + pub fn merged_rows(&self, as_aggregate: bool) -> crate::Result> { + let merged = merge_payload_rows(&self.payloads)?; + Ok(if as_aggregate { + finalize_aggregate(&merged) + } else { + merged + }) + } + + /// The payload in the shape a single core produces. + /// + /// No payload is an empty msgpack array. One payload returns verbatim, + /// whatever its shape. Two or more merge as row arrays through + /// [`Self::merged_array`]. + pub fn single_core_shape(&self) -> crate::Result> { + match self.payloads.as_slice() { + [] => Ok(encode_msgpack_array(&[])), + [one] => Ok(one.clone()), + _ => self.merged_array(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::server::payload_merge::extract_msgpack_elements; + + /// One flat row, the plain msgpack map `{"k": value}`. + fn row(value: u8) -> Vec { + assert!(value < 0x80, "a positive fixint holds 0..=127"); + vec![0x81, 0xa1, b'k', value] + } + + #[test] + fn row_arrays_merge_in_answer_order() { + let outcome = GatherOutcome::unversioned(vec![ + encode_msgpack_array(&[row(1), row(2)]), + Vec::new(), + encode_msgpack_array(&[row(3)]), + ]); + assert_eq!(outcome.payloads().len(), 2, "empty payloads are dropped"); + let merged = outcome.merged_array().expect("merge"); + assert_eq!( + extract_msgpack_elements(&merged).expect("extract"), + vec![row(1), row(2), row(3)] + ); + } + + /// A per-core map envelope is not an array element: the array merge + /// refuses it. + #[test] + fn a_non_array_payload_is_not_merged_as_an_element() { + let outcome = GatherOutcome::unversioned(vec![encode_msgpack_array(&[row(1)]), row(2)]); + assert!(matches!( + outcome.merged_array(), + Err(crate::Error::Codec { .. }) + )); + } + + /// A routed point read answers its one row as a bare map. + #[test] + fn a_point_read_row_merges_as_one_relational_row() { + let outcome = GatherOutcome::unversioned(vec![encode_msgpack_array(&[row(1)]), row(2)]); + let merged = outcome.merged_rows(false).expect("rows"); + assert_eq!( + extract_msgpack_elements(&merged).expect("extract"), + vec![row(1), row(2)] + ); + } + + /// A payload in no row shape fails the relational merge. + #[test] + fn a_scalar_payload_is_not_a_relational_row() { + let outcome = GatherOutcome::unversioned(vec![vec![0x2a]]); + assert!(matches!( + outcome.merged_rows(true), + Err(crate::Error::Codec { .. }) + )); + } + + /// The per-core payloads stay as their cores wrote them. + #[test] + fn payloads_keep_each_core_shape() { + let outcome = GatherOutcome::unversioned(vec![row(1), row(2)]); + assert_eq!(outcome.payloads(), &[row(1), row(2)]); + } + + #[test] + fn no_payload_is_an_empty_array() { + let outcome = GatherOutcome::unversioned(vec![Vec::new()]); + let shape = outcome.single_core_shape().expect("shape"); + assert!( + extract_msgpack_elements(&shape) + .expect("extract") + .is_empty() + ); + assert!( + extract_msgpack_elements(&outcome.merged_array().expect("merge")) + .expect("extract") + .is_empty() + ); + } + + /// One core's payload is the single-core shape as is, a map included. + #[test] + fn one_payload_returns_verbatim() { + let outcome = GatherOutcome::unversioned(vec![row(7)]); + assert_eq!(outcome.single_core_shape().expect("shape"), row(7)); + } + + #[test] + fn several_row_arrays_merge_into_one() { + let outcome = GatherOutcome::unversioned(vec![ + encode_msgpack_array(&[row(1)]), + encode_msgpack_array(&[row(2)]), + ]); + let shape = outcome.single_core_shape().expect("shape"); + assert_eq!( + extract_msgpack_elements(&shape).expect("extract"), + vec![row(1), row(2)] + ); + } + + /// Two cores that each answer a map have no single-core shape. + #[test] + fn several_non_array_payloads_are_refused() { + let outcome = GatherOutcome::unversioned(vec![row(1), row(2)]); + assert!(matches!( + outcome.single_core_shape(), + Err(crate::Error::Codec { .. }) + )); + } +} diff --git a/nodedb/src/control/server/exchange/mod.rs b/nodedb/src/control/server/exchange/mod.rs index 3aa897f1d..59ed2804e 100644 --- a/nodedb/src/control/server/exchange/mod.rs +++ b/nodedb/src/control/server/exchange/mod.rs @@ -13,6 +13,7 @@ mod cluster_leaf; pub mod core_outcome; pub mod full_scan; pub mod gather; +pub mod gather_outcome; pub mod owner_dispatch; pub mod owning_core; pub mod read_scope; @@ -25,8 +26,9 @@ pub use all_cores::NodeLevelResult; pub(crate) use all_cores::{ capture_base_on_local_cores, execute_plan_all_local_cores, snapshot_tenant_on_local_cores, }; +pub use gather::finalize_aggregate; pub(crate) use gather::gather_all_cores; -pub use gather::{GatherOutcome, finalize_aggregate}; +pub use gather_outcome::GatherOutcome; pub use read_scope::ReadScope; pub(crate) use received::execute_received_plan; pub use resolve::{ diff --git a/nodedb/src/control/server/exchange/owning_core.rs b/nodedb/src/control/server/exchange/owning_core.rs index 3701141e9..8fb839687 100644 --- a/nodedb/src/control/server/exchange/owning_core.rs +++ b/nodedb/src/control/server/exchange/owning_core.rs @@ -12,11 +12,10 @@ use crate::bridge::envelope::PhysicalPlan; use crate::control::local_dispatch::reject_data_plane_error; use crate::control::server::dispatch_utils::dispatch_routed_read_to_data_plane; -use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; use crate::control::state::SharedState; use crate::types::{DatabaseId, TenantId, TraceId, TxnId, VShardId}; -use super::gather::GatherOutcome; +use super::gather_outcome::GatherOutcome; /// Dispatch `plan` to the single Data-Plane core that owns `vshard_id` and /// gather the one bounded response into a [`GatherOutcome`]. @@ -30,9 +29,10 @@ use super::gather::GatherOutcome; /// `read_versions` and exactly one `shard_watermarks` entry keyed to the /// collection's vShard — matching the cluster `dispatch_local` path so an /// in-transaction read records the same OCC read-set entry the write-set uses -/// (writes home to the same `CollectionKey` vShard). Aggregate -/// finalization (`finalize_aggregate`) is a passthrough over the merged array, -/// so one complete aggregate row in yields one row out. +/// (writes home to the same `CollectionKey` vShard). The outcome holds the +/// core's payload as the core wrote it. Aggregate finalization +/// (`finalize_aggregate`) is a passthrough over the merged array, so one +/// complete aggregate row in yields one row out. pub async fn gather_single_owning_core( state: &SharedState, tenant_id: TenantId, @@ -52,13 +52,14 @@ pub async fn gather_single_owning_core( txn_id, ) .await?; - let payload_bytes: &[u8] = resp.payload.as_ref(); - let all_elements = extract_msgpack_elements(payload_bytes); - let merged_array = encode_msgpack_array(&all_elements); + let payloads = if resp.payload.is_empty() { + Vec::new() + } else { + vec![resp.payload.to_vec()] + }; Ok(GatherOutcome { - raw: payload_bytes.to_vec(), - merged_array, + payloads, watermark_lsn: resp.watermark_lsn, read_versions: resp.read_versions.clone(), shard_watermarks: vec![(vshard_id, resp.watermark_lsn)], diff --git a/nodedb/src/control/server/exchange/resolve/array_prefilter.rs b/nodedb/src/control/server/exchange/resolve/array_prefilter.rs new file mode 100644 index 000000000..b35f632ab --- /dev/null +++ b/nodedb/src/control/server/exchange/resolve/array_prefilter.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Resolve a vector search's `ARRAY_SLICE` prefilter across every shard. +//! +//! The planner lowers the prefilter to an `ArrayOp::SurrogateBitmapScan` in +//! the search's `inline_prefilter_plan`. The search runs on the core that +//! owns its collection, and that core holds only its own buckets of the +//! array, so the inline scan would admit only the cells that core holds. +//! This pass reads the slice's surrogates from every shard through the +//! array coordinator, intersects them into the search's `filter_bitmap`, +//! and clears the inline plan. A search with no prefilter, and a node with +//! no cluster transport, keeps its plan as it is. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use nodedb_physical::physical_plan::{ArrayOp, ExchangeOp, PhysicalPlan, QueryOp, VectorOp}; + +use crate::control::cluster::ClusterArrayExecutor; +use crate::control::state::SharedState; + +/// Replace every array prefilter of a vector search in `plan` with the +/// cross-shard surrogate bitmap of its slice. Walks the coordinator-local +/// wrappers a search can sit under. +pub(super) fn resolve_array_prefilters<'a>( + state: &'a SharedState, + plan: &'a mut PhysicalPlan, +) -> Pin> + Send + 'a>> { + Box::pin(async move { + match plan { + PhysicalPlan::Vector(op) => resolve_search(state, op).await?, + PhysicalPlan::Query(QueryOp::Exchange(ExchangeOp { child, .. })) + | PhysicalPlan::Query(QueryOp::PostProcess { input: child, .. }) + | PhysicalPlan::Query(QueryOp::LateralTopK { + outer_plan: child, .. + }) + | PhysicalPlan::Query(QueryOp::LateralLoop { + outer_plan: child, .. + }) => resolve_array_prefilters(state, child).await?, + PhysicalPlan::Query( + QueryOp::Aggregate { input, .. } | QueryOp::PartialAggregateState { input, .. }, + ) => { + if let Some(child) = input.as_deref_mut() { + resolve_array_prefilters(state, child).await?; + } + } + PhysicalPlan::Query(QueryOp::HashJoin { + left_input, + right_input, + left_bitmap, + right_bitmap, + .. + }) => { + for side in [left_input, right_input, left_bitmap, right_bitmap] { + if let Some(child) = side.as_deref_mut() { + resolve_array_prefilters(state, child).await?; + } + } + } + PhysicalPlan::Query( + QueryOp::NestedLoopJoin { + left_input, + right_input, + .. + } + | QueryOp::SortMergeJoin { + left_input, + right_input, + .. + }, + ) => { + for side in [left_input, right_input] { + if let Some(child) = side.as_deref_mut() { + resolve_array_prefilters(state, child).await?; + } + } + } + PhysicalPlan::Query(QueryOp::SetOp { inputs, .. }) => { + for child in inputs.iter_mut() { + resolve_array_prefilters(state, child).await?; + } + } + // No child plan and no vector search. + PhysicalPlan::Query( + QueryOp::ProviderScan { .. } + | QueryOp::PartialAggregate { .. } + | QueryOp::ShuffleJoinConsume { .. } + | QueryOp::ShuffleAggregateConsume { .. } + | QueryOp::FacetCounts { .. } + | QueryOp::RecursiveScan { .. } + | QueryOp::RecursiveValue { .. }, + ) + | PhysicalPlan::Graph(_) + | PhysicalPlan::Document(_) + | PhysicalPlan::Kv(_) + | PhysicalPlan::Text(_) + | PhysicalPlan::Columnar(_) + | PhysicalPlan::Timeseries(_) + | PhysicalPlan::Spatial(_) + | PhysicalPlan::Crdt(_) + | PhysicalPlan::Meta(_) + | PhysicalPlan::Array(_) + | PhysicalPlan::ClusterArray(_) + | PhysicalPlan::ClusterEvent(_) => {} + } + Ok(()) + }) +} + +/// Replace the array prefilter of a vector search with the cross-shard +/// bitmap. Every other vector op has no prefilter. +async fn resolve_search(state: &SharedState, op: &mut VectorOp) -> crate::Result<()> { + let VectorOp::Search { + inline_prefilter_plan, + filter_bitmap, + .. + } = op + else { + return Ok(()); + }; + let Some(PhysicalPlan::Array(ArrayOp::SurrogateBitmapScan { + array_id, + slice_msgpack, + })) = inline_prefilter_plan.as_deref() + else { + return Ok(()); + }; + let Some(bitmap) = cluster_bitmap(state, array_id, slice_msgpack).await? else { + return Ok(()); + }; + *filter_bitmap = Some(match filter_bitmap.take() { + Some(existing) => existing.intersect(&bitmap), + None => bitmap, + }); + *inline_prefilter_plan = None; + Ok(()) +} + +/// The cross-shard surrogate bitmap of `array_id` inside `slice_msgpack`, +/// or `None` on a node with no cluster transport. +async fn cluster_bitmap( + state: &SharedState, + array_id: &nodedb_array::types::ArrayId, + slice_msgpack: &[u8], +) -> crate::Result> { + let (Some(transport), Some(routing)) = ( + state.cluster_transport.as_ref(), + state.cluster_routing.as_ref(), + ) else { + return Ok(None); + }; + let prefix_bits = { + let catalog = state + .array_catalog + .read() + .map_err(|_| crate::Error::Internal { + detail: "array catalog lock poisoned".into(), + })?; + catalog + .lookup_by_name_in_database(array_id.tenant_id, array_id.database_id, &array_id.name) + .ok_or_else(|| crate::Error::PlanError { + detail: format!("array prefilter: array '{}' not found", array_id.name), + })? + .prefix_bits + }; + let executor = ClusterArrayExecutor::new( + Arc::clone(transport), + Arc::clone(routing), + state.node_id, + state.self_arc()?, + ); + executor + .surrogate_bitmap(array_id, slice_msgpack, prefix_bits) + .await + .map(Some) +} diff --git a/nodedb/src/control/server/exchange/resolve/array_read_rows.rs b/nodedb/src/control/server/exchange/resolve/array_read_rows.rs new file mode 100644 index 000000000..779309669 --- /dev/null +++ b/nodedb/src/control/server/exchange/resolve/array_read_rows.rs @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Which gathered reads are cluster array reads. +//! +//! A cluster array read (`ClusterArrayOp::Slice` / `Agg` / `Elementwise`) +//! reaches the coordinator as the bare relational rows its array coordinator +//! merged (`cluster_leaf`): a slice row holds one field per dim and attr, +//! keyed by name, an aggregate row is `{group, result}` or `{result}`, and +//! an elementwise row is `{coords, attrs}`. The storage +//! flatten that unwraps `{id, data}` document rows must not run over them: it +//! will read an attr named `data` as a row envelope. +//! +//! A local `PhysicalPlan::Array` read never yields gathered rows. +//! `gather_all_vshards` hands every array plan to `cluster_leaf`, which +//! inlines cluster array leaves and refuses a local one. + +use nodedb_physical::physical_plan::{ClusterArrayOp, PhysicalPlan}; + +/// Whether `plan` is a cluster array read, whose gathered rows are already +/// relational. +pub(super) fn is_cluster_array_read(plan: &PhysicalPlan) -> bool { + matches!( + plan, + PhysicalPlan::ClusterArray( + ClusterArrayOp::Slice { .. } + | ClusterArrayOp::Agg { .. } + | ClusterArrayOp::Elementwise { .. } + ) + ) +} diff --git a/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs index 0e7300bed..363056920 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/aggregate_input_arm.rs @@ -30,6 +30,7 @@ pub(super) struct AggregateFields { pub sub_aggregates: Vec, pub grouping_sets: Vec>, pub sort_keys: Vec, + pub group_key_aliases: Vec>, } /// Resolve an input-sourced `QueryOp::Aggregate`. @@ -59,6 +60,7 @@ pub(super) async fn resolve_aggregate_input( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, } = fields; let rebuild = |input: Box| { @@ -74,6 +76,7 @@ pub(super) async fn resolve_aggregate_input( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }))) }; diff --git a/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs b/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs index a99481c6a..4aa6201eb 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/dispatch.rs @@ -22,7 +22,7 @@ use super::{ /// /// - Root-level `Gather` → gather all vShards, return `Resolved::Gathered`. /// - `Broadcast` nested inside a `HashJoin` input → gather the child, embed -/// the `merged_array` as `ProviderScan{None, rows}`, return `Resolved::Plan`. +/// its rows as `ProviderScan{None, rows}`, return `Resolved::Plan`. /// - Root-level `Shuffle` wrapping a `HashJoin` → orchestrate a cross-node /// grace hash join, return `Resolved::Gathered`. `Shuffle` as a join input is /// a typed error. @@ -187,6 +187,7 @@ pub(super) async fn resolve_exchange( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }) => { aggregate_input_arm::resolve_aggregate_input( state, @@ -204,6 +205,7 @@ pub(super) async fn resolve_exchange( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }, ) .await diff --git a/nodedb/src/control/server/exchange/resolve/exchange/entry.rs b/nodedb/src/control/server/exchange/resolve/exchange/entry.rs index 63e04be4a..93b6afa86 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/entry.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/entry.rs @@ -10,6 +10,7 @@ use crate::control::server::exchange::read_scope::ReadScope; use crate::control::state::SharedState; use crate::types::{Lsn, VShardId}; +use crate::control::server::exchange::resolve::array_prefilter::resolve_array_prefilters; use crate::control::server::exchange::resolve::capture::DistributedReadCapture; use crate::control::server::exchange::resolve::materialize::materialize_providers; use crate::control::server::result_stream::ResultStream; @@ -64,7 +65,8 @@ pub async fn resolve_and_materialize( scope: ReadScope, ) -> crate::Result { // Pass 1: fill empty ProviderScan rows (identity-scoped, per-request). - let plan = materialize_providers(state, identity, scope.database_id, plan).await?; + let mut plan = materialize_providers(state, identity, scope.database_id, plan).await?; + resolve_array_prefilters(state, &mut plan).await?; // Pass 2: resolve Exchange nodes. The captures accumulator is filled at every // base-collection gather point beneath the plan root and consumed (taken) @@ -85,9 +87,10 @@ pub async fn resolve_and_materialize( /// See `resolve_and_materialize` for the `scope` semantics. pub async fn resolve_exchange_in_plan( state: &SharedState, - plan: PhysicalPlan, + mut plan: PhysicalPlan, scope: ReadScope, ) -> crate::Result { + resolve_array_prefilters(state, &mut plan).await?; let mut captures = Vec::new(); resolve_exchange(state, scope, plan, &mut captures).await } diff --git a/nodedb/src/control/server/exchange/resolve/exchange/gather_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/gather_arm.rs index 30d9fee21..17fa29775 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/gather_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/gather_arm.rs @@ -6,9 +6,8 @@ use nodedb_physical::physical_plan::PhysicalPlan; use crate::control::server::exchange::full_scan::{ScanSide, full_scan_plan_for_collection}; -use crate::control::server::exchange::gather::{ - GatherOutcome, finalize_aggregate, gather_all_vshards, outcome_to_response, -}; +use crate::control::server::exchange::gather::{gather_all_vshards, outcome_to_response}; +use crate::control::server::exchange::gather_outcome::GatherOutcome; use crate::control::server::exchange::resolve::capture::DistributedReadCapture; use crate::control::state::SharedState; @@ -125,11 +124,7 @@ pub(super) async fn resolve_gather( }); } - let payload = if as_aggregate { - finalize_aggregate(&outcome.merged_array) - } else { - outcome.merged_array - }; + let payload = outcome.merged_rows(as_aggregate)?; Ok(Resolved::Gathered( outcome_to_response(payload, outcome.watermark_lsn, outcome.read_versions), outcome.shard_watermarks, @@ -148,7 +143,7 @@ pub(super) async fn resolve_broadcast( let outcome = gather_all_vshards(state, child, ctx).await?; Ok(Resolved::Gathered( outcome_to_response( - outcome.merged_array, + outcome.merged_rows(false)?, outcome.watermark_lsn, outcome.read_versions, ), diff --git a/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs index d25c7ffeb..1de45ffee 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/post_process_arm.rs @@ -8,14 +8,17 @@ use nodedb_physical::physical_plan::{ plan_contains_cluster_partitioned_leaf, }; +use super::super::array_read_rows::is_cluster_array_read; + use crate::control::server::exchange::full_scan::{ScanSide, full_scan_plan_for_collection}; -use crate::control::server::exchange::gather::{ - GatherOutcome, finalize_aggregate, gather_all_vshards, -}; +use crate::control::server::exchange::gather::gather_all_vshards; +use crate::control::server::exchange::gather_outcome::GatherOutcome; use crate::control::server::exchange::owning_core::gather_single_owning_core; use crate::control::server::exchange::resolve::capture::DistributedReadCapture; use crate::control::server::response_translate::hit_key::parse_surrogate_hex; -use crate::control::server::response_translate::vector::resolve_surrogate_pk; +use crate::control::server::response_translate::vector::{ + filter_vector_hits, resolve_surrogate_pk, +}; use crate::control::server::shared::authorization::gathered_sources; use crate::control::state::SharedState; use crate::data::executor::response_codec::{ @@ -52,6 +55,9 @@ enum HitShape { /// Hybrid (RRF) fusion hits: `{doc_id: , , /// ...}` with no body. Resolve `doc_id` to the user PK as `id`. Hybrid, + /// A cluster array read, whose rows are already relational (see + /// [`is_cluster_array_read`]). + ArrayRows, /// Flat storage rows (`{id, data}` document / text, columnar, spatial) or /// computed rows — already fully columned after the storage flatten. None, @@ -59,6 +65,9 @@ enum HitShape { /// Classify a resolved `PostProcess` child by the row shape its engine emits. fn classify_hit_shape(plan: &PhysicalPlan) -> HitShape { + if is_cluster_array_read(plan) { + return HitShape::ArrayRows; + } match plan { PhysicalPlan::Vector( VectorOp::Search { .. } @@ -73,6 +82,33 @@ fn classify_hit_shape(plan: &PhysicalPlan) -> HitShape { } } +/// The residual filter and `top_k` bound a vector search's hits are cut to +/// on the Control Plane, for a body that carries them. +/// +/// The Data Plane ranks a wider window while a residual filter applies and +/// leaves the filter and the cut to the Control Plane. A body read under a +/// tail reaches the gather path, not the response translator, so the cut +/// runs here. +fn vector_hit_cut(plan: &PhysicalPlan) -> Option<(Vec, usize)> { + if let PhysicalPlan::Vector( + VectorOp::Search { + rls_filters, top_k, .. + } + | VectorOp::MultiSearch { + rls_filters, top_k, .. + } + | VectorOp::SparseSearch { + filters: rls_filters, + top_k, + .. + }, + ) = plan + { + return Some((rls_filters.clone(), *top_k)); + } + None +} + /// Collection a `PostProcess` child reads, for the surrogate→PK resolver. /// /// The search ops that emit surrogate-keyed hits carry their collection in a @@ -191,6 +227,7 @@ pub(super) async fn materialize_child_rows( // `collection` and `hit_kind` are captured before the gather // consumes `child`. let hit_kind = classify_hit_shape(&child); + let hit_cut = vector_hit_cut(&child); // Extract the collection from the hit op directly: `collection()` // has no arm for sparse / multi-vector search, so it will yield // `None` and the PK resolver will be handed an empty collection. @@ -241,31 +278,37 @@ pub(super) async fn materialize_child_rows( }); } - let merged = if as_aggregate { - finalize_aggregate(&outcome.merged_array) - } else { - outcome.merged_array - }; - // Flatten to the bare relational row shape the `ProviderScan` tail // consumes, resolving surrogate→PK for hit-shaped bodies via the - // catalog so `SELECT id` returns the user PK, not the surrogate. + // catalog so `SELECT id` returns the user PK, not the surrogate. A + // cluster array read's rows pass through unflattened. let coll = hit_collection.unwrap_or_default(); let rows = match hit_kind { - HitShape::Vector => flatten_vector_hits_to_relational_rows(&merged, |surrogate| { - resolve_surrogate_pk( - state, - database_id, - tenant_id, - &coll, - nodedb_types::Surrogate::new(surrogate), - ) - }), - HitShape::Hybrid => flatten_hybrid_hits_to_relational_rows(&merged, |key| { - let surrogate = parse_surrogate_hex(key)?; - resolve_surrogate_pk(state, database_id, tenant_id, &coll, surrogate) - }), - HitShape::None => flatten_to_relational_rows(&merged), + HitShape::Vector => { + let merged = outcome.merged_rows(as_aggregate)?; + let merged = match &hit_cut { + Some((rls_filters, top_k)) => filter_vector_hits(&merged, rls_filters, *top_k)?, + None => merged, + }; + flatten_vector_hits_to_relational_rows(&merged, |surrogate| { + resolve_surrogate_pk( + state, + database_id, + tenant_id, + &coll, + nodedb_types::Surrogate::new(surrogate), + ) + }) + } + HitShape::Hybrid => { + let merged = outcome.merged_rows(as_aggregate)?; + flatten_hybrid_hits_to_relational_rows(&merged, |key| { + let surrogate = parse_surrogate_hex(key)?; + resolve_surrogate_pk(state, database_id, tenant_id, &coll, surrogate) + }) + } + HitShape::ArrayRows => outcome.merged_array()?, + HitShape::None => flatten_to_relational_rows(&outcome.merged_rows(as_aggregate)?), }; Ok(ChildRows::Rows { rows, sources }) } diff --git a/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs index 92f365dca..a50ca6947 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/set_op_arm.rs @@ -6,7 +6,7 @@ use nodedb_physical::physical_plan::{PhysicalPlan, SetOpKind}; use crate::control::server::exchange::resolve::capture::DistributedReadCapture; -use crate::control::server::payload_merge::merge_msgpack_arrays; +use crate::control::server::payload_merge::merge_payload_rows; use crate::control::server::set_op_merge::{ SetMergeMode, dedup_union_payloads, merge_set_op_payloads, }; @@ -60,7 +60,7 @@ pub(super) async fn resolve_set_op( } // The merged rows carry the requirements of every branch. - let merged = merge_set_op_rows(&payloads, op); + let merged = merge_set_op_rows(&payloads, op)?; Ok(Resolved::Plan(Box::new(provider_scan_of_rows( merged, sources, )))) @@ -70,10 +70,10 @@ pub(super) async fn resolve_set_op( /// array of flat row maps. The mapping from kind to merge mirrors the pgwire /// per-task set-op path: `INTERSECT [ALL]` and `EXCEPT [ALL]` share one /// value-keyed merge each, `UNION` dedups on raw bytes, and `UNION ALL` -/// concatenates. -fn merge_set_op_rows(payloads: &[Vec], op: SetOpKind) -> Vec { +/// concatenates. A payload in no row shape fails the merge. +fn merge_set_op_rows(payloads: &[Vec], op: SetOpKind) -> crate::Result> { match op { - SetOpKind::UnionAll => merge_msgpack_arrays(payloads), + SetOpKind::UnionAll => merge_payload_rows(payloads), SetOpKind::UnionDistinct => dedup_union_payloads(payloads), SetOpKind::Intersect | SetOpKind::IntersectAll => { merge_set_op_payloads(payloads, SetMergeMode::Intersect) @@ -101,7 +101,7 @@ mod tests { fn union_all_keeps_every_row_in_branch_order() { let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); let right = encode_array(&[serde_json::json!({"x": 2})]); - let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionAll); + let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionAll).unwrap(); assert_eq!(decode(&merged), r#"[{"x":1},{"x":2},{"x":2}]"#); } @@ -109,7 +109,7 @@ mod tests { fn union_distinct_drops_duplicates() { let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); let right = encode_array(&[serde_json::json!({"x": 2})]); - let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionDistinct); + let merged = merge_set_op_rows(&[left, right], SetOpKind::UnionDistinct).unwrap(); assert_eq!(decode(&merged), r#"[{"x":1},{"x":2}]"#); } @@ -117,7 +117,7 @@ mod tests { fn intersect_keeps_rows_present_in_every_branch() { let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); let right = encode_array(&[serde_json::json!({"x": 2}), serde_json::json!({"x": 3})]); - let merged = merge_set_op_rows(&[left, right], SetOpKind::Intersect); + let merged = merge_set_op_rows(&[left, right], SetOpKind::Intersect).unwrap(); assert_eq!(decode(&merged), r#"[{"x":2}]"#); } @@ -125,7 +125,24 @@ mod tests { fn except_drops_rows_present_in_later_branches() { let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); let right = encode_array(&[serde_json::json!({"x": 2})]); - let merged = merge_set_op_rows(&[left, right], SetOpKind::Except); + let merged = merge_set_op_rows(&[left, right], SetOpKind::Except).unwrap(); assert_eq!(decode(&merged), r#"[{"x":1}]"#); } + + /// A branch payload in no row shape fails the merge in every mode. + #[test] + fn a_scalar_branch_fails_every_merge() { + let left = encode_array(&[serde_json::json!({"x": 1})]); + for op in [ + SetOpKind::UnionAll, + SetOpKind::UnionDistinct, + SetOpKind::Intersect, + SetOpKind::Except, + ] { + assert!(matches!( + merge_set_op_rows(&[left.clone(), vec![0x2a]], op), + Err(crate::Error::Codec { .. }) + )); + } + } } diff --git a/nodedb/src/control/server/exchange/resolve/exchange/txn_source_arm.rs b/nodedb/src/control/server/exchange/resolve/exchange/txn_source_arm.rs index e19a12d74..c43f34dbe 100644 --- a/nodedb/src/control/server/exchange/resolve/exchange/txn_source_arm.rs +++ b/nodedb/src/control/server/exchange/resolve/exchange/txn_source_arm.rs @@ -54,6 +54,7 @@ pub(super) async fn txn_aggregate( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }) = plan else { return Ok(TxnRewrite::Unchanged(Box::new(plan))); @@ -79,6 +80,7 @@ pub(super) async fn txn_aggregate( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }, )))); }; @@ -99,6 +101,7 @@ pub(super) async fn txn_aggregate( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }, ) .await?; diff --git a/nodedb/src/control/server/exchange/resolve/join_input.rs b/nodedb/src/control/server/exchange/resolve/join_input.rs index 56f6ab7e3..13de544b1 100644 --- a/nodedb/src/control/server/exchange/resolve/join_input.rs +++ b/nodedb/src/control/server/exchange/resolve/join_input.rs @@ -10,8 +10,10 @@ use crate::control::state::SharedState; use crate::data::executor::response_codec::flatten_to_relational_rows; use crate::control::server::exchange::full_scan::{ScanSide, full_scan_plan_for_collection}; -use crate::control::server::exchange::gather::{finalize_aggregate, gather_all_vshards}; +use crate::control::server::exchange::gather::gather_all_vshards; +use crate::control::server::exchange::gather_outcome::GatherOutcome; +use super::array_read_rows::is_cluster_array_read; use super::capture::DistributedReadCapture; use super::exchange::provider_scan_of_rows; use crate::control::server::shared::authorization::gathered_sources; @@ -19,7 +21,7 @@ use crate::control::server::shared::authorization::gathered_sources; /// Resolve a `HashJoin` input slot. /// /// When the slot contains an `Exchange{Broadcast}` child, gathers the child to -/// the coordinator and replaces the slot with a `ProviderScan{None, merged_array}`. +/// the coordinator and replaces the slot with a `ProviderScan{None, rows}`. /// When the slot is already a `ProviderScan{None, ..}` or `None`, it is /// returned unchanged. /// @@ -47,23 +49,22 @@ pub(super) async fn resolve_join_input( child, mode: ExchangeMode::Broadcast, })) => { - // Gather the broadcast child to the coordinator, then embed the - // merged msgpack array as an inline ProviderScan. + // Gather the broadcast child to the coordinator, then embed its + // rows as one msgpack array in an inline ProviderScan. // - // We use `merged_array` (not `raw`) because `merged_array` is a - // single well-formed msgpack array. The Data-Plane executor - // materialises the ProviderScan via `response_with_payload(rows)`, - // producing a Response whose payload is exactly `merged_array`. - // `decode_response_to_docs` in `hash_handlers.rs` then reads that - // Response as a msgpack array — so the two shapes match. + // The Data-Plane executor materialises the ProviderScan via + // `response_with_payload(rows)`, producing a Response whose + // payload is exactly those rows. `decode_response_to_docs` in + // `hash_handlers.rs` then reads that Response as a msgpack + // array, so the two shapes match. // The gather reads every vShard of the child from its owner, not // only the groups this node replicates, and confirms each leg as // `scope` requires. let sources = gathered_sources(&child, database_id)?; + let array_rows = is_cluster_array_read(&child); let outcome = Box::pin(gather_all_vshards(state, *child, scope)).await?; - let provider_scan = - provider_scan_of_rows(flatten_to_relational_rows(&outcome.merged_array), sources); - Ok(Some(Box::new(provider_scan))) + let rows = relational_rows(array_rows, &outcome, false)?; + Ok(Some(Box::new(provider_scan_of_rows(rows, sources)))) } // Exchange{Shuffle} inside a join input is never a shape the emit @@ -108,6 +109,7 @@ pub(super) async fn resolve_join_input( None }; let sources = gathered_sources(&child, database_id)?; + let array_rows = is_cluster_array_read(&child); // Read from every vShard's owner, as the Broadcast arm does. let outcome = Box::pin(gather_all_vshards(state, *child, scope)).await?; if let Some(coll) = child_collection @@ -123,13 +125,8 @@ pub(super) async fn resolve_join_input( read_versions: outcome.read_versions.clone(), }); } - let merged = if as_aggregate { - finalize_aggregate(&outcome.merged_array) - } else { - outcome.merged_array - }; - let provider_scan = provider_scan_of_rows(flatten_to_relational_rows(&merged), sources); - Ok(Some(Box::new(provider_scan))) + let rows = relational_rows(array_rows, &outcome, as_aggregate)?; + Ok(Some(Box::new(provider_scan_of_rows(rows, sources)))) } // Already resolved (ProviderScan{None, ..} or any other plan): @@ -218,7 +215,24 @@ pub(super) async fn gather_join_build_side( } Ok(Some(Box::new(provider_scan_of_rows( - flatten_to_relational_rows(&outcome.merged_array), + flatten_to_relational_rows(&outcome.merged_rows(false)?), sources, )))) } + +/// The relational rows of a gathered join input. A cluster array read +/// (`array_rows`) keeps its rows as it emits them. Any other plan's merged +/// rows pass through the storage flatten. `as_aggregate` finalizes the +/// merged rows of a `Gather{as_aggregate}`. +fn relational_rows( + array_rows: bool, + outcome: &GatherOutcome, + as_aggregate: bool, +) -> crate::Result> { + if array_rows { + return outcome.merged_array(); + } + Ok(flatten_to_relational_rows( + &outcome.merged_rows(as_aggregate)?, + )) +} diff --git a/nodedb/src/control/server/exchange/resolve/materialize.rs b/nodedb/src/control/server/exchange/resolve/materialize.rs index 2c21e534e..21a60623c 100644 --- a/nodedb/src/control/server/exchange/resolve/materialize.rs +++ b/nodedb/src/control/server/exchange/resolve/materialize.rs @@ -106,6 +106,7 @@ pub(super) async fn materialize_providers( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, }) => { let input = Box::pin(materialize_providers(state, identity, database_id, *input)).await?; @@ -121,6 +122,7 @@ pub(super) async fn materialize_providers( sub_aggregates, grouping_sets, sort_keys, + group_key_aliases, })) } diff --git a/nodedb/src/control/server/exchange/resolve/mod.rs b/nodedb/src/control/server/exchange/resolve/mod.rs index 890be7b89..866e18ea5 100644 --- a/nodedb/src/control/server/exchange/resolve/mod.rs +++ b/nodedb/src/control/server/exchange/resolve/mod.rs @@ -29,7 +29,12 @@ //! per-part consumers (one per partition), and concatenate + global-cap the //! finalized rows into a `Resolved::Gathered` response (distributed GROUP BY). //! - No Exchange / no empty ProviderScan → `Resolved::Plan` unchanged. +//! +//! Before pass 2, `array_prefilter` replaces a vector search's array +//! prefilter scan with the slice's surrogate bitmap read from every shard. +mod array_prefilter; +mod array_read_rows; mod cagg_rows; mod capture; pub mod exchange; diff --git a/nodedb/src/control/server/exchange/resolve/shuffle.rs b/nodedb/src/control/server/exchange/resolve/shuffle.rs index af0f31fb1..b85fb871e 100644 --- a/nodedb/src/control/server/exchange/resolve/shuffle.rs +++ b/nodedb/src/control/server/exchange/resolve/shuffle.rs @@ -354,7 +354,7 @@ pub async fn resolve_shuffle_join( // 9. Merge every part-owner's msgpack-array rows into one array (reuse the // shared gather merge helper — never hand-roll msgpack framing). - let merged = merge_msgpack_arrays(&per_part_rows); + let merged = merge_msgpack_arrays(&per_part_rows)?; // The per-node consumer registries self-clean their inboxes once the part's // staged sides are consumed; there is no coordinator-side cleanup hook to diff --git a/nodedb/src/control/server/exchange/resolve/shuffle_aggregate.rs b/nodedb/src/control/server/exchange/resolve/shuffle_aggregate.rs index 1ad4e28d6..0a520a261 100644 --- a/nodedb/src/control/server/exchange/resolve/shuffle_aggregate.rs +++ b/nodedb/src/control/server/exchange/resolve/shuffle_aggregate.rs @@ -90,6 +90,7 @@ pub async fn resolve_shuffle_aggregate( having, limit, sort_keys, + group_key_aliases, .. }) = child else { @@ -325,12 +326,17 @@ pub async fn resolve_shuffle_aggregate( // ORDER BY). Reuse the shared extract/encode helpers — no hand-rolled framing. let mut elements: Vec> = Vec::new(); for rows in &per_part_rows { - elements.extend(extract_msgpack_elements(rows)); + elements.extend(extract_msgpack_elements(rows)?); } if elements.len() > limit { elements.truncate(limit); } - let merged = encode_msgpack_array(&elements); + // Each part finalizes its groups under the keys' row names. The aliases + // apply to the merged rows, as they do on a single-node aggregate. + let merged = crate::util::row_rename::rename_row_keys( + &encode_msgpack_array(&elements), + &crate::util::row_rename::group_key_renames(&group_by, &group_key_aliases), + )?; // The producers report the source collection's versions at read time on // their `ShuffleProduceResponse`; the fold above is what this aggregate diff --git a/nodedb/src/control/server/graph_dispatch/keyed_read.rs b/nodedb/src/control/server/graph_dispatch/keyed_read.rs index 3bdedf5a9..19430b1fc 100644 --- a/nodedb/src/control/server/graph_dispatch/keyed_read.rs +++ b/nodedb/src/control/server/graph_dispatch/keyed_read.rs @@ -120,7 +120,7 @@ pub async fn read_on_vshard( let payloads = outcome.payloads; Payload::from_vec(match payloads.len() { 1 => payloads.into_iter().next().unwrap_or_default(), - _ => merge_msgpack_arrays(&payloads), + _ => merge_msgpack_arrays(&payloads)?, }) } RouteDecision::LeaderUnknown { vshard_id } => { diff --git a/nodedb/src/control/server/graph_dispatch/match_broadcast.rs b/nodedb/src/control/server/graph_dispatch/match_broadcast.rs index 96de6264e..8875bef47 100644 --- a/nodedb/src/control/server/graph_dispatch/match_broadcast.rs +++ b/nodedb/src/control/server/graph_dispatch/match_broadcast.rs @@ -132,7 +132,7 @@ fn decode_match_envelope(payload: &[u8]) -> crate::Result detail: "match envelope: missing or malformed 'frontier' field".into(), })?; - let row_elements = extract_msgpack_elements(rows_bytes); + let row_elements = extract_msgpack_elements(rows_bytes)?; let frontier: Vec = zerompk::from_msgpack(frontier_bytes).map_err(|e| crate::Error::Codec { detail: format!("match envelope: invalid frontier: {e}"), diff --git a/nodedb/src/control/server/http/auth.rs b/nodedb/src/control/server/http/auth.rs index ffb663fad..74ce7dec1 100644 --- a/nodedb/src/control/server/http/auth.rs +++ b/nodedb/src/control/server/http/auth.rs @@ -16,6 +16,7 @@ use nodedb_types::error::sqlstate; use crate::control::server::pgwire::types::error_map::numeric_code_to_sqlstate; use crate::config::auth::AuthMode; +use crate::control::security::bearer::VerifiedBearer; use crate::control::security::identity::AuthenticatedIdentity; use crate::control::server::session_auth; use crate::control::state::SharedState; @@ -50,17 +51,23 @@ pub struct AppState { /// provider's key set over the network, and the whole HTTP request path runs on /// Tokio worker threads, so blocking here will stall a worker at best and /// abort the request at worst. +/// +/// The request's session id is minted before verification. +/// The login's `AuthSuccess` entry and the request's `$auth.session_id` carry the same id. async fn try_validate_jwt( state: &AppState, token: &str, -) -> Option<( - AuthenticatedIdentity, - crate::control::security::jwks::registry::VerifiedJwtClaims, -)> { + peer_addr: &str, +) -> Result, ApiError> { if token.matches('.').count() != 2 { - return None; + return Ok(None); } - session_auth::authenticate_bearer_jwt(&state.shared, token).await + let session_id = crate::control::security::auth_context::generate_session_id()?; + Ok( + session_auth::authenticate_bearer_jwt(&state.shared, token, peer_addr, &session_id) + .await + .map(|(identity, claims)| (identity, VerifiedBearer { claims, session_id })), + ) } /// Enforce the TLS policy for a resolved HTTP identity. @@ -93,20 +100,23 @@ pub async fn resolve_identity( peer_addr: &str, transport: crate::control::security::tls_policy::TransportSecurity, ) -> Result { - let identity = resolve_identity_unchecked(headers, state, peer_addr).await?; + let (identity, _) = resolve_credentials_unchecked(headers, state, peer_addr).await?; enforce_transport_policy(state, &identity, transport, peer_addr)?; Ok(identity) } /// The credential half of [`resolve_identity`], without the transport check. /// +/// Each credential kind is tried once, so one refused JWT leaves one `AuthFailure` entry. +/// Returns the verified bearer token when the identity came from a JWT. +/// /// Private on purpose: an identity that has not been through /// [`enforce_transport_policy`] must not leave this module. -async fn resolve_identity_unchecked( +async fn resolve_credentials_unchecked( headers: &HeaderMap, state: &AppState, peer_addr: &str, -) -> Result { +) -> Result<(AuthenticatedIdentity, Option), ApiError> { if let Some(auth_header) = headers.get("authorization") { let auth_str = auth_header .to_str() @@ -116,15 +126,15 @@ async fn resolve_identity_unchecked( let token = token.trim(); // Try JWT first (token has 2 dots = JWT format). - if let Some((identity, _)) = try_validate_jwt(state, token).await { - return Ok(identity); + if let Some((identity, bearer)) = try_validate_jwt(state, token, peer_addr).await? { + return Ok((identity, Some(bearer))); } // Try API key. if let Some(identity) = session_auth::verify_api_key_identity(&state.shared, token, peer_addr, "HTTP") { - return Ok(identity); + return Ok((identity, None)); } return Err(ApiError::Unauthorized("invalid bearer token".into())); @@ -132,9 +142,11 @@ async fn resolve_identity_unchecked( } if state.auth_mode == AuthMode::Trust { - return session_auth::configured_trust_identity(&state.shared).ok_or_else(|| { - ApiError::Unauthorized("configured trust identity is unavailable".into()) - }); + return session_auth::configured_trust_identity(&state.shared) + .map(|identity| (identity, None)) + .ok_or_else(|| { + ApiError::Unauthorized("configured trust identity is unavailable".into()) + }); } Err(ApiError::Unauthorized( @@ -184,11 +196,13 @@ pub async fn resolve_auth( ), ApiError, > { - use crate::control::security::auth_context::{AuthContext, generate_session_id}; + use crate::control::security::auth_context::AuthContext; let (identity, verified_jwt) = resolve_auth_parts(headers, state, peer_addr, transport).await?; - let auth_ctx = match &verified_jwt { - Some(claims) => AuthContext::from_verified_jwt(claims, &identity, generate_session_id()?), + let auth_ctx = match verified_jwt { + Some(bearer) => { + AuthContext::from_verified_jwt(&bearer.claims, &identity, bearer.session_id) + } None => session_auth::build_auth_context(&identity)?, }; let auth_ctx = apply_on_deny_header(headers, auth_ctx); @@ -198,37 +212,19 @@ pub async fn resolve_auth( /// Resolve identity plus the raw ingredients a caller needs to build a /// [`RequestAuthScope`](crate::control::security::request_scope::RequestAuthScope) /// directly, instead of a pre-assembled `AuthContext`: the verified JWT -/// claims when authentication went through a JWT bearer token (so -/// `RequestAuthScope::builder().with_verified_jwt()` reproduces the exact -/// claim-derived enrichment `resolve_auth` gives `AuthContext::from_verified_jwt`), -/// or `None` for API key / trust / password auth. +/// claims and the request's session id when authentication went through a JWT +/// bearer token (so `RequestAuthScope::builder().with_verified_jwt()` reproduces +/// the exact claim-derived enrichment `resolve_auth` gives +/// `AuthContext::from_verified_jwt`), or `None` for API key / trust / password auth. pub(crate) async fn resolve_auth_parts( headers: &HeaderMap, state: &AppState, peer_addr: &str, transport: crate::control::security::tls_policy::TransportSecurity, -) -> Result< - ( - AuthenticatedIdentity, - Option, - ), - ApiError, -> { - if let Some(auth_header) = headers.get("authorization") - && let Ok(auth_str) = auth_header.to_str() - && let Some(token) = auth_str.strip_prefix("Bearer ") - { - let token = token.trim(); - if let Some((identity, verified_claims)) = try_validate_jwt(state, token).await { - // The JWT path returns before `resolve_identity`, so it enforces - // the transport policy itself rather than inheriting it. - enforce_transport_policy(state, &identity, transport, peer_addr)?; - return Ok((identity, Some(verified_claims))); - } - } - - let identity = resolve_identity(headers, state, peer_addr, transport).await?; - Ok((identity, None)) +) -> Result<(AuthenticatedIdentity, Option), ApiError> { + let (identity, bearer) = resolve_credentials_unchecked(headers, state, peer_addr).await?; + enforce_transport_policy(state, &identity, transport, peer_addr)?; + Ok((identity, bearer)) } /// Parse the `X-On-Deny` header into a `DenyMode`, if present and valid. @@ -276,20 +272,25 @@ pub(crate) fn on_deny_header_mode( /// different one. pub(crate) fn build_request_scope<'a, 'p>( identity: &'a AuthenticatedIdentity, - verified_jwt: Option<&'a crate::control::security::jwks::registry::VerifiedJwtClaims>, + verified_jwt: Option<&'a VerifiedBearer>, headers: &HeaderMap, state: &'a AppState, database_id: crate::types::DatabaseId, peer_addr: &'p str, ) -> crate::Result> { - crate::control::security::request_scope::RequestAuthScope::builder( + let builder = crate::control::security::request_scope::RequestAuthScope::builder( identity, state.shared.auth_stores(), ) .with_session_database(Some(database_id)) .with_on_deny(on_deny_header_mode(headers)) - .with_optional_verified_jwt(verified_jwt) - .build_for_client(peer_addr) + .with_optional_verified_jwt(verified_jwt.map(|bearer| &bearer.claims)); + // A bearer request keeps the session id its `AuthSuccess` entry recorded. + let builder = match verified_jwt { + Some(bearer) => builder.with_session_id(bearer.session_id.clone()), + None => builder, + }; + builder.build_for_client(peer_addr) } /// Apply [`on_deny_header_mode`] to an already-built `AuthContext`. diff --git a/nodedb/src/control/server/http/routes/promql/source/scan.rs b/nodedb/src/control/server/http/routes/promql/source/scan.rs index 6becb8847..c517a4405 100644 --- a/nodedb/src/control/server/http/routes/promql/source/scan.rs +++ b/nodedb/src/control/server/http/routes/promql/source/scan.rs @@ -50,6 +50,7 @@ pub(crate) async fn scan_metric( valid_at_ms: None, having: Vec::new(), tiering: metric.tiering.clone(), + output_columns: Vec::new(), }); let payload = dispatch_for_identity(DispatchRequest { state: scope.state, diff --git a/nodedb/src/control/server/native/dispatch/admission_op.rs b/nodedb/src/control/server/native/dispatch/admission_op.rs index e439540b9..5a0237666 100644 --- a/nodedb/src/control/server/native/dispatch/admission_op.rs +++ b/nodedb/src/control/server/native/dispatch/admission_op.rs @@ -56,7 +56,8 @@ pub(crate) fn admission_operation(op: OpCode) -> &'static str { | OpCode::DocumentRegister | OpCode::DocumentDropIndex | OpCode::DocumentBatchInsert - | OpCode::ColumnarInsert => "point_put", + | OpCode::ColumnarInsert + | OpCode::ArrayPutBatch => "point_put", // Scans. OpCode::RangeScan diff --git a/nodedb/src/control/server/native/dispatch/array_batch.rs b/nodedb/src/control/server/native/dispatch/array_batch.rs new file mode 100644 index 000000000..4b7658475 --- /dev/null +++ b/nodedb/src/control/server/native/dispatch/array_batch.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Native transport for `ArrayPutBatch`: decode the columnar cells and hand +//! them to the shared array bulk-load core, which checks, chunks and +//! dispatches them as `COPY FROM STDIN` does. The frame cap bounds +//! one request. +//! +//! Outside a transaction block each chunk of the request commits on its +//! own, and a bad cell stops the request with the committed row count. +//! Inside a block the request stages whole or aborts the block. + +use nodedb_types::protocol::{ArrayPutBatch, NativeResponse, TextFields}; + +use crate::control::server::native::sqlstate_code::sqlstate_error; +use crate::control::server::shared::ddl::DdlResult; +use crate::control::server::shared::ddl::neutral::collection::copy_from::{ + ArrayLoadCtx, ArrayLoader, LoadInput, +}; +use crate::control::server::shared::session::DmlTxnCtx; + +use super::{DispatchCtx, ddl_result_to_native}; + +/// Load the request's cells into the array `collection` names. Answers the +/// row count as a `COPY` status. +pub(crate) async fn handle_array_put_batch( + ctx: &DispatchCtx<'_>, + seq: u64, + fields: &TextFields, +) -> NativeResponse { + let Some(array) = fields.collection.as_deref() else { + return sqlstate_error(seq, "42601", "ArrayPutBatch: missing 'collection' field"); + }; + let Some(payload) = fields.payload.as_deref() else { + return sqlstate_error(seq, "42601", "ArrayPutBatch: missing 'payload' field"); + }; + let batch: ArrayPutBatch = match zerompk::from_msgpack(payload) { + Ok(batch) => batch, + Err(e) => { + return sqlstate_error( + seq, + "22P02", + format!("ArrayPutBatch: payload is not an ArrayPutBatch: {e}"), + ); + } + }; + let txn_ctx = DmlTxnCtx { + sessions: ctx.sessions, + session_id: ctx.peer_addr.into(), + }; + let load = ArrayLoadCtx { + state: ctx.state, + identity: ctx.identity, + database_id: ctx.database_id(), + txn_ctx: &txn_ctx, + }; + let _request = ctx.state.tenant_request_guard(ctx.tenant_id()); + let result = async { + let mut loader = ArrayLoader::open(&load, array, LoadInput::Cells)?; + loader.feed_cells(&load, &batch.dims, &batch.attrs).await?; + loader.finish(&load).await + } + .await + .map(|rows| { + vec![DdlResult::Status { + command: "COPY".to_string(), + rows_affected: Some(rows), + }] + }); + ddl_result_to_native(seq, result) +} diff --git a/nodedb/src/control/server/native/dispatch/auth.rs b/nodedb/src/control/server/native/dispatch/auth.rs index 30bcf5efe..0c4ceb18b 100644 --- a/nodedb/src/control/server/native/dispatch/auth.rs +++ b/nodedb/src/control/server/native/dispatch/auth.rs @@ -4,8 +4,8 @@ use nodedb_types::protocol::{AuthMethod as ProtoAuth, NativeResponse}; +use crate::control::security::bearer::VerifiedBearer; use crate::control::security::identity::AuthenticatedIdentity; -use crate::control::security::jwks::registry::VerifiedJwtClaims; use crate::control::state::SharedState; /// Result of a native-protocol authentication attempt. @@ -16,12 +16,13 @@ use crate::control::state::SharedState; /// groups, permissions, metadata) via `AuthContext::from_verified_jwt` /// without re-deriving authority from the token: authority always comes from /// `identity`, never from the claims. +/// It also carries the session id the login's `AuthSuccess` entry recorded. pub(crate) struct NativeAuthOutcome { pub(crate) identity: AuthenticatedIdentity, /// Non-empty when the account is in a password grace period or /// `must_change_password` is set. pub(crate) warning: Option, - pub(crate) verified_jwt: Option, + pub(crate) verified_jwt: Option, } /// Authenticate a native protocol client. @@ -36,22 +37,20 @@ pub(crate) async fn handle_auth( peer_addr: &str, ) -> crate::Result { if let ProtoAuth::OidcBearer { token, .. } = auth { - let (identity, verified_jwt) = - crate::control::security::oidc::verify_bearer_token(state, token).await?; - state.audit_record( - crate::control::security::audit::AuditEvent::AuthSuccess, - Some(identity.tenant_id), + // The session id exists before verification, so the `AuthSuccess` entry and + // `$auth.session_id` carry the same id. + let session_id = crate::control::security::auth_context::generate_session_id()?; + let (identity, claims) = crate::control::security::oidc::verify_bearer_token( + state, + token, peer_addr, - &format!( - "OIDC bearer login: sub={} method=oidc_bearer", - identity.username - ), - ); - state.auth_metrics.record_auth_success("oidc_bearer"); + &session_id, + ) + .await?; return Ok(NativeAuthOutcome { identity, warning: None, - verified_jwt: Some(verified_jwt), + verified_jwt: Some(VerifiedBearer { claims, session_id }), }); } diff --git a/nodedb/src/control/server/native/dispatch/direct_ops.rs b/nodedb/src/control/server/native/dispatch/direct_ops.rs index 5a1d0878b..c78caa762 100644 --- a/nodedb/src/control/server/native/dispatch/direct_ops.rs +++ b/nodedb/src/control/server/native/dispatch/direct_ops.rs @@ -96,6 +96,21 @@ pub(crate) async fn handle_direct_op( }; let vshard_id = ctx.task_vshard(&plan, fields.document_id.as_deref(), &collection); + // Refuse a read of a redacted column outside the projection, before RLS + // adds the read policy's own predicates. + if let Err(e) = crate::control::planner::redaction_refusal::refuse_redacted_expression_reads( + &plan, + tenant_id, + ctx.database_id(), + &crate::control::planner::redaction_refusal::RedactionReadCheck { + auth: ctx.auth_context(), + store: &ctx.state.redaction, + catalog: ctx.state.credentials.catalog(), + }, + ) { + return error_to_native_with_sqlstate(seq, "42501", &e); + } + // Apply RLS before any special Control-Plane orchestration can observe the plan. if let Err(e) = crate::control::planner::rls_injection::inject_rls_for_single_plan( tenant_id.as_u64(), @@ -116,7 +131,7 @@ pub(crate) async fn handle_direct_op( ctx.auth_context(), &ctx.state.redaction, ) { - return error_to_native_with_sqlstate(seq, "0A000", &e); + return error_to_native_with_sqlstate(seq, "42501", &e); } // Extracted before `plan` moves; a no-op when metering is disabled (the default). diff --git a/nodedb/src/control/server/native/dispatch/graph_match.rs b/nodedb/src/control/server/native/dispatch/graph_match.rs index 7e7ad061d..94d404ab7 100644 --- a/nodedb/src/control/server/native/dispatch/graph_match.rs +++ b/nodedb/src/control/server/native/dispatch/graph_match.rs @@ -65,7 +65,7 @@ pub(crate) async fn handle_graph_match( ctx.auth_context(), &ctx.state.redaction, ) { - return error_to_native_with_sqlstate(seq, "0A000", &error); + return error_to_native_with_sqlstate(seq, "42501", &error); } // Stamp the active transaction id so MATCH reads resolve this connection's diff --git a/nodedb/src/control/server/native/dispatch/mod.rs b/nodedb/src/control/server/native/dispatch/mod.rs index c691fa0ef..a6327a320 100644 --- a/nodedb/src/control/server/native/dispatch/mod.rs +++ b/nodedb/src/control/server/native/dispatch/mod.rs @@ -3,6 +3,7 @@ //! Per-opcode dispatch handlers for the native protocol. mod admission_op; +mod array_batch; mod auth; mod cluster_array; mod conversion; @@ -32,6 +33,7 @@ mod transaction; mod transaction_savepoint; pub(crate) use admission_op::admission_operation; +pub(crate) use array_batch::handle_array_put_batch; pub(crate) use auth::{NativeAuthOutcome, handle_auth, handle_ping}; pub(crate) use conversion::{ apply_dml_outcome, ddl_result_to_native, dml_fold_error_to_native, error_code_to_native, diff --git a/nodedb/src/control/server/native/dispatch/plan_builder/spatial.rs b/nodedb/src/control/server/native/dispatch/plan_builder/spatial.rs index e89731417..062c3592a 100644 --- a/nodedb/src/control/server/native/dispatch/plan_builder/spatial.rs +++ b/nodedb/src/control/server/native/dispatch/plan_builder/spatial.rs @@ -21,12 +21,13 @@ pub(crate) async fn build_scan( detail: "missing 'query_geometry'".to_string(), })?; - let geometry: nodedb_types::geometry::Geometry = + // GeoJSON, with the optional `crs` member naming a supported SRID. + let geometry: nodedb_types::geometry::SridGeometry = crate::util::bounded_json::from_slice(raw_bytes).map_err(|e| crate::Error::BadRequest { detail: format!("invalid query geometry: {e}"), })?; - let issues = nodedb_spatial::validate::validate_geometry(&geometry); + let issues = nodedb_spatial::validate::validate_geometry(&geometry.geometry); if !issues.is_empty() { return Err(crate::Error::BadRequest { detail: format!("invalid query geometry: {}", issues.join("; ")), diff --git a/nodedb/src/control/server/native/dispatch/plan_builder/timeseries.rs b/nodedb/src/control/server/native/dispatch/plan_builder/timeseries.rs index dbed7d20c..f9e434343 100644 --- a/nodedb/src/control/server/native/dispatch/plan_builder/timeseries.rs +++ b/nodedb/src/control/server/native/dispatch/plan_builder/timeseries.rs @@ -50,6 +50,7 @@ pub(crate) async fn build_scan( valid_at_ms: None, having: Vec::new(), tiering: None, + output_columns: Vec::new(), })) } diff --git a/nodedb/src/control/server/native/session/auth.rs b/nodedb/src/control/server/native/session/auth.rs index 1761a125f..d7f559ad2 100644 --- a/nodedb/src/control/server/native/session/auth.rs +++ b/nodedb/src/control/server/native/session/auth.rs @@ -136,13 +136,15 @@ impl NativeSession { // Every other auth method keeps the identity-only context. // The context is built before admission capacity is acquired, // so a session-id generation error leaves no state behind. + // A bearer login reuses the session id its `AuthSuccess` entry recorded. let auth_context_result = match &verified_jwt { - Some(claims) => crate::control::security::auth_context::generate_session_id() - .map(|session_id| { - crate::control::security::auth_context::AuthContext::from_verified_jwt( - claims, &identity, session_id, - ) - }), + Some(bearer) => Ok( + crate::control::security::auth_context::AuthContext::from_verified_jwt( + &bearer.claims, + &identity, + bearer.session_id.clone(), + ), + ), None => super::super::super::session_auth::build_auth_context(&identity), }; let mut auth_context = match auth_context_result { @@ -213,7 +215,7 @@ impl NativeSession { // rebuild its per-request `RequestAuthScope` with the same // claim-derived enrichment (see `session::request::handle_request`), // not just the auth frame's own response. - self.verified_jwt = verified_jwt; + self.verified_jwt = verified_jwt.map(|bearer| bearer.claims); self.cleanup.publish_identity(identity.clone()); self.identity = Some(identity); resp diff --git a/nodedb/src/control/server/native/session/request.rs b/nodedb/src/control/server/native/session/request.rs index 31c53c79c..b8b68e80d 100644 --- a/nodedb/src/control/server/native/session/request.rs +++ b/nodedb/src/control/server/native/session/request.rs @@ -451,6 +451,9 @@ async fn route_opcode( dispatch::handle_direct_op(&ctx, seq, op, fields).await } + // Array bulk load: the shared array bulk-load core. + OpCode::ArrayPutBatch => dispatch::handle_array_put_batch(&ctx, seq, fields).await, + // Copy from file. OpCode::CopyFrom => { let sql = match &fields.sql { diff --git a/nodedb/src/control/server/payload_merge.rs b/nodedb/src/control/server/payload_merge.rs index 6ce98a092..5ec2e7497 100644 --- a/nodedb/src/control/server/payload_merge.rs +++ b/nodedb/src/control/server/payload_merge.rs @@ -9,50 +9,94 @@ //! a separate trailing array that the decoder silently ignores (truncating the //! result to the first chunk). //! +//! [`merge_msgpack_arrays`] takes only whole msgpack arrays. A relational +//! read also answers a point read's one row as a bare msgpack map, which +//! [`payload_rows`] and [`merge_payload_rows`] read as one row. A payload of +//! any other shape is a typed error: wrapping it as one element will hand +//! the caller rows that no core produced. +//! //! Used by both `local_dispatch::collect_bounded_response` (the per-request -//! bounded collector) and `exchange::gather` (the cross-core/vShard gather). +//! bounded collector) and `exchange::gather_outcome` (the cross-core/vShard +//! gather). use nodedb_query::msgpack_scan; /// Extract individual msgpack elements from one msgpack-array payload. /// -/// If the payload is not a valid msgpack array, it is returned as a single -/// element with a warning logged (matching the gather fallback). -pub(crate) fn extract_msgpack_elements(payload: &[u8]) -> Vec> { +/// An empty payload holds no elements. A payload that is not one whole +/// msgpack array is a `Codec` error: a non-array header, an element that +/// does not decode, or bytes after the last element. +pub(crate) fn extract_msgpack_elements(payload: &[u8]) -> crate::Result>> { if payload.is_empty() { - return Vec::new(); + return Ok(Vec::new()); } let Some((count, mut pos)) = msgpack_scan::array_header(payload, 0) else { - tracing::warn!( - payload_len = payload.len(), - "payload_merge: payload is not a msgpack array; treating as single element" - ); - return vec![payload.to_vec()]; + return Err(crate::Error::Codec { + detail: format!( + "payload_merge: a {}-byte payload starting with byte 0x{:02x} is not a \ + msgpack array; only a row-array payload can merge into rows", + payload.len(), + payload[0] + ), + }); }; - let mut rows = Vec::with_capacity(count); - for _ in 0..count { - if pos >= payload.len() { - break; - } + // Every element takes at least one byte, so the header cannot reserve + // more elements than the payload has bytes left. + let mut rows = Vec::with_capacity(count.min(payload.len() - pos)); + for index in 0..count { let start = pos; - match msgpack_scan::skip_value(payload, pos) { - Some(next) => { - rows.push(payload[start..next].to_vec()); - pos = next; - } - None => { - tracing::warn!( - pos, - payload_len = payload.len(), - "payload_merge: could not skip msgpack element; stopping early" - ); - break; - } - } - } - rows + let Some(next) = msgpack_scan::skip_value(payload, pos) else { + return Err(crate::Error::Codec { + detail: format!( + "payload_merge: element {index} of a {count}-element msgpack array does \ + not decode at offset {start} of a {}-byte payload", + payload.len() + ), + }); + }; + rows.push(payload[start..next].to_vec()); + pos = next; + } + if pos != payload.len() { + return Err(crate::Error::Codec { + detail: format!( + "payload_merge: {} bytes follow the {count}-element msgpack array in a \ + {}-byte payload; a payload must hold exactly one array", + payload.len() - pos, + payload.len() + ), + }); + } + Ok(rows) +} + +/// The rows one relational response payload holds, each as its raw msgpack +/// bytes. +/// +/// A scan answers a msgpack array of rows. A point read answers its one row +/// as a bare msgpack map, and that map is the payload's one row. An empty +/// payload holds no rows. Any other payload is a `Codec` error from +/// [`extract_msgpack_elements`]. +pub(crate) fn payload_rows(payload: &[u8]) -> crate::Result>> { + let single_row = msgpack_scan::map_header(payload, 0).is_some() + && msgpack_scan::skip_value(payload, 0) == Some(payload.len()); + if single_row { + return Ok(vec![payload.to_vec()]); + } + extract_msgpack_elements(payload) +} + +/// Merge the rows of several relational response payloads into one msgpack +/// array, in order. Each payload holds its rows in a shape [`payload_rows`] +/// reads. +pub(crate) fn merge_payload_rows(payloads: &[Vec]) -> crate::Result> { + let mut rows: Vec> = Vec::new(); + for payload in payloads { + rows.extend(payload_rows(payload)?); + } + Ok(encode_msgpack_array(&rows)) } /// Encode a list of pre-extracted msgpack elements into a single msgpack array. @@ -79,12 +123,14 @@ pub(crate) fn encode_msgpack_array(rows: &[Vec]) -> Vec { /// Merge several standalone msgpack-array payloads into one msgpack array /// containing every element from every input array, in order. -pub(crate) fn merge_msgpack_arrays(payloads: &[Vec]) -> Vec { +/// +/// Fails with the first payload [`extract_msgpack_elements`] refuses. +pub(crate) fn merge_msgpack_arrays(payloads: &[Vec]) -> crate::Result> { let mut elements: Vec> = Vec::new(); for payload in payloads { - elements.extend(extract_msgpack_elements(payload)); + elements.extend(extract_msgpack_elements(payload)?); } - encode_msgpack_array(&elements) + Ok(encode_msgpack_array(&elements)) } #[cfg(test)] @@ -98,6 +144,10 @@ mod tests { encode_msgpack_array(&rows) } + fn is_codec_error(result: crate::Result>>) -> bool { + matches!(result, Err(crate::Error::Codec { .. })) + } + #[test] fn merge_concatenates_all_elements() { let chunks = vec![ @@ -105,8 +155,8 @@ mod tests { int_array(1000, 1000), int_array(2000, 500), ]; - let merged = merge_msgpack_arrays(&chunks); - let elements = extract_msgpack_elements(&merged); + let merged = merge_msgpack_arrays(&chunks).expect("merge"); + let elements = extract_msgpack_elements(&merged).expect("extract"); assert_eq!( elements.len(), 2500, @@ -117,13 +167,118 @@ mod tests { #[test] fn single_array_round_trips() { let one = int_array(0, 3); - let merged = merge_msgpack_arrays(std::slice::from_ref(&one)); - assert_eq!(extract_msgpack_elements(&merged).len(), 3); + let merged = merge_msgpack_arrays(std::slice::from_ref(&one)).expect("merge"); + assert_eq!(extract_msgpack_elements(&merged).expect("extract").len(), 3); } #[test] fn empty_input_is_empty_array() { - let merged = merge_msgpack_arrays(&[]); - assert_eq!(extract_msgpack_elements(&merged).len(), 0); + let merged = merge_msgpack_arrays(&[]).expect("merge"); + assert_eq!(extract_msgpack_elements(&merged).expect("extract").len(), 0); + } + + #[test] + fn an_empty_payload_holds_no_elements() { + assert!(extract_msgpack_elements(&[]).expect("extract").is_empty()); + } + + /// A msgpack map, such as an array slice envelope, is not a row array. + #[test] + fn a_map_payload_is_refused() { + // fixmap of 1: "k" => 1. + assert!(is_codec_error(extract_msgpack_elements(&[ + 0x81, 0xa1, b'k', 1 + ]))); + } + + /// Bytes that are not msgpack at all, such as a raw little-endian LSN, + /// are refused rather than wrapped as one element. + #[test] + fn raw_bytes_are_refused() { + assert!(is_codec_error(extract_msgpack_elements( + &42u64.to_le_bytes() + ))); + } + + /// A header that claims more elements than the payload holds is refused, + /// never cut short at the bytes that are there. + #[test] + fn a_truncated_array_is_refused() { + let mut payload = int_array(0, 3); + payload.pop(); + assert!(is_codec_error(extract_msgpack_elements(&payload))); + } + + /// Two arrays back to back are two payloads, not one. + #[test] + fn bytes_after_the_array_are_refused() { + let mut payload = int_array(0, 2); + payload.extend_from_slice(&int_array(2, 2)); + assert!(is_codec_error(extract_msgpack_elements(&payload))); + } + + #[test] + fn a_non_array_payload_fails_the_merge() { + let chunks = vec![int_array(0, 2), vec![0x81, 0xa1, b'k', 1]]; + assert!(matches!( + merge_msgpack_arrays(&chunks), + Err(crate::Error::Codec { .. }) + )); + } + + /// fixmap of 1: "k" => `value`. + fn map_row(value: u8) -> Vec { + vec![0x81, 0xa1, b'k', value] + } + + #[test] + fn a_row_array_holds_its_rows() { + let payload = encode_msgpack_array(&[map_row(1), map_row(2)]); + assert_eq!( + payload_rows(&payload).expect("rows"), + vec![map_row(1), map_row(2)] + ); + } + + #[test] + fn a_point_read_map_is_one_row() { + assert_eq!(payload_rows(&map_row(3)).expect("rows"), vec![map_row(3)]); + } + + #[test] + fn an_empty_payload_holds_no_rows() { + assert!(payload_rows(&[]).expect("rows").is_empty()); + } + + /// A map followed by more bytes is not one row. + #[test] + fn a_map_with_trailing_bytes_is_not_a_row() { + let mut payload = map_row(1); + payload.push(0x01); + assert!(is_codec_error(payload_rows(&payload))); + } + + #[test] + fn a_scalar_payload_holds_no_rows() { + assert!(is_codec_error(payload_rows(&[0x2a]))); + } + + #[test] + fn row_payloads_merge_point_reads_and_scans_in_order() { + let payloads = vec![encode_msgpack_array(&[map_row(1)]), map_row(2), Vec::new()]; + let merged = merge_payload_rows(&payloads).expect("merge"); + assert_eq!( + extract_msgpack_elements(&merged).expect("extract"), + vec![map_row(1), map_row(2)] + ); + } + + #[test] + fn a_scalar_payload_fails_the_row_merge() { + let payloads = vec![encode_msgpack_array(&[map_row(1)]), vec![0x2a]]; + assert!(matches!( + merge_payload_rows(&payloads), + Err(crate::Error::Codec { .. }) + )); } } diff --git a/nodedb/src/control/server/pgwire/catalog/audit_log.rs b/nodedb/src/control/server/pgwire/catalog/audit_log.rs index 77550dc3c..c32eab588 100644 --- a/nodedb/src/control/server/pgwire/catalog/audit_log.rs +++ b/nodedb/src/control/server/pgwire/catalog/audit_log.rs @@ -2,15 +2,20 @@ //! `_system.audit_log` row producer. //! -//! Materializes the durable audit log as msgpack rows. There is no SQL-level +//! Materializes the audit log as msgpack rows. There is no SQL-level //! pushdown here — the caller's row budget is enforced by an unconditional //! materialize cap (`MATERIALIZE_LIMIT`) that bounds memory regardless of the -//! client query. +//! client query. The cap keeps the newest entries. +//! +//! Rows come from [`SharedState::recent_audit_entries`]. +//! That reader merges the in-memory buffer with the catalog. +//! A scan therefore returns every entry recorded before it started. //! //! Columns: //! - `seq` — monotonic sequence number (int8). //! - `timestamp_us` — UTC microseconds since epoch (int8). -//! - `event` — event discriminant name (text, e.g. "AuthSuccess"). +//! - `event` — snake_case event name (text, e.g. "auth_success"). +//! It matches the `event` column of every `SHOW AUDIT` command. //! - `tenant_id` — tenant identifier, 0 if not applicable (int8). //! - `source` — source IP or node identifier (text). //! - `detail` — human-readable event detail (text). @@ -44,82 +49,29 @@ pub fn audit_log( }); } - let mut rows: Vec> = Vec::new(); - - // Merge catalog-persisted entries with the in-memory tail. The audit log - // is flushed from memory to the catalog by a periodic background timer - // (`SharedState::flush_audit_log`), so the in-memory log always holds the - // most recent entries that have not yet been persisted. Reading only the - // catalog would hide those entries from operators querying - // `_system.audit_log` between flush ticks. Dedupe by `seq`. - let mut seen: std::collections::HashSet = std::collections::HashSet::new(); - - { - let catalog = state.credentials.catalog(); - let entries = catalog - .load_audit_entries_ranged(1, u64::MAX, 0, u64::MAX, MATERIALIZE_LIMIT) - .map_err(|e| crate::Error::Storage { - engine: "catalog".to_string(), - detail: e.to_string(), - })?; - for e in entries { - if !seen.insert(e.seq) { - continue; - } - let mut r: HashMap = HashMap::with_capacity(7); - r.insert("seq".into(), Value::Integer(e.seq as i64)); - r.insert("timestamp_us".into(), Value::Integer(e.timestamp_us as i64)); - r.insert("event".into(), Value::String(e.event)); - r.insert( - "tenant_id".into(), - Value::Integer(e.tenant_id.unwrap_or(0) as i64), - ); - r.insert("source".into(), Value::String(e.source)); - r.insert("detail".into(), Value::String(e.detail)); - r.insert( - "prev_hash".into(), - if e.prev_hash.is_empty() { - Value::Null - } else { - Value::String(e.prev_hash) - }, - ); - rows.push(encode_row(r)?); - } - } + let entries = state.recent_audit_entries(MATERIALIZE_LIMIT, |_| true)?; - let log = match state.audit.lock() { - Ok(l) => l, - Err(p) => p.into_inner(), - }; - let all = log.all(); - let skip = all.len().saturating_sub(MATERIALIZE_LIMIT); - for entry in all.iter().skip(skip) { - if !seen.insert(entry.seq) { - continue; - } - if rows.len() >= MATERIALIZE_LIMIT { - break; - } + let mut rows: Vec> = Vec::with_capacity(entries.len()); + for e in entries { let mut r: HashMap = HashMap::with_capacity(7); - r.insert("seq".into(), Value::Integer(entry.seq as i64)); + r.insert("seq".into(), Value::Integer(e.seq as i64)); + r.insert("timestamp_us".into(), Value::Integer(e.timestamp_us as i64)); r.insert( - "timestamp_us".into(), - Value::Integer(entry.timestamp_us as i64), + "event".into(), + Value::String(e.event.snake_name().to_string()), ); - r.insert("event".into(), Value::String(format!("{:?}", entry.event))); r.insert( "tenant_id".into(), - Value::Integer(entry.tenant_id.map_or(0i64, |t| t.as_u64() as i64)), + Value::Integer(e.tenant_id.map_or(0i64, |t| t.as_u64() as i64)), ); - r.insert("source".into(), Value::String(entry.source.clone())); - r.insert("detail".into(), Value::String(entry.detail.clone())); + r.insert("source".into(), Value::String(e.source)); + r.insert("detail".into(), Value::String(e.detail)); r.insert( "prev_hash".into(), - if entry.prev_hash.is_empty() { + if e.prev_hash.is_empty() { Value::Null } else { - Value::String(entry.prev_hash.clone()) + Value::String(e.prev_hash) }, ); rows.push(encode_row(r)?); diff --git a/nodedb/src/control/server/pgwire/factory/server.rs b/nodedb/src/control/server/pgwire/factory/server.rs index 49a9a6389..f2e0271c0 100644 --- a/nodedb/src/control/server/pgwire/factory/server.rs +++ b/nodedb/src/control/server/pgwire/factory/server.rs @@ -17,7 +17,7 @@ use crate::control::state::SharedState; use super::super::connection_identity::PgConnectionContext; use super::super::connection_registry::{ConnectionRegistry, ConnectionRegistryError}; -use super::super::handler::{NodeDbCopyHandler, NodeDbPgHandler}; +use super::super::handler::{ArrayCopyState, NodeDbCopyHandler, NodeDbPgHandler}; use super::super::wire_safe_error::WireSafeErrorHandler; use super::auth::NodeDbAuthSource; use super::provider::nodedb_parameter_provider; @@ -65,6 +65,7 @@ pub struct NodeDbPgHandlerFactory { state: Arc, sessions: Arc, restore_state: Arc, + array_copy: Arc, registry: Arc, } @@ -76,6 +77,7 @@ impl NodeDbPgHandlerFactory { state, sessions: Arc::new(SessionStore::new()), restore_state: Arc::new(crate::control::backup::RestoreState::new()), + array_copy: Arc::new(ArrayCopyState::new()), registry: Arc::new(ConnectionRegistry::new()), } } @@ -106,6 +108,7 @@ impl NodeDbPgHandlerFactory { self.auth_mode.clone(), Arc::clone(&self.sessions), Arc::clone(&self.restore_state), + Arc::clone(&self.array_copy), Arc::clone(&self.registry), context.id, )); @@ -133,6 +136,8 @@ impl NodeDbPgHandlerFactory { copy: Arc::new(NodeDbCopyHandler { state: Arc::clone(&self.state), restore_state: Arc::clone(&self.restore_state), + array_copy: Arc::clone(&self.array_copy), + sessions: Arc::clone(&self.sessions), connection_id: context.id, }), cancel: Arc::new(NoopHandler), @@ -146,6 +151,7 @@ impl NodeDbPgHandlerFactory { self.auth_mode.clone(), Arc::clone(&self.sessions), Arc::clone(&self.restore_state), + Arc::clone(&self.array_copy), Arc::clone(&self.registry), id, ); @@ -200,6 +206,10 @@ impl NodeDbPgHandlerFactory { { warn!(connection_id = id.get(), %peer_addr, "pgwire restore cleanup panicked"); } + if std::panic::catch_unwind(AssertUnwindSafe(|| self.array_copy.cancel(id.get()))).is_err() + { + warn!(connection_id = id.get(), %peer_addr, "pgwire array COPY cleanup panicked"); + } if std::panic::catch_unwind(AssertUnwindSafe(|| { self.sessions .cleanup_listen_on_disconnect(id.into(), &self.state.notify_bus); diff --git a/nodedb/src/control/server/pgwire/factory/startup.rs b/nodedb/src/control/server/pgwire/factory/startup.rs index d7279c10e..3b3fd9ef2 100644 --- a/nodedb/src/control/server/pgwire/factory/startup.rs +++ b/nodedb/src/control/server/pgwire/factory/startup.rs @@ -22,7 +22,7 @@ use crate::control::server::session_auth::identity::stored_user_identity; use crate::control::state::SharedState; use super::super::handler::NodeDbPgHandler; -use super::super::types::{error_to_pg_in_context, sqlstate_error}; +use super::super::types::error_to_sqlstate; use super::provider::NodeDbParameterProvider; /// Enum dispatch for startup handler — avoids dyn trait object issues. @@ -53,6 +53,10 @@ pub(crate) enum AuthStartup { /// /// A catalog read failure is refused for the same reason — it says nothing /// about whether the database exists, so it cannot become a fallback. +/// +/// Every refusal is `FATAL`. A non-fatal error leaves the connection open, and +/// on the SCRAM path, where ReadyForQuery is already sent, the client reads the +/// error as the reply to its first statement. fn bind_startup_database( client: &C, session_id: crate::control::server::shared::session::SessionId, @@ -68,8 +72,11 @@ fn bind_startup_database( .credentials .catalog() .get_database_id_by_name(&db_name) - .map_err(|e| error_to_pg_in_context("catalog lookup failed", &e))? - .ok_or_else(|| sqlstate_error("3D000", &format!("database '{db_name}' does not exist")))?; + .map_err(|e| { + let (_, code, message) = error_to_sqlstate(&e); + startup_refusal(code, format!("catalog lookup failed: {message}")) + })? + .ok_or_else(|| startup_refusal("3D000", format!("database '{db_name}' does not exist")))?; handler.sessions.set_current_database(session_id, db_id); Ok(()) @@ -108,9 +115,9 @@ fn admit_connection( .sessions .set_admission_permit(handler.session_id, permit) { - return Err(sqlstate_error( + return Err(startup_refusal( nodedb_types::error::sqlstate::INTERNAL_ERROR, - "internal error: connection session is not registered", + "internal error: connection session is not registered".to_owned(), )); } Ok(()) @@ -122,10 +129,18 @@ fn admit_connection( /// native protocol reports for the same condition, so a cap refusal reads the /// same on both protocols. fn admission_refusal(error: &AdmissionError) -> PgWireError { + startup_refusal( + nodedb_types::error::sqlstate::QUOTA_EXCEEDED, + error.to_string(), + ) +} + +/// A `FATAL` startup refusal. pgwire closes the connection after sending it. +fn startup_refusal(code: &str, message: String) -> PgWireError { PgWireError::UserError(Box::new(ErrorInfo::new( "FATAL".to_owned(), - nodedb_types::error::sqlstate::QUOTA_EXCEEDED.to_owned(), - error.to_string(), + code.to_owned(), + message, ))) } diff --git a/nodedb/src/control/server/pgwire/handler/array_copy.rs b/nodedb/src/control/server/pgwire/handler/array_copy.rs new file mode 100644 index 000000000..3acef04ac --- /dev/null +++ b/nodedb/src/control/server/pgwire/handler/array_copy.rs @@ -0,0 +1,272 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! pgwire transport for `COPY FROM STDIN`. +//! +//! The statement opens a COPY-in stream and an array load for the +//! connection. Each `CopyData` message streams through the shared array +//! bulk-load core, which checks its rows and dispatches each chunk they +//! fill. `CopyDone` loads the last record and the pending chunks. A +//! connection holds one incomplete record and a few chunks, whatever the +//! input size. + +use std::collections::HashMap; +use std::sync::Mutex; + +use futures::stream; +use nodedb_sql::ddl_ast::statement::{MiscStmt, NodedbStatement}; +use pgwire::api::results::{CopyResponse, Response}; +use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; + +use crate::control::planner::context::{BulkLoadRoute, QueryContext}; +use crate::control::security::identity::AuthenticatedIdentity; +use crate::control::server::shared::ddl::DdlError; +use crate::control::server::shared::ddl::neutral::collection::copy_from::{ + ArrayLoadCtx, ArrayLoader, LoadInput, TextFormat, +}; +use crate::control::server::shared::session::{ + DmlTxnCtx, SessionId, SessionStore, TransactionState, +}; +use crate::control::state::SharedState; +use crate::types::DatabaseId; + +use super::core::NodeDbPgHandler; + +/// One connection's open `COPY FROM STDIN`. +pub struct ArrayCopyPending { + identity: AuthenticatedIdentity, + database_id: DatabaseId, + loader: ArrayLoader, +} + +/// Connection-keyed map of open array COPY loads. +#[derive(Default)] +pub struct ArrayCopyState { + inner: Mutex>, +} + +impl ArrayCopyState { + pub fn new() -> Self { + Self::default() + } + + fn put(&self, conn_id: u64, pending: ArrayCopyPending) { + let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + g.insert(conn_id, pending); + } + + pub fn take(&self, conn_id: u64) -> Option { + let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); + g.remove(&conn_id) + } + + pub fn cancel(&self, conn_id: u64) { + let _ = self.take(conn_id); + } +} + +impl NodeDbPgHandler { + /// The COPY-stream response for `sql` when it is a streaming COPY shape: + /// backup, restore, or a load into an array. `None` for any other + /// statement, which runs on the regular path. + pub(super) async fn copy_stream_response( + &self, + identity: &AuthenticatedIdentity, + session_id: SessionId, + sql: &str, + ) -> Option> { + if let Some(intent) = crate::control::backup::detect(sql) { + return Some(self.intent_to_response(identity, session_id, intent).await); + } + self.array_copy_in(identity, session_id, sql).await + } + + /// Open `COPY FROM STDIN`. `None` when `sql` is not one, or + /// names no array: that statement runs on the regular path, which + /// answers a non-array target as before. + async fn array_copy_in( + &self, + identity: &AuthenticatedIdentity, + session_id: SessionId, + sql: &str, + ) -> Option> { + if !sql.trim_start().get(..5)?.eq_ignore_ascii_case("COPY ") { + return None; + } + let Some(Ok(NodedbStatement::Misc(MiscStmt::CopyFromStdin { + collection, + format, + delimiter, + header, + }))) = nodedb_sql::ddl_ast::parse(sql) + else { + return None; + }; + let database_id = self + .sessions + .get_current_database(session_id) + .unwrap_or(DatabaseId::DEFAULT); + let route = QueryContext::for_state(&self.state).bulk_load_route( + &collection, + identity.tenant_id, + database_id, + ); + let Ok(BulkLoadRoute::ArrayCells(view)) = route else { + return None; + }; + let open = ArrayCopyOpen { + database_id, + array: &collection, + columns: view.dims.len() + view.attrs.len(), + format: TextFormat { + format, + delimiter, + header, + }, + }; + Some(self.open_array_copy(identity, session_id, open).await) + } + + async fn open_array_copy( + &self, + identity: &AuthenticatedIdentity, + session_id: SessionId, + open: ArrayCopyOpen<'_>, + ) -> PgWireResult { + let ArrayCopyOpen { + database_id, + array, + columns, + format, + } = open; + self.admit_statement(identity, session_id, database_id) + .await?; + self.state + .check_tenant_quota(identity.tenant_id) + .map_err(|e| super::super::types::error_to_pg(&e))?; + let SessionId::Connection(connection_id) = session_id else { + return Err(user_error( + nodedb_types::error::sqlstate::INTERNAL_ERROR, + "COPY FROM STDIN requires a typed connection".to_owned(), + )); + }; + let txn_ctx = DmlTxnCtx { + sessions: &self.sessions, + session_id, + }; + let ctx = ArrayLoadCtx { + state: &self.state, + identity, + database_id, + txn_ctx: &txn_ctx, + }; + let loader = ArrayLoader::open(&ctx, array, LoadInput::Text(format)) + .map_err(|e| refused(&self.sessions, session_id, e))?; + self.array_copy.put( + connection_id.get(), + ArrayCopyPending { + identity: identity.clone(), + database_id, + loader, + }, + ); + // Text format, one column per dim and attr. The client starts + // sending rows. + Ok(Response::CopyIn(CopyResponse::new( + 0, + columns, + stream::empty(), + ))) + } +} + +/// What an array COPY opens with. +struct ArrayCopyOpen<'a> { + database_id: DatabaseId, + array: &'a str, + /// One per dim and attr. + columns: usize, + format: TextFormat, +} + +/// The connection's session store, session and array COPY state. +pub(super) struct ArrayCopyConn<'a> { + pub state: &'a SharedState, + pub sessions: &'a SessionStore, + pub session_id: SessionId, + pub copies: &'a ArrayCopyState, + pub conn_id: u64, +} + +/// Load one `CopyData` message of the connection's open array COPY. +/// `Ok(false)` when no array COPY is open. A refusal drops the load. +pub(super) async fn feed_array_copy(conn: &ArrayCopyConn<'_>, data: &[u8]) -> PgWireResult { + let Some(mut pending) = conn.copies.take(conn.conn_id) else { + return Ok(false); + }; + let txn_ctx = DmlTxnCtx { + sessions: conn.sessions, + session_id: conn.session_id, + }; + let ctx = ArrayLoadCtx { + state: conn.state, + identity: &pending.identity, + database_id: pending.database_id, + txn_ctx: &txn_ctx, + }; + let _request = conn.state.tenant_request_guard(pending.identity.tenant_id); + match pending.loader.feed_text(&ctx, data).await { + Ok(()) => { + conn.copies.put(conn.conn_id, pending); + Ok(true) + } + Err(e) => Err(refused(conn.sessions, conn.session_id, e)), + } +} + +/// Finish an array COPY: load its last record and pending chunks. Returns +/// the rows loaded. +pub(super) async fn finish_array_copy( + state: &SharedState, + sessions: &SessionStore, + session_id: SessionId, + pending: ArrayCopyPending, +) -> PgWireResult { + let ArrayCopyPending { + identity, + database_id, + loader, + } = pending; + let txn_ctx = DmlTxnCtx { + sessions, + session_id, + }; + let ctx = ArrayLoadCtx { + state, + identity: &identity, + database_id, + txn_ctx: &txn_ctx, + }; + let _request = state.tenant_request_guard(identity.tenant_id); + loader + .finish(&ctx) + .await + .map_err(|e| refused(sessions, session_id, e)) +} + +/// The pgwire error of a refused load. A refusal inside a transaction block +/// aborts the block, as a failed statement does, so the staged chunks roll +/// back. +fn refused(sessions: &SessionStore, session_id: SessionId, e: DdlError) -> PgWireError { + if sessions.transaction_state(session_id) == TransactionState::InBlock { + sessions.fail_transaction(session_id); + } + user_error(&e.sqlstate, e.message) +} + +fn user_error(code: &str, message: String) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".to_owned(), + code.to_owned(), + message, + ))) +} diff --git a/nodedb/src/control/server/pgwire/handler/copy_handler.rs b/nodedb/src/control/server/pgwire/handler/copy_handler.rs index 2d75ca08f..f9588cb8c 100644 --- a/nodedb/src/control/server/pgwire/handler/copy_handler.rs +++ b/nodedb/src/control/server/pgwire/handler/copy_handler.rs @@ -1,13 +1,15 @@ // SPDX-License-Identifier: BUSL-1.1 -//! pgwire COPY-protocol bridge for backup/restore wire surface. +//! pgwire COPY-protocol bridge for the backup/restore wire surface and for +//! `COPY FROM STDIN`, whose load lives in `array_copy`. //! //! - `intent_to_response()` runs in `SimpleQueryHandler::do_query` //! when `backup::detect()` recognises a wire COPY shape. It either //! returns a `Response::CopyOut` whose stream pulls bytes from //! `backup_tenant`, or registers per-connection state and returns //! `Response::CopyIn` so the client begins streaming bytes. -//! - `CopyHandler::on_copy_data` accumulates client bytes (size-capped). +//! - `CopyHandler::on_copy_data` streams an array COPY's bytes into its +//! load, or accumulates a restore's bytes (size-capped). //! - `CopyHandler::on_copy_done` validates the envelope and dispatches //! `restore::restore_tenant`. @@ -35,10 +37,11 @@ use crate::control::security::request_scope::RequestAuthScope; use crate::control::server::shared::backup_metering::{ admit_backup_restore_quota, meter_backup_restore, }; -use crate::control::server::shared::session::{ConnectionId, SessionId}; +use crate::control::server::shared::session::{ConnectionId, SessionId, SessionStore}; use crate::control::state::SharedState; use crate::types::TenantId; +use super::array_copy::{ArrayCopyConn, ArrayCopyState, feed_array_copy, finish_array_copy}; use super::core::NodeDbPgHandler; /// Hard cap on accumulated COPY IN bytes for one restore. 16 GiB matches @@ -182,6 +185,10 @@ impl NodeDbPgHandler { pub struct NodeDbCopyHandler { pub state: Arc, pub restore_state: Arc, + /// Open `COPY FROM STDIN` loads, keyed by connection. + pub array_copy: Arc, + /// The session store, for the transaction an array load runs in. + pub sessions: Arc, pub connection_id: ConnectionId, } @@ -194,6 +201,16 @@ impl CopyHandler for NodeDbCopyHandler { PgWireError: From<>::Error>, { let id = self.connection_id.get(); + let array_conn = ArrayCopyConn { + state: &self.state, + sessions: &self.sessions, + session_id: SessionId::Connection(self.connection_id), + copies: &self.array_copy, + conn_id: id, + }; + if feed_array_copy(&array_conn, ©_data.data).await? { + return Ok(()); + } match self.restore_state.append(id, ©_data.data) { Ok(()) => Ok(()), Err(e @ AppendError::NotPending) => { @@ -212,6 +229,9 @@ impl CopyHandler for NodeDbCopyHandler { C::Error: Debug, PgWireError: From<>::Error>, { + if self.array_copy.take(self.connection_id.get()).is_some() { + return sqlstate(ss::QUERY_CANCELED.0, "COPY aborted by the client"); + } cancel_restore(&self.restore_state, self.connection_id); sqlstate(ss::QUERY_CANCELED.0, "COPY restore aborted") } @@ -223,6 +243,11 @@ impl CopyHandler for NodeDbCopyHandler { PgWireError: From<>::Error>, { let id = self.connection_id.get(); + if let Some(pending) = self.array_copy.take(id) { + let session_id = SessionId::Connection(self.connection_id); + let rows = finish_array_copy(&self.state, &self.sessions, session_id, pending).await?; + return complete_copy_in(client, Tag::new("COPY").with_rows(rows as usize)).await; + } let pending = self.restore_state.take(id).ok_or_else(|| { sqlstate( ss::FEATURE_NOT_SUPPORTED, @@ -265,24 +290,33 @@ impl CopyHandler for NodeDbCopyHandler { // back in `intent_to_response`, carried through `RestorePending` // since this per-connection `CopyHandler` has no identity of its own. meter_backup_restore(&self.state, &scope, pending.tenant_id, Some(rows as u64)); - let tag = Tag::new("RESTORE TENANT").with_rows(rows); - client - .send(PgWireBackendMessage::CommandComplete(tag.into())) - .await - .map_err(|e| { - sqlstate( - ss::INTERNAL_ERROR, - &format!("CommandComplete send failed: {e:?}"), - ) - })?; - // Leave the COPY-in-progress state so the next Sync from the - // client gets dispatched normally. pgwire's `process_message` - // only routes Sync via the `AwaitingSync` arm. - client.set_state(PgWireConnectionState::AwaitingSync); - Ok(()) + complete_copy_in(client, Tag::new("RESTORE TENANT").with_rows(rows)).await } } +/// Send the COPY's CommandComplete `tag` and leave the COPY-in state. +async fn complete_copy_in(client: &mut C, tag: Tag) -> PgWireResult<()> +where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, +{ + client + .send(PgWireBackendMessage::CommandComplete(tag.into())) + .await + .map_err(|e| { + sqlstate( + ss::INTERNAL_ERROR, + &format!("CommandComplete send failed: {e:?}"), + ) + })?; + // Leave the COPY-in-progress state so the next Sync from the + // client gets dispatched normally. pgwire's `process_message` + // only routes Sync via the `AwaitingSync` arm. + client.set_state(PgWireConnectionState::AwaitingSync); + Ok(()) +} + fn cancel_restore(state: &backup::RestoreState, connection_id: ConnectionId) { state.cancel(connection_id.get()); } diff --git a/nodedb/src/control/server/pgwire/handler/core.rs b/nodedb/src/control/server/pgwire/handler/core.rs index 3e2f962a2..1fe4b6307 100644 --- a/nodedb/src/control/server/pgwire/handler/core.rs +++ b/nodedb/src/control/server/pgwire/handler/core.rs @@ -51,6 +51,8 @@ pub struct NodeDbPgHandler { pub(crate) sessions: Arc, /// Per-connection in-flight COPY IN restore accumulators. pub(crate) restore_state: Arc, + /// Per-connection open `COPY FROM STDIN` loads. + pub(crate) array_copy: Arc, pub(crate) registry: Arc, pub(crate) session_id: SessionId, } @@ -78,6 +80,7 @@ impl NodeDbPgHandler { auth_mode, sessions, restore_state: Arc::new(crate::control::backup::RestoreState::new()), + array_copy: Arc::new(super::array_copy::ArrayCopyState::new()), registry: Arc::new(ConnectionRegistry::new()), session_id: SessionId::LegacySocket(([0, 0, 0, 0], 0).into()), } @@ -93,6 +96,7 @@ impl NodeDbPgHandler { auth_mode: AuthMode, sessions: Arc, restore_state: Arc, + array_copy: Arc, registry: Arc, connection_id: ConnectionId, ) -> Self { @@ -109,6 +113,7 @@ impl NodeDbPgHandler { auth_mode, sessions, restore_state, + array_copy, registry, session_id, } diff --git a/nodedb/src/control/server/pgwire/handler/mod.rs b/nodedb/src/control/server/pgwire/handler/mod.rs index 6f4c275fc..59e49378f 100644 --- a/nodedb/src/control/server/pgwire/handler/mod.rs +++ b/nodedb/src/control/server/pgwire/handler/mod.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: BUSL-1.1 +pub mod array_copy; mod auth; mod connection_admin; mod copy_handler; @@ -30,5 +31,6 @@ mod transaction_cmds; mod transaction_savepoint; mod trust_auth; +pub use self::array_copy::ArrayCopyState; pub use self::copy_handler::NodeDbCopyHandler; pub use self::core::NodeDbPgHandler; diff --git a/nodedb/src/control/server/pgwire/handler/prepared/mod.rs b/nodedb/src/control/server/pgwire/handler/prepared/mod.rs index be1f106ae..b8c01f705 100644 --- a/nodedb/src/control/server/pgwire/handler/prepared/mod.rs +++ b/nodedb/src/control/server/pgwire/handler/prepared/mod.rs @@ -3,6 +3,7 @@ pub mod describe; pub mod execute; pub mod param_bind; +pub mod parse_refusal; pub mod parser; pub mod statement; diff --git a/nodedb/src/control/server/pgwire/handler/prepared/parse_refusal.rs b/nodedb/src/control/server/pgwire/handler/prepared/parse_refusal.rs new file mode 100644 index 000000000..1b8101d6e --- /dev/null +++ b/nodedb/src/control/server/pgwire/handler/prepared/parse_refusal.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Which planning errors a Parse message reports. +//! +//! PostgreSQL runs parse analysis at Parse. A statement that names an +//! unknown relation, column or function fails there with its SQLSTATE, so a +//! driver's prepare fails. Planning and execution errors, such as a division +//! by zero folded from constants, surface at Bind or Execute. +//! +//! Parse plans the statement with every `$N` replaced by the text `NULL`. +//! Execute binds the client's values into the parsed statement instead. So +//! with placeholders, Parse plans a different statement than Execute does. +//! It reports only the errors no value can change: an unknown table, column +//! or function, an ambiguous column, and a column outside GROUP BY. Every +//! other error can come from the substituted `NULL`. Examples are a `LIMIT` +//! or a function argument that must be a literal value, or a type check +//! against the placeholder. Those errors defer to Execute, which plans with +//! the bound values. + +use nodedb_types::error::sqlstate; + +/// SQLSTATE class `42`: syntax error or access rule violation. +const SYNTAX_OR_ACCESS_CLASS: &str = "42"; + +/// Whether a Parse reports the planning error with SQLSTATE `state`. +/// +/// `has_placeholders` says whether the statement text holds a `$N`. +/// +/// A statement without placeholders plans at Parse exactly as Execute plans +/// it, so a class-`42` error and `0A000` (feature not supported) report at +/// Parse. Three kinds defer to Execute, as PostgreSQL does: +/// - `42501` (insufficient privilege): Execute checks privileges under the +/// session's adopted authorization, which Parse does not hold. +/// - A data error of class `22`: it comes from evaluating a value. PostgreSQL +/// folds `1 / 0` while planning, after Parse, and raises `22012` there. +/// - A retryable error of class `40`: Execute retries the statement. +pub(super) fn reported_at_parse(state: &str, has_placeholders: bool) -> bool { + if has_placeholders { + return matches!( + state, + sqlstate::UNDEFINED_TABLE + | sqlstate::UNDEFINED_COLUMN + | sqlstate::AMBIGUOUS_COLUMN + | sqlstate::UNDEFINED_FUNCTION + | sqlstate::GROUPING_ERROR + ); + } + if state == sqlstate::INSUFFICIENT_PRIVILEGE { + return false; + } + state.starts_with(SYNTAX_OR_ACCESS_CLASS) || state == sqlstate::FEATURE_NOT_SUPPORTED +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn name_resolution_reports_with_or_without_placeholders() { + for state in [ + sqlstate::UNDEFINED_TABLE, + sqlstate::UNDEFINED_COLUMN, + sqlstate::AMBIGUOUS_COLUMN, + sqlstate::UNDEFINED_FUNCTION, + sqlstate::GROUPING_ERROR, + ] { + assert!(reported_at_parse(state, false), "{state}"); + assert!(reported_at_parse(state, true), "{state}"); + } + } + + #[test] + fn a_value_dependent_error_defers_with_placeholders() { + for state in [ + sqlstate::SYNTAX_ERROR, + sqlstate::DATATYPE_MISMATCH, + sqlstate::FEATURE_NOT_SUPPORTED, + sqlstate::INVALID_LIMIT_VALUE, + ] { + assert!(!reported_at_parse(state, true), "{state}"); + } + assert!(reported_at_parse(sqlstate::SYNTAX_ERROR, false)); + assert!(reported_at_parse(sqlstate::FEATURE_NOT_SUPPORTED, false)); + } + + #[test] + fn privilege_data_and_retryable_errors_defer() { + for state in [ + sqlstate::INSUFFICIENT_PRIVILEGE, + sqlstate::DIVISION_BY_ZERO, + sqlstate::NUMERIC_VALUE_OUT_OF_RANGE, + sqlstate::SERIALIZATION_FAILURE, + ] { + assert!(!reported_at_parse(state, false), "{state}"); + } + } +} diff --git a/nodedb/src/control/server/pgwire/handler/prepared/parser.rs b/nodedb/src/control/server/pgwire/handler/prepared/parser.rs index d6e98432c..7958ee328 100644 --- a/nodedb/src/control/server/pgwire/handler/prepared/parser.rs +++ b/nodedb/src/control/server/pgwire/handler/prepared/parser.rs @@ -16,6 +16,8 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use crate::config::auth::AuthMode; use crate::control::security::audit::ArcAuditEmitter; +use crate::control::server::pgwire::types::error_map::error_to_pg; +use crate::control::server::pgwire::types::error_to_sqlstate; use crate::control::server::pgwire::types::wire_type::{TEXT_RESULTS, pg_type}; use crate::control::server::shared::authorization::{authorize_database, authorize_task_set}; use crate::control::server::shared::returning; @@ -161,16 +163,24 @@ impl NodeDbQueryParser { } } + /// Plan `sql` as Execute plans it and authorize the planned tasks. + /// + /// The outer error is a refusal Parse always reports: authorization or + /// the permission view. The inner result is the planning outcome. The + /// caller decides whether a planning error reports at Parse. async fn authorize_plannable_sql( &self, sql: &str, identity: &crate::control::security::identity::AuthenticatedIdentity, database_id: crate::types::DatabaseId, emitter: &ArcAuditEmitter, - ) -> PgWireResult { - let (sql_without_returning, _) = match returning::strip_returning(sql) { + ) -> PgWireResult> { + // Execute strips a per-query `ON DENY` clause and then `RETURNING` + // before it plans, so Parse does too. + let (sql_without_on_deny, _) = crate::control::server::session_auth::extract_on_deny(sql); + let (sql_without_returning, _) = match returning::strip_returning(&sql_without_on_deny) { Ok(parts) => parts, - Err(_) => return Ok(false), + Err(error) => return Ok(Err(error)), }; let sql_for_planning = substitute_placeholders_with_null(&sql_without_returning); let query_ctx = @@ -203,7 +213,7 @@ impl NodeDbQueryParser { &self.state.permission_cache, ), }; - let Ok((tasks, _)) = query_ctx + let tasks = match query_ctx .plan_sql_with_rls_metadata(crate::control::planner::context::PlanSqlWithRlsParams { sql: &sql_for_planning, tenant_id: identity.tenant_id, @@ -211,8 +221,9 @@ impl NodeDbQueryParser { sec: &security, }) .await - else { - return Ok(false); + { + Ok((tasks, _)) => tasks, + Err(error) => return Ok(Err(error)), }; // Parse/Describe is metadata-only: authorize the original task set, @@ -225,7 +236,7 @@ impl NodeDbQueryParser { emitter, ) .map_err(pgwire_authorization_error)?; - Ok(true) + Ok(Ok(())) } /// The tenant-scoped catalog a single Parse message resolves against, so a @@ -265,9 +276,12 @@ impl NodeDbQueryParser { // pass below and survives that pass failing. let param_types = Self::param_types_with_inference(sql, client_types, catalog); - // Strip RETURNING from DML before planning. The item text is kept so - // the result fields for Describe are built from the resolved clause. - let (sql_stripped, returning_items) = match returning::strip_returning(sql) { + // Strip a per-query `ON DENY` clause, then RETURNING, before planning, + // as Execute does. The RETURNING item text is kept so the result + // fields for Describe are built from the resolved clause. + let (sql_without_on_deny, _) = crate::control::server::session_auth::extract_on_deny(sql); + let (sql_stripped, returning_items) = match returning::strip_returning(&sql_without_on_deny) + { Ok(pair) => pair, Err(_) => return (param_types, Vec::new()), }; @@ -398,9 +412,24 @@ impl QueryParser for NodeDbQueryParser { })); } - let can_infer_schema = self + // A DSL or NodeDB DDL statement has no plan at Parse: Execute routes + // it to its own dispatcher. Any other statement plans at Execute as it + // plans here, so a parse-analysis error reports now, as in PostgreSQL. + let can_infer_schema = match self .authorize_plannable_sql(sql, &identity, database_id, &emitter) - .await?; + .await? + { + Ok(()) => true, + Err(error) => { + let state = error_to_sqlstate(&error).1; + if !is_dsl_statement(sql) + && super::parse_refusal::reported_at_parse(state, count_placeholders(sql) > 0) + { + return Err(error_to_pg(&error)); + } + false + } + }; // One catalog per Parse, shared by parameter-type inference and schema // planning. The unplannable branch still needs it: inference resolves // `WHERE col = $1` from the catalog whether or not the statement as a diff --git a/nodedb/src/control/server/pgwire/handler/routing/planning.rs b/nodedb/src/control/server/pgwire/handler/routing/planning.rs index 3651b71fc..b13ce0620 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/planning.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/planning.rs @@ -217,6 +217,14 @@ impl NodeDbPgHandler { .map_err(StatementSetupError::from)? .tenant_version(tenant_id.as_u64()); let current_rls_version = self.state.rls.tenant_version(tenant_id.as_u64()); + // A cached plan's column reads were checked against redaction policy + // before row-level security was injected. A policy change evicts the + // entry, so the next plan is checked against the current policy. + let policies = crate::control::planner::descriptor_set::PolicyVersions { + permission_tree: current_permission_tree_version, + rls: current_rls_version, + redaction: self.state.redaction.version(), + }; let cached_tasks = if bypass_cache { None @@ -232,8 +240,7 @@ impl NodeDbPgHandler { session_id, &statement_sql, move |id| current_descriptor_version(&state, tenant, db, id), - current_permission_tree_version, - current_rls_version, + policies, ) }; diff --git a/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs b/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs index 3a2e72dbb..b841f95e5 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/set_ops.rs @@ -17,7 +17,7 @@ use crate::control::server::set_op_merge::{ SetMergeMode, dedup_union_payloads, merge_set_op_payloads, }; -use super::super::super::types::shape_error_to_pg; +use super::super::super::types::{error_to_pg, shape_error_to_pg}; use super::super::plan::{PlanKind, multirow_payload_to_response}; use super::super::shape_encode; @@ -42,7 +42,8 @@ pub(super) fn apply_set_ops( merge_set_op_payloads(dedup_payloads, SetMergeMode::Except) } _ => dedup_union_payloads(dedup_payloads), - }; + } + .map_err(|e| error_to_pg(&e))?; Ok( match compose::shape_payload_no_plan( &merged, diff --git a/nodedb/src/control/server/pgwire/handler/sql_exec.rs b/nodedb/src/control/server/pgwire/handler/sql_exec.rs index 09fd882c0..1740c25a7 100644 --- a/nodedb/src/control/server/pgwire/handler/sql_exec.rs +++ b/nodedb/src/control/server/pgwire/handler/sql_exec.rs @@ -117,12 +117,12 @@ impl NodeDbPgHandler { return result; } - // ── Wire-streaming COPY shapes for backup/restore ───────────── - if let Some(intent) = crate::control::backup::detect(sql_trimmed) { - return self - .intent_to_response(identity, session_id, intent) - .await - .map(|r| vec![r]); + // ── Wire-streaming COPY shapes: backup, restore, array load ─── + if let Some(r) = self + .copy_stream_response(identity, session_id, sql_trimmed) + .await + { + return r.map(|r| vec![r]); } // ── Cursor commands ─────────────────────────────────────────── diff --git a/nodedb/src/control/server/pgwire/types/error_map.rs b/nodedb/src/control/server/pgwire/types/error_map.rs index a80267c07..efee02dfd 100644 --- a/nodedb/src/control/server/pgwire/types/error_map.rs +++ b/nodedb/src/control/server/pgwire/types/error_map.rs @@ -118,6 +118,9 @@ pub fn error_to_sqlstate(err: &crate::Error) -> (&'static str, &'static str, Str sqlstate::AMBIGUOUS_COLUMN, format!("column reference \"{column}\" is ambiguous"), ), + crate::Error::UngroupedColumn { .. } => { + ("ERROR", sqlstate::GROUPING_ERROR, err.to_string()) + } crate::Error::UnknownStrictField { .. } => { ("ERROR", sqlstate::UNDEFINED_COLUMN, err.to_string()) } @@ -125,6 +128,7 @@ pub fn error_to_sqlstate(err: &crate::Error) -> (&'static str, &'static str, Str crate::Error::DataException { detail } => { ("ERROR", sqlstate::DATA_EXCEPTION, detail.clone()) } + crate::Error::MixedSrid { .. } => ("ERROR", sqlstate::DATA_EXCEPTION, err.to_string()), crate::Error::NumericValueOutOfRange { detail } => ( "ERROR", sqlstate::NUMERIC_VALUE_OUT_OF_RANGE, diff --git a/nodedb/src/control/server/resp/gateway_dispatch.rs b/nodedb/src/control/server/resp/gateway_dispatch.rs index ce39d1809..1b66cbedc 100644 --- a/nodedb/src/control/server/resp/gateway_dispatch.rs +++ b/nodedb/src/control/server/resp/gateway_dispatch.rs @@ -226,6 +226,19 @@ fn authorize_resp_task( // // Operations that cannot carry a filter (`BatchGet`, `FieldGet`) fail // closed here with a typed error rather than executing unfiltered. + // + // A read of a redacted column outside the projection is refused first, + // before RLS adds the read policy's own predicates. + crate::control::planner::redaction_refusal::refuse_redacted_expression_reads( + &plan, + session.tenant_id, + database_id, + &crate::control::planner::redaction_refusal::RedactionReadCheck { + auth: scope.auth(), + store: &state.redaction, + catalog: state.credentials.catalog(), + }, + )?; crate::control::planner::rls_injection::inject_rls_for_single_plan( session.tenant_id.as_u64(), database_id, diff --git a/nodedb/src/control/server/response_shape/compose/array_slice.rs b/nodedb/src/control/server/response_shape/compose/array_slice.rs index c61bc6482..d31ead499 100644 --- a/nodedb/src/control/server/response_shape/compose/array_slice.rs +++ b/nodedb/src/control/server/response_shape/compose/array_slice.rs @@ -2,9 +2,11 @@ //! Shaping for an `ArrayOp::Slice` response. +use crate::control::sequence::SequenceAccess; use crate::data::executor::response_codec::{ArraySliceResponse, decode_payload_value}; use super::super::redaction::RedactionCtx; +use super::super::schema::OutputSchema; use super::super::types::ShapedRows; use super::kernel::{empty_shaped, shape_decoded_rows}; @@ -15,17 +17,19 @@ pub(crate) const TRUNCATED_BEFORE_HORIZON_NOTICE: &str = "AS OF SYSTEM TIME cuto results may be incomplete"; /// Shape an `ArrayOp::Slice` response: decode the `ArraySliceResponse` -/// envelope (falling back to a plain payload decode for legacy shapes), -/// unwrap the row envelope, and surface `truncated_before_horizon` as a -/// notice. +/// envelope (a cluster slice answers its rows bare), project the rows to +/// `projection`, and surface `truncated_before_horizon` as a notice. /// -/// Array slices never carry a SELECT-list projection, so `shape_decoded_rows` -/// is always called with a `None` projection and no sequence access here — -/// but redaction still applies to the cells. A payload that decodes to no -/// value shapes as an empty result set. +/// Each row holds one field per dim and per attr, keyed by name. The +/// projection the planner announces for the slice's relation names those +/// fields in schema order with their types, so the client receives typed +/// columns in that order. Redaction applies to the cells. A payload that +/// decodes to no value shapes as an empty result set. pub(super) fn shape_array_slice( payload: &[u8], + projection: Option<&OutputSchema>, redaction: Option>, + sequences: Option<&dyn SequenceAccess>, ) -> crate::Result { if payload.is_empty() { return Ok(empty_shaped()); @@ -41,7 +45,7 @@ pub(super) fn shape_array_slice( let notice = truncated.then(|| TRUNCATED_BEFORE_HORIZON_NOTICE.to_string()); let mut shaped = match rows { - Ok(value) => shape_decoded_rows(value, None, redaction, None)?, + Ok(value) => shape_decoded_rows(value, projection, redaction, sequences)?, Err(_) => empty_shaped(), }; shaped.notice = notice; diff --git a/nodedb/src/control/server/response_shape/compose/materialized.rs b/nodedb/src/control/server/response_shape/compose/materialized.rs index 959db2b15..5ed6840f6 100644 --- a/nodedb/src/control/server/response_shape/compose/materialized.rs +++ b/nodedb/src/control/server/response_shape/compose/materialized.rs @@ -84,7 +84,7 @@ pub fn shape_response_materialized( let translated = translate_search_response(&wrapped, plan, state, database_id, tenant_id); let shaped = match plan_kind { - PlanKind::ArraySlice => shape_array_slice(&translated, redaction)?, + PlanKind::ArraySlice => shape_array_slice(&translated, projection, redaction, sequences)?, // `RETURNING` rows are held to the columns already announced to the // client, when any were — see `super::returning`. PlanKind::ReturningRows => { @@ -127,7 +127,9 @@ pub fn shape_payload_no_plan( PlanKind::Execution | PlanKind::DmlResult(_) | PlanKind::DmlResultByOp => { ShapeOutcome::Passthrough } - PlanKind::ArraySlice => ShapeOutcome::Rows(shape_array_slice(payload, redaction)?), + PlanKind::ArraySlice => ShapeOutcome::Rows(shape_array_slice( + payload, projection, redaction, sequences, + )?), PlanKind::ReturningRows => ShapeOutcome::Rows(shape_returning_rows( payload, projection, redaction, sequences, )?), diff --git a/nodedb/src/control/server/response_shape/redaction/query.rs b/nodedb/src/control/server/response_shape/redaction/query.rs index 934663d9b..abcf4a9e7 100644 --- a/nodedb/src/control/server/response_shape/redaction/query.rs +++ b/nodedb/src/control/server/response_shape/redaction/query.rs @@ -311,9 +311,14 @@ fn collect_sources( } return; } - // Every remaining relational op reads at most the single - // collection `extract_collection` reports, resolved below. - _ => {} + // Each reads at most the single collection `extract_collection` + // reports, resolved below. + QueryOp::PartialAggregate { .. } + | QueryOp::ShuffleJoinConsume { .. } + | QueryOp::ShuffleAggregateConsume { .. } + | QueryOp::FacetCounts { .. } + | QueryOp::RecursiveScan { .. } + | QueryOp::RecursiveValue { .. } => {} } } @@ -559,6 +564,39 @@ mod tests { assert!(redaction.field_has_rule(&store, "ssn")); } + /// Every similarity search is a source of its collection, so a rule on + /// that collection masks its rows and refuses its filters. + #[test] + fn a_similarity_search_is_a_source_of_its_collection() { + use nodedb_physical::physical_plan::VectorOp; + + let docs = || QualifiedCollection::new(DatabaseId::DEFAULT, "docs"); + let plans = [ + PhysicalPlan::Vector(VectorOp::SparseSearch { + collection: docs(), + field_name: "terms".into(), + query_entries: vec![(1, 1.0)], + top_k: 1, + filters: Vec::new(), + }), + PhysicalPlan::Vector(VectorOp::MultiVectorScoreSearch { + collection: docs(), + field_name: "tokens".into(), + query_vector: vec![0.0], + top_k: 1, + ef_search: 2, + mode: "max_sim".into(), + }), + ]; + for plan in plans { + assert_eq!( + plan_source_collections(&plan, DatabaseId::DEFAULT), + vec![(String::new(), docs().as_str().to_string())], + "{plan:?}" + ); + } + } + /// A masked column of a provider relation is rewritten in a SELECT /// result, as a collection's is. #[test] diff --git a/nodedb/src/control/server/response_shape/types/plan_kind/array.rs b/nodedb/src/control/server/response_shape/types/plan_kind/array.rs index a8b39a0cc..4063492cc 100644 --- a/nodedb/src/control/server/response_shape/types/plan_kind/array.rs +++ b/nodedb/src/control/server/response_shape/types/plan_kind/array.rs @@ -11,9 +11,10 @@ use super::kind::PlanKind; pub(super) fn describe_cluster_array(op: &ClusterArrayOp) -> PlanKind { match op { ClusterArrayOp::Slice { .. } => PlanKind::ArraySlice, - ClusterArrayOp::Agg { .. } => PlanKind::MultiRow, + ClusterArrayOp::Agg { .. } | ClusterArrayOp::Elementwise { .. } => PlanKind::MultiRow, ClusterArrayOp::Put { .. } => PlanKind::DmlResult("INSERT"), ClusterArrayOp::Delete { .. } => PlanKind::DmlResult("DELETE"), + ClusterArrayOp::Flush { .. } | ClusterArrayOp::Compact { .. } => PlanKind::SingleDocument, } } @@ -22,7 +23,7 @@ pub(super) fn describe_array(op: &ArrayOp) -> PlanKind { ArrayOp::Slice { .. } => PlanKind::ArraySlice, // JSON-array payloads: each row streams as its own pgwire field. - ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } => { + ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. } => { PlanKind::MultiRow } diff --git a/nodedb/src/control/server/response_translate/dispatch.rs b/nodedb/src/control/server/response_translate/dispatch.rs index c8695ec7b..d4e666e05 100644 --- a/nodedb/src/control/server/response_translate/dispatch.rs +++ b/nodedb/src/control/server/response_translate/dispatch.rs @@ -11,7 +11,7 @@ use crate::control::state::SharedState; use nodedb_physical::physical_plan::{TextOp, VectorOp}; use super::text_hybrid::{translate_hybrid_search_payload, translate_text_search_payload}; -use super::vector::translate_vector_search_payload; +use super::vector::{VectorHitCut, translate_vector_search_payload}; /// Inspect the executed plan and apply the surrogate→PK translation that /// matches its shape: vector search hits, full-text search hits, or hybrid @@ -31,13 +31,28 @@ pub fn translate_search_response( } match plan { + // A slow-path search attached each body for its column projection, + // so the body stays after the filter. The fast path attached it only + // for the predicate. PhysicalPlan::Vector(VectorOp::Search { collection, rls_filters, top_k, + skip_payload_fetch, .. - }) - | PhysicalPlan::Vector(VectorOp::MultiSearch { + }) => translate_vector_search_payload( + payload, + state, + database_id, + tenant_id, + collection.as_str(), + VectorHitCut { + rls_filters: rls_filters.as_slice(), + top_k: *top_k, + strip_bodies: *skip_payload_fetch, + }, + ), + PhysicalPlan::Vector(VectorOp::MultiSearch { collection, rls_filters, top_k, @@ -48,22 +63,44 @@ pub fn translate_search_response( database_id, tenant_id, collection.as_str(), - rls_filters.as_slice(), - *top_k, + VectorHitCut { + rls_filters: rls_filters.as_slice(), + top_k: *top_k, + strip_bodies: true, + }, + ), + // A sparse search attaches each body for its column projection, so + // the body stays after the filter. + PhysicalPlan::Vector(VectorOp::SparseSearch { + collection, + filters, + top_k, + .. + }) => translate_vector_search_payload( + payload, + state, + database_id, + tenant_id, + collection.as_str(), + VectorHitCut { + rls_filters: filters.as_slice(), + top_k: *top_k, + strip_bodies: false, + }, ), PhysicalPlan::Vector(VectorOp::MultiVectorScoreSearch { collection, top_k, .. - }) - | PhysicalPlan::Vector(VectorOp::SparseSearch { - collection, top_k, .. }) => translate_vector_search_payload( payload, state, database_id, tenant_id, collection.as_str(), - &[], - *top_k, + VectorHitCut { + rls_filters: &[], + top_k: *top_k, + strip_bodies: false, + }, ), PhysicalPlan::Text( TextOp::Search { collection, .. } diff --git a/nodedb/src/control/server/response_translate/vector.rs b/nodedb/src/control/server/response_translate/vector.rs index f12731050..02d73e944 100644 --- a/nodedb/src/control/server/response_translate/vector.rs +++ b/nodedb/src/control/server/response_translate/vector.rs @@ -42,16 +42,41 @@ struct Hit { body: Option>, } -/// Apply RLS post-filter at the Control-Plane security boundary, then -/// truncate to `top_k` and strip the body bytes that the Data Plane -/// attached purely for the predicate evaluation. A no-op when -/// `rls_filters` is empty. -fn apply_rls_filter(hits: &mut Vec, rls_filters: &[u8], top_k: usize) { - // When filters are empty the Data Plane never attaches a body, so there - // is nothing to evaluate or strip — return immediately. +/// The residual filter and row bound a vector search's hits are cut to. +pub struct VectorHitCut<'a> { + /// Serialized `Vec`: the statement's `WHERE` conjuncts and + /// any read policy. Empty when none applies. + pub rls_filters: &'a [u8], + pub top_k: usize, + /// Drop each body after the filter. `true` when the Data Plane attached + /// the body only for the predicate, as the pure-vector fast path does. + pub strip_bodies: bool, +} + +/// Apply the residual filter, then truncate to `top_k`. With +/// `strip_bodies`, drop the body bytes the Data Plane attached only for the +/// predicate. A no-op when `rls_filters` is empty. +fn apply_rls_filter(hits: &mut Vec, rls_filters: &[u8], top_k: usize, strip_bodies: bool) { + // When filters are empty the Data Plane ranks exactly `top_k` and + // attaches a body only for projection, so there is nothing to do. if rls_filters.is_empty() { return; } + admit_hits(hits, rls_filters, top_k); + if strip_bodies { + for h in hits.iter_mut() { + h.body = None; + } + } +} + +/// Keep the hits whose body passes `rls_filters`, nearest first, cut to +/// `top_k`. +/// +/// The Data Plane ranks a wider window when a residual filter applies: the +/// statement's `WHERE` conjuncts and any read policy. The filter and the cut +/// run here, on the Control Plane. +fn admit_hits(hits: &mut Vec, rls_filters: &[u8], top_k: usize) { let filters: Vec = match zerompk::from_msgpack(rls_filters) { Ok(f) => f, Err(_) => { @@ -76,9 +101,33 @@ fn apply_rls_filter(hits: &mut Vec, rls_filters: &[u8], top_k: usize) { if hits.len() > top_k { hits.truncate(top_k); } - for h in hits.iter_mut() { - h.body = None; +} + +/// Apply a vector search's residual filter and `top_k` cut to its raw hit +/// payload, keeping each body for the rows built from it. +/// +/// A search read as the body of a post-processing tail reaches the gather +/// path, not the response translator. The gather path calls this so the +/// tail sees the rows a root search returns. A payload that does not decode +/// fails the statement while a filter applies: no hit can be vouched for. +pub(crate) fn filter_vector_hits( + payload: &[u8], + rls_filters: &[u8], + top_k: usize, +) -> crate::Result> { + if rls_filters.is_empty() { + return Ok(payload.to_vec()); } + let mut hits: Vec = + zerompk::from_msgpack(payload).map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("vector search hits do not decode: {e}"), + })?; + admit_hits(&mut hits, rls_filters, top_k); + zerompk::to_msgpack_vec(&hits).map_err(|e| crate::Error::Serialization { + format: "msgpack".into(), + detail: format!("vector search hits do not encode: {e}"), + }) } /// Resolve one surrogate to its user-facing primary key via the catalog. @@ -115,8 +164,7 @@ pub fn translate_vector_search_payload( database_id: DatabaseId, tenant_id: TenantId, collection: &str, - rls_filters: &[u8], - top_k: usize, + cut: VectorHitCut<'_>, ) -> Vec { if payload.is_empty() { return payload.to_vec(); @@ -131,7 +179,7 @@ pub fn translate_vector_search_payload( Err(_) => return payload.to_vec(), }; - apply_rls_filter(&mut hits, rls_filters, top_k); + apply_rls_filter(&mut hits, cut.rls_filters, cut.top_k, cut.strip_bodies); for hit in &mut hits { if hit.doc_id.is_some() { diff --git a/nodedb/src/control/server/result_stream.rs b/nodedb/src/control/server/result_stream.rs index 95d111e65..c2c9d660f 100644 --- a/nodedb/src/control/server/result_stream.rs +++ b/nodedb/src/control/server/result_stream.rs @@ -129,7 +129,7 @@ pub(crate) async fn materialize(mut stream: ResultStream) -> crate::Result<(Vec< } payloads.push(batch.payload); } - Ok((merge_msgpack_arrays(&payloads), max_lsn)) + Ok((merge_msgpack_arrays(&payloads)?, max_lsn)) } #[cfg(test)] @@ -221,7 +221,7 @@ mod tests { ); let (merged, _lsn) = materialize(stream).await.unwrap(); assert_eq!( - extract_msgpack_elements(&merged).len(), + extract_msgpack_elements(&merged).unwrap().len(), 2500, "three array batches must materialize into one array of all rows" ); @@ -242,6 +242,28 @@ mod tests { assert!(matches!(err, crate::Error::ExecutionLimitExceeded { .. })); } + /// A batch that is not a msgpack array fails the materialize instead of + /// reading as one row. + #[tokio::test] + async fn a_non_array_batch_is_refused() { + let (tx, rx) = mpsc::channel(8); + tx.send(partial(2)).await.unwrap(); + let mut last = final_frame(0); + // fixmap of 1: "k" => 1. + last.payload = Payload::from_vec(vec![0x81, 0xa1, b'k', 1]); + tx.send(last).await.unwrap(); + drop(tx); + let stream = stream_response_channel( + crate::control::ResponseReceiver::from_channel(rx), + 1 << 20, + false, + ); + match materialize(stream).await { + Err(crate::Error::Codec { .. }) => {} + other => panic!("expected a codec error, got {other:?}"), + } + } + /// A terminal error frame ends the stream with the typed error its code /// maps to. Collapsing every code into one generic variant is what made a /// client's own timeout arrive as an internal failure. @@ -318,7 +340,7 @@ mod tests { ); let (merged, _lsn) = materialize(stream).await.unwrap(); assert_eq!( - extract_msgpack_elements(&merged).len(), + extract_msgpack_elements(&merged).unwrap().len(), 0, "tolerated NotFound yields an empty result, not an error" ); diff --git a/nodedb/src/control/server/session_auth/bearer_jwt.rs b/nodedb/src/control/server/session_auth/bearer_jwt.rs index afef1e7b9..a974a6498 100644 --- a/nodedb/src/control/server/session_auth/bearer_jwt.rs +++ b/nodedb/src/control/server/session_auth/bearer_jwt.rs @@ -4,6 +4,7 @@ use tracing::{debug, warn}; +use crate::control::security::bearer::{BearerAttempt, BearerMethod}; use crate::control::security::identity::AuthenticatedIdentity; use crate::control::security::jwks::registry::VerifiedJwtClaims; use crate::control::state::SharedState; @@ -24,19 +25,30 @@ use crate::control::state::SharedState; /// assertion must be refused, never downgraded to an anonymous or trust /// identity. The reason stays server-side so a caller cannot probe which /// providers exist or why a token was refused. +/// +/// Every refusal records one `AuthFailure` with the reason, and every login +/// records one `AuthSuccess`, through [`BearerAttempt`]. `peer_addr` is the +/// audit source. `session_id` is the id the caller's session or request scope +/// carries, or empty for a sessionless route. pub async fn authenticate_bearer_jwt( state: &SharedState, token: &str, + peer_addr: &str, + session_id: &str, ) -> Option<(AuthenticatedIdentity, VerifiedJwtClaims)> { + let mut attempt = BearerAttempt::new(BearerMethod::JwtBearer, peer_addr, session_id); let Some(registry) = state.jwks_registry.as_ref() else { warn!("bearer token presented but no [auth.jwt] provider is configured; refusing"); + attempt.refuse(state, "no [auth.jwt] provider is configured"); return None; }; + registry.label_unverified(token, &mut attempt); let (identity, verified) = match registry.validate_with_claims(token).await { Ok(verified) => verified, Err(error) => { debug!(%error, "bearer token rejected by the JWKS registry"); + attempt.refuse(state, &error.to_string()); return None; } }; @@ -47,8 +59,10 @@ pub async fn authenticate_bearer_jwt( &identity, ) { debug!(%error, "bearer token refused by auth.jwt policy"); + attempt.refuse(state, &error.to_string()); return None; } + attempt.accept(state, &identity, &verified.claims().sub); Some((identity, verified)) } diff --git a/nodedb/src/control/server/session_auth/native.rs b/nodedb/src/control/server/session_auth/native.rs index cf869b148..05adcd1a9 100644 --- a/nodedb/src/control/server/session_auth/native.rs +++ b/nodedb/src/control/server/session_auth/native.rs @@ -6,7 +6,9 @@ //! - `{"method": "trust"}` — trust mode (only if configured) //! - `{"method": "password", "username": "...", "password": "..."}` //! - `{"method": "api_key", "token": "ndb_..."}` -//! - `{"method": "oidc_bearer", "token": "eyJ..."}` +//! +//! An OIDC bearer token uses the `OidcBearer` auth frame, which +//! `native::dispatch::handle_auth` verifies before any JSON body is built. use crate::config::auth::AuthMode; use crate::control::security::audit::{ @@ -299,40 +301,12 @@ pub async fn authenticate( .map(|id| (id, None)) } - "oidc_bearer" => { - let token = body["token"] - .as_str() - .ok_or_else(|| crate::Error::BadRequest { - detail: "missing 'token' for oidc_bearer auth".into(), - })?; - - // This dispatcher's `"oidc_bearer"` method string is unreachable from - // the native protocol: `dispatch::handle_auth` intercepts - // `ProtoAuth::OidcBearer` before building a JSON body for this - // function, and no other caller passes this method name. The - // verified-claims proof this returns is therefore discarded here - // rather than threaded further — the reachable OIDC path is - // `dispatch::handle_auth`. - let (identity, _verified_claims) = - crate::control::security::oidc::verify_bearer_token(state, token).await?; - - state.audit_record( - AuditEvent::AuthSuccess, - Some(identity.tenant_id), - peer_addr, - &format!( - "OIDC bearer login: sub={} method=oidc_bearer", - identity.username - ), - ); - state.auth_metrics.record_auth_success("oidc_bearer"); - - Ok((identity, None)) - } - + // An OIDC bearer token never reaches this dispatcher. + // `dispatch::handle_auth` verifies it and keeps the verified claims for the session. other => Err(crate::Error::BadRequest { detail: format!( - "unknown auth method: '{other}'. Use 'trust', 'password', 'api_key', or 'oidc_bearer'." + "unknown auth method: '{other}'. Use 'trust', 'password', or 'api_key'; \ + an OIDC bearer token uses the native OidcBearer auth frame." ), }), } diff --git a/nodedb/src/control/server/set_op_merge/intersect_except.rs b/nodedb/src/control/server/set_op_merge/intersect_except.rs index 0dc9c11b0..71438d84b 100644 --- a/nodedb/src/control/server/set_op_merge/intersect_except.rs +++ b/nodedb/src/control/server/set_op_merge/intersect_except.rs @@ -3,7 +3,7 @@ //! `INTERSECT` / `EXCEPT` merge over msgpack row payloads, comparing rows //! by logical column values. -use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; +use crate::control::server::payload_merge::{encode_msgpack_array, payload_rows}; use super::row_key::{extract_values_key, logical_row_bytes, rows_match}; @@ -19,19 +19,24 @@ pub(crate) enum SetMergeMode { /// /// Rows compare by value (see `row_key::rows_match`). The output holds the /// logical row bytes (the `{id, data}` wrapper is stripped), deduplicated by -/// values key, as one msgpack array. No payloads yields an empty array. -pub(crate) fn merge_set_op_payloads(payloads: &[Vec], mode: SetMergeMode) -> Vec { - if payloads.is_empty() { - return vec![0x90]; - } +/// values key, as one msgpack array. No payloads yields an empty array. Each +/// payload holds its rows in a shape [`payload_rows`] reads, and a payload it +/// refuses fails the merge. +pub(crate) fn merge_set_op_payloads( + payloads: &[Vec], + mode: SetMergeMode, +) -> crate::Result> { + let Some((first, others)) = payloads.split_first() else { + return Ok(encode_msgpack_array(&[])); + }; - let first_rows = extract_msgpack_elements(&payloads[0]); + let first_rows = payload_rows(first)?; let mut result_rows: Vec> = match mode { SetMergeMode::Intersect => { - let other_rows: Vec>> = payloads[1..] + let other_rows: Vec>> = others .iter() - .map(|p| extract_msgpack_elements(p)) - .collect(); + .map(|p| payload_rows(p)) + .collect::>()?; first_rows .into_iter() .filter(|row| { @@ -43,10 +48,10 @@ pub(crate) fn merge_set_op_payloads(payloads: &[Vec], mode: SetMergeMode) -> .collect() } SetMergeMode::Except => { - let other_rows: Vec> = payloads[1..] - .iter() - .flat_map(|p| extract_msgpack_elements(p)) - .collect(); + let mut other_rows: Vec> = Vec::new(); + for payload in others { + other_rows.extend(payload_rows(payload)?); + } first_rows .into_iter() .filter(|row| !other_rows.iter().any(|other| rows_match(row, other))) @@ -58,7 +63,7 @@ pub(crate) fn merge_set_op_payloads(payloads: &[Vec], mode: SetMergeMode) -> let mut seen = std::collections::HashSet::with_capacity(result_rows.len()); result_rows.retain(|r| seen.insert(extract_values_key(r))); - encode_msgpack_array(&result_rows) + Ok(encode_msgpack_array(&result_rows)) } #[cfg(test)] @@ -80,7 +85,7 @@ mod tests { serde_json::json!({"id":"doc-2","data":{"user_id":"u3"}}), ]); - let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Intersect); + let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Intersect).unwrap(); let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); assert_eq!(json, r#"[{"id":"u1","name":"Alice"}]"#); @@ -94,9 +99,29 @@ mod tests { ]); let right = encode_array(&[serde_json::json!({"id":"doc-1","data":{"user_id":"u1"}})]); - let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Except); + let merged = merge_set_op_payloads(&[left, right], SetMergeMode::Except).unwrap(); let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); assert_eq!(json, r#"[{"id":"u2"}]"#); } + + #[test] + fn no_payloads_is_an_empty_array() { + assert_eq!( + merge_set_op_payloads(&[], SetMergeMode::Intersect).unwrap(), + vec![0x90] + ); + } + + /// A branch payload in no row shape is refused in either mode. + #[test] + fn a_scalar_branch_is_refused() { + let left = encode_array(&[serde_json::json!({"id":"u1"})]); + for mode in [SetMergeMode::Intersect, SetMergeMode::Except] { + assert!(matches!( + merge_set_op_payloads(&[left.clone(), vec![0x2a]], mode), + Err(crate::Error::Codec { .. }) + )); + } + } } diff --git a/nodedb/src/control/server/set_op_merge/union.rs b/nodedb/src/control/server/set_op_merge/union.rs index 4bcb9da77..777009ffd 100644 --- a/nodedb/src/control/server/set_op_merge/union.rs +++ b/nodedb/src/control/server/set_op_merge/union.rs @@ -3,19 +3,19 @@ //! `UNION DISTINCT` merge: every row from every payload, deduplicated on //! its raw msgpack bytes, in encounter order. -use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; +use crate::control::server::payload_merge::{encode_msgpack_array, payload_rows}; /// Merge multiple row payloads and drop duplicate rows. /// -/// Each payload is a msgpack array of rows. A row's raw msgpack bytes are -/// its dedup key, so no decode round-trip runs. A payload that is not a -/// msgpack array is treated as one row. The output is one msgpack array of -/// the unique rows in encounter order. -pub(crate) fn dedup_union_payloads(payloads: &[Vec]) -> Vec { - let rows: Vec> = payloads - .iter() - .flat_map(|payload| extract_msgpack_elements(payload)) - .collect(); +/// Each payload holds its rows in a shape [`payload_rows`] reads. A row's +/// raw msgpack bytes are its dedup key, so no decode round-trip runs. A +/// payload [`payload_rows`] refuses fails the merge. The output is one +/// msgpack array of the unique rows in encounter order. +pub(crate) fn dedup_union_payloads(payloads: &[Vec]) -> crate::Result> { + let mut rows: Vec> = Vec::new(); + for payload in payloads { + rows.extend(payload_rows(payload)?); + } let mut seen: std::collections::HashSet> = std::collections::HashSet::with_capacity(rows.len()); let mut unique_rows: Vec> = Vec::with_capacity(rows.len()); @@ -26,7 +26,7 @@ pub(crate) fn dedup_union_payloads(payloads: &[Vec]) -> Vec { } } - encode_msgpack_array(&unique_rows) + Ok(encode_msgpack_array(&unique_rows)) } #[cfg(test)] @@ -42,7 +42,7 @@ mod tests { let left = encode_array(&[serde_json::json!({"x": 1}), serde_json::json!({"x": 2})]); let right = encode_array(&[serde_json::json!({"x": 2}), serde_json::json!({"x": 3})]); - let merged = dedup_union_payloads(&[left, right]); + let merged = dedup_union_payloads(&[left, right]).unwrap(); let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); assert_eq!(json, r#"[{"x":1},{"x":2},{"x":3}]"#); @@ -50,6 +50,26 @@ mod tests { #[test] fn union_distinct_of_no_payloads_is_empty_array() { - assert_eq!(dedup_union_payloads(&[]), vec![0x90]); + assert_eq!(dedup_union_payloads(&[]).unwrap(), vec![0x90]); + } + + /// A point-get branch answers its one row as a bare map. + #[test] + fn a_point_get_branch_is_one_row() { + let left = encode_array(&[serde_json::json!({"x": 1})]); + let map = nodedb_types::json_to_msgpack(&serde_json::json!({"x": 2})).unwrap(); + let merged = dedup_union_payloads(&[left, map]).unwrap(); + let json = crate::data::executor::response_codec::decode_payload_to_json(&merged); + assert_eq!(json, r#"[{"x":1},{"x":2}]"#); + } + + /// A branch payload in no row shape is refused, never read as one row. + #[test] + fn a_scalar_branch_is_refused() { + let left = encode_array(&[serde_json::json!({"x": 1})]); + assert!(matches!( + dedup_union_payloads(&[left, vec![0x2a]]), + Err(crate::Error::Codec { .. }) + )); } } diff --git a/nodedb/src/control/server/shared/authorization/requirements.rs b/nodedb/src/control/server/shared/authorization/requirements.rs index 2352cfd73..f95b2f066 100644 --- a/nodedb/src/control/server/shared/authorization/requirements.rs +++ b/nodedb/src/control/server/shared/authorization/requirements.rs @@ -204,6 +204,7 @@ mod tests { sub_aggregates: Vec::new(), grouping_sets: Vec::new(), sort_keys: Vec::new(), + group_key_aliases: Vec::new(), }); let orders = vec![AuthorizationRequirement::collection( "orders", diff --git a/nodedb/src/control/server/shared/cluster_array_dispatch.rs b/nodedb/src/control/server/shared/cluster_array_dispatch.rs index a816b659d..50da36f14 100644 --- a/nodedb/src/control/server/shared/cluster_array_dispatch.rs +++ b/nodedb/src/control/server/shared/cluster_array_dispatch.rs @@ -101,8 +101,11 @@ pub(crate) async fn run_trusted_cluster_array( let op = match &cluster_op { ClusterArrayOp::Slice { .. } => "slice", ClusterArrayOp::Agg { .. } => "agg", + ClusterArrayOp::Elementwise { .. } => "elementwise", ClusterArrayOp::Put { .. } => "put", ClusterArrayOp::Delete { .. } => "delete", + ClusterArrayOp::Flush { .. } => "flush", + ClusterArrayOp::Compact { .. } => "compact", }; tracing::debug!( op, diff --git a/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/chunker.rs b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/chunker.rs new file mode 100644 index 000000000..438648d56 --- /dev/null +++ b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/chunker.rs @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Chunking of a bulk load's rows into cell batches, one batch per vShard. + +use std::collections::HashMap; + +use nodedb_sql::types_array::ArrayInsertRow; + +/// The encoded size one load chunk targets: 4 MiB of `ArrayPutCell` +/// msgpack, by the per-row upper bound. +/// +/// One chunk is one `ArrayOp::Put`: one plan, one WAL record and one Data +/// Plane batch. 4 MiB stays far under the 64 MiB WAL record cap, so a chunk +/// never fails at the WAL append. It keeps the Control Plane's per-chunk +/// buffers small, so a load holds its input plus a few chunks. It is large +/// enough that the per-chunk costs (one plan, one WAL append, one SPSC round +/// trip) spread over tens of thousands of cells. +pub(crate) const ARRAY_COPY_CHUNK_BYTES: usize = 4 * 1024 * 1024; + +/// How many chunk sizes of rows the chunker holds across all its vShard +/// buckets. Past it the fullest bucket dispatches, so a load that touches +/// many shards still holds bounded memory. +const PENDING_CHUNKS: usize = 4; + +/// One chunk to dispatch: its rows, all bound for one vShard, in input +/// order, the 1-based number of its first row, and the upper bound of its +/// cells' encoded bytes. +pub(super) struct Chunk { + pub rows: Vec, + pub first_row: usize, + pub bytes: usize, +} + +struct Bucket { + rows: Vec, + bytes: usize, + first_row: usize, +} + +/// Buckets rows by the vShard their tile routes to and hands out a bucket as +/// a chunk once it reaches the chunk size. +/// +/// A repeated coordinate routes to the same vShard, so it lands in the same +/// bucket after its first write. Chunks of one bucket dispatch in input +/// order, so the last write of a coordinate wins. +pub(super) struct ArrayChunker { + chunk_bytes: usize, + pending_cap: usize, + pending: usize, + buckets: HashMap, +} + +impl ArrayChunker { + pub(super) fn new(chunk_bytes: usize) -> Self { + Self { + chunk_bytes, + pending_cap: chunk_bytes.saturating_mul(PENDING_CHUNKS), + pending: 0, + buckets: HashMap::new(), + } + } + + /// Add row `row_no`, bound for `vshard`, whose cell encodes to at most + /// `bound` bytes. Returns a chunk to dispatch now: the row's bucket once + /// it reaches the chunk size, else the fullest bucket once the pending + /// total reaches its cap. + pub(super) fn push( + &mut self, + vshard: u32, + row_no: usize, + row: ArrayInsertRow, + bound: usize, + ) -> Option { + let bucket = self.buckets.entry(vshard).or_insert_with(|| Bucket { + rows: Vec::new(), + bytes: 0, + first_row: row_no, + }); + if bucket.rows.is_empty() { + bucket.first_row = row_no; + } + bucket.rows.push(row); + bucket.bytes += bound; + self.pending += bound; + if bucket.bytes >= self.chunk_bytes { + return self.take(vshard); + } + if self.pending >= self.pending_cap { + let fullest = self + .buckets + .iter() + .max_by_key(|(_, b)| b.bytes) + .map(|(&v, _)| v)?; + return self.take(fullest); + } + None + } + + /// The buckets left at the end of the input, as chunks in the order of + /// their first rows. + pub(super) fn finish(&mut self) -> Vec { + self.pending = 0; + let mut chunks: Vec = self + .buckets + .drain() + .map(|(_, bucket)| bucket) + .filter(|b| !b.rows.is_empty()) + .map(|b| Chunk { + rows: b.rows, + first_row: b.first_row, + bytes: b.bytes, + }) + .collect(); + chunks.sort_by_key(|c| c.first_row); + chunks + } + + fn take(&mut self, vshard: u32) -> Option { + let bucket = self.buckets.remove(&vshard)?; + self.pending -= bucket.bytes; + Some(Chunk { + rows: bucket.rows, + first_row: bucket.first_row, + bytes: bucket.bytes, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_sql::types_array::{ArrayAttrLiteral, ArrayCoordLiteral}; + + fn row(n: i64) -> ArrayInsertRow { + ArrayInsertRow { + coords: vec![ArrayCoordLiteral::Int64(n)], + attrs: vec![ArrayAttrLiteral::Int64(n)], + } + } + + #[test] + fn one_bucket_chunks_by_size_in_order() { + let mut chunker = ArrayChunker::new(100); + let mut chunks = Vec::new(); + for n in 1..=10 { + chunks.extend(chunker.push(0, n as usize, row(n), 30)); + } + chunks.extend(chunker.finish()); + let sizes: Vec = chunks.iter().map(|c| c.rows.len()).collect(); + assert_eq!(sizes, vec![4, 4, 2]); + let firsts: Vec = chunks.iter().map(|c| c.first_row).collect(); + assert_eq!(firsts, vec![1, 5, 9]); + } + + #[test] + fn the_pending_cap_flushes_the_fullest_bucket() { + let mut chunker = ArrayChunker::new(100); + // Buckets of 95, 90, 90 and 90 bytes: 365 pending, under the 400 cap. + assert!(chunker.push(1, 1, row(1), 95).is_none()); + for v in 2..=4u32 { + assert!(chunker.push(v, v as usize, row(i64::from(v)), 90).is_none()); + } + // 40 bytes on a fifth bucket reach 405. No bucket is full, so the + // fullest, bucket 1, dispatches. + let chunk = chunker.push(5, 5, row(5), 40).expect("flush"); + assert_eq!(chunk.first_row, 1); + assert_eq!(chunk.rows.len(), 1); + let left: Vec = chunker.finish().iter().map(|c| c.first_row).collect(); + assert_eq!(left, vec![2, 3, 4, 5]); + } + + #[test] + fn many_small_buckets_flush_before_the_cap_is_passed() { + let mut chunker = ArrayChunker::new(100); + let mut flushed = 0; + for v in 0..20u32 { + if chunker + .push(v, v as usize + 1, row(i64::from(v)), 25) + .is_some() + { + flushed += 1; + } + } + assert!(flushed > 0, "pending rows across buckets stay bounded"); + } +} diff --git a/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/decode.rs b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/decode.rs new file mode 100644 index 000000000..3ecc9572a --- /dev/null +++ b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/decode.rs @@ -0,0 +1,539 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Typed decode and validation of one bulk-load row against an array schema. +//! +//! The columns are the array's dims in schema order, then its attrs in +//! schema order. Each value parses by its declared type and nothing else: +//! no type is guessed from the text. A row is checked for arity, type, +//! NOT NULL and the dim domain before any row of the load is written. + +use std::borrow::Cow; + +use nodedb_array::types::cell_value::value::CellValue; +use nodedb_sql::catalog::ArrayCatalogView; +use nodedb_sql::types_array::{ + ArrayAttrAst, ArrayAttrLiteral, ArrayAttrType, ArrayCoordLiteral, ArrayDimType, + ArrayDomainBound, ArrayInsertRow, +}; +use nodedb_types::Value; + +use crate::control::planner::sql_plan_convert::array_convert::engine_attr_type; + +/// A row a load refuses: its 1-based number, SQLSTATE and message. +#[derive(Debug)] +pub(super) struct RowError { + pub row: usize, + pub sqlstate: &'static str, + pub message: String, +} + +impl RowError { + pub(super) fn new(row: usize, sqlstate: &'static str, message: String) -> Self { + Self { + row, + sqlstate, + message, + } + } +} + +/// Upper bound of the msgpack bytes of one cell's fixed fields: the +/// six-field array header, the coordinate and attribute array headers, the +/// surrogate, and three timestamps. +const CELL_FIXED_BOUND: usize = 48; + +/// Upper bound of the msgpack bytes of one coordinate or attribute value, +/// beyond its string or byte payload. zerompk writes an enum value as a +/// two-element array of the variant name and the payload. The longest name +/// is `TimestampMs`, and an `i64` or `f64` takes 9 bytes. +const SCALAR_BOUND: usize = 24; + +/// Check a header record against the schema: dims, then attrs, by name. +pub(super) fn check_header( + fields: &[Option>], + view: &ArrayCatalogView, + row: usize, +) -> Result<(), RowError> { + let expected: Vec<&str> = view + .dims + .iter() + .map(|d| d.name.as_str()) + .chain(view.attrs.iter().map(|a| a.name.as_str())) + .collect(); + if fields.len() != expected.len() { + return Err(RowError::new( + row, + "42703", + format!( + "header has {} columns, array '{}' has {} ({})", + fields.len(), + view.name, + expected.len(), + expected.join(", ") + ), + )); + } + for (index, (field, want)) in fields.iter().zip(&expected).enumerate() { + let got = field.as_deref().unwrap_or(""); + if !got.trim().eq_ignore_ascii_case(want) { + return Err(RowError::new( + row, + "42703", + format!( + "header column {} is '{got}', array '{}' has '{want}' there", + index + 1, + view.name + ), + )); + } + } + Ok(()) +} + +/// Decode one text record into a typed row. +pub(super) fn decode_text_row( + fields: &[Option>], + view: &ArrayCatalogView, + row: usize, +) -> Result { + check_arity(fields.len(), view, row)?; + let (dim_fields, attr_fields) = fields.split_at(view.dims.len()); + let mut coords = Vec::with_capacity(view.dims.len()); + for (field, dim) in dim_fields.iter().zip(&view.dims) { + let text = field.as_deref().ok_or_else(|| null_dim(&dim.name, row))?; + let coord = parse_coord(text, dim.dtype) + .ok_or_else(|| invalid(&dim.name, dim_type_name(dim.dtype), text, row))?; + check_domain(&coord, &dim.lo, &dim.hi, &dim.name, row)?; + coords.push(coord); + } + let mut attrs = Vec::with_capacity(view.attrs.len()); + for (field, attr) in attr_fields.iter().zip(&view.attrs) { + let value = match field.as_deref() { + None if attr.nullable => ArrayAttrLiteral::Null, + None => return Err(not_null(&attr.name, row)), + Some(text) => { + let literal = parse_attr(text, attr.dtype) + .ok_or_else(|| invalid(&attr.name, attr.dtype.sql_name(), text, row))?; + check_width(&literal, attr, row)?; + literal + } + }; + attrs.push(value); + } + Ok(ArrayInsertRow { coords, attrs }) +} + +/// Decode row `index` of native columns into a typed row. `dims` and +/// `attrs` hold one column per dim and per attr, in schema order. +pub(super) fn decode_value_row( + dims: &[Vec], + attrs: &[Vec], + index: usize, + view: &ArrayCatalogView, +) -> Result { + let row = index + 1; + let mut coords = Vec::with_capacity(view.dims.len()); + for (column, dim) in dims.iter().zip(&view.dims) { + let value = column.get(index).unwrap_or(&Value::Null); + let coord = match (value, dim.dtype) { + (Value::Null, _) => return Err(null_dim(&dim.name, row)), + (Value::Integer(n), ArrayDimType::Int64) => ArrayCoordLiteral::Int64(*n), + (Value::Integer(n), ArrayDimType::TimestampMs) => ArrayCoordLiteral::Int64(*n), + (Value::Float(f), ArrayDimType::Float64) => ArrayCoordLiteral::Float64(*f), + (Value::Integer(n), ArrayDimType::Float64) => ArrayCoordLiteral::Float64(*n as f64), + (Value::String(s), ArrayDimType::String) => ArrayCoordLiteral::String(s.clone()), + (other, dtype) => { + return Err(mismatch(&dim.name, dim_type_name(dtype), other, row)); + } + }; + check_domain(&coord, &dim.lo, &dim.hi, &dim.name, row)?; + coords.push(coord); + } + let mut out = Vec::with_capacity(view.attrs.len()); + for (column, attr) in attrs.iter().zip(&view.attrs) { + let value = column.get(index).unwrap_or(&Value::Null); + let literal = match (value, attr.dtype) { + (Value::Null, _) if attr.nullable => ArrayAttrLiteral::Null, + (Value::Null, _) => return Err(not_null(&attr.name, row)), + (Value::Integer(n), dtype) if dtype.is_integer() => ArrayAttrLiteral::Int64(*n), + (Value::Integer(n), dtype) if dtype.is_float() => ArrayAttrLiteral::Float64(*n as f64), + (Value::Float(f), dtype) if dtype.is_float() => ArrayAttrLiteral::Float64(*f), + (Value::String(s), ArrayAttrType::String) => ArrayAttrLiteral::String(s.clone()), + (Value::Bytes(b), ArrayAttrType::Bytes) => ArrayAttrLiteral::Bytes(b.clone()), + (other, dtype) => { + return Err(mismatch(&attr.name, dtype.sql_name(), other, row)); + } + }; + check_width(&literal, attr, row)?; + out.push(literal); + } + Ok(ArrayInsertRow { coords, attrs: out }) +} + +/// Upper bound of the msgpack bytes the row's `ArrayPutCell` encodes to. +pub(super) fn wire_bound(row: &ArrayInsertRow) -> usize { + let coords: usize = row + .coords + .iter() + .map(|c| { + SCALAR_BOUND + + match c { + ArrayCoordLiteral::String(s) => s.len(), + ArrayCoordLiteral::Int64(_) | ArrayCoordLiteral::Float64(_) => 0, + } + }) + .sum(); + let attrs: usize = row + .attrs + .iter() + .map(|a| { + SCALAR_BOUND + + match a { + ArrayAttrLiteral::String(s) => s.len(), + ArrayAttrLiteral::Bytes(b) => b.len(), + ArrayAttrLiteral::Int64(_) + | ArrayAttrLiteral::Float64(_) + | ArrayAttrLiteral::Null => 0, + } + }) + .sum(); + CELL_FIXED_BOUND + coords + attrs +} + +fn check_arity(got: usize, view: &ArrayCatalogView, row: usize) -> Result<(), RowError> { + let want = view.dims.len() + view.attrs.len(); + if got == want { + return Ok(()); + } + Err(RowError::new( + row, + "22P02", + format!( + "row has {got} columns, array '{}' has {want} ({} dims, {} attrs)", + view.name, + view.dims.len(), + view.attrs.len() + ), + )) +} + +fn parse_coord(text: &str, dtype: ArrayDimType) -> Option { + match dtype { + ArrayDimType::Int64 | ArrayDimType::TimestampMs => { + text.trim().parse().ok().map(ArrayCoordLiteral::Int64) + } + ArrayDimType::Float64 => text.trim().parse().ok().map(ArrayCoordLiteral::Float64), + ArrayDimType::String => Some(ArrayCoordLiteral::String(text.to_string())), + } +} + +fn parse_attr(text: &str, dtype: ArrayAttrType) -> Option { + match dtype { + ArrayAttrType::Int64 + | ArrayAttrType::UInt8 + | ArrayAttrType::Int16 + | ArrayAttrType::UInt16 + | ArrayAttrType::Int32 + | ArrayAttrType::UInt32 => text.trim().parse().ok().map(ArrayAttrLiteral::Int64), + ArrayAttrType::Float64 | ArrayAttrType::Float32 => { + text.trim().parse().ok().map(ArrayAttrLiteral::Float64) + } + ArrayAttrType::String => Some(ArrayAttrLiteral::String(text.to_string())), + ArrayAttrType::Bytes => parse_bytea_hex(text).map(ArrayAttrLiteral::Bytes), + } +} + +/// A number outside its attr's declared width is `22003`, checked here so +/// the refusal names its row. +fn check_width( + literal: &ArrayAttrLiteral, + attr: &ArrayAttrAst, + row: usize, +) -> Result<(), RowError> { + let dtype = engine_attr_type(attr.dtype); + let narrowed = match literal { + ArrayAttrLiteral::Int64(n) => CellValue::from_int_literal(*n, dtype), + ArrayAttrLiteral::Float64(f) => CellValue::from_float_literal(*f, dtype), + ArrayAttrLiteral::String(_) | ArrayAttrLiteral::Bytes(_) | ArrayAttrLiteral::Null => { + return Ok(()); + } + }; + match narrowed { + Ok(_) => Ok(()), + Err(e) if e.is_out_of_range() => Err(RowError::new( + row, + "22003", + format!("attr '{}': {e}", attr.name), + )), + Err(e) => Err(RowError::new( + row, + "42804", + format!("attr '{}': {e}", attr.name), + )), + } +} + +/// PostgreSQL `bytea` hex form: `\x` then pairs of hex digits. +fn parse_bytea_hex(text: &str) -> Option> { + let hex = text.trim().strip_prefix("\\x")?; + if !hex.len().is_multiple_of(2) { + return None; + } + (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(hex.get(i..i + 2)?, 16).ok()) + .collect() +} + +/// A coordinate outside its dim's inclusive domain is `22003`. A bound of +/// another type than the coordinate never holds it. +fn check_domain( + coord: &ArrayCoordLiteral, + lo: &ArrayDomainBound, + hi: &ArrayDomainBound, + dim: &str, + row: usize, +) -> Result<(), RowError> { + let inside = match (coord, lo, hi) { + ( + ArrayCoordLiteral::Int64(v), + ArrayDomainBound::Int64(lo) | ArrayDomainBound::TimestampMs(lo), + ArrayDomainBound::Int64(hi) | ArrayDomainBound::TimestampMs(hi), + ) => lo <= v && v <= hi, + ( + ArrayCoordLiteral::Float64(v), + ArrayDomainBound::Float64(lo), + ArrayDomainBound::Float64(hi), + ) => *lo <= *v && *v <= *hi, + ( + ArrayCoordLiteral::String(v), + ArrayDomainBound::String(lo), + ArrayDomainBound::String(hi), + ) => lo <= v && v <= hi, + _ => false, + }; + if inside { + return Ok(()); + } + Err(RowError::new( + row, + "22003", + format!( + "coordinate {} for dim '{dim}' is outside its domain [{}..{}]", + coord_text(coord), + bound_text(lo), + bound_text(hi) + ), + )) +} + +fn null_dim(dim: &str, row: usize) -> RowError { + RowError::new(row, "23502", format!("dim '{dim}' cannot be NULL")) +} + +fn not_null(attr: &str, row: usize) -> RowError { + RowError::new(row, "23502", format!("attr '{attr}' is NOT NULL")) +} + +fn invalid(column: &str, dtype: &str, text: &str, row: usize) -> RowError { + RowError::new( + row, + "22P02", + format!("invalid input for {dtype} column '{column}': '{text}'"), + ) +} + +fn mismatch(column: &str, dtype: &str, value: &Value, row: usize) -> RowError { + RowError::new( + row, + "22P02", + format!("column '{column}' is {dtype}, the cell holds {value:?}"), + ) +} + +fn dim_type_name(dtype: ArrayDimType) -> &'static str { + match dtype { + ArrayDimType::Int64 => "INT64", + ArrayDimType::Float64 => "FLOAT64", + ArrayDimType::TimestampMs => "TIMESTAMP", + ArrayDimType::String => "STRING", + } +} + +fn coord_text(coord: &ArrayCoordLiteral) -> String { + match coord { + ArrayCoordLiteral::Int64(v) => v.to_string(), + ArrayCoordLiteral::Float64(v) => v.to_string(), + ArrayCoordLiteral::String(v) => format!("'{v}'"), + } +} + +fn bound_text(bound: &ArrayDomainBound) -> String { + match bound { + ArrayDomainBound::Int64(v) | ArrayDomainBound::TimestampMs(v) => v.to_string(), + ArrayDomainBound::Float64(v) => v.to_string(), + ArrayDomainBound::String(v) => format!("'{v}'"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use nodedb_sql::types_array::{ArrayAttrAst, ArrayDimAst}; + + fn view() -> ArrayCatalogView { + ArrayCatalogView { + name: "scene".into(), + dims: vec![ArrayDimAst { + name: "y".into(), + dtype: ArrayDimType::Int64, + lo: ArrayDomainBound::Int64(0), + hi: ArrayDomainBound::Int64(255), + }], + attrs: vec![ + ArrayAttrAst { + name: "b1".into(), + dtype: ArrayAttrType::Float64, + nullable: false, + }, + ArrayAttrAst { + name: "label".into(), + dtype: ArrayAttrType::String, + nullable: true, + }, + ], + tile_extents: vec![64], + } + } + + fn fields(values: &[Option<&'static str>]) -> Vec>> { + values.iter().map(|v| v.map(Cow::Borrowed)).collect() + } + + #[test] + fn a_row_decodes_by_declared_type() { + let row = + decode_text_row(&fields(&[Some("7"), Some("1.5"), None]), &view(), 1).expect("decode"); + assert_eq!(row.coords, vec![ArrayCoordLiteral::Int64(7)]); + assert_eq!( + row.attrs, + vec![ArrayAttrLiteral::Float64(1.5), ArrayAttrLiteral::Null] + ); + assert!(wire_bound(&row) >= CELL_FIXED_BOUND + 3 * SCALAR_BOUND); + } + + #[test] + fn refusals_carry_their_sqlstate() { + let cases: [(&[Option<&'static str>], &str); 4] = [ + (&[Some("7"), Some("1.5")], "22P02"), + (&[Some("7"), Some("abc"), None], "22P02"), + (&[Some("7"), None, None], "23502"), + (&[Some("256"), Some("1.0"), None], "22003"), + ]; + for (input, sqlstate) in cases { + let err = decode_text_row(&fields(input), &view(), 9).expect_err("refused"); + assert_eq!(err.sqlstate, sqlstate, "{input:?}"); + assert_eq!(err.row, 9); + } + } + + #[test] + fn a_header_must_name_dims_then_attrs() { + assert!(check_header(&fields(&[Some("y"), Some("B1"), Some("label")]), &view(), 1).is_ok()); + let err = check_header(&fields(&[Some("y"), Some("label"), Some("b1")]), &view(), 1) + .expect_err("mismatch"); + assert_eq!(err.sqlstate, "42703"); + } + + #[test] + fn native_cells_decode_and_check_types() { + let dims = vec![vec![Value::Integer(3), Value::Integer(4)]]; + let attrs = vec![ + vec![Value::Float(0.5), Value::String("x".into())], + vec![Value::Null, Value::Null], + ]; + let row = decode_value_row(&dims, &attrs, 0, &view()).expect("decode"); + assert_eq!(row.coords, vec![ArrayCoordLiteral::Int64(3)]); + let err = decode_value_row(&dims, &attrs, 1, &view()).expect_err("string in float"); + assert_eq!(err.sqlstate, "22P02"); + assert_eq!(err.row, 2); + } + + fn narrow_view() -> ArrayCatalogView { + let attr = |name: &str, dtype| ArrayAttrAst { + name: name.into(), + dtype, + nullable: true, + }; + ArrayCatalogView { + name: "scene".into(), + dims: vec![ArrayDimAst { + name: "y".into(), + dtype: ArrayDimType::Int64, + lo: ArrayDomainBound::Int64(0), + hi: ArrayDomainBound::Int64(255), + }], + attrs: vec![ + attr("q", ArrayAttrType::UInt8), + attr("b1", ArrayAttrType::UInt16), + attr("t", ArrayAttrType::Int16), + attr("ndvi", ArrayAttrType::Float32), + ], + tile_extents: vec![64], + } + } + + /// Text and native rows narrow by declared width: an in-range value + /// loads, and one past the width is `22003` with its row number. + #[test] + fn narrow_attrs_load_in_range_and_refuse_past_it() { + let ok = decode_text_row( + &fields(&[ + Some("1"), + Some("255"), + Some("65535"), + Some("-32768"), + Some("0.1"), + ]), + &narrow_view(), + 1, + ) + .expect("in range"); + assert_eq!(ok.attrs[1], ArrayAttrLiteral::Int64(65535)); + let cases: [&[Option<&'static str>]; 5] = [ + &[Some("1"), Some("256"), None, None, None], + &[Some("1"), Some("-1"), None, None, None], + &[Some("1"), None, Some("65536"), None, None], + &[Some("1"), None, None, Some("32768"), None], + &[Some("1"), None, None, None, Some("1e39")], + ]; + for input in cases { + let err = decode_text_row(&fields(input), &narrow_view(), 4).expect_err("out of range"); + assert_eq!(err.sqlstate, "22003", "{input:?}"); + assert_eq!(err.row, 4); + } + let float_into_int = decode_text_row( + &fields(&[Some("1"), Some("1.5"), None, None, None]), + &narrow_view(), + 2, + ) + .expect_err("float text into UINT8"); + assert_eq!(float_into_int.sqlstate, "22P02"); + + let dims = vec![vec![Value::Integer(3)]]; + let attrs = vec![ + vec![Value::Integer(300)], + vec![Value::Null], + vec![Value::Null], + vec![Value::Float(0.5)], + ]; + let err = decode_value_row(&dims, &attrs, 0, &narrow_view()).expect_err("300 > UINT8"); + assert_eq!(err.sqlstate, "22003"); + } + + #[test] + fn bytea_hex_parses() { + assert_eq!(parse_bytea_hex("\\x0aFF"), Some(vec![0x0a, 0xff])); + assert_eq!(parse_bytea_hex("0aff"), None); + } +} diff --git a/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/dispatch.rs b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/dispatch.rs new file mode 100644 index 000000000..32e84df8e --- /dev/null +++ b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/dispatch.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Chunk dispatch and the refusals of an array bulk load. + +use nodedb_sql::SqlError; +use nodedb_sql::catalog::SqlCatalog; +use nodedb_sql::engine_rules::{BulkLoadParams, BulkLoadPlan}; +use nodedb_wal::record::MAX_WAL_PAYLOAD_SIZE; + +use crate::control::security::identity::AuthenticatedIdentity; +use crate::control::server::shared::ddl::neutral::collection::dml::{ + PlanBuilder, plan_and_dispatch_built, +}; +use crate::control::server::shared::ddl::result::DdlError; +use crate::control::server::shared::session::DmlTxnCtx; +use crate::control::state::SharedState; +use crate::types::DatabaseId; + +use super::super::entry::wrap_row_error; +use super::chunker::Chunk; +use super::decode::RowError; + +/// Upper bound of the redo bytes one chunk's sub-record adds beyond its +/// cells and the array name: the record type, the version byte, the +/// payload's array headers, the array id's ids, and the provenance. +const REDO_CHUNK_OVERHEAD: usize = 256; + +/// Who loads, into which database, in which transaction scope. +pub(crate) struct ArrayLoadCtx<'a> { + pub state: &'a SharedState, + pub identity: &'a AuthenticatedIdentity, + pub database_id: DatabaseId, + pub txn_ctx: &'a DmlTxnCtx<'a>, +} + +/// Dispatch one chunk through the write funnel. +pub(super) async fn dispatch_chunk( + ctx: &ArrayLoadCtx<'_>, + array: &str, + chunk: Chunk, + label: &str, + committed: u64, +) -> Result<(), DdlError> { + let Chunk { + rows, first_row, .. + } = chunk; + let collection = array.to_string(); + let build: PlanBuilder<'_> = Box::new(move |catalog: &dyn SqlCatalog| { + let params = BulkLoadParams { + collection: collection.clone(), + cells: rows, + }; + match nodedb_sql::planner::dml::plan_bulk_load(params, catalog)? { + BulkLoadPlan::CellChunks(plans) => Ok(plans), + BulkLoadPlan::RowInserts => Err(SqlError::Unsupported { + detail: format!("COPY: '{collection}' is no longer an array"), + }), + } + }); + plan_and_dispatch_built( + ctx.state, + ctx.identity, + ctx.identity.tenant_id, + ctx.database_id, + build, + ctx.txn_ctx, + true, + ) + .await + .map_err(|e| refuse(e, first_row, label, committed))?; + Ok(()) +} + +/// The most bytes one incomplete record of a COPY stream holds: the array +/// engine's memory budget headroom when the load opens. +pub(super) fn input_cap(state: &SharedState) -> usize { + state + .governor + .budget(nodedb_mem::EngineId::Array) + .available() +} + +/// The `54000` refusal of an incomplete record past [`input_cap`]. +pub(super) fn record_over_cap( + len: usize, + cap: usize, + row: usize, + label: &str, + committed: u64, +) -> DdlError { + let e = DdlError::new( + "54000", + format!( + "{len} bytes hold no record end, over the array engine's {cap}-byte memory \ + headroom; check the input's line breaks and quotes" + ), + ); + refuse(e, row, label, committed) +} + +/// Inside a block, every staged chunk's redo joins the one commit record. +/// Refuse the chunk that would take the staged redo of `chunks` chunks and +/// `cell_bytes` cell bytes past the WAL record cap, before it stages. +pub(super) fn check_redo_cap( + array: &str, + rows: u64, + cell_bytes: usize, + chunks: usize, +) -> Result<(), DdlError> { + let redo = cell_bytes + chunks * (REDO_CHUNK_OVERHEAD + array.len()); + if redo <= MAX_WAL_PAYLOAD_SIZE { + return Ok(()); + } + Err(DdlError::new( + "54000", + format!( + "COPY: {rows} rows need up to {redo} bytes of transaction redo, over the \ + {MAX_WAL_PAYLOAD_SIZE}-byte WAL record cap; load them outside a transaction \ + block, where each chunk commits on its own (0 rows committed)" + ), + )) +} + +/// A refused row. Row 0 is the input as a whole, such as a bad delimiter, +/// so its message names no row. +pub(super) fn refuse_row(e: RowError, label: &str, committed: u64) -> DdlError { + if e.row == 0 { + return DdlError::new( + e.sqlstate, + format!("COPY: {label}: {} ({committed} rows committed)", e.message), + ); + } + let row = e.row; + refuse(DdlError::new(e.sqlstate, e.message), row, label, committed) +} + +/// `COPY: