From 37d00861f9f451e43afa5c9cd6d54b7db40cd215 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Fri, 16 Oct 2026 15:29:59 +0800 Subject: [PATCH 01/29] docs: align README and guides with current behavior and array function names Document psql superuser authentication, the 0.6 release status, and the engines that carry bitemporal columns. Clarify that AS OF SYSTEM TIME NULL works on ARRAY_SLICE but not ARRAY_AGG. Rename the stale NDARRAY_* references in array comments and the ARRAY_AGG planner error message to the ARRAY_* function names. --- README.md | 10 ++++------ docs/bitemporal.md | 2 +- docs/getting-started.md | 4 +++- .../sql_plan_convert/array_fn_convert/aggregate.rs | 2 +- .../sql_plan_convert/array_fn_convert/helpers.rs | 6 +++--- nodedb/src/data/executor/array_checkpoint.rs | 2 +- .../handlers/control/checkpoint_durable_lsn.rs | 2 +- nodedb/src/data/executor/handlers/control/snapshot.rs | 2 +- 8 files changed, 15 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index a1e77b106..13f57ac27 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,7 @@ 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 +psql -h localhost -p 6432 -U nodedb # 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/bitemporal.md b/docs/bitemporal.md index 3bd2faf5f..e6ab2414c 100644 --- a/docs/bitemporal.md +++ b/docs/bitemporal.md @@ -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`. 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/getting-started.md b/docs/getting-started.md index b95b51d3c..bbc1bb872 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -734,9 +734,11 @@ You get a full TUI with syntax highlighting, tab completion, and history search. ### With psql ```bash -psql -h localhost -p 6432 +psql -h localhost -p 6432 -U nodedb ``` +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`. + 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/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..cb7adb672 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 @@ -216,7 +216,7 @@ 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, 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..637f29cbc 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 @@ -129,7 +129,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 +140,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 ARRAY_AGG; \ + use a point-in-time AS OF SYSTEM TIME , or ARRAY_SLICE for the audit log" ), }); } diff --git a/nodedb/src/data/executor/array_checkpoint.rs b/nodedb/src/data/executor/array_checkpoint.rs index cd70fb59f..7d983a6ab 100644 --- a/nodedb/src/data/executor/array_checkpoint.rs +++ b/nodedb/src/data/executor/array_checkpoint.rs @@ -10,7 +10,7 @@ //! calls `note_write` with the Control Plane's `wal_lsn`). The periodic //! checkpoint reports those writes as durable, and the manager then removes //! the `ArrayPut` / `ArrayDelete` records below that LSN. The checkpoint must -//! therefore flush every memtable first. An explicit `NDARRAY_FLUSH` and the +//! therefore flush every memtable first. An explicit `ARRAY_FLUSH` and the //! `flush_cell_threshold` flush in `dispatch::array::mutate` are not ordered //! against that removal, so they cannot stand in for it. //! diff --git a/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs b/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs index 453884f67..d0316977f 100644 --- a/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs +++ b/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs @@ -193,7 +193,7 @@ impl CoreLoop { /// Unlike the engines above this writes no checkpoint file of its own: the /// array engine's durable form IS its on-disk segments, and /// `checkpoint_array_engines` calls the same `ArrayEngine::flush` an - /// `NDARRAY_FLUSH` does. The bug it closes is that the flush was reachable + /// `ARRAY_FLUSH` does. The bug it closes is that the flush was reachable /// ONLY by that explicit command, so every cell written since the last one a /// user happened to run sat in a memory-only memtable while its `ArrayPut` /// records were truncated out from under it. diff --git a/nodedb/src/data/executor/handlers/control/snapshot.rs b/nodedb/src/data/executor/handlers/control/snapshot.rs index fc04eeae0..6bcbec055 100644 --- a/nodedb/src/data/executor/handlers/control/snapshot.rs +++ b/nodedb/src/data/executor/handlers/control/snapshot.rs @@ -359,7 +359,7 @@ impl CoreLoop { // 3d. Flush the array engines. Same memory-only memtable as columnar, // but the durable form already exists: this calls the very - // `ArrayEngine::flush` an `NDARRAY_FLUSH` calls, whose segments the + // `ArrayEngine::flush` an `ARRAY_FLUSH` calls, whose segments the // boot path already mmaps. The hazard was that nothing but that // explicit user command ever invoked it. durable_lsns.push(self.checkpoint_array_durable_lsn()); From 017d8f11c35092a32a3d28e08193598c1ef33247 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Fri, 16 Oct 2026 23:26:59 +0800 Subject: [PATCH 02/29] perf(codec): read double-delta bit fields a byte at a time BitReader::read_bits now takes whole byte runs instead of one bit per step, rejects fields wider than 64 bits as corrupt, and leaves the cursor untouched when a read runs past the buffer. --- nodedb-codec/src/double_delta.rs | 70 +++++++++++++++++++++++++++++++- 1 file changed, 68 insertions(+), 2 deletions(-) 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 { .. }))); + } } From df0dbba615d70c29e642acf3f810f31ef32ca9b2 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Fri, 16 Oct 2026 23:27:07 +0800 Subject: [PATCH 03/29] feat(spatial): tag geometries with an SRID and add CRS transforms Value::Geometry now carries an SRID alongside its shape. GeoJSON reads and writes a `crs` member for any SRID other than 4326, and the msgpack form gains an SRID field. nodedb-spatial gains a projection module (transverse Mercator and oblique Mercator families) and planar measures for projected SRIDs. WKB/EWKB handling moves into a module that preserves the SRID. R-tree bounding boxes are computed in the geometry's own space. SQL geometry functions read the SRID of each operand. Operands in different SRIDs, an unsupported SRID, or a coordinate outside a projection's domain raise a data exception (SQLSTATE 22000) rather than transforming implicitly; ST_Transform converts explicitly. A spatial index fixes its SRID at its first entry and refuses geometries in another SRID. --- fuzz/src/targets/strict_tuple.rs | 6 +- .../src/distributed_spatial/coordinator.rs | 11 +- nodedb-physical/src/physical_plan/spatial.rs | 12 +- nodedb-query/src/expr/eval.rs | 23 +- nodedb-query/src/functions/eval.rs | 25 +- nodedb-query/src/geo_functions/accessors.rs | 71 +- .../src/geo_functions/catalog/spec.rs | 4 +- .../src/geo_functions/catalog/table.rs | 3 + .../src/geo_functions/constructors.rs | 345 +++++++--- nodedb-query/src/geo_functions/crs.rs | 156 +++++ nodedb-query/src/geo_functions/dispatch.rs | 46 +- nodedb-query/src/geo_functions/helpers.rs | 151 ++++- nodedb-query/src/geo_functions/indexing.rs | 17 +- nodedb-query/src/geo_functions/measures.rs | 93 ++- nodedb-query/src/geo_functions/mod.rs | 1 + nodedb-query/src/geo_functions/predicates.rs | 126 +++- nodedb-spatial/src/crs/ellipsoid.rs | 69 ++ nodedb-spatial/src/crs/error.rs | 130 ++++ nodedb-spatial/src/crs/hom.rs | 284 ++++++++ nodedb-spatial/src/crs/mod.rs | 18 + nodedb-spatial/src/crs/params.rs | 103 +++ nodedb-spatial/src/crs/tm.rs | 292 ++++++++ nodedb-spatial/src/crs/transform.rs | 206 ++++++ nodedb-spatial/src/hybrid.rs | 48 +- nodedb-spatial/src/lib.rs | 14 +- nodedb-spatial/src/measures.rs | 72 +- nodedb-spatial/src/operations.rs | 134 +++- nodedb-spatial/src/persist.rs | 49 +- nodedb-spatial/src/planar/area.rs | 96 +++ nodedb-spatial/src/planar/buffer.rs | 37 + nodedb-spatial/src/planar/distance.rs | 265 ++++++++ nodedb-spatial/src/planar/mod.rs | 11 + nodedb-spatial/src/predicates/contains.rs | 16 +- nodedb-spatial/src/predicates/edge.rs | 8 + nodedb-spatial/src/predicates/intersects.rs | 16 +- nodedb-spatial/src/rtree/bulk_load.rs | 45 +- nodedb-spatial/src/rtree/delete.rs | 13 +- nodedb-spatial/src/rtree/insert.rs | 30 +- nodedb-spatial/src/rtree/mod.rs | 2 +- nodedb-spatial/src/rtree/node.rs | 9 +- nodedb-spatial/src/rtree/search.rs | 41 +- nodedb-spatial/src/rtree/split.rs | 65 +- nodedb-spatial/src/rtree/tree.rs | 145 +++- nodedb-spatial/src/spatial_join.rs | 32 +- nodedb-spatial/src/wkb.rs | 634 ------------------ nodedb-spatial/src/wkb/ewkb.rs | 62 ++ nodedb-spatial/src/wkb/format.rs | 25 + nodedb-spatial/src/wkb/mod.rs | 21 + nodedb-spatial/src/wkb/read.rs | 387 +++++++++++ nodedb-spatial/src/wkb/write.rs | 226 +++++++ .../builtins/scalars/spatial/arg_shape.rs | 2 + nodedb-sql/src/planner/const_fold.rs | 8 +- .../src/planner/geometry_expr/resolve.rs | 52 +- nodedb-sql/src/types/plan/variants/plan.rs | 4 +- nodedb-sql/src/visitor/plan_visitor/args.rs | 2 +- nodedb-strict/src/decode/value.rs | 29 +- nodedb-strict/src/encode/value.rs | 1 + nodedb-types/src/bbox.rs | 127 +++- nodedb-types/src/geometry.rs | 245 ------- nodedb-types/src/geometry/math.rs | 97 +++ nodedb-types/src/geometry/mod.rs | 17 + nodedb-types/src/geometry/shape.rs | 221 ++++++ nodedb-types/src/geometry/srid.rs | 169 +++++ nodedb-types/src/geometry/tagged.rs | 368 ++++++++++ nodedb-types/src/lib.rs | 6 +- nodedb-types/src/sync/wire/spatial.rs | 7 +- nodedb-types/src/value/convert.rs | 9 +- nodedb-types/src/value/core.rs | 9 +- nodedb-types/src/value/display.rs | 2 +- nodedb-types/src/value/msgpack.rs | 35 +- nodedb-types/src/value/sql_literal.rs | 17 +- nodedb-wal/src/record/fts_spatial.rs | 5 +- nodedb/src/bridge/envelope/error_code_from.rs | 8 +- .../control/cluster/array_executor/refusal.rs | 1 + .../control/cluster/execution_error_wire.rs | 1 + .../control/cluster/metadata_applier/wedge.rs | 1 + .../error_map/class_parity/error_index.rs | 3 +- .../error_map/class_parity/error_samples.rs | 6 + nodedb/src/control/gateway/error_map/resp.rs | 1 + .../planner/sql_plan_convert/scan_params.rs | 2 +- .../native/dispatch/plan_builder/spatial.rs | 5 +- .../control/server/pgwire/types/error_map.rs | 1 + nodedb/src/control/server/shared/retry.rs | 1 + .../predicate/txn_buffering/classify.rs | 6 +- .../sync/async_dispatch/delta/compensation.rs | 1 + nodedb/src/control/server/sync/refusal.rs | 2 + .../control/server/sync/spatial_handler.rs | 6 +- .../control/server/sync/spatial_session.rs | 6 +- .../src/control/server/wal_dispatch/core.rs | 10 +- .../server/wal_dispatch_fts_spatial.rs | 8 +- .../wal_replication/decode_sync_engines.rs | 11 +- .../wal_replication/encode/columnar.rs | 2 +- .../columnar_checkpoint/geometry_restore.rs | 6 +- .../handlers/columnar_mutation_apply.rs | 28 +- .../handlers/columnar_write/geometry_index.rs | 98 ++- .../columnar_write/geometry_remove.rs | 12 +- .../handlers/columnar_write/insert.rs | 13 +- .../handlers/columnar_write/schema.rs | 4 +- .../control/checkpoint_durable_lsn.rs | 3 +- .../handlers/point/apply_put/index.rs | 142 +++- .../executor/handlers/spatial/full_scan.rs | 13 +- .../executor/handlers/spatial/rtree_scan.rs | 73 +- .../data/executor/handlers/spatial_refine.rs | 95 ++- .../data/executor/handlers/spatial_sync.rs | 67 +- .../transaction/overlay/spatial_merge.rs | 11 +- .../redo_apply/install_refusal_tests.rs | 6 +- .../handlers/transaction/resolve/entry.rs | 8 +- .../handlers/transaction/resolve/spatial.rs | 8 +- .../transaction/stage_write/stage_spatial.rs | 4 +- .../handlers/transaction/undo/entry.rs | 4 +- .../handlers/transaction/undo/spatial.rs | 34 +- .../handlers/transaction/undo/spatial_row.rs | 37 +- .../transaction/write_version_parity_tests.rs | 6 +- nodedb/src/data/executor/snapshot/capture.rs | 8 +- .../data/executor/spatial_checkpoint/load.rs | 1 + .../data/executor/spatial_checkpoint/write.rs | 10 +- .../src/data/executor/wal_replay_spatial.rs | 12 +- nodedb/src/engine/spatial/mod.rs | 1 + nodedb/src/error/types.rs | 14 + nodedb/src/error_classify/public.rs | 1 + nodedb/src/error_classify/unclassified.rs | 1 + nodedb/src/error_from.rs | 21 +- nodedb/src/error_from_eval.rs | 59 ++ nodedb/tests/wire/cases/spatial_crs.rs | 201 ++++++ 124 files changed, 5697 insertions(+), 1513 deletions(-) create mode 100644 nodedb-query/src/geo_functions/crs.rs create mode 100644 nodedb-spatial/src/crs/ellipsoid.rs create mode 100644 nodedb-spatial/src/crs/error.rs create mode 100644 nodedb-spatial/src/crs/hom.rs create mode 100644 nodedb-spatial/src/crs/mod.rs create mode 100644 nodedb-spatial/src/crs/params.rs create mode 100644 nodedb-spatial/src/crs/tm.rs create mode 100644 nodedb-spatial/src/crs/transform.rs create mode 100644 nodedb-spatial/src/planar/area.rs create mode 100644 nodedb-spatial/src/planar/buffer.rs create mode 100644 nodedb-spatial/src/planar/distance.rs create mode 100644 nodedb-spatial/src/planar/mod.rs delete mode 100644 nodedb-spatial/src/wkb.rs create mode 100644 nodedb-spatial/src/wkb/ewkb.rs create mode 100644 nodedb-spatial/src/wkb/format.rs create mode 100644 nodedb-spatial/src/wkb/mod.rs create mode 100644 nodedb-spatial/src/wkb/read.rs create mode 100644 nodedb-spatial/src/wkb/write.rs delete mode 100644 nodedb-types/src/geometry.rs create mode 100644 nodedb-types/src/geometry/math.rs create mode 100644 nodedb-types/src/geometry/mod.rs create mode 100644 nodedb-types/src/geometry/shape.rs create mode 100644 nodedb-types/src/geometry/srid.rs create mode 100644 nodedb-types/src/geometry/tagged.rs create mode 100644 nodedb/src/error_from_eval.rs create mode 100644 nodedb/tests/wire/cases/spatial_crs.rs 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-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-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-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/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/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-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/types/plan/variants/plan.rs b/nodedb-sql/src/types/plan/variants/plan.rs index 12369b8a9..dd883d35f 100644 --- a/nodedb-sql/src/types/plan/variants/plan.rs +++ b/nodedb-sql/src/types/plan/variants/plan.rs @@ -296,7 +296,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, diff --git a/nodedb-sql/src/visitor/plan_visitor/args.rs b/nodedb-sql/src/visitor/plan_visitor/args.rs index 23c31efac..059af432a 100644 --- a/nodedb-sql/src/visitor/plan_visitor/args.rs +++ b/nodedb-sql/src/visitor/plan_visitor/args.rs @@ -214,7 +214,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, 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-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/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/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/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/bridge/envelope/error_code_from.rs b/nodedb/src/bridge/envelope/error_code_from.rs index a15b5be85..71b9c13a6 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. diff --git a/nodedb/src/control/cluster/array_executor/refusal.rs b/nodedb/src/control/cluster/array_executor/refusal.rs index b14d384af..4a98d3fae 100644 --- a/nodedb/src/control/cluster/array_executor/refusal.rs +++ b/nodedb/src/control/cluster/array_executor/refusal.rs @@ -110,6 +110,7 @@ pub(super) fn execution_error(context: &str, error: crate::Error) -> ClusterErro | 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/execution_error_wire.rs b/nodedb/src/control/cluster/execution_error_wire.rs index 2bfd23196..567de9d6b 100644 --- a/nodedb/src/control/cluster/execution_error_wire.rs +++ b/nodedb/src/control/cluster/execution_error_wire.rs @@ -125,6 +125,7 @@ pub(crate) fn execution_error_to_typed(err: crate::Error) -> TypedClusterError { | 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/wedge.rs b/nodedb/src/control/cluster/metadata_applier/wedge.rs index 761b40fea..6a1adda88 100644 --- a/nodedb/src/control/cluster/metadata_applier/wedge.rs +++ b/nodedb/src/control/cluster/metadata_applier/wedge.rs @@ -119,6 +119,7 @@ pub fn classify(error: &crate::Error) -> ApplyFailureClass { | 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/gateway/error_map/class_parity/error_index.rs b/nodedb/src/control/gateway/error_map/class_parity/error_index.rs index 7a3acc263..392098e35 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 = 125; /// 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,6 @@ pub(crate) fn error_variant_index(err: &crate::Error) -> usize { E::InvalidDatetimeFormat { .. } => 119, E::DatetimeFieldOverflow { .. } => 120, E::WrongObjectType { .. } => 123, + E::MixedSrid { .. } => 124, } } 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..37b1c1e36 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 @@ -214,6 +214,12 @@ pub(crate) fn error_samples() -> Vec { }, 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..5015d650b 100644 --- a/nodedb/src/control/gateway/error_map/resp.rs +++ b/nodedb/src/control/gateway/error_map/resp.rs @@ -95,6 +95,7 @@ impl GatewayErrorMap { | Error::UnknownStrictField { .. } | Error::DivisionByZero | Error::DataException { .. } + | Error::MixedSrid { .. } | Error::NumericValueOutOfRange { .. } | Error::InvalidTextRepresentation { .. } | Error::DatatypeMismatch { .. } 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..09cee8e75 100644 --- a/nodedb/src/control/planner/sql_plan_convert/scan_params.rs +++ b/nodedb/src/control/planner/sql_plan_convert/scan_params.rs @@ -185,7 +185,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/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/pgwire/types/error_map.rs b/nodedb/src/control/server/pgwire/types/error_map.rs index a80267c07..a16efb6eb 100644 --- a/nodedb/src/control/server/pgwire/types/error_map.rs +++ b/nodedb/src/control/server/pgwire/types/error_map.rs @@ -125,6 +125,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/shared/retry.rs b/nodedb/src/control/server/shared/retry.rs index bb6b7967f..e2e636742 100644 --- a/nodedb/src/control/server/shared/retry.rs +++ b/nodedb/src/control/server/shared/retry.rs @@ -122,6 +122,7 @@ impl RetryableSchemaChange for Error { | Error::UnknownStrictField { .. } | Error::DivisionByZero | Error::DataException { .. } + | Error::MixedSrid { .. } | Error::NumericValueOutOfRange { .. } | Error::InvalidTextRepresentation { .. } | Error::DatatypeMismatch { .. } diff --git a/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs index f3e21378f..91e1907aa 100644 --- a/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs +++ b/nodedb/src/control/server/shared/write_admission/predicate/txn_buffering/classify.rs @@ -434,7 +434,7 @@ mod tests { use nodedb_types::calvin::{ EngineKeySet, EngineTag, PassiveReadKey, ReadKeyIdent, SortedVec, VersionedReadEntry, }; - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::id::{RequestId, TxnId, VShardId}; use nodedb_types::timeseries::continuous_agg::{ContinuousAggregateDef, RefreshPolicy}; use nodedb_types::vector_distance::DistanceMetric; @@ -1584,7 +1584,7 @@ mod tests { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "c"), field: "f".into(), surrogate: Surrogate::new(1), - geometry: Geometry::point(0.0, 0.0), + geometry: SridGeometry::wgs84(Geometry::point(0.0, 0.0)), provenance: None, }), PhysicalPlan::Spatial(SpatialOp::Delete { @@ -1597,7 +1597,7 @@ mod tests { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "c"), field: "f".into(), predicate: SpatialPredicate::Contains, - query_geometry: Geometry::point(0.0, 0.0), + query_geometry: SridGeometry::wgs84(Geometry::point(0.0, 0.0)), distance_meters: 0.0, attribute_filters: Vec::new(), limit: 0, diff --git a/nodedb/src/control/server/sync/async_dispatch/delta/compensation.rs b/nodedb/src/control/server/sync/async_dispatch/delta/compensation.rs index 0af19406e..8e7a854d6 100644 --- a/nodedb/src/control/server/sync/async_dispatch/delta/compensation.rs +++ b/nodedb/src/control/server/sync/async_dispatch/delta/compensation.rs @@ -117,6 +117,7 @@ pub(super) fn compensation_hint_for_dispatch_error(e: &crate::Error) -> Compensa | 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/server/sync/refusal.rs b/nodedb/src/control/server/sync/refusal.rs index 4b5289af6..8a1beacbf 100644 --- a/nodedb/src/control/server/sync/refusal.rs +++ b/nodedb/src/control/server/sync/refusal.rs @@ -133,6 +133,7 @@ pub(super) fn retryable_refusal_reason(error: &crate::Error) -> Option<&str> { | crate::Error::UnknownStrictField { .. } | crate::Error::DivisionByZero | crate::Error::DataException { .. } + | crate::Error::MixedSrid { .. } | crate::Error::NumericValueOutOfRange { .. } | crate::Error::InvalidTextRepresentation { .. } | crate::Error::DatatypeMismatch { .. } @@ -299,6 +300,7 @@ fn is_indeterminate(error: &crate::Error) -> bool { | 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/server/sync/spatial_handler.rs b/nodedb/src/control/server/sync/spatial_handler.rs index 0a2ff98b7..d9aa8665f 100644 --- a/nodedb/src/control/server/sync/spatial_handler.rs +++ b/nodedb/src/control/server/sync/spatial_handler.rs @@ -15,7 +15,7 @@ use async_trait::async_trait; use nodedb_types::Surrogate; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use crate::control::server::dispatch_utils::RecordOwner; use crate::types::{DatabaseId, TenantId, VShardId}; @@ -32,7 +32,7 @@ pub struct SpatialInsertTarget { pub collection: String, pub field: String, pub surrogate: Surrogate, - pub geometry: Geometry, + pub geometry: SridGeometry, } /// Encapsulates async Data Plane dispatch for spatial insert/delete. @@ -380,7 +380,7 @@ mod tests { } fn make_point_geometry_bytes() -> Vec { - let geom = nodedb_types::geometry::Geometry::point(10.0, 20.0); + let geom = SridGeometry::wgs84(nodedb_types::geometry::Geometry::point(10.0, 20.0)); zerompk::to_msgpack_vec(&geom).unwrap() } diff --git a/nodedb/src/control/server/sync/spatial_session.rs b/nodedb/src/control/server/sync/spatial_session.rs index 489fa348f..83f0d18ed 100644 --- a/nodedb/src/control/server/sync/spatial_session.rs +++ b/nodedb/src/control/server/sync/spatial_session.rs @@ -7,7 +7,7 @@ use tracing::{debug, error}; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use nodedb_types::sync::wire::AckStatus; use super::session::SyncSession; @@ -42,8 +42,8 @@ impl SyncSession { return SyncFrame::try_encode(SyncMessageType::SpatialInsertAck, &ack); } - // Deserialise the geometry from MessagePack bytes. - let geometry: Geometry = match zerompk::from_msgpack(&msg.geometry_bytes) { + // Deserialise the geometry and its SRID from MessagePack bytes. + let geometry: SridGeometry = match zerompk::from_msgpack(&msg.geometry_bytes) { Ok(g) => g, Err(e) => { error!( diff --git a/nodedb/src/control/server/wal_dispatch/core.rs b/nodedb/src/control/server/wal_dispatch/core.rs index 4fe3a1be7..138deee83 100644 --- a/nodedb/src/control/server/wal_dispatch/core.rs +++ b/nodedb/src/control/server/wal_dispatch/core.rs @@ -217,7 +217,7 @@ mod tests { use nodedb_physical::physical_plan::{SpatialOp, TextOp}; use nodedb_types::QualifiedCollection; use nodedb_types::Surrogate; - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, SridGeometry}; fn open_wal(dir: &std::path::Path) -> WalManager { WalManager::open_for_testing(&dir.join("test.wal")).expect("open wal") @@ -361,7 +361,7 @@ mod tests { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "places"), field: "loc".to_string(), surrogate: Surrogate::new(9), - geometry: Geometry::point(10.0, 20.0), + geometry: SridGeometry::wgs84(Geometry::point(10.0, 20.0)), provenance: None, }); @@ -383,9 +383,9 @@ mod tests { nodedb_wal::record::SpatialPutPayload::from_bytes(&record.payload).expect("decode"); assert_eq!(decoded.collection, "places"); assert_eq!(decoded.field, "loc"); - let geometry: Geometry = + let geometry: SridGeometry = zerompk::from_msgpack(&decoded.geometry_bytes).expect("decode geometry"); - assert_eq!(geometry, Geometry::point(10.0, 20.0)); + assert_eq!(geometry, SridGeometry::wgs84(Geometry::point(10.0, 20.0))); } #[test] @@ -443,7 +443,7 @@ mod tests { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "docs"), field: "loc".to_string(), surrogate: Surrogate::new(1), - geometry: Geometry::point(1.0, 2.0), + geometry: SridGeometry::wgs84(Geometry::point(1.0, 2.0)), provenance: None, }); wal_append_if_write( diff --git a/nodedb/src/control/server/wal_dispatch_fts_spatial.rs b/nodedb/src/control/server/wal_dispatch_fts_spatial.rs index 7f2847052..15772b96f 100644 --- a/nodedb/src/control/server/wal_dispatch_fts_spatial.rs +++ b/nodedb/src/control/server/wal_dispatch_fts_spatial.rs @@ -7,14 +7,14 @@ //! runs Data-Plane-side at the apply handler. use nodedb_types::Surrogate; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use nodedb_types::sync::wire::SyncProvenance; use crate::types::{DatabaseId, TenantId, VShardId}; use crate::wal::manager::WalAppender; /// Build a `SpatialPutPayload` from raw op fields, msgpack-encoding the -/// geometry. +/// geometry with its SRID. /// /// This is the ONE builder for the shape: both the autocommit WAL append /// and the transaction-resolve serializer call it so @@ -25,7 +25,7 @@ pub(crate) fn encode_spatial_put_payload( collection: &str, field: &str, surrogate: Surrogate, - geometry: &Geometry, + geometry: &SridGeometry, provenance: &SyncProvenance, ) -> crate::Result { let doc_id = crate::engine::document::store::StorageKey::for_surrogate(surrogate).to_string(); @@ -59,7 +59,7 @@ pub(crate) fn encode_spatial_delete_payload( /// Append a spatial put (insert) to the WAL and return the assigned LSN. /// -/// The `payload` carries provenance and the msgpack-encoded `Geometry` +/// The `payload` carries provenance and the msgpack-encoded `SridGeometry` /// (identical to what `SpatialInsertMsg.geometry_bytes` carries). pub fn wal_append_spatial_put( wal: WalAppender<'_>, diff --git a/nodedb/src/control/wal_replication/decode_sync_engines.rs b/nodedb/src/control/wal_replication/decode_sync_engines.rs index f746e7ffc..b8cacb94e 100644 --- a/nodedb/src/control/wal_replication/decode_sync_engines.rs +++ b/nodedb/src/control/wal_replication/decode_sync_engines.rs @@ -152,7 +152,7 @@ pub fn spatial_insert( geometry_bytes: &[u8], prov_bytes: &Option>, ) -> crate::Result { - let geometry = zerompk::from_msgpack::(geometry_bytes) + let geometry = zerompk::from_msgpack::(geometry_bytes) .map_err(|e| crate::Error::Internal { detail: format!("SpatialInsert geometry decode failed: {e}"), })?; @@ -706,9 +706,12 @@ mod tests { stream_id: 2, seq: 44, }; - let geometry = Geometry::Point { - coordinates: [-73.985, 40.758], - }; + let geometry = nodedb_types::geometry::SridGeometry::new( + nodedb_types::geometry::Srid::PeninsularRso, + Geometry::Point { + coordinates: [412_834.0, 349_472.0], + }, + ); let plan = PhysicalPlan::Spatial(SpatialOp::Insert { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "places"), diff --git a/nodedb/src/control/wal_replication/encode/columnar.rs b/nodedb/src/control/wal_replication/encode/columnar.rs index 42dbe7a86..5927a136d 100644 --- a/nodedb/src/control/wal_replication/encode/columnar.rs +++ b/nodedb/src/control/wal_replication/encode/columnar.rs @@ -94,7 +94,7 @@ pub(super) fn spatial_insert( collection: &str, field: &str, surrogate: u32, - geometry: &nodedb_types::geometry::Geometry, + geometry: &nodedb_types::geometry::SridGeometry, provenance: Option>, ) -> crate::Result { // A failed encode refuses the write rather than replicate bytes that diff --git a/nodedb/src/data/executor/columnar_checkpoint/geometry_restore.rs b/nodedb/src/data/executor/columnar_checkpoint/geometry_restore.rs index 94912b2e4..49d5b63d5 100644 --- a/nodedb/src/data/executor/columnar_checkpoint/geometry_restore.rs +++ b/nodedb/src/data/executor/columnar_checkpoint/geometry_restore.rs @@ -145,8 +145,10 @@ impl CoreLoop { ); let indexed = docs.len(); - // Boot-time rebuild: nothing to roll back, so the delta is dropped. - let _ = self.index_columnar_geometry_columns(&task, &schema, collection, &docs); + // Boot-time rebuild: nothing to roll back, so the delta is dropped. A + // restored row set in mixed SRIDs refuses the restore, as an + // unreadable row does. + self.index_columnar_geometry_columns(&task, &schema, collection, &docs)?; Ok(indexed) } diff --git a/nodedb/src/data/executor/handlers/columnar_mutation_apply.rs b/nodedb/src/data/executor/handlers/columnar_mutation_apply.rs index d5d3c43b7..19f21cec1 100644 --- a/nodedb/src/data/executor/handlers/columnar_mutation_apply.rs +++ b/nodedb/src/data/executor/handlers/columnar_mutation_apply.rs @@ -99,6 +99,15 @@ impl CoreLoop { // names the R-tree entries the old row owns. let old_pks = rows.iter().map(|r| &r.0); let pre_images = self.read_geometry_pre_images(key, has_geometry, old_pks)?; + // Every post-image geometry's SRID is checked before the first row + // changes, so a mixed-SRID statement changes nothing. + if has_geometry { + let post_images: Vec = rows + .iter() + .map(|(_, new_row)| row_values_to_object(schema, new_row)) + .collect(); + self.check_columnar_geometry_srids(task, schema, collection, &post_images)?; + } for ((old_pk, new_row), pre_image) in rows.iter().zip(pre_images) { let old_pk_bytes = encode_pk(old_pk); @@ -167,12 +176,27 @@ impl CoreLoop { key.0, key.1, collection, schema, row, )); } - let delta = self.index_columnar_geometry_columns( + let delta = match self.index_columnar_geometry_columns( task, schema, collection, &[row_values_to_object(schema, new_row)], - ); + ) { + Ok(delta) => delta, + Err(error) => { + let mut undo = Vec::new(); + push_removed_spatial_undo(&mut undo, spatial_removed); + Self::push_geometry_index_undo(&mut undo, geometry_delta); + undo.push(UndoEntry::ColumnarUpdate { + collection_key: key.clone(), + row_count_before, + inserted_pks: outcome.inserted_pks, + displaced: outcome.displaced, + restored: outcome.restored, + }); + return Err(self.reverse_columnar_statement(key, undo, error)); + } + }; geometry_delta.removed.extend(delta.removed); geometry_delta.inserted.extend(delta.inserted); } diff --git a/nodedb/src/data/executor/handlers/columnar_write/geometry_index.rs b/nodedb/src/data/executor/handlers/columnar_write/geometry_index.rs index 07afe4776..5163414eb 100644 --- a/nodedb/src/data/executor/handlers/columnar_write/geometry_index.rs +++ b/nodedb/src/data/executor/handlers/columnar_write/geometry_index.rs @@ -4,6 +4,7 @@ //! spatial predicates can find newly-inserted rows. use nodedb_types::columnar::{ColumnType, ColumnarSchema}; +use nodedb_types::geometry::{Srid, SridGeometry}; use nodedb_types::value::Value; use crate::data::executor::core_loop::CoreLoop; @@ -22,10 +23,74 @@ pub(in crate::data::executor) struct GeometryIndexDelta { pub inserted: Vec<(SpatialIndexKey, u64)>, } +/// The geometry a columnar cell holds: a native geometry or GeoJSON text. +/// Any other cell holds none. +fn cell_geometry(value: &Value) -> Option { + match value { + Value::Geometry(g) => Some(g.clone()), + Value::String(s) => nodedb_types::geometry::from_geojson_str(s), + _ => None, + } +} + impl CoreLoop { + /// Check that every geometry cell of `ndb_rows` can join its column's + /// R-tree. Changes nothing. + /// + /// A non-empty index admits only its own SRID. An empty or absent index + /// takes the SRID of the batch's first geometry, and every later + /// geometry of the column must match it. The first refusal returns + /// `Error::MixedSrid`. + pub(in crate::data::executor) fn check_columnar_geometry_srids( + &self, + task: &ExecutionTask, + schema: &ColumnarSchema, + collection: &str, + ndb_rows: &[nodedb_types::Value], + ) -> crate::Result<()> { + let db_id = task.request.database_id; + let tid = task.request.tenant_id; + for col_def in schema + .columns + .iter() + .filter(|c| c.column_type == ColumnType::Geometry) + { + let index_key = (db_id, tid, collection.to_string(), col_def.name.clone()); + let mut fixed: Option = self + .spatial_indexes + .get(&index_key) + .filter(|rtree| !rtree.is_empty()) + .map(|rtree| rtree.srid()); + for row in ndb_rows { + let Value::Object(obj) = row else { + continue; + }; + let Some(geom) = obj.get(&col_def.name).and_then(cell_geometry) else { + continue; + }; + match fixed { + Some(srid) if srid != geom.srid => { + return Err(crate::Error::MixedSrid { + collection: collection.to_string(), + field: col_def.name.clone(), + index_srid: srid.code(), + row_srid: geom.srid.code(), + }); + } + Some(_) => {} + None => fixed = Some(geom.srid), + } + } + } + Ok(()) + } + /// Populate the R-tree for any `Geometry` columns on `schema`, from the /// rows just inserted (`ndb_rows`). Returns what it changed. /// + /// Runs [`Self::check_columnar_geometry_srids`] first, so a mixed-SRID + /// batch fails with `Error::MixedSrid` before any index changes. + /// /// `schema` is threaded in from the caller rather than re-read from /// `columnar_engines` here: the caller already resolved it, and taking it /// as a parameter means there is no "engine missing" case to invent an @@ -38,7 +103,7 @@ impl CoreLoop { schema: &ColumnarSchema, collection: &str, ndb_rows: &[nodedb_types::Value], - ) -> GeometryIndexDelta { + ) -> crate::Result { let mut delta = GeometryIndexDelta::default(); let db_id = task.request.database_id; let tid = task.request.tenant_id; @@ -51,8 +116,9 @@ impl CoreLoop { .collect(); if geom_cols.is_empty() { - return delta; + return Ok(delta); } + self.check_columnar_geometry_srids(task, schema, collection, ndb_rows)?; for row in ndb_rows { let obj = match row { nodedb_types::Value::Object(m) => m, @@ -91,15 +157,10 @@ impl CoreLoop { // See `nodedb_types::geometry::from_geojson_str` — shared with the // document index path (`apply_point_put_spatial`) and the read path // (`extract_geometry`); keep all three in sync. - let geom: nodedb_types::geometry::Geometry = match field_val { - Value::Geometry(g) => g.clone(), - Value::String(s) => match nodedb_types::geometry::from_geojson_str(s) { - Some(g) => g, - None => continue, - }, - _ => continue, + let Some(geom) = cell_geometry(field_val) else { + continue; }; - let bbox = nodedb_types::bbox::geometry_bbox(&geom); + let bbox = geom.bbox(); let index_key = (db_id, tid, collection.to_string(), col_def.name.clone()); let entry_id = spatial_entry_id.as_u64(); let memory = nodedb_mem::ScopedMemory::new( @@ -111,7 +172,17 @@ impl CoreLoop { let rtree = self .spatial_indexes .entry(index_key.clone()) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)); + .or_insert_with(|| crate::engine::spatial::RTree::new(memory, geom.srid)); + // The batch check above admitted every SRID here; removing a + // re-indexed row's prior entry can only empty the tree. + rtree + .bind_srid(geom.srid) + .map_err(|m| crate::Error::MixedSrid { + collection: collection.to_string(), + field: col_def.name.clone(), + index_srid: m.index.code(), + row_srid: m.entry.code(), + })?; rtree.insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); self.spatial_doc_map.insert( ( @@ -126,7 +197,7 @@ impl CoreLoop { delta.inserted.push((index_key, entry_id)); } } - delta + Ok(delta) } /// Record a geometry-index pass on `undo_log` so a rollback reverses @@ -138,12 +209,13 @@ impl CoreLoop { delta: GeometryIndexDelta, ) { use crate::data::executor::handlers::transaction::undo::UndoEntry; - for (key, entry_id, bbox, document_id) in delta.removed { + for (key, entry_id, bbox, document_id, srid) in delta.removed { undo_log.push(UndoEntry::SpatialDelete { key, entry_id, bbox, document_id, + srid, }); } for (key, entry_id) in delta.inserted { diff --git a/nodedb/src/data/executor/handlers/columnar_write/geometry_remove.rs b/nodedb/src/data/executor/handlers/columnar_write/geometry_remove.rs index b57bef9c8..2f3c877cc 100644 --- a/nodedb/src/data/executor/handlers/columnar_write/geometry_remove.rs +++ b/nodedb/src/data/executor/handlers/columnar_write/geometry_remove.rs @@ -12,9 +12,15 @@ use crate::data::executor::handlers::point::apply_put::SpatialEntryId; use crate::data::executor::spatial_key::SpatialIndexKey; /// One R-tree entry the cascade removed: enough to re-insert it on undo -/// (`UndoEntry::SpatialDelete`). -pub(in crate::data::executor) type RemovedSpatialEntry = - (SpatialIndexKey, u64, nodedb_types::BoundingBox, String); +/// (`UndoEntry::SpatialDelete`): the index key, entry id, bbox, reverse-map +/// document id, and the SRID of the tree it left. +pub(in crate::data::executor) type RemovedSpatialEntry = ( + SpatialIndexKey, + u64, + nodedb_types::BoundingBox, + String, + nodedb_types::geometry::Srid, +); /// Whether `schema` declares at least one `Geometry` column, i.e. whether /// the collection's rows can carry R-tree entries at all. diff --git a/nodedb/src/data/executor/handlers/columnar_write/insert.rs b/nodedb/src/data/executor/handlers/columnar_write/insert.rs index b992b8539..d3c1f1dac 100644 --- a/nodedb/src/data/executor/handlers/columnar_write/insert.rs +++ b/nodedb/src/data/executor/handlers/columnar_write/insert.rs @@ -147,6 +147,12 @@ impl CoreLoop { Err(e) => return self.response_error(task, ErrorCode::from(e)), }; + // A geometry in an SRID its column's R-tree refuses fails the insert + // before any row is written. + if let Err(e) = self.check_columnar_geometry_srids(task, &schema, collection, &ndb_rows) { + return self.response_error(task, ErrorCode::from(e)); + } + let outcome = match self.insert_columnar_rows( task, RowIngestParams { @@ -172,8 +178,13 @@ impl CoreLoop { } // Populate R-tree for geometry columns so spatial predicates work. + // The SRID check above ran on the same rows, so this changes the + // index or fails before touching it. let geometry_delta = - self.index_columnar_geometry_columns(task, &schema, collection, &ndb_rows); + match self.index_columnar_geometry_columns(task, &schema, collection, &ndb_rows) { + Ok(delta) => delta, + Err(e) => return self.response_error(task, ErrorCode::from(e)), + }; if let Some(undo_log) = spatial_undo { Self::push_geometry_index_undo(undo_log, geometry_delta); } diff --git a/nodedb/src/data/executor/handlers/columnar_write/schema.rs b/nodedb/src/data/executor/handlers/columnar_write/schema.rs index 5f2362236..df6232811 100644 --- a/nodedb/src/data/executor/handlers/columnar_write/schema.rs +++ b/nodedb/src/data/executor/handlers/columnar_write/schema.rs @@ -328,7 +328,9 @@ mod tests { ("id".to_string(), Value::String("a".into())), ( "geom".to_string(), - Value::Geometry(nodedb_types::geometry::Geometry::point(1.0, 2.0)), + Value::Geometry(nodedb_types::geometry::SridGeometry::wgs84( + nodedb_types::geometry::Geometry::point(1.0, 2.0), + )), ), ( "emb".to_string(), diff --git a/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs b/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs index d0316977f..c0fda2c6f 100644 --- a/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs +++ b/nodedb/src/data/executor/handlers/control/checkpoint_durable_lsn.rs @@ -809,7 +809,8 @@ mod tests { nodedb_types::TenantId::new(1), nodedb_mem::EngineId::Spatial, ); - let mut rtree = crate::engine::spatial::RTree::new(memory); + let mut rtree = + crate::engine::spatial::RTree::new(memory, nodedb_types::geometry::Srid::Wgs84); rtree.insert(crate::engine::spatial::RTreeEntry { id: 1, bbox: nodedb_types::BoundingBox::new(0.0, 0.0, 1.0, 1.0), diff --git a/nodedb/src/data/executor/handlers/point/apply_put/index.rs b/nodedb/src/data/executor/handlers/point/apply_put/index.rs index a1ddb1f07..a7785db76 100644 --- a/nodedb/src/data/executor/handlers/point/apply_put/index.rs +++ b/nodedb/src/data/executor/handlers/point/apply_put/index.rs @@ -11,8 +11,8 @@ use crate::data::executor::core_loop::CoreLoop; use crate::data::executor::doc_format; +use crate::data::executor::handlers::columnar_write::RemovedSpatialEntry; use crate::data::executor::handlers::transaction::undo::UndoEntry; -use crate::data::executor::spatial_key::SpatialIndexKey; use crate::engine::document::store::StorageKey; /// Entry id of one row in a spatial R-tree: `fnv1a_hash` of the row's identity text. @@ -55,6 +55,11 @@ impl CoreLoop { /// of the document it removes, then a `SpatialInsert` for each entry it /// inserts. Rollback runs the log in reverse, so it removes the new /// entries before it puts the prior ones back. + /// + /// Every geometry's SRID is checked against its field's R-tree before + /// anything changes. A geometry in an SRID other than a non-empty + /// index's fails the put with `Error::MixedSrid` and leaves every index + /// untouched; the caller's abort path then drops the row write. pub(in crate::data::executor) fn apply_point_put_spatial( &mut self, database_id: u64, @@ -63,7 +68,7 @@ impl CoreLoop { storage_key: StorageKey, value: &[u8], undo: &mut Vec, - ) { + ) -> crate::Result<()> { // Rendered once here; every reverse-map / hash use below shares it. let document_id = storage_key.to_string(); let document_id = document_id.as_str(); @@ -92,7 +97,7 @@ impl CoreLoop { for (field_name, field_value) in obj { let parsed_geom = match field_value { serde_json::Value::String(s) => nodedb_types::geometry::from_geojson_str(s), - _ => serde_json::from_value::( + _ => serde_json::from_value::( field_value.clone(), ) .ok(), @@ -103,6 +108,23 @@ impl CoreLoop { } } + let db_id = nodedb_types::DatabaseId::new(database_id); + let tid_id = crate::types::TenantId::new(tid); + // Check every SRID before any mutation: a refused put changes no index. + for (field_name, geom) in &geometries { + let spatial_key = (db_id, tid_id, collection.to_string(), field_name.clone()); + if let Some(rtree) = self.spatial_indexes.get(&spatial_key) + && !rtree.accepts_srid(geom.srid) + { + return Err(crate::Error::MixedSrid { + collection: collection.to_string(), + field: field_name.clone(), + index_srid: rtree.srid().code(), + row_srid: geom.srid.code(), + }); + } + } + // Re-indexing a document must REPLACE, not append: `RTree::insert` // blindly pushes a fresh entry even when one with this `entry_id` // already exists, so a live geometry UPDATE, a WAL replay, or the @@ -116,11 +138,9 @@ impl CoreLoop { spatial_entry_id, undo, ); - let db_id = nodedb_types::DatabaseId::new(database_id); - let tid_id = crate::types::TenantId::new(tid); let entry_id = spatial_entry_id.as_u64(); for (field_name, geom) in geometries { - let bbox = nodedb_types::bbox::geometry_bbox(&geom); + let bbox = geom.bbox(); let spatial_key = (db_id, tid_id, collection.to_string(), field_name.clone()); let memory = nodedb_mem::ScopedMemory::new( self.governor.clone(), @@ -131,7 +151,17 @@ impl CoreLoop { let rtree = self .spatial_indexes .entry(spatial_key.clone()) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)); + .or_insert_with(|| crate::engine::spatial::RTree::new(memory, geom.srid)); + // The pre-check above admitted this SRID. The prior-entry removal + // can only empty the tree, so binding cannot refuse here. + rtree + .bind_srid(geom.srid) + .map_err(|m| crate::Error::MixedSrid { + collection: collection.to_string(), + field: field_name.clone(), + index_srid: m.index.code(), + row_srid: m.entry.code(), + })?; rtree.insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); // Maintain reverse map: entry_id → document_id. self.spatial_doc_map.insert( @@ -143,6 +173,7 @@ impl CoreLoop { entry_id, }); } + Ok(()) } /// [`Self::remove_document_spatial_indexes`], pushing a `SpatialDelete` @@ -156,12 +187,13 @@ impl CoreLoop { undo: &mut Vec, ) { let removed = self.remove_document_spatial_indexes(database_id, tid, collection, entry_id); - for (key, entry_id, bbox, document_id) in removed { + for (key, entry_id, bbox, document_id, srid) in removed { undo.push(UndoEntry::SpatialDelete { key, entry_id, bbox, document_id, + srid, }); } } @@ -178,15 +210,15 @@ impl CoreLoop { /// removed geometry) so a transactional caller can push /// `UndoEntry::SpatialDelete` re-insert reversals — the reverse /// `spatial_doc_map` stores only the doc id. Returns the removed - /// `(spatial_index_key, entry_id, bbox, document_id)` tuples; empty when the - /// document had no spatial fields. + /// `(spatial_index_key, entry_id, bbox, document_id, srid)` tuples, `srid` + /// being the tree's; empty when the document had no spatial fields. pub(in crate::data::executor) fn remove_document_spatial_indexes( &mut self, database_id: u64, tid: u64, collection: &str, entry_id: SpatialEntryId, - ) -> Vec<(SpatialIndexKey, u64, nodedb_types::BoundingBox, String)> { + ) -> Vec { let mut spatial_deletes = Vec::new(); let entry_id = entry_id.as_u64(); let db_id = nodedb_types::DatabaseId::new(database_id); @@ -202,11 +234,13 @@ impl CoreLoop { // Read the bbox BEFORE deleting — the R-tree `delete` does not // return the removed geometry, so a reversible undo must capture // it here (the reverse `spatial_doc_map` stores only the doc id). - let bbox = self - .spatial_indexes - .get(&skey) - .and_then(|rtree| rtree.entries().into_iter().find(|e| e.id == entry_id)) - .map(|e| e.bbox); + let found = self.spatial_indexes.get(&skey).and_then(|rtree| { + rtree + .entries() + .into_iter() + .find(|e| e.id == entry_id) + .map(|e| (e.bbox, rtree.srid())) + }); if let Some(rtree) = self.spatial_indexes.get_mut(&skey) { rtree.delete(entry_id); } @@ -217,8 +251,8 @@ impl CoreLoop { field, entry_id, )); - if let (Some(bbox), Some(doc)) = (bbox, removed_doc) { - spatial_deletes.push((skey, entry_id, bbox, doc)); + if let (Some((bbox, srid)), Some(doc)) = (found, removed_doc) { + spatial_deletes.push((skey, entry_id, bbox, doc, srid)); } } spatial_deletes @@ -423,6 +457,78 @@ mod tests { assert_eq!(scanned.len(), 4, "every document must be scanned"); } + /// The index SRID is fixed by its first entry. A later row in another + /// SRID is refused before anything changes: the R-tree keeps its single + /// entry and no row is stored, for a new row and for an overwrite alike. + #[test] + fn a_row_in_another_srid_is_refused_and_changes_nothing() { + use crate::engine::document::store::StorageKey; + + let dir = tempfile::tempdir().unwrap(); + let (mut core, _tx, _rx) = make_core_with_dir(dir.path()); + let coll = "geo_srid"; + let put = |core: &mut crate::data::executor::core_loop::CoreLoop, + id: &str, + surrogate: u32, + doc: &[u8]| { + let task = point_put_task(coll, id, doc); + core.execute_point_put( + &task, + PointPutExec { + tid: 1, + collection: coll, + document_id: id, + surrogate: Surrogate::new(surrogate), + value: doc, + returning: None, + rls_filters: &[], + resolved_sum_targets: &[], + }, + ) + }; + + let wgs84 = br#"{"loc":{"type":"Point","coordinates":[101.7,3.1]}}"#; + assert_eq!(put(&mut core, "a", 1, wgs84).status, Status::Ok); + + let utm = br#"{"loc":{"type":"Point","coordinates":[134445.5,349624.7],"crs":{"type":"name","properties":{"name":"EPSG:32648"}}}}"#; + let key = ( + DatabaseId::DEFAULT, + TenantId::new(1), + coll.to_string(), + "loc".to_string(), + ); + let before = core.spatial_indexes.get(&key).unwrap().entries()[0].bbox; + + // A new row in another SRID. + assert_eq!(put(&mut core, "b", 2, utm).status, Status::Error); + // An overwrite of the first row into another SRID. + assert_eq!(put(&mut core, "a", 1, utm).status, Status::Error); + + let rtree = core.spatial_indexes.get(&key).unwrap(); + assert_eq!(rtree.srid(), nodedb_types::geometry::Srid::Wgs84); + let entries = rtree.entries(); + assert_eq!(entries.len(), 1, "the refused puts added no entry"); + assert_eq!(entries[0].bbox, before, "the first entry is unchanged"); + + let db = DatabaseId::DEFAULT.as_u64(); + let stored_b = core + .sparse + .get(db, 1, coll, &StorageKey::for_surrogate(Surrogate::new(2))) + .unwrap(); + assert!(stored_b.is_none(), "the refused new row was not stored"); + let stored_a = core + .sparse + .get(db, 1, coll, &StorageKey::for_surrogate(Surrogate::new(1))) + .unwrap() + .expect("the first row stays stored"); + let doc = crate::data::executor::doc_format::decode_document_value(&stored_a).unwrap(); + let geom = nodedb_types::geometry::SridGeometry::from_value( + doc.get("loc").expect("geometry field"), + ) + .expect("stored geometry"); + assert_eq!(geom.srid, nodedb_types::geometry::Srid::Wgs84); + } + /// A put that overwrites a geometry row and then aborts puts the row's /// prior R-tree entry and its `spatial_doc_map` record back. #[test] diff --git a/nodedb/src/data/executor/handlers/spatial/full_scan.rs b/nodedb/src/data/executor/handlers/spatial/full_scan.rs index fd60c3e7f..1861d7dd1 100644 --- a/nodedb/src/data/executor/handlers/spatial/full_scan.rs +++ b/nodedb/src/data/executor/handlers/spatial/full_scan.rs @@ -7,7 +7,7 @@ use crate::bridge::scan_filter::ScanFilter; use crate::data::executor::core_loop::CoreLoop; use crate::data::executor::doc_format; use crate::data::executor::handlers::spatial_refine::{ - SpatialHit, apply_predicate, extract_geometry, hit_rows, project_doc, + SpatialHit, apply_predicate, extract_geometry, hit_rows, mixed_srid_error, project_doc, }; use crate::data::executor::handlers::transaction::overlay::SpatialOverlayMergeParams; use crate::data::executor::response_codec; @@ -24,7 +24,7 @@ pub(super) struct SpatialFullScanParams<'a> { pub collection: &'a str, pub field: &'a str, pub predicate: &'a SpatialPredicate, - pub query_geom: &'a nodedb_types::geometry::Geometry, + pub query_geom: &'a nodedb_types::geometry::SridGeometry, pub distance_meters: f64, pub limit: usize, pub projection: &'a [String], @@ -102,8 +102,13 @@ impl CoreLoop { None => continue, }; - if !apply_predicate(predicate, query_geom, &doc_geom, distance_meters) { - continue; + match apply_predicate(predicate, query_geom, &doc_geom, distance_meters) { + Ok(true) => {} + Ok(false) => continue, + Err(mismatch) => { + return self + .response_error(task, mixed_srid_error(collection, field, mismatch)); + } } match ScanFilter::all_match_value(attr_filters, &doc) { diff --git a/nodedb/src/data/executor/handlers/spatial/rtree_scan.rs b/nodedb/src/data/executor/handlers/spatial/rtree_scan.rs index c032bb7ac..498ca9a25 100644 --- a/nodedb/src/data/executor/handlers/spatial/rtree_scan.rs +++ b/nodedb/src/data/executor/handlers/spatial/rtree_scan.rs @@ -8,7 +8,8 @@ use crate::data::executor::core_loop::CoreLoop; use crate::data::executor::doc_format; use crate::data::executor::handlers::columnar_read::filter::decode_rls_filters; use crate::data::executor::handlers::spatial_refine::{ - SpatialHit, apply_predicate, expand_bbox, extract_geometry, hit_rows, project_doc, + SpatialHit, apply_predicate, expand_bbox, extract_geometry, hit_rows, mixed_srid_error, + project_doc, }; use crate::data::executor::handlers::transaction::overlay::SpatialOverlayMergeParams; use crate::data::executor::response_codec; @@ -26,7 +27,7 @@ pub(in crate::data::executor) struct SpatialScanParams<'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 [u8], pub limit: usize, @@ -100,10 +101,12 @@ impl CoreLoop { Err(e) => return self.response_error(task, e), }; - // 3. Compute search bbox (expand by distance for ST_DWithin). - let query_bbox = nodedb_types::bbox::geometry_bbox(query_geom); + // 3. Compute search bbox (expand by distance for ST_DWithin) in the + // query SRID's space. + let space = query_geom.srid.bbox_space(); + let query_bbox = query_geom.bbox(); let search_bbox = if distance_meters > 0.0 { - expand_bbox(&query_bbox, distance_meters) + expand_bbox(&query_bbox, distance_meters, space) } else { query_bbox }; @@ -112,6 +115,18 @@ impl CoreLoop { let tid_id = crate::types::TenantId::new(tid); let spatial_key = (db_id, tid_id, collection.to_string(), field.to_string()); let has_index = self.spatial_indexes.contains_key(&spatial_key); + + // A query in an SRID other than the index's is refused before the + // scan. No implicit transform runs. + if let Some(rtree) = self.spatial_indexes.get(&spatial_key) + && !rtree.accepts_srid(query_geom.srid) + { + let mismatch = nodedb_spatial::SridMismatch { + index: rtree.srid(), + entry: query_geom.srid, + }; + return self.response_error(task, mixed_srid_error(collection, field, mismatch)); + } // `limit` is the statement's LIMIT, `usize::MAX` without one. A // `LIMIT 0` returns no rows. // No R-tree: full scan with predicate post-filter. @@ -283,8 +298,13 @@ impl CoreLoop { None => continue, }; - if !apply_predicate(predicate, query_geom, &doc_geom, distance_meters) { - continue; + match apply_predicate(predicate, query_geom, &doc_geom, distance_meters) { + Ok(true) => {} + Ok(false) => continue, + Err(mismatch) => { + return self + .response_error(task, mixed_srid_error(collection, field, mismatch)); + } } match ScanFilter::all_match_value(&attr_filters, &doc) { @@ -438,7 +458,9 @@ mod tests { let rtree = core .spatial_indexes .entry(spatial_key.clone()) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)); + .or_insert_with(|| { + crate::engine::spatial::RTree::new(memory, nodedb_types::geometry::Srid::Wgs84) + }); rtree.insert(RTreeEntry { id: entry_id, bbox }); core.spatial_doc_map.insert( ( @@ -455,8 +477,10 @@ mod tests { } /// GeoJSON Point query centred on (0.0, 0.0) for DWithin. - fn origin_point() -> nodedb_types::geometry::Geometry { - nodedb_types::geometry::Geometry::point(0.0, 0.0) + fn origin_point() -> nodedb_types::geometry::SridGeometry { + nodedb_types::geometry::SridGeometry::wgs84(nodedb_types::geometry::Geometry::point( + 0.0, 0.0, + )) } fn dummy_spatial_plan() -> PhysicalPlan { @@ -519,4 +543,33 @@ mod tests { zerompk::from_msgpack(resp.payload.as_bytes()).unwrap_or_default(); assert!(decoded.is_empty(), "empty prefilter must return no results"); } + + /// A query in an SRID other than the index's is refused before the scan. + #[test] + fn query_in_another_srid_is_refused() { + let (mut core, _resp_rx, _dir) = make_core(); + let tid = 1u64; + insert_spatial_doc(&mut core, tid, "geo", "loc", Surrogate(10), 0.0, 0.0); + + let query = nodedb_types::geometry::SridGeometry::new( + nodedb_types::geometry::Srid::Utm48N, + nodedb_types::geometry::Geometry::point(500_000.0, 0.0), + ); + let task = make_task(dummy_spatial_plan()); + let resp = core.execute_spatial_scan(SpatialScanParams { + task: &task, + tid, + collection: "geo", + field: "loc", + predicate: &nodedb_physical::physical_plan::SpatialPredicate::DWithin, + query_geometry: &query, + distance_meters: 1_000.0, + attribute_filters: &[], + limit: 100, + projection: &[], + rls_filters: &[], + prefilter: None, + }); + assert_eq!(resp.status, Status::Error); + } } diff --git a/nodedb/src/data/executor/handlers/spatial_refine.rs b/nodedb/src/data/executor/handlers/spatial_refine.rs index 67e9faba3..3903f4e6f 100644 --- a/nodedb/src/data/executor/handlers/spatial_refine.rs +++ b/nodedb/src/data/executor/handlers/spatial_refine.rs @@ -10,7 +10,9 @@ //! `spatial` / columnar-family row (fetched from columnar) alike. use nodedb_physical::physical_plan::SpatialPredicate; -use nodedb_types::Value; +use nodedb_spatial::SridMismatch; +use nodedb_types::geometry::SridGeometry; +use nodedb_types::{BboxSpace, Value}; /// Extract geometry from a document field. /// @@ -18,41 +20,65 @@ use nodedb_types::Value; /// - `Value::Geometry(g)` — native geometry (columnar path preserves type) /// - `Value::String(s)` — GeoJSON string (from SQL ST_Point → serialized) /// - `Value::Object(_)` — GeoJSON object (from schemaless doc storage) +/// +/// GeoJSON keeps the SRID its `crs` member names; no member is WGS 84. pub(in crate::data::executor) fn extract_geometry( doc: &Value, field: &str, -) -> Option { - let field_val = doc.get(field)?; - match field_val { - Value::Geometry(g) => Some(g.clone()), - Value::String(s) => nodedb_types::geometry::from_geojson_str(s), - Value::Object(map) => { - // GeoJSON object stored as Value::Object — serialize to JSON then parse. - let json = serde_json::Value::from(Value::Object(map.clone())); - serde_json::from_value(json).ok() - } - _ => None, - } +) -> Option { + SridGeometry::from_value(doc.get(field)?) } /// Apply the spatial predicate. +/// +/// The stored and query geometries must share an SRID: a mismatch is +/// [`SridMismatch`] (`index` the stored SRID, `entry` the query's), which +/// the caller raises as `Error::MixedSrid` instead of skipping the row. +/// `ST_DWithin` measures geodesic metres for 4326 and planar metres for a +/// projected SRID. pub(in crate::data::executor) fn apply_predicate( predicate: &SpatialPredicate, - query: &nodedb_types::geometry::Geometry, - doc: &nodedb_types::geometry::Geometry, + query: &SridGeometry, + doc: &SridGeometry, distance_meters: f64, -) -> bool { - match predicate { +) -> Result { + if query.srid != doc.srid { + return Err(SridMismatch { + index: doc.srid, + entry: query.srid, + }); + } + let (q, d) = (&query.geometry, &doc.geometry); + Ok(match predicate { SpatialPredicate::DWithin => { - crate::engine::spatial::st_dwithin(query, doc, distance_meters) + if doc.srid.is_geographic() { + crate::engine::spatial::st_dwithin(q, d, distance_meters) + } else { + nodedb_spatial::planar_dwithin(q, d, distance_meters) + } } // `ST_Contains(loc, q)` asks whether the *stored* geometry contains // the query geometry — the geofencing shape, where `loc` is a zone // polygon and `q` a point. `ST_Within(loc, q)` is its converse. Both // pass the stored geometry in the position SQL named first. - SpatialPredicate::Contains => crate::engine::spatial::st_contains(doc, query), - SpatialPredicate::Intersects => crate::engine::spatial::st_intersects(query, doc), - SpatialPredicate::Within => crate::engine::spatial::st_within(doc, query), + SpatialPredicate::Contains => crate::engine::spatial::st_contains(d, q), + SpatialPredicate::Intersects => crate::engine::spatial::st_intersects(q, d), + SpatialPredicate::Within => crate::engine::spatial::st_within(d, q), + }) +} + +/// The `Error::MixedSrid` a spatial scan of `collection.field` raises for +/// `mismatch`. +pub(in crate::data::executor) fn mixed_srid_error( + collection: &str, + field: &str, + mismatch: SridMismatch, +) -> crate::Error { + crate::Error::MixedSrid { + collection: collection.to_string(), + field: field.to_string(), + index_srid: mismatch.index.code(), + row_srid: mismatch.entry.code(), } } @@ -116,14 +142,33 @@ pub(in crate::data::executor) fn project_doc( } } -/// Expand a bounding box by a distance in meters. +/// Metres per degree used to widen a geographic search box: the WGS 84 +/// equatorial degree of longitude, rounded. Using it for latitude too keeps +/// the box at least as wide as the true distance. +const METRES_PER_DEGREE: f64 = 111_320.0; + +/// Floor on cos(latitude) so a box near a pole stays finite. +const MIN_COS_LAT: f64 = 0.001; + +/// Expand a bounding box by a distance in meters, in `space`. +/// +/// Geographic space converts metres to degrees at the box's mean latitude. +/// Planar space adds the metres directly. pub(in crate::data::executor) fn expand_bbox( bbox: &nodedb_types::BoundingBox, meters: f64, + space: BboxSpace, ) -> nodedb_types::BoundingBox { - let lat_delta = meters / 111_320.0; - let avg_lat = ((bbox.min_lat + bbox.max_lat) / 2.0).to_radians(); - let lng_delta = meters / (111_320.0 * avg_lat.cos().max(0.001)); + let (lng_delta, lat_delta) = match space { + BboxSpace::Geographic => { + let avg_lat = ((bbox.min_lat + bbox.max_lat) / 2.0).to_radians(); + ( + meters / (METRES_PER_DEGREE * avg_lat.cos().max(MIN_COS_LAT)), + meters / METRES_PER_DEGREE, + ) + } + BboxSpace::Planar => (meters, meters), + }; nodedb_types::BoundingBox::new( bbox.min_lng - lng_delta, diff --git a/nodedb/src/data/executor/handlers/spatial_sync.rs b/nodedb/src/data/executor/handlers/spatial_sync.rs index fcbe5e9e0..b49f1ce6c 100644 --- a/nodedb/src/data/executor/handlers/spatial_sync.rs +++ b/nodedb/src/data/executor/handlers/spatial_sync.rs @@ -36,8 +36,7 @@ use crate::engine::spatial::{RTree, RTreeEntry}; use crate::types::TenantId; use crate::util::fnv1a_hash; use nodedb_types::Surrogate; -use nodedb_types::bbox::geometry_bbox; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::sync::wire::{AckStatus, SyncProvenance}; /// Parameters for a spatial geometry insert on the Data Plane. @@ -50,7 +49,7 @@ pub(in crate::data::executor) struct SpatialInsertExec<'a> { pub collection: &'a str, pub field: &'a str, pub surrogate: Surrogate, - pub geometry: &'a Geometry, + pub geometry: &'a SridGeometry, pub provenance: Option<&'a SyncProvenance>, } @@ -113,6 +112,23 @@ impl CoreLoop { } } + // ── SRID check, before any write ──────────────────────────────────── + let db_id = task.request.database_id; + let tenant_id = TenantId::new(tid); + let spatial_key = (db_id, tenant_id, collection.to_string(), field.to_string()); + if let Some(rtree) = self.spatial_indexes.get(&spatial_key) + && !rtree.accepts_srid(geometry.srid) + { + let mismatch = crate::engine::spatial::SridMismatch { + index: rtree.srid(), + entry: geometry.srid, + }; + return self.response_error( + task, + super::spatial_refine::mixed_srid_error(collection, field, mismatch), + ); + } + // ── 1. Serialise + write minimal geometry document to sparse store ── let geom_value = geometry_to_value(geometry); let mut doc_map = std::collections::HashMap::new(); @@ -164,9 +180,6 @@ impl CoreLoop { } // ── 2. Update R-tree ──────────────────────────────────────────────── - let db_id = task.request.database_id; - let tenant_id = TenantId::new(tid); - let spatial_key = (db_id, tenant_id, collection.to_string(), field.to_string()); let entry_id = fnv1a_hash(doc_id.as_bytes()); let doc_map_key = ( db_id, @@ -184,7 +197,7 @@ impl CoreLoop { rtree.delete(entry_id); } - let bbox = geometry_bbox(geometry); + let bbox = geometry.bbox(); let memory = nodedb_mem::ScopedMemory::new( self.governor.clone(), db_id, @@ -194,7 +207,15 @@ impl CoreLoop { let rtree = self .spatial_indexes .entry(spatial_key) - .or_insert_with(|| RTree::new(memory)); + .or_insert_with(|| RTree::new(memory, geometry.srid)); + // The SRID check above admitted this SRID; the stale-entry delete can + // only empty the tree. + if let Err(mismatch) = rtree.bind_srid(geometry.srid) { + return self.response_error( + task, + super::spatial_refine::mixed_srid_error(collection, field, mismatch), + ); + } rtree.insert(RTreeEntry { id: entry_id, bbox }); self.spatial_doc_map.insert(doc_map_key, doc_id); @@ -308,16 +329,38 @@ impl CoreLoop { } } -/// Convert a `Geometry` into a `nodedb_types::Value` for msgpack serialisation. +/// Convert a `SridGeometry` into a `nodedb_types::Value` for msgpack +/// serialisation. /// /// Produces a JSON-compatible nested structure matching what a pgwire INSERT -/// with a GeoJSON geometry value would produce. +/// with a GeoJSON geometry value produces. A non-4326 SRID adds the +/// GeoJSON `crs` member, so the read path recovers the SRID. /// /// Shared with the in-transaction staging path /// (`transaction::stage_write::stage_spatial`), which stages the identical /// `{field: geometry, "id": surrogate_hex}` document shape as an overlay /// `Put` body. -pub(in crate::data::executor) fn geometry_to_value(geometry: &Geometry) -> nodedb_types::Value { +pub(in crate::data::executor) fn geometry_to_value(geometry: &SridGeometry) -> nodedb_types::Value { + use nodedb_types::Value; + let mut shape = shape_to_value(&geometry.geometry); + if !geometry.srid.is_geographic() + && let Value::Object(map) = &mut shape + { + let mut properties = std::collections::HashMap::new(); + properties.insert( + "name".to_string(), + Value::String(format!("EPSG:{}", geometry.srid.code())), + ); + let mut crs = std::collections::HashMap::new(); + crs.insert("type".to_string(), Value::String("name".to_string())); + crs.insert("properties".to_string(), Value::Object(properties)); + map.insert("crs".to_string(), Value::Object(crs)); + } + shape +} + +/// The GeoJSON object of one shape, as a `Value`. +fn shape_to_value(geometry: &Geometry) -> nodedb_types::Value { use nodedb_types::Value; use nodedb_types::geometry::Geometry::*; @@ -415,7 +458,7 @@ mod tests { collection: "places", field: "loc", surrogate, - geometry: &Geometry::point(10.0, 20.0), + geometry: &SridGeometry::wgs84(Geometry::point(10.0, 20.0)), provenance: None, }) } diff --git a/nodedb/src/data/executor/handlers/transaction/overlay/spatial_merge.rs b/nodedb/src/data/executor/handlers/transaction/overlay/spatial_merge.rs index b21810cf7..c4d0718d8 100644 --- a/nodedb/src/data/executor/handlers/transaction/overlay/spatial_merge.rs +++ b/nodedb/src/data/executor/handlers/transaction/overlay/spatial_merge.rs @@ -32,14 +32,14 @@ use std::collections::HashSet; use nodedb_physical::physical_plan::SpatialPredicate; use nodedb_types::Surrogate; use nodedb_types::columnar::ColumnarSchema; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use nodedb_types::value::Value; use crate::bridge::scan_filter::ScanFilter; use crate::data::executor::core_loop::CoreLoop; use crate::data::executor::handlers::columnar_read::convert::row_to_projected_value; use crate::data::executor::handlers::spatial_refine::{ - SpatialHit, apply_predicate, extract_geometry, project_doc, + SpatialHit, apply_predicate, extract_geometry, mixed_srid_error, project_doc, }; use crate::data::executor::handlers::transaction::overlay::Staged; use crate::engine::document::store::StorageKey; @@ -51,7 +51,7 @@ pub(in crate::data::executor) struct SpatialOverlayMergeParams<'a> { pub coll_key: &'a (DatabaseId, TenantId, String), pub field: &'a str, pub predicate: &'a SpatialPredicate, - pub query_geom: &'a Geometry, + pub query_geom: &'a SridGeometry, pub distance_meters: f64, pub projection: &'a [String], pub attr_filters: &'a [ScanFilter], @@ -124,7 +124,10 @@ impl CoreLoop { let Some(doc_geom) = extract_geometry(doc, field) else { return Ok(false); }; - if !apply_predicate(predicate, query_geom, &doc_geom, distance_meters) { + // A staged row in another SRID fails the scan, as a base row does. + if !apply_predicate(predicate, query_geom, &doc_geom, distance_meters) + .map_err(|m| mixed_srid_error(&coll_key.2, field, m))? + { return Ok(false); } Ok(ScanFilter::all_match_value(attr_filters, doc)? diff --git a/nodedb/src/data/executor/handlers/transaction/redo_apply/install_refusal_tests.rs b/nodedb/src/data/executor/handlers/transaction/redo_apply/install_refusal_tests.rs index 9cf6d9dac..c40dea1ad 100644 --- a/nodedb/src/data/executor/handlers/transaction/redo_apply/install_refusal_tests.rs +++ b/nodedb/src/data/executor/handlers/transaction/redo_apply/install_refusal_tests.rs @@ -16,7 +16,7 @@ use nodedb_array::types::domain::{Domain, DomainBound}; use nodedb_physical::physical_plan::{ ArrayOp, CrdtOp, CrdtWriteVerb, DocumentOp, KvOp, PhysicalPlan, SpatialOp, VectorOp, }; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::{DatabaseId, QualifiedCollection, RlsWriteCheck, Surrogate, TenantId, Value}; use crate::bridge::envelope::{ErrorCode, Response, Status}; @@ -230,9 +230,9 @@ fn spatial_insert() -> PhysicalPlan { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "places"), field: "geom".to_string(), surrogate: Surrogate::new(81), - geometry: Geometry::Point { + geometry: SridGeometry::wgs84(Geometry::Point { coordinates: [1.0, 2.0], - }, + }), provenance: None, }) } diff --git a/nodedb/src/data/executor/handlers/transaction/resolve/entry.rs b/nodedb/src/data/executor/handlers/transaction/resolve/entry.rs index 785bba145..da7a0ae96 100644 --- a/nodedb/src/data/executor/handlers/transaction/resolve/entry.rs +++ b/nodedb/src/data/executor/handlers/transaction/resolve/entry.rs @@ -3506,7 +3506,7 @@ mod tests { // directly, mirroring the vector/columnar plan-driven tests above. use nodedb_physical::physical_plan::SpatialOp; - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, SridGeometry}; fn spatial_prov(seq: u64) -> SyncProvenance { SyncProvenance { @@ -3517,15 +3517,15 @@ mod tests { } } - fn spatial_point(x: f64, y: f64) -> Geometry { - Geometry::point(x, y) + fn spatial_point(x: f64, y: f64) -> SridGeometry { + SridGeometry::wgs84(Geometry::point(x, y)) } fn spatial_insert_plan( collection: &str, field: &str, surrogate: u32, - geometry: Geometry, + geometry: SridGeometry, seq: u64, ) -> PhysicalPlan { PhysicalPlan::Spatial(SpatialOp::Insert { diff --git a/nodedb/src/data/executor/handlers/transaction/resolve/spatial.rs b/nodedb/src/data/executor/handlers/transaction/resolve/spatial.rs index 7d062aba6..92148b291 100644 --- a/nodedb/src/data/executor/handlers/transaction/resolve/spatial.rs +++ b/nodedb/src/data/executor/handlers/transaction/resolve/spatial.rs @@ -101,7 +101,7 @@ mod tests { use super::*; use nodedb_types::Surrogate; - use nodedb_types::geometry::Geometry; + use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::sync::wire::SyncProvenance; use nodedb_types::{DatabaseId, QualifiedCollection}; @@ -114,8 +114,8 @@ mod tests { } } - fn point(x: f64, y: f64) -> Geometry { - Geometry::point(x, y) + fn point(x: f64, y: f64) -> SridGeometry { + SridGeometry::wgs84(Geometry::point(x, y)) } #[test] @@ -136,7 +136,7 @@ mod tests { .expect("decode SpatialPutPayload"); assert_eq!(decoded.collection, "places"); assert_eq!(decoded.field, "loc"); - let geom: Geometry = + let geom: SridGeometry = zerompk::from_msgpack(&decoded.geometry_bytes).expect("decode geometry"); assert_eq!(geom, point(10.0, 20.0)); } diff --git a/nodedb/src/data/executor/handlers/transaction/stage_write/stage_spatial.rs b/nodedb/src/data/executor/handlers/transaction/stage_write/stage_spatial.rs index 43de1809a..4308b5b3a 100644 --- a/nodedb/src/data/executor/handlers/transaction/stage_write/stage_spatial.rs +++ b/nodedb/src/data/executor/handlers/transaction/stage_write/stage_spatial.rs @@ -24,7 +24,7 @@ //! branch). COMMIT installs the spatial write from the transaction's redo //! record, which serializes it from the buffered `SpatialOp` plan node. -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::SridGeometry; use nodedb_types::{RowIdentity, Surrogate}; use super::context::StageCtx; @@ -43,7 +43,7 @@ pub(in crate::data::executor) struct StageSpatialInsertParams<'a> { pub collection: &'a str, pub field: &'a str, pub surrogate: Surrogate, - pub geometry: &'a Geometry, + pub geometry: &'a SridGeometry, } impl CoreLoop { diff --git a/nodedb/src/data/executor/handlers/transaction/undo/entry.rs b/nodedb/src/data/executor/handlers/transaction/undo/entry.rs index 43fc796e6..04aa889e6 100644 --- a/nodedb/src/data/executor/handlers/transaction/undo/entry.rs +++ b/nodedb/src/data/executor/handlers/transaction/undo/entry.rs @@ -205,12 +205,14 @@ pub(in crate::data::executor) enum UndoEntry { /// /// `bbox` is the entry's geometry captured BEFORE the forward `delete` /// (the R-tree `delete` does not return it); `document_id` is the reverse - /// map value removed by the forward cascade. + /// map value removed by the forward cascade. `srid` is the SRID of the + /// tree the entry left, which the re-insert binds the tree back to. SpatialDelete { key: SpatialIndexKey, entry_id: u64, bbox: nodedb_types::BoundingBox, document_id: String, + srid: nodedb_types::geometry::Srid, }, /// Undo a graph edge write: remove the version it added and put the CSR /// back. diff --git a/nodedb/src/data/executor/handlers/transaction/undo/spatial.rs b/nodedb/src/data/executor/handlers/transaction/undo/spatial.rs index ebcf268ba..69064d8e1 100644 --- a/nodedb/src/data/executor/handlers/transaction/undo/spatial.rs +++ b/nodedb/src/data/executor/handlers/transaction/undo/spatial.rs @@ -17,7 +17,7 @@ use super::{UndoEntry, UndoError}; impl CoreLoop { pub(super) fn apply_undo_spatial( &mut self, - _entry_index: usize, + entry_index: usize, entry: UndoEntry, ) -> Result<(), UndoError> { match entry { @@ -37,10 +37,13 @@ impl CoreLoop { entry_id, bbox, document_id, + srid, } => { // Reverse a forward spatial removal: re-insert the entry with // its captured bbox and re-populate the reverse map, matching - // the forward `apply_point_put_spatial` insert shape. + // the forward `apply_point_put_spatial` insert shape. Undo + // runs in reverse, so every later insert into this tree is + // already gone and the tree binds back to `srid`. let memory = nodedb_mem::ScopedMemory::new( self.governor.clone(), key.0, @@ -50,7 +53,19 @@ impl CoreLoop { let rtree = self .spatial_indexes .entry(key.clone()) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)); + .or_insert_with(|| crate::engine::spatial::RTree::new(memory, srid)); + rtree.bind_srid(srid).map_err(|m| { + UndoError::failed( + entry_index, + format!("re-inserting a spatial entry into '{}.{}'", key.2, key.3), + crate::Error::MixedSrid { + collection: key.2.clone(), + field: key.3.clone(), + index_srid: m.index.code(), + row_srid: m.entry.code(), + }, + ) + })?; rtree.insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); self.spatial_doc_map .insert((key.0, key.1, key.2, key.3, entry_id), document_id); @@ -102,10 +117,9 @@ mod tests { key.1, nodedb_mem::EngineId::Spatial, ); - let rtree = core - .spatial_indexes - .entry(key.clone()) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)); + let rtree = core.spatial_indexes.entry(key.clone()).or_insert_with(|| { + crate::engine::spatial::RTree::new(memory, nodedb_types::geometry::Srid::Wgs84) + }); rtree.insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); core.spatial_doc_map.insert( (key.0, key.1, key.2.clone(), key.3.clone(), entry_id), @@ -145,6 +159,7 @@ mod tests { entry_id, bbox, document_id: "d1".to_string(), + srid: nodedb_types::geometry::Srid::Utm48N, }; core.apply_undo_spatial(0, undo).unwrap(); @@ -157,6 +172,11 @@ mod tests { restored.bbox, bbox, "restored bbox must match captured bbox" ); + assert_eq!( + core.spatial_indexes.get(&key).map(|rt| rt.srid()), + Some(nodedb_types::geometry::Srid::Utm48N), + "the re-created tree takes the captured SRID" + ); assert_eq!( core.spatial_doc_map .get(&(key.0, key.1, key.2, key.3, entry_id)) diff --git a/nodedb/src/data/executor/handlers/transaction/undo/spatial_row.rs b/nodedb/src/data/executor/handlers/transaction/undo/spatial_row.rs index 304e2ccf7..c28bda4f7 100644 --- a/nodedb/src/data/executor/handlers/transaction/undo/spatial_row.rs +++ b/nodedb/src/data/executor/handlers/transaction/undo/spatial_row.rs @@ -24,9 +24,9 @@ pub(in crate::data::executor) struct SpatialRowUndo { pub prior_row: Option>, /// The R-tree entry id the write keys on. pub entry_id: u64, - /// The R-tree entry and its reverse-map value before the write, `None` - /// when the field held no entry for the row. - pub prior_entry: Option<(BoundingBox, String)>, + /// The R-tree entry, its reverse-map value and the tree's SRID before the + /// write, `None` when the field held no entry for the row. + pub prior_entry: Option<(BoundingBox, String, nodedb_types::geometry::Srid)>, } impl CoreLoop { @@ -56,14 +56,13 @@ impl CoreLoop { .spatial_doc_map .get(&doc_map_key) .and_then(|document_id| { - let bbox = self - .spatial_indexes - .get(&key)? + let rtree = self.spatial_indexes.get(&key)?; + let bbox = rtree .entries() .into_iter() .find(|entry| entry.id == entry_id) .map(|entry| entry.bbox)?; - Some((bbox, document_id.clone())) + Some((bbox, document_id.clone(), rtree.srid())) }); Ok(UndoEntry::SpatialRow(Box::new(SpatialRowUndo { key, @@ -112,17 +111,33 @@ impl CoreLoop { } let doc_map_key = (key.0, key.1, key.2.clone(), key.3.clone(), entry_id); match prior_entry { - Some((bbox, document_id)) => { + Some((bbox, document_id, srid)) => { let memory = nodedb_mem::ScopedMemory::new( self.governor.clone(), key.0, key.1, nodedb_mem::EngineId::Spatial, ); - self.spatial_indexes + let (collection, field) = (key.2.clone(), key.3.clone()); + let rtree = self + .spatial_indexes .entry(key) - .or_insert_with(|| crate::engine::spatial::RTree::new(memory)) - .insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); + .or_insert_with(|| crate::engine::spatial::RTree::new(memory, srid)); + // Undo runs in reverse, so the tree holds no later entry in + // another SRID and binds back to the row's. + rtree.bind_srid(srid).map_err(|m| { + UndoError::failed( + entry_index, + format!("restoring spatial entry in '{collection}.{field}'"), + crate::Error::MixedSrid { + collection: collection.clone(), + field: field.clone(), + index_srid: m.index.code(), + row_srid: m.entry.code(), + }, + ) + })?; + rtree.insert(crate::engine::spatial::RTreeEntry { id: entry_id, bbox }); self.spatial_doc_map.insert(doc_map_key, document_id); } None => { diff --git a/nodedb/src/data/executor/handlers/transaction/write_version_parity_tests.rs b/nodedb/src/data/executor/handlers/transaction/write_version_parity_tests.rs index f9ecc7ba4..bddce48d2 100644 --- a/nodedb/src/data/executor/handlers/transaction/write_version_parity_tests.rs +++ b/nodedb/src/data/executor/handlers/transaction/write_version_parity_tests.rs @@ -12,7 +12,7 @@ use nodedb_physical::physical_plan::{ CrdtWriteVerb, DocumentOp, GraphOp, KvOp, MetaOp, PhysicalPlan, RedoOrigin, SpatialOp, TextOp, TimeseriesOp, VectorOp, }; -use nodedb_types::geometry::Geometry; +use nodedb_types::geometry::{Geometry, SridGeometry}; use nodedb_types::{DatabaseId, QualifiedCollection, RlsWriteCheck, Surrogate, Value}; use crate::bridge::envelope::{Response, Status}; @@ -266,9 +266,9 @@ fn a_spatial_slice_records_its_collection_floor() { collection: QualifiedCollection::new(DatabaseId::DEFAULT, "places"), field: "geom".to_string(), surrogate: Surrogate::new(81), - geometry: Geometry::Point { + geometry: SridGeometry::wgs84(Geometry::Point { coordinates: [1.0, 2.0], - }, + }), provenance: None, }); staged_slice_parity(&[plan]); diff --git a/nodedb/src/data/executor/snapshot/capture.rs b/nodedb/src/data/executor/snapshot/capture.rs index e962d6546..284548451 100644 --- a/nodedb/src/data/executor/snapshot/capture.rs +++ b/nodedb/src/data/executor/snapshot/capture.rs @@ -506,7 +506,8 @@ mod tests { tid(), nodedb_mem::EngineId::Spatial, ); - let mut rtree = crate::engine::spatial::RTree::new(memory); + let mut rtree = + crate::engine::spatial::RTree::new(memory, nodedb_types::geometry::Srid::Utm47N); rtree.insert(crate::engine::spatial::RTreeEntry { id: 1, bbox: nodedb_types::BoundingBox::new(0.0, 0.0, 1.0, 1.0), @@ -516,6 +517,11 @@ mod tests { let (_dir, restored) = restore_and_boot(&capture(core)); assert_eq!(restored.spatial_indexes.get(&key).map(|t| t.len()), Some(1)); + assert_eq!( + restored.spatial_indexes.get(&key).map(|t| t.srid()), + Some(nodedb_types::geometry::Srid::Utm47N), + "the SRID survives the snapshot" + ); } #[test] diff --git a/nodedb/src/data/executor/spatial_checkpoint/load.rs b/nodedb/src/data/executor/spatial_checkpoint/load.rs index 902c67fd8..903efb7f5 100644 --- a/nodedb/src/data/executor/spatial_checkpoint/load.rs +++ b/nodedb/src/data/executor/spatial_checkpoint/load.rs @@ -179,6 +179,7 @@ mod tests { bbox: BoundingBox::new(0.0, 0.0, 1.0, 1.0), }], memory, + nodedb_types::geometry::Srid::Wgs84, ); rtree.checkpoint_to_bytes(None).expect("encode R-tree") } diff --git a/nodedb/src/data/executor/spatial_checkpoint/write.rs b/nodedb/src/data/executor/spatial_checkpoint/write.rs index f68a6cc4a..63686a5e7 100644 --- a/nodedb/src/data/executor/spatial_checkpoint/write.rs +++ b/nodedb/src/data/executor/spatial_checkpoint/write.rs @@ -216,7 +216,10 @@ mod tests { fn seed_one_entry(core: &mut CoreLoop) { let (db, tid, coll, field) = key(); - let mut rtree = RTree::new(test_rtree_memory(db, tid)); + let mut rtree = RTree::new( + test_rtree_memory(db, tid), + nodedb_types::geometry::Srid::Wgs84, + ); rtree.insert(RTreeEntry { id: 1, bbox: BoundingBox::new(0.0, 0.0, 1.0, 1.0), @@ -313,7 +316,10 @@ mod tests { // Geometry with no doc_map entries at all: the shape a columnar-family // index takes before any document row maps into it. let (db, tid, _, _) = key(); - let mut rtree = RTree::new(test_rtree_memory(db, tid)); + let mut rtree = RTree::new( + test_rtree_memory(db, tid), + nodedb_types::geometry::Srid::Wgs84, + ); rtree.insert(RTreeEntry { id: 1, bbox: BoundingBox::new(0.0, 0.0, 1.0, 1.0), diff --git a/nodedb/src/data/executor/wal_replay_spatial.rs b/nodedb/src/data/executor/wal_replay_spatial.rs index e6c9dc85c..afc274ff8 100644 --- a/nodedb/src/data/executor/wal_replay_spatial.rs +++ b/nodedb/src/data/executor/wal_replay_spatial.rs @@ -16,9 +16,10 @@ //! ## Geometry decode on replay //! //! `SpatialPutPayload.geometry_bytes` carries msgpack-encoded -//! `nodedb_types::geometry::Geometry` (the same format stored in -//! `SpatialInsertMsg.geometry_bytes`). On replay we decode it with -//! `zerompk::from_msgpack` and pass the `&Geometry` to `execute_spatial_insert`. +//! `nodedb_types::geometry::SridGeometry` (the same format stored in +//! `SpatialInsertMsg.geometry_bytes`). On replay we decode it with +//! `zerompk::from_msgpack` and pass the `&SridGeometry` to +//! `execute_spatial_insert`. //! //! ## Why there is no replay floor or watermark here //! @@ -204,7 +205,7 @@ impl CoreLoop { }; // Decode geometry from msgpack bytes stored in the WAL payload. - let geometry: nodedb_types::geometry::Geometry = + let geometry: nodedb_types::geometry::SridGeometry = match zerompk::from_msgpack(&payload.geometry_bytes) { Ok(g) => g, Err(e) => { @@ -480,7 +481,8 @@ mod tests { /// idempotency gate, so these tests exercise the engine's OWN idempotency /// rather than the gate's — the property replay actually depends on. fn put_record(lsn: u64, geometry: &Geometry) -> nodedb_wal::WalRecord { - let geometry_bytes = zerompk::to_msgpack_vec(geometry).expect("encode geometry"); + let tagged = nodedb_types::geometry::SridGeometry::wgs84(geometry.clone()); + let geometry_bytes = zerompk::to_msgpack_vec(&tagged).expect("encode geometry"); let payload = nodedb_wal::record::SpatialPutPayload::new( SyncProvenance::default(), COLLECTION, diff --git a/nodedb/src/engine/spatial/mod.rs b/nodedb/src/engine/spatial/mod.rs index 83959cb00..77def4414 100644 --- a/nodedb/src/engine/spatial/mod.rs +++ b/nodedb/src/engine/spatial/mod.rs @@ -6,6 +6,7 @@ pub use nodedb_spatial::GeohashIndex; pub use nodedb_spatial::RTree; pub use nodedb_spatial::RTreeEntry; +pub use nodedb_spatial::SridMismatch; pub use nodedb_spatial::geo_meta; pub use nodedb_spatial::geohash; pub use nodedb_spatial::geohash_index; diff --git a/nodedb/src/error/types.rs b/nodedb/src/error/types.rs index cfed4ab61..2f6fb6660 100644 --- a/nodedb/src/error/types.rs +++ b/nodedb/src/error/types.rs @@ -374,6 +374,20 @@ pub enum Error { #[error("{detail}")] DataException { detail: String }, + /// A geometry's SRID differs from the SRID of the spatial index it + /// writes to or queries. The index SRID is fixed by its first entry, and + /// no implicit transform runs. Rendered as SQLSTATE `22000` + /// (data_exception) at the pgwire layer. + #[error( + "spatial index on {collection}.{field} is SRID {index_srid}; geometry is SRID {row_srid}" + )] + MixedSrid { + collection: String, + field: String, + index_srid: u32, + row_srid: u32, + }, + /// A value does not fit its numeric type: an integer past a column's /// declared width, a float that overflows `REAL`, a literal past the exact /// numeric range, or arithmetic that overflows. Rendered as SQLSTATE diff --git a/nodedb/src/error_classify/public.rs b/nodedb/src/error_classify/public.rs index 71481b9ce..0a29c5cd7 100644 --- a/nodedb/src/error_classify/public.rs +++ b/nodedb/src/error_classify/public.rs @@ -196,6 +196,7 @@ pub(crate) fn classify(e: &Error) -> NodeDbError { Error::UnknownStrictField { column, .. } => NodeDbError::undefined_column(column.clone()), Error::DivisionByZero => NodeDbError::division_by_zero(), Error::DataException { detail } => NodeDbError::data_exception(detail.clone()), + e @ Error::MixedSrid { .. } => NodeDbError::data_exception(e.to_string()), Error::NumericValueOutOfRange { detail } => { NodeDbError::numeric_value_out_of_range(detail.clone()) } diff --git a/nodedb/src/error_classify/unclassified.rs b/nodedb/src/error_classify/unclassified.rs index 599df3586..2e747e6cd 100644 --- a/nodedb/src/error_classify/unclassified.rs +++ b/nodedb/src/error_classify/unclassified.rs @@ -92,6 +92,7 @@ pub(crate) fn is_unclassified_failure(e: &Error) -> bool { | Error::UnknownStrictField { .. } | Error::DivisionByZero | Error::DataException { .. } + | Error::MixedSrid { .. } | Error::NumericValueOutOfRange { .. } | Error::InvalidTextRepresentation { .. } | Error::DatatypeMismatch { .. } diff --git a/nodedb/src/error_from.rs b/nodedb/src/error_from.rs index 84fc650b6..1578d224d 100644 --- a/nodedb/src/error_from.rs +++ b/nodedb/src/error_from.rs @@ -16,26 +16,6 @@ impl From for Error { } } -/// The `match` is exhaustive rather than a `_ =>` fallback, so a new -/// evaluator error picks its own `crate::Error` mapping instead of silently -/// inheriting one. -impl From for Error { - fn from(e: nodedb_query::EvalError) -> Self { - match e { - nodedb_query::EvalError::DivisionByZero => Self::DivisionByZero, - nodedb_query::EvalError::UnknownFunction { name } => Self::UndefinedFunction { name }, - e @ (nodedb_query::EvalError::VectorDimensionMismatch { .. } - | nodedb_query::EvalError::ArgumentType { .. } - | nodedb_query::EvalError::InvalidJsonPath { .. }) => Self::DataException { - detail: e.to_string(), - }, - e @ nodedb_query::EvalError::NumericOverflow { .. } => Self::NumericValueOutOfRange { - detail: e.to_string(), - }, - } - } -} - /// A non-scalar offered where a raw KV body is required is a request /// shape error, not an internal fault. impl From for Error { @@ -427,6 +407,7 @@ impl From for nodedb_cluster::rpc_codec::TypedClusterError { | Error::UnknownStrictField { .. } | Error::DivisionByZero | Error::DataException { .. } + | Error::MixedSrid { .. } | Error::NumericValueOutOfRange { .. } | Error::InvalidLimitValue { .. } | Error::RetryableLeaderChange { .. } diff --git a/nodedb/src/error_from_eval.rs b/nodedb/src/error_from_eval.rs new file mode 100644 index 000000000..bd7ed6c15 --- /dev/null +++ b/nodedb/src/error_from_eval.rs @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! `nodedb_query::EvalError` into the crate error. + +use crate::Error; + +/// The `match` is exhaustive rather than a `_ =>` fallback, so a new +/// evaluator error picks its own `crate::Error` mapping instead of silently +/// inheriting one. +impl From for Error { + fn from(e: nodedb_query::EvalError) -> Self { + match e { + nodedb_query::EvalError::DivisionByZero => Self::DivisionByZero, + nodedb_query::EvalError::UnknownFunction { name } => Self::UndefinedFunction { name }, + e @ (nodedb_query::EvalError::VectorDimensionMismatch { .. } + | nodedb_query::EvalError::ArgumentType { .. } + | 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 { + detail: e.to_string(), + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// An SRID fault fails the statement with SQLSTATE `22000`. + #[test] + fn srid_faults_are_data_exceptions() { + for e in [ + nodedb_query::EvalError::UnsupportedSrid { + function: "st_transform", + srid: 2154, + }, + nodedb_query::EvalError::MixedSrid { + function: "st_intersects", + left: 4326, + right: 32648, + }, + nodedb_query::EvalError::ProjectionDomain { + function: "st_transform", + srid: 32648, + }, + ] { + let text = e.to_string(); + match Error::from(e) { + Error::DataException { detail } => assert_eq!(detail, text), + other => panic!("expected Error::DataException, got {other:?}"), + } + } + } +} diff --git a/nodedb/tests/wire/cases/spatial_crs.rs b/nodedb/tests/wire/cases/spatial_crs.rs new file mode 100644 index 000000000..3d19b4f57 --- /dev/null +++ b/nodedb/tests/wire/cases/spatial_crs.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Coordinate reference systems over pgwire: `ST_SRID`, `ST_SetSRID`, +//! `ST_Transform`, the SRID argument of the text constructors, planar +//! measures in a projected SRID, and the one-SRID-per-index rule. +//! +//! Reference coordinates come from PROJ 9.8.1 via pyproj (EPSG:4326 to the +//! target, always_xy). + +use crate::harness::TestServer; + +/// Largest error accepted against a PROJ reference value, metres. +const MM: f64 = 1e-3; + +async fn scalar(srv: &TestServer, sql: &str) -> String { + let rows = srv + .query_rows(sql) + .await + .unwrap_or_else(|e| panic!("`{sql}` failed: {e}")); + assert_eq!(rows.len(), 1, "`{sql}` must return one row, got {rows:?}"); + rows[0][0].clone() +} + +async fn number(srv: &TestServer, sql: &str) -> f64 { + let text = scalar(srv, sql).await; + text.parse() + .unwrap_or_else(|_| panic!("`{sql}` must return a number, got {text:?}")) +} + +#[tokio::test] +async fn setsrid_then_srid_reads_the_new_srid() { + let srv = TestServer::start().await; + let srid = scalar(&srv, "SELECT ST_SRID(ST_SetSRID(ST_Point(1, 2), 32648))").await; + assert_eq!(srid, "32648"); +} + +#[tokio::test] +async fn geomfromtext_srid_argument_sets_the_srid() { + let srv = TestServer::start().await; + let srid = scalar(&srv, "SELECT ST_SRID(ST_GeomFromText('POINT(1 2)', 3375))").await; + assert_eq!(srid, "3375"); + let plain = scalar(&srv, "SELECT ST_SRID(ST_GeomFromText('POINT(1 2)'))").await; + assert_eq!(plain, "4326", "WKT without an SRID argument is WGS 84"); +} + +#[tokio::test] +async fn mixed_srid_predicate_is_an_error() { + let srv = TestServer::start().await; + srv.expect_error( + "SELECT ST_Intersects(ST_Point(1, 2), ST_SetSRID(ST_Point(1, 2), 32648))", + "different SRIDs", + ) + .await; + srv.expect_error( + "SELECT ST_Intersects(ST_Point(1, 2), ST_SetSRID(ST_Point(1, 2), 32648))", + "22000", + ) + .await; +} + +#[tokio::test] +async fn transform_to_an_unsupported_srid_names_every_supported_code() { + let srv = TestServer::start().await; + let sql = "SELECT ST_Transform(ST_Point(101.7, 3.1), 2154)"; + srv.expect_error(sql, "unsupported SRID 2154").await; + for code in ["4326", "32647", "32648", "3375", "3376"] { + srv.expect_error(sql, code).await; + } +} + +#[tokio::test] +async fn transform_klcc_matches_proj() { + let srv = TestServer::start().await; + let cases = [ + (32647, 801_401.794_9, 349_439.510_6), + (32648, 134_445.592_4, 349_624.727_8), + (3375, 412_834.468_1, 349_472.917_7), + ]; + for (srid, east, north) in cases { + let geom = format!("ST_Transform(ST_Point(101.7117, 3.1579), {srid})"); + let x = number(&srv, &format!("SELECT ST_X({geom})")).await; + let y = number(&srv, &format!("SELECT ST_Y({geom})")).await; + assert!( + (x - east).abs() < MM && (y - north).abs() < MM, + "SRID {srid}: got ({x}, {y}), want ({east}, {north})" + ); + let srid_text = scalar(&srv, &format!("SELECT ST_SRID({geom})")).await; + assert_eq!(srid_text, srid.to_string()); + } + let x = number( + &srv, + "SELECT ST_X(ST_Transform(ST_Point(116.0735, 5.9804), 3376))", + ) + .await; + assert!( + (x - 708_697.150_6).abs() < MM, + "3376 Kota Kinabalu: got {x}" + ); +} + +#[tokio::test] +async fn projected_area_is_planar_square_metres() { + let srv = TestServer::start().await; + let area = number( + &srv, + "SELECT ST_Area(ST_GeomFromText(\ + 'POLYGON((500000 100000, 500100 100000, 500100 100100, 500000 100100, 500000 100000))', \ + 32648))", + ) + .await; + assert!((area - 10_000.0).abs() < 1e-6, "got {area}"); +} + +#[tokio::test] +async fn geojson_round_trip_keeps_the_srid() { + let srv = TestServer::start().await; + let json = scalar( + &srv, + "SELECT ST_AsGeoJSON(ST_SetSRID(ST_Point(1, 2), 3376))", + ) + .await; + assert!(json.contains("EPSG:3376"), "{json}"); + let srid = scalar( + &srv, + &format!("SELECT ST_SRID(ST_GeomFromGeoJSON('{json}'))"), + ) + .await; + assert_eq!(srid, "3376"); + + let plain = scalar(&srv, "SELECT ST_AsGeoJSON(ST_Point(1, 2))").await; + assert!(!plain.contains("crs"), "4326 writes no crs member: {plain}"); +} + +/// The index SRID is fixed by its first row. A row in another SRID is +/// refused, and the collection keeps only the first row. +#[tokio::test] +async fn second_srid_on_an_indexed_field_is_refused() { + let srv = TestServer::start().await; + srv.exec( + "CREATE COLLECTION crs_mixed \ + COLUMNS (id TEXT, loc GEOMETRY) \ + WITH (engine='spatial')", + ) + .await + .unwrap(); + srv.exec("INSERT INTO crs_mixed (id, loc) VALUES ('a', ST_Point(101.7, 3.1))") + .await + .unwrap(); + srv.expect_error( + "INSERT INTO crs_mixed (id, loc) \ + VALUES ('b', ST_SetSRID(ST_Point(500000, 0), 32648))", + "SRID", + ) + .await; + + let rows = srv.query_rows("SELECT id FROM crs_mixed").await.unwrap(); + assert_eq!(rows.len(), 1, "the refused row was not stored: {rows:?}"); + assert_eq!(rows[0][0], "a"); +} + +/// A projected index answers `ST_DWithin` in grid metres, and refuses a +/// query geometry in another SRID. +#[tokio::test] +async fn projected_index_is_queried_in_its_own_srid() { + let srv = TestServer::start().await; + srv.exec( + "CREATE COLLECTION crs_utm \ + COLUMNS (id TEXT, loc GEOMETRY) \ + WITH (engine='spatial')", + ) + .await + .unwrap(); + srv.exec( + "INSERT INTO crs_utm (id, loc) VALUES \ + ('near', ST_SetSRID(ST_Point(500300, 400), 32648))", + ) + .await + .unwrap(); + srv.exec( + "INSERT INTO crs_utm (id, loc) VALUES \ + ('far', ST_SetSRID(ST_Point(510000, 0), 32648))", + ) + .await + .unwrap(); + + let rows = srv + .query_rows( + "SELECT id FROM crs_utm \ + WHERE ST_DWithin(loc, ST_SetSRID(ST_Point(500000, 0), 32648), 501)", + ) + .await + .unwrap(); + let ids: Vec<&str> = rows.iter().map(|r| r[0].as_str()).collect(); + assert_eq!(ids, ["near"], "planar 500 m lies within 501 m only"); + + srv.expect_error( + "SELECT id FROM crs_utm WHERE ST_DWithin(loc, ST_Point(101.7, 3.1), 500)", + "SRID", + ) + .await; +} From a20e333f11d45e6c1326e85a9005733591344d20 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Fri, 16 Oct 2026 23:27:23 +0800 Subject: [PATCH 04/29] feat(array): load and read arrays as SQL relations ARRAY_SLICE and ARRAY_PROJECT now return a typed relation, one row per cell, with the dimensions then the attributes as columns. WHERE, computed columns, GROUP BY, aggregates, ORDER BY, window functions and CTEs apply to those rows. Slice bounds push down into tile selection, and tile reads skip tiles by their bounding boxes. CTE inlining moves from the control-plane converter into the SQL planner. Arrays bulk-load through COPY ... FROM STDIN (text and CSV, optional header) and through a new ArrayPutBatch native opcode exposed as NodeDb::array_put_cells. Rows are checked as they arrive and written in bounded chunks, so memory stays flat regardless of input size. Other engines refuse bulk load with a typed error where it does not apply. Surrogate reservation for batched array writes is vectorised, array WAL records and the catalog are reworked for batch appends, and auto flush keeps freshly written cells visible to reads. Adds ordered-map msgpack encoding so rows have a stable column order. --- Cargo.lock | 3 + docs/array.md | 44 ++ docs/bitemporal.md | 2 +- docs/query-language.md | 18 + nodedb-array/Cargo.toml | 1 + nodedb-array/src/codec/tile_decode.rs | 81 ++- nodedb-array/src/codec/tile_encode.rs | 7 + nodedb-array/src/query/mod.rs | 4 +- nodedb-array/src/query/slice.rs | 63 ++ nodedb-array/src/segment/reader.rs | 37 +- nodedb-array/src/tile/mbr.rs | 55 +- nodedb-array/src/tile/sparse_tile.rs | 74 +- nodedb-client-tests/Cargo.toml | 3 + .../tests/array_load_measure.rs | 104 +++ .../tests/array_put_cells_round_trip.rs | 162 +++++ nodedb-client/src/native/client/array.rs | 74 ++ nodedb-client/src/native/client/dispatch.rs | 8 + nodedb-client/src/native/client/mod.rs | 1 + .../src/traits/core/default_impls/batch.rs | 12 + .../src/traits/core/default_impls/mod.rs | 4 +- nodedb-client/src/traits/core/trait_def.rs | 25 +- .../cases/bitemporal_array_cluster.rs | 34 +- .../tests/common_suite/cases/cluster_array.rs | 26 +- .../cluster_array_cell_raft_replication.rs | 15 +- nodedb-cluster/src/distributed_array/merge.rs | 3 +- nodedb-physical/src/physical_plan/array.rs | 9 +- nodedb-sql/src/ddl_ast/parse/copy_from.rs | 25 +- .../src/ddl_ast/parse/copy_from_stdin.rs | 147 ++++ nodedb-sql/src/ddl_ast/parse/mod.rs | 1 + .../src/ddl_ast/statement/maintenance.rs | 11 + .../src/ddl_ast/statement/types/misc.rs | 17 +- nodedb-sql/src/engine_rules/array.rs | 47 ++ nodedb-sql/src/engine_rules/bulk_load.rs | 27 + nodedb-sql/src/engine_rules/columnar.rs | 5 + .../src/engine_rules/continuous_aggregate.rs | 5 + .../src/engine_rules/document_schemaless.rs | 5 + .../src/engine_rules/document_strict.rs | 5 + nodedb-sql/src/engine_rules/kv.rs | 5 + nodedb-sql/src/engine_rules/mod.rs | 2 + nodedb-sql/src/engine_rules/rules.rs | 5 + nodedb-sql/src/engine_rules/spatial.rs | 10 + nodedb-sql/src/engine_rules/timeseries.rs | 10 + nodedb-sql/src/planner/array_fn/mod.rs | 4 +- nodedb-sql/src/planner/array_fn/pushdown.rs | 484 ++++++++++++++ nodedb-sql/src/planner/array_fn/table_fn.rs | 107 ++- .../src/planner/cte/inline.rs | 52 +- nodedb-sql/src/planner/cte/mod.rs | 2 + nodedb-sql/src/planner/dml/bulk_load.rs | 104 +++ nodedb-sql/src/planner/dml/mod.rs | 2 + nodedb-sql/src/planner/join/array_arm.rs | 12 +- nodedb-sql/src/planner/select/array_from.rs | 630 ++++++++++++++++++ nodedb-sql/src/planner/select/cte_catalog.rs | 68 +- nodedb-sql/src/planner/select/derived_from.rs | 16 +- nodedb-sql/src/planner/select/helpers.rs | 23 + nodedb-sql/src/planner/select/limit.rs | 25 +- nodedb-sql/src/planner/select/mod.rs | 1 + nodedb-sql/src/planner/select/post_process.rs | 39 +- nodedb-sql/src/planner/select/select_stmt.rs | 12 +- nodedb-sql/src/planner/select/where_search.rs | 2 +- nodedb-sql/src/resolver/array_tvf.rs | 134 +++- nodedb-sql/src/resolver/columns.rs | 8 +- nodedb-types/src/json_msgpack/mod.rs | 4 +- nodedb-types/src/json_msgpack/writer.rs | 45 ++ nodedb-types/src/protocol/array_batch.rs | 116 ++++ nodedb-types/src/protocol/mod.rs | 2 + nodedb-types/src/protocol/opcodes.rs | 14 + .../src/control/array_sync/raft_apply/op.rs | 9 +- .../cluster/array_cluster_exec/executor.rs | 10 + nodedb/src/control/planner/context/mod.rs | 2 +- .../src/control/planner/context/query/mod.rs | 2 + .../control/planner/context/query/planning.rs | 85 +-- .../planner/context/query/planning_built.rs | 142 ++++ .../sql_plan_convert/array_convert/dml.rs | 76 ++- .../sql_plan_convert/array_convert/mod.rs | 6 +- .../array_convert/tile_route.rs | 27 + .../planner/sql_plan_convert/dml/mod.rs | 2 +- .../sql_plan_convert/dml/surrogate_keys.rs | 43 +- .../planner/sql_plan_convert/expr/mod.rs | 4 +- .../output_schema/build/schema.rs | 81 ++- .../output_schema/cte_types.rs | 14 +- .../planner/sql_plan_convert/set_ops.rs | 3 +- .../surrogate_prefetch/bind.rs | 6 +- .../surrogate_prefetch/cache.rs | 105 ++- .../surrogate_prefetch/resolve.rs | 44 +- .../control/server/exchange/cluster_leaf.rs | 42 +- .../exchange/resolve/array_read_rows.rs | 145 ++++ .../resolve/exchange/post_process_arm.rs | 9 + .../server/exchange/resolve/join_input.rs | 21 +- .../control/server/exchange/resolve/mod.rs | 1 + .../server/native/dispatch/admission_op.rs | 3 +- .../server/native/dispatch/array_batch.rs | 70 ++ .../src/control/server/native/dispatch/mod.rs | 2 + .../control/server/native/session/request.rs | 3 + .../control/server/pgwire/factory/server.rs | 12 +- .../server/pgwire/handler/array_copy.rs | 272 ++++++++ .../server/pgwire/handler/copy_handler.rs | 70 +- .../src/control/server/pgwire/handler/core.rs | 5 + .../src/control/server/pgwire/handler/mod.rs | 2 + .../control/server/pgwire/handler/sql_exec.rs | 12 +- .../response_shape/compose/array_slice.rs | 20 +- .../response_shape/compose/materialized.rs | 6 +- .../copy_from/array_load/chunker.rs | 186 ++++++ .../collection/copy_from/array_load/decode.rs | 431 ++++++++++++ .../copy_from/array_load/dispatch.rs | 182 +++++ .../collection/copy_from/array_load/loader.rs | 362 ++++++++++ .../collection/copy_from/array_load/mod.rs | 13 + .../copy_from/array_load/records.rs | 513 ++++++++++++++ .../collection/copy_from/array_load/router.rs | 58 ++ .../collection/copy_from/array_load/source.rs | 182 +++++ .../ddl/neutral/collection/copy_from/entry.rs | 120 ++-- .../ddl/neutral/collection/copy_from/mod.rs | 31 +- .../shared/ddl/neutral/collection/dml/mod.rs | 3 +- .../ddl/neutral/collection/dml/parse.rs | 6 +- .../neutral/collection/dml/parse/dispatch.rs | 134 ++-- .../collection/dml/parse/plan_statement.rs | 110 +++ .../shared/ddl/neutral/router/dispatch.rs | 14 +- .../shared/ddl/neutral/router/typed_misc.rs | 13 +- .../src/control/server/wal_dispatch/array.rs | 38 +- .../surrogate/assign/cluster_reserve.rs | 36 +- .../surrogate/assign/core/batch_ops.rs | 108 ++- .../src/control/surrogate/registry/cluster.rs | 42 ++ nodedb/src/data/executor/array_checkpoint.rs | 8 +- .../data/executor/dispatch/array/aggregate.rs | 2 +- .../data/executor/dispatch/array/encode.rs | 74 +- .../src/data/executor/dispatch/array/mod.rs | 1 + .../src/data/executor/dispatch/array/read.rs | 226 +++---- .../data/executor/dispatch/array/slice_row.rs | 207 ++++++ .../executor/handlers/point/apply_put/core.rs | 2 +- .../handlers/transaction/resolve/array.rs | 45 +- .../executor/wal_replay_columnar_image.rs | 4 +- nodedb/src/engine/array/flush.rs | 24 + nodedb/src/engine/array/store/catalog/mod.rs | 1 + .../src/engine/array/store/catalog/resolve.rs | 290 ++++++++ nodedb/src/engine/array/store/catalog/scan.rs | 107 +-- .../engine/array/store/catalog/versions.rs | 63 +- .../src/engine/array/store/segment_handle.rs | 24 +- nodedb/src/engine/array/wal.rs | 116 +++- nodedb/src/lib.rs | 1 + .../cases/array_auto_flush_visibility.rs | 106 +++ nodedb/tests/inproc/cases/array_scan_cost.rs | 97 +++ nodedb/tests/inproc/cases/array_sql.rs | 49 +- .../tests/inproc/cases/array_sql_relation.rs | 270 ++++++++ nodedb/tests/inproc/cases/mod.rs | 3 + .../inproc/cases/surrogate_round_trip.rs | 21 +- nodedb/tests/wire/cases/array_copy_from.rs | 456 +++++++++++++ .../tests/wire/cases/bitemporal_array_sql.rs | 38 +- .../tests/wire/cases/engine_surface_array.rs | 118 ++++ nodedb/tests/wire/cases/mod.rs | 2 + .../tests/wire/cases/native_cluster_array.rs | 93 ++- 149 files changed, 8368 insertions(+), 1120 deletions(-) create mode 100644 nodedb-client-tests/tests/array_load_measure.rs create mode 100644 nodedb-client-tests/tests/array_put_cells_round_trip.rs create mode 100644 nodedb-client/src/native/client/array.rs create mode 100644 nodedb-sql/src/ddl_ast/parse/copy_from_stdin.rs create mode 100644 nodedb-sql/src/engine_rules/bulk_load.rs create mode 100644 nodedb-sql/src/planner/array_fn/pushdown.rs rename nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs => nodedb-sql/src/planner/cte/inline.rs (95%) create mode 100644 nodedb-sql/src/planner/dml/bulk_load.rs create mode 100644 nodedb-sql/src/planner/select/array_from.rs create mode 100644 nodedb-types/src/protocol/array_batch.rs create mode 100644 nodedb/src/control/planner/context/query/planning_built.rs create mode 100644 nodedb/src/control/planner/sql_plan_convert/array_convert/tile_route.rs create mode 100644 nodedb/src/control/server/exchange/resolve/array_read_rows.rs create mode 100644 nodedb/src/control/server/native/dispatch/array_batch.rs create mode 100644 nodedb/src/control/server/pgwire/handler/array_copy.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/chunker.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/decode.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/dispatch.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/loader.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/mod.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/records.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/router.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/source.rs create mode 100644 nodedb/src/control/server/shared/ddl/neutral/collection/dml/parse/plan_statement.rs create mode 100644 nodedb/src/data/executor/dispatch/array/slice_row.rs create mode 100644 nodedb/src/engine/array/store/catalog/resolve.rs create mode 100644 nodedb/tests/inproc/cases/array_auto_flush_visibility.rs create mode 100644 nodedb/tests/inproc/cases/array_scan_cost.rs create mode 100644 nodedb/tests/inproc/cases/array_sql_relation.rs create mode 100644 nodedb/tests/wire/cases/array_copy_from.rs 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/docs/array.md b/docs/array.md index 77ff4be44..6477e82f2 100644 --- a/docs/array.md +++ b/docs/array.md @@ -137,6 +137,50 @@ SELECT * FROM ARRAY_SLICE( | `attrs` | No | ARRAY[STRING] | Attributes to project. `NULL` = all attributes | | `limit` | No | INT64 | Max cells returned. `NULL` = 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 | +| ------------- | ------------------ | +| `INT64` | `bigint` | +| `FLOAT64` | `double precision` | +| `TIMESTAMP` | `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: diff --git a/docs/bitemporal.md b/docs/bitemporal.md index e6ab2414c..ace37704f 100644 --- a/docs/bitemporal.md +++ b/docs/bitemporal.md @@ -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, and on `ARRAY_SLICE`. 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 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/query-language.md b/docs/query-language.md index 570c5db60..ca06f4724 100644 --- a/docs/query-language.md +++ b/docs/query-language.md @@ -1153,6 +1153,24 @@ COPY users FROM STDIN; COPY users FROM STDIN 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 +``` + +- 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`. +- 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 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/tile_decode.rs b/nodedb-array/src/codec/tile_decode.rs index 8ba84ecb4..6635fae27 100644 --- a/nodedb-array/src/codec/tile_decode.rs +++ b/nodedb-array/src/codec/tile_decode.rs @@ -16,8 +16,10 @@ use crate::codec::limits::{ }; 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 +51,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 +94,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 +107,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() { @@ -174,12 +191,6 @@ fn decode_structural(body: &[u8]) -> ArrayResult { 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 +201,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 +222,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 +438,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..011d13a99 100644 --- a/nodedb-array/src/codec/tile_encode.rs +++ b/nodedb-array/src/codec/tile_encode.rs @@ -212,6 +212,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 +288,9 @@ 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 derived result tile holds no identity, so it is never written. diff --git a/nodedb-array/src/query/mod.rs b/nodedb-array/src/query/mod.rs index 68608f76f..9722c3175 100644 --- a/nodedb-array/src/query/mod.rs +++ b/nodedb-array/src/query/mod.rs @@ -23,4 +23,6 @@ pub use elementwise::{BinaryOp, elementwise}; pub use project::{Projection, project_sparse}; pub use rechunk::rechunk_sparse; pub use retention::{RetentionMergeResult, merge_for_retention}; -pub use slice::{DimRange, Slice, slice_sparse, tile_overlaps_slice}; +pub use slice::{ + DimRange, Slice, slice_sparse, slice_sparse_owned, tile_overlaps_slice, tile_within_slice, +}; diff --git a/nodedb-array/src/query/slice.rs b/nodedb-array/src/query/slice.rs index 7ee98c72b..4fe0ebddb 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. @@ -278,4 +316,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/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/tile/mbr.rs b/nodedb-array/src/tile/mbr.rs index c23429db1..6d29d1a35 100644 --- a/nodedb-array/src/tile/mbr.rs +++ b/nodedb-array/src/tile/mbr.rs @@ -94,7 +94,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 +124,25 @@ 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(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()); } } - self.nnz += 1; } pub fn build(self) -> TileMBR { 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-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/src/distributed_array/merge.rs b/nodedb-cluster/src/distributed_array/merge.rs index 9d11396e7..0ecf6e274 100644 --- a/nodedb-cluster/src/distributed_array/merge.rs +++ b/nodedb-cluster/src/distributed_array/merge.rs @@ -156,7 +156,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 diff --git a/nodedb-physical/src/physical_plan/array.rs b/nodedb-physical/src/physical_plan/array.rs index b0e3606dd..f07c8a4f4 100644 --- a/nodedb-physical/src/physical_plan/array.rs +++ b/nodedb-physical/src/physical_plan/array.rs @@ -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 { 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/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/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..9389e1e59 100644 --- a/nodedb-sql/src/engine_rules/timeseries.rs +++ b/nodedb-sql/src/engine_rules/timeseries.rs @@ -115,6 +115,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 { diff --git a/nodedb-sql/src/planner/array_fn/mod.rs b/nodedb-sql/src/planner/array_fn/mod.rs index c815f2e8b..30dca86a7 100644 --- a/nodedb-sql/src/planner/array_fn/mod.rs +++ b/nodedb-sql/src/planner/array_fn/mod.rs @@ -16,7 +16,9 @@ mod helpers; mod maint_fn; +mod pushdown; mod table_fn; 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..8a9048325 100644 --- a/nodedb-sql/src/planner/array_fn/table_fn.rs +++ b/nodedb-sql/src/planner/array_fn/table_fn.rs @@ -18,22 +18,57 @@ use crate::temporal::TemporalScope; use crate::types::{SqlCatalog, SqlPlan}; use crate::types_array::{ArrayBinaryOpAst, ArrayReducerAst, ArraySliceAst, NamedDimRange}; +/// 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. +#[derive(Debug, Clone)] +pub enum ArrayTableFn { + Slice(ArraySlicePlan), + Project(ArrayProjectPlan), + 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) => SqlPlan::ArraySlice(plan), + ArrayTableFn::Project(plan) => SqlPlan::ArrayProject(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 +93,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, + )?))), + "array_agg" => Ok(Some(ArrayTableFn::Agg(plan_agg( + &arg_exprs, catalog, temporal, + )?))), + "array_elementwise" => Ok(Some(ArrayTableFn::Elementwise(plan_elementwise( + &arg_exprs, catalog, + )?))), _ => Ok(None), } } @@ -70,7 +113,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!( @@ -146,16 +189,16 @@ fn plan_slice( 0 }; - Ok(SqlPlan::ArraySlice(ArraySlicePlan { + Ok(ArraySlicePlan { name, slice: ArraySliceAst { dim_ranges }, attr_projection, limit, temporal, - })) + }) } -fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { +fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { if args.len() != 2 { return Err(SqlError::Unsupported { detail: format!( @@ -183,17 +226,17 @@ fn plan_project(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result }); } } - Ok(SqlPlan::ArrayProject(ArrayProjectPlan { + Ok(ArrayProjectPlan { name, attr_projection, - })) + }) } fn plan_agg( args: &[ast::Expr], catalog: &dyn SqlCatalog, temporal: TemporalScope, -) -> Result { +) -> Result { if args.len() < 3 || args.len() > 4 { return Err(SqlError::Unsupported { detail: format!( @@ -233,16 +276,16 @@ fn plan_agg( None }; - Ok(SqlPlan::ArrayAgg(ArrayAggPlan { + Ok(ArrayAggPlan { name, attr, reducer, group_by_dim, temporal, - })) + }) } -fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { +fn plan_elementwise(args: &[ast::Expr], catalog: &dyn SqlCatalog) -> Result { if args.len() != 4 { return Err(SqlError::Unsupported { detail: format!( @@ -280,12 +323,12 @@ 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,7 +468,7 @@ mod tests { #[test] fn project_happy() { let p = plan_one("SELECT * FROM ARRAY_PROJECT('g', ['qual', 'variant'])").unwrap(); - match p { + match relation_body(p) { SqlPlan::ArrayProject(ArrayProjectPlan { name, attr_projection, diff --git a/nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs b/nodedb-sql/src/planner/cte/inline.rs similarity index 95% rename from nodedb/src/control/planner/sql_plan_convert/expr/inline_cte.rs rename to nodedb-sql/src/planner/cte/inline.rs index 4cc6ee722..41d468faf 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 { @@ -416,7 +416,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 +425,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 +448,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -528,13 +528,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 +603,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 +629,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 +642,7 @@ mod tests { offset: 0, distinct: false, window_functions, - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -735,7 +735,7 @@ mod tests { offset: 0, distinct: false, window_functions: Vec::new(), - temporal: nodedb_sql::TemporalScope::default(), + temporal: crate::TemporalScope::default(), } } @@ -769,7 +769,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 +798,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 +809,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 +842,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..c063c8bf2 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; 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/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/select/array_from.rs b/nodedb-sql/src/planner/select/array_from.rs new file mode 100644 index 000000000..e1178a4d8 --- /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_relation_columns, 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(mut slice) => { + let view = array_view(catalog, &slice.name)?; + let read = RelationRead::new(&ctx, "ARRAY_SLICE", &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); + read.into_cte(SqlPlan::ArraySlice(slice), outer) + } + ArrayTableFn::Project(project) => { + let view = array_view(catalog, &project.name)?; + let read = RelationRead::new(&ctx, "ARRAY_PROJECT", &view.name)?; + let columns = array_relation_columns(&view, &project.attr_projection); + let outer = read.plan_outer(&ctx, columns)?; + read.into_cte(SqlPlan::ArrayProject(project), outer) + } + call @ (ArrayTableFn::Agg(_) | ArrayTableFn::Elementwise(_)) => { + plan_non_relational(&ctx, temporal, call)? + } + }; + Ok(Some(planned)) +} + +/// 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..4894058cf 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,74 @@ 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), + } + } + + 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), } - self.inner.get_collection(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/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/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..7ec3b56c1 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`. diff --git a/nodedb-sql/src/planner/select/where_search.rs b/nodedb-sql/src/planner/select/where_search.rs index 51f9cf682..19574b0b1 100644 --- a/nodedb-sql/src/planner/select/where_search.rs +++ b/nodedb-sql/src/planner/select/where_search.rs @@ -390,7 +390,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("; ")), diff --git a/nodedb-sql/src/resolver/array_tvf.rs b/nodedb-sql/src/resolver/array_tvf.rs index 7d74b3e60..b566f3c32 100644 --- a/nodedb-sql/src/resolver/array_tvf.rs +++ b/nodedb-sql/src/resolver/array_tvf.rs @@ -3,23 +3,31 @@ //! 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, + ArrayCatalogView, ArraySlicePlan, CollectionInfo, ColumnInfo, EngineType, SqlCatalog, + SqlDataType, }; use crate::types_array::{ArrayAttrType, ArrayDimType}; +/// 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"; + /// 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 +65,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 +93,97 @@ 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(column(&a.name, attr_type_to_sql(a.dtype), a.nullable)); + } + } else { + for name in attrs { + if let Some(a) = view.attrs.iter().find(|a| &a.name == name) { + cols.push(column(&a.name, attr_type_to_sql(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 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) => { + Ok(array_slice_columns(&view_of(catalog, &slice.name)?, slice)) + } + ArrayTableFn::Project(project) => Ok(array_relation_columns( + &view_of(catalog, &project.name)?, + &project.attr_projection, + )), + ArrayTableFn::Agg(agg) => { + let view = view_of(catalog, &agg.name)?; + 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)); + } + cols.push(column("result", SqlDataType::Float64, true)); + Ok(cols) + } + 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, 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-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/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/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/cluster/array_cluster_exec/executor.rs b/nodedb/src/control/cluster/array_cluster_exec/executor.rs index ea3680a59..8d6fc2fbe 100644 --- a/nodedb/src/control/cluster/array_cluster_exec/executor.rs +++ b/nodedb/src/control/cluster/array_cluster_exec/executor.rs @@ -347,6 +347,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) } 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..ec42d2169 --- /dev/null +++ b/nodedb/src/control/planner/context/query/planning_built.rs @@ -0,0 +1,142 @@ +// 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()); + + 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); + Ok(()) + } +} 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..ce3d84569 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 @@ -5,6 +5,7 @@ use nodedb_array::schema::ArraySchema; use nodedb_array::tile::cell_tile_prefix; use nodedb_array::types::ArrayId; +use nodedb_array::types::coord::value::CoordValue; use nodedb_sql::types_array::{ArrayCoordLiteral, ArrayInsertRow}; use crate::bridge::envelope::PhysicalPlan; @@ -47,18 +48,50 @@ 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>, +} + +/// 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 }) +} + +/// Row `index`'s coerced coordinate and identity key: the bind step's when +/// it kept them, else coerced and encoded here. +fn cell_key<'k>( + kept: Option<&'k ArrayCellKeys>, + index: usize, + row: &ArrayInsertRow, + schema: &ArraySchema, +) -> crate::Result<(Vec, std::borrow::Cow<'k, [u8]>)> { + if let Some(keys) = kept + && let (Some(coord), Some(pk)) = (keys.coords.get(index), keys.pks.get(index)) + { + return Ok((coord.clone(), std::borrow::Cow::Borrowed(pk.as_slice()))); + } + let coord = coerce_coords(&row.coords, schema)?; + let pk = coord_pk(&coord)?; + Ok((coord, std::borrow::Cow::Owned(pk))) } pub(in super::super) fn convert_insert_array( @@ -70,20 +103,17 @@ 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)?; + for (index, row) in rows.iter().enumerate() { + let (coord, pk_bytes) = cell_key(kept, index, row, &schema)?; 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)?; + let surrogate = ctx.surrogate_for_pk(key, &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 = @@ -131,15 +161,10 @@ 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, pk_bytes) = cell_key(kept, index, row, &schema)?; 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)?; + let surrogate = ctx.surrogate_for_pk(key, &pk_bytes)?; cells.push(ArrayPutCell { coord, attrs, @@ -284,8 +309,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/mod.rs b/nodedb/src/control/planner/sql_plan_convert/array_convert/mod.rs index 08333b733..a28fd6451 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 @@ -3,6 +3,10 @@ mod ddl; mod dml; mod helpers; +mod tile_route; 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..7ce008786 --- /dev/null +++ b/nodedb/src/control/planner/sql_plan_convert/array_convert/tile_route.rs @@ -0,0 +1,27 @@ +// 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::cell_tile_prefix; +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. +pub(crate) fn cell_vshard( + coords: &[ArrayCoordLiteral], + schema: &ArraySchema, + prefix_bits: u8, +) -> crate::Result { + let coord = coerce_coords(coords, schema)?; + let tile_key = cell_tile_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/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/schema.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/build/schema.rs index cd2858c1e..87a62efde 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 @@ -14,14 +14,16 @@ use nodedb_sql::SqlError; use nodedb_sql::catalog::SqlCatalog; 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, + ArrayProjectPlan, 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, +}; use crate::control::server::response_shape::types::DdlColType; use super::super::columns::{ @@ -263,24 +265,24 @@ 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, + )) + } + SqlPlan::ArrayProject(ArrayProjectPlan { + name, + attr_projection, + }) => { + let view = array_view(catalog, name)?; + array_columns_schema(&nodedb_sql::resolver::array_tvf::array_relation_columns( + &view, + attr_projection, + )) + } // 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 @@ -328,6 +330,37 @@ pub fn build_output_schema( }) } +/// 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. +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(&column.data_type), + }) + .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/cte_types.rs b/nodedb/src/control/planner/sql_plan_convert/output_schema/cte_types.rs index a596dea24..697332b01 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. @@ -93,6 +97,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/set_ops.rs b/nodedb/src/control/planner/sql_plan_convert/set_ops.rs index 14e9eb07e..5cd1e57b5 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) } 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..ddd55976a 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 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..8d62fc61a 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 @@ -13,19 +13,36 @@ use std::sync::Mutex; use nodedb_types::{CollectionKey, DatabaseId, Surrogate}; +use super::super::array_convert::ArrayCellKeys; + /// One collection: its database and bare catalog name. 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 = HashMap, 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: HashMap>, + /// 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 +123,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 = HashMap::new(); + 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 +259,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 +331,34 @@ 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()], + }, + ); + + let kept = prefetched.array_keys(scene, &rows).expect("kept"); + assert_eq!(kept.pks[1], b"b".to_vec()); + 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..7ac348029 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 @@ -8,8 +8,9 @@ use std::sync::Arc; use nodedb_sql::types::SqlPlan; use nodedb_types::CollectionKey; +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,34 @@ 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()) - } else { - (Vec::new(), pks) - }; - let request = KeyRequest { - key: ctx.collection_key(batch.collection), - binds: &binds, - lookups: &lookups, - fresh: batch.fresh, - }; - resolver.resolve(request, &mut prefetched).await?; + let KeyBatch { + collection, + pks, + binds: binding, + fresh, + array_rows, + } = batch; + let key = ctx.collection_key(collection); + { + 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, &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 }); + } } Ok(prefetched) } diff --git a/nodedb/src/control/server/exchange/cluster_leaf.rs b/nodedb/src/control/server/exchange/cluster_leaf.rs index 074d30195..ab2781338 100644 --- a/nodedb/src/control/server/exchange/cluster_leaf.rs +++ b/nodedb/src/control/server/exchange/cluster_leaf.rs @@ -25,7 +25,7 @@ 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::data::executor::response_codec::ArraySliceResponse; use crate::types::{Lsn, TxnId}; use super::gather::GatherOutcome; @@ -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(outcome_of(rows)); } let mut local = plan; inline_cluster_arrays(state, &mut local, scope).await?; @@ -62,6 +62,11 @@ 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, and an aggregate row is `{group, result}` or +/// `{result}`. 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, @@ -139,7 +144,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() { @@ -236,17 +241,30 @@ mod tests { } #[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"), - ); - assert_eq!(extract_msgpack_elements(&flat), rows); + let out = cluster_array_rows(&slice_op(), envelope).expect("slice rows"); + assert_eq!(extract_msgpack_elements(&out), 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), vec![named]); } #[tokio::test] @@ -271,11 +289,9 @@ 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), rows); } } 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..3ed597f8b --- /dev/null +++ b/nodedb/src/control/server/exchange/resolve/array_read_rows.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: BUSL-1.1 + +//! Relational rows of a coordinator-gathered array read. +//! +//! An array read emits rows that are already relational: a slice or project +//! cell holds one field per dim and attr, keyed by name, an aggregate row is +//! `{group, result}`, an elementwise row is one cell. The storage flatten +//! that unwraps `{id, data}` document rows must not run over them: it would +//! read an attr named `data` as a row envelope. +//! +//! A local `ArrayOp::Slice` answers with one `ArraySliceResponse` envelope +//! per core it ran on, not with a bare row array, so its rows are read out +//! of the envelopes. A cluster slice reaches the coordinator as the bare +//! rows its array coordinator merged (`cluster_leaf`). + +use nodedb_physical::physical_plan::{ArrayOp, ClusterArrayOp, PhysicalPlan}; + +use crate::control::server::payload_merge::{encode_msgpack_array, extract_msgpack_elements}; +use crate::control::server::response_shape::compose::array_slice::TRUNCATED_BEFORE_HORIZON_NOTICE; +use crate::control::server::shared::session::statement_notice; +use crate::data::executor::response_codec::ArraySliceResponse; + +/// How a gathered array read's payload holds its rows. +#[derive(Clone, Copy)] +pub(super) enum ArrayReadShape { + /// One `ArraySliceResponse` envelope per core. + SliceEnvelopes, + /// A bare array of relational rows. + Rows, +} + +impl ArrayReadShape { + /// The shape `plan`'s gathered payload takes, or `None` when `plan` is + /// not an array read. + pub(super) fn of(plan: &PhysicalPlan) -> Option { + if matches!(plan, PhysicalPlan::Array(ArrayOp::Slice { .. })) { + return Some(Self::SliceEnvelopes); + } + let bare_rows = matches!( + plan, + PhysicalPlan::Array( + ArrayOp::Project { .. } | ArrayOp::Aggregate { .. } | ArrayOp::Elementwise { .. }, + ) | PhysicalPlan::ClusterArray( + ClusterArrayOp::Slice { .. } | ClusterArrayOp::Agg { .. } + ) + ); + bare_rows.then_some(Self::Rows) + } + + /// The relational rows of `merged`, the gathered payload, as one msgpack + /// array. + pub(super) fn rows(self, merged: Vec) -> crate::Result> { + match self { + Self::SliceEnvelopes => local_slice_rows(&merged), + Self::Rows => Ok(merged), + } + } +} + +/// The rows of every slice envelope in `merged`, as one msgpack array. +/// +/// `merged` holds one element per envelope. An envelope cut below its +/// retention horizon raises the same client notice a slice shaped on its +/// own reports, through the statement notice slot. An element that is not +/// an envelope is an error: it is not slice rows, and reading it as rows +/// would answer the query from the wrong data. +fn local_slice_rows(merged: &[u8]) -> crate::Result> { + let mut rows = Vec::new(); + let mut truncated = false; + for envelope in extract_msgpack_elements(merged) { + let response: ArraySliceResponse = + zerompk::from_msgpack(&envelope).map_err(|e| crate::Error::Codec { + detail: format!("array slice response decode: {e}"), + })?; + truncated |= response.truncated_before_horizon; + rows.extend(extract_msgpack_elements(&response.rows_msgpack)); + } + if truncated { + statement_notice::raise(TRUNCATED_BEFORE_HORIZON_NOTICE.to_string()); + } + Ok(encode_msgpack_array(&rows)) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 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] + } + + fn envelope(rows: &[Vec], truncated: bool) -> Vec { + zerompk::to_msgpack_vec(&ArraySliceResponse { + rows_msgpack: encode_msgpack_array(rows), + truncated_before_horizon: truncated, + }) + .expect("encode slice envelope") + } + + #[test] + fn rows_of_every_envelope_concatenate() { + let merged = encode_msgpack_array(&[ + envelope(&[row(1), row(2)], false), + envelope(&[row(3)], false), + ]); + let rows = ArrayReadShape::SliceEnvelopes + .rows(merged) + .expect("slice rows"); + assert_eq!( + extract_msgpack_elements(&rows), + vec![row(1), row(2), row(3)] + ); + } + + #[test] + fn a_bare_row_is_not_an_envelope() { + let merged = encode_msgpack_array(&[row(1)]); + assert!(ArrayReadShape::SliceEnvelopes.rows(merged).is_err()); + } + + #[test] + fn bare_rows_pass_through() { + let merged = encode_msgpack_array(&[row(1), row(2)]); + let rows = ArrayReadShape::Rows.rows(merged.clone()).expect("rows"); + assert_eq!(rows, merged); + } + + #[tokio::test] + async fn a_truncated_envelope_raises_the_horizon_notice() { + use crate::control::server::shared::session::conn_scope; + conn_scope::scoped(async { + let merged = encode_msgpack_array(&[envelope(&[row(1)], true)]); + ArrayReadShape::SliceEnvelopes + .rows(merged) + .expect("slice rows"); + assert_eq!( + statement_notice::take(), + vec![TRUNCATED_BEFORE_HORIZON_NOTICE.to_string()] + ); + }) + .await; + } +} 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..734df89bd 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,6 +8,8 @@ use nodedb_physical::physical_plan::{ plan_contains_cluster_partitioned_leaf, }; +use super::super::array_read_rows::ArrayReadShape; + 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, @@ -52,6 +54,9 @@ enum HitShape { /// Hybrid (RRF) fusion hits: `{doc_id: , , /// ...}` with no body. Resolve `doc_id` to the user PK as `id`. Hybrid, + /// An array read, whose rows are already relational (see + /// [`ArrayReadShape`]). + Array(ArrayReadShape), /// Flat storage rows (`{id, data}` document / text, columnar, spatial) or /// computed rows — already fully columned after the storage flatten. None, @@ -59,6 +64,9 @@ enum HitShape { /// Classify a resolved `PostProcess` child by the row shape its engine emits. fn classify_hit_shape(plan: &PhysicalPlan) -> HitShape { + if let Some(shape) = ArrayReadShape::of(plan) { + return HitShape::Array(shape); + } match plan { PhysicalPlan::Vector( VectorOp::Search { .. } @@ -265,6 +273,7 @@ pub(super) async fn materialize_child_rows( let surrogate = parse_surrogate_hex(key)?; resolve_surrogate_pk(state, database_id, tenant_id, &coll, surrogate) }), + HitShape::Array(shape) => shape.rows(merged)?, HitShape::None => flatten_to_relational_rows(&merged), }; Ok(ChildRows::Rows { rows, sources }) diff --git a/nodedb/src/control/server/exchange/resolve/join_input.rs b/nodedb/src/control/server/exchange/resolve/join_input.rs index 56f6ab7e3..72e67da5d 100644 --- a/nodedb/src/control/server/exchange/resolve/join_input.rs +++ b/nodedb/src/control/server/exchange/resolve/join_input.rs @@ -12,6 +12,7 @@ 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 super::array_read_rows::ArrayReadShape; use super::capture::DistributedReadCapture; use super::exchange::provider_scan_of_rows; use crate::control::server::shared::authorization::gathered_sources; @@ -60,10 +61,10 @@ pub(super) async fn resolve_join_input( // only the groups this node replicates, and confirms each leg as // `scope` requires. let sources = gathered_sources(&child, database_id)?; + let array_shape = ArrayReadShape::of(&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_shape, outcome.merged_array)?; + 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_shape = ArrayReadShape::of(&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 @@ -128,8 +130,8 @@ pub(super) async fn resolve_join_input( } 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_shape, merged)?; + Ok(Some(Box::new(provider_scan_of_rows(rows, sources)))) } // Already resolved (ProviderScan{None, ..} or any other plan): @@ -222,3 +224,12 @@ pub(super) async fn gather_join_build_side( sources, )))) } + +/// The relational rows of a gathered join input: an array read's rows as +/// it emits them, any other plan's rows through the storage flatten. +fn relational_rows(array_shape: Option, merged: Vec) -> crate::Result> { + match array_shape { + Some(shape) => shape.rows(merged), + None => Ok(flatten_to_relational_rows(&merged)), + } +} diff --git a/nodedb/src/control/server/exchange/resolve/mod.rs b/nodedb/src/control/server/exchange/resolve/mod.rs index 890be7b89..3824dd4a1 100644 --- a/nodedb/src/control/server/exchange/resolve/mod.rs +++ b/nodedb/src/control/server/exchange/resolve/mod.rs @@ -30,6 +30,7 @@ //! finalized rows into a `Resolved::Gathered` response (distributed GROUP BY). //! - No Exchange / no empty ProviderScan → `Resolved::Plan` unchanged. +mod array_read_rows; mod cagg_rows; mod capture; pub mod exchange; 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/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/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/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/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/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/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/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..35d7dd1c0 --- /dev/null +++ b/nodedb/src/control/server/shared/ddl/neutral/collection/copy_from/array_load/decode.rs @@ -0,0 +1,431 @@ +// 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_sql::catalog::ArrayCatalogView; +use nodedb_sql::types_array::{ + ArrayAttrLiteral, ArrayAttrType, ArrayCoordLiteral, ArrayDimType, ArrayDomainBound, + ArrayInsertRow, +}; +use nodedb_types::Value; + +/// 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) => parse_attr(text, attr.dtype) + .ok_or_else(|| invalid(&attr.name, attr_type_name(attr.dtype), text, row))?, + }; + 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), ArrayAttrType::Int64) => ArrayAttrLiteral::Int64(*n), + (Value::Integer(n), ArrayAttrType::Float64) => ArrayAttrLiteral::Float64(*n as f64), + (Value::Float(f), ArrayAttrType::Float64) => 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, attr_type_name(dtype), other, 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 => text.trim().parse().ok().map(ArrayAttrLiteral::Int64), + ArrayAttrType::Float64 => text.trim().parse().ok().map(ArrayAttrLiteral::Float64), + ArrayAttrType::String => Some(ArrayAttrLiteral::String(text.to_string())), + ArrayAttrType::Bytes => parse_bytea_hex(text).map(ArrayAttrLiteral::Bytes), + } +} + +/// 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 attr_type_name(dtype: ArrayAttrType) -> &'static str { + match dtype { + ArrayAttrType::Int64 => "INT64", + ArrayAttrType::Float64 => "FLOAT64", + ArrayAttrType::String => "STRING", + ArrayAttrType::Bytes => "BYTES", + } +} + +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); + } + + #[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: